Files
2026-05-16 17:16:51 +08:00

12 KiB
Raw Permalink Blame History

title, draft, tags
title draft tags
2-工程要求 false
Transformer
工程实践
深度学习

下面,我们将跳出“注意力机制就像人类注意力”这种高层隐喻,直接将视角下切到张量、算子、反向传播与内存总线上。 !Pasted image 20260512213048.png

1. 词嵌入Word Embedding的极度细节

数学本质Look-up Table 的物理实现

在代码层面Embedding 并非是一个复杂的“网络层”,而是一个参数矩阵 $E \in \mathbb{R}^{V \times d_{model}}$(其中 V 是词表大小)。

当输入一个离散的 Token ID序列\text{[Batch\_size, Seq\_len]} 的整型张量Embedding 层本质上是在执行非连续内存索引Gather 操作)

在数学上,这等价于将一个 One-hot 向量 x \in \{0, 1\}^V 与矩阵 E 相乘:


\text{Embed}(x) = x \cdot E

物理意义它将互相正交的高维稀疏空间One-hot投影到了一个稠密的、具有内积度量意义的低维连续几何空间中。在这个连续空间里向量的方向和模长编码了语义和语法关系。

训练机制与 \sqrt{d_{model}} 的神秘缩放

Embedding 矩阵 E 的更新完全依赖于反向传播。在 Pytorch 等框架中为了提升计算效率通常将计算梯度视为稀疏更新Sparse Update即每次只更新当前 Batch 中出现的词对应的行向量。

为什么要乘以 $\sqrt{d_{model}}$ 在工程实现中,输出张量 \text{[Batch\_size, Seq\_len, d\_model]} 需要与绝对位置编码Positional Encoding, PE相加。

  • 通常 Embedding 矩阵在初始化时采用标准正态分布(或均匀分布),其激活值的方差近似为 1但权重的绝对数值较小。
  • PE 是通过 \sin\cos 生成的,其数值稳定分布在 [-1, 1] 之间。

如果不进行缩放,随着 d_{model} 的增大Embedding 向量的 L2 范数(模长)会相对较小。乘以 \sqrt{d_{model}} 将 Embedding 向量的尺度拉大,防止语义信息在相加时被固定的位置编码信号所淹没Drown out

共享机制Weight Tying 策略

在工业级模型中,底层的输入 Embedding 矩阵 E_{in} 与顶层 Softmax 前的输出线性映射矩阵 W_{proj} 共享完全相同的内存指针。

  • 数学意义:输出层计算的是上下文向量 H 与候选词向量的内积(相似度)。共享权重意味着“将词编码为特征”和“从特征解码回词”是在完全相同的语义度量空间中进行的。
  • 工程好处:参数量骤降。对于词表 V=50000, d_{model}=1024 的配置,节省了约 50MB 的显存。同时,这也作为一种极强的正则化手段,防止了模型在大量低频词上过拟合。

2. 参数矩阵的构造与训练($W^Q, W^K, W^V, W^O$

参数初始化Xavier/Glorot 的工程考量

在代码实现中,W^Q, W^K, W^V \in \mathbb{R}^{d_{model} \times d_{model}} 通常采用 Xavier Uniform 初始化。

为什么? 在未经训练的深度线性网络中方差会随着层数呈指数级放大或缩小。Xavier 初始化通过绑定输入维度fan_in和输出维度fan_out使得前向传播时的激活值方差和反向传播时的梯度方差在每一层保持恒定


W \sim U\left(-\sqrt{\frac{6}{fan\_in + fan\_out}}, \sqrt{\frac{6}{fan\_in + fan\_out}}\right)

这防止了初始状态下 Softmax 极速陷入饱和区。

多头并行的代码级映射(d_k 与 $d_{model}$

d_k = d_{model} / h 不是需要训练的参数矩阵的维度,而是张量重塑Reshape后的逻辑维度

在工业实践中,绝对不会用 h 个独立的 d_k 矩阵去做循环计算。

  1. 输入 X 形状:\text{[Batch\_size, Seq\_len, d\_model]}
  2. 单一矩阵乘法:计算 $Q = X \cdot W^Q$,结果形状仍为 $\text{[Batch_size, Seq_len, d_model]}$。
  3. 多头切分Reshape & Transpose:将张量 Reshape 为 $\text{[Batch_size, Seq_len, h, d_k]}$然后进行维度转置Transpose/Permute为 $\text{[Batch_size, h, Seq_len, d_k]}$。

此时,多头($h$)被转移到了 Batch 维度的后方,从而可以利用高度优化的 Batch Matrix Multiplication (BMM) 算子一次性完成所有的注意力评分。

梯度流向与 Bias 选择

  • 梯度流(链式法则):前向 Loss 标量通过 Softmax 的雅可比矩阵回传,变成注意力矩阵 A 上的梯度矩阵。随后梯度分支:一路通过 A^T \cdot \nabla \text{Output} 传向 $V$,另一路通过 \nabla A \cdot K 传向 $Q$。由于 $Q=XW^Q$$\nabla W^Q = X^T \cdot \nabla Q$。最终所有梯度汇聚回张量 $X$,再沿残差路径直达 Embedding 层。
  • Bias 项考量:在 BERT 时代Linear 层普遍带有 Bias。但在现代 LLM如 LLaMA, GPT-3 的部分实现)中,W^Q, W^K 通常直接干掉 Bias。原因是1) Bias 在极高维度的内积空间中对模型表达能力的提升微乎其微2) 去除 Bias 可以提升算子融合Kernel Fusion的效率减少显存读写带宽。

