跳转到内容
理解大语言模型

第 16 章 掩码与多头注意力

输入是 red key .，red 上方的输出要预测 key。第 15 章却允许这个位置直接读取 key 自己的表示。训练可能因此奖励一条捷径，而生成新 token 时根本没有这条捷径可走。即使把验证数据另存一份文件，只要验证时仍允许偷看未来，问题就没有解决。

我们需要先规定每次预测时有哪些信息，再在混合值向量之前落实这一规定。可见性正确以后，另一个限制才值得处理：一个头接收的所有值坐标，都共用一套位置权重。如果并行计算几套各自投影的混合结果，就能增加选择的余地。本章把这两步都算出来，位置顺序和行向量约定继续沿用第 15 章。

16.1 先规定可见关系，再选掩码写法

Section titled “16.1 先规定可见关系，再选掩码写法”

例 16.1：限制同一张三乘三分数表。 复用第 15 章的 X、Q、K、V 和缩放后的分数，记 s=1/2s=1/\sqrt{2}：

S=[s0sss2s2ss3s].S=\begin{bmatrix}s&0&s\\s&s&2s\\2s&s&3s\end{bmatrix}.

行仍是接收查询，列仍是提供信息的键。预测下一个 token 时，位置 1 知道自己的当前输入，还不知道未来输入；位置 2 知道输入 1—2；位置 3 知道输入 1—3。这叫 因果可见性（causal visibility）。允许读取当前输入，是因为目标是 下一个 token。对于 [BOS,red,key,.]，各位置对应的目标是 [red,key,.,EOS]。目标已经错开一位，不能靠删除对角线来补救没有错位的目标。

一批样本中还可能有长短不同的记录。如果第三个槽位是补齐长度用的占位符，它就不代表真实来源，不能向有效查询提供内容。这需要 填充可见性（padding visibility） 规则。表 16.1 先把关系说清楚，不依赖某个 API 用 0、1、True 还是 False 表示。

表 16.1：每个查询允许读取哪些键位置。 填充示例中，只有前两个输入是真实内容，末尾一个槽位是占位符。

查询槽位仅限制未来仅排除填充键同时施加两项限制该查询是否参与损失？
111、21有待评分目标时参与
21、21、21、2有待评分目标时参与
31、2、31、21、2不参与：该查询是填充位置

最后一列针对填充示例。排除填充键，管的是查询能读什么；排除填充查询，管的是是否使用该位置的输出；损失掩码管的是哪些目标需要评分。三者有关联，却不能互相代替。只从损失中排除填充目标，并不能阻止双向注意力里的真实查询读取填充键。

16.2 归一化以前，先排除不允许的来源

Section titled “16.2 归一化以前，先排除不允许的来源”

把可见关系表示为计算中的规则，就得到了 掩码（mask）。本书函数约定 visible=True 表示允许。也可以用加法表达：允许的位置加零，禁止的位置加负无穷。负无穷的作用是让取指数后的结果为零，并非“虽然很不受欢迎，但仍有可能被用到”。

只施加因果限制时，掩码后的分数和相应权重是：

Scausal=[s−∞−∞ss−∞2ss3s],Acausal≈[1000.50.500.2839950.1400290.575975].S_{\text{causal}}=\begin{bmatrix}s&-\infty&-\infty\\s&s&-\infty\\2s&s&3s\end{bmatrix},\qquad A_{\text{causal}}\approx\begin{bmatrix}1&0&0\\0.5&0.5&0\\0.283995&0.140029&0.575975\end{bmatrix}.

第一行只剩一个非零指数值，所以把全部权重 1 分给它。第二行剩下两个相同的指数值，各自除以两者之和，得到各一半。第三行与第 15 章相同。再乘以没有改动的 V：

Ocausal≈[10111.5759751.432009].O_{\text{causal}}\approx\begin{bmatrix}1&0\\1&1\\1.575975&1.432009\end{bmatrix}.

第二行恰好是 0.5[1,0]+0.5[1,2]=[1,1]0.5[1,0]+0.5[1,2]=[1,1]。原来没有掩码时的 [1.503490,1.503490]，用到了现在禁止读取的第三个位置。如果先做 Softmax，之后 只乘一个 0/1 掩码就结束，余下权重的和通常不再是 1。正确顺序是先屏蔽分数，再针对允许的键重算归一化分母。

查询 1 → 下一个 token
  1. 处理步骤键 1：允许；2：禁止；3：禁止
  2. 一组数字权重 [1, 0, 0] → 输出 [1, 0]

当前输入已知，要预测的是它后面的 token。

