什么是注意力中的因果掩码（Causal Masking）？为什么 LLM 需要它？

注意力中的因果掩码Causal Masking是一种机制它把未来的 token 挡在外面让每个 token 只能看见自己以及自己之前的 token永远看不到自己之后的 token。本文将围绕注意力中的因果掩码展开先介绍因果掩码是什么用一个例子说明看见未来 token到底会带来什么问题然后再走一遍实现过程看看掩码注意力是如何阻止模型访问未来 token 的。众所周知注意力层是 Transformer 架构中最重要的组件也是 LLM大语言模型背后的关键构建块。在注意力层内部我们使用因果掩码。它保证一个 token 只能关注自己和过去的 token永远不会关注未来的 token。因果掩码是一种借助掩码矩阵mask matrix来屏蔽注意力分数矩阵中未来位置的机制从而保证每个 token 只关注自己和过去的 token永远不关注未来的 token。正因如此Transformer 的 Decoder 只能向后看而 Encoder 则可以双向读取整个输入。我们用一个例子来理解。取这句话“I love teaching AI”token 位置[ I | love | teaching | AI ] 1 2 3 4下面看看不用因果掩码会发生什么。不用因果掩码会怎样在预测 “love” 时模型可以关注到 “teaching” 和 “AI”。在预测 “teaching” 时它已经能看见 “AI”。这不是我们想要的结果。我们不应该让模型看见未来因为这样做会引入信息泄漏。如果未来 token 可见模型就能抄近路——把它本来应该预测的词直接当作输入。这会让训练表现好得不真实却破坏了语言建模的基本假设只用过去的上下文来预测下一个 token。而在推理阶段未来 token 根本不存在。如果模型在训练时依赖了未来信息那么当它被要求一步步生成文本时就会崩掉。因果掩码保证训练条件和推理条件一致迫使模型去学真实的语言规律而不是背答案。所以允许模型看见未来训练会变简单但生成会变得不正确。再看看用了因果掩码会怎样。用了因果掩码之后在预测 “love” 时模型关注 “I”。在预测 “teaching” 时它关注 “I” 和 “love”。在预测 “AI” 时它关注 “I”、“love” 和 “teaching”。这里我们不允许模型看见未来因此它被迫表现得和推理时完全一样。屏蔽未来 token 是限制一个 token 能关注什么的一种方式。限制一个 token 能往回看多远是另一种方式即滑动窗口注意力Sliding Window Attention只允许 token 关注最近一段固定长度的上下文更早的内容全部屏蔽。接下来看看如何用掩码矩阵实现这一点。因果掩码的实现在注意力层中我们会计算一个注意力分数矩阵。注意力分数是一个数值表示某个 token 与另一个 token 的相关程度。每个 token 都会和其他所有 token 做比较。注意力分数衡量的是应该给每个 token 分配多少关注度。分数越高关注越多。分数越低关注越少。因此每个 token 都会与其余每个 token 算出一个分数最终构成一个二维矩阵行代表 query token列代表 key tokenquery 与 key 向量是如何算出来的见「注意力背后的数学Q、K、V」。由于句子 “I love teaching AI” 包含 4 个 token注意力分数构成一个 4 × 4 的矩阵。每一行对应一个 query token每一列对应一个 key token。矩阵中的每个值表示在做归一化之前一个 token 应该关注另一个 token 多少。注意力分数矩阵示例掩码前I love teaching AI -------------------------------- I | 2.0 | 1.0 | 3.0 | 0.5 | love | 1.5 | 2.5 | 0.5 | 1.0 | teaching | 0.2 | 1.0 | 2.0 | 1.5 | AI | 0.1 | 0.3 | 0.7 | 2.5 |理想情况下掩码应该在 softmax 之前施加。我们先看看跳过掩码、直接做 softmax 会怎样。Softmax 之后注意力权重I love teaching AI -------------------------------- I | 0.23 | 0.09 | 0.63 | 0.05 | love | 0.21 | 0.58 | 0.08 | 0.13 | teaching | 0.08 | 0.17 | 0.47 | 0.28 | AI | 0.07 | 0.08 | 0.12 | 0.73 |在这个阶段每个 token 都能关注其他任何 token包括未来的 token。在预测 “love” 时模型不只关注过去的 token还关注了未来的 token它给 “teaching” 分配了 0.08 的注意力权重、给 “AI” 分配了 0.13这说明信息已经泄漏。在预测 “teaching” 时模型已经用 0.28 的注意力权重关注到了未来的 token “AI”尽管在这个阶段 “AI” 本不该可见。正是因为这个毛病掩码是必需的为此我们使用一个因果掩码矩阵。因果掩码矩阵我们创建一个同样大小的掩码矩阵。这个矩阵规定了哪些位置被允许1、哪些位置被屏蔽0。I love teaching AI -------------------------- I | 1 | 0 | 0 | 0 | love | 1 | 1 | 0 | 0 | teaching | 1 | 1 | 1 | 0 | AI | 1 | 1 | 1 | 1 |这就是我们的掩码矩阵接下来要把它作用到分数矩阵上以屏蔽未来的 token。读下文时请把这个掩码矩阵记在脑子里因为后面就用它来掩码注意力分数矩阵。1→ 允许关注0→ 屏蔽关注未来位置应用掩码之后被屏蔽的位置会被置为 −∞。因果掩码正是通过给这些位置赋 −∞来阻断对未来 token 的访问。稍后我们会明白为什么要用 −∞。注意力分数矩阵示例掩码前I love teaching AI -------------------------------- I | 2.0 | 1.0 | 3.0 | 0.5 | love | 1.5 | 2.5 | 0.5 | 1.0 | teaching | 0.2 | 1.0 | 2.0 | 1.5 | AI | 0.1 | 0.3 | 0.7 | 2.5 |掩码后的注意力分数矩阵用上面的掩码矩阵作用之后I love teaching AI -------------------------------- I | 2.0 | -∞ | -∞ | -∞ | love | 1.5 | 2.5 | -∞ | -∞ | teaching | 0.2 | 1.0 | 2.0 | -∞ | AI | 0.1 | 0.3 | 0.7 | 2.5 |掩码生效时掩码矩阵中标记为 0 的位置其注意力分数被置为 −∞而标记为 1 的位置保留注意力分数矩阵中的原值。为什么要用 −∞因为后面要做 softmax而按 softmax 的数学性质−∞ 经过 softmax 之后恰好变成 0。如果改用 0softmax 仍然会给被屏蔽的位置分配非零概率模型就能部分地关注未来 token因果掩码的目的也就落空了。Softmax 之后注意力权重I love teaching AI -------------------------------- I | 1.00 | 0.00 | 0.00 | 0.00 | love | 0.27 | 0.73 | 0.00 | 0.00 | teaching | 0.11 | 0.24 | 0.65 | 0.00 | AI | 0.07 | 0.08 | 0.12 | 0.73 |现在未来的 token 拿到的关注度是零在预测 “love” 时模型给 “I” 分配 0.27 的注意力权重给自己“love”分配 0.73。它不关注未来的 token “teaching” 和 “AI”两者都拿到 0.00。在预测 “teaching” 时模型关注 “I”0.11和 “love”0.24并给自己“teaching”分配 0.65。它不关注未来的 token “AI”后者注意力权重为 0.00。在预测 “AI” 时模型关注 “I”0.07、“love”0.08、“teaching”0.12并给自己“AI”分配 0.73。注意每一行都仍然给第一个 token “I” 留了一些权重因为它是任何位置都始终能看见的那个 token。这正是开头几个 token 最终会成为注意力沉淀池attention sinks的原因——模型会把那些在别处用不上的注意力统统倒在这里。以上就是因果掩码阻止模型看见未来 token 的方式。常见问题为什么掩码用负无穷而不是 0负无穷经过 softmax 之后恰好变成 0。如果改用 0softmax 仍然会给被屏蔽的位置分配非零概率。模型于是能部分关注未来 token因果掩码就失效了。因果掩码是在 softmax 之前还是之后施加在 softmax 之前。注意力分数矩阵中被屏蔽的位置被置为负无穷随后 softmax 把它们变成 0 注意力权重。如果跳过掩码直接做 softmax每个 token 都能关注未来的 token。推理时未来 token 本来就不存在为什么还需要因果掩码因果掩码让训练条件和推理条件保持一致。没有它模型会在训练时学会抄近路——直接使用它本该预测的词。这样训练指标好得不真实但模型一步步生成文本时就会失败。Transformer 的 Encoder 会用因果掩码吗不会。Transformer 的 Decoder 因为因果掩码只能向后看。Encoder 则双向读取整个输入。因果掩码和滑动窗口注意力有什么区别因果掩码屏蔽未来 token因此每个 token 只关注自己和之前的 token。滑动窗口注意力限制的是一个 token 能往回看多远。两者都在限制 token 能关注什么只是限制的方式不同。