模块 6: Transformer

第 04 模块 最后介绍了注意力（attention）：它如何缓解编码器-解码器的瓶颈。在每个输出步骤，解码器读取编码器状态的加权和；权重由解码器当前状态与各编码器状态之间的匹配分数决定。Bahdanau 的方法用一个小网络计算分数：

其中，\mathbf{s}_{t-1} 是解码器状态，\mathbf{h}_j 是位置 j 的编码器状态，\mathbf{a}_t 是解码器在步骤 t 读取的上下文向量。Luong 等人用乘法形式取代小网络：或者直接计算点积 \mathbf{s}^\top\mathbf{h}_j，或者使用双线性形式 \mathbf{s}^\top\mathbf{W}\mathbf{h}_j，从而将所有位置的打分合并为一次矩阵乘法。两种方法都是在两个循环网络之上添加注意力。Transformer（Vaswani 等，2017）则彻底去掉循环计算，使注意力成为位置之间传递信息的唯一途径。

循环计算的代价

循环计算有两项代价，无论怎样优化工程实现都无法消除。

时间维度上的计算必须串行进行。 步骤 t 需要 \mathbf{h}_{t-1}，而后者又需要 \mathbf{h}_{t-2}，依此类推，直到第一个 token。因此，无论使用什么硬件，一条包含 10,000 个 token 的训练序列在每一层都需要 10,000 个相互依赖的步骤。GPU 可以并行处理多条序列，但同一条序列中的每一步都必须等待前一步。教师强制虽然让所有输入提前已知，也无法解决这个问题，因为非线性运算位于循环内部。

信息传播路径很长。 位置 j 的信息要经过 t - j 次循环更新才能到达位置 t。每次更新都将信息压入相同宽度的状态向量，可能丢失一部分信息。第 04 模块 中的梯度消失，正是沿这条长路径反向传播时出现的问题。

自注意力（self-attention）消除了这两项代价。每个位置都对所有允许看到的位置计算加权和，而且所有位置可以通过一次矩阵乘法同时完成计算。一条包含 10,000 个 token 的序列，因而对应一次大型矩阵运算，而非 10,000 次相互依赖的小型运算。无论相距多远，任意两个位置之间的传播路径都只经过一层。

代价与权衡

这种改进也有代价。Vaswani 等（2017，表 1）比较了两种层处理长度为 T、向量宽度为 d 的序列时的成本：

循环层在 T 个步骤中的每一步，都将一个 d 维向量乘以一个 d \times d 矩阵；自注意力则通过长度为 d 的点积，为 T^2 对位置打分。两者的成本比为 T^2 d / (T d^2) = T/d，因此在 T < d 时，自注意力的成本较低。原始翻译模型的句子通常只有几十个 token，宽度为 d = 512，正属于这种情况；T = 50 时，成本比约为 0.1。超过 T = d 后，自注意力就要承担关于序列长度的二次方成本：在 T = 32{,}768、d = 4{,}096 时，成本比为 8，并随上下文长度线性增长。表中只计算位置之间的信息混合；自注意力还需要 O(Td^2) 的投影运算（第 4 节），循环层的权重运算也有同样的成本。第 04 模块，第 12 节 还介绍了卷积这个替代方案，以及循环网络仍保有的一项优势：每生成一个 token，只需恒定大小的状态内存。

T^2 也体现在内存需求上，因为注意力权重构成一个 T \times T 矩阵。第 10 节 说明 FlashAttention 如何在不保存完整矩阵的情况下计算注意力：它消除了 T^2 的内存需求，但没有消除 T^2 的算术运算量。第 04 模块，第 13 节 的状态空间模型采用另一条路线：保留循环，但使递推关系成为线性的，从而能够并行计算。

注意力作为软字典查找

Python 字典是一种硬查找。给定 {"pump": v1, "valve": v2, "tank": v3} 和查询 "tank"，它将查询与键比较，找到相等的键，并精确返回该键对应的值 v3。其他查询都找不到结果。

注意力放宽了这三个步骤的要求。它用点积计算查询与每个键的匹配分数，允许部分匹配；再用 softmax 将分数转为权重，使每个键都获得正权重，且权重总和为 1；最后返回值的加权平均，因此输出混合了多个值，匹配越好的值占比越大（图 6.1）。

