🔄 Transformer 架构深度拆解

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

class SelfAttention(nn.Module):
    """
    自注意力机制的 PyTorch 实现

    数学公式: Attention(Q, K, V) = softmax(QK^T / √d_k) V

    参数:
        d_model (int): 模型维度，即每个 token 向量的长度
        n_heads  (int): 注意力头的数量
    """

    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度: 512/8 = 64

        # Q, K, V 的线性投影层 (可学习参数)
        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 scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        缩放点积注意力

        Step 1: QK^T — 查询与键的点积，得到注意力分数
        Step 2: / √d_k — 缩放，防止梯度消失
        Step 3: softmax — 归一化为概率分布
        Step 4: × V — 加权求和 Value
        """
        # Step 1 & 2: 计算注意力分数并缩放
        # Q: (batch, n_heads, seq_len, d_k)
        # K^T: (batch, n_heads, d_k, seq_len)
        # scores: (batch, n_heads, seq_len, seq_len)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 应用掩码（如因果掩码）: 将需要屏蔽的位置设为 -∞
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Step 3: softmax 归一化
        attention_weights = F.softmax(scores, dim=-1)

        # Step 4: 用注意力权重加权 Value
        # output: (batch, n_heads, seq_len, d_k)
        output = torch.matmul(attention_weights, V)

        return output, attention_weights

    def forward(self, x, mask=None):
        """
        前向传播

        输入 x: (batch_size, seq_len, d_model)
        """
        batch_size, seq_len, _ = x.size()

        # Step 1: 线性投影得到 Q, K, V
        Q = self.W_q(x)  # (batch, seq_len, d_model)
        K = self.W_k(x)
        V = self.W_v(x)

        # Step 2: 拆分为多头
        # reshape 为 (batch, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Step 3: 计算缩放点积注意力
        attn_output, attn_weights = self.scaled_dot_product_attention(Q, K, V, mask)

        # Step 4: 拼接多头
        # (batch, n_heads, seq_len, d_k) → (batch, seq_len, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        # Step 5: 输出线性变换
        output = self.W_o(attn_output)

        return output, attn_weights


# ===== 使用示例 =====
if __name__ == "__main__":
    # 模型参数
    d_model = 512    # 模型维度
    n_heads = 8      # 注意力头数
    seq_len = 10     # 序列长度（如 10 个词）
    batch_size = 4  # 批量大小

    # 创建模型
    self_attn = SelfAttention(d_model=d_model, n_heads=n_heads)

    # 模拟输入: 4 个句子，每句 10 个词，每个词 512 维
    x = torch.randn(batch_size, seq_len, d_model)

    # 创建因果掩码（下三角矩阵，防止看到未来）
    mask = torch.tril(torch.ones(seq_len, seq_len))

    # 前向传播
    output, weights = self_attn(x, mask)

    print(f"输入形状:  {x.shape}")          # (4, 10, 512)
    print(f"输出形状:  {output.shape}")      # (4, 10, 512)
    print(f"注意力权重: {weights.shape}")    # (4, 8, 10, 10)
    print(f"每个样本有 {n_heads} 个头, 每个头是 {seq_len}×{seq_len} 的注意力矩阵")