- Causal Mask(上三角):挡住未来位置,自回归生成必须遵守时序;
- Padding Mask:挡住PAD 占位符,过滤无效填充位;
- 合并规则:
final_mask = causal_mask | pad_mask只要其中一个掩码标记为
True(需要遮挡),该位置就统一填充-1e9。 -
import torch# 1. 模拟批次输入,0=PAD input_ids = torch.tensor([[1,2,3,0,0], [4,5,0,0,0]]) batch, seq_len = input_ids.shape# ---------------------- # 第一步:生成 Padding Mask # ---------------------- pad_mask = (input_ids == 0) # [B, L] # 扩维到 [B, L, L],适配注意力分数矩阵 [B, L, L] pad_mask = pad_mask.unsqueeze(1).repeat(1, seq_len, 1) # [2,5,5]# ---------------------- # 第二步:生成 Causal 前瞻掩码(上三角) # ---------------------- # [L, L] 上三角,diagonal=1 表示主对角线右侧全部遮挡 causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() # [5,5]# ---------------------- # 第三步:合并掩码 逻辑或 # ---------------------- # 利用广播:[B,L,L] | [L,L] → [B,L,L] full_mask = pad_mask | causal_mask# ---------------------- # 第四步:作用到注意力分数 # ---------------------- attn_score = torch.randn(batch, seq_len, seq_len) # 模拟注意力分数 [B,L,L] attn_score = attn_score.masked_fill(full_mask, -1e9) attn_weight = torch.softmax(attn_score, dim=-1)print(attn_weight)
三、逐位置解释遮挡逻辑
以序列单条样本[1,2,3,0,0](后两位是 PAD)为例:- 未来位置(j > i):被
causal_mask挡住; - PAD 位置(原序列为 0):被
pad_mask挡住; - 既是未来又是 PAD:依然遮挡(或逻辑天然覆盖);
- 只有前文有效 token:正常计算注意力。
四、分模块使用总结
-
Transformer Encoder / BERT 类双向模型只需要
Padding Mask,不用 Causal Mask,因为允许双向互看。 -
Transformer Decoder / GPT/LLaMA 类自回归模型必须两个掩码混合使用,流程固定:生成 Causal Mask → 生成 Padding Mask → 逻辑或合并 → 作用于注意力分数。
五、补充:适配多头注意力的小改动
如果是多头注意力(形状[B, heads, L, L]),只需提前把两个掩码都扩出 head 维度再合并: - 未来位置(j > i):被
n_heads = 8 # Padding Mask 扩维 [B, heads, L, L] pad_mask = pad_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # Causal Mask 扩维 [1, heads, L, L] 广播使用 causal_mask = causal_mask.unsqueeze(0).unsqueeze(0) # 合并 full_mask = pad_mask | causal_mask