这一行为有两个极限。将所有分数乘以因子 c。随着 c 增大，最大分数支配 softmax：c = 10 时，权重为 (0.001, 0.001, 0.998)，输出为 (2.996, 2.997)，几乎完全是 \mathbf{v}_3。在极限下，softmax 变成 argmax，软查找变为硬查找。当 c \to 0，或各分数相等时，权重均为 1/3，输出为简单平均值 (1.333, 1.667)。注意力介于两者之间，由分数的尺度决定具体位置；这也是 第 2 节 讨论分数尺度的原因。

import numpy as np

table = {"pump": (1, 0), "valve": (0, 2), "tank": (3, 3)}
print(table["tank"])                                 # hard: one exact match

q = np.array([1.0, 1.0])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])   # keys: pump, valve, tank
V = np.array([[1.0, 0.0], [0.0, 2.0], [3.0, 3.0]])   # values
s = K @ q / np.sqrt(2)                               # scaled scores
w = np.exp(s - s.max())                              # subtracting the max avoids overflow
w = w / w.sum()                                      # softmax weights
print(np.round(w, 3), np.round(w @ V, 3))            # soft: a weighted average

(3, 3)
[0.248 0.248 0.503] [1.759 2.007]

软查找的平滑性使它能够学习。字典查找不能提供有用的梯度：查询的微小变化要么不改变结果，要么突然跳到另一个键。软查找的输出则是查询、每个键和每个值的平滑函数，损失梯度能够指出它们各自应向什么方向移动。因此，可以通过可学习的投影从 token 生成查询、键和值（第 2 节），让训练决定每个位置寻找什么、提供什么。

必须添加回来的内容

去掉循环计算，也去掉了它自然提供的两项特性。第一项是顺序。加权和本身没有位置概念：打乱输入 token 后，每个输出仍是原来的向量，只是移动到对应 token 的新位置。因此，“压力超过限值”和“限值超过压力”得到相同的输出集合。第 6 节 加入位置信息，练习 4 证明这种对称性。第二项对生成而言是时间方向。循环网络无法看到未来，因为步骤 t 的状态只由步骤 1 到 t 构成；注意力层同时看到所有位置，所以预测下一个 token 的模型需要掩码，阻止每个位置看到未来（第 2 节）。

将包含 T 个 token 向量的序列堆叠成矩阵 \mathbf{X} \in \R^{T\times d} 的行。通过可学习矩阵，将每个向量投影到三个不同角色：

其中 \mathbf{W}_Q, \mathbf{W}_K \in \R^{d\times d_k}，\mathbf{W}_V \in \R^{d\times d_v}。第 i 行是 token i 的投影，即 \mathbf{q}_i = \mathbf{x}_i\mathbf{W}_Q，\mathbf{k}_i 和 \mathbf{v}_i 同理。查询（query）表示该位置寻找什么；键（key）表示该位置提供什么供匹配；值（value）表示匹配后传递什么。三个角色不同，所以需要三个矩阵：一个单词寻找的内容（如动词寻找主语）不必与它提供的内容相同，传递的内容也不必与匹配的内容相同。查询和键共享宽度 d_k，以便比较；值可以采用任意宽度 d_v。于是

\mathbf{S} = \mathbf{Q}\mathbf{K}^\top/\sqrt{d_k} 是 T \times T 矩阵，每一对（查询位置，键位置）对应一个分数 S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}。逐行应用 softmax，得到 P_{ij} = \exp(S_{ij})/\sum_k \exp(S_{ik})。因此 \mathbf{P} = \softmax(\mathbf{S}) 的各元素非负，各行之和为 1，每个查询都得到键位置上的概率分布。输出的第 i 行为 \mathbf{o}_i = \sum_j P_{ij}\mathbf{v}_j，是值向量的凸组合，因而位于它们的凸包内（第 3 节 给出图示）。图 6.2 标出了各张量的形状。

逐行应用 softmax 雅可比矩阵

注意力通过 softmax 学习，因此其导数决定学习速度。第 02 模块，第 3 节 已推导 \mathbf{p} = \softmax(\mathbf{s}) 的雅可比矩阵：

