Attention Mask：Transformer如何通过一个矩阵控制信息流向

在Transformer架构中，注意力机制让每个token都能"看见"序列中的所有其他token。但有时候，我们恰恰需要某些token"看不见"其他token——这正是Attention Mask存在的意义。这个看似简单的矩阵，承担着因果性保证、填充处理、计算优化等多重职责，是理解Transformer工作方式的关键入口。

一个无法回避的设计问题

假设你要训练一个语言模型，任务是预测"那只猫坐在垫子上"这句话中每个词的下一个词。当模型预测"坐"这个词时，它应该只能参考"那只猫"，而不应该看到后面的"在垫子上"——否则训练就变成了作弊。

问题在于，Transformer的注意力机制天然是"全连接"的：每个位置都能 attending to 所有其他位置。这种设计在理解任务中是优势（比如BERT可以同时看到上下文），但在生成任务中却成了致命缺陷。

解决方案看起来很直接：在计算注意力权重时，把不该看到的位置屏蔽掉。但具体怎么做？直接把权重设为0？还是乘以一个掩码矩阵？这里涉及到softmax函数的特殊性质，稍有不慎就会导致数值不稳定。

更复杂的是，实际训练中还有一个并行的需求：批处理时，一个batch里可能有长度不同的句子，短的句子需要用特殊符号填充到相同长度。这些填充位置不应该参与注意力计算——否则模型会从无意义的填充中"学习"到错误的模式。

因果性和填充处理，这两个看似独立的需求，最终都汇聚到了同一个机制上：Attention Mask。

为什么是负无穷大

先看因果掩码的具体实现。对于一个长度为n的序列，我们需要一个n×n的矩阵，其中位置(i,j)表示第i个token是否能attend到第j个token。

如果是双向注意力（比如BERT），这个矩阵全是1。但因果掩码要求：第i个位置只能看到第0到第i-1个位置。这意味着矩阵应该是下三角形的。

关键的问题是：这个矩阵的值应该是什么？

一个直觉的想法是用0和1：1表示可以attend，0表示不能。然后在计算注意力权重后，直接乘以这个掩码矩阵。

但这个方案行不通。注意力权重的计算公式是：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Softmax函数会对每一行做归一化：

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}$$

如果先算softmax再乘以掩码，被屏蔽位置的权重确实变成0了，但剩下的权重之和不再是1，破坏了概率分布的性质。如果先乘以掩码再算softmax，0值在指数运算后变成1，被屏蔽的位置仍然会贡献权重。

正确的方法是：在softmax之前，把需要屏蔽的位置设为负无穷大。

这是因为：

$$\lim_{x \to -\infty} e^x = 0$$

所以：

$$\text{softmax}(-\infty) = \frac{e^{-\infty}}{\sum_j e^{x_j}} = \frac{0}{\text{正数}} = 0$$

被屏蔽的位置权重精确地变成了0，同时softmax的归一化性质得以保持。这就是为什么在PyTorch实现中，你会看到这样的代码：

scores = scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)

在实际工程中，float('-inf')有时会被替换成一个很大的负数，比如-1e9。这两种方式在大多数情况下等价，但后者在某些硬件上数值稳定性更好。

flowchart TB
    A[输入序列] --> B[计算 Q, K, V]
    B --> C[计算注意力分数<br/>Q × K^T / √d]
    C --> D{应用掩码}
    D --> E[被屏蔽位置<br/>设为 -∞]
    D --> F[保留位置<br/>保持原值]
    E --> G[Softmax 归一化]
    F --> G
    G --> H[-∞ 位置权重变为 0]
    H --> I[加权求和<br/>Attention × V]

下三角矩阵的几何意义

因果掩码的经典形式是一个下三角矩阵。对于序列长度为5的情况：

$$M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty & -\infty \\ 0 & 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 & 0 \end{bmatrix}$$

这个矩阵的几何结构精确地编码了"因果性"：第i行从第0列到第i列都是0，第i+1列到最后一列都是负无穷大。