查询 2 → 下一个 token
  1. 处理步骤键 1、2：允许；键 3：禁止
  2. 一组数字权重 [0.5, 0.5, 0] → 输出 [1, 1]

归一化分母包含两个允许的键。

查询 3：真实位置还是填充？
  1. 处理步骤真实位置：键 1、2、3 都允许
  2. 处理步骤填充位置：查询不用；键 3 被排除

真实的全遮蔽行会报错；未使用的填充行按明确约定返回零。

图 16.1 — 保留已知输入所在的对角线，禁止通向未来的连接。 本书原创教学图，CC BY-SA 4.0。数值与规则见正文及共享数据。

图 16.1 明确标出了允许的对角线和禁止的未来连接。如果第三个槽位还是填充，就进一步排除这个键。该位置的查询输出按约定置零，也不计入损失；位置 1、2 的输出仍为 [1,0]、[1,1]。这里的零向量只是未使用行的处理约定，不是一行“全零概率”。

如果一条 真实查询 没有任何可见键，会怎样？此时没有可以归一化的混合：所有指数值都是零，分母也是零。若先减最大值，全为负无穷的行还可能在中途出现无定义的运算。因此，实现会抛出 ValueError("a valid query has no visible key")。填充查询走另一条安全处理路径，返回零向量。计算损失时，如果整批都没有待评分目标，也会报错。给真实的全遮蔽行偷偷补一份均匀分布，等于凭空放进了原规则禁止的来源。

16.3 偷看答案，也能得到漂亮的损失

Section titled “16.3 偷看答案，也能得到漂亮的损失”

例 16.2：故意构造一个无效的未来复制器。 固定候选顺序 [red,blue]。第一个输入位置是已知的 BOS，给定值向量 [0.5,0.5]。第二个位置若是 red，给定值为 [1,0]；若是 blue，则为 [0,1]。第一个位置的任务恰好是预测第二个 token。现在故意给它分数行 [0,ln(99)]。这个例子没有做任何参数学习。

如果允许读取两个位置，指数值为 [1,99]，权重为 [0.01,0.99]。对 red 记录，输出是 0.01[0.5,0.5]+0.99[1,0]=[0.995,0.005]0.01[0.5,0.5]+0.99[1,0]=[0.995,0.005]；对 blue 记录，则是 [0.005,0.995]。这里特意把 V 的各行设成候选分布，混合后才可以直接评分；普通注意力的值向量并不是词表概率。这样构造，是为了直接暴露泄漏，不把它藏进语言模型的输出层。

表 16.2：答案可见时，不学习也能获得很低的损失。 损失是正确候选概率的负自然对数，单位是每次首位置预测的奈特数（nat）。

可见性red 记录的正确标签概率blue 记录的正确标签概率平均损失
未来可见0.9950.9950.005013
因果限制，只能看 BOS0.50.50.693147

即使换成另存的留出记录，这个复制器也会觉得很容易：未来输入里依然放着待预测的答案。低损失只说明我们让答案沿着不允许的路径流了进来，既没有证明预测器有用，也没有证明它会泛化。

一种有效的实现检查是：固定前缀，只改未来输入。在参数固定、没有训练模式随机性的条件下，正确因果注意力为这个前缀算出的全部 logits 都应保持不变。配套测试还会故意去掉或反转三角掩码，并要求所用测试输入的前缀输出发生变化。错误版本违反了应有的不变量，才能说明测试确实抓得住这条泄漏路径；单靠 shape 检查，错误版本也能过关。

16.4 为什么一套混合比例可能不够

Section titled “16.4 为什么一套混合比例可能不够”

一个头接收的每个值坐标，都共用同一行位置权重。不同的输入坐标组合却可能需要不同的比较方式和混合比例。为每一组分别设置 Q、K、V 投影，就可以并行完成多次注意力计算。每一次称为一个 头（head），合起来就是 多头注意力（multi-head attention）。头由它的参数和计算定义，并没有预先承诺某种人能命名的专长。

例 16.3：两个小头，再用输出投影组合它们。 仍用第 15 章 X 中的三个 [2] 向量。这次不再用一个宽度为 2 的头，而是用两个宽度为 1 的头。第一个头的三个投影矩阵都取列向量 [1,0]；第二个头都取 [0,1]：

WQ,1=WK,1=WV,1=[10],WQ,2=WK,2=WV,2=[01].W_{Q,1}=W_{K,1}=W_{V,1}=\begin{bmatrix}1\\0\end{bmatrix},\qquad W_{Q,2}=W_{K,2}=W_{V,2}=\begin{bmatrix}0\\1\end{bmatrix}.

