第三章:Transformer架构详解

第三章:Transformer架构详解

本章学习目标

  • 理解Transformer架构诞生的背景和动机
  • 掌握Transformer的整体架构(Encoder-Decoder结构)
  • 深入理解自注意力机制(Self-Attention)的数学原理
  • 掌握多头注意力(Multi-Head Attention)的设计哲学
  • 理解位置编码的原理和不同方案
  • 了解Transformer的各种变体(BERT、GPT、T5等)
  • 掌握大模型训练的关键技术(分布式训练、混合精度等)

3.1 Transformer诞生背景与整体架构

3.1.1 从Seq2Seq到Transformer的演进

概念引入:Seq2Seq的瓶颈

在Transformer出现之前,序列到序列(Seq2Seq)任务主要使用**循环神经网络(RNN/LSTM/GRU)**来建模。

类比理解:想象你在阅读一本长篇小说。RNN就像是一个只能记住最近几页内容的读者------当小说很长时,前面章节的细节会逐渐被遗忘。虽然LSTM通过门控机制缓解了这个问题,但当序列长度超过几百个词时,性能仍然会显著下降。

RNN/LSTM的核心问题

  1. 序列依赖导致无法并行训练

    复制代码
    处理序列 [x₁, x₂, x₃, x₄, ...]
    RNN必须按顺序计算:h₁ → h₂ → h₃ → h₄ → ...
    无法并行!训练速度慢
  2. 长期依赖问题(Long-Term Dependency)

    • 梯度消失/爆炸使得模型难以捕捉长距离依赖
    • 例如:在 "The cat, which was ... (100个词后) ... was black." 中,模型很难建立 "cat" 和 "was" 之间的语法关系
  3. 计算效率低下

    • 每个时间步都需要完整的矩阵运算
    • GPU的并行计算能力无法被充分利用
Attention机制的引入

Bahdanau等人在2014年将**注意力机制(Attention Mechanism)**引入Seq2Seq模型,使得解码器在生成每个词时,能够"关注"编码器隐藏状态的不同部分。

Attention的改进

  • ✅ 解决了瓶颈问题(不再依赖单一上下文向量)
  • ✅ 提升了长距离依赖的建模能力
  • 但仍然是RNN-based,无法完全并行化
Transformer的突破:完全基于Attention

Vaswani等人在2017年发表了里程碑论文《Attention is All You Need》,提出了Transformer架构,其革命性在于:

核心思想 :完全抛弃RNN,仅使用注意力机制前馈神经网络来构建序列模型。

优势

  1. 完全可并行化:不同位置的注意力计算互不影响
  2. 长距离依赖建模能力强:任意两个位置的距离都是O(1)
  3. 训练速度快:充分利用GPU并行计算能力
  4. 可解释性强:注意力权重可视化能展示模型的"关注点"
python 复制代码
# ===== 代码实战3.1:对比RNN和Transformer的计算模式 =====
import torch
import torch.nn as nn
import time

def compare_rnn_transformer_speed():
    """
    对比RNN和Transformer的处理速度
    展示Transformer的并行化优势
    """
    print("=== RNN vs Transformer 计算模式对比 ===\n")
    
    # 模拟输入 (batch_size=32, seq_len=128, hidden_dim=512)
    batch_size = 32
    seq_len = 128
    hidden_dim = 512
    
    # ===== RNN模式(序列依赖)=====
    print("1. RNN/LSTM计算模式(无法并行):")
    print("   时间步计算顺序: t₁ → t₂ → t₃ → ... → tₙ")
    print("   每个时间步必须等待前一步完成")
    
    # 模拟RNN计算(简化)
    rnn = nn.LSTM(input_size=hidden_dim, hidden_size=hidden_dim, batch_first=True)
    
    x = torch.randn(batch_size, seq_len, hidden_dim)
    start_time = time.time()
    output, (h_n, c_n) = rnn(x)
    rnn_time = time.time() - start_time
    print(f"   RNN处理时间: {rnn_time:.4f}秒\n")
    
    # ===== Transformer模式(完全并行)=====
    print("2. Transformer计算模式(完全并行):")
    print("   所有位置同时计算: t₁, t₂, t₃, ..., tₙ 并行")
    print("   通过注意力机制直接建模任意位置间的关系")
    
    # 模拟Self-Attention计算(简化)
    class SimpleSelfAttention(nn.Module):
        def __init__(self, hidden_dim):
            super(SimpleSelfAttention, self).__init__()
            self.W_q = nn.Linear(hidden_dim, hidden_dim)
            self.W_k = nn.Linear(hidden_dim, hidden_dim)
            self.W_v = nn.Linear(hidden_dim, hidden_dim)
        
        def forward(self, x):
            # x: (batch, seq_len, hidden_dim)
            Q = self.W_q(x)  # (batch, seq_len, hidden_dim)
            K = self.W_k(x)
            V = self.W_v(x)
            
            # 计算注意力分数(完全并行!)
            scores = torch.bmm(Q, K.transpose(1, 2)) / torch.sqrt(torch.tensor(hidden_dim, dtype=torch.float32))
            attention_weights = torch.softmax(scores, dim=-1)
            output = torch.bmm(attention_weights, V)
            
            return output
    
    attention = SimpleSelfAttention(hidden_dim)
    
    start_time = time.time()
    output = attention(x)
    transformer_time = time.time() - start_time
    print(f"   Self-Attention处理时间: {transformer_time:.4f}秒")
    print(f"   加速比: {rnn_time/transformer_time:.2f}x")
    print("\n   说明: 实际加速效果取决于序列长度和硬件配置")

compare_rnn_transformer_speed()

3.1.2 Transformer整体架构概览

Encoder-Decoder结构

Transformer遵循经典的**编码器-解码器(Encoder-Decoder)**结构:

复制代码
输入序列: [x₁, x₂, ..., xₙ]
    ↓
[Encoder Stack] (N层)
    ↓
中间表示: [z₁, z₂, ..., zₙ]
    ↓
[Decoder Stack] (N层)
    ↓
输出序列: [y₁, y₂, ..., yₘ]

架构图(文字版)

复制代码
Encoder:
    Input → Embedding → Positional Encoding
                                      ↓
                            ┌────────────────────┐
                            │   Multi-Head       │
                            │   Attention        │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Add & Norm       │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Feed Forward     │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Add & Norm       │
                            └────────────────────┘
                                      ↓
                            (重复N次)

Decoder:
    Output → Embedding → Positional Encoding
                                      ↓
                            ┌────────────────────┐
                            │ Masked Multi-Head  │
                            │ Attention          │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Add & Norm       │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Multi-Head       │
                            │   (Encoder-Decoder)│
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Add & Norm       │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Feed Forward     │
                            └────────────────────┘
                                      ↓
                            ┌────────────────────┐
                            │   Add & Norm       │
                            └────────────────────┘
                                      ↓
                            (重复N次)
                                      ↓
                            Linear + Softmax
核心组件概览
组件 位置 作用
Input/Output Embedding Encoder/Decoder输入 将词转换为向量表示
Positional Encoding Embedding之后 为序列引入位置信息
Multi-Head Attention Encoder中1次,Decoder中2次 捕捉序列内部的依赖关系
Add & Norm 每个子层之后 残差连接 + 层归一化
Feed Forward Network 每个Attention之后 非线性变换
Linear & Softmax Decoder输出 生成最终词概率分布

关键洞察 :Transformer的成功不仅在于Attention机制,更在于残差连接层归一化位置编码等设计的精妙组合。这些组件共同使得训练极深的网络成为可能。


3.1.3 输入嵌入与位置编码

词嵌入(Word Embedding)

将离散的词映射到连续的向量空间。

数学表达

给定词汇表大小为 VVV,嵌入维度为 dmodeld_{model}dmodel,嵌入层为一个矩阵 E∈RV×dmodelE \in \mathbb{R}^{V \times d_{model}}E∈RV×dmodel。

对于输入词 xix_ixi(表示为词汇表中的索引),其嵌入向量为:

ei=Exi∈Rdmodel \mathbf{e}i = Ex_i \in \mathbb{R}^{d{model}} ei=Exi∈Rdmodel

缩放因子

在原始Transformer中,嵌入向量会乘以 dmodel\sqrt{d_{model}}dmodel :

ei′=ei×dmodel \mathbf{e}_i' = \mathbf{e}i \times \sqrt{d{model}} ei′=ei×dmodel

原因 :如果没有缩放,嵌入向量的量级会随着 dmodeld_{model}dmodel 增大而增大,导致点积注意力的值过大。

位置编码(Positional Encoding)

问题:Attention机制是**排列不变(Permutation Invariant)**的------即,打乱输入序列的顺序,Attention的输出不变。

类比理解:想象你在读一句话 "The cat sat on the mat"。如果Transformer不知道每个词的位置,它就无法理解 "cat" 是主语,"sat" 是谓语。位置编码就是给每个词贴上"我是第几个词"的标签。

解决方案 :为每个位置 pospospos 添加一个位置编码向量 ppos∈Rdmodel\mathbf{p}{pos} \in \mathbb{R}^{d{model}}ppos∈Rdmodel。

最终输入表示:

xi=ei+pi \mathbf{x}_i = \mathbf{e}_i + \mathbf{p}_i xi=ei+pi

绝对位置编码(Sinusoidal Positional Encoding)

原始Transformer使用正弦/余弦函数生成位置编码:

PE(pos,2i)=sin⁡(pos100002i/dmodel) PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) PE(pos,2i)=sin(100002i/dmodelpos)

PE(pos,2i+1)=cos⁡(pos100002i/dmodel) PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) PE(pos,2i+1)=cos(100002i/dmodelpos)

其中:

  • pospospos:位置索引(0, 1, 2, ...)
  • iii:维度索引(0, 1, ..., dmodel/2d_{model}/2dmodel/2)

特性

  1. 确定性:不需要学习,可以直接计算
  2. 外推性:理论上可以处理比训练时更长的序列
  3. 相对位置信息 :对于任意固定偏移 kkk,PEpos+kPE_{pos+k}PEpos+k 可以表示为 PEposPE_{pos}PEpos 的线性组合
python 复制代码
# ===== 代码实战3.2:实现Sinusoidal位置编码 =====
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

