LLM令牌遮蔽技术详解：从因果掩码到滑动窗口的PyTorch实践

1. 项目概述：为什么LLM需要“看不见”某些词？

在大型语言模型（LLM）的训练和应用中，我们常常希望模型能“选择性失明”——不是真的看不见，而是有策略地忽略输入序列中的某些部分。这种技术就是令牌遮蔽（Token Masking）。乍一听，你可能觉得这和BERT等模型预训练时的掩码语言模型（MLM）任务很像，都是为了学习上下文。但实际上，在LLM（特指GPT这类自回归模型）的语境下，令牌遮蔽的应用场景、技术目标和实现方式都更加多样和精细。

简单来说，LLM的令牌遮蔽核心目标是控制模型的注意力范围。想象一下，你正在写一篇长文，但规定自己不能回头看已经写过的某个段落，或者被要求必须忽略文章里的所有数字。这种限制会迫使你改变思考和生成的方式。对LLM而言，遮蔽技术就是施加这种限制的工具。它能让模型在生成下一个词时，不去“注意”某些特定的、我们不想让它参考的令牌（Token），比如敏感信息、无关上下文、或者答案本身。这直接关系到模型的安全性、可控性、推理能力以及训练效率。

从搜索热词来看，大家关注点很集中：LLM框架（如如何搭建、使用）、具体实现（尤其是PyTorch）、以及高级应用（如RAG、Agent）。这反映出社区已经从“惊叹模型能力”进入到“深入控制与优化模型行为”的实践阶段。掌握几种核心的遮蔽技术，就等于拿到了精细调控LLM的“手术刀”。接下来，我会结合PyTorch，拆解五种最常用、也最具代表性的令牌遮蔽技术，不仅告诉你“怎么做”，更重点剖析“为什么这么做”以及“实践中会遇到什么坑”。

2. 核心需求与场景解析：不止于预训练

在深入代码之前，我们必须厘清：在LLM的哪些环节，我们需要动用遮蔽技术？这决定了我们选择哪种方法。

2.1 训练阶段的因果注意力遮蔽

这是LLM训练的基石。在标准的自回归语言模型训练中（比如训练一个GPT），我们必须确保模型在预测位置i的令牌时，只能看到位置0到i-1的令牌，而不能“偷看”未来的信息。这就是因果遮蔽。它通过一个下三角矩阵（对角线及以下为1，以上为负无穷）来实现，强制注意力机制具有因果性。没有它，模型就失去了“预测未来”的意义，因为答案已经摆在眼前了。

2.2 推理与部署阶段的输入控制

模型训练好后，在推理时我们同样需要遮蔽。例如：

• 防止信息泄露：在对话系统中，当模型生成回复时，我们不应让它看到自己即将生成的回复内容。这通常通过动态扩展的因果掩码来实现。
• 上下文管理：在处理超长文本时（超出模型上下文窗口），我们需要有策略地遮蔽掉一部分历史信息，比如滑动窗口遮蔽，只让模型关注最近N个令牌。
• 安全与合规：主动遮蔽输入中的敏感词或特定实体，防止模型基于这些信息生成不当内容。

2.3 高级任务中的结构化遮蔽

对于一些复杂任务，遮蔽模式不再是简单的三角或窗口，而是根据任务结构定制。

• 填充生成：在文本摘要、翻译等任务中，我们可能先遮蔽掉待生成的部分，让模型根据上下文进行填充。
• 思维链提示：为了引导模型进行分步推理，我们可能会在提示中遮蔽中间推理步骤，让模型自己推导出来。
• 多模态对齐：在视觉-语言模型中，可能需要遮蔽掉文本序列中的某些视觉标记，以研究模态间的依赖关系。

理解这些场景后，我们就能明白，遮蔽不是一个单一的“开关”，而是一套用于塑造模型信息流的“模具”。下面，我们就从最基础的开始，用PyTorch逐一实现。

3. 五种核心令牌遮蔽技术详解与PyTorch实现

