Transformer实战

TransformerBlock

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)   # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

附录知识

Loss Function 损失函数

用于量化模型预测值与真实值之间差异的函数，其核心目标是通过最小化该差异来优化模型参数。

评估性能：衡量模型预测的准确性，值越小表示预测越接近真实值。
指导优化：通过梯度下降等算法调整模型参数，使损失最小化。
问题类型：回归选MSE/MAE，分类选交叉熵/Hinge损失。

通过合理选择损失函数，可显著提升模型性能。例如，图像分类常用交叉熵，房价预测多用MSE。

Cross Entropy是信息论中用于衡量两个概率分布差异的重要指标

在分类任务中：

真实分布 P：通常是标签的 one-hot 编码（如真实类别为3，则 P=[0,0,0,1,0,…]）。
预测分布 Q：模型输出的概率分布（如 Softmax 后的结果）。

交叉熵损失函数的目标是让预测分布 Q 逼近真实分布 P。其数学形式为：

C：类别总数。
yi：真实标签的第 i 个分量（0 或 1）。
pi：模型预测的第 i 个类别的概率。（需满足 ∑pi=1）

Activation Function激活函数

Activation Function是神经网络中的核心组件，用于向网络引入非线性能力，使模型能够学习复杂的数据模式。激活函数的作用：

引入非线性：若无激活函数，多层神经网络等价于单层线性变换（W2(W1x+b1)+b2=W′x+b′），无法拟合复杂函数。
控制输出范围：如Sigmoid将输出压缩到(0,1)，适合概率场景；ReLU将负值归零，增强稀疏性。
梯度传递：激活函数的导数影响反向传播的梯度，决定参数更新效率（如ReLU的梯度在正区间恒为1，缓解梯度消失）。

Sigmoid

输出范围：(0, 1)
优点：输出可解释为概率，适用于二分类输出层。
缺点：梯度消失：当输入绝对值较大时，导数接近0，导致深层网络难以训练。
应用场景：二分类输出层

Tanh（双曲正切）

ReLU（Rectified Linear Unit）

Leaky ReLU

Softmax

输出范围：(0,1)，且和为1
特点：将输出转换为概率分布，用于多分类输出层（配合交叉熵损失）。
注意：输入值较大时可能数值溢出，需配合Log-Softmax或稳定化技巧。

表格：Softmax 计算示例

输入向量	最大值	调整后向量	指数值	总和	Softmax 输出
[1, 2, 3]	3	[-2, -1, 0]	[0.1353, 0.3679, 1.0]	1.5032	[0.09, 0.2447, 0.6653]
[4, 5, 6]	6	[-2, -1, 0]	[0.1353, 0.3679, 1.0]	1.5032	[0.09, 0.2447, 0.6653]

def softmax(x):
    x_max = np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(x - x_max)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

GELU

GELU核心思想是将确定性的阈值判断（如ReLU）转化为概率化的平滑决策。GELU在整个坐标系中是平滑的，特别是在 x<0时有轻微的弯曲，这使得它在神经网络中更有利于梯度流动。GELU的数学表达式为：

其中，Φ(x) 是标准正态分布的累积分布函数（Cumulative Distribution Function, CDF）。实际应用中常使用近似公式：

class GELU(nn.Module):
    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2 / torch.pi)) * (x + 0.044715 * x**3)))

激活函数选择指南

场景	推荐激活函数	理由
二分类输出层	Sigmoid	输出概率值，适配交叉熵损失。
多分类输出层	Softmax	输出归一化概率分布。
隐藏层（通用）	ReLU/Leaky ReLU/Swish	计算高效，缓解梯度消失，适合深层网络。
RNN/LSTM隐藏层	Tanh	零中心化，梯度幅度适中，避免输出漂移。
对抗死亡神经元问题	Leaky ReLU/PReLU	允许负区间梯度，保持神经元活性。
轻量级模型	ReLU	计算简单，节省资源。

优化器

用于调整模型参数以最小化损失函数的核心组件，其核心作用是通过梯度信息指导参数更新方向和步长。

SGD（随机梯度下降）

公式：θ = θ - lr * ∇J(θ)
特点：每次用单个样本或小批量样本计算梯度，更新参数。
优点：简单高效，适合大规模数据。

缺点：易陷入局部极小值，震荡明显。

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

Adam（Adaptive Moment Estimation）

为每个参数维护独立的自适应学习率：

# 使用Adam优化器（自动调整各参数的学习率）
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 初始lr仅作为参考基准

# 训练循环中，参数更新方式与SGD一致
optimizer.step()

内部机制：

Adam 根据梯度的一阶矩（均值）和二阶矩（方差）动态调整每个参数的学习率。
实际更新公式： param = param - lr * m_t / (sqrt(v_t) + eps) 其中 m_t 和 v_t 是动量项。

参数默认值：β1=0.9,β2=0.999,ϵ=1e−8
特点：
- 结合Momentum和RMSprop，自适应调整学习率。
- 适用于大多数深度学习任务，尤其是默认选择。

选择策略

优化器类型	适用场景	注意事项
SGD/Momentum	简单模型、需精细调参	需配合学习率调度器使用
Adam/AdamW	大多数深度学习任务（默认首选）	注意权重衰减参数设置
自适应优化器	稀疏数据、特征分布差异大	避免用于RNN等动态网络结构
二阶优化器	小规模模型、需快速收敛	计算成本高，工业级应用较少

学习率

更新速度：高学习率（如0.1）可能导致训练不稳定甚至发散，低学习率（如0.001）会减缓收敛速度，易陷入局部最优。
稳定性：合理的学习率能帮助模型平稳收敛到全局最优或较好的局部最优

训练和推理过程中Attention的差异

1. 计算模式与并行性

训练阶段：

全序列并行处理：输入是完整的序列（如一个句子或段落），所有位置的 token 同时计算注意力权重，无需逐步生成。例如，输入序列长度为 ( n )，模型一次性计算所有 ( n \times n ) 的注意力分数。
完全利用硬件加速：通过矩阵乘法并行计算所有注意力头（Multi-Head），最大化 GPU/TPU 的并行计算能力。

推理阶段：

自回归生成（逐token生成）：输出 token 逐个生成（如生成文本时，每次预测下一个 token）。每次生成新 token 时，只能计算当前 token 与之前所有 token 的注意力权重（因果注意力，Causal Attention）。
KV 缓存优化：为避免重复计算，已生成的 token 的 Key（K）和 Value（V）会被缓存。每次生成新 token 时，只需计算当前 token 的 Q，并与缓存的 K、V 进行注意力计算，从而减少计算量。

2. 内存与计算复杂度

训练阶段：

内存需求高：需要存储完整的注意力矩阵（( n \times n )），以及中间梯度信息（如 Q、K、V 矩阵），导致显存占用随序列长度平方增长（( O(n^2) )）。
计算复杂度：注意力计算复杂度为 ( O(n^2 d) )（( d ) 为模型维度），对长序列（如 4096 token）显存压力极大。

推理阶段：

内存需求低：通过 KV 缓存，每次只需存储历史 token 的 K、V（缓存大小为 ( O(n d) )），无需保存完整的注意力矩阵。
增量式计算：每次生成新 token 的计算复杂度为 ( O(n d) )，整体复杂度仍为 ( O(n^2 d) )，但通过缓存减少重复计算。

3. 注意力掩码（Masking）

训练阶段：

因果掩码（Causal Mask）：在自回归任务（如 GPT）中，使用下三角矩阵掩码，防止当前 token 关注未来信息。
填充掩码（Padding Mask）：在批处理中，忽略无效的填充 token（如 [PAD]）。

推理阶段：

仅需因果掩码：生成时天然满足因果性（未来 token 尚未生成），无需显式掩码，但需保证每次生成的 token 仅关注历史信息。

4. 参数更新与梯度计算

训练阶段：

反向传播：注意力权重（Q、K、V 的线性变换参数）需计算梯度并更新，需保留中间结果（如注意力分数、Softmax 输出）以支持梯度计算。
梯度检查点（Gradient Checkpointing）：为节省显存，可能选择性丢弃部分中间结果，需要时重新计算。

推理阶段：

无梯度计算：参数固定，仅需前向传播，无需存储中间变量，显存占用更低。

5. 优化策略

训练阶段：

混合精度训练：使用 FP16/BF16 加速计算，结合梯度缩放（Gradient Scaling）避免数值下溢。
分布式训练：通过数据并行、模型并行或流水线并行拆分大模型。

推理阶段：

量化（Quantization）：将模型权重从 FP32 转换为 INT8/INT4，减少内存占用和计算延迟。
稀疏注意力（Sparse Attention）：对长序列裁剪部分注意力连接（如局部窗口注意力），降低计算复杂度。

6. 示例对比

训练阶段（输入序列长度 ( n=4 )）：

输入：[A, B, C, D]
注意力计算：一次性计算所有 token 之间的注意力（包括 A→B, A→C, A→D, B→C, B→D 等）。

推理阶段（生成序列 `[A, B, C, D]`）：

生成 A：无历史信息，仅自身注意力（或初始提示）。
生成 B：计算 B 与 A 的注意力，缓存 A 的 K、V。
生成 C：计算 C 与 [A, B] 的注意力，复用缓存的 K、V。
生成 D：计算 D 与 [A, B, C] 的注意力，复用缓存。

总结

训练和推理的核心差异在于：

训练追求高吞吐量和参数优化，需处理完整序列并保留中间结果；
推理追求低延迟和资源效率，需增量生成并复用历史信息。