class SinusoidalPositionalEncoding(nn.Module):
    """
    正弦位置编码(原始Transformer使用)
    """
    def __init__(self, d_model, max_len=5000):
        """
        初始化位置编码
        
        Args:
            d_model: 模型维度(嵌入维度)
            max_len: 支持的最大序列长度
        """
        super(SinusoidalPositionalEncoding, self).__init__()
        
        # 创建位置编码矩阵 (max_len, d_model)
        pe = torch.zeros(max_len, d_model)
        
        # 位置索引 (max_len, 1)
        position = torch.arange(0, max_len).unsqueeze(1)
        
        # 分母项 1 / 10000^(2i/d_model)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model)
        )
        
        # 偶数维度使用sin
        pe[:, 0::2] = torch.sin(position * div_term)
        # 奇数维度使用cos
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # 添加batch维度 (1, max_len, d_model)
        pe = pe.unsqueeze(0)
        
        # 注册为buffer(不作为模型参数,但会被保存)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """
        添加位置编码到输入嵌入
        
        Args:
            x: 输入嵌入,形状为 (batch, seq_len, d_model)
        
        Returns:
            添加位置编码后的表示
        """
        # 只取前seq_len个位置编码
        return x + self.pe[:, :x.size(1), :]


def visualize_positional_encoding():
    """可视化位置编码"""
    d_model = 64
    max_len = 100
    
    pe = SinusoidalPositionalEncoding(d_model, max_len)
    
    # 获取位置编码矩阵
    pe_matrix = pe.pe.squeeze(0).numpy()  # (max_len, d_model)
    
    # 绘制热力图
    plt.figure(figsize=(12, 6))
    plt.imshow(pe_matrix, aspect='auto', cmap='viridis')
    plt.colorbar(label='Encoding Value')
    plt.xlabel('Dimension', fontsize=12)
    plt.ylabel('Position', fontsize=12)
    plt.title('Sinusoidal Positional Encoding Heatmap', fontsize=14)
    plt.tight_layout()
    plt.savefig('positional_encoding_heatmap.png', dpi=150)
    plt.show()
    
    # 展示不同位置在某些维度上的值
    plt.figure(figsize=(12, 4))
    positions = [0, 10, 20, 30]
    for i, pos in enumerate(positions):
        plt.plot(pe_matrix[pos, :16], label=f'Position {pos}', 
                marker='o', markersize=4)
    plt.xlabel('Dimension (first 16)', fontsize=12)
    plt.ylabel('Encoding Value', fontsize=12)
    plt.title('Positional Encoding Values for Different Positions', fontsize=14)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('positional_encoding_values.png', dpi=150)
    plt.show()


def demonstrate_positional_encoding_usage():
    """演示位置编码的实际使用"""
    print("=== 位置编码使用示例 ===\n")
    
    batch_size = 2
    seq_len = 10
    d_model = 512
    
    # 模拟输入嵌入(通常来自词嵌入层)
    embeddings = torch.randn(batch_size, seq_len, d_model)
    print(f"输入嵌入形状: {embeddings.shape}")
    
    # 添加位置编码
    pe_layer = SinusoidalPositionalEncoding(d_model)
    encoded = pe_layer(embeddings)
    
    print(f"添加位置编码后形状: {encoded.shape}")
    print(f"\n说明:")
    print(f"  - 位置编码通过与嵌入向量相加来注入位置信息")
    print(f"  - 不同位置的编码向量不同")
    print(f"  - 模型可以通过这些值感知词在序列中的位置")

if __name__ == "__main__":
    demonstrate_positional_encoding_usage()
    visualize_positional_encoding()
相对位置编码(Relative Positional Encoding)

绝对位置编码的问题

  • 无法直接建模词之间的相对位置关系
  • 例如:"cat" 和 "sat" 之间的距离应该是相同的,无论它们在句子的哪个位置

相对位置编码的核心思想

不直接编码每个位置的绝对值,而是编码位置之间的差异

主流方案

方案 提出者 核心思想
Relative Position Representations Shaw et al. (2018) 在Attention计算中显式添加相对位置偏置
RoPE (Rotary Position Embedding) Su et al. (2021) 通过旋转矩阵编码相对位置
ALiBi (Attention with Linear Biases) Press et al. (2021) 在Attention分数上添加线性衰减偏置

RoPE(当前主流方案)

将查询向量 q\mathbf{q}q 和键向量 k\mathbf{k}k 通过旋转矩阵进行变换:

qm′=qm⋅(cos⁡mθ0−sin⁡mθ0⋯sin⁡mθ0cos⁡mθ0⋯⋮⋮⋱) \mathbf{q}_m' = \mathbf{q}_m \cdot \begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & \cdots \\ \sin m\theta_0 & \cos m\theta_0 & \cdots \\ \vdots & \vdots & \ddots \end{pmatrix} qm′=qm⋅ cosmθ0sinmθ0⋮−sinmθ0cosmθ0⋮⋯⋯⋱

kn′=kn⋅(cos⁡nθ0−sin⁡nθ0⋯sin⁡nθ0cos⁡nθ0⋯⋮⋮⋱) \mathbf{k}_n' = \mathbf{k}_n \cdot \begin{pmatrix} \cos n\theta_0 & -\sin n\theta_0 & \cdots \\ \sin n\theta_0 & \cos n\theta_0 & \cdots \\ \vdots & \vdots & \ddots \end{pmatrix} kn′=kn⋅ cosnθ0sinnθ0⋮−sinnθ0cosnθ0⋮⋯⋯⋱

关键性质

⟨qm′,kn′⟩=⟨qm,kn⟩⋅cos⁡((m−n)θ0) \langle \mathbf{q}_m', \mathbf{k}_n' \rangle = \langle \mathbf{q}_m, \mathbf{k}_n \rangle \cdot \cos((m-n)\theta_0) ⟨qm′,kn′⟩=⟨qm,kn⟩⋅cos((m−n)θ0)

即,旋转后的内积仅依赖于相对位置 m−nm-nm−n!

为什么RoPE在LLM时代成为主流?

  1. 优秀的外推性:可以在推理时处理比训练时更长的序列
  2. 计算高效:可以通过复数乘法快速实现
  3. 理论优雅:完美编码相对位置信息

当前主流大模型(LLaMA、ChatGLM、Qwen等)都使用RoPE。


3.1.4 编码器与解码器堆栈结构

编码器堆栈(Encoder Stack)

编码器由 NNN 个相同的层堆叠而成(原始Transformer中 N=6N=6N=6)。

每一层包含两个子层

  1. Multi-Head Self-Attention
  2. Position-wise Feed-Forward Network

每个子层之后都有残差连接(Residual Connection)层归一化(Layer Normalization)

数学表达

对于第 lll 层编码器:

z(l)=LayerNorm(x+MultiHeadAttention(x)) \mathbf{z}^{(l)} = \text{LayerNorm}(\mathbf{x} + \text{MultiHeadAttention}(\mathbf{x})) z(l)=LayerNorm(x+MultiHeadAttention(x))

z(l+1)=LayerNorm(z(l)+FFN(z(l))) \mathbf{z}^{(l+1)} = \text{LayerNorm}(\mathbf{z}^{(l)} + \text{FFN}(\mathbf{z}^{(l)})) z(l+1)=LayerNorm(z(l)+FFN(z(l)))

其中 x\mathbf{x}x 为输入,FFN\text{FFN}FFN 为前馈神经网络。

数据流转过程

复制代码
输入序列: [x₁, x₂, ..., xₙ]
    ↓ (词嵌入 + 位置编码)
Embedded: [e₁, e₂, ..., eₙ]
    ↓ (Encoder Layer 1)
Encoder Layer 1 Output: [h₁¹, h₂¹, ..., hₙ¹]
    ↓ (Encoder Layer 2)
Encoder Layer 2 Output: [h₁², h₂², ..., hₙ²]
    ↓
    ...
    ↓ (Encoder Layer N)
Encoder Output: [z₁, z₂, ..., zₙ]
解码器堆栈(Decoder Stack)

解码器也由 NNN 个相同的层堆叠而成。

每一层包含三个子层

  1. Masked Multi-Head Self-Attention(带掩码的自注意力)
  2. Multi-Head Encoder-Decoder Attention(编码器-解码器注意力)
  3. Position-wise Feed-Forward Network

Masked Self-Attention的作用

在训练时,解码器不应该"看到"未来的词(因为生成是自回归的)。通过**遮蔽(Mask)**未来位置来实现:

Maskij={0,i≥j(可见)−∞,i<j(遮蔽) \text{Mask}_{ij} = \begin{cases} 0, & i \geq j \quad \text{(可见)} \\ -\infty, & i < j \quad \text{(遮蔽)} \end{cases} Maskij={0,−∞,i≥j(可见)i<j(遮蔽)

Encoder-Decoder Attention的作用

解码器的每个位置都可以"关注"编码器输出的所有位置。这是Seq2Seq模型的核心------使得解码器能够参考输入序列的信息。

python 复制代码
# ===== 代码实战3.3:实现完整的Encoder Layer =====
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制
    """
    def __init__(self, d_model, num_heads):
        """
        初始化多头注意力
        
        Args:
            d_model: 模型维度
            num_heads: 注意力头数
        """
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_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)
    
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        缩放点积注意力
        
        Args:
            Q: 查询矩阵 (batch, num_heads, seq_len, d_k)
            K: 键矩阵 (batch, num_heads, seq_len, d_k)
            V: 值矩阵 (batch, num_heads, seq_len, d_k)
            mask: 可选的掩码
        
        Returns:
            注意力输出和注意力权重
        """
        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        
        # 应用掩码(如果提供)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # Softmax归一化
        attention_weights = F.softmax(scores, dim=-1)
        
        # 加权求和
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights
    
    def split_heads(self, x):
        """
        将最后一维分割成多个头
        
        Args:
            x: (batch, seq_len, d_model)
        
        Returns:
            (batch, num_heads, seq_len, d_k)
        """
        batch_size, seq_len, _ = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def combine_heads(self, x):
        """
        将多个头合并回最后一维
        
        Args:
            x: (batch, num_heads, seq_len, d_k)
        
        Returns:
            (batch, seq_len, d_model)
        """
        batch_size, num_heads, seq_len, d_k = x.size()
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, self.d_model)
    
    def forward(self, query, key, value, mask=None):
        """
        前向传播
        
        Args:
            query: (batch, seq_len, d_model)
            key: (batch, seq_len, d_model)
            value: (batch, seq_len, d_model)
            mask: 可选的掩码
        
        Returns:
            注意力输出 (batch, seq_len, d_model)
        """
        batch_size = query.size(0)
        
        # 线性变换
        Q = self.split_heads(self.W_q(query))
        K = self.split_heads(self.W_k(key))
        V = self.split_heads(self.W_v(value))
        
        # 计算注意力
        attention_output, attention_weights = self.scaled_dot_product_attention(Q, K, V, mask)
        
        # 合并多头
        output = self.combine_heads(attention_output)
        
        # 输出线性变换
        output = self.W_o(output)
        
        return output, attention_weights


class PositionWiseFFN(nn.Module):
    """
    位置级前馈神经网络
    对每个位置独立地应用相同的FFN
    """
    def __init__(self, d_model, d_ff):
        """
        初始化FFN
        
        Args:
            d_model: 输入和输出维度
            d_ff: 中间层维度(通常设为d_model的4倍)
        """
        super(PositionWiseFFN, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        """
        前向传播
        
        Args:
            x: (batch, seq_len, d_model)
        
        Returns:
            (batch, seq_len, d_model)
        """
        return self.fc2(self.relu(self.fc1(x)))


class EncoderLayer(nn.Module):
    """
    Transformer编码器单层
    """
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        """
        初始化编码器层
        
        Args:
            d_model: 模型维度
            num_heads: 注意力头数
            d_ff: FFN中间层维度
            dropout: Dropout比例
        """
        super(EncoderLayer, self).__init__()
        
        # 多头自注意力
        self.self_attention = MultiHeadAttention(d_model, num_heads)
        
        # 前馈神经网络
        self.ffn = PositionWiseFFN(d_model, d_ff)
        
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        """
        前向传播
        
        Args:
            x: 输入 (batch, seq_len, d_model)
            mask: 可选的注意力掩码
        
        Returns:
            输出 (batch, seq_len, d_model)
        """
        # ===== 子层1: Multi-Head Self-Attention (带残差连接和LayerNorm) =====
        # Pre-Norm架构: LayerNorm在子层之前
        normed_x = self.norm1(x)
        attention_output, _ = self.self_attention(normed_x, normed_x, normed_x, mask)
        x = x + self.dropout(attention_output)  # 残差连接
        
        # ===== 子层2: FFN (带残差连接和LayerNorm) =====
        normed_x = self.norm2(x)
        ffn_output = self.ffn(normed_x)
        x = x + self.dropout(ffn_output)  # 残差连接
        
        return x


class Encoder(nn.Module):
    """
    Transformer编码器(多层堆叠)
    """
    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1):
        """
        初始化编码器
        
        Args:
            num_layers: 编码器层数
            d_model: 模型维度
            num_heads: 注意力头数
            d_ff: FFN中间层维度
            dropout: Dropout比例
        """
        super(Encoder, self).__init__()
        
        # 堆叠多个编码器层
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        # 最后的LayerNorm
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        """
        前向传播(通过所有编码器层)
        
        Args:
            x: 输入 (batch, seq_len, d_model)
            mask: 可选的注意力掩码
        
        Returns:
            编码器输出 (batch, seq_len, d_model)
        """
        for layer in self.layers:
            x = layer(x, mask)
        
        # 最后的归一化
        return self.norm(x)


def test_encoder_implementation():
    """测试编码器实现"""
    print("=== Transformer编码器测试 ===\n")
    
    # 参数设置
    batch_size = 2
    seq_len = 10
    d_model = 512
    num_heads = 8
    d_ff = 2048
    num_layers = 6
    
    # 创建编码器
    encoder = Encoder(num_layers, d_model, num_heads, d_ff)
    
    # 构造输入
    x = torch.randn(batch_size, seq_len, d_model)
    print(f"输入形状: {x.shape}")
    
    # 前向传播
    output = encoder(x)
    print(f"输出形状: {output.shape}")
    print(f"\n说明:")
    print(f"  - 编码器保持了序列长度和维度")
    print(f"  - 每个位置都能'看到'所有其他位置(通过Self-Attention)")
    print(f"  - 残差连接使得梯度可以直接传播到早期层")

if __name__ == "__main__":
    test_encoder_implementation()

3.2 核心机制深度解析

3.2.1 自注意力机制(Self-Attention)数学原理

什么是自注意力?

**自注意力(Self-Attention)**是指序列中的每个元素都与序列中的所有其他元素(包括自己)计算注意力权重。

类比理解:想象一个团队合作项目。在会议中,每个成员都需要考虑其他所有成员的意见(包括自己的),来决定自己下一步该说什么。自注意力就是让序列中的每个词都能够"参考"所有其他词的信息。

数学原理详解

给定输入序列 X=x1,x2,...,xn\mathbf{X} = \\mathbf{x}_1, \\mathbf{x}_2, \\dots, \\mathbf{x}_nX=x1,x2,...,xn,其中 xi∈Rdmodel\mathbf{x}i \in \mathbb{R}^{d{model}}xi∈Rdmodel。

步骤1:生成Query、Key、Value

通过三个可学习的线性变换矩阵,将每个输入向量转换为三个表示:

Q=XWQ,K=XWK,V=XWV \mathbf{Q} = \mathbf{X}W^Q, \quad \mathbf{K} = \mathbf{X}W^K, \quad \mathbf{V} = \mathbf{X}W^V Q=XWQ,K=XWK,V=XWV

其中:

  • WQ∈Rdmodel×dkW^Q \in \mathbb{R}^{d_{model} \times d_k}WQ∈Rdmodel×dk
  • WK∈Rdmodel×dkW^K \in \mathbb{R}^{d_{model} \times d_k}WK∈Rdmodel×dk
  • WV∈Rdmodel×dvW^V \in \mathbb{R}^{d_{model} \times d_v}WV∈Rdmodel×dv

通常 dk=dv=dmodel/num_headsd_k = d_v = d_{model} / \text{num\_heads}dk=dv=dmodel/num_heads。

步骤2:计算注意力分数

计算Query和Key之间的相似度(使用点积):

scoresij=qiTkjdk \text{scores}_{ij} = \frac{\mathbf{q}_i^T \mathbf{k}_j}{\sqrt{d_k}} scoresij=dk qiTkj

为什么要除以 dk\sqrt{d_k}dk ?

假设 qqq 和 kkk 的元素是独立同分布的,均值为0,方差为1。则点积 qTk=∑i=1dkqiki\mathbf{q}^T\mathbf{k} = \sum_{i=1}^{d_k} q_i k_iqTk=∑i=1dkqiki 的方差为 dkd_kdk。除以 dk\sqrt{d_k}dk 可以将方差重新缩放到1,避免点积过大导致Softmax函数进入梯度很小的饱和区。

步骤3:Softmax归一化

αij=exp⁡(scoresij)∑k=1nexp⁡(scoresik) \alpha_{ij} = \frac{\exp(\text{scores}{ij})}{\sum{k=1}^{n} \exp(\text{scores}_{ik})} αij=∑k=1nexp(scoresik)exp(scoresij)

步骤4:加权求和

yi=∑j=1nαijvj \mathbf{y}i = \sum{j=1}^{n} \alpha_{ij} \mathbf{v}_j yi=j=1∑nαijvj

矩阵形式(整个序列并行计算)

Attention(Q,K,V)=Softmax(QKTdk)V \text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=Softmax(dk QKT)V

自注意力的复杂度分析
操作 时间复杂度 空间复杂度
计算 Q,K,VQ, K, VQ,K,V O(ndmodel2)O(nd_{model}^2)O(ndmodel2) O(ndmodel)O(nd_{model})O(ndmodel)
计算 QKTQK^TQKT O(n2dk)O(n^2 d_k)O(n2dk) O(n2)O(n^2)O(n2)
Softmax O(n2)O(n^2)O(n2) -
加权求和 O(n2dv)O(n^2 d_v)O(n2dv) -

关键问题 :Self-Attention的时间和空间复杂度都是 O(n2)O(n^2)O(n2)(nnn 为序列长度)。当 nnn 很大时(如 n=10000n=10000n=10000),计算和存储Attention矩阵会变得非常昂贵。

这也是为什么需要高效注意力机制(如稀疏注意力、Linear Attention等)的原因! 我们将在3.3.4节详细讨论。

python 复制代码
# ===== 代码实战3.4:从零实现Self-Attention =====
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

def manual_self_attention(x):
    """
    手动实现Self-Attention(用于理解原理)
    
    Args:
        x: 输入序列,形状为 (batch, seq_len, d_model)
    
    Returns:
        output: 注意力输出 (batch, seq_len, d_model)
        attention_weights: 注意力权重 (batch, seq_len, seq_len)
    """
    batch_size, seq_len, d_model = x.shape
    d_k = d_model  # 简化:假设d_k = d_model
    
    # 步骤1:生成Q、K、V(使用简单的线性变换)
    W_q = nn.Parameter(torch.randn(d_model, d_model))
    W_k = nn.Parameter(torch.randn(d_model, d_model))
    W_v = nn.Parameter(torch.randn(d_model, d_model))
    
    Q = torch.matmul(x, W_q)  # (batch, seq_len, d_model)
    K = torch.matmul(x, W_k)
    V = torch.matmul(x, W_v)
    
    # 步骤2:计算注意力分数
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    # scores形状: (batch, seq_len, seq_len)
    
    # 步骤3:Softmax归一化
    attention_weights = F.softmax(scores, dim=-1)
    
    # 步骤4:加权求和
    output = torch.matmul(attention_weights, V)
    # output形状: (batch, seq_len, d_model)
    
    return output, attention_weights


def compare_with_pytorch():
    """对比手动实现和PyTorch内置的Self-Attention"""
    print("=== Self-Attention实现对比 ===\n")
    
    batch_size = 2
    seq_len = 5
    d_model = 8
    
    # 构造输入
    x = torch.randn(batch_size, seq_len, d_model)
    
    # 手动实现
    output_manual, weights_manual = manual_self_attention(x)
    print(f"手动实现输出形状: {output_manual.shape}")
    print(f"注意力权重形状: {weights_manual.shape}")
    print(f"\n注意力权重和(应为1.0): {weights_manual[0, 0, :].sum().item():.4f}")
    
    # PyTorch内置实现(使用nn.MultiheadAttention)
    # 注意:PyTorch的MultiheadAttention将d_model分割到多个头
    mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=2, batch_first=True)
    
    output_pytorch, weights_pytorch = mha(x, x, x)
    print(f"\nPyTorch实现输出形状: {output_pytorch.shape}")
    print(f"PyTorch注意力权重形状: {weights_pytorch.shape}")
    
    print("\n说明:")
    print(f"  - 手动实现是单头的Self-Attention")
    print(f"  - PyTorch实现是多头的(num_heads=2)")
    print(f"  - 多头可以让模型同时关注不同位置的不同表示子空间")

if __name__ == "__main__":
    compare_with_pytorch()

3.2.2 多头注意力(Multi-Head Attention)设计哲学

为什么需要多头注意力?

单头注意力的问题

单个注意力头只能捕捉一种类型的依赖关系。例如,在处理句子 "The animal didn't cross the street because it was too tired" 时:

  • 一个注意力头可能关注语法关系(如 "it" 指代 "animal")
  • 另一个注意力头可能关注语义关系(如 "cross" 和 "street" 的搭配)

如果只有一个头,模型需要在同一种表示中平衡所有这些关系,这会降低模型的表达能力。

多头注意力的设计

核心思想 :将 dmodeld_{model}dmodel 维的Query、Key、Value并行地 通过 hhh 个不同的注意力头进行处理,然后将结果拼接起来。

数学表达

给定 hhh 个注意力头,第 iii 个头的输出为:

headi=Attention(QWiQ,KWiK,VWiV) \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) headi=Attention(QWiQ,KWiK,VWiV)

其中 WiQ∈Rdmodel×dkW_i^Q \in \mathbb{R}^{d_{model} \times d_k}WiQ∈Rdmodel×dk, WiK∈Rdmodel×dkW_i^K \in \mathbb{R}^{d_{model} \times d_k}WiK∈Rdmodel×dk, WiV∈Rdmodel×dkW_i^V \in \mathbb{R}^{d_{model} \times d_k}WiV∈Rdmodel×dk,且 dk=dmodel/hd_k = d_{model} / hdk=dmodel/h。

将所有头的输出拼接

MultiHead(Q,K,V)=Concat(head1,head2,...,headh)WO \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) W^O MultiHead(Q,K,V)=Concat(head1,head2,...,headh)WO

其中 WO∈Rhdk×dmodelW^O \in \mathbb{R}^{h d_k \times d_{model}}WO∈Rhdk×dmodel 为输出线性变换矩阵。

多头注意力的优势
  1. 表示子空间的多样性:不同的头可以学习关注不同类型的模式
  2. 增加模型的表达能力 :相当于将 dmodeld_{model}dmodel 维的空间分解为 hhh 个 dkd_kdk 维的子空间
  3. 并行计算:所有头的计算是完全独立的,可以利用GPU并行能力

可视化洞察:通过对训练好的Transformer的注意力权重进行可视化,可以发现:

  • 某些头关注语法关系(如主谓一致)
  • 某些头关注共指关系(如代词指代)
  • 某些头关注位置信息(如相邻词)

这种"分工"是自发产生的,没有显式地告诉模型应该如何分配注意力头。

python 复制代码
# ===== 代码实战3.5:可视化多头注意力的不同头 =====
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention_heads():
    """可视化不同注意力头的注意力权重"""
    print("=== 可视化多头注意力的不同头 ===\n")
    
    # 构造一个简单的句子(用于解释)
    sentence = ["The", "cat", "sat", "on", "the", "mat", "."]
    seq_len = len(sentence)
    
    # 模拟多头注意力的输出
    # 假设我们有4个头,每个头关注不同的模式
    num_heads = 4
    d_model = 64
    
    # 创建模拟的注意力权重
    attention_weights = torch.zeros(num_heads, seq_len, seq_len)
    
    # 头1:关注相邻词
    for i in range(seq_len):
        for j in range(max(0, i-1), min(seq_len, i+2)):
            attention_weights[0, i, j] = 1.0 / min(3, seq_len)
    
    # 头2:关注第一个词(可能是主语)
    attention_weights[1, :, 0] = 0.5
    attention_weights[1, 0, 0] = 0.5
    
    # 头3:关注最后一个词(可能是标点或从句结束)
    attention_weights[2, :, -1] = 0.5
    attention_weights[2, -1, -1] = 0.5
    
    # 头4:均匀关注所有词
    attention_weights[3, :, :] = 1.0 / seq_len
    
    # 可视化
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle('Different Attention Heads Focus on Different Patterns', fontsize=14)
    
    for idx in range(num_heads):
        ax = axes[idx // 2, idx % 2]
        
        # 绘制热力图
        sns.heatmap(
            attention_weights[idx].numpy(),
            annot=True,
            fmt='.2f',
            cmap='YlOrRd',
            xticklabels=sentence,
            yticklabels=sentence,
            ax=ax
        )
        
        ax.set_title(f'Head {idx+1}', fontsize=12)
        ax.set_xlabel('Key (attended to)', fontsize=10)
        ax.set_ylabel('Query (attending)', fontsize=10)
    
    plt.tight_layout()
    plt.savefig('multi_head_attention_visualization.png', dpi=150)
    plt.show()
    
    print("说明:")
    print("  - Head 1: 关注相邻词(局部依赖)")
    print("  - Head 2: 关注第一个词(可能是主语)")
    print("  - Head 3: 关注最后一个词(可能是从句结束)")
    print("  - Head 4: 均匀关注所有词(全局依赖)")
    print("\n  实际训练中,这些模式是模型自动学习的!")

if __name__ == "__main__":
    visualize_attention_heads()

3.2.3 前馈神经网络(FFN)与残差连接

位置级前馈神经网络(Position-wise FFN)

在Self-Attention之后,Transformer对每个位置独立地应用一个全连接前馈神经网络。

数学表达

FFN(x)=max⁡(0,xW1+b1)W2+b2 \text{FFN}(\mathbf{x}) = \max(0, \mathbf{x}W_1 + \mathbf{b}_1)W_2 + \mathbf{b}_2 FFN(x)=max(0,xW1+b1)W2+b2

或者写成两层线性变换 + ReLU激活:

FFN(x)=ReLU(xW1+b1)W2+b2 \text{FFN}(\mathbf{x}) = \text{ReLU}(\mathbf{x}W_1 + \mathbf{b}_1)W_2 + \mathbf{b}_2 FFN(x)=ReLU(xW1+b1)W2+b2

维度变化

  • 输入:x∈Rdmodel\mathbf{x} \in \mathbb{R}^{d_{model}}x∈Rdmodel
  • 中间层:W1∈Rdmodel×dffW_1 \in \mathbb{R}^{d_{model} \times d_{ff}}W1∈Rdmodel×dff
  • 输出:W2∈Rdff×dmodelW_2 \in \mathbb{R}^{d_{ff} \times d_{model}}W2∈Rdff×dmodel

通常 dff=4×dmodeld_{ff} = 4 \times d_{model}dff=4×dmodel(如 dmodel=512d_{model}=512dmodel=512, dff=2048d_{ff}=2048dff=2048)。

为什么需要FFN?

Self-Attention是一个线性变换 (虽然有Softmax,但本质上还是在做加权求和)。如果没有FFN,多层Transformer就等价于单层。FFN引入了非线性(通过ReLU/GELU),使得模型能够学习更复杂的函数。

位置级(Position-wise)的含义

FFN对序列中的每个位置独立地应用相同的变换。这意味着:

FFN(x1),FFN(x2),...,FFN(xn) \text{FFN}(\mathbf{x}_1), \text{FFN}(\mathbf{x}_2), \dots, \text{FFN}(\mathbf{x}_n) FFN(x1),FFN(x2),...,FFN(xn)

这些计算是完全独立的,可以并行化。

残差连接(Residual Connection)

问题:深层网络训练困难(梯度消失/爆炸)。

解决方案 :在每一个子层(Self-Attention和FFN)之后添加残差连接

y=LayerNorm(x+SubLayer(x)) \mathbf{y} = \text{LayerNorm}(\mathbf{x} + \text{SubLayer}(\mathbf{x})) y=LayerNorm(x+SubLayer(x))

其中 SubLayer\text{SubLayer}SubLayer 可以是Self-Attention或FFN。

作用

  1. 梯度可以直接传播 :即使某一层的梯度很小,x\mathbf{x}x 也可以直接传播到前面的层
  2. 允许训练更深的网络:ResNet(He et al., 2016)证明了残差连接可以训练上千层的网络
  3. 提供"身份路径":如果某一层没有必要,模型可以学习将其输出设为0,直接传递输入

Post-Norm vs Pre-Norm

  • 原始Transformer(Post-Norm)x → SubLayer → Add → LayerNorm
  • 现代Transformer(Pre-Norm)x → LayerNorm → SubLayer → Add

Pre-Norm在训练稳定性上表现更好,是现代大模型(如GPT-3、LLaMA)的默认选择。

python 复制代码
# ===== 代码实战3.6:实现带残差连接的FFN =====
import torch
import torch.nn as nn

class FeedForwardNetwork(nn.Module):
    """
    前馈神经网络(带残差连接和LayerNorm)
    """
    def __init__(self, d_model, d_ff, dropout=0.1, pre_norm=True):
        """
        初始化FFN
        
        Args:
            d_model: 模型维度
            d_ff: 中间层维度
            dropout: Dropout比例
            pre_norm: 是否使用Pre-Norm架构
        """
        super(FeedForwardNetwork, self).__init__()
        
        self.pre_norm = pre_norm
        
        # FFN的两层线性变换
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        
        # 激活函数(现代Transformer常用GELU)
        self.activation = nn.GELU()  # 也可以用nn.ReLU()
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
        
        # LayerNorm
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x):
        """
        前向传播(带残差连接)
        
        Args:
            x: 输入 (batch, seq_len, d_model)
        
        Returns:
            输出 (batch, seq_len, d_model)
        """
        if self.pre_norm:
            # ===== Pre-Norm架构 =====
            # 先归一化,再FFN,最后残差连接
            normed_x = self.norm(x)
            
            # FFN
            ffn_output = self.fc2(self.dropout(self.activation(self.fc1(normed_x))))
            
            # 残差连接
            output = x + ffn_output
        
        else:
            # ===== Post-Norm架构(原始Transformer)=====
            # 先FFN,再残差连接,最后归一化
            ffn_output = self.fc2(self.dropout(self.activation(self.fc1(x))))
            
            # 残差连接
            residual = x + ffn_output
            
            # LayerNorm
            output = self.norm(residual)
        
        return output


def compare_pre_norm_post_norm():
    """对比Pre-Norm和Post-Norm的训练稳定性"""
    print("=== Pre-Norm vs Post-Norm ===\n")
    
    d_model = 512
    d_ff = 2048
    batch_size = 2
    seq_len = 10
    
    # 创建Pre-Norm FFN
    ffn_pre_norm = FeedForwardNetwork(d_model, d_ff, pre_norm=True)
    
    # 创建Post-Norm FFN
    ffn_post_norm = FeedForwardNetwork(d_model, d_ff, pre_norm=False)
    
    # 构造输入
    x = torch.randn(batch_size, seq_len, d_model)
    
    # 前向传播
    output_pre = ffn_pre_norm(x)
    output_post = ffn_post_norm(x)
    
    print(f"输入形状: {x.shape}")
    print(f"Pre-Norm输出形状: {output_pre.shape}")
    print(f"Post-Norm输出形状: {output_post.shape}")
    
    print("\n说明:")
    print(f"  - Pre-Norm: 梯度更稳定,适合训练深层模型")
    print(f"  - Post-Norm: 原始Transformer使用,需要仔细的warmup策略")
    print(f"  - 现代大模型(GPT-3、LLaMA)都使用Pre-Norm")

if __name__ == "__main__":
    compare_pre_norm_post_norm()

3.2.4 层归一化(Layer Normalization)的作用与变体

什么是Layer Normalization?

Batch Normalization(BN)的问题

BN在CNN中非常有效,但在Transformer(尤其是RNN/Transformer)中效果不佳,因为:

  1. 依赖于批量大小:小批量时统计量不准确
  2. 训练和推理行为不一致:训练时使用当前batch的统计量,推理时使用滑动平均
  3. 不适合变长序列:不同样本的序列长度可能不同

Layer Normalization(LN)的解决方案

LN对每个样本的特征维度进行归一化,而不依赖于batch。

数学表达

给定输入 x=x1,x2,...,xH\mathbf{x} = x_1, x_2, \\dots, x_Hx=x1,x2,...,xH(一个样本的所有特征),LN计算:

μ=1H∑i=1Hxi \mu = \frac{1}{H}\sum_{i=1}^{H} x_i μ=H1i=1∑Hxi

σ2=1H∑i=1H(xi−μ)2 \sigma^2 = \frac{1}{H}\sum_{i=1}^{H} (x_i - \mu)^2 σ2=H1i=1∑H(xi−μ)2

x^i=xi−μσ2+ϵ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} x^i=σ2+ϵ xi−μ

yi=γix^i+βi y_i = \gamma_i \hat{x}_i + \beta_i yi=γix^i+βi

其中 γ\gammaγ 和 β\betaβ 是可学习的缩放和偏移参数。

与Batch Normalization的对比

特性 Batch Norm Layer Norm
归一化维度 跨batch的同一特征 每个样本的所有特征
训练/推理一致性 需要维护滑动平均 完全一致
适用于变长序列
适用于RNN/Transformer
Layer Normalization的变体

1. RMSNorm(Root Mean Square Layer Normalization)

GPT-2之后的一些模型(如GPT-J、LLaMA)使用RMSNorm,它是LayerNorm的简化版本。

数学表达

x^i=xi1H∑j=1Hxj2+ϵ \hat{x}i = \frac{x_i}{\sqrt{\frac{1}{H}\sum{j=1}^{H} x_j^2 + \epsilon}} x^i=H1∑j=1Hxj2+ϵ xi

yi=γix^i y_i = \gamma_i \hat{x}_i yi=γix^i

优势

  • 不需要计算均值 μ\muμ,计算更快
  • 在实践中性能与LayerNorm相当

2. Adaptive Layer Norm(AdaLN)

在扩散模型(如Stable Diffusion)中使用的变体,允许条件信息(如文本嵌入)动态调整LayerNorm的参数。

为什么LayerNorm在Transformer中如此重要?
  1. 稳定训练:归一化使得每一层的输入分布保持稳定,加速收敛
  2. 允许更大的学习率:没有LN,大学习率会导致训练不稳定
  3. 减少梯度消失/爆炸:归一化使得梯度保持在合理范围内

⚠️ 企业级避坑1:LayerNorm的位置

在原始Transformer中,LayerNorm位于残差连接之后(Post-Norm):

复制代码
x → Attention → Add → LayerNorm → FFN → Add → LayerNorm

在现代Transformer中,LayerNorm位于子层之前(Pre-Norm):

复制代码
x → LayerNorm → Attention → Add → LayerNorm → FFN → Add

Pre-Norm的训练稳定性显著更好,是现代大模型的默认选择。如果你在训练自己的Transformer时遇到不稳定问题,首先检查是否使用了Pre-Norm!

python 复制代码
# ===== 代码实战3.7:实现RMSNorm =====
import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization
    简化版的LayerNorm,计算更快
    """
    def __init__(self, d_model, eps=1e-6):
        """
        初始化RMSNorm
        
        Args:
            d_model: 模型维度
            eps: 数值稳定常数
        """
        super(RMSNorm, self).__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(d_model))  # 可学习的缩放参数
    
    def forward(self, x):
        """
        前向传播
        
        Args:
            x: 输入 (batch, seq_len, d_model)
        
        Returns:
            归一化后的输出
        """
        # 计算RMS
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        
        # 归一化
        normalized_x = x / rms
        
        # 缩放
        output = self.gamma * normalized_x
        
        return output


