Transformer 架构详解与完整实现

2026-03-10

基于 Google 2017年论文《Attention Is All You Need》的完整解析
实现版本：PyTorch，包含 Encoder-Decoder 完整架构

1. Transformer 总体架构

1.1 架构图

                Transformer
┌─────────────────────────────────────────────────┐
│                                                 │
│  Input Embedding + Positional Encoding          │
│         │                                       │
│         ▼                                       │
│  ┌─────────────┐    ┌─────────────┐            │
│  │   Encoder   │    │   Decoder   │            │
│  │   xN层      │    │   xN层      │            │
│  │             │    │             │            │
│  │  Self-Attn  │    │ MaskedSA    │            │
│  │  FFN        │    │ CrossAttn   │            │
│  │             │    │ FFN         │            │
│  └─────────────┘    └─────────────┘            │
│         │                                       │
│         ▼                                       │
│  Linear + Softmax → Output                      │
└─────────────────────────────────────────────────┘

1.2 数据流

源语言句子: "I love learning"
    ↓
词嵌入 (Embedding) + 位置编码 (Positional Encoding)
    ↓
编码器 (Encoder) — 6层堆叠 — 输出 memory
    ↓                          ↕
解码器 (Decoder) — 6层堆叠 — 交叉注意力
    ↓
线性投影 (Linear) + Softmax
    ↓
目标语言句子: "Ich lerne"

1.3 关键特性

特性	说明
完全基于注意力	不使用 RNN/CNN，全部由注意力机制构成
并行计算	Encoder 可以完全并行；Decoder 在训练时可并行（Teacher Forcing）
长距离依赖	任意两个位置之间的路径长度为 O(1)，远优于 RNN 的 O(n)
多尺度信息	多头注意力允许模型同时关注不同子空间的信息

2. 核心组件详解

2.1 位置编码 (Positional Encoding)

问题背景

Transformer 本身是排列不变的 (permutation invariant)。Self-Attention 对输入序列的顺序不敏感——打乱输入顺序，输出也会相应打乱。这与 RNN 天然处理序列顺序不同，因此需要显式注入位置信息。

数学公式

对于位置 pos 和第 i 个维度：

1 2	PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

为什么选择正弦函数？

能够表达相对位置：对于固定偏移 k，PE(pos + k) 可以表示为 PE(pos) 的线性函数，这使得模型可以学习关注相对位置而非绝对位置。
不同频率：低频分量捕捉全局结构，高频分量捕捉局部细节。
泛化到更长序列：由于 sin/cos 的定义域是整个实数轴，可以处理比训练时更长的序列。

实现要点

pe = torch.zeros(max_len, d_model)
pos = torch.arange(0, max_len).float().unsqueeze(1)           # [max_len, 1]
div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                     (-math.log(10000.0) / d_model))          # [d_model/2]
pe[:, 0::2] = torch.sin(pos * div_term)   # 偶数索引
pe[:, 1::2] = torch.cos(pos * div_term)   # 奇数索引

2.2 缩放点积注意力 (Scaled Dot-Product Attention)

公式

1	Attention(Q, K, V) = softmax(QK^T / √d_k) V

其中：

Q (Query): 查询向量，代表”我在寻找什么”
K (Key): 键向量，代表”我有什么”
V (Value): 值向量，代表”实际内容”
d_k: Q 和 K 的维度

为什么除以 √d_k？

当 d_k 较大时，QK^T 的点积结果的方差会增大（约为 d_k 倍）。这会导致 softmax 进入梯度极小的饱和区域（接近 0 或 1），使得反向传播时梯度消失。除以 √d_k 将方差缩放到 1，保持梯度的稳定性。

直观理解：

d_k = 64:  点积范围 [-5, 5]  → softmax 梯度适中
d_k = 512: 点积范围 [-15, 15] → softmax 梯度接近 0

除以 √d_k 后，两种情况的行为一致。

2.3 多头注意力 (Multi-Head Attention)

为什么要多头？

单个注意力头可能无法同时捕获多种类型的依赖关系。多头注意力让模型在不同的子空间中并行地学习不同的关注模式：