在PyTorch中，遮蔽的核心是操作一个与注意力权重矩阵形状相同的布尔张量或浮点数张量（mask），其中被遮蔽的位置（需要忽略的位置）通常被设置为True或一个极大的负值（如-1e9），然后在计算注意力分数后，将这个掩码加到分数上（scores = scores + mask）。被遮蔽位置的分数经过Softmax后会趋近于0，从而使其对应的注意力权重为0。

我们先定义一个通用的注意力函数，以便后续演示：

import torch import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, mask=None): """ 计算缩放点积注意力。 Args: query: [batch_size, num_heads, seq_len_q, head_dim] key: [batch_size, num_heads, seq_len_k, head_dim] value: [batch_size, num_heads, seq_len_v, head_dim] mask: [batch_size, num_heads, seq_len_q, seq_len_k] 或广播兼容的形状。 True/1表示需要遮蔽的位置。 Returns: 注意力输出，注意力权重 """ d_k = query.size(-1) scores = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32)) if mask is not None: # 将布尔掩码转换为分数掩码：True的位置置为负无穷 scores = scores.masked_fill(mask, float('-inf')) attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, value) return output, attn_weights

3.1 基础技术：因果遮蔽

这是自回归模型的命脉，确保生成过程是单向的。

原理与实现： 因果掩码是一个下三角矩阵，形状为[1, 1, target_len, source_len]。对于目标序列中的每个位置i，它只能看到源序列中j <= i的位置。在训练时，target_len和source_len通常是相等的，都是输入序列的长度。

def create_causal_mask(seq_len, device='cpu'): """ 创建因果遮蔽矩阵。 Args: seq_len: 序列长度 Returns: mask: [1, 1, seq_len, seq_len], 下三角为False（不遮蔽），上三角为True（遮蔽） """ # 创建一个上三角矩阵（不包括对角线），作为需要遮蔽的区域 mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() # 调整维度以适配注意力头： [1, 1, seq_len, seq_len] mask = mask.unsqueeze(0).unsqueeze(0) return mask.to(device) # 使用示例 batch_size = 2 num_heads = 4 seq_len = 10 head_dim = 16 query = torch.randn(batch_size, num_heads, seq_len, head_dim) key = value = query # 简化示例 causal_mask = create_causal_mask(seq_len, query.device) # 注意：这个掩码对所有批次和头都是相同的，所以可以广播 output, attn_weights = scaled_dot_product_attention(query, key, value, causal_mask) # 验证：查看第一个批次，第一个头，最后一个查询位置的注意力权重 print(attn_weights[0, 0, -1, :]) # 输出应该只有最后几个位置有非零值（因为因果遮蔽），前面的位置权重应为0。

实操心得与注意事项：

1. 对角线问题：torch.triu(..., diagonal=1)中的diagonal=1是关键。它确保了对角线元素（当前位置自身）不被遮蔽。在一些严格的定义中，预测当前位置时也不应看到当前位置的key，所以使用diagonal=1。如果你希望包含自身，则使用diagonal=0，但这在标准的自回归预测中不常见。
2. 广播机制：我们创建的掩码形状是[1, 1, seq_len, seq_len]。当与形状为[batch_size, num_heads, seq_len, seq_len]的注意力分数相加时，PyTorch的广播机制会自动将其扩展，非常高效。避免为每个批次和头都创建独立的掩码张量。
3. 推理时的动态掩码：在自回归生成（如GPT推理）时，序列是逐步生成的。常见的做法是缓存之前时间步的键值对（KV Cache），并为当前新生成的令牌计算注意力。此时，掩码需要动态扩展。通常我们会维护一个全局的因果掩码，随着生成步骤增加一列一行（新的行对应新的查询，新的列对应新的键）。

3.2 实用技术：填充遮蔽

在实际任务中，批次内的序列长度往往不一致。我们需要将短序列填充到同一长度，并在计算注意力时遮蔽这些填充位置，防止模型从无意义的填充符中学习。

