FFN 与 SwiGLU:Attention 之后,每个 Token 还要做什么?

从逐位置加工、门控选择,到“知识存储”这个比喻的边界

目录

  Attention 已经让 Token 互相交换过信息了,为什么 Transformer Block 里还要紧跟一个 FFN?

  可以先用一句话区分它们:

Attention 负责“从其他 Token 那里取回什么”,FFN 负责“拿到这些信息以后,当前 Token 要把它加工成什么”。

FFN 与 SwiGLU 的工作流程

Attention 先交流,FFN 再逐 Token 加工,点击查看原始 SVG

Attention 已经交流过了,为什么还需要 FFN

  假设输入是“我喜欢苹果”。经过 Attention 后,“苹果”这个位置已经能吸收“我”和“喜欢”的信息。此时它不再只是最初的词向量,而是带着上下文的表示。

  但“收集信息”和“加工信息”是两件事。Attention 主要把不同位置的信息按权重混合;FFN 则对混合后的表示做非线性变换,识别并组合出更有用的特征。

  更重要的是,FFN 是逐 Token工作的:

我     ──同一个 FFN──> 新的“我”表示
喜欢   ──同一个 FFN──> 新的“喜欢”表示
苹果   ──同一个 FFN──> 新的“苹果”表示

  三个位置使用同一套 FFN 参数,但每个位置独立计算。FFN 本身不会让“我”再去读取“苹果”;跨 Token 的交流已经由 Attention 完成。

FFN 究竟做了什么

  经典 Transformer 在每一层中使用逐位置前馈网络:所有 Token 共享同一套参数,但各位置独立计算。[1] 它可以写成:

\[\operatorname{FFN}(x)=W_2\,\sigma(W_1x+b_1)+b_2\]

  它包含三个动作:

  1. 升维:把 $d_{model}$ 维的输入投影到更宽的 $d_{ff}$。
  2. 激活:用 ReLU、GELU 或 SiLU 加入非线性。
  3. 降维:把结果投影回 $d_{model}$,以便接回残差流。

  如果 $d_{model}=4096$,中间宽度可能达到一万多维。可以把这片更宽的空间理解成一排“特征探测器”:有些响应语法,有些响应实体或语义模式。输入不同,亮起的组合也不同。

为什么中间必须有非线性

  如果拿掉激活函数,连续两个线性层仍然只是一个线性层:

\[W_2(W_1x)=(W_2W_1)x\]

  无论中间扩得多宽,都可以合并成一个矩阵,模型并没有获得新的表达能力。

  激活函数改变了这件事。它让不同输入走出不同的响应模式:某些中间维度被压低,某些被保留。于是 FFN 不只是“换一个坐标系”,而能根据输入选择和组合特征。

从 GLU 到 SwiGLU:多加一扇门

  SwiGLU 是 Shazeer 对 Transformer FFN 中 GLU 变体进行系统比较后提出并验证的结构之一。[2] 它不再只计算一条升维支路,而是把输入送进两条支路:

\[\operatorname{SwiGLU}(x)=W_{down}\left(\operatorname{SiLU}(W_{gate}x)\odot(W_{up}x)\right)\]

  两条支路分工如下:

  • $W_{up}x$ 产生候选内容:这里有什么特征?
  • $\operatorname{SiLU}(W_{gate}x)$ 产生门控信号:这些特征该通过多少?
  • $\odot$ 做逐元素相乘,把门应用到候选内容上。
  • $W_{down}$ 把筛选后的宽向量投影回 $d_{model}$。

  这里的“门”并不是只能开或关。SiLU 输出的是连续数值,因此某个特征可以被加强、减弱,也可能改变符号。把它理解成一排可调节的旋钮,比“0 或 1 的开关”更准确。

动手试试:不同 Token 会打开不同的门