def compare_ln_rmsnorm():
    """对比LayerNorm和RMSNorm"""
    print("=== LayerNorm vs RMSNorm ===\n")
    
    d_model = 512
    batch_size = 2
    seq_len = 10
    
    # 创建LayerNorm和RMSNorm
    ln = nn.LayerNorm(d_model)
    rms_norm = RMSNorm(d_model)
    
    # 构造输入
    x = torch.randn(batch_size, seq_len, d_model)
    
    # 前向传播
    output_ln = ln(x)
    output_rms = rms_norm(x)
    
    print(f"输入形状: {x.shape}")
    print(f"LayerNorm输出形状: {output_ln.shape}")
    print(f"RMSNorm输出形状: {output_rms.shape}")
    
    # 检查均值和方差
    print(f"\nLayerNorm输出均值: {output_ln.mean(dim=-1).mean().item():.6f} (接近0)")
    print(f"LayerNorm输出方差: {output_ln.var(dim=-1).mean().item():.6f} (接近1)")
    
    print(f"\nRMSNorm输出RMS: {torch.sqrt(torch.mean(output_rms**2, dim=-1)).mean().item():.6f}")
    print(f"  (RMSNorm不强制均值为0,只归一化RMS)")
    
    print("\n说明:")
    print(f"  - LayerNorm: 计算均值和方差,更准确")
    print(f"  - RMSNorm: 只计算RMS,计算更快")
    print(f"  - 在大模型中,RMSNorm的速度优势更明显")