本模块会用到它的三个性质。各行之和为零，\sum_j p_i(\delta_{ij} - p_j) = p_i - p_i = 0，因为对所有分数加上相同常数不会改变 \mathbf{p}。当 \mathbf{p} 接近独热分布时，\mathbf{J} 趋于零：每个元素都包含趋于零的因子（p_i、p_j 或 1 - p_i）。当 \mathbf{p} 分布较均匀时，\mathbf{J} 较大：它的迹 1 - \sum_i p_i^2 在均匀分布处达到最大值。

一行的梯度。 设一行的输出为 \mathbf{o} = \sum_j p_j\mathbf{v}_j，\mathcal{L} 是由该输出计算的任意标量。由 \partial\mathbf{o}/\partial p_j = \mathbf{v}_j 可得，各权重的梯度为 g_j = \partial\mathcal{L}/\partial p_j = (\partial\mathcal{L}/\partial\mathbf{o})\cdot\mathbf{v}_j。通过 \mathbf{J} 应用链式法则，得到

g_j 表示将权重移向值 j 时损失的变化率，\bar g 是当前混合下的平均变化率。若某键的值比当前加权平均更有利于降低损失（g_j < \bar g），梯度下降就提高该键的分数；反之则降低分数。变化量与 p_j 成比例，所以权重已经很小的键几乎不会移动。分数通过 S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k} 将梯度传给查询和键；第 3 节 用具体数字演示这一过程。

为什么分数会被缩放

假设 \mathbf{q} 和 \mathbf{k} 的条目是独立的，均值为 0，方差为 1。得分 \mathbf{q}\cdot\mathbf{k} = \sum_{i=1}^{d_k} q_i k_i 是 d_k 项的总和。因为 q_i 和 k_i 是独立的，所以它的均值是

项 q_ik_i 彼此独立，因此它们的方差相加，并且每一项的均值为 0，因此其方差是其二阶矩，通过独立性进行因式分解：

标准差为 \sqrt{d_k}；在常见的头宽度 d_k = 128 下，约为 11.3。除以 \sqrt{d_k} 后，无论头宽度如何，都恢复单位方差：\operatorname{Var}(\mathbf{q}\cdot\mathbf{k}/\sqrt{d_k}) = d_k/d_k = 1。

这一假设描述的是初始化状态：归一化的输入经过保留方差的投影初始化（第 02 模块，第 6 节）后，各元素的方差大致为 1。训练可以改变尺度：增大 \mathbf{W}_Q 和 \mathbf{W}_K 会使注意力分布更集中，模型可以据此获益。因此，缩放因子设定的是 softmax 的初始温度，并不限制后续学习。

没有它会出现什么问题。 d_k = 128 处的未缩放分数的标准偏差为 11.3，因此最佳键与其余键之间的差距通常为 10 或更多。然后 softmax 饱和，接近 one-hot，其中 \mathbf{J} 几乎为零。每个到达 \mathbf{W}_Q 和 \mathbf{W}_K 的梯度都会经过 \partial\mathcal{L}/\partial\mathbf{s} = \mathbf{J}\mathbf{g} （\mathbf{J} 是对称的，因此它是它自己的转置），所以在初始化时，无论 \mathbf{g} 是什么，这些梯度都很小：模型以任意的、几乎独热的注意力模式开始，并且只能慢慢地学习改变它。

实验 1 测量了随机行的效果。在 d_k = 128 下，使用 16 个键、5,000 次随机抽样，未缩放 softmax 的最大权重中位数为 0.978，平均熵为 0.28 奈特；缩放后分别为 0.224 和 2.36 奈特，而均匀行的熵为 \ln 16 = 2.77（图 6.3；末位数字随随机抽样变化）。缩放并非唯一措施：一些大型训练还在点积之前归一化 \mathbf{q} 和 \mathbf{k}，即 QK 归一化，直接控制分数尺度；第 08 模块，第 7 节 将它作为稳定性措施讨论。

因果掩码

从左到右生成时，模型不能看到将要预测的内容，因此位置 i 只能关注 j \le i 的位置。在 softmax 之前，将满足 j > i 的条目设为 S_{ij} = -\infty。由于 \exp(-\infty) = 0，未来位置的权重精确为零，并非仅仅很小，也没有梯度流过这些条目。掩码屏蔽对角线上方的 T(T-1)/2 个条目；第 i 行保留 i 个分数。