3. 训练全流程的数值稳定性(工程核心)

LayerNorm 的工程实现与 \epsilon

在张量层面LayerNorm 并非在 Batch 或 Seq 维度计算,而是严格在特征维度 d_{model} 上计算统计量。


\mu = \frac{1}{d_{model}}\sum_{i=1}^{d_{model}} x_i \quad \sigma^2 = \frac{1}{d_{model}}\sum_{i=1}^{d_{model}} (x_i - \mu)^2

LN(x) = \gamma \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

\epsilon 的致命作用\epsilon 通常取 $1e-5$。在 FP16 混合精度训练中,如果某个 Token 的各维度特征非常接近,\sigma^2 会极小。由于 FP16 的精度下限限制,如果没有 $\epsilon$,分母极易出现 0 或下溢出Underflow导致整个张量被 NaN 污染,直接使训练崩溃。

Pre-LN vs Post-LN 的架构之争

特性 Post-LN (原版 Transformer) Pre-LN (现代 LLM如 GPT/LLaMA)
公式 x_{l+1} = LN(x_l + SubLayer(x_l)) x_{l+1} = x_l + SubLayer(LN(x_l))
恒等路径 被 LN 打断 完美的 x_{l+1} = x_l + \dots 恒等路径
梯度流动 顶部梯度向底层传递时被 LN 缩放,容易引起深层梯度弥散/爆炸。 梯度可以无损穿透至第一层。
Warmup依赖 极度依赖,缺少 Warmup 必炸。 相对鲁棒,收敛更平稳。

Scaled Dot-Product 的数学必要性(方差推导)

假设 QK 的分量 q_i, k_i 是相互独立、均值为 0、方差为 1 的随机变量。

在计算点积 q \cdot k = \sum_{i=1}^{d_k} q_i k_i 时:

  • 均值:\mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\mathbb{E}[k_i] = 0
  • 方差:\text{Var}(q_i k_i) = \mathbb{E}[q_i^2 k_i^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] - 0 = 1 \times 1 = 1

因此,点积总和的方差为:


\text{Var}\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k

工程灾难:当 d_k 很大(如 64 或 128点积结果的分布方差达到了 $d_k$,意味着会产生极大的绝对值。经过 \text{Softmax}(z_i) = \frac{e^{z_i}}{\sum e^{z_j}} 时,最大值所在的指数项将占据绝对统治地位,输出的概率分布变为类似 One-hot 的 $[0, 0, \dots, 1, \dots, 0]$。此时Softmax 的局部梯度 \frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j) 几乎全部为 0反向传播的梯度被掐断梯度消失

解法:除以 $\sqrt{d_k}$,使得方差重新拉回 1维持在 Softmax 的非饱和区。

4. 损失函数与收敛策略

Label Smoothing 细节

通过软化真实标签分布,缓解模型在特定词上的过度自信:


y_{true}^{LS} = (1 - \alpha) \cdot y_{true} + \frac{\alpha}{V}

(例如 $\alpha=0.1$。正确标签的概率要求变为 $0.9 + 0.1/V$,其余标签要求为 $0.1/V$)。