第一行只有第一个位置是0——第一个token只能attend到自己。第二行前两个位置是0——第二个token可以attend到自己和第一个token。依此类推，形成了一个严格的"因果链"。

这个设计有一个重要推论：训练时，模型可以并行处理整个序列。这是Transformer相比RNN的核心优势之一。RNN必须顺序处理，每个时间步依赖前一个时间步的隐藏状态。而Transformer通过因果掩码，可以在一次前向传播中计算所有位置的输出，同时保证每个位置不会"偷看"未来。

但要理解一个关键区别：训练和推理的掩码使用方式不同。训练时，整个序列一次性输入，因果掩码确保信息单向流动。推理时，token逐个生成，每次新增一个token，就需要扩展掩码矩阵——这正是KV Cache优化发挥作用的地方。

填充掩码：批处理的无声助手

因果掩码解决了"未来信息泄露"的问题，但实际训练中还有一个同样棘手的问题：批处理。

一个batch中的句子长度往往不同。为了并行计算，需要把所有句子填充到相同长度。假设一个batch包含三个句子：

“今天天气很好”（长度4）
“我爱自然语言处理”（长度6）
“你好”（长度2）

假设填充到长度6，用<PAD>表示填充符号：

今天天气很好<PAD><PAD>
我爱自然语言处理
你好<PAD><PAD><PAD><PAD>

问题来了：当模型处理"你好"这个短句时，它不应该attend到后面的四个填充符号——这些符号没有语义意义，不应该影响模型的判断。

填充掩码的作用就是标记哪些位置是真实的token，哪些是填充。它的形式通常是：

$$M_{pad} = \begin{bmatrix} 1 & 1 & 1 & 1 & 0 & 0 \\ 1 & 1 & 1 & 1 & 1 & 1 \\ 1 & 1 & 0 & 0 & 0 & 0 \end{bmatrix}$$

其中1表示真实token，0表示填充。在注意力计算中，填充位置同样被设为负无穷大，最终权重变为0。

填充掩码的形状需要注意。对于批处理，掩码形状通常是(batch_size, 1, 1, seq_len)或(batch_size, seq_len)，通过广播机制与注意力分数矩阵相匹配。形状不匹配是初学者最常遇到的错误之一。

还有一个容易被忽视的细节：填充的位置。大多数模型采用右填充（padding_side=‘right’），即填充符号放在句子末尾。但某些场景需要左填充，比如某些生成任务。填充位置会影响掩码的构造方式，也会影响某些位置编码的工作方式。

因果掩码与填充掩码的组合

在自回归生成模型（如GPT）中，因果掩码和填充掩码经常需要同时使用。一个训练样本可能包含不同长度的序列，每个序列既要遵守因果性，又要忽略填充。

组合方式很简单：两个掩码相加。因果掩码限制每个位置只能看过去，填充掩码限制所有位置不能看填充。相加后，同时满足两个约束。

# 组合因果掩码和填充掩码
combined_mask = causal_mask + padding_mask
# 或者使用逻辑运算
combined_mask = torch.minimum(causal_mask, padding_mask)

在PyTorch的MultiheadAttention模块中，这两个掩码通过不同的参数传入：attn_mask用于因果掩码或自定义掩码，key_padding_mask用于填充掩码。模块内部会自动处理它们的组合。

不同注意力模式的掩码设计

Attention Mask的设计直接决定了模型的注意力模式。不同的任务需要不同的注意力模式，也就需要不同的掩码。

双向注意力（BERT风格）：没有因果掩码，每个token可以attend到所有其他token。这种模式适合理解任务，比如文本分类、命名实体识别、问答系统。BERT通过掩码语言模型（MLM）预训练，随机遮盖一些token让模型预测，但注意力本身是全连接的。

因果注意力（GPT风格）：严格的下三角掩码，每个token只能attend到自己之前的token。这种模式适合生成任务，比如文本续写、对话生成。训练时并行计算，推理时自回归生成。

