attention 的mask 的简单实现

• 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）为例：
   
  1. 未来位置（j > i）：被 causal_mask 挡住；
  2. PAD 位置（原序列为 0）：被 pad_mask 挡住；
  3. 既是未来又是 PAD：依然遮挡（或逻辑天然覆盖）；
  4. 只有前文有效 token：正常计算注意力。
   

  ---

   

  四、分模块使用总结

   
  1. Transformer Encoder / BERT 类双向模型

      

     只需要 Padding Mask，不用 Causal Mask，因为允许双向互看。
      
  2. Transformer Decoder / GPT/LLaMA 类自回归模型

      

     必须两个掩码混合使用，流程固定：

      

     生成 Causal Mask → 生成 Padding Mask → 逻辑或合并 → 作用于注意力分数。
      
   

  ---

   

  五、补充：适配多头注意力的小改动

   
  如果是多头注意力（形状 [B, heads, L, L]），只需提前把两个掩码都扩出 head 维度再合并：

　　

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