原理与实现： 填充掩码通常是一个二维张量[batch_size, seq_len]，指示哪些位置是真实的令牌（False），哪些是填充的（True）。在注意力中，我们需要将其扩展为四维[batch_size, 1, 1, seq_len]，这样对于每个查询位置，都会遮蔽所有填充的键位置。

def create_padding_mask(padding_indices, seq_len): """ 根据填充索引创建掩码。 Args: padding_indices: 一个列表的列表，每个内层列表是一个序列中填充位置的索引。 例如，batch_size=2, seq_len=5: [[3,4], [4]] 表示第一个序列索引3,4是填充；第二个序列索引4是填充。 seq_len: 序列长度 Returns: mask: [batch_size, 1, 1, seq_len], True对应填充位置。 """ batch_size = len(padding_indices) mask = torch.zeros(batch_size, seq_len, dtype=torch.bool) for i, indices in enumerate(padding_indices): if indices: mask[i, indices] = True # 调整维度: [batch_size, 1, 1, seq_len] # 这样，对于该批次每个序列的所有查询，都会遮蔽这些填充键。 mask = mask.unsqueeze(1).unsqueeze(2) return mask # 更常见的做法是从tokenizer的attention_mask生成 def create_padding_mask_from_attention_mask(attention_mask): """ 从标准的attention_mask（1表示真实token，0表示填充）生成用于遮蔽的mask。 Args: attention_mask: [batch_size, seq_len], 1为真实token，0为填充。 Returns: mask: [batch_size, 1, 1, seq_len], True对应填充位置。 """ # 将attention_mask反转：真实token为0（不遮蔽），填充为1（遮蔽） mask = (attention_mask == 0) mask = mask.unsqueeze(1).unsqueeze(2) return mask # 使用示例 attention_mask = torch.tensor([[1, 1, 1, 0, 0], # 序列1，后两个位置是填充 [1, 1, 1, 1, 0]]) # 序列2，最后一个位置是填充 padding_mask = create_padding_mask_from_attention_mask(attention_mask) print(padding_mask.shape) # torch.Size([2, 1, 1, 5]) print(padding_mask) # 输出：第一个序列的索引3,4为True，第二个序列的索引4为True。

组合使用：在训练中，我们通常需要同时应用因果遮蔽和填充遮蔽。实现方式是将两个掩码用逻辑或（|）合并。

def create_combined_mask(seq_len, attention_mask, device='cpu'): """ 创建用于训练自回归模型的组合掩码（因果+填充）。 """ causal_mask = create_causal_mask(seq_len, device) # [1,1,seq_len,seq_len] padding_mask = create_padding_mask_from_attention_mask(attention_mask).to(device) # [batch,1,1,seq_len] # 扩展因果掩码以匹配批次大小（如果需要） if causal_mask.size(0) != padding_mask.size(0): # 通常因果掩码是[1,1,...]，可以直接广播。但为了逻辑清晰，我们显式扩展。 causal_mask = causal_mask.expand(padding_mask.size(0), -1, -1, -1) # 合并：如果一个位置是填充（True）或者是未来的token（True），则遮蔽。 # 注意：因果掩码中True表示未来（需遮蔽），填充掩码中True表示填充（需遮蔽）。 combined_mask = causal_mask | padding_mask return combined_mask

避坑指南：

1. 掩码类型：确保你的掩码是布尔类型（torch.bool）。如果使用float(‘-inf’)初始化的张量，直接使用逻辑运算可能会出错。
2. 维度对齐：padding_mask的形状是[batch_size, 1, 1, seq_len]，这意味着它对一个批次内所有序列的所有查询位置，遮蔽的键位置是相同的（即所有填充位置）。这是正确的，因为无论查询位置在哪，都不应该去关注填充符。
3. 解码器的交叉注意力：在Seq2Seq模型（如T5、BART）的解码器中，除了自注意力的因果掩码，其交叉注意力（关注编码器输出）通常只需要填充掩码，而不需要因果掩码，因为解码器可以同时看到编码器的所有输出。

