第 17 章 拼出完整的 Transformer 块

注意力输出现在与输入一样宽，因果掩码也能阻止读取未来。把这个操作重复几次，就能叫 Transformer 了吗？还有几件事没有解决。不提供位置信息时，不受可见性限制的注意力面对重新排序的列表，只会相应地重排输出。混合结果还可能替换掉原本有用的输入表示。而且，位置之间交换信息，与一个位置内部形成新的非线性坐标组合，是不同的工作。

完整的 Transformer 块（Transformer block） 需要几种操作相互配合。我们先看各操作补了什么缺口，再沿一个位置走完一份规则完整的计算。投影和两个头的约定来自第 15—16 章，前向计算与学习的区别沿用第 7 章。下面所有权重均为给定的教学数值。

17.1 让相同 token 能处在不同位置

例 17.1：没有位置信息时，交换行会怎样？ 使用第 15 章不受限制的注意力，把输入第三行移到最前面，原第一、第二行依次跟在后面。由于所有位置共用投影矩阵，Q、K、V 也会一起重排。输出就变成原第三行、原第一行、原第二行。计算从未收到“原先排第几”这个信息，自然无从仅靠它恢复原来的位置。

这个性质叫 置换等变性（permutation equivariance）：输入做什么重排，输出就跟着做同样的重排。它与“输出不变”的不变性不同。这里有一个前提：所有位置都可见，或者可见关系也随位置一起重排。如果固定因果三角关系，只移动 token，那么哪些 token 可以相互作用就变了，上述结论不再适用。因果可见性确实提供了顺序限制，但不能据此保证各种有用的位置关系都能得到良好表示。

提供位置的一种直接办法，是给每个槽位再配一个 C 坐标向量，与 token 表示相加。把第 15 章的 token 矩阵改记为 E，以便与送入块的 X 区分，并给定位置表 P：

$$E=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix},\qquad P=\begin{bmatrix}0&0\\0&1\\1&0\end{bmatrix},\qquad X=E+P=\begin{bmatrix}1&0\\0&2\\2&1\end{bmatrix}.$$

三份矩阵都是 [3,2]。第二个位置变成 [0+0,1+1]=[0,2]。相加仍保留两个坐标；若改成拼接，宽度就翻倍了。这样，相同 token 放在不同槽位，可以得到不同的输入向量。但模型还得学会利用这一区别；给定坐标本身并没有表达一种普遍适用的距离含义。

P 可以是一张学习得到的表，也可以由固定规则计算。本书 Mini GPT 使用有限长度的学习表。32 行的表只有索引 0—31；其他代码即使能接收更多 token，也不会凭空多出训练过的第 32 号位置。即使某种位置规则在更大索引上仍有定义，也不能单凭这一点证明长上下文有效。现代替代方法留到第 25 章。

17.2 给原输入留一条直接相加的路径

设注意力子层算出 U，shape 与 X 相同。可以不直接用 U 替换 X，而是把它加上去：$R=X+U$。这叫 残差连接（residual connection）。例如，一个位置的计算是 [1,0]+[0,2]=[1,2]。原来的两个数有一条直接相加的路径通向 R，子层则贡献一份调整量。如果子层返回零，这条路径就完整保留 X。

参与相加的两个数组，必须对应同一批样本、同一组位置和相同的坐标宽度。维度相等是必要条件，却不是充分条件：误加了另一条样本的表示，shape 仍可能完全正确。很多库还允许把 [1,2] 的一行广播到三个位置，但那并不是规定的“各位置加上自己的一份更新”。配套 residual 函数因此拒绝不完全相等的 shape。多头注意力的输出投影先把宽度恢复过来，才能做这次相加。

残差路径并不保证原来的每条事实都保存下来，也不保证梯度永不消失。更新量可以抵消输入的一部分，后面的操作也可以丢掉差别。它提供的是一条加法路径，不是一份可靠性保证。

17.3 在一个位置内部归一化，不把整句混在一起

进入各次变换的数值幅度可能不同。归一化可以先整理行内的数值尺度，再交给下一子层，并保留可学习的调整。先看 [1,3]：平均值是 2，偏差是 [-1,1]，偏差平方的平均值为 1。除以这个方差的平方根，就得到 [-1,1]。但如果输入是 [10,10]，方差为零，还必须规定怎样避免除以零。

层归一化（Layer normalization，简称 LayerNorm） 会在方差中加入一个很小的正数 epsilon，再开平方；随后为每个坐标乘以可学习的系数、加上可学习的偏移。对于有 C 个坐标的一行 x，本书使用的精确规则是：