if __name__ == "__main__":
    compare_ln_rmsnorm()

3.3 Transformer变体与优化

3.3.1 Encoder-only架构:BERT系列

BERT:双向编码器表征

BERT(Bidirectional Encoder Representations from Transformers)是Encoder-only架构的代表。

核心特点

  1. 双向上下文:同时利用左侧和右侧的上下文(通过Self-Attention)
  2. 预训练目标
    • Masked Language Modeling(MLM):随机遮盖15%的词,预测被遮盖的词
    • Next Sentence Prediction(NSP):预测两个句子是否连续
  3. 适合理解任务:文本分类、命名实体识别、问答系统等

架构图

复制代码
输入: [CLS] 句子1 [SEP] 句子2 [SEP]
    ↓
Embedding + Segment Embedding + Position Embedding
    ↓
[Encoder Layer 1] → [Encoder Layer 2] → ... → [Encoder Layer N]
    ↓
输出: [h_CLS, h_1, h_2, ..., h_n]
    ↓
[CLS] token的表示用于分类任务

CLS token的作用 :BERT在输入序列前添加一个特殊的 [CLS] token,其对应的输出向量被用作整个序列的聚合表示,用于分类任务。

BERT家族
模型 参数量 层数 隐藏维度 特点
BERT-base 110M 12 768 原始BERT基准模型
BERT-large 340M 24 1024 更大的BERT
RoBERTa 125M 12 768 去掉NSP,使用更多数据训练
ALBERT 12M 12 768 参数共享,减少参数量
ELECTRA 110M 12 768 使用判别式预训练目标
python 复制代码
# ===== 代码实战3.8:使用Hugging Face加载BERT =====
from transformers import BertTokenizer, BertModel
import torch