3.3 进阶技术：滑动窗口注意力遮蔽

对于超长序列，完全的自注意力计算复杂度是序列长度的平方（O(n²)），无法承受。滑动窗口注意力限制每个令牌只能关注其前后一定窗口大小内的令牌，将复杂度降至O(n * w)，其中w是窗口大小。这是像Longformer、BigBird等模型处理长文本的核心。

原理与实现： 为序列中的每个位置i，创建一个掩码，仅允许关注位置在[i - window_size, i + window_size]范围内的j，同时还要结合因果限制（在自回归模型中，j不能大于i）。

def create_sliding_window_mask(seq_len, window_size, is_causal=True, device='cpu'): """ 创建滑动窗口注意力掩码。 Args: seq_len: 序列长度 window_size: 单侧窗口大小。实际关注范围为 [i-window_size, i+window_size]（如果非因果）。 is_causal: 是否为因果（自回归）模型。如果是，则j不能>i。 Returns: mask: [1, 1, seq_len, seq_len], True表示需要遮蔽的位置。 """ # 创建全1矩阵，然后挖出窗口内的区域（设为0/False） mask = torch.ones(seq_len, seq_len, dtype=torch.bool) for i in range(seq_len): start = max(0, i - window_size) end = i + window_size + 1 if not is_causal else i + 1 # 因果模式下，不能看未来 end = min(end, seq_len) mask[i, start:end] = False # 如果是因果的，还需要确保上三角（j>i）被遮蔽，上面循环中的`end=i+1`已经实现了这一点。 # 但为了通用性，我们可以显式地应用一个因果掩码。 if is_causal: causal_part = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() mask = mask | causal_part # 合并：窗口外或未来的位置都遮蔽 mask = mask.unsqueeze(0).unsqueeze(0).to(device) return mask # 使用示例 seq_len = 15 window_size = 3 sw_mask = create_sliding_window_mask(seq_len, window_size, is_causal=True) print(sw_mask[0,0]) # 查看掩码矩阵 # 你会发现它是一个带状矩阵，只有主对角线附近的带状区域和左下角是False（可关注）。

性能与实现考量：

1. 高效实现：上述循环实现仅用于演示。在实际模型中（如Longformer），会使用更高效的带状矩阵乘法或自定义CUDA内核来实现滑动窗口注意力，避免构建巨大的显式掩码矩阵，尤其当seq_len很大时。
2. 全局注意力：许多滑动窗口注意力模型会为某些特殊位置（如序列开头[CLS]、结尾[SEP]）或用户指定的位置添加“全局注意力”，让这些位置可以看到整个序列，反之亦然。这需要更复杂的掩码逻辑。
3. 结合填充：同样需要与填充掩码结合，确保模型不关注填充符。

3.4 高级技术：随机令牌遮蔽

这直接来源于BERT的MLM任务，但在LLM训练中也有其用途，例如用于数据增强、提高模型鲁棒性，或在T5等“编码器-解码器”架构的预训练中。

原理与实现： 随机选择输入序列中一定比例（如15%）的令牌，将其替换为一个特殊的[MASK]令牌，或者随机替换为其他令牌，然后训练模型预测被遮蔽的原始令牌。在LLM的自回归训练中，我们需要小心地整合这种遮蔽，因为标准的自回归损失是预测下一个令牌，而MLM是预测被遮蔽的任意位置。

