读回上下文之后，一个 token 的表示怎样更新？

English · 中文 · Français

First Principles of LLMs & RL · Article 07 · 上一篇

写作说明：想法与初始材料来自我本人。ChatGPT 协助了草稿整理、编辑、翻译、数值核查，以及原创图示与动画源码的制作。

上一篇算到最后，第二个位置手里有两个向量。一个是它进入这一层时已有的表示，另一个是它从可见位置读回来的内容：

假如现在要写下一行代码，把某个向量交给下一层，该交哪一个？

直接交出第二个，程序完全可以运行。两个向量的形状相同。但这行代码已经做了一个决定：用刚读回来的内容替换当前位置原本的表示。还有一种写法：把读回来的内容当作对这个位置的更新。一个 Transformer block 的结构，可以从这两行代码的差别讲起。

既然能读取自己，为什么还要加回自己？

先回忆第二个向量是怎么算出来的。位置 2 可以读取位置 1，也可以读取自己，不能读取未来的位置 3。在上一篇的构造例子里，它给两个可见位置各分了一半权重：

这里确实已经有一份“自己的信息”。但请看它经过了哪条路：原表示先经过 Value 投影，再乘以分给自身的权重，最后和另一个位置的贡献混在一起。这个例子碰巧让第二个位置的 Value 和原输入数值相同；换一组 Value 投影，两者就可能不同。即使投影不变，如果这个 head 把更多注意力给了位置 1，自己的那份贡献也会随之减小。

这就给出一个很实际的设计理由：一个位置可以把注意力给别处，同时让自己的原表示继续传下去。

我们可以给输入另留一条直接通路。先取最简单的情形：只用一个 head，读出宽度恰好等于输入宽度，输出投影取单位矩阵，暂不加入归一化。令 u₂ 表示这次更新后的位置表示：

这条绕过读取分支、直接进入加法的通路，就是这里的 residual connection（残差连接）。为什么叫“残差”？如果希望把输入从某个值变成另一个值，分支只需提供两者的差。上面第一坐标原本为零，分支把它增加了 1；第二坐标原本为 1，分支只需再增加半个单位。换成别的参数，更新也可以为负。¹

有一个不用猜测词义的检验：把分支的输出设为零。 若用读取结果替换原表示，下一层拿到零；若将它加回原表示，下一层仍拿到原输入。反过来，如果分支恰好给出原输入的相反数，两者也能完全抵消。所谓“保留”，精确地说，是给输入保留一条直接通路，并不意味着在最终的和里面总能找回原来的每一项。

两条路径的输入和分支输出都相同。把分支输出设为零，再看下一层收到什么。

一个 head 只有一组位置权重，这限制了什么？

再看一眼刚才那次读取。乘在第一个 Value 前面的二分之一，作用于这个向量的所有坐标。如果想从位置 1 多取一点第一坐标，同时从位置 2 多取一点第二坐标，一个 head 能分别做这两件事吗？

把位置 1 的权重记为 a，位置 2 的权重就是 1−a。固定刚才的两个 Values，所有可能的读取结果都具有同一种形式：

用手把 a 调大一点：第一坐标增大，第二坐标就减小。把两个 Values 画成平面上的两个点，读出的向量只能沿着连接它们的线段移动。想得到第一坐标 3/2，就要把 a 调到 3/4；想让第二坐标也达到 3/4，却要把同一个 a 调到 1/4。差的不是一个更精确的权重，而是第二个可以独立调节的权重。

现在改用两个 head。为看清差别，给它们各一个标量输出：A 的 Value 投影取上例的第一坐标，B 的 Value 投影取第二坐标。两个 head 各自计算匹配，这里人为指定它们在位置 2 的读取权重：

现在两个想要的数可以同时得到。A 不必迁就 B，因为每个 head 都有自己的匹配分数与 Softmax。这里改变了读取规则，表中的权重并非上一篇模型的测量结果；为了看清机制才令每个头挑选一个坐标，实际的 Value 投影可以学习坐标的线性组合。在这两个固定的 Values 上，多头多出来的自由，是让不同的特征按不同的比例从各位置取值。 如果连 Value 投影也允许更换，单头同样可能在这一个输入上凑出这两个数；上面的比较说的是读取方式。

多头读完以后，还差一步。拼接只是把两个结果并排放好，模型还需要决定它们分别怎样影响原表示。沿用上一篇的记号，单个位置的向量写成列向量，放进序列矩阵时转置成行。把拼接结果记作 s₂，选一个便于手算的输出投影 W_O：

这个矩阵的作用是：第一坐标写入 A 与 B 的差，第二坐标写入 B。于是 attention 分支真正交给残差加法的更新，以及加完之后的位置表示，分别是：

可以把这次计算里的三件事分得很清楚：每个头决定从哪些位置取多少；输出投影决定各头读回的数怎样组合；残差加法才把组合结果写入当前位置。 例如 4 个头各读出 2 个数，拼接后就有 8 个数。输出投影把这 8 个数映射回原表示的宽度，两条路径才可以逐坐标相加。原始 Transformer 的多头构造采用的正是拼接后做输出投影的形式。²

两组权重均为教学设定。先看各头读出什么，再分别跟踪拼接、输出投影和最后的残差加法。

信息到了这个位置，然后呢？

刚才更新完，位置 2 手里有了两个数：3/4 和 7/4。现在交给它一个小任务：如果第二坐标比第一坐标大，就把超出的部分加到第一坐标；否则保持原样。

这一步还需要再看一遍上下文吗？不用。做决定所需的两个数都已经在这个位置了。把一个位置上的两个坐标记为 p、q，只需一个隐藏单元计算差值，把负数截成零，再把结果写回第一坐标：