def demonstrate_bert_usage():
    """演示BERT的使用"""
    print("=== BERT使用示例 ===\n")
    
    # 加载预训练的BERT模型和分词器
    model_name = 'bert-base-uncased'
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertModel.from_pretrained(model_name)
    
    print(f"模型: {model_name}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M\n")
    
    # 准备输入
    sentences = ["I love natural language processing.", "Transformer is amazing!"]
    
    # 分词
    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
    print(f"分词结果:")
    print(f"  input_ids形状: {encoded_input['input_ids'].shape}")
    print(f"  attention_mask形状: {encoded_input['attention_mask'].shape}")
    
    # 前向传播
    with torch.no_grad():
        output = model(**encoded_input)
    
    # 获取最后一层的隐藏状态
    last_hidden_states = output.last_hidden_state
    print(f"\n隐藏状态形状: {last_hidden_states.shape}")
    print(f"  解释: (batch_size={last_hidden_states.shape[0]}, "
          f"seq_len={last_hidden_states.shape[1]}, "
          f"hidden_dim={last_hidden_states.shape[2]})")
    
    # 使用[CLS] token的输出进行分类
    cls_output = last_hidden_states[:, 0, :]  # [CLS]是第一个token
    print(f"\n[CLS]输出形状: {cls_output.shape}")
    print(f"  可用于下游分类任务")

if __name__ == "__main__":
    demonstrate_bert_usage()

3.3.2 Decoder-only架构:GPT系列

GPT:生成式预训练Transformer

