The Annotated Transformer 源码剖析与复现笔记

💡 项目本质与复现动机

“Paper is cheap, show me the code.”
《Attention Is All You Need》虽然是开山之作，但忽略了许多工程细节。哈佛 NLP 团队的这个项目将论文原句与 PyTorch 代码逐行对应，是理解 Transformer 底层张量流转（Tensor Flow）和 Mask 机制的最佳学习材料。此笔记记录了我的本地复现过程与关键代码的理解。

🛠️ 个人复现感悟与思考

1. 宏观架构：Model Architecture

整体大框架：核心 Encoder-Decoder 骨架 + 独立的 Generator 任务头

  • Encoder-Decoder 骨架： 负责核心的张量流转与上下文语义特征提取，输出高维隐状态（Hidden States）。
  • Generator 任务头： 本质是一个线性分类器 + Softmax。它负责将 Decoder 输出的抽象向量，映射为词表空间上的概率分布（Logits -> Probabilities），从而为后续的贪婪解码或采样策略提供依据。这种骨架与头部解耦的设计，正是如今大模型 lm_head 的雏形。

2. 核心架构与工程细节探索

2.1 避免权重共享的“克隆陷阱” (clones 模块)

在构建多层同构网络（如 6 层 Encoder）时，绝不能使用简单的列表乘法（如 [layer] * 6），这会在底层内存中导致 6 层网络指向同一组权重参数指针，使得网络强行同质化。源码中通过 copy.deepcopy 在内存中开辟独立空间，确保每一层拥有独立可训练的大脑；同时必须使用 nn.ModuleList 进行外层包装，这是向底层的 Autograd 引擎显式“注册参数”的必要步骤，以保证反向传播时能正确追踪梯度。

2.2 LayerNorm (层归一化) 的“反悔机制”

深度网络极易因矩阵连乘导致数值失控（梯度消失/爆炸）。LayerNorm 的核心公式为：

$$ y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta $$

这里的精髓在于 $\gamma$ (Scale) 和 $\beta$ (Shift) 这两个可训练参数。归一化虽然稳定了数值分布，但也可能破坏了网络好不容易提取到的特征表达。这两个参数赋予了网络“反悔”的权利：如果归一化效果不好，网络可以通过学习将 $\gamma$ 设为方差，$\beta$ 设为均值，从而还原原始特征。

2.3 Attention 机制的本质与缩放因子 (ScaledDotProductAttention)

自注意力的底层运算不过是一堆矩阵乘法，其核心公式为：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

源码中特别强调了除以 $\sqrt{d_k}$ 这个缩放因子的重要性。当 $d_k$（Key的维度）很大时，点积结果的方差会变得极大，导致输入到 Softmax 的值变得极其陡峭，梯度接近于 0（梯度消失）。除以根号 $d_k$ 是为了将分布拉回平缓区间，保证反向传播的顺畅。

2.4 两种 Mask 的巧妙实现 (掩码机制)

Transformer 在张量运算时是高度并行的，必须通过 Mask 来屏蔽不该看的信息。源码中展现了极其优雅的 Tensor 操作：

  • Padding Mask (填充掩码)： 为了让不同长度的句子对齐，补全了 <pad> 标签。通过生成布尔矩阵，将对应位置的值替换为 -1e9（极小的负数），这样经过 Softmax 后权重就变成了 0，模型彻底“无视”了填充字符。
  • Subsequent Mask (因果/自回归掩码)： 在 Decoder 阶段，为了防止模型“偷看”未来的词，使用 torch.tril 生成了一个下三角矩阵。这奠定了后来所有 GPT 系列自回归生成任务的底层基石。

2.5 Positional Encoding (位置编码) 的工程落地

由于 Transformer 完全没有 RNN 的时序概念，必须人为注入位置信息。这里最容易踩坑的工程细节是：位置编码矩阵不参与梯度更新（不可训练）。
在 PyTorch 中，必须使用 self.register_buffer('pe', pe) 将其注册为 Buffer，而不是 Parameter。这样不仅能保证它不被 Optimizer 更新，还能确保在调用 .cuda() 时，它能跟随模型一起正确移动到 GPU 显存中。