喜欢
上下文化表示 x
→
SwiGLU 的门控响应
动作关系
主体线索
实体类别
→
Wdown
回到 dmodel

  这只是为了帮助理解的示意,不代表真实模型中存在三个写好名字的神经元。真实 FFN 的中间维度通常更宽,特征也会纠缠、叠加,一个概念可能由许多维度共同表示。

FFN 真的在“存知识”吗

  “FFN 像知识库”来自一个很有启发性的观察。对于经典两层 FFN,第一层的行向量可以看成一组模式匹配器;输入与某个模式越相似,对应激活越强。第二层的列向量则把这些激活重新写回模型表示。

  因此可以用近似的 Key-Value Memory 视角理解:

\[a_i=\sigma(k_i^\top x), \qquad \operatorname{FFN}(x)=\sum_i a_i v_i\]

  这里的 $a_i$ 是输入与第 $i$ 个模式的匹配强度,$v_i$ 是它被激活后写回模型表示的方向。

  Geva 等人的实验发现,一些“Key”会响应可解释的文本模式,低层往往偏向浅层模式,高层更偏语义;对应的“Value”会推动某些词的输出概率。[3] 后续 ROME 模型编辑研究也发现,中间层 FFN 的计算对部分事实关联十分重要。[4]

  但“知识存储”是解释视角,不是硬盘地址。不能把它理解成“巴黎是法国首都”完整地放在某一个神经元里:

  • 一个事实可能分布在多个神经元、多个层和多种参数中。
  • 同一个神经元可能参与多个不相干的模式。
  • 是否能召回某个知识,还依赖 Attention 带来的上下文和残差流。
  • 能编辑某层的事实,不等于知识只存放在那一层。

  更稳妥的说法是:FFN 参数参与保存和调用模型学到的模式与事实关联,但知识是分布式的。

为什么 FFN 又大又贵

  经典 FFN 有两个大矩阵;SwiGLU 有 gate、up、down 三个投影。如果直接沿用相同的中间宽度,参数和计算都会增加,所以使用 SwiGLU 的模型通常会相应缩小中间维度,以控制总参数量。

结构 主要投影 核心特点
ReLU/GELU FFN up、down 一条内容支路,结构简单
SwiGLU FFN gate、up、down 用门控支路筛选候选内容

  FFN 的大矩阵乘法也解释了为什么 MoE 通常改造 FFN:准备很多组 FFN“专家”,但每个 Token 只经过少数专家。这样可以增加参数容量,而不让每次计算同比例增长。

容易混淆的四件事

  1. FFN 不负责跨 Token 通信。 它逐位置使用同一套参数。
  2. 升维不是目的本身。 更宽的空间配合非线性,才带来丰富的特征组合。
  3. SwiGLU 的门不是二进制开关。 它是连续的、逐元素的调节信号。
  4. “FFN 存知识”不是精确地址表。 知识分布在多个组件与层中。

90 秒讲解视频

参考文献(References)

  1. Vaswani, A., Shazeer, N., Parmar, N., et al. “Attention Is All You Need.” Advances in Neural Information Processing Systems 30, 2017. arXiv:1706.03762.↩
  2. Shazeer, N. “GLU Variants Improve Transformer.” arXiv preprint arXiv:2002.05202, 2020.↩
  3. Geva, M., Schuster, R., Berant, J., and Levy, O. “Transformer Feed-Forward Layers Are Key-Value Memories.” Proceedings of EMNLP 2021, pp. 5484–5495, 2021. doi:10.18653/v1/2021.emnlp-main.446.↩
  4. Meng, K., Bau, D., Andonian, A., and Belinkov, Y. “Locating and Editing Factual Associations in GPT.” Advances in Neural Information Processing Systems 35, 2022. arXiv:2202.05262.↩

第十四课复习总图

FFN、SwiGLU 与知识存储复习总图

从逐 Token 加工到门控与分布式知识,点击查看原始 SVG

  下一课看 MoE:为什么准备很多专家,却只让每个 Token 经过其中几个?