Head 1: 可能关注语法关系（主谓一致）
Head 2: 可能关注语义相似性
Head 3: 可能关注位置关系
…

计算流程

Input: [B, Seq, d_model]
    │
    ├→ W_Q → Q: [B, Seq, d_model] → [B, n_heads, Seq, d_k]
    ├→ W_K → K: [B, Seq, d_model] → [B, n_heads, Seq, d_k]
    ├→ W_V → V: [B, Seq, d_model] → [B, n_heads, Seq, d_k]
    │
    │  对每个 head i (共 h 个头):
    │    head_i = Attention(Q_i, K_i, V_i)     # [B, Seq, d_k]
    │
    Concat(head_1, ..., head_h) → [B, Seq, d_model]
    │
    └→ W_O → Output: [B, Seq, d_model]

参数计算

W_Q, W_K, W_V: 每个都是 [d_model, d_model] 的矩阵
W_O: [d_model, d_model] 的矩阵

虽然看起来有 4 个线性层，但总参数量与单头相同，
因为 d_k × n_heads = d_model。

2.4 位置前馈网络 (Position-wise FFN)

结构

Input: [B, Seq, d_model]
    │
    ├→ Linear(d_model → d_ffn)     # 通常 d_ffn = 4 × d_model
    ├→ ReLU
    ├→ Dropout
    ├→ Linear(d_ffn → d_model)
    └→ Dropout
    │
Output: [B, Seq, d_model]

数学表达

1	FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

作用

引入非线性变换（ReLU）
在每个位置上独立地变换特征表示
与注意力机制互补：注意力负责跨位置的交互，FFN 负责特征的局部变换

2.5 残差连接与 LayerNorm

结构

SubLayer(Input)
    │
    ├──→ Dropout
    ├──→ + (残差连接)
    └──→ LayerNorm

为什么用 LayerNorm 而不是 BatchNorm？

对比项	BatchNorm	LayerNorm
归一化维度	跨 batch	跨特征维度
序列模型适用性	不稳定（batch 间序列长度不同）	稳定
依赖关系	batch 统计量	单个样本

LayerNorm 对每个样本的特征维度独立归一化，不依赖 batch 统计量，更适合变长序列。

Post-LN vs Pre-LN

Post-LN（原始论文，我们使用的）：

1	Output = LayerNorm(Input + SubLayer(Input))

Pre-LN（后续改进）：

1	Output = Input + SubLayer(LayerNorm(Input))

Pre-LN 在更深层次的模型中更稳定，但需要调整学习率策略。

3. 训练与推理模式

3.1 训练：Teacher Forcing

源序列:     [BOS] "I" "love" "learning" [EOS]
目标序列:   [BOS] "Ich" "liebe" "lernen" [EOS]
                              ↓
编码器输出 (memory) ──────────┘
                              ↓
解码器输入: [BOS] "Ich" "liebe" "lernen" [EOS]
                              ↓
预测输出:   "Ich" "liebe" "lernen" [EOS] "UNK"
                              ↓
损失: CrossEntropy(pred[:, :-1], tgt[:, 1:])

关键特点：

一次性输入完整的目标序列
因果掩码确保每个位置只能看到之前的 token
损失计算时对齐：预测第 t 个位置的 token，与真实标签的第 t+1 个位置比较

3.2 推理：自回归生成 (Autoregressive)

Step 1:  输入: [BOS]                    → 输出: "Ich"
Step 2:  输入: [BOS] "Ich"              → 输出: "liebe"
Step 3:  输入: [BOS] "Ich" "liebe"      → 输出: "lernen"
Step 4:  输入: [BOS] "Ich" "liebe" "lernen" → 输出: [EOS]

关键特点：

编码器只需运行一次
每一步只输入一个新生成的 token
因果掩码动态增长（每一步序列长度 +1）
可以使用 Beam Search、Top-k 采样等策略

3.3 两种模式的差异

方面	训练	推理
解码器输入	完整序列	逐步生成
效率	高度并行	串行（每一步依赖前一步）
掩码	静态下三角	动态增长
错误累积	无（使用真实 token）	有（错误会传播）