置信度影响:这使得推理时输出的 Logits 差异变得平滑,模型不会为了追求 1.0 的概率而产生极其极端的 Logit 峰值,从而使得 Beam Search 等解码算法的评分更具有泛化意义。

Teacher Forcing 与 Masking 的矩阵魔法

在工程上,为了避免像 RNN 那样串行循环我们将整个目标序列右移Shift Right后一并输入 Decoder依赖掩码Mask实现绝对的时间隔离。

合并掩码机制:

  1. Padding Mask:形状为 $\text{[Batch_size, 1, 1, Seq_len]}$值为布尔型True 代表有效False 代表 Pad
  2. Sequence (Look-ahead) Mask:形状为 \text{[1, 1, Seq\_len, Seq\_len]} 的下三角矩阵(主对角线及以下为 True

张量操作将两者做逻辑与Logical AND得到 \text{[Batch\_size, 1, Seq\_len, Seq\_len]} 的综合 Mask矩阵。

接着生成掩码偏置矩阵:通过 torch.where 或类似函数,将 True 映射为 0.0,将 False 映射为一个极小的负数(如 $-1e9$)。


Attention\_Scores = \frac{Q K^T}{\sqrt{d_k}} + Mask\_Bias

在此矩阵中,未来的信息由于加上了 $-1e9$,在 Softmax 后权重彻底变为 0.0,从而在 O(1) 时间内以高度并行的矩阵乘法完成了自回归的前向阻断。

5. 优化器与学习率工程

Adam 优化器参数配置

在 Transformer 训练中,由于参数更新的极端不稳定性,标准配置通常为:

  • $\beta_1 = 0.9$:控制动量(一阶矩)。
  • $\beta_2 = 0.98 \sim 0.999$:控制 RMSprop二阶矩。在处理稀疏梯度和剧烈波动时稍微调低 $\beta_2$(如 0.98)有时能防止二阶矩历史信息阻碍快速步进,但大多数实现使用 0.999。

Learning Rate Warmup 的不可替代性

为何先升再降? 在训练初期,模型参数是随机初始化的,输出层 Logits 方差极大。由于 Adam 使用历史梯度的平方(二阶矩)来归一化更新步长,在最初几个 Step 内二阶矩非常小且不稳定。如果此时给予巨大的学习率会导致参数更新的步长越界Overshoot直接摧毁精心初始化的参数分布。

如果没有 Warmup 会怎样? 由于残差连接的存在和方差积累,底层 Attention 层的输出会产生巨大的突变,导致 Softmax 急剧收敛到单个 Token 上,产生**“注意力坍缩Attention Collapse”**。梯度立刻消失,网络表现为 Loss 挂在一个极大值永远不下降。Warmup 给出了一个缓冲期,让 Adam 的一二阶统计量收集到足够多的稳定数据后,再全速推进。

🛠️ 工业避坑指南Top 3 致命陷阱)

  1. 半精度FP16Mask 数值溢出: 在 FP32 中Mask 的无效值填入 -1e9 是合理的但在混合精度AMP或纯 FP16 中FP16 的能表示的最小负数下界约是 -65504。如果强行写入 $-1e9$,张量会变成 -Inf 甚至 NaN反向传播第一步直接全网 NaN。 解法:在代码中必须根据 dtype 动态获取:mask_val = torch.finfo(q.dtype).min

  2. 忽略 Adam 优化器的 \epsilon 精度灾难: 在混合精度训练中,如果分母(二阶矩均方根)加上极小值 \epsilon 时,\epsilon 的默认值(如 $1e-8$)超出了 FP16 的精度范围FP16的最小值阈值通常在 1e-7 量级)。这会导致除以零。 解法:在用 FP16 训练 Transformer 时,需手动将 Adam 的 eps 参数调整为 $1e-5$。

  3. Weight Decay (L2 正则化) 误伤 LayerNorm 参数: 将所有模型参数无脑送入带 Weight Decay 的 Optimizer 中。LayerNorm 中的 \gamma\beta 参数,以及所有 Linear 层的 Bias 项,其本身维度极低且作用是保证仿射变换的平移和缩放能力。如果对其施加 Weight Decay会限制归一化层的表达力严重阻碍收敛。 解法:在构建优化器组时,必须使用脚本分离出所有的 LayerNorm.weightLayerNorm.bias*.bias 等一维张量,对它们应用 weight_decay = 0.0