def random_token_masking(input_ids, mask_token_id, vocab_size, mask_prob=0.15, replace_prob=0.1, random_token_prob=0.1): """ 对输入序列进行随机令牌遮蔽，遵循BERT的原始策略。 Args: input_ids: [batch_size, seq_len] mask_token_id: 用于替换的[MASK]令牌的id vocab_size: 词表大小，用于随机令牌替换 mask_prob: 被选中进行遮蔽的令牌比例 replace_prob: 在被选中的令牌中，有多大比例直接替换为[MASK] random_token_prob: 在被选中的令牌中，有多大比例替换为随机令牌 Returns: masked_input_ids: 遮蔽后的输入 labels: 用于计算损失的真实标签，未被遮蔽的位置通常设为-100（在CrossEntropyLoss中忽略） """ labels = input_ids.clone() # 创建概率矩阵 prob_matrix = torch.full_like(input_ids, float(mask_prob), dtype=torch.float32) # 决定哪些位置被选中进行遮蔽操作 masked_indices = torch.bernoulli(prob_matrix).bool() # 确保特殊令牌（如[CLS], [SEP]）不被遮蔽，这里假设pad_id=0，也需要排除 # 假设pad_id=0, cls_id=101, sep_id=102 (根据具体tokenizer调整) special_tokens_mask = (input_ids == 0) | (input_ids == 101) | (input_ids == 102) masked_indices.masked_fill_(special_tokens_mask, False) # 将被选中的位置，在labels中保持不变（用于计算损失），在input_ids中进行替换 labels[~masked_indices] = -100 # 忽略未被遮蔽位置的损失 # 对于被遮蔽的位置，决定是替换为[MASK]、随机词还是保持不变 replace_mask = torch.bernoulli(torch.full_like(input_ids, float(replace_prob))).bool() & masked_indices random_mask = torch.bernoulli(torch.full_like(input_ids, float(random_token_prob))).bool() & masked_indices & ~replace_mask keep_mask = masked_indices & ~replace_mask & ~random_mask # 执行替换 masked_input_ids = input_ids.clone() masked_input_ids[replace_mask] = mask_token_id # 替换为[MASK] masked_input_ids[random_mask] = torch.randint(5, vocab_size-5, random_mask.shape, device=input_ids.device) # 替换为随机令牌（避免极特殊id） # keep_mask的位置，input_ids保持不变 return masked_input_ids, labels # 使用示例（假设一个简单的环境） batch_size = 4 seq_len = 20 vocab_size = 30522 mask_token_id = 103 input_ids = torch.randint(100, 30000, (batch_size, seq_len)) masked_input, labels = random_token_masking(input_ids, mask_token_id, vocab_size) print("原始输入片段:", input_ids[0, :8]) print("遮蔽后输入:", masked_input[0, :8]) print("损失标签 (忽略-100):", labels[0, :8]) # 可以看到部分位置被替换为103([MASK])或随机id，labels中对应位置为原始id，其余为-100。

在自回归LLM中的整合： 对于纯解码器LLM，直接应用上述MLM会破坏自回归性质。一种变通方法是“前缀语言模型”或“跨度遮蔽”：随机遮蔽一个连续的令牌跨度，然后让模型自回归地预测这个跨度内的所有令牌。这需要更复杂的掩码生成和损失计算逻辑。

3.5 策略性技术：定制化模式遮蔽

这是最灵活的一类，掩码模式完全由下游任务定义。例如，在文本填充任务中，我们遮蔽掉句子中间的一段；在文档级翻译中，为了保持段落连贯性，可能遮蔽其他段落的信息。

原理与实现： 核心是根据任务规则生成一个二进制矩阵。这里以实现一个简单的“文本中间挖空”任务为例。