GPT(Generative Pre-trained Transformer)是Decoder-only架构的代表。

核心特点

  1. 单向上下文:只能看到左侧的上下文(通过Masked Self-Attention)
  2. 预训练目标Causal Language Modeling(CLM)------预测下一个词
  3. 适合生成任务:文本生成、对话系统、代码补全等

架构图

复制代码
输入: [s_1, s_2, ..., s_n]
    ↓
Embedding + Positional Encoding
    ↓
[Decoder Layer 1] → [Decoder Layer 2] → ... → [Decoder Layer N]
    ↓ (只使用Masked Self-Attention和FFN,没有Encoder-Decoder Attention)
输出: [h_1, h_2, ..., h_n]
    ↓
Linear + Softmax → 预测下一个词

Masked Self-Attention

通过上三角矩阵掩码,确保每个位置只能看到它之前的词:

Maskij={0,i≥j(可见)−∞,i<j(遮蔽) \text{Mask}_{ij} = \begin{cases} 0, & i \geq j \quad \text{(可见)} \\ -\infty, & i < j \quad \text{(遮蔽)} \end{cases} Maskij={0,−∞,i≥j(可见)i<j(遮蔽)

GPT家族演进
模型 参数量 层数 发布时间 特点
GPT-1 117M 12 2018.06 首个GPT模型,证明了预训练-微调范式
GPT-2 1.5B 48 2019.02 零样本学习能力的初步展现
GPT-3 175B 96 2020.05 少样本学习,涌现能力
GPT-3.5 (InstructGPT) 175B 96 2022.03 引入RLHF,对话能力大幅提升
GPT-4 ~1.8T (MoE) - 2023.03 多模态,推理能力大幅提升

缩放定律(Scaling Law):GPT系列展示了模型性能随参数量、数据量、计算量的幂律增长。这为大模型的发展提供了理论指导。

python 复制代码
# ===== 代码实战3.9:使用Hugging Face加载GPT-2 =====
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

def demonstrate_gpt2_generation():
    """演示GPT-2的文本生成"""
    print("=== GPT-2文本生成示例 ===\n")
    
    # 加载预训练的GPT-2模型和分词器
    model_name = 'gpt2'
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)
    
    print(f"模型: {model_name}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M\n")
    
    # 准备输入
    prompt = "Once upon a time"
    input_ids = tokenizer.encode(prompt, return_tensors='pt')
    print(f"输入: {prompt}")
    print(f"输入IDs形状: {input_ids.shape}\n")
    
    # 生成文本
    print("生成结果:")
    output_ids = model.generate(
        input_ids,
        max_length=50,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        early_stopping=True
    )
    
    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    print(f"  {generated_text}")

if __name__ == "__main__":
    demonstrate_gpt2_generation()

3.3.3 Encoder-Decoder架构:T5、BART等

T5(Text-to-Text Transfer Transformer)

T5将所有NLP任务统一为文本到文本的格式。

核心思想

复制代码
输入: "translate English to German: Hello, how are you?"
输出: "Hallo, wie geht es dir?"

输入: "summarize: [文章文本]"
输出: "[摘要文本]"

输入: "cola acceptability: First, the car drove down the street."
输出: "acceptable"

架构:完整的Encoder-Decoder Transformer。

预训练目标Masked Language Modeling(类似BERT,但是Encoder-Decoder架构)。

BART(Bidirectional and Auto-Regressive Transformer)

BART结合了BERT的双向编码器和GPT的单向解码器。

预训练目标噪声函数(如token masking、token deletion、text infilling等)。

适合任务:文本生成、摘要、翻译等。

Encoder-Decoder vs Encoder-only vs Decoder-only
架构 代表模型 适合任务 计算成本
Encoder-only BERT, RoBERTa 理解任务(分类、NER、QA)
Decoder-only GPT系列, LLaMA 生成任务(对话、续写、代码)
Encoder-Decoder T5, BART, mT5 序列到序列任务(翻译、摘要)

为什么GPT成为大模型时代的主流?

虽然Encoder-Decoder架构在序列到序列任务上表现出色,但Decoder-only架构在以下方面更具优势:

  1. 统一的生成框架:所有任务都可以表示为"预测下一个词"
  2. 更好的零样本/少样本学习能力:通过Prompt可以适配各种任务
  3. 更适合大规模训练:架构简单,易于并行化
  4. 推理效率高:自回归生成可以逐步进行,易于控制

这也是为什么GPT-3、LLaMA、ChatGLM等主流大模型都采用Decoder-only架构。


3.3.4 高效注意力机制

问题:Self-Attention的O(n²)复杂度

标准Self-Attention的时间和空间复杂度都是 O(n2)O(n^2)O(n2)(nnn 为序列长度)。

当 nnn 很大时(如长文档理解、基因组序列分析等),标准Attention会变得非常昂贵。

解决方案 :设计高效注意力机制 ,将复杂度降低到 O(nlog⁡n)O(n \log n)O(nlogn) 或 O(n)O(n)O(n)。

稀疏注意力(Sparse Attention)

核心思想 :每个token只关注一部分其他token,而非全部。

方案

  1. Fixed Pattern Sparse Attention(固定模式稀疏注意力)

    • 每个token只关注固定窗口内的token(如前后各64个token)
    • 复杂度:O(n×w)O(n \times w)O(n×w),其中 www 为窗口大小
  2. Longformer(Beltagy et al., 2020)

    • 结合局部窗口注意力和全局token
    • 适合长文档理解
  3. BigBird(Zaheer et al., 2020)

    • 理论证明了稀疏注意力可以近似全注意力
    • 使用随机图和全局token
线性注意力(Linear Attention)

核心思想 :重新排列Attention计算的顺序,避免显式计算 n×nn \times nn×n 的注意力矩阵。

数学原理

标准Attention:

Attention(Q,K,V)=Softmax(QKTdk)V \text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=Softmax(dk QKT)V

问题:Softmax(QKTdk)\text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)Softmax(dk QKT) 的形状是 n×nn \times nn×n,计算和存储都很昂贵。

Linear Attention的解决方案

使用一个核函数(Kernel Function) ϕ(⋅)\phi(\cdot)ϕ(⋅) 来近似Softmax:

LinearAttention(Q,K,V)=ϕ(Q)(ϕ(K)TV)ϕ(Q)(ϕ(K)T1) \text{LinearAttention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)(\phi(K)^T \mathbf{1})} LinearAttention(Q,K,V)=ϕ(Q)(ϕ(K)T1)ϕ(Q)(ϕ(K)TV)

关键技巧 :通过重新排列矩阵乘法顺序,将复杂度从 O(n2d)O(n^2 d)O(n2d) 降低到 O(nd2)O(n d^2)O(nd2):

复制代码
标准Attention: (n×d) × (d×n) → (n×n) × (n×d) = O(n²d)
Linear Attention: (n×d) × ((d×n) × (n×d)) × (d×n) = O(nd²)

当 d≪nd \ll nd≪n 时,Linear Attention显著更快。

Flash Attention

核心思想:通过**IO感知(IO-Aware)**的算法设计,减少GPU高带宽内存(HBM)的读写次数。

关键技术

  1. 分块计算(Tiling):将Attention矩阵分块计算,每个块适合GPU SRAM
  2. 重新计算(Recomputation):在反向传播时重新计算注意力矩阵,而非存储(节省显存)

效果

  • 显存占用从 O(n2)O(n^2)O(n2) 降低到 O(n)O(n)O(n)
  • 实际运行速度提升2-4倍

Flash Attention已经成为现代大模型训练/推理的标配! PyTorch 2.0+已经内置了Flash Attention支持。

python 复制代码
# ===== 代码实战3.10:使用Flash Attention =====
import torch
import torch.nn as nn

def demonstrate_flash_attention():
    """演示Flash Attention的使用"""
    print("=== Flash Attention示例 ===\n")
    
    # 检查PyTorch是否支持Flash Attention
    if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
        print("✅ PyTorch版本支持Flash Attention!")
        
        # 构造输入
        batch_size = 2
        seq_len = 1024  # 长序列
        num_heads = 8
        d_k = 64
        
        Q = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
        K = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
        V = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
        
        # 使用PyTorch内置的scaled_dot_product_attention
        # 它会自动选择最优实现(包括Flash Attention)
        output = torch.nn.functional.scaled_dot_product_attention(Q, K, V)
        
        print(f"输入形状: Q={Q.shape}")
        print(f"输出形状: {output.shape}")
        print("\n说明:")
        print(f"  - PyTorch会自动检测是否支持Flash Attention")
        print(f"  - 如果支持,会自动使用(无需额外代码)")
        print(f"  - Flash Attention可以显著减少显存占用并加速训练")
    else:
        print("❌ 当前PyTorch版本不支持Flash Attention")
        print("   请升级到PyTorch 2.0+")

if __name__ == "__main__":
    demonstrate_flash_attention()

3.4 大模型训练技术

3.4.1 预训练目标函数

语言建模目标

1. Causal Language Modeling(CLM)

目标:给定前面的词,预测下一个词。

LCLM=−∑t=1Tlog⁡P(wt∣w1,...,wt−1) \mathcal{L}{\text{CLM}} = -\sum{t=1}^{T} \log P(w_t | w_1, \dots, w_{t-1}) LCLM=−t=1∑TlogP(wt∣w1,...,wt−1)

使用模型:GPT系列、LLaMA等Decoder-only模型。

2. Masked Language Modeling(MLM)

目标:随机遮盖部分词,预测被遮盖的词。

LMLM=−∑i∈Mlog⁡P(wi∣w∖M) \mathcal{L}{\text{MLM}} = -\sum{i \in \mathcal{M}} \log P(w_i | \mathbf{w}_{\setminus \mathcal{M}}) LMLM=−i∈M∑logP(wi∣w∖M)

其中 M\mathcal{M}M 为被遮盖的位置集合。

使用模型:BERT等Encoder-only模型。

3. Prefix Language Modeling(PLM)

目标:结合CLM和MLM,部分词可见,部分词需要预测。

使用模型:T5、BART等Encoder-Decoder模型。

去噪自编码目标

1. Denoising Autoencoder(DAE)

目标:向输入添加噪声(如随机遮盖、删除、打乱),让模型重建原始输入。

使用模型:BART、mBART。

2. Permutation Language Modeling(PLM)

目标:随机排列输入序列的顺序,然后预测被排列位置的词。