$$\mu=\frac{1}{C}\sum_{c=1}^{C}x_c,\qquad \sigma^2=\frac{1}{C}\sum_{c=1}^{C}(x_c-\mu)^2,\qquad \operatorname{LN}(x)_c=\gamma_c\frac{x_c-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_c.$$

下标 c 选的是坐标，不是 token 位置。分母取 C，因为这里定义的是 眼前这 C 个值 的平均平方偏差，不是在用样本估计总体方差，因此不用 C−1 修正。gamma 和 beta 各有 C 个可学习参数，各位置共用。例中 gamma 全为 1，beta 全为 0，epsilon=0.00001。于是 [1,3] 变成约 [-0.999995,0.999995]，[10,10] 变成 [0,0]。如果 beta 不是零，后一行的结果就是 beta。

对于 [B,T,C]，只沿最后的 C 轴归一化，每条样本、每个位置分别计算。如果把未来位置也混进统计量，即使注意力掩码正确，仍可能泄漏信息。这种 LayerNorm 在训练和 模型运行（inference） 时，都根据当前行计算统计量，不保存整句的滑动平均值。epsilon、可学习的缩放与偏移还意味着，最终结果的方差不一定严格为 1，均值也不一定为零。LayerNorm 的公式与计算轴。

17.4 在每个位置内部形成新的组合

注意力负责位置之间的通信，我们还想在一个位置内部对坐标做非线性变换。第 7 章已经用过这样的结构：线性变换、把负数改成零的 ReLU、再一次线性变换。现在，让每个位置分别通过同一个小网络，共用权重。这叫 逐位置前馈网络（position-wise feed-forward network，简称 FFN）。

给出一个完整的 2→3→2 算例：输入行 [-1,1]，偏置全为零，矩阵为：

$$W_1=\begin{bmatrix}1&-1&0\\0&1&1\end{bmatrix},\qquad W_2=\begin{bmatrix}1&0\\0&1\\1&-1\end{bmatrix}.$$

第一次乘法得到 [-1,2,1]。例如，中间坐标来自 $(-1)(-1)+1\times1=2$。ReLU 把这一行变成 [0,2,1]，第二次乘法得到 [1,1]，其中最后一个坐标来自 $0\times0+2\times1+1\times(-1)=1$。宽度先由 2 扩成 3，再回到 2，以便做残差相加。一般形式为：

$$\operatorname{FFN}(x)=\operatorname{ReLU}(xW_1+b_1)W_2+b_2.$$

设模型宽度为 C，中间宽度为 F，那么 W_1 是 [C,F]，b_1 是 [F]，W_2 是 [F,C]，b_2 是 [C]。各位置用同一组矩阵，但输入值是自己的。如果没有中间的非线性操作，两次变换可以合成一次矩阵乘法加偏置，也就是一次仿射变换。有了 ReLU，不同输入中被清零的坐标可以不同。这个 FFN 没有读取另一个位置的操作，只加工当前行已经带来的信息。

17.5 沿一个位置走完整个 Pre-Norm 块

例 17.2：第一个位置的每一步。 输入使用前面的 X=E+P。注意力使用例 16.3 的两个一坐标因果头和 W_O，FFN 使用刚给出的矩阵；两次 LayerNorm 都设 epsilon=0.00001、gamma 为 1、beta 为 0，不使用 dropout。把归一化放在对应子层 之前，就是 前置归一化（pre-norm） 块：

$$N_1=\operatorname{LN}_1(X),\quad U=\operatorname{MHA}(N_1),\quad R=X+U,\quad N_2=\operatorname{LN}_2(R),\quad Y=R+\operatorname{FFN}(N_2).$$

两次 LN 的下标用来区分参数集，不表示不同的归一化轴。X 第一行是 [1,0]，均值 0.5、方差 0.25，归一化后为 [a,-a]，其中 $a=0.5/\sqrt{0.25001}\approx0.999980$。第一个位置受因果限制，只能看到自己的键。每个头唯一允许的权重都恰好为 1，输出分别是 a 和 −a。拼接后乘 W_O，得到 [0,2a]，第一次残差相加得到 [1,2a]。

这一行的均值是 $(1+2a)/2$。记 $u=(2a-1)/(2\sqrt{((2a-1)/2)^2+0.00001})\approx0.999980$，归一化结果就是 [-u,u]。FFN 先得到 [-u,2u,u]，ReLU 保留 [0,2u,u]，最后返回 [u,u]。因此，第一行最终为 [1+u,2a+u]，约为 [1.999980,2.999940]。两次残差相加、两次归一化都参与了计算，已经不是前面单独的注意力输出。

表 17.1：第一行的完整运行记录，保留六位小数。 独立标量计算与 NumPy 路径的结果之差，小于绝对容差 0.000000000001。

阶段	第一个位置	整个数组的 shape
token 与位置相加，X	[1.000000,0.000000]	[3,2]
第一次归一化，N_1	[0.999980,-0.999980]	[3,2]
注意力输出，U	[0.000000,1.999960]	[3,2]
第一次残差相加，R	[1.000000,1.999960]	[3,2]
第二次归一化，N_2	[-0.999980,0.999980]	[3,2]
FFN 的 ReLU 之后	[0.000000,1.999960,0.999980]	[3,3]
FFN 输出	[0.999980,0.999980]	[3,2]
块的最终输出，Y	[1.999980,2.999940]	[3,2]

包含读取多个键的位置在内，完整输出是：

$$Y\approx\begin{bmatrix}1.999980&2.999940\\0.999912&1.476742\\2.999964&3.746396\end{bmatrix}.$$

进入块之前

1. 模型保存的数值token 表 [V, C]；位置表 [T_max, C]
2. 处理步骤取出相应行，再相加 → X [B, T, C]
3. 一组数字本例 X₁ = [1, 0]

此处提供位置；不是往序列末尾再加一个 token。

第一条残差路径

1. 处理步骤沿 C 做 LN₁ → N₁ [B, T, C]
2. 模型保存的数值Q/K/V 投影权重 [C, C]
3. 处理步骤拆头 → [B, H, T, d]；QKᵀ/√d → [B, H, T, T]
4. 处理步骤掩码 → 逐行 Softmax → 混合 V [B, H, T, d]
5. 处理步骤拼接 → W_O → U [B, T, C]
6. 一组数字相加 X + U → R；R₁ = [1, 1.999960]

X 一路经过归一化和注意力，另一路直接参与相加。

第二条残差路径

1. 处理步骤沿 C 做 LN₂ → N₂ [B, T, C]
2. 模型保存的数值W₁ [C, F]、b₁ [F]；W₂ [F, C]、b₂ [C]
3. 处理步骤N₂W₁+b₁ → ReLU [B, T, F] → W₂、b₂ [B, T, C]
4. 一组数字相加 R + FFN → Y；Y₁ = [1.999980, 2.999940]

R 同样有一路直接参与第二次相加。Mini GPT 在此块之后再做最终归一化。

图 17.1 — 每个子层都回到残差相加所需的宽度。 本书原创教学图，CC BY-SA 4.0。数值与规则见正文及共享数据。

图 17.1 分开画出了两条残差路径。位置相加在块之前完成；注意力中途引入头轴和来源位置轴，最后返回 [B,T,C]；FFN 扩展的是最后的坐标轴，不是位置轴。这样就能把图与 code/part-iv/mechanisms.py 的 small_block 对上：norm1、attention、residual1、norm2、hidden、ffn、output 都是明确命名的记录项，不是藏在框架里的步骤。reference.py 另用标量循环计算同一个块，不依赖 NumPy 的向量化实现。

norm1 = layer_norm(x)
attention = two_heads(norm1)["output"]
r = residual(x, attention)
norm2 = layer_norm(r)
hidden = np.maximum(norm2 @ w1, 0)
y = residual(r, hidden @ w2)
assert y.shape == x.shape

这段代码使用例中为零的偏置。完整函数提供了上述矩阵，并检查残差两端的 shape。运行记录见 data/part-iv/numpy-run.json。以后加上批轴，各行仍分别做逐位置计算，各条样本也仍保留自己的注意力关系，不会把归一化或 FFN 变成跨样本混合。

17.6 换了顺序，就换了一个块

原始 Transformer 在每个子层与残差相加 之后 才做归一化，这叫 后置归一化（post-norm）。继续用相同名字表示各种操作，它的顺序是：

$$R_{\text{post}}=\operatorname{LN}_1\!\left(X+\operatorname{MHA}(X)\right),\qquad Y_{\text{post}}=\operatorname{LN}_2\!\left(R_{\text{post}}+\operatorname{FFN}(R_{\text{post}})\right).$$

Post-Norm 把 X 本身送进注意力，Pre-Norm 则送入 LN(X)。Post-Norm 还在返回结果前，对第二次残差相加做归一化。因此，在图上移动同样几个框，会改变数值结果；两者不是同一算法的两种写法。顺序还会影响优化过程，但一个小型前向算例无法证明哪种在各种任务上训练得更好。原始后置顺序、前后置对照，表 1。

本书 Mini GPT 选择 Pre-Norm，并在块之后、词表输出层之前，再加一次最终 LayerNorm。它位于本块两条残差路径之外。原始 Transformer 包含 dropout，本书的有界教学项目将其设为零。这里的仅解码器块有因果自注意力和 FFN；编码器—解码器结构中的解码器，还需要交叉注意力子层来读取编码后的源序列。“块”表示一种组合方式，不代表所有 Transformer 的内部都完全相同。

17.7 现在能在图中找到什么

token 表示现在能够加上位置向量，在允许的范围内交换信息，保留直接相加的路径，归一化自身坐标，再经过共用的非线性变换。shape 相容让这些操作接得上，轴的含义则决定接得对不对。增加这些结构，既不等于把一次未训练的前向计算变成了学会语言的模型，也不能证明它能可靠使用远处的钥匙线索。

架构还剩下一组选择：哪些位置应该可见，哪些输出应该成为训练目标？同一套基本组件，可以用于阅读完整输入、续写前缀，也可以在另一条序列的条件下生成新序列。第 18 章先比较这些选择，再为我们的小生成器确定接口。

17.8 习题

1. E 的两个位置都为 [2,1]，对应的位置向量分别是 [0,0]、[0,1]。计算送入块的两行。它们不同，是否证明网络学会了顺序？
2. 使用本章 epsilon、gamma、beta，归一化 [2,4]，再归一化 [20,20]。解释方差分母为什么取坐标数，以及 epsilon 的作用。
3. 张量是 [2,5,16]。有人让 LayerNorm 沿五个位置求平均，让 FFN 把 16 个坐标扩为 64 再缩回 16。哪项操作违反本章的块约定？因果掩码能修复它吗？
4. 用给定 W_1、W_2，计算输入 [1,-1] 的 FFN 输出，与 [-1,1] 比较。指出各自被 ReLU 清零的项。
5. 把 [2,5,16] 的注意力输出与 [1,5,16] 相加，库通过广播算成功了。为什么本块的残差断言仍应该拒绝？
6. 一份检查点采用 Pre-Norm，有人因为所有 shape 都相同，就把矩阵加载到 Post-Norm 实现中。指出两项改变的计算，以及 shape 以外还必须核对什么。

17.9 参考解答

1. 两行分别是 [2,1]、[2,2]。模型收到了数值差别，但尚未证明它学会使用。它可以忽略发生变化的坐标，也可以在后面抹去差别。需要在受控任务中改变顺序，再检查行为，才能获得相应证据。
2. 均值 3、方差 1，结果为 [-1,1]/sqrt(1.00001)，约 [-0.999995,0.999995]。[20,20] 的两个偏差都为零，因此输出 [0,0]。本行有两个坐标，所以偏差平方和除以 2；epsilon 保证方差为零时，分母仍为正。
3. 沿位置求平均违反了只归一化坐标的规定，还可能让未来 token 改变较早位置。注意力掩码挡不住这条独立路径。FFN 的 [2,5,16] → [2,5,64] → [2,5,16] 是合适的，前提是每个位置分别使用同一组权重。
4. [1,-1] @ W_1 得到 [1,-2,-1]，ReLU 留下 [1,0,0]，再乘 W_2 得到 [1,0]。[-1,1] 依次得到 [-1,2,1]、[0,2,1]、[1,1]。不同输入被清零的项不同，所以整体不能化为一次线性映射。
5. 第二个张量只有一条样本，没有为两条样本分别计算更新量。广播可能把一条样本的表示误用到另一条样本中。完整 shape 相等是本块残差约定的一部分；库算成功了，不能证明样本对应正确。
6. Pre-Norm 的注意力收到 LN(X)，Post-Norm 收到 X；Post-Norm 还会归一化最后一次残差相加。必须核对操作顺序、参数身份、epsilon、位置规则、掩码，并对照一份已知前向记录。仅有一组大小相容的张量，无法复现检查点原来的函数。

17.10 参考文献

• Vaswani 等（2017），Attention Is All You Need。3.1、3.3、3.5 节分别说明原始架构中的残差与后置归一化、逐位置 FFN 和位置相加。
• Ba、Kiros、Hinton（2016），Layer Normalization。第 3 节与式 5 解释样本内部的统计量以及可学习的缩放、偏移；本章另行明确了 epsilon 和行内计算。
• Xiong 等（2020），On Layer Normalization in the Transformer Architecture。表 1 给出两种运算顺序，以及 Pre-Norm 堆叠之后的最终归一化。本章没有复现其优化实验。
• PyTorch，LayerNorm，2.7 版。可核对显式实现所用的末轴统计量、方差分母和仿射参数。