def create_span_masking_mask(seq_len, mask_span_start, mask_span_length, device='cpu'): """ 创建遮蔽一个连续区间的掩码。 Args: seq_len: 序列长度 mask_span_start: 遮蔽区间开始索引 mask_span_length: 遮蔽区间长度 Returns: mask: [1, 1, seq_len, seq_len], True表示需要遮蔽的位置。 这个掩码用于自注意力，确保对于所有查询位置，被遮蔽的键位置都不可见。 但更常见的做法是直接修改输入（将对应位置替换为[MASK]），并调整损失函数。 这里展示如何生成一个注意力掩码来“阻止”关注该区间。 """ mask = torch.zeros(seq_len, seq_len, dtype=torch.bool) # 我们想遮蔽掉键（Key）中位于[mask_span_start, mask_span_start+mask_span_length)的位置 # 对于任何查询（Query），都不应关注这些键。 mask[:, mask_span_start:mask_span_start+mask_span_length] = True # 如果是因果模型，还需要叠加因果掩码 causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() mask = mask | causal_mask mask = mask.unsqueeze(0).unsqueeze(0).to(device) return mask # 更实用的：直接生成用于损坏输入的掩码和标签 def create_span_masking_inputs(input_ids, mask_span_start, mask_span_length, mask_token_id): """ 对输入进行区间遮蔽，生成损坏的输入和标签。 Args: input_ids: [batch_size, seq_len] mask_span_start: 开始索引（列表，长度为batch_size） mask_span_length: 遮蔽长度（列表，长度为batch_size） mask_token_id: [MASK] token id Returns: masked_input_ids: 遮蔽后的输入 labels: 标签，遮蔽区间外为-100 """ batch_size, seq_len = input_ids.shape masked_input_ids = input_ids.clone() labels = torch.full_like(input_ids, -100) # 默认全部忽略 for i in range(batch_size): start = mask_span_start[i] length = mask_span_length[i] end = min(start + length, seq_len) # 保存原始标签 labels[i, start:end] = input_ids[i, start:end] # 将输入中的该区间替换为[MASK] masked_input_ids[i, start:end] = mask_token_id return masked_input_ids, labels

应用场景： 这种定制化遮蔽是构建指令微调或特定任务适配数据的关键。例如，为了训练模型完成“根据上文填充下文”的任务，我们可以随机选择文章的一个位置进行遮蔽。在RAG系统中，当模型生成答案时，我们可以遮蔽掉检索到的文档中的某些无关段落，迫使模型更依赖核心证据。

4. 综合应用与避坑实战

理解了单项技术，如何将它们组合并应用到真实的LLM训练或推理流程中？这里以构建一个简单的、支持因果和填充遮蔽的自回归模型训练批次为例。

4.1 完整训练步骤中的掩码集成

假设我们使用Hugging Face的Transformers库中的GPT-2模型。

from transformers import AutoTokenizer, AutoModelForCausalLM import torch model_name = "gpt2" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) # 设置pad_token，如果tokenizer没有的话 if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 准备一个批次数据 texts = ["Hello, how are you?", "I'm fine, thank you. And you?"] inputs = tokenizer(texts, return_tensors='pt', padding=True, truncation=True, max_length=10) input_ids = inputs['input_ids'] attention_mask = inputs['attention_mask'] # 这是标准的attention mask (1 for real tokens) # 模型内部已经实现了因果遮蔽。我们只需要传入attention_mask。 # 在transformers库中，attention_mask会自动被转换为模型需要的格式。 outputs = model(input_ids, attention_mask=attention_mask, labels=input_ids) # 使用labels进行训练 loss = outputs.loss

关键点：transformers库的模型内部已经集成了因果逻辑。我们提供的attention_mask主要用于处理填充。模型内部的forward方法会调用_prepare_decoder_attention_mask函数来合并因果掩码和传入的填充掩码。

4.2 自定义注意力中的掩码处理

如果你想在自己的注意力层中实现这些掩码，一个完整的多头注意力模块可能如下：

