技术博客

残差流视角:信息在层间怎么流

第 64 / 147 篇 · 2025-04-13 · ai · 残差连接可解释性架构分析

前面几篇把注意力的内部拆开看了,这篇把镜头拉远一层:一整个 $L$ 层的 Transformer block 堆叠,信息到底怎么流动?答案是出人意料地简单——只有一条"总线"。

一条加法链

pre-LN 架构下,每一层做的是

$$ x_{l+1} = x_l + f_l\big(\mathrm{LN}(x_l)\big),\qquad f_l \in \{\mathrm{Attn}_l,\ \mathrm{FFN}_l\}, $$

递归展开就是一条纯加法链:

$$ x_L = x_0 + \sum_{l=0}^{L-1} f_l\big(\mathrm{LN}(x_l)\big). $$

这个式子是理解一切的钥匙:$x_l$ 不是某个层的"私有输出",而是一个所有层共享、逐层累加的公共状态。Elhage et al.(2021,A Mathematical Framework for Transformer Circuits)给它起了个贴切的名字——残差流(residual stream)。注意力头、FFN、乃至最后的 unembedding,都只是挂在这条总线上的设备:读总线当前内容,算出一个增量,写回去。没有任何层能改写别人的写入,只能往上加。

为什么这个视角有用

三个此前散落的观察,在总线视角下变成了同一件事的推论。

logit lens。 既然 $x_l$ 是"半程的 $x_L$",那任意中间层的表示都可以直接乘上 unembedding 矩阵 $U$ 看预测分布:$\hat p = \mathrm{softmax}(U\,\mathrm{LN}(x_l))$。实测发现靠后的层已经能给出接近最终答案的分布——残差流里确实流动着"正在成形的答案",而不是只有抽象特征。这个便宜实验(不需要任何训练)只有在"所有层共享同一坐标系"的前提下才说得通。

梯度高速公路。 对总线求梯度,$\partial \mathcal L/\partial x_l = \partial \mathcal L/\partial x_L$,干干净净直通输入。每个子层的贡献梯度还要乘上它自己的雅可比,但那条贯穿全网络的主干道不经过任何非线性、任何 LN——这就是深层网络还能训练的直接原因。代价留到 pre-LN 那篇细说。

叠加态。 总线只有 $d$ 维(LLaMA-3 是 4096),但要同时承载:embedding 本身、几十个注意力头的写入、几十层 FFN 的写入、最后被 unembedding 读出来。容量不够分怎么办?做法是叠加(superposition):远多于 $d$ 个的特征以非正交方向压缩进 $d$ 维空间,靠高维几何(近乎随机的方向几乎正交)减少互相干扰。这和 FFN 的键值记忆假说(下一篇)正好衔接——每个知识神经元在总线上占一个方向,读写互不干扰的前提是方向足够"分散"。

一个不能忽略的副作用:范数膨胀

逐层加增量,$x_l$ 的范数几乎必然增长。粗略估计:若各层写入方向近似随机(内积期望为零),

$$ \|x_L\|^2 \approx \|x_0\|^2 + \sum_l \|f_l\|^2 \;\Rightarrow\; \|x_L\| \sim \sqrt{L}\ \text{量级}. $$

70 层堆下来,末层范数可以是 embedding 的数倍。这件事本身不是 bug——总线本来就该越写越"满"——但它和归一化的交互很微妙:LN 按 $\|x_l\|$ 缩放输入,深层 LN 的输入范数大,意味着每个子层实际拿到的"步长"也在变。这解释了为什么 pre-LN 架构里 residual 分支的初始化常带 $1/\sqrt{N}$ 缩放,也是"归一化摆放位置"那两篇要展开的问题之一。

残差流:所有层读写同一条总线,增量逐层叠加 残差流 x Attn(LN(x)) + FFN(LN(x)) + Attn(LN(x)) + 读出:U·LN(x) 嵌入 x₀

小结

残差流的等式 $x_L = x_0 + \sum_l f_l$ 把 Transformer 从"层层传递的黑箱"改写成了"一条总线上的并发读写":logit lens 可行是因为各层共享坐标系,深层可训练是因为梯度有一条不经非线性的高速公路,特征能塞进 $d$ 维是因为叠加。总线视角的代价问题是范数随 $\sqrt L$ 膨胀——子层写入的相对尺度怎么控,这就轮到归一化出场了。