使用模型:XLNet。

对比学习目标

1. Contrastive Learning

目标:拉近正样本对的距离,推远负样本对的距离。

使用模型:SimCSE、Sentence-BERT。


3.4.2 分布式训练策略

问题:大模型无法放入单GPU显存

以GPT-3(175B参数)为例:

  • FP32精度:~700GB显存
  • 即使是A100(80GB显存),也无法放入单卡

解决方案分布式训练,将模型或数据分布到多个GPU/机器上。

数据并行(Data Parallelism, DP)

核心思想 :将数据分割成多个子集,每个GPU都有完整的模型副本,处理不同的数据子集。

流程

复制代码
输入数据: [Batch 1, Batch 2, ..., Batch N]
    ↓ (分割)
GPU 1: [Batch 1] → 模型副本1 → 梯度1
GPU 2: [Batch 2] → 模型副本2 → 梯度2
    ↓ (梯度聚合)
All-Reduce: 平均梯度 → 更新所有模型副本

优势 :实现简单,适合模型较小但数据量大的场景

局限:模型必须能放入单GPU显存

模型并行(Model Parallelism, MP)

核心思想 :将模型分割成多个部分,每个GPU负责一部分模型参数。

流程(以两层Transformer为例):

复制代码
输入: [x₁, x₂, ..., xₙ]
    ↓
GPU 1: Encoder Layer 1 → 中间结果
    ↓ (通过网络传输)
GPU 2: Encoder Layer 2 → 输出

优势 :可以训练超大模型

局限:GPU之间需要频繁通信,训练速度慢

流水线并行(Pipeline Parallelism, PP)

核心思想 :将模型的不同层 分配到不同GPU,并通过流水线调度提高GPU利用率。

流程

复制代码
时间步1: GPU1处理Batch1的Layer1-4
时间步2: GPU1处理Batch2的Layer1-4, GPU2处理Batch1的Layer5-8
时间步3: GPU1处理Batch3, GPU2处理Batch2, GPU3处理Batch1

优势 :提高了GPU利用率

局限:需要仔细设计流水线调度策略

张量并行(Tensor Parallelism, TP)

核心思想 :将每一层的参数矩阵分割成多个分片,分布到不同GPU上。

示例(以线性层为例):

Y=XW Y = XW Y=XW

其中 WWW 可以按列分割:

W=W1,W2,Y=XW1,XW2 W = W_1, W_2, \quad Y = XW_1, XW_2 W=W1,W2,Y=XW1,XW2

GPU1计算 XW1XW_1XW1,GPU2计算 XW2XW_2XW2,最后拼接结果。

优势 :细粒度并行,适合超大模型

局限:GPU间通信开销大

3D并行(DP + PP + TP)

现代大模型训练通常组合使用多种并行策略:

复制代码
总GPU数 = DP度数 × PP度数 × TP度数

示例(GPT-3训练配置):

  • 数据并行:64路
  • 流水线并行:4路
  • 张量并行:8路
  • 总GPU数:64 × 4 × 8 = 2048块GPU
python 复制代码
# ===== 代码实战3.11:使用PyTorch实现数据并行 =====
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import os

def setup_distributed_training():
    """设置分布式训练环境(简化示例)"""
    print("=== 分布式训练示例 ===\n")
    
    # 检查是否有多GPU
    if torch.cuda.device_count() > 1:
        print(f"✅ 检测到 {torch.cuda.device_count()} 个GPU!")
        print("   可以使用DataParallel或DistributedDataParallel\n")
        
        # 简单的DataParallel示例
        model = nn.Transformer(
            d_model=512,
            nhead=8,
            num_encoder_layers=6,
            num_decoder_layers=6
        )
        
        if torch.cuda.is_available():
            model = model.cuda()
            model = nn.DataParallel(model)  # 自动使用所有可用GPU
            
            print(f"模型已包装为DataParallel")
            print(f"  使用GPU数量: {torch.cuda.device_count()}")
    else:
        print("❌ 只有1个GPU或没有GPU")
        print("   无法演示分布式训练")

if __name__ == "__main__":
    setup_distributed_training()

3.4.3 混合精度训练与梯度累积

混合精度训练(Mixed Precision Training)

核心思想 :在训练过程中混合使用不同数值精度(如FP16和FP32),以加速训练并减少显存占用。

为什么需要混合精度?

  1. FP32训练慢:FP32矩阵运算的吞吐量是FP16的1/2~1/8
  2. FP16易溢出:FP16的表示范围有限,容易出现下溢(梯度为0)或上溢(梯度为inf)
  3. 解决方案 :使用FP16进行计算 ,但使用FP32存储权重和梯度

技术细节

  1. FP16前向传播和反向传播:加速计算
  2. FP32权重副本(Master Weights):存储高精度权重,用于更新
  3. Loss Scaling:将损失乘以一个缩放因子(如1024),避免梯度下溢

效果

  • 训练速度提升:1.5~3倍
  • 显存占用减少:~50%

AMP(Automatic Mixed Precision) :PyTorch提供了 torch.cuda.amp 自动混合精度训练工具,无需手动管理精度转换。

python 复制代码
# ===== 代码实战3.12:使用AMP进行混合精度训练 =====
import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler

def demonstrate_mixed_precision_training():
    """演示混合精度训练"""
    print("=== 混合精度训练示例 ===\n")
    
    # 创建模型
    model = nn.Transformer(
        d_model=512,
        nhead=8,
        num_encoder_layers=6,
        num_decoder_layers=6
    )
    
    if torch.cuda.is_available():
        model = model.cuda()
        print("✅ 模型已移动到GPU\n")
        
        # 定义优化器
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        
        # 创建GradScaler(用于Loss Scaling)
        scaler = GradScaler()
        
        # 模拟训练循环
        print("训练循环中...")
        
        # 构造假数据
        src = torch.randn(32, 10, 512).cuda()
        tgt = torch.randn(32, 20, 512).cuda()
        
        for epoch in range(3):
            optimizer.zero_grad()
            
            # ===== 使用autocast进行混合精度前向传播 =====
            with autocast():
                output = model(src, tgt)
                loss = output.mean()  # 假损失
            
            # ===== 使用GradScaler进行混合精度反向传播 =====
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            
            print(f"  Epoch {epoch+1}: Loss = {loss.item():.4f}")
        
        print("\n说明:")
        print(f"  - autocast(): 自动将计算转换为FP16")
        print(f"  - GradScaler: 自动进行Loss Scaling,避免梯度下溢")
        print(f"  - 混合精度训练可以显著加速训练并减少显存占用")
    else:
        print("❌ 没有GPU,无法演示")

if __name__ == "__main__":
    demonstrate_mixed_precision_training()
梯度累积(Gradient Accumulation)

问题:当GPU显存不足以支持大Batch Size时,如何模拟大Batch训练?

解决方案梯度累积------在多次前向-反向传播中累积梯度,然后一次性更新参数。

数学原理

假设目标Batch Size为 BBB,但GPU显存只能支持 B/kB/kB/k。

标准训练 (Batch Size = BBB):

g=1B∑i=1B∇L(xi,yi) \mathbf{g} = \frac{1}{B} \sum_{i=1}^{B} \nabla \mathcal{L}(x_i, y_i) g=B1i=1∑B∇L(xi,yi)

wt+1=wt−ηg \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \mathbf{g} wt+1=wt−ηg

梯度累积 (累积 kkk 步):

gj=1B/k∑i∈mini-batchj∇L(xi,yi),j=1,...,k \mathbf{g}j = \frac{1}{B/k} \sum{i \in \text{mini-batch}_j} \nabla \mathcal{L}(x_i, y_i), \quad j = 1, \dots, k gj=B/k1i∈mini-batchj∑∇L(xi,yi),j=1,...,k

gaccumulated=1k∑j=1kgj \mathbf{g}{\text{accumulated}} = \frac{1}{k} \sum{j=1}^{k} \mathbf{g}_j gaccumulated=k1j=1∑kgj

wt+1=wt−ηgaccumulated \mathbf{w}_{t+1} = \mathbf{w}t - \eta \mathbf{g}{\text{accumulated}} wt+1=wt−ηgaccumulated

效果 :等效于使用Batch Size = BBB 的训练,但显存占用只有 B/kB/kB/k。

python 复制代码
# ===== 代码实战3.13:实现梯度累积 =====
import torch
import torch.nn as nn

def demonstrate_gradient_accumulation():
    """演示梯度累积"""
    print("=== 梯度累积示例 ===\n")
    
    # 创建模型
    model = nn.Linear(10, 2)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    # 目标Batch Size: 128
    # 实际Batch Size: 32
    # 累积步数: 4
    target_batch_size = 128
    actual_batch_size = 32
    accumulation_steps = target_batch_size // actual_batch_size
    
    print(f"目标Batch Size: {target_batch_size}")
    print(f"实际Batch Size: {actual_batch_size}")
    print(f"累积步数: {accumulation_steps}\n")
    
    # 模拟训练循环
    print("训练循环中...")
    model.train()
    
    for epoch in range(3):
        optimizer.zero_grad()
        
        # 累积梯度
        total_loss = 0.0
        for step in range(accumulation_steps):
            # 构造假数据
            x = torch.randn(actual_batch_size, 10)
            y = torch.randint(0, 2, (actual_batch_size,))
            
            # 前向传播
            logits = model(x)
            loss = nn.functional.cross_entropy(logits, y)
            
            # 反向传播(不立即更新参数)
            loss = loss / accumulation_steps  # 平均损失
            loss.backward()
            
            total_loss += loss.item()
        
        # 所有累积步完成后,更新参数
        optimizer.step()
        optimizer.zero_grad()
        
        print(f"  Epoch {epoch+1}: Loss = {total_loss:.4f}")
    
    print("\n说明:")
    print(f"  - 梯度累积允许使用大Batch Size训练,即使显存有限")
    print(f"  - 注意:损失需要除以累积步数,以保持数值等价性")

if __name__ == "__main__":
    demonstrate_gradient_accumulation()

3.4.4 大模型训练的稳定性挑战与解决方案

挑战1:梯度爆炸/消失

问题:在深层网络中,梯度可能指数级增长(爆炸)或指数级衰减(消失)。

解决方案

  1. 梯度裁剪(Gradient Clipping)
python 复制代码
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
  1. 使用Pre-Norm架构:如前所述,Pre-Norm比Post-Norm更稳定。

  2. 合适的权重初始化:如使用Xavier初始化或Kaiming初始化。

挑战2:Loss Spike(损失尖峰)

问题:在训练大模型时,Loss会突然出现很大的尖峰,然后恢复。

原因:可能是梯度爆炸,或者是某些batch的数据分布异常。

解决方案

  1. 降低学习率
  2. 使用Warmup:在训练初期逐渐提高学习率
  3. 使用AdamW优化器:AdamW对超参数的选择更鲁棒
挑战3:显存不足(OOM)

问题:GPU显存不足以支持训练。

解决方案

  1. 使用梯度累积(如前所述)
  2. 使用混合精度训练(如前所述)
  3. 使用梯度检查点(Gradient Checkpointing)
python 复制代码
from torch.utils.checkpoint import checkpoint

# 在前向传播时不保存中间激活值,在反向传播时重新计算
output = checkpoint(model, input)
  1. 减少模型规模 :如使用更小的 dmodeld_{model}dmodel 或层数

⚠️ 企业级避坑2:大模型训练的调试技巧

训练大模型时,由于训练时间长(可能数周或数月),调试非常困难。以下是一些实用技巧:

  1. 使用小规模实验验证代码:在训练全规模模型前,先用小模型和小数据验证代码正确性
  2. 保存检查点(Checkpoint):定期保存模型参数,以便从故障中恢复
  3. 监控训练指标:使用TensorBoard或Wandb监控Loss、梯度范数、学习率等
  4. 使用确定性算法:设置随机种子,确保实验可复现
  5. 渐进式增加规模:先训练小模型,验证无误后再扩展到大规模

企业级考量与避坑指南

避坑1:位置编码选择不当导致长序列性能下降

问题描述

不同的位置编码方案对长序列(如长度 > 2048)的处理能力不同。

解决方案

位置编码方案 外推性 适合场景
Sinusoidal ✅ 好 训练时长度固定,推理时可能需要更长序列
RoPE ✅ 好(当前主流) 大模型(LLaMA、ChatGLM等)
ALiBi ✅ 非常好 需要极强外推能力的场景
可学习位置编码 ❌ 差 训练和推理长度必须相同

经验法则

  • 如果需要在推理时处理比训练时更长的序列,使用RoPEALiBi
  • 如果训练和推理长度相同,可以使用可学习位置编码

避坑2:注意力头数设置不当导致计算效率低下

问题描述

多头注意力的头数 hhh 必须能够整除 dmodeld_{model}dmodel。但即使满足这个条件,不同的 hhh 选择也会影响计算效率。

解决方案

  1. 头数选择 :通常 h=8,12,16h = 8, 12, 16h=8,12,16 等
  2. 每个头的维度 :dk=dmodel/hd_k = d_{model} / hdk=dmodel/h,通常 dk≥32d_k \geq 32dk≥32(否则表示能力不足)
  3. 硬件效率 :某些GPU对特定的 hhh 有更好的支持(如Tensor Core要求 hhh 是8的倍数)

经验法则

  • 对于 dmodel=512d_{model} = 512dmodel=512,可以选择 h=8h = 8h=8(dk=64d_k = 64dk=64)或 h=16h = 16h=16(dk=32d_k = 32dk=32)
  • 对于 dmodel=768d_{model} = 768dmodel=768,可以选择 h=12h = 12h=12(dk=64d_k = 64dk=64)

避坑3:未考虑推理时的内存优化

问题描述

训练时可以通过梯度累积、混合精度等技术减少显存占用,但推理时没有梯度,这些技术不再适用。如果推理时显存不足,无法部署模型。

解决方案

  1. 使用KV Cache:在自回归生成时,缓存已计算的Key和Value,避免重复计算
  2. 使用量化:将模型权重量化到INT8或INT4
  3. 使用模型并行:将模型分布到多个GPU上推理
  4. 使用推理加速框架:如vLLM、TensorRT-LLM等

⚠️ 企业级避坑3:推理吞吐量与延迟的权衡

在部署大模型时,需要在吞吐量 (Throughput)和延迟(Latency)之间进行权衡:

  • 高吞吐量:使用更大的Batch Size,但延迟更高
  • 低延迟:使用更小的Batch Size,但吞吐量更低

解决方案

  1. 使用Continuous Batching:动态地将请求组合成Batch,提高GPU利用率
  2. 使用推测解码(Speculative Decoding):使用一个小的草稿模型快速生成多个候选词,然后用大模型验证
  3. 使用量化:减少计算量,降低延迟

本章小结

核心Takeaways

  1. Transformer通过Attention机制完全替代了RNN,解决了序列依赖问题,实现了完全并行化训练。

  2. Self-Attention是Transformer的核心,通过Query、Key、Value的计算,使得序列中的每个位置都能"关注"所有其他位置。

  3. 多头注意力通过并行计算多个注意力头,让模型能够同时关注不同类型的依赖关系,提高了模型的表达能力。

  4. 位置编码为序列引入位置信息,Sinusoidal编码是原始Transformer的选择,RoPE是现代大模型的主流方案。

  5. 残差连接和LayerNorm是训练深层Transformer的关键,Pre-Norm架构比Post-Norm更稳定。

  6. FFN通过引入非线性,使得Transformer能够学习复杂的函数。位置级FFN对每个位置独立地应用相同的变换。

  7. Transformer有三种主要变体:Encoder-only(BERT,适合理解任务)、Decoder-only(GPT,适合生成任务)、Encoder-Decoder(T5,适合序列到序列任务)。

  8. **高效注意力机制(稀疏注意力、Linear Attention、Flash Attention)**解决了标准Attention的O(n²)复杂度问题,使得处理长序列成为可能。

  9. **分布式训练(数据并行、模型并行、流水线并行、张量并行)**使得训练超大模型成为可能。

  10. 混合精度训练和梯度累积是训练大模型的关键技术,可以加速训练并减少显存占用。

  11. 大模型训练面临稳定性挑战(梯度爆炸/消失、Loss Spike、OOM等),需要通过梯度裁剪、Warmup、梯度检查点等技术来解决。


思考题

思考题1:为什么Self-Attention的计算复杂度是O(n²),有哪些方法可以降低这个复杂度?请详细解释其中一种方法的原理。

参考答案

Self-Attention的O(n²)复杂度来源

标准Self-Attention需要计算一个 n×nn \times nn×n 的注意力矩阵(其中 nnn 为序列长度),每个元素都需要计算Query和Key的点积。因此,时间和空间复杂度都是 O(n2d)O(n^2 d)O(n2d)(其中 ddd 为向量维度)。

降低复杂度的方法

  1. 稀疏注意力(Sparse Attention) :每个token只关注固定窗口内的token,复杂度降至 O(nw)O(nw)O(nw)(www 为窗口大小)
  2. 线性注意力(Linear Attention) :通过重新排列矩阵乘法顺序,将复杂度降至 O(nd2)O(nd^2)O(nd2)
  3. Flash Attention :通过IO感知的算法设计,减少显存读写,虽然理论复杂度仍是 O(n2)O(n^2)O(n2),但实际运行更快且显存占用降至 O(n)O(n)O(n)

详细解释Linear Attention

Linear Attention使用一个核函数 ϕ(⋅)\phi(\cdot)ϕ(⋅) 来近似Softmax:

Attention(Q,K,V)=ϕ(Q)(ϕ(K)TV)ϕ(Q)(ϕ(K)T1) \text{Attention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)(\phi(K)^T \mathbf{1})} Attention(Q,K,V)=ϕ(Q)(ϕ(K)T1)ϕ(Q)(ϕ(K)TV)

通过首先计算 ϕ(K)TV\phi(K)^T Vϕ(K)TV(复杂度 O(nd2)O(nd^2)O(nd2)),然后再与 ϕ(Q)\phi(Q)ϕ(Q) 相乘,避免了显式计算 n×nn \times nn×n 的注意力矩阵。

思考题2:在训练一个大模型时,你发现Loss出现了异常的尖峰(Spike),应该如何排查和解决这个问题?

参考答案

排查步骤

  1. 检查数据:查看出现Spike的batch是否包含异常数据(如极长序列、特殊字符等)
  2. 检查梯度:打印梯度的范数(Gradient Norm),查看是否出现梯度爆炸
  3. 检查学习率:查看学习率调度是否正确,是否在某些步骤学习率突然变大
  4. 检查数值稳定性:查看是否有NaN或Inf出现

解决方案

  1. 降低学习率:如果Spike频繁出现,可能是学习率过大
  2. 使用Warmup:在训练初期逐渐提高学习率,避免初期的不稳定
  3. 使用梯度裁剪:限制梯度的最大范数,避免梯度爆炸
  4. 使用混合精度训练:确保Loss Scaling因子设置正确
  5. 数据清洗:移除或修正导致Spike的异常数据
  6. 使用更稳定的优化器:如AdamW(相比SGD更稳定)

预防措施

  1. 使用小规模实验验证:在训练全规模模型前,先用小模型和小数据验证训练稳定性
  2. 监控训练指标:使用TensorBoard或Wandb实时监控Loss、梯度范数、学习率等
  3. 使用确定性算法:设置随机种子,确保实验可复现

本章已生成完毕。请回复【继续生成第四章】或提出您对当前章节的疑问。


参考文献

  1. Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention is all you need. Advances in Neural Information Processing Systems (NeurIPS), 30.
  2. Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of deep bidirectional transformers for language understanding. Proceedings of NAACL-HLT, 4171-4186.
  3. Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving language understanding by generative pre-training.
  4. Radford, A., Wu, J., Child, R., et al. (2019). Language models are unsupervised multitask learners. OpenAI Blog, 1(8), 9.
  5. Brown, T., Mann, B., Ryder, N., et al. (2020). Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 33, 1877-1901.
  6. Su, J., Lu, Y., Pan, S., et al. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864.
  7. Press, O., Smith, N. A., & Lewis, M. (2021). Train short, test long: Attention with linear biases enables input length extrapolation. International Conference on Learning Representations (ICLR).
  8. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems (NeurIPS), 35, 16344-16359.
  9. Shoeybi, M., Patwary, M., Puri, R., et al. (2019). Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.
  10. Raffel, C., Shazeer, N., Roberts, A., et al. (2020). Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140), 1-67.