4. 掩码机制

4.1 三种掩码类型

掩码类型	形状	用途
源填充掩码	`[B, 1, 1, src_len]`	编码器自注意力、解码器交叉注意力
因果掩码	`[B, 1, tgt_len, tgt_len]`	解码器自注意力（下三角）
组合掩码	`因果 × 填充`	解码器自注意力

4.2 掩码实现

# 源填充掩码: 非填充位置为 1，填充位置为 0
def make_src_mask(src):
    return (src != 0).unsqueeze(1).unsqueeze(2)  # [B, 1, 1, Src]

# 因果掩码: 下三角矩阵
def make_tgt_mask(tgt):
    batch, length = tgt.size()
    causal = torch.tril(torch.ones(length, length))  # [L, L]
    causal = causal.unsqueeze(0).unsqueeze(0)          # [1, 1, L, L]
    pad = (tgt != 0).unsqueeze(1).unsqueeze(2)         # [B, 1, 1, L]
    pad = pad.expand(batch, 1, length, length)
    return causal * pad  # [B, 1, L, L]

4.3 掩码如何工作

在注意力分数计算后、softmax 之前：

1
2
3

scores = QK^T / √d_k
scores = scores.masked_fill(mask == 0, -1e9)  # 掩码位置设为极大负值
attn_weights = softmax(scores)                  # exp(-1e9) ≈ 0

这样，被掩码的位置的注意力权重接近于 0，相当于”看不见”这些位置。

5. 完整代码实现

"""
Transformer 完整实现 (PyTorch)
================================
从零实现完整的 Encoder-Decoder Transformer 架构。
包含：位置编码、多头注意力、前馈网络、残差连接、LayerNorm、掩码处理。
"""

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