3. 其他核心组件：FFN、Embedding与位置编码

在 Attention 实现了全局上下文的交互之后，模型还需要针对局部特征进行非线性映射，并为无序的序列注入位置信息。

3.1 前馈神经网络 (Position-wise FFN)

Attention 层负责跨 Token 的信息交互，而 FFN 层则负责对每个位置的 Token 进行独立的非线性变换。两者交替进行。

  • 升维与降维： 源码中，第一层线性变换 $W_1$ 将 512 维的隐向量映射到 2048 维，经过 ReLU 激活函数引入非线性后，再由第二层 $W_2$ 映射回 512 维。这种结构有助于模型在更高维空间中提取和解耦复杂的语义特征。

3.2 词嵌入与权重共享 (Embeddings & Softmax)

模型通过 Embedding 层将离散的文本 Token 映射为连续的高维向量。

  • 权重共享 (Weight Tying)： 原论文和源码中，Encoder 的输入 Embedding、Decoder 的输入 Embedding，以及 Decoder 最后的 Pre-softmax 线性映射层共享了同一组权重矩阵。这不仅大幅减少了模型参数量，还使得输入和输出能够对齐在同一个语义空间中。
  • Scale 缩放 (* math.sqrt(d_model))： 源码在获取词向量后，乘以了 $\sqrt{512}$。因为 Embedding 初始化的方差较小，而接下来要相加的位置编码（正余弦函数）值域在 $[-1, 1]$ 之间。放大词向量的量级，是为了防止自身的语义特征被位置信号掩盖。

3.3 位置编码 (Positional Encoding)

由于 Self-Attention 自身不具备处理序列顺序的能力，必须显式地引入位置信息。论文使用了不同频率的正弦和余弦函数：

  • 对数空间计算： 在 PyTorch 实现中，直接计算大底数幂容易导致下溢或速度较慢。源码利用 $a^{-x} = e^{-x \cdot \ln(a)}$，在 Log 空间内（利用 torch.exp 和 math.log）计算分母，是常见的工程加速与防溢出技巧。
  • 相对位置特性： 对于固定的偏移量 $k$，$PE{pos+k}$ 可以表示为 $PE{pos}$ 的线性组合。这使得 Attention 机制在点积时更容易学习到词与词之间的相对距离关系。
  • Buffer 注册： 位置编码是固定的数学特征，不需要计算梯度。源码通过 self.register_buffer("pe", pe) 将其注册为模型的常量状态，避免被优化器更新。

4. 模型组装与推理逻辑简述

在完成所有子模块后，make_model 和 inference_test 展示了模型的完整组装和前向推理过程。

  • 深拷贝隔离 (copy.deepcopy)： 在堆叠 6 层 Encoder 和 Decoder 时，必须使用深拷贝来实例化每一层，否则 PyTorch 底层会让多层网络共享同一块内存参数，导致网络层同质化。
  • 参数初始化： 组装完成后，源码对维度大于 1 的参数统一使用了 Xavier Uniform 初始化（Glorot 初始化），以保证前向传播时输入输出方差一致，维持训练早期的数值稳定性。
  • 自回归推理逻辑： 在文本生成时，Encoder 只需要对源序列执行一次前向传播生成 memory。随后 Decoder 进入循环：每次接收历史 Token 序列和 memory，经过 subsequent_mask 掩码处理，预测下一个 Token，并将其拼接到历史序列中送入下一轮循环。

5. 模型训练与优化机制

模型的训练不仅依赖前向传播，还需要合理的批处理策略、学习率调度以及正则化手段来保证收敛的稳定性和泛化能力。

5.1 数据批处理与掩码合并 (Batching & Masking)

