Attention Mask 笔记

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


def make_padding_mask(token_ids: torch.Tensor, pad_token_id: int = 0) -> torch.Tensor:
    """
    生成 Padding Mask。
    输入：token_ids [B, T]
    输出：mask [B, 1, 1, T]，可广播到 [B, num_heads, T_q, T_k]
    值：PAD 位置为 False（屏蔽），非 PAD 位置为 True（保留）
    """
    return (token_ids != pad_token_id).unsqueeze(1).unsqueeze(2)  # [B, 1, 1, T]


def make_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
    """
    生成 Causal Mask（下三角矩阵）。
    输出：mask [1, 1, T, T]，可广播到 [B, num_heads, T, T]
    值：下三角为 True（可以看），上三角为 False（屏蔽未来）
    """
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
                      ).unsqueeze(0).unsqueeze(0)  # [1, 1, T, T]


def make_decoder_mask(token_ids: torch.Tensor, pad_token_id: int = 0) -> torch.Tensor:
    """
    GPT 训练时的完整 Mask = Causal Mask AND Padding Mask。
    两者取交集：既不能看未来，也不能看 PAD。
    """
    B, T = token_ids.shape
    causal = make_causal_mask(T, token_ids.device)           # [1, 1, T, T]
    padding = make_padding_mask(token_ids, pad_token_id)     # [B, 1, 1, T]
    return causal & padding                                   # [B, 1, T, T]


class MultiHeadAttentionWithMask(nn.Module):
    """带 Mask 支持的 MHA"""

    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.d_model = d_model
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        B, T, _ = x.shape
        Q = self.split_heads(self.W_Q(x))  # [B, h, T, d_k]
        K = self.split_heads(self.W_K(x))
        V = self.split_heads(self.W_V(x))

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, h, T, T]

        if mask is not None:
            # mask=False 的位置设为 -inf，softmax 后权重变为 0
            scores = scores.masked_fill(mask == False, float('-inf'))

        weights = F.softmax(scores, dim=-1)
        attended = torch.matmul(weights, V)
        attended = attended.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_O(attended)


# ── 使用示例 ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    d_model = 32
    num_heads = 4
    pad_token_id = 0

    model = MultiHeadAttentionWithMask(d_model=d_model, num_heads=num_heads)
    model.eval()

    # ── 场景1：训练时（GPT/Decoder），有 PAD，需要 Causal + Padding Mask ──
    token_ids_train = torch.tensor([
        [3, 7, 2, 0, 0],   # 句子1，后两个是 PAD
        [5, 1, 8, 4, 6],   # 句子2，无 PAD
    ])
    x_train = torch.randn(2, 5, d_model)
    decoder_mask = make_decoder_mask(token_ids_train, pad_token_id)

    print("=== 训练场景（GPT Decoder）===")
    print(f"token_ids shape: {token_ids_train.shape}")
    print(f"decoder_mask shape: {decoder_mask.shape}")   # [2, 1, 5, 5]
    print("句子1的 mask（下三角 + 屏蔽PAD列）：")
    print(decoder_mask[0, 0].int())
    # [[1, 0, 0, 0, 0],
    #  [1, 1, 0, 0, 0],
    #  [1, 1, 1, 0, 0],   ← PAD列（第4、5列）被屏蔽
    #  [1, 1, 1, 0, 0],   ← PAD行也只能看到有效列
    #  [1, 1, 1, 0, 0]]

    with torch.no_grad():
        output_train = model(x_train, mask=decoder_mask)
    print(f"输出 shape: {output_train.shape}\n")  # [2, 5, 32]

    # ── 场景2：训练时（BERT/Encoder），有 PAD，只需 Padding Mask ──
    token_ids_bert = torch.tensor([
        [3, 7, 2, 0, 0],
        [5, 1, 8, 4, 6],
    ])
    padding_mask = make_padding_mask(token_ids_bert, pad_token_id)  # [2, 1, 1, 5]

    print("=== 训练场景（BERT Encoder）===")
    print(f"padding_mask shape: {padding_mask.shape}")
    print("句子1的 padding_mask：", padding_mask[0, 0, 0].int())  # [1, 1, 1, 0, 0]

    with torch.no_grad():
        output_bert = model(x_train, mask=padding_mask)
    print(f"输出 shape: {output_bert.shape}\n")

    # ── 场景3：推理时，单条输入，无 PAD，mask=None ──
    single_input = torch.randn(1, 8, d_model)  # batch=1，序列长度=8

    print("=== 推理场景（单条，无PAD）===")
    with torch.no_grad():
        output_infer = model(single_input, mask=None)  # 直接传 None
    print(f"输出 shape: {output_infer.shape}")  # [1, 8, 32]

scores.masked_fill(mask == False, float('-inf'))  # True=保留，False=屏蔽
scores.masked_fill(mask == 0,     float('-inf'))  # 1=保留，0=屏蔽（等价）