# ============================================================
# 1. 位置编码 (Positional Encoding)
# ============================================================
class PositionalEncoding(nn.Module):
    """
    正弦/余弦位置编码
    
    原理：
    - Transformer 本身是排列不变的(permutation invariant)，没有捕获token顺序的机制
    - 位置编码为每个位置的token添加唯一的位置信息
    - 使用不同频率的正弦和余弦函数，使模型能够关注相对位置
    
    公式：
      PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
      PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    
    关键性质：
    - 对于固定偏移 k，PE(pos+k) 可以表示为 PE(pos) 的线性函数
    - 这使得模型能够学习关注相对位置
    """
    def __init__(self, d_model: int, max_len: int = 5000, dropout_p: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout_p)
        
        # 创建形状为 (1, max_len, d_model) 的位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # [max_len, 1]
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )  # [d_model/2]
        
        # 偶数索引用sin，奇数索引用cos
        pe[:, 0::2] = torch.sin(pos * div_term)   # 偶数位
        pe[:, 1::2] = torch.cos(pos * div_term)   # 奇数位
        
        # 添加 batch 维度: (1, max_len, d_model)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 词嵌入张量 [batch_size, seq_len, d_model]
        Returns:
            加上位置编码后的张量
        """
        # 只取序列长度对应的位置编码
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


# ============================================================
# 2. 多头注意力 (Multi-Head Attention)
# ============================================================
class MultiHeadAttention(nn.Module):
    """
    多头自注意力机制
    
    原理：
    - 将 Q, K, V 分别投影到多个子空间（head）
    - 在每个子空间中独立计算缩放点积注意力
    - 将所有头的输出拼接后做最终投影
    
    公式：
      MultiHead(Q,K,V) = Concat(head_1, ..., head_h) W^O
      where head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
      
      Attention(Q,K,V) = softmax(QK^T / sqrt(d_k)) V
    
    为什么需要多头？
    - 每个head可以学习不同的关注模式（如语法、语义、位置等）
    - 扩大模型的表达能力
    
    为什么除以 sqrt(d_k)？
    - 当 d_k 较大时，点积结果的方差变大，softmax 进入梯度极小的饱和区
    - 缩放后保持梯度稳定
    """
    def __init__(self, d_model: int, n_heads: int, dropout_p: float = 0.1):
        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  # 每个头的维度
        
        # 四个线性变换：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)
        
        self.dropout = nn.Dropout(dropout_p)

    def _scaled_dot_product_attention(
        self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
        mask: torch.Tensor = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        缩放点积注意力计算
        
        Args:
            Q: [batch, n_heads, seq_q, d_k]
            K: [batch, n_heads, seq_k, d_k]
            V: [batch, n_heads, seq_k, d_v]
            mask: 可选的掩码张量
            
        Returns:
            output: 注意力输出 [batch, n_heads, seq_q, d_v]
            attn_weights: 注意力权重 [batch, n_heads, seq_q, seq_k]
        """
        # 计算注意力分数: Q @ K^T / sqrt(d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, H, Sq, Sk]
        
        # 应用掩码（如果提供）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)  # 将掩码位置设为极大负值
        
        # Softmax 归一化 → 注意力权重
        attn_weights = F.softmax(scores, dim=-1)  # [B, H, Sq, Sk]
        attn_weights = self.dropout(attn_weights)
        
        # 加权求和 V
        output = torch.matmul(attn_weights, V)  # [B, H, Sq, d_v]
        
        return output, attn_weights

    def forward(
        self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
        mask: torch.Tensor = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            Q: query [batch_size, seq_q, d_model]
            K: key   [batch_size, seq_k, d_model]
            V: value [batch_size, seq_k, d_model]
            mask: 掩码 [batch_size, 1, 1, seq_k] 或 [batch_size, 1, seq_q, seq_q]
        Returns:
            output: [batch_size, seq_q, d_model]
            attn_weights: [batch_size, n_heads, seq_q, seq_k]
        """
        batch_size = Q.size(0)
        
        # 1. 线性投影
        Q_t = self.W_Q(Q)  # [B, Sq, d_model]
        K_t = self.W_K(K)  # [B, Sk, d_model]
        V_t = self.W_V(V)  # [B, Sk, d_model]
        
        # 2. 重塑为多头: [B, Seq, d_model] → [B, Seq, n_heads, d_k] → [B, n_heads, Seq, d_k]
        Q_t = Q_t.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K_t = K_t.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V_t = V_t.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 3. 缩放点积注意力
        attn_output, attn_weights = self._scaled_dot_product_attention(Q_t, K_t, V_t, mask)
        
        # 4. 拼接所有头并做最终投影
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.W_O(attn_output)
        
        return output, attn_weights


# ============================================================
# 3. 位置前馈网络 (Position-wise Feed-Forward Network)
# ============================================================
class PositionwiseFFN(nn.Module):
    """
    位置-wise 前馈神经网络
    
    在Transformer中，每个位置的向量独立地通过同一个FFN处理。
    
    结构：
      FFN(x) = max(0, xW1 + b1)W2 + b2
    
    参数：
      - 内层维度 d_ffn 通常是 d_model 的 4 倍
      - 两层线性变换，中间用ReLU激活
    
    为什么叫"position-wise"？
    - 相同的FFN在所有位置共享
    - 每个位置的变换是独立的（不跨位置交互）
    - 跨位置的交互完全由注意力机制完成
    """
    def __init__(self, d_model: int, d_ffn: int, dropout_p: float = 0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ffn)  # 升维
        self.linear2 = nn.Linear(d_ffn, d_model)  # 降维
        self.dropout = nn.Dropout(dropout_p)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear2(self.dropout(F.relu(self.linear1(x))))


# ============================================================
# 4. 编码器层 (Encoder Layer)
# ============================================================
class EncoderLayer(nn.Module):
    """
    单个编码器层
    
    结构：
      Input → Multi-Head Self-Attention → Add & LayerNorm
                → Position-wise FFN → Add & LayerNorm → Output
    
    每个子层（Attention 和 FFN）都包裹着残差连接和LayerNorm。
    这是原论文的实现（Post-LN）。
    """
    def __init__(self, d_model: int, n_heads: int, d_ffn: int, dropout_p: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout_p)
        self.ffn = PositionwiseFFN(d_model, d_ffn, dropout_p)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop1 = nn.Dropout(dropout_p)
        self.drop2 = nn.Dropout(dropout_p)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            x: 编码器输入 [batch_size, seq_len, d_model]
            mask: 源序列填充掩码
        Returns:
            编码器层输出
        """
        # 子层1: 自注意力 + 残差连接 + LayerNorm
        attn_out, _ = self.self_attn(x, x, x, mask=mask)
        x = self.norm1(x + self.drop1(attn_out))
        
        # 子层2: 前馈网络 + 残差连接 + LayerNorm
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.drop2(ffn_out))
        
        return x