Prefix LM（UL2、T5风格）：一种混合模式。序列分为prefix和target两部分：prefix部分使用双向注意力，target部分使用因果注意力。这种设计让模型在理解给定上下文（prefix）时能充分利用双向信息，在生成回答（target）时保持因果性。

graph LR
    subgraph 双向注意力
        A1[Token 1] <-.-> A2[Token 2]
        A1 <-.-> A3[Token 3]
        A2 <-.-> A3
    end
    
    subgraph 因果注意力
        B1[Token 1] --> B2[Token 2]
        B1 --> B3[Token 3]
        B2 --> B3
    end
    
    subgraph Prefix LM
        C1[Prefix 1] <-.-> C2[Prefix 2]
        C1 <-.-> C3[Target 1]
        C2 <-.-> C3
        C3 --> C4[Target 2]
        C1 --> C4
        C2 --> C4
    end

滑动窗口注意力：每个token只能attend到固定窗口内的邻居。这种设计将计算复杂度从O(n²)降低到O(n·w)，其中w是窗口大小。Mistral等模型使用这种模式处理长序列。掩码在这里不仅是下三角，还是带状的——主对角线附近的一条带。

稀疏注意力：更复杂的掩码模式，比如Longformer和BigBird。它们结合了局部注意力（滑动窗口）和全局注意力（特定token可以attend到所有位置）。这种设计适合处理超长文档，比如整本书或长篇论文。

多头注意力中的掩码广播

多头注意力将输入分割成多个头并行计算。一个自然的问题是：每个头需要单独的掩码吗？

通常不需要。在标准实现中，所有头共享同一个掩码，通过广播机制自动扩展。掩码的形状从(batch_size, seq_len, seq_len)或(seq_len, seq_len)扩展到(batch_size, num_heads, seq_len, seq_len)。

但在某些特殊场景，可能需要每个头有不同的掩码。比如某些多头注意力变体，不同的头关注不同类型的关系。这种情况下，需要显式构造(batch_size, num_heads, seq_len, seq_len)形状的掩码。

PyTorch的MultiheadAttention支持多种掩码形状：

二维掩码(seq_len, seq_len)：所有batch共享同一个掩码
三维掩码(batch_size, seq_len, seq_len)：每个batch有自己的掩码
四维掩码(batch_size*num_heads, seq_len, seq_len)：每个头有自己的掩码（需要手动reshape）

Encoder-Decoder架构中的三层掩码

原始Transformer采用Encoder-Decoder架构，包含三种不同的注意力层，每种使用不同的掩码策略。

Encoder自注意力：双向注意力，无因果掩码。但有填充掩码，忽略batch中的填充token。编码器需要完整理解输入序列，所以每个token都能看到所有其他token。

Decoder自注意力：因果掩码加填充掩码。因果掩码确保生成的token不能"偷看"未来，填充掩码忽略batch中的填充。解码器逐个生成token，必须严格遵守因果性。

Encoder-Decoder交叉注意力：这是连接编码器和解码器的桥梁。Query来自解码器，Key和Value来自编码器。这里只有填充掩码，没有因果掩码。解码器的每个位置都可以attend到编码器输出的所有位置——因为编码器输出已经完整计算好了，不存在"未来信息"的概念。

flowchart TB
    subgraph Encoder
        E1[输入序列] --> E2[自注意力<br/>双向 + 填充掩码]
        E2 --> E3[前馈网络]
        E3 --> E4[编码器输出]
    end
    
    subgraph Decoder
        D1[目标序列] --> D2[自注意力<br/>因果 + 填充掩码]
        D2 --> D3[交叉注意力<br/>仅填充掩码]
        E4 --> D3
        D3 --> D4[前馈网络]
        D4 --> D5[输出]
    end

这个设计反映了Encoder-Decoder架构的核心分工：编码器负责理解输入（双向），解码器负责生成输出（因果），交叉注意力让解码器能够"参考"编码器的理解。

KV Cache与增量生成