掩码使训练能够并行进行。位置 i 的输出只依赖 token 1, \dots, i，因此可以用它预测 token i + 1；所有位置同时这样做。对包含 T 个 token 的序列进行一次前向传播，就得到 T 个下一个 token 的预测，各自对应一个损失。第 04 模块，第 10 节 的解码器即使使用教师强制，也必须逐步计算，因为循环依赖仍然存在。

代码中的掩码可以是加性或布尔形式。加性掩码是 T \times T 矩阵：允许关注的位置为 0，禁止关注的位置为 -\infty（或该 dtype 最小的有限值），再将其加到 \mathbf{S} 上。布尔掩码则在允许关注的位置取 True。PyTorch 的 F.scaled_dot_product_attention 接受 attn_mask，也可以用 is_causal=True 自动构造因果掩码，使融合 kernel 跳过被屏蔽的一半计算（第 10 节）。

填充掩码。 不等长序列组成 batch 时，会填充到共同长度 T。填充的键位置必须对所有查询不可见，因此填充掩码是形状为 (B, 1, 1, T) 的布尔张量，在头和查询维度上广播；与因果掩码结合后成为 (B, 1, T, T)。填充的查询位置仍会计算输出，因为张量是矩形的，但这些输出没有意义，必须从损失中排除；通常将对应目标设为损失函数的 ignore_index。

import math, torch, torch.nn.functional as F

B, h, T, dk = 2, 8, 5, 32
q, k, v = (torch.randn(B, h, T, dk) for _ in range(3))
lengths = torch.tensor([5, 3])                         # sequence 2 ends in two pads
key_real = torch.arange(T) < lengths[:, None]          # (B, T)
pad = key_real[:, None, None, :]                       # (B, 1, 1, T)
causal = torch.ones(T, T, dtype=torch.bool).tril()     # (T, T), True = may attend
mask = causal & pad                                    # broadcasts to (B, 1, T, T)

y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)        # fused
s = q @ k.transpose(-2, -1) / math.sqrt(dk)                        # the same, by hand
y_hand = s.masked_fill(~mask, float("-inf")).softmax(dim=-1) @ v

数值安全。 用 \exp(s_j - m)/\sum_k \exp(s_k - m) 计算 softmax，其中 m = \max_k s_k。减去相同常数不改变结果，同时使所有指数的自变量不大于零，避免溢出；第 10 节 展示了不这样做的后果。另一种陷阱是所有键都被屏蔽的行：所有指数均为 0，softmax 出现 0 除以 0。手写实现返回 NaN，并向后续各层传播；某些融合 kernel 返回零（本模块所用 PyTorch 版本如此）；若使用最小有限值构成加性掩码，则所有键得到相同分数，该行悄然变成所有值（包括填充值）的平均。三种行为都不能依赖。因果掩码下的左填充 batch 很容易出现这种情况：开头的填充查询只能看到填充键。应确保每个查询至少有一个有效键，或显式将这些行置零。

取三个 token，d_k = d_v = 2。为清楚展示算术过程，直接给出投影向量，不再从嵌入计算：

分数，以及分数除以 \sqrt 2 = 1.414：

仅出现三个不同的缩放分数，因此三个指数适用于以下所有内容：e^{0} = 1、e^{0.707} = 2.028 和 e^{1.414} = 4.113。

使用因果掩码

使用因果掩码，第 1 行仅看到第 1 列，第 2 行看到第 1-2 列，第 3 行看到全部三列。

token 3 的查询 (1, 1) 与自己的键匹配最好，所以输出偏向自己的值。token 1 则只能复制 \mathbf{v}_1，无论其查询和键是什么。

去掉掩码与缩放