# ============================================================
# 5. 解码器层 (Decoder Layer)
# ============================================================
class DecoderLayer(nn.Module):
    """
    单个解码器层
    
    结构：
      Input → Masked Self-Attention → Add & LayerNorm
                → Cross-Attention (Q=tgt, K/V=memory) → Add & LayerNorm
                → Position-wise FFN → Add & LayerNorm → Output
    
    关键区别：
    - 第一个注意力层使用因果掩码（masked），确保只能看到当前位置之前的token
    - 第二个注意力层是交叉注意力，查询来自解码器，键/值来自编码器输出
    """
    def __init__(self, d_model: int, n_heads: int, d_ffn: int, dropout_p: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout_p)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout_p)
        self.ffn = PositionwiseFFN(d_model, d_ffn, dropout_p)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.drop1 = nn.Dropout(dropout_p)
        self.drop2 = nn.Dropout(dropout_p)
        self.drop3 = nn.Dropout(dropout_p)

    def forward(
        self, tgt: torch.Tensor, memory: torch.Tensor,
        tgt_mask: torch.Tensor = None, memory_mask: torch.Tensor = None
    ) -> torch.Tensor:
        """
        Args:
            tgt: 目标序列 [batch_size, tgt_len, d_model]
            memory: 编码器输出 [batch_size, src_len, d_model]
            tgt_mask: 目标序列因果掩码
            memory_mask: 源序列填充掩码
        Returns:
            解码器层输出
        """
        # 子层1: 掩码自注意力（防止看到未来token）
        attn1, _ = self.self_attn(tgt, tgt, tgt, mask=tgt_mask)
        tgt = self.norm1(tgt + self.drop1(attn1))
        
        # 子层2: 交叉注意力（decoder查询encoder的key/value）
        attn2, _ = self.cross_attn(tgt, memory, memory, mask=memory_mask)
        tgt = self.norm2(tgt + self.drop2(attn2))
        
        # 子层3: 前馈网络
        ffn_out = self.ffn(tgt)
        tgt = self.norm3(tgt + self.drop3(ffn_out))
        
        return tgt