因此，第一个头的 Q、K、V 都是列向量 [1,0,1]，第二个头的三者都是 [0,1,1]。投影矩阵是 [2,1]，投影结果是 [3,1]，分数矩阵各为 [3,3]。缩放分母是 1=1\sqrt{1}=1。两个头使用相同的因果三角关系，但分别计算分数：

D1=[101000101],D2=[000011011].D_1=\begin{bmatrix}1&0&1\\0&0&0\\1&0&1\end{bmatrix},\qquad D_2=\begin{bmatrix}0&0&0\\0&1&1\\0&1&1\end{bmatrix}.

看第二条查询。头 1 允许的分数为 [0,0]，所以输出是 (1+0)/2=0.5(1+0)/2=0.5。头 2 允许的分数为 [0,1]，权重为 [1,e]/(1+e)[1,e]/(1+e)，输出为 e/(1+e)≈0.731059e/(1+e)\approx0.731059。到第三个位置，两个头的输出恰好都为 2e/(2e+1)≈0.8446382e/(2e+1)\approx0.844638，但权重分配给的位置不同：头 1 是 [e,1,e]/(2e+1)，头 2 是 [1,e,e]/(2e+1)。输出相等，不代表读取方式相同。

把头 1 的输出放在前面，头 2 的放在后面。这一步叫 拼接（concatenation），它把坐标连起来，并不求平均。然后再乘一份保存的输出矩阵：

J≈[100.50.7310590.8446380.844638],WO=[111−1],Y=JWO≈[111.231059−0.2310591.6892750].J\approx\begin{bmatrix}1&0\\0.5&0.731059\\0.844638&0.844638\end{bmatrix},\qquad W_O=\begin{bmatrix}1&1\\1&-1\end{bmatrix},\qquad Y=JW_O\approx\begin{bmatrix}1&1\\1.231059&-0.231059\\1.689275&0\end{bmatrix}.

J 是 [3,2]，W_O 是 [2,2]，Y 是 [3,2]。第二行输出的第一个坐标把两个头的结果相加，第二个坐标则相减。输出投影使下一层能使用不同头坐标的组合。它的数值通常也像输入投影一样通过学习得到；本例为便于核对，全部直接给定。

头 1，宽度 1
  1. 模型保存的数值W_Q = W_K = W_V = [1, 0]ᵀ
  2. 一组数字允许分数 [0, 0] → [0.5, 0.5]
  3. 处理步骤0.5×1 + 0.5×0 = 0.5

三个位置都要计算；此处只追踪位置 2。

头 2，宽度 1
  1. 模型保存的数值W_Q = W_K = W_V = [0, 1]ᵀ
  2. 一组数字允许分数 [0, 1] → [0.268941, 0.731059]
  3. 处理步骤0.268941×0 + 0.731059×1 = 0.731059

允许的键相同，用于比较的坐标不同。

先拼接，再投影
  1. 一组数字J₂ = [0.5, 0.731059]
  2. 模型保存的数值W_O 各行：[1, 1]；[1, -1]
  3. 一组数字Y₂ = [1.231059, -0.231059]

拼接顺序必须与 W_O 输入行的含义一致。

图 16.2 — 两份混合结果，组合成一个和与一个差。 本书原创教学图，CC BY-SA 4.0。数值与规则见正文及共享数据。

图 16.2 追踪了位置 2 在两个头中的计算。只换拼接顺序、不改 W_O，就会改变 W_O 各行所接收坐标的含义。如果交换两个头后还想保留原函数，需要把 W_O 对应的行也一起交换。

16.5 在代码里保住每条轴的含义

Section titled “16.5 在代码里保住每条轴的含义”

常见实现一次处理 B 条样本，每条 T 个位置，表示宽度为 C，共有 H 个头，每头宽度为 d，选择 C=HdC=Hd。例如 C=16、H=2 时，d=8。先投影成 [B,T,C]，再 reshape 为 [B,T,H,d]，然后交换 T、H 两条轴，得到 [B,H,T,d]。reshape 在这里重新分组坐标，transpose 则移动轴的位置；两者做的不是同一件事。

只交换 K 的最后两条轴，再与 Q 相乘，得到 [B,H,T,T] 分数。[B,1,T,T] 的可见关系可以广播到各个头：规则共用，分数仍各不相同。权重乘以 V 后，各头输出为 [B,H,T,d]。交换回 [B,T,H,d]，把最后两条轴拼成一条，再乘 [C,C] 的 W_O，输出就是 [B,T,C]。对四维 NumPy 数组直接使用 .T 会把所有轴的顺序反过来，不是这里所需的键转置。