第 1 行和第 2 行的权重相互镜像，前两个条目交换，因为 \mathbf{Q} = \mathbf{K} 和交换 token 1 和 2 交换了它们的查询和键： \mathbf{q}_1\cdot\mathbf{k}_1 = \mathbf{q}_2\cdot\mathbf{k}_2、\mathbf{q}_1\cdot\mathbf{k}_2 = \mathbf{q}_2\cdot\mathbf{k}_1 和 \mathbf{q}_1\cdot\mathbf{k}_3 = \mathbf{q}_2\cdot\mathbf{k}_3。输出不同，因为 \mathbf{v}_1 和 \mathbf{v}_2 不同。第 3 行根本没有改变：掩码不会从最后一行中删除任何内容，而最后一行已经看到了每个键。

去掉缩放后，每个分数差距增大 \sqrt 2 倍（第 3 行中，键 3 比其他键高 1，而非 0.707），权重更集中，输出更靠近最佳匹配的值：第 3 行中 \mathbf{v}_3 的权重从 0.503 升至 0.576，输出从 (1.759, 2.007) 移至 (1.940, 2.152)。即使 d_k = 2 时，效果也可见。d_k = 128 时，未缩放分数的分布范围是 d_k = 2 时的八倍（11.3 对 1.4），导致 第 2 节 所讨论的饱和。

几何形状

每个输出都是该查询可见的值向量的凸组合：权重非负且总和为 1。因此 \mathbf{o}_1 就是 \mathbf{v}_1 本身；\mathbf{o}_2 = 0.330\,\mathbf{v}_1 + 0.670\,\mathbf{v}_2 位于 \mathbf{v}_1 到 \mathbf{v}_2 的线段上，约为全程的三分之二；\mathbf{o}_3 位于三角形 \mathbf{v}_1\mathbf{v}_2\mathbf{v}_3 内，最靠近 \mathbf{v}_3（图 6.4）。分数只决定输出落在该区域的哪个位置。注意力能够混合值，却不能离开它们的凸包：它无法外推、放大值，或产生可见值所不包含的方向。因此，每个 Transformer 块都配有独立变换各位置的前馈网络，以及保留每个 token 自身向量的残差路径（第 5 节）。

手算反向传播

取标量 \mathcal{L} = o_{3,2}，即 token 3 输出的第二个分量（2.007），沿因果、缩放后的计算反向求梯度。一般规则来自 \mathbf{O} = \mathbf{P}\mathbf{V} 的逐行推导，以及 第 2 节 和 S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}：

其中 g_{ij} = (\partial\mathcal{L}/\partial\mathbf{o}_i)\cdot\mathbf{v}_j、\bar g_i = \sum_k P_{ik}g_{ik}。被屏蔽条目的 P_{ij} = 0，因此不接收梯度。这里 \partial\mathcal{L}/\partial\mathbf{O} 只有第 3 行第 2 列为 1，其余为零，所以只有第 3 行贡献梯度。

数字说明了公式所承诺的内容。提高 s_{33} 将权重移向 \mathbf{v}_3，其第二个分量 (3) 高于当前 2.007；提高 s_{31} 将其移向 \mathbf{v}_1，其第二个分量为 0；键 2 的值 2 几乎恰好位于当前平均值，因此它的分数几乎不重要。在查询中，提高第二个分量会增加与键 2 和 3 的匹配，其值具有较大的第二分量，但代价是键 1：梯度 0.352。提高第一个分量会增加与键 1 和 3 的匹配，其拉力（-0.498 和 +0.500）几乎完全抵消：梯度 0.001。 实验 1 通过有限差异和 PyTorch 的 autograd 确认每个数字。

这就是注意力的核心：可学习的软查找。其余设计是在组织多个这样的计算：并行的头（第 4 节）、与前馈网络交替组成的块（第 5 节）、位置信息（第 6 节），以及分块计算（第 10 节）。

单个注意力头为每个查询提供一种键位置上的概率分布，捕获一种关系。多头注意力（multi-head attention）并行计算 h 个头，每个头有独立的投影，将输入映射到较窄的空间 d_k = d_v = d/h；然后拼接各头的输出，再用另一个矩阵混合：

其中 \mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)}, \mathbf{W}_V^{(i)} \in \R^{d\times d_k}，每个 \text{head}_i \in \R^{T\times d_k}，分号表示沿特征维度拼接，\mathbf{W}_O \in \R^{d\times d}。若 d = 4{,}096、h = 32，每个头就在 128 维空间中工作。