import torch.nn as nn import math class MultiHeadAttentionWithMasking(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads self.wq = nn.Linear(d_model, d_model) self.wk = nn.Linear(d_model, d_model) self.wv = nn.Linear(d_model, d_model) self.wo = nn.Linear(d_model, d_model) def forward(self, query, key, value, mask=None): batch_size = query.size(0) # 1. 线性投影并分头 Q = self.wq(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) K = self.wk(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) V = self.wv(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # 2. 计算缩放点积注意力 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: # mask形状应为 [batch_size, 1, 1, seq_len] 或 [batch_size, 1, seq_len_q, seq_len_k] # 确保mask能广播到scores的形状 scores = scores.masked_fill(mask, float('-1e9')) # 使用负无穷遮蔽 attn_weights = F.softmax(scores, dim=-1) context = torch.matmul(attn_weights, V) # 3. 合并多头并输出 context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output = self.wo(context) return output, attn_weights

4.3 常见陷阱与调试技巧

1. 掩码形状错误：这是最常见的问题。务必记住，注意力分数的形状是[batch_size, num_heads, seq_len_q, seq_len_k]。你的掩码必须能广播成这个形状。padding_mask通常为[batch_size, 1, 1, seq_len_k]，causal_mask为[1, 1, seq_len_q, seq_len_k]。使用mask.unsqueeze(1).unsqueeze(2)是增加维度的常用技巧。

2. 遮蔽值的选择：被遮蔽的位置在加到注意力分数上时，应设置为一个很大的负数。float(‘-inf’)是理论上的选择，但在某些硬件或软件环境下可能不稳定。通常使用-1e9或-1e4等足够大的负数。在应用Softmax之前添加掩码。

3. 梯度问题：被遮蔽的位置由于Softmax后权重为0，其梯度也为0。这通常是我们期望的。但要确保你的掩码操作不会意外地阻断需要梯度的路径。

4. 验证掩码有效性：编写简单的测试用例来验证掩码是否正确工作。

   def test_causal_mask(): seq_len = 5 mask = create_causal_mask(seq_len) print("Causal Mask (True=掩碼):") print(mask[0,0]) # 模拟一个均匀分数 scores = torch.zeros(1,1,seq_len,seq_len) scores_masked = scores.masked_fill(mask, float('-inf')) attn = F.softmax(scores_masked, dim=-1) print("注意力权重（最后一行）:") print(attn[0,0,-1]) # 应该只有最后一个元素是1（因为前面都是-inf），验证因果性。

5. 与KV Cache的配合：在自回归生成中，KV Cache极大地提升了效率。此时，因果掩码需要动态更新。通常做法是维护一个全局的掩码，每次生成新token时，在右侧和下侧扩展一行一列（新token不能看到未来的key，未来的query也不能看到它）。

5. 性能优化与高级话题

当序列长度很长时，显式的全尺寸掩码矩阵（[seq_len, seq_len]）会消耗大量内存。例如，seq_len=8192时，一个布尔掩码矩阵就要占用 8192*8192/8 ≈ 8MB 内存（如果是float则更大）。对于批处理和多头，这个开销会倍增。

优化策略：

1. 使用带状矩阵：对于滑动窗口这类稀疏掩码，可以使用PyTorch的带状矩阵函数（如torch.band）或稀疏张量来隐式表示。
2. Flash Attention：现代高效的注意力实现（如Flash Attention）将掩码计算融合到核函数中，避免了在HBM（显存）中实例化庞大的中间矩阵，包括掩码矩阵。如果你的模型支持，优先使用这些优化后的注意力实现。
3. 自定义CUDA内核：对于极其复杂的掩码模式（如BigBird中的随机+窗口+全局注意力），可能需要编写自定义的CUDA内核来实现高效的掩码计算和注意力。

选择哪种技术？

• 训练标准自回归LM：因果遮蔽 + 填充遮蔽是标配。
• 处理长文本：考虑滑动窗口遮蔽（如Longformer）或其变种。
• 数据增强或特定预训练：可尝试随机令牌遮蔽（但要注意与自回归损失的兼容性）。
• 实现特定任务逻辑：使用定制化模式遮蔽。

令牌遮蔽是连接LLM理论能力与实际可控行为的关键桥梁。从确保模型不乱说（因果性），到让它忽略无关信息（填充、窗口），再到引导它完成特定任务（定制掩码），每一种遮蔽技术都对应着一种对模型“注意力”的约束和引导。理解并熟练运用它们，是进行LLM二次开发、模型优化乃至安全对齐的必备技能。在实际编码时，多画图理解掩码矩阵的形状和含义，从小例子开始测试，能帮你避开大多数坑。