# ============================================================
# 6. 嵌入层 (Embedding)
# ============================================================
class EmbeddingLayer(nn.Module):
    """
    词嵌入层
    
    包含：
    - 词表查找 (Embedding)
    - 缩放因子 sqrt(d_model) — 原论文约定，稳定初始梯度
    """
    def __init__(self, vocab_size: int, d_model: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 词嵌入 × sqrt(d_model)
        return self.embedding(x) * math.sqrt(self.d_model)


# ============================================================
# 7. 完整 Transformer
# ============================================================
class Transformer(nn.Module):
    """
    完整的 Encoder-Decoder Transformer 模型
    
    整体流程：
    1. 输入序列 → 词嵌入 + 位置编码 → 编码器
    2. 编码器输出 (memory) + 目标序列 → 解码器
    3. 解码器输出 → 线性投影 + Softmax → 词概率分布
    
    训练 vs 推理：
    - 训练时：使用Teacher Forcing，一次性输入完整的目标序列
    - 推理时：自回归生成，逐步预测下一个token
    """
    def __init__(
        self,
        src_vocab_size: int,
        tgt_vocab_size: int,
        d_model: int = 512,
        n_heads: int = 8,
        d_ffn: int = 2048,
        n_encoder_layers: int = 6,
        n_decoder_layers: int = 6,
        max_len: int = 5000,
        dropout_p: float = 0.1,
        tie_embeddings: bool = False
    ):
        super().__init__()
        
        self.d_model = d_model
        self.tie_embeddings = tie_embeddings
        
        # 嵌入层
        self.src_embedding = EmbeddingLayer(src_vocab_size, d_model)
        self.tgt_embedding = EmbeddingLayer(tgt_vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout_p)
        
        # 编码器
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(d_model, n_heads, d_ffn, dropout_p)
            for _ in range(n_encoder_layers)
        ])
        
        # 解码器
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, n_heads, d_ffn, dropout_p)
            for _ in range(n_decoder_layers)
        ])
        
        # 输出投影
        if tie_embeddings:
            self.output_proj = nn.Linear(d_model, tgt_vocab_size, bias=False)
            # 共享权重
            self.output_proj.weight = self.tgt_embedding.embedding.weight
        else:
            self.output_proj = nn.Linear(d_model, tgt_vocab_size)
            self.dropout = nn.Dropout(dropout_p)
        
        # Xavier 初始化
        self._init_parameters()

    def _init_parameters(self):
        """Xavier均匀初始化所有线性层参数"""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def make_src_mask(self, src: torch.Tensor) -> torch.Tensor:
        """
        创建源序列填充掩码
        
        Args:
            src: 源序列 [batch_size, src_len]
        Returns:
            掩码 [batch_size, 1, 1, src_len]，非填充位置为1，填充位置为0
        """
        return (src != 0).unsqueeze(1).unsqueeze(2)  # [B, 1, 1, Src]

    def make_tgt_mask(self, tgt: torch.Tensor) -> torch.Tensor:
        """
        创建目标序列因果掩码 + 填充掩码
        
        Args:
            tgt: 目标序列 [batch_size, tgt_len]
        Returns:
            组合掩码 [batch_size, 1, tgt_len, tgt_len]
        """
        batch_size, tgt_len = tgt.size()
        
        # 因果掩码：下三角矩阵（只能看到当前位置及之前的token）
        tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len, device=tgt.device)).unsqueeze(0).unsqueeze(0)
        # [1, 1, tgt_len, tgt_len]
        
        # 填充掩码
        pad_mask = (tgt != 0).unsqueeze(1).unsqueeze(2)  # [B, 1, 1, tgt_len]
        
        # 扩展填充掩码到因果掩码的形状
        pad_mask = pad_mask.expand(batch_size, 1, tgt_len, tgt_len)
        
        # 组合：两个位置都为1才有效
        return tgt_mask * pad_mask  # [B, 1, tgt_len, tgt_len]

    def encode(self, src: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        """
        编码阶段
        
        Args:
            src: 源序列 [batch_size, src_len]
            src_mask: 源序列掩码
        Returns:
            编码器输出 (memory) [batch_size, src_len, d_model]
        """
        # 嵌入 + 位置编码
        x = self.src_embedding(src)
        x = self.pos_encoding(x)
        
        # 通过所有编码器层
        for layer in self.encoder_layers:
            x = layer(x, src_mask)
        
        return x

    def decode(
        self, tgt: torch.Tensor, memory: torch.Tensor,
        tgt_mask: torch.Tensor, src_mask: torch.Tensor
    ) -> torch.Tensor:
        """
        解码阶段
        
        Args:
            tgt: 目标序列 [batch_size, tgt_len]
            memory: 编码器输出
            tgt_mask: 目标序列掩码（因果+填充）
            src_mask: 源序列掩码
        Returns:
            解码器输出 [batch_size, tgt_len, d_model]
        """
        # 嵌入 + 位置编码
        x = self.tgt_embedding(tgt)
        x = self.pos_encoding(x)
        
        # 通过所有解码器层
        for layer in self.decoder_layers:
            x = layer(x, memory, tgt_mask, src_mask)
        
        return x

    def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        
        Args:
            src: 源序列 [batch_size, src_len]
            tgt: 目标序列 [batch_size, tgt_len]
        Returns:
            输出logits [batch_size, tgt_len, tgt_vocab_size]
        """
        src_mask = self.make_src_mask(src)
        tgt_mask = self.make_tgt_mask(tgt)
        
        memory = self.encode(src, src_mask)
        output = self.decode(tgt, memory, tgt_mask, src_mask)
        
        if self.tie_embeddings:
            logits = self.output_proj(output)
        else:
            logits = self.output_proj(self.dropout(output))
        
        return logits

    @torch.no_grad()
    def generate(
        self, src: torch.Tensor, max_length: int = 50, sos_token: int = 1, eos_token: int = 2
    ) -> torch.Tensor:
        """
        自回归生成（推理模式）
        
        逐步生成目标序列，每次只输入已生成的token。
        
        Args:
            src: 源序列 [1, src_len] （假设batch_size=1）
            max_length: 最大生成长度
            sos_token: 开始token ID
            eos_token: 结束token ID
        Returns:
            生成的token序列 [1, generated_len]
        """
        self.eval()
        batch_size = src.size(0)
        
        # 先编码（只需一次）
        src_mask = self.make_src_mask(src)
        memory = self.encode(src, src_mask)
        
        # 初始化：从SOS token开始
        generated = torch.full((batch_size, 1), sos_token, dtype=torch.long, device=src.device)
        
        for _ in range(max_length):
            tgt_len = generated.size(1)
            tgt_mask = torch.tril(
                torch.ones(tgt_len, tgt_len, device=src.device)
            ).unsqueeze(0).unsqueeze(0)  # [1, 1, tgt_len, tgt_len]
            
            # 解码当前生成的序列
            output = self.decode(generated, memory, tgt_mask, src_mask)
            
            # 取最后一个位置的logits
            logits = self.output_proj(output[:, -1, :])  # [B, tgt_vocab_size]
            probs = F.softmax(logits, dim=-1)
            
            # 贪心选择（也可以改用采样或beam search）
            next_token = probs.argmax(dim=-1, keepdim=True)  # [B, 1]
            
            # 附加到新序列
            generated = torch.cat([generated, next_token], dim=1)
            
            # 如果生成了EOS，提前停止
            if (next_token == eos_token).all():
                break
        
        self.train()
        return generated


# ============================================================
# 辅助函数：统计模型参数量
# ============================================================
def count_parameters(model: nn.Module) -> int:
    """统计模型的可学习参数总数"""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# ============================================================
# 示例：创建并测试模型
# ============================================================
if __name__ == "__main__":
    # 超参数
    SRC_VOCAB = 10000
    TGT_VOCAB = 10000
    D_MODEL = 256
    N_HEADS = 8
    D_FFN = 512
    N_LAYERS = 3
    MAX_LEN = 200
    DROPOUT = 0.1
    
    # 创建模型
    model = Transformer(
        src_vocab_size=SRC_VOCAB,
        tgt_vocab_size=TGT_VOCAB,
        d_model=D_MODEL,
        n_heads=N_HEADS,
        d_ffn=D_FFN,
        n_encoder_layers=N_LAYERS,
        n_decoder_layers=N_LAYERS,
        max_len=MAX_LEN,
        dropout_p=DROPOUT
    )
    
    print("=" * 60)
    print("Transformer 模型架构")
    print("=" * 60)
    print(f"  词表大小 (源/目标): {SRC_VOCAB} / {TGT_VOCAB}")
    print(f"  隐藏层维度 (d_model): {D_MODEL}")
    print(f"  头数 (n_heads): {N_HEADS}")
    print(f"  FFN 维度 (d_ffn): {D_FFN}")
    print(f"  编码器/解码器层数: {N_LAYERS} / {N_LAYERS}")
    print(f"  Dropout: {DROPOUT}")
    print(f"  总参数量: {count_parameters(model):,}")
    print()
    
    # 测试前向传播
    batch_size = 2
    src_len = 10
    tgt_len = 8
    
    src_seq = torch.randint(1, SRC_VOCAB, (batch_size, src_len))  # [2, 10]
    tgt_seq = torch.randint(1, TGT_VOCAB, (batch_size, tgt_len))  # [2, 8]
    
    print("输入:")
    print(f"  src: {tuple(src_seq.shape)}")
    print(f"  tgt: {tuple(tgt_seq.shape)}")
    
    output = model(src_seq, tgt_seq)
    print(f"\n输出: {tuple(output.shape)}")
    print(f"  期望: ([{batch_size}, {tgt_len}, {TGT_VOCAB}])")
    print()
    
    # 测试生成
    print("自回归生成测试:")
    generated = model.generate(src_seq[:1], max_length=15)
    print(f"  生成序列: {tuple(generated.shape)}")
    print(f"  生成token: {generated.tolist()}")
    print()
    
    # 打印模型结构
    print("模型结构:")
    print(model)

5.1 文件结构

transformer_full.py
├── PositionalEncoding      # 位置编码模块
├── MultiHeadAttention      # 多头注意力模块
├── PositionwiseFFN         # 位置前馈网络
├── EncoderLayer            # 编码器单层
├── DecoderLayer            # 解码器单层
├── EmbeddingLayer          # 词嵌入层
└── Transformer             # 完整模型
    ├── encode()            # 编码方法
    ├── decode()            # 解码方法
    ├── forward()           # 前向传播
    └── generate()          # 自回归生成

5.2 类层次关系

Transformer
├── src_embedding: EmbeddingLayer
├── tgt_embedding: EmbeddingLayer
├── pos_encoding: PositionalEncoding
├── encoder_layers: [EncoderLayer] × N
│   └── MultiHeadAttention (self-attn)
│   └── PositionwiseFFN
├── decoder_layers: [DecoderLayer] × N
│   └── MultiHeadAttention (masked self-attn)
│   └── MultiHeadAttention (cross-attn)
│   └── PositionwiseFFN
└── output_proj: Linear

6. 关键设计决策

6.1 为什么不用 RNN？

对比项	RNN/LSTM	Transformer
并行性	串行（必须按顺序处理）	完全并行
长距离依赖	梯度消失/爆炸问题	O(1) 路径长度
训练速度	慢	快（GPU 并行）
推理速度	快（可流式）	慢（自回归）

6.2 为什么用 Self-Attention 而不是普通 Attention？

Self-Attention: Q, K, V 来自同一序列 → 捕获序列内部关系
Cross-Attention (普通 Attention): Q 来自一个序列，K/V 来自另一个 → 跨序列信息传递

Transformer 两者都用：

Encoder 内部：Self-Attention
Decoder 内部：Self-Attention + Cross-Attention

6.3 超参数选择指南

参数	推荐值	说明
d_model	256-512	小型模型可用 128，大型模型可用 1024
n_heads	8	必须能整除 d_model
d_ffn	4 × d_model	通常设为 d_model 的 4 倍
n_layers	6	原始论文使用 6 层
dropout	0.1	防止过拟合
max_len	5000	根据任务需求调整

6.4 常见变体

变体	架构	典型应用
Encoder-only	仅编码器	BERT, 文本分类, NER
Decoder-only	仅解码器	GPT, 文本生成
Encoder-Decoder	完整结构	机器翻译, 摘要

7. 运行结果

以上代码已成功编译并运行，测试结果如下：

============================================================
Transformer 模型架构
============================================================
  词表大小 (源/目标): 10000 / 10000
  隐藏层维度 (d_model): 256
  头数 (n_heads): 8
  FFN 维度 (d_ffn): 512
  编码器/解码器层数: 3 / 3
  Dropout: 0.1
  总参数量: 11,643,664

输入:
  src: (2, 10)
  tgt: (2, 8)

输出: (2, 8, 10000)
  期望: ([2, 8, 10000])

自回归生成测试:
  生成序列: (1, 16)
  生成token: [[1, 6074, 6074, 9395, ...]]

模型结构: [完整模型架构输出]

结果解读：

✅ 模型成功构建，总参数量约 1160 万
✅ 前向传播输出形状正确：[batch, tgt_len, vocab_size]
✅ 自回归生成正常工作，输出形状 [1, generated_len]
✅ 模型结构清晰，包含所有必要组件

附录：参考资源

原始论文: Vaswani et al., “Attention Is All You Need”, NeurIPS 2017
可视化教程: The Illustrated Transformer by Jay Alammar
PyTorch 官方实现: torch.nn.Transformer
HuggingFace Transformers: transformers 库

目录