固定 C 时，增加 H 不会自动把总投影宽度放大几倍，因为每个 d 会随之变小。但显式保存的分数数组仍有 BHT2BHT^2 个数。B=1、H=2、T=3 时，共有 18 个分数；若每个 float64 数占 8 字节，单是分数就占 144 字节，还不包括权重、梯度和其他数组。T 翻倍后，有 72 个分数，占 576 字节。这是在已知条件下计算数组大小，不是运行速度或完整训练内存的实测。

code/part-iv/mechanisms.py 中的实践代码把每一步都展开了，masked_softmax 约定 True 表示允许。不能把这个约定不加检查地搬进别的接口：配套验证使用的 PyTorch 版本中，functional.scaled_dot_product_attention 的 True 表示允许；nn.MultiheadAttention 的布尔注意力掩码和填充键掩码却用 True 表示禁止或忽略。具体转换和 CPU 测试与实践代码放在一起；数学上的可见性并不依赖这些接口名称。函数接口、模块接口。

16.6 限制可见性、增加头数，仍没有解决什么

Section titled “16.6 限制可见性、增加头数，仍没有解决什么”

因果掩码把未来位置排除在允许的计算之外，却修复不了已经抄进前面输入字段的答案、受到污染的数据划分，或预处理中使用未来标签的问题。多头允许不同投影和不同混合，但不保证一个头专管颜色、另一个专管语法，也不保证每个头都有用。

我们现在能控制信息交换，输出也回到了原来的模型宽度。但位置本身怎样表示、原输入怎样保留下来、每个位置内部还需要什么计算，都没有交代完整。这些正是第 17 章要补上的 Transformer 块组成部分。

16.7 习题

Section titled “16.7 习题”
  1. 一条四槽位记录为 [BOS,red,key,PAD]，对应目标是 [red,key,EOS,ignored]。在因果限制和填充限制下，列出允许的键位置及待评分目标。第三个槽位的对角线允许吗？
  2. 有人先对第 15 章第一行做 Softmax，再应用因果掩码，得到 [0.401112,0,0]。计算它与 V 的混合，解释为何不同于本章规定的因果结果。
  3. 把例 16.2 指数值比例中的 99 改成 9。计算泄漏条件下正确标签的概率和损失。禁止未来以后，什么结果保持不变？
  4. 交换例 16.3 的两个头，W_O 不动。位置 2 的新输出是什么？怎样修改 W_O 才能恢复原结果？
  5. 某实现使用严格下三角，连对角线也排除了。解释第一条真实查询会遇到什么问题，以及为什么默默返回均匀分布会改变任务。
  6. 一份验证报告只展示 shape 一致、损失下降，没有改动未来输入的测试。指出这两项观察无法排除的一种错误，以及即使未来扰动测试通过也仍存在的一项限制。

16.8 参考解答

Section titled “16.8 参考解答”
  1. 三条真实查询允许的键分别是 [1]、[1,2]、[1,2,3]。第四条查询不使用，按约定返回零向量。red、key、EOS 三个目标计入损失，填充目标不计入。第三个对角线允许：预测 EOS 时已经知道 key。占位符不是要让模型学会预测的一次词表事件。
  2. 因为 V 第一行是 [1,0]，混合结果约为 [0.401112,0]，权重总和只有 0.401112。正确的因果归一化得到 [1,0,0]，输出为 [1,0]。本例也可以把剩余正权重再除以其总和来补救，但先屏蔽分数再做 Softmax，能直接表达所需规则，对数值尺度的处理也更稳妥。
  3. 权重变为 [0.1,0.9]，正确概率是 0.1×0.5+0.9×1=0.950.1\times0.5+0.9\times1=0.95，损失为 −ln⁡(0.95)≈0.051293-\ln(0.95)\approx0.051293。因果限制下仍只能使用 BOS，正确概率为 0.5，损失为 0.693147。削弱不允许的捷径，不会让它变得合法。
  4. 拼接行变为 [0.7310585786300049,0.5]，输出是 [1.2310585786300048,0.2310585786300049]，第二个坐标变号。把 W_O 的两行交换为 [1,-1]、[1,1]，就能补偿这次交换，让同一个头的值仍乘以原来的输出系数。
  5. 第一条查询没有任何允许的键，概率归一化分母为零，所以真实查询检查会报错。均匀权重反而让它读到了原先禁止的位置。对于错开一位的下一个 token 目标，包含对角线才是本章规定的关系；若采用另一种建模方式，需要另行说明输入与可见性。
  6. 因果掩码反向或缺失，都可能保留相同 shape，并靠未来信息得到很低的损失。固定前缀、改动未来输入，能检查相应的前缀不变量；但不能证明标签没有抄进更早的字段，也不能证明数据划分独立，或模型能够泛化到新任务。

16.9 参考文献

Section titled “16.9 参考文献”