FFN 是 key-value memory

Transformer 里 FFN 层占了约三分之二的参数，却长期被当作「升维-非线性-降维」的黑盒。Geva 等人 2021 年的工作（Transformer Feed-Forward Layers Are Key-Value Memories）给出一个至今最有解释力的视角：FFN 就是一个显式的 key-value memory——第一层权重是 key，检测输入模式；第二层权重是 value，向残差流写入知识。这个视角是理解混合专家模型 MoE「知识容量假说」、模型知识编辑等方向的地基。

从矩阵乘法到记忆读写

标准 FFN 是两层 MLP：

$FFN (x) = W_{2} f (W_{1} x)$

其中 $f$ 是 GELU/ReLU， $W_{1} \in R^{d_{ff} \times d}$ ， $W_{2} \in R^{d \times d_{ff}}$ 。换个记号，把 $W_{1}$ 的每一行记作 key $k_{i}$ 、 $W_{2}$ 的每一列记作 value $v_{i}$ ，上式改写为：

$FFN (x) = \sum_{i = 1}^{d_{ff}} f (k_{i}^{⊤} x) \cdot v_{i}$

这个形式与 attention 的相似度加权求和同构，只是「检索」发生在输入与 key 的内积上：

key 匹配： $k_{i}^{⊤} x$ 度量输入 $x$ 与模式 $k_{i}$ 的匹配度，过非线性后大部分项归零、少数显著存活
value 写入：存活 key 对应的 $v_{i}$ 加权求和，直接加进残差流，改变当前 token 的表示

也就是说，FFN 的前向就是一次内容寻址的内存读取（content-addressed memory read）： $d_{ff}$ 个记忆槽，每个槽由（模式检测器，写入内容）组成。d_model=4096、d_ff=12288 的 Llama-7B FFN 就有 12288 个槽，每层两份参数。

key：少数派模式检测器

论文对 GPT-2 的 key 激活做了统计，两个核心发现：

激活天然稀疏。 对大量输入统计每个 key 的激活强度，分布呈重尾——绝大多数 key 对绝大多数输入接近零，每个输入只强激活一小撮 key。ReLU/GELU 的零区只是表层原因，更本质的是训练把 key 推向了「专精少数模式」的解。

key 检测的模式浅层但可解释。 取每个 key 激活最强的输入集合，人眼看大多能归纳出共同点：某种词法后缀、标点组合、数字串、特定主题的词。论文称这种性质为 few-shot pattern detection——每个 key 只被训练数据中一小撮输入激活，因此可以腾出容量编码罕见模式（小语种拼写、生僻搭配），不会被高频模式淹没。这正是「记忆」而非「通用变换」的特征：一个分布平滑的变换器记不住长尾。

value：向输出分布写入残差

value 的分析更直接：把单个 $v_{i}$ 加到残差流上，经过后续层到 unembedding，观察输出词分布的变化。

每个 value 促进一个语义相关的 token 簇，而非单个词——把 value 直接投影到词表，top tokens 呈主题聚类
中间层的 value 编码的是「中间贡献」：单看某一层，value 增强的 token 与最终预测还有距离；但随着层数加深，value 促进的分布逐步向最终输出对齐——每层的写入都以前面所有层的写入为条件，知识是逐层聚合上去的
末层的 value 已经几乎直接对应输出 token，最后几层 FFN 承担最多的「事实检索」职能

这与 attention 形成清晰分工：attention 在位置间搬运已有信息，FFN 向表示写入新信息。一句话——attention 是路由，FFN 是记忆。

这个视角解释了什么

MoE 的知识容量假说。 如果 FFN 是记忆，MoE 就是把记忆槽乘以 $N$ ：总容量线性涨，但每 token 只读写被路由到的 k 份，访问成本不变。混合专家模型 MoE 里「同样训练 FLOPs 下 MoE 的 loss 更低」的实证，在 memory 视角下有了机制解释。

知识编辑与事实定位。 Geva 等人 2023 年的后续工作（ROME）把「北京-中国」这类事实三元组的召回定位到最后几个 FFN 层：subject 最后一个 token 位置的 FFN key 匹配「关系模式」，对应的 value 注入 object。既然知识有明确的 key-value 地址，就能按地址改写——模型编辑（model editing）整个方向由此展开。

剪枝与冗余。 记忆槽视角预言：删掉低频激活的 key 对常见输入影响极小。FFN 结构化剪枝（按行/列删）确实比 attention 剪枝耐性好得多，与该视角一致。

验证实验

用 GPT-2 复现论文两个核心观察：key 激活的稀疏性，以及 value 投影到词表的语义聚类。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
 
model = AutoModelForCausalLM.from_pretrained("gpt2", torch_dtype=torch.float32)
tok = AutoTokenizer.from_pretrained("gpt2")
layer = model.transformer.h[8].mlp                    # 任取一层 FFN
W1, W2 = layer.c_fc.weight, layer.c_proj.weight        # (d_ff, d) 与 (d, d_ff)
 
text = "The capital of France is"
x = model.transformer(tok(text, return_tensors="pt").input_ids).last_hidden_state[0, -1]
 
# 观察 1：key 激活稀疏性 —— 该 token 只激活极少数 key
act = torch.nn.functional.gelu(W1 @ x)                 # (d_ff,)
ratio = (act > act.max() * 0.1).float().mean()
print(f"激活超过峰值 10% 的 key 占比：{ratio:.2%}")     # GPT-2 该设置下通常 < 5%
 
# 观察 2：value 的词表分布 —— 单个 value 直接投影到词表，看语义聚类
top_k = act.topk(5).indices
unembed = model.lm_head.weight                          # (vocab, d)
for i in top_k:
    logits = unembed @ (W2[:, i] * act[i])              # 该 value 引导的输出分布
    words = [tok.decode(j) for j in logits.topk(5).indices]
    print(f"key {i}: 激活 {act[i]:.2f} -> {words}")     # 同一 key 的 top 词常呈主题聚类

局限

分析基于 GPT-2 级别的小模型与相关性统计，「value 写入 = 知识存储」是功能解释，不构成因果证明（ROME 式因果干预补了一部分）
memory 视角解释事实性知识最有效，对推理、组合泛化等能力覆盖有限
d_ff 的「槽」是连续向量而非离散条目，说「一个 key 存一条知识」是过度解读——粒度在 key 簇与层间组合层面

但作为一级近似，key-value memory 已是 FFN 研究的默认镜头：MoE 把它做成分布式柜子，模型编辑按它寻址，剪枝按它评估冗余。

知识花园

探索