在刚才的例子中，差值恰好为 1，加回之后得到：

这个小计算就是一个逐位置前馈网络（position-wise FFN）能做的事。第一次线性变换提取“第二坐标减第一坐标”；ReLU 让这个特征只在差值为正时起作用；第二次线性变换指定它应该写到哪里。若去掉 ReLU，两次线性变换可以合并成一次，也就失去了这个“只在一侧起作用”的性质。没有理由期待真实模型的两个坐标恰好带有这个解释；这里给的是一组可以用手检查的参数。

对比一下序列中的两个位置。Attention 可以让位置 2 的结果依赖位置 1 的内容；FFN 拿到位置 2 更新后的向量，就只变换这个向量。把同一 FFN 应用到位置 1 时，用的是同一套参数，得到的结果却可以不同，因为输入不同。说 FFN “逐位置”不意味着它只能处理孤立的词：上下文已经通过前面的 attention 到达了这里。²

原始 Transformer 使用带 ReLU 的 FFN。我的手写实现用了 SwiGLU：在每个位置上，一路投影产生候选特征，另一路投影经过激活后调节这些特征，再投回模型宽度。这改变了一个位置内部的计算方式；它依然不会自行去读取下一个位置。³

归一化放在哪条路上？

到这里，我们为了把加法算清楚，一直没有做归一化。把这样的层摞起来，新的问题就出现了：每一层都在写入更新，后面分支收到的向量，数值尺度未必和前面一样。

这个变化会影响实际的读取。想象一次不带偏置的点积 attention：所有输入向量都放大两倍，其他条件不变，线性投影得到的 Queries 和 Keys 也放大两倍。它们的点积于是放大四倍。哪怕原先只是某个候选略占优势，Softmax 也可能因此变得更集中。上一篇除以 √dₖ 处理的是点积维数带来的尺度问题；这里连维数都没变，变化来自输入本身的幅度。

一种设计是：让残差流继续携带现有表示，但每个计算分支先用归一化后的表示工作。具体选用 RMSNorm。对单个位置，记输入 z 的宽度为 d，γ 为逐坐标的可学习系数，ε 是一个很小的正数，防止分母为零。其计算为：

分母由这个位置内部的各个坐标计算，不跨 token 求平均。当 ε 相对于输入的均方值可以忽略时，把整根向量乘以同一个正数，分子与分母会近似按相同比例变化。分支因而不会仅仅因为整根输入变长，就自动给出更极端的匹配分数。这里没有减去均值，这一点也区别于 LayerNorm。⁴

现在把完整 block 写出来。用 X 表示输入序列，n 行对应 n 个位置，每行宽度为 d_model；U 是 attention 写入后的序列，Y 是整个 block 的输出。N₁ 和 N₂ 是两组参数独立、分别作用于每一行的 RMSNorm；MHA 包含 causal 多头读取及输出投影。省略 batch 维和 dropout：

注意两个加号的左边：第一处加回的是 X 本身，第二处加回的是 刚刚更新过的 U。归一化发生在两个分支入口处，直接通路不经过它。这就是本文采用的 Pre-Norm。前面的数值是为了分开观察各部件而算的局部例子；完整 block 改用归一化后的分支输入，不应照抄那些数值作为新的前向结果。对应的实现只有两次残差写入，去掉缓存逻辑后，核心可以读成这两行：

u = x + attention(attention_norm(x))
y = u + feed_forward(feed_forward_norm(u))

完整 block 的分支使用归一化后的输入，前面的局部手算则省略了归一化。FFN 分别作用于各位置；更新后的表示为下一层的投影提供输入。⁵

现在回到最开始的检验：若两个分支对各自的输入都给出零更新，第一行得到原输入，第二行仍得到原输入。漏掉第一行的 x +，即使所有张量的形状都正确，结果也变成零。这里检验的是当前这个 block，不包括模型末尾另设的归一化层。在原始 Transformer 的 Post-LN 写法中，归一化位于残差相加之后；即使子层输出为零，输入还要经过 LayerNorm，不能原封不动地搬用同一结论。²⁶

这也能接回 Article 03 的反向传播。把一个残差子层的整个输入展平成列向量 x，将其分支记为 F，输出记为 y。设损失对输出的梯度为列向量 g，在可微处，由链式法则：

J_F 是分支的雅可比矩阵。右边第一份 g 不必穿过分支；它来自前向那个直接的加号。另一项仍可能放大或抵消它，因而这条式子说明了梯度有哪条路可走，并不保证梯度永远足够大。¹

下一层的 Query 从哪里来？

回到 I like 和 They like。在上一篇的设定中，like 的 token ID 和位置相同，所以它在第一层可以从同样的表示出发，提出同样的读取请求。前面的词不同，同一个请求也可能读回不同内容。

经过这一层，like 的 token ID 并没有改变。变的是它所在位置交给下一层的那一行：原来的表示已经接收了 attention 的写入，FFN 又根据写入后的特征计算了一次更新。下一层从这一行的新值出发，用自己的归一化和投影参数生成 Query。于是，前文的差异有机会改变的不只是本次读到了什么，还包括下一次据以查找信息的线索。投影仍可能忽略这些差异；结构给出的是这条依赖关系，并没有预先规定模型会学到哪一种语法解释。

这就是堆叠两层比重复播放一次 attention 动画多出来的东西：第二次读取的起点，已经可以包含第一次读回的内容。

到了最后一层，还有一件事没有发生。最后一个位置交出来的依然是一个表示向量，而我们等的是一个词。下一篇就从这一步走到真正的生成：怎样从这个向量得到词表上的分布？选出的 token 一旦追加到句尾，哪些计算需要重做，哪些可以保留下来？