为什么使用多个头。 一个 softmax 只有总量为 1 的权重可分配。如果某位置需要两处信息（用于语法判断的前一个单词、此前提到的名字），单个头只能返回二者的加权平均；第 3 节 表明，平均值位于两者之间，并不等于任何一个值。使用 h 个头后，该位置可因 h 种不同原因关注 h 个位置。成本不增加：由于 hd_k = d，无论 h 取何值，投影参数共 d \times d，分数计算共 h \cdot T^2 \cdot d_k = T^2 d 次乘加，与一个全宽头相同。区别是每个头在 d_k 维而非 d 维空间中比较查询和键。

形状，一步一步

代码不会为每个头创建独立对象。一个 d \times d 矩阵同时计算所有头的查询（各 \mathbf{W}_Q^{(i)} 是它的列块），随后 reshape 将结果拆分到各头，头索引成为一个 batch 维度。下面的显式实现逐行对应 第 12 节 的注意力类，只是直接构造分数，而非调用融合的 F.scaled_dot_product_attention：

import math, torch, torch.nn as nn

B, T, d, h = 2, 16, 256, 8
dk = d // h                                                    # 32
x = torch.randn(B, T, d)
wq, wk, wv, wo = (nn.Linear(d, d, bias=False) for _ in range(4))

q = wq(x).view(B, T, h, dk).transpose(1, 2)                    # (B, h, T, dk)
k = wk(x).view(B, T, h, dk).transpose(1, 2)
v = wv(x).view(B, T, h, dk).transpose(1, 2)
s = q @ k.transpose(-2, -1) / math.sqrt(dk)                    # (B, h, T, T)
future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
p = s.masked_fill(future, float("-inf")).softmax(dim=-1)       # (B, h, T, T)
y = (p @ v).transpose(1, 2)                                    # (B, T, h, dk), not contiguous
out = wo(y.reshape(B, T, d))                                   # (B, T, d)

逐步，对于 B 序列的 batch：

• wq(x) 将 (B, T, d) 映射到 (B, T, d)：所有头的查询并排。
• .view(B, T, h, dk) 将最后一个维度分割为 h 个 d_k 块；没有数据移动。
• .transpose(1, 2) 给出 (B, h, T, d_k)：每个头现在是一个单独的 batch 条目。
• q @ k.transpose(-2, -1) 给出分数，(B, h, T, T)。矩阵乘法将每个前导维度视为 batch，因此所有 B \times h 分数矩阵都来自一次批量乘法。
• p @ v 给出每个头的输出，(B, h, T, d_k)。
• .transpose(1, 2) 返回到 (B, T, h, d_k)，并且 .reshape(B, T, d) 连接头。它必须是 reshape，而不是 view：转置后，内存仍然是逐个头布局的，并且合并头尺寸违反了 view 的步幅兼容性条件（此处引发 RuntimeError）。非连续张量仍然可以支持其他兼容的视图。 reshape 在可能的情况下返回视图，并在必要时进行复制。
• wo 将 (B, T, d) 映射到 (B, T, d)。

连接，然后投影：每个头写入的总和

将 \mathbf{W}_O 拆分为 h 个 d_k 连续行的块，\mathbf{W}_O^{(1)}, \dots, \mathbf{W}_O^{(h)}，每个 d_k \times d。块矩阵乘法给出

因为 \text{head}_i 的元素只乘以 \mathbf{W}_O 中与之对应的行。拼接只是组织数据的方式：每个头通过自己的切片 \mathbf{W}_O^{(i)}，将 d_k 维结果写入 d 维输出，然后将各头的写入相加。

每个头的两个低秩电路

根据层的输入，写入位置 t 处的查询和位置 s 处的键之间头 i 的分数：

这是关于两个输入的双线性形式，矩阵 \mathbf{W}_Q^{(i)}\mathbf{W}_K^{(i)\top} 的形状为 d \times d。由于中间经过 d_k 维空间，其秩至多为 d_k。Elhage 等（2021）称其为头的 QK 电路，它决定该头关注哪里。同理，头对位置 t 的写入为

即将加权输入通过形状为 d \times d 的矩阵 \mathbf{W}_V^{(i)}\mathbf{W}_O^{(i)}，其秩至多为 d_k。这就是 OV 电路，它决定关注之后写入什么。因此，每个头从输入的 d_k 维子空间读取，并写入输出的 d_k 维子空间：d = 4{,}096、h = 32 时，只有 4,096 维中的 128 维。两个电路独立，所以头可根据 token 的一种属性（如位置）选择关注位置，再复制另一种属性（如身份）；下文的归纳头正是如此。

计数

四个投影中的每一个（所有头的 \mathbf{W}_Q、\mathbf{W}_K 和 \mathbf{W}_V 以及 \mathbf{W}_O）均为 d \times d，因此注意力层具有 4d^2 参数，无论 h 是什么。对于每个 token，将 d 向量乘以 d \times d 矩阵需要 d^2 次乘加，或 2d^2 次浮点运算（各一次乘法和一次加法），因此四个投影的成本为 8d^2 次浮点运算。上下文 t 处的 token，关注 t 位置，为其 \mathbf{Q}\mathbf{K}^\top 行支付 2td 更多费用（每个 h 头中长度 d_k 的 t 点积， thd_k = td 乘加）和 2td 用于混合值：4td。对 T token 序列求和，即 O(T^2 d)，其长度是二次方，是 Transformer 最著名的成本。 第 11 节 将这些计数转换为该系列使用的 FLOP 约定。

残差流

从一个位置的角度看整个模型。它的向量以 token 嵌入开始，\mathbf{x}_0。每个注意力层和每个前馈网络通过归一化从该向量读取，并将其输出添加回来：

其中 F_l 表示第 l 个子层，可为注意力或前馈网络（注意力还读取其所关注位置的向量）。最后一层之后，将最终向量归一化并线性读出 logits。Elhage 等（2021）称这一持续累加的向量为残差流（residual stream）。展开得 \mathbf{x}_L = \mathbf{x}_0 + \sum_l F_l(\cdot)：最终状态是嵌入加上各层写入的全部更新。

有两个结果决定了如何理解 Transformer。各层仅通过流进行通信：注意力头从它所关注的位置的流中读取数据并将其写入到其自己位置的流中，前馈网络在一个位置上进行读取和写入，并且没有其他内容在层之间传递。流的 d 维度是共享资源：每个层的写入必须适合每个位置相同的 d 数字，以及前面层写入和后面层需要的所有内容。 第 5 节将残差连接写为块的方程，并解释了为什么归一化位于分支而不是流上。

头学习什么

训练后的模型包含行为可辨认的头。前一 token 头在每个位置从 t 关注到 t - 1。对训练模型的分析还发现，有的头关注匹配括号，有的头从单词关注到句子主语。理解这些头的行为，本身就是一个研究领域。

最容易理解的例子是归纳头（induction head）（Olsson 等，2022）：它将“… A B … A”补全为 B。例如读到“P-104 pump … P-104”后预测“pump”。这一电路需要分属两层的两个头。较早层的前一 token 头，将前一个 token 写入各位置的残差流；于是 B 位置包含“我的前一个 token 是 A”。在后面再次出现的 A 处，归纳头的查询寻找“前一个 token 为 A”的位置，并通过 QK 电路匹配第一个头写入 B 位置的键。它关注该位置，再通过 OV 电路将 B 的身份复制到残差流，提高 B 的 logit。第二个头的键依赖第一个头搬运的信息。这一特定电路需要连续两个注意力阶段，但不能据此断言一层无法完成任何复制任务。

归纳头是从上下文复制名称、标识符和重复短语的通用机制。Olsson 等报告，它们在训练早期的一个狭窄窗口中突然形成，损失曲线出现一个凸起，上下文学习也在同一时刻改善。实验 6 在重复序列上训练两层模型，找到这两类头（图 6.7）；第 12 节 的字符级模型学习复制标识符时，也出现类似的突然变化。

谨慎解读注意力图。 注意力图显示头从哪里读取，无法单独说明为什么读取，也无法说明上层如何使用这些信息。Jain 和 Wallace（2019）发现，注意力权重往往与其他输入重要性指标不一致，而且差异很大的注意力模式可以产生相同预测。判断某个头的作用需要干预：将其输出置零，重新测量损失，检查归因于该头的行为是否消失。实验 6 就进行了这种测试。