最近在看 MIT 6.S191 Lecture 9 讲 Context Parallelism 的时候，碰到了一个词：Causal Attention。
我对 Self-Attention 的公式并不陌生——之前那篇已经把 QKV 和 routing score table 拆得比较清楚了。但 Causal Attention 这个概念我一直没有认真想过。回头看，我写 Attention 公式、写 Multi-Head、写 FlashAttention 的时候，脑子里跑的都是"所有 token 都能看到所有 token"的 Full Self-Attention。
然后一个问题冒出来：如果 Self-Attention 允许每个 token 访问整个 sequence，那在训练语言模型做 next-token prediction 的时候，当前位置岂不是可以直接看到它要预测的那个正确答案？
训练的时候，模型能偷看答案吗#
先把问题具体化。假设训练数据是一句话：
1我 今天 学习 AI之前关于 Next-Token Prediction 的笔记里已经讲过，语言模型的训练目标是在每个位置预测下一个 token：
| 位置 | 输入（已知） | 预测目标 |
|---|---|---|
| 1 | 我 | 今天 |
| 2 | 我 今天 | 学习 |
| 3 | 我 今天 学习 | AI |
位置 1 应该根据「我」预测「今天」，位置 2 应该根据「我 今天」预测「学习」。
但如果用 Full Self-Attention，位置 1 的 Query 可以关注位置 2、3、4 的 Key——也就是说，「我」这个 token 可以直接读取「今天」的信息。而「今天」恰好就是位置 1 的预测目标。
这不是模型"太聪明"的问题。问题出在计算图允许的信息依赖上：如果 architecture 不限制，\(P(x_2 \mid x_1)\) 的计算过程中就可以访问 \(x_2\) 本身。模型拿着答案做预测，训练 loss 可以降得很低，但学到的不是 context 和 next token 之间的统计规律，而是一条捷径——直接把目标位置的信息抄过来。
这就是 Future Information Leakage。它不会在推理时起作用（推理时未来 token 根本还不存在），所以模型在训练时表现很好、推理时表现很差——训练和推理之间出现了根本性的 mismatch。
Causal Mask 到底做了什么#
解决方法其实不复杂。QKV 那篇里，Attention 的完整公式是：
$$ \operatorname{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$这个公式里，\(QK^T\) 是一张 \(n \times n\) 的 routing score table，每个 Query 可以和所有 Key 匹配。Causal Mask 要做的事情就是：在 softmax 之前，对这张表里不该看的位置做手术。
加了 Mask 之后：
$$ \operatorname{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V $$\(M\) 是一个 \(n \times n\) 的 Mask 矩阵。规则很简单：如果位置 \(i\) 允许关注位置 \(j\)（即 \(j \le i\)），\(M_{ij} = 0\)，不影响原来的 score；如果位置 \(i\) 不该看到位置 \(j\)（即 \(j \gt i\)），\(M_{ij} = -\infty\)。
拿 4 个 token 的 sequence 来看，\(M\) 长这样：
1 位置1 位置2 位置3 位置4
2位置1 [ 0 -∞ -∞ -∞ ]
3位置2 [ 0 0 -∞ -∞ ]
4位置3 [ 0 0 0 -∞ ]
5位置4 [ 0 0 0 0 ]下三角和对角线是 0（允许），上三角是 \(-\infty\)（禁止）。
\(-\infty\) 加到 logit 上之后，经过 softmax 的 \(\exp(\cdot)\)，\(\exp(-\infty) = 0\)——这些位置的 attention weight 变成精确的 0。不是很小，是 0。对应的 Value 完全不会被读取。
有一点容易想歪：Causal Mask 不是把某些 token 从 sequence 中删掉了，而是阻止了特定方向的 Query-Key 信息交互。 Token 本身还在那里，它的 embedding、它对其他位置的贡献都没有消失——只是它作为 Key 和 Value 的身份，对某些 Query 是不可见的。位置 3 的 token 在 sequence 中仍然存在，位置 4 可以正常关注它，只是位置 1 和位置 2 看不到它。
顺带区分几个容易混的 Mask。Attention 层里可能出现的 Mask 不止一种：
- Causal Mask：阻止关注未来位置，保证自回归的因果约束。
- Padding Mask：对 batch 中长度不足的 sequence，把 padding token 遮蔽掉，不让它们参与 attention 计算。
- BERT 的 Masked Language Modeling：这不是 attention mask，而是训练策略——随机遮盖部分 token，让模型从上下文预测被遮盖的词。
三者解决的问题完全不同，不要搞混。
Positional Encoding 和 Causal Mask 不是一回事#
我在想 Causal Mask 的时候，一度和 Positional Encoding 搞混过。两者看起来都和"位置"有关——一个管位置编码，一个管可见范围，它们之间是什么关系？
之前写 Positional Encoding 那篇的核心问题是：去掉 RNN 之后，纯 Self-Attention 是 permutation equivariant 的，模型不知道 token 的先后顺序，所以需要把位置信息加回去。
但想清楚之后发现，Positional Encoding 和 Causal Mask 回答的是两个不同的问题：
- Positional Encoding：模型如何知道每个 token 的位置？——给 token 的 representation 里注入位置信息，让模型在计算时能利用「谁在前、谁在后、相隔多远」这些结构。
- Causal Mask：当前 token 允许关注哪些位置？——在 attention score 上设置硬性的可见性边界，位置 \(i\) 只能看到 \(j \le i\)。
一个是提供位置信息，一个是约束信息访问范围。两者的职责不同，不能互相替代。
举个具体例子：假设 Token 5 经过 Positional Encoding 之后，模型知道它处在 Position 5。但"知道自己在 Position 5"并不等于"不能关注 Position 6"。如果没有 Causal Mask，Positional Encoding 只是告诉模型"Position 6 在你后面一格"，但 attention 依然可以把 weight 分配过去。Positional Encoding 提供的是位置感知能力，不是访问控制。
反过来也有一个微妙的地方：Causal Mask 的下三角结构本身会带来一种不对称性——位置 1 只能看到 1 个 token，位置 4 能看到 4 个 token。这种可见 context 长度的差异，确实会让不同位置在计算上扮演不同的角色。但这和 Positional Encoding 提供的位置信息是不同层面的东西。Causal Mask 产生的是一种由可见性约束带来的隐式位置不对称，而不是像 Positional Encoding 那样显式地把位置坐标编码进 representation。
原始 Transformer 的 Encoder 不需要 Causal Mask#
我第二个疑问是：Causal Mask 是后来补上的东西，还是 Transformer 从一开始就有的？
回去翻了 2017 年的 Attention Is All You Need（Vaswani et al.），发现原始 Transformer 的 Decoder 就使用了 Masked Self-Attention。论文第 3.1 节明确写了 “we need to prevent leftward information flow in the decoder to preserve the auto-regressive property”。所以 Causal Mask 不是后来的补丁，而是一开始就在设计里的。
但 Encoder 端不需要它。原因要从原始 Transformer 的 Encoder-Decoder 结构说起。
拿机器翻译举例。Source sentence 是英文 “The cat sat on the mat”，Target sentence 是中文 “那只猫坐在垫子上”。
Encoder Self-Attention：处理 Source Sequence。翻译"cat"的时候，知道后面有"mat"不是信息泄漏——整句英文是完整给定的输入，不是需要预测的目标。Encoder 可以双向读取，让每个 source token 和所有其他 source token 交互。
Decoder Self-Attention：处理 Target Sequence。生成"坐在"的时候，不应该看到"垫子上"——因为后面的词是 Decoder 自己要生成的。这里就需要 Causal Mask。
Decoder Cross-Attention：Decoder 的每个位置去读取 Encoder 的输出。这里 Decoder 可以看到 Encoder 的全部输出——因为 Encoder 输出的是完整的 source 表示，不是 target 的未来信息。
GPT 系列采用的是 Decoder-only 架构，没有 Encoder 也没有 Cross-Attention，只有一个自回归的 Decoder stack。所以整个模型的 Self-Attention 都是 Causal 的。
BERT 走了另一条路：它的训练目标不是自回归地预测下一个 token，而是 Masked Language Modeling——随机遮盖一些 token，从周围的上下文去预测它们。因为预测目标不在 sequence 的"未来"，而是被随机散落在各处，所以 BERT 用的是 Bidirectional Attention，不需要 Causal Mask。这里点到为止，两种架构的取舍不是这篇的主线。
限制了信息流，训练为什么还能并行#
这个问题困扰了我一小会儿。
Causal Mask 限制了每个位置只能看到过去，直觉上似乎意味着计算必须一步步来——先算位置 1，再算位置 2，因为位置 2 要看位置 1 的结果。这不就退回了 RNN 的串行模式吗？
但实际上不是。区别在于：RNN 的串行依赖来自 state 递推——\(h_2\) 的计算需要 \(h_1\) 的输出作为输入。没有 \(h_1\)，\(h_2\) 根本算不出来。这是 computation 层面的依赖。
Transformer 的 Causal Mask 限制的是 information flow——位置 2 只能 attend to 位置 1 和位置 2 自身。但关键是：在训练时，所有 token 的 embedding 从一开始就是已知的。 输入 sequence 是完整给定的，位置 1 的 Key、Value 不需要"先算出来"才能给位置 2 用——它们在 forward 的开头就已经通过 embedding + projection 产生了。
Causal Mask 只是在做矩阵乘法 \(QK^T\) 之后、softmax 之前，把上三角的值遮掉。整个操作仍然是矩阵运算，所有位置可以并行计算。
这里需要区分三种场景，它们的依赖关系不同：
Training（Teacher Forcing）：输入 sequence 完整给定。所有位置的 \(Q, K, V\) 在 forward 开始时就已经准备好。Causal Mask 限制的是 attention weight 的分布，不是 computation order。所有位置的 attention 可以通过一次矩阵运算并行完成。每个位置的预测目标是下一个 token——输入是 \([x_1, x_2, \ldots, x_{T-1}]\)，target 是 \([x_2, x_3, \ldots, x_T]\)，错开一位。正因为整个 sequence 已知，一次 forward 就能同时产生所有位置的预测和 loss。
Prefill（推理阶段处理 Prompt）：和 Training 类似，完整 Prompt 已知，所有位置可以并行计算 attention。
Decode（自回归生成）：每一步只生成一个新 token。新 token 的 Query 需要和之前所有 token 的 Key / Value 做 attention，但"之前所有 token"包括上一步刚生成的那个——而那个 token 在这一步之前并不存在。所以 Decode 阶段不可避免地需要逐 token 生成。
总结一下：Causal Mask 约束的是"允许看到什么"，不是"什么时候可以计算"。 训练时输入完整，因此即使有因果约束，计算依然可以并行。生成时输出尚未产生，串行是因为数据本身的依赖，不是因为 Mask。
Context Parallelism 和 Causal Mask 的关系#
回到最初的起点——我是在看 Context Parallelism 的时候接触到 Causal Attention 的。
Context Parallelism 是把一个很长的 sequence 切分到多个 GPU 上并行处理。最简单的情况：sequence length 是 8，分配给 2 个 GPU。
1GPU 0: Token 1, 2, 3, 4
2GPU 1: Token 5, 6, 7, 8如果是 Full Self-Attention（没有 Causal Mask），两个 GPU 之间是完全对称的——GPU 0 上的 token 需要看到 GPU 1 上的 token，GPU 1 上的 token 也需要看到 GPU 0 上的 token。双向通信，对称负载。
但加上 Causal Mask 之后，情况就变了。
GPU 1 上的 Token 5 需要关注 Token 1–5，其中 Token 1–4 在 GPU 0 上。所以 GPU 1 需要从 GPU 0 拿到 Token 1–4 的 Key 和 Value。
GPU 0 上的 Token 4 只需要关注 Token 1–4，全部在本地。GPU 0 在 forward attention 中不需要 GPU 1 上的任何信息。
这就产生了非对称的数据依赖：在 forward 阶段，GPU 1 依赖 GPU 0 的 K/V，但 GPU 0 不依赖 GPU 1。通信变成了单向的。
需要注意：这里说的是 forward attention 中的因果可见性。完整的训练过程还包括 backward pass，backward 的梯度流向和 forward 的数据依赖不一定对称。 不能直接从"forward 只有单向通信"推导出"整个训练只需要单向通信"。
Attention 的信息依赖关系，会影响 Context Parallelism 的通信方向、通信量和计算调度。 Causal Mask 带来的非对称可见性，使得不同 GPU 上的工作负载和通信需求不再均等——这正是 Context Parallelism 实现中需要认真处理的工程问题。
具体到 Ring Attention 的通信调度、分布式 softmax 的数值一致性、以及如何平衡不同 chunk 之间的计算量，这些适合后续单独写一篇。
写在最后#
Causal Mask 是自回归 Transformer 保证因果可见性约束的核心机制。 它不是后来添加的补丁——2017 年 Attention Is All You Need 的 Decoder 就内置了 Masked Self-Attention。
Positional Encoding 和 Causal Mask 职责不同。 前者回答"模型如何知道位置"，后者回答"当前位置允许看到谁"。位置感知能力不等于访问控制，两者不能互相替代。
Causal Mask 限制的是信息流方向，不是计算顺序。 训练时完整 sequence 已知，所有位置的 attention 可以并行计算；生成时逐 token 串行，是因为输出尚未产生，不是因为 Mask 要求串行。
参考资料#
- Vaswani et al., Attention Is All You Need, 2017 — https://arxiv.org/abs/1706.03762 （特别是 Section 3.1 Encoder-Decoder 结构和 Section 3.2.3 Masked Multi-Head Attention）
- PyTorch
scaled_dot_product_attention文档 — https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html - MIT 6.S191, Introduction to Deep Learning — https://introtodeeplearning.com/
- 从动态读取到 QKV：Self-Attention 的公式是怎么长出来的
- 去掉 RNN 之后，顺序去哪了？
- Next-Token Prediction：语言模型的训练目标从哪来