在处理变长序列时，合理的批处理能够大幅提升 GPU 算力利用率。

  • 动态批处理 (Dynamic Batching)： 原论文在训练时没有采用固定的句子数量作为 Batch Size，而是将长度相近的句子打包，使每个 Batch 包含固定数量的 Token（约 25000 个）。这种策略极大地减少了 <pad> 填充符的数量，提升了矩阵运算的有效载荷。
  • Teacher Forcing 偏移输入： 在构造 Decoder 的输入和标签时，通过切片实现错位预测。tgt[:, :-1] 去掉最后一个词作为输入，tgt[:, 1:] 去掉第一个词作为预测目标。
  • 掩码的逻辑交集： Decoder 在自回归生成时，必须同时满足“不能看填充符”和“不能看未来词”两个条件。源码中通过逻辑与运算 tgt_mask & subsequent_mask 获取两种掩码的交集，确保模型只关注合法的历史真实 Token。

5.2 训练流水线与显存管理 (Training Loop)

run_epoch 封装了标准的深度学习训练循环，并引入了适用于大模型训练的工程优化：

  • 梯度累加 (Gradient Accumulation)： 受限于显存大小，模型通常无法一次性处理极大的 Batch。通过设置 accum_iter，模型在多个 Step 内仅执行反向传播 .backward() 累加梯度，达到指定次数后再调用 optimizer.step() 更新权重并清空梯度。这在工程上等效模拟了大 Batch Size 的平稳训练效果。
  • 防止显存泄漏： 每次计算完 Loss 并在反向传播结束后，源码通过 del loss 和 del loss_node 显式释放变量。这强制解除了对庞大计算图的引用，避免了随着循环推进导致显存溢出 (OOM)。

5.3 动态学习率调度 (Noam Schedule)

为了防止模型在随机初始化的早期因步长过大而导致梯度爆炸，作者提出了一种包含预热（Warmup）和衰减的动态学习率公式：

  • 两阶段调整：
    1. 线性预热： 在前 warmup_steps 步（如 4000 步）内，公式选中后半部分，学习率呈线性稳步上升，帮助模型在初始阶段平稳建立优化方向。
    2. 反平方根衰减： 超过预热期后，公式选中前半部分，学习率按 $1/\sqrt{x}$ 的规律平滑下降，在接近最优解时进行更细致的微调。
  • 维度惩罚： 公式前方的 $d_{model}^{-0.5}$ 表明，模型的隐层维度越高（参数空间越复杂敏感），其整体学习率的峰值上限就会被压制得越低。代码中通过 PyTorch 的 LambdaLR 调度器实现了该策略。

5.4 正则化与标签平滑 (Label Smoothing)

在标准的交叉熵损失中，使用 One-Hot 编码要求模型将目标类别的概率推向 1.0。这容易导致模型过度自信（Overconfident）并在训练集上过拟合。

  • 平滑概率分布： Label Smoothing 机制不再追求 100% 的绝对确信。通过引入平滑参数 $\epsilon$（如 0.1），将正确目标词的置信度降低到 $1 - \epsilon$（0.9），并将剩余的 $\epsilon$（0.1）概率质量均匀分配给词表中所有其他错误的 Token。
  • KL 散度约束： 源码使用 nn.KLDivLoss（KL 散度）代替普通交叉熵，强制模型预测的输出概率去拟合这个人工调配的平滑分布。
  • 工程效果： 这种机制迫使模型去学习类间更丰富的软标签特征，而不是盲目放大正确类别的 Logit 权重。它容忍了一定的不确定性，从而提高了模型在面对未知数据时的泛化能力和 BLEU 评分。

总结： 跑通 The Annotated Transformer，最大的收获不是记住了几个公式，而是彻底理解了从“纸面数学公式”到“GPU 张量流转”之间那道工程鸿沟是如何跨越的。这些底层认知，在如今微调、部署 LLaMA 或 GPT 架构时，依然是极其宝贵的财富。

Tech-Practice
#Transformer #PyTorch #Source-Code
The Annotated Transformer 源码剖析与复现笔记
https://cutesxy.github.io/2026/03/03/Transformer-Code-Note/
作者
hxy
发布于
2026年3月3日
许可协议