在自回归生成中，每次只生成一个新token。理论上，每次生成都需要重新计算整个序列的注意力，时间复杂度O(n²)。但KV Cache技术通过缓存Key和Value，将复杂度降低到O(n)。

KV Cache的工作原理：第一次生成时，计算并存储所有token的Key和Value。后续每次生成新token，只需要：

计算新token的Query、Key、Value
把新的Key和Value追加到缓存中
用新token的Query与所有缓存的Key计算注意力
加权求和得到输出

在这个过程中，因果掩码只需要扩展一行——新增的token只能attend到自己和之前所有token。由于之前token之间的注意力权重已经计算并缓存，不需要重复计算。

# 增量生成时的掩码扩展
# 原始掩码: [seq_len, seq_len]
# 新token的掩码: 新增一行，全为0（可以attend到所有之前token）
new_mask_row = torch.zeros(1, current_seq_len + 1)
causal_mask = torch.cat([causal_mask, new_mask_row.unsqueeze(0)], dim=0)

这就是为什么LLM推理的第一次生成（prefill阶段）比较慢——需要处理整个prompt，计算完整的KV Cache。后续生成（decode阶段）每个token只需要常数时间的计算（相对于已有序列长度）。

Flash Attention中的掩码优化

Flash Attention是近年来最重要的注意力优化技术之一，它通过重新组织计算顺序，将内存访问从O(n²)降低到O(n)。在掩码处理上，Flash Attention也有特殊优化。

传统的掩码实现需要显式构造一个n×n的矩阵，这在长序列场景下开销巨大。对于序列长度100,000的输入，掩码矩阵需要40GB内存（float32）。

Flash Attention针对因果掩码做了专门优化：当is_causal=True时，它不显式构造掩码矩阵，而是在计算过程中隐式应用因果约束。具体来说，在分块计算注意力时，只计算下三角部分的块，跳过上三角部分。

# PyTorch 2.0+ 的高效因果注意力
import torch.nn.functional as F

# 传统方式：显式掩码
# mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
# output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

# Flash Attention方式：隐式因果
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

FlashMask是Flash Attention的扩展，支持更灵活的掩码模式。它采用列稀疏表示，避免显式存储完整的n×n掩码矩阵，同时保持Flash Attention的计算效率。

Attention Sink：第一个token的特殊角色

2023年的一项研究发现了一个有趣的现象：大语言模型在推理时，几乎所有的位置都会对第一个token（通常是BOS token或句首）分配大量注意力权重，形成所谓的"Attention Sink"。

这个现象与因果掩码的设计有微妙的关系。因果掩码要求第一个token只能attend到自己，这迫使第一个token必须承载足够的信息来"解释"自己。在训练过程中，模型学会了将全局信息聚合到第一个token上，让其他位置在需要时可以"参考"它。

StreamingLLM正是利用这个现象，通过保留第一个token的KV Cache，实现了无限长度的流式生成。即使显存有限，只要保留第一个token和最近几个token的KV Cache，模型就能持续生成而不会崩溃。

graph LR
    subgraph 正常因果注意力
        T1[Token 1<br/>Sink] --> T2[Token 2]
        T1 --> T3[Token 3]
        T1 --> T4[Token 4]
        T2 --> T3
        T2 --> T4
        T3 --> T4
    end

这个发现也解释了为什么简单地丢弃早期token会导致模型崩溃：不仅是位置信息的丢失，更是全局语义锚点的消失。

工程实践中的常见陷阱

理论理解之后，实际编码时仍有很多细节容易出错。

掩码形状不匹配是最常见的问题。注意力分数的形状是(batch_size, num_heads, seq_len, seq_len)，但掩码可能是(seq_len, seq_len)、(batch_size, seq_len)或(batch_size, seq_len, seq_len)。形状不匹配会导致广播错误或静默地产生错误结果。

掩码值的选择也有讲究。理论上应该用-inf，但float('-inf')在某些操作中可能产生NaN。工程中常用-1e9或torch.finfo(dtype).min作为替代。选择过小的值（比如-1）会导致掩码不完全生效，softmax后仍有非零权重。

训练和推理的不一致是另一个常见陷阱。训练时使用填充掩码忽略填充token，但推理时如果忘记传入相同的掩码，模型可能会从填充中"读取"错误信息。HuggingFace的tokenizer会自动生成attention_mask，但某些场景下需要手动处理。

左填充和右填充的混淆也会导致问题。大多数预训练模型期望右填充，但某些任务（如某些生成场景）需要左填充。填充位置不同，掩码的构造方式也要相应调整。

# 常见的掩码构造错误

# 错误1：掩码值太小
wrong_mask = torch.zeros(seq_len, seq_len)  # 0不是有效的掩码值
wrong_mask = torch.triu(wrong_mask, diagonal=1) * (-1)  # -1不够小

# 正确的掩码
correct_mask = torch.triu(torch.full((seq_len, seq_len), float('-inf')), diagonal=1)

# 错误2：形状不匹配
# 注意力分数: [batch, heads, seq, seq]
# 掩码: [seq, seq] - 需要扩展维度
mask = correct_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq, seq]

序列打包与文档掩码

为了提高训练效率，一个常见做法是将多个短序列"打包"成一个长序列，而不是用填充符号填充。这种技术叫做Sequence Packing或Document Packing。

假设你有三个短句子：

“你好”（长度2）
“今天天气很好”（长度4）
“早上好”（长度3）

传统方法会填充到长度4，产生大量填充token。而打包方法将它们拼接成一个长度9的序列：“你好今天天气很好早上好”。

但这里有一个关键问题：不同句子之间不应该相互attend。“你好"不应该看到"今天天气很好”。这就需要特殊的文档掩码（Document Mask），确保每个token只能attend到同一文档内的token。

文档掩码的结构比因果掩码更复杂：在整体下三角的基础上，增加了文档边界的约束。不同文档之间即使时间上相邻，也不能相互attend。

graph TD
    subgraph 文档掩码结构
        D1["文档1: 你好"]
        D2["文档2: 今天天气很好"]
        D3["文档3: 早上好"]
        
        D1 -.->|"不能attend"| D2
        D1 -.->|"不能attend"| D3
        D2 -.->|"不能attend"| D3
    end

现代框架如HuggingFace的DataCollatorWithFlattening支持这种打包方式，自动生成相应的文档掩码。这种技术可以将训练效率提升20-30%，因为减少了填充符号的计算开销。

掩码的内存开销与优化

对于长序列模型，掩码本身可能成为内存瓶颈。一个长度为100K的序列，其因果掩码需要10^10个元素，约40GB内存（float32）。

幸运的是，因果掩码有特殊的优化空间。由于其结构是固定的下三角矩阵，可以：

动态生成：不存储完整矩阵，而是在需要时动态计算i >= j判断
稀疏表示：只存储非零元素的位置，或使用块稀疏表示
融合计算：在Flash Attention等优化实现中，因果约束直接融入计算核心

滑动窗口注意力的掩码更紧凑：只有窗口内的位置需要存储，内存复杂度从O(n²)降低到O(n·w)。

对于填充掩码，通常不需要完整的二维矩阵，一维掩码(batch_size, seq_len)就足够了——它可以广播到所有Query位置。

训练与推理：掩码的不同生命周期

掩码在训练和推理阶段的使用方式存在重要差异，这种差异往往被忽视。

训练阶段，掩码服务于两个目的：因果性和填充处理。整个序列一次性输入，因果掩码确保自回归特性，填充掩码处理变长序列。训练时的并行性是Transformer的核心优势——GPU可以同时计算所有位置的损失，通过因果掩码保证训练的有效性。

推理阶段，场景更加复杂。对于批量推理，填充掩码仍然需要，因为batch中可能有不同长度的prompt。但对于单个请求的流式生成，每次只生成一个token，因果掩码动态扩展。

flowchart LR
    subgraph 训练阶段
        TR1[完整序列输入] --> TR2[应用因果掩码]
        TR2 --> TR3[应用填充掩码]
        TR3 --> TR4[并行计算所有位置]
        TR4 --> TR5[计算总损失]
    end
    
    subgraph 推理阶段
        IN1[Prompt输入] --> IN2[Prefill: 完整计算]
        IN2 --> IN3[Decode: 增量生成]
        IN3 --> IN4[扩展掩码一行]
        IN4 --> IN5[计算新token]
        IN5 --> IN6{是否结束?}
        IN6 -->|否| IN3
        IN6 -->|是| IN7[输出结果]
    end

关键区别在于：训练时掩码形状固定，推理时掩码动态增长。这种差异在实现KV Cache时尤其重要——缓存的Key和Value不断增长，但因果掩码只需要新增一行，不需要重新计算整个矩阵。

总结：掩码作为注意力的"交通规则"

回到最初的问题：为什么Attention Mask如此重要？

本质上，注意力机制是一种信息路由——让每个位置决定从其他位置获取多少信息。如果没有约束，信息可以自由流动，这在某些场景（如理解任务）是优势，但在其他场景（如生成任务）是缺陷。

掩码就是这套"交通规则"。它规定了哪些方向是单行道（因果），哪些位置是禁区（填充），哪些区域需要特别关注（全局token）。通过精心设计的掩码模式，同样的注意力机制可以适应完全不同的任务需求。

从工程角度看，掩码处理是Transformer实现中最容易出错的环节之一。形状不匹配、值的选择、训练推理不一致——每一个细节都可能导致难以察觉的问题。理解掩码的数学原理，不仅能帮助调试，更能为选择合适的注意力模式提供理论基础。

从研究角度看，掩码设计仍然是活跃的研究领域。如何设计更高效的掩码模式处理超长序列？如何让模型自动学习最优的掩码模式？这些问题的答案，可能会定义下一代注意力机制的形态。

参考资料

Vaswani, A., et al. “Attention Is All You Need.” NeurIPS 2017.
“A Gentle Introduction to Attention Masking in Transformer Models.” MachineLearningMastery, 2026.
“Masked and Causal Attention.” Abhik Sarkar, 2025.
“Why do LLMs attend to the first token?” arXiv:2504.02732, 2025.
“FlashMask: Efficient and Rich Mask Extension of FlashAttention.” arXiv:2410.01359, 2024.
“Efficient LLM Pretraining: Packed Sequences and Masked Attention.” HuggingFace Blog, 2024.
PyTorch Documentation: torch.nn.MultiheadAttention and F.scaled_dot_product_attention.
“Sliding Window Attention: Efficient Long-Context Modeling.” DigitalOcean, 2026.
“UL2: Unifying Language Learning Paradigms.” arXiv:2205.05131, 2023.
“Prefix Linear Attention Can Outspeed Causal Linear Attention.” Hailey Schoelkopf, 2024.

一个无法回避的设计问题#

为什么是负无穷大#

下三角矩阵的几何意义#

填充掩码：批处理的无声助手#

因果掩码与填充掩码的组合#

不同注意力模式的掩码设计#

多头注意力中的掩码广播#

Encoder-Decoder架构中的三层掩码#

KV Cache与增量生成#

Flash Attention中的掩码优化#

Attention Sink：第一个token的特殊角色#

工程实践中的常见陷阱#

序列打包与文档掩码#

掩码的内存开销与优化#

训练与推理：掩码的不同生命周期#

总结：掩码作为注意力的"交通规则"#

更多精彩内容

KV Cache：为什么这个"缓存"决定了大模型推理的速度和成本

序列长度增加一倍，推理时间翻四倍？Transformer注意力复杂度的技术真相

从输入文本到输出：大模型推理的完整流程解析

变长序列处理：大模型如何应对长短不一的输入

相对位置偏置如何改变Transformer的序列理解能力：从Shaw到ALiBi的七年技术演进