12 KiB
title, draft, tags
| title | draft | tags | |||
|---|---|---|---|---|---|
| 2-工程要求 | false |
|
下面,我们将跳出“注意力机制就像人类注意力”这种高层隐喻,直接将视角下切到张量、算子、反向传播与内存总线上。
!
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 矩阵去做循环计算。
- 输入 X 形状:
\text{[Batch\_size, Seq\_len, d\_model]} - 单一矩阵乘法:计算 $Q = X \cdot W^Q$,结果形状仍为 $\text{[Batch_size, Seq_len, d_model]}$。
- 多头切分(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 的数学必要性(方差推导)
假设 Q 和 K 的分量 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)实现绝对的时间隔离。
合并掩码机制:
- Padding Mask:形状为 $\text{[Batch_size, 1, 1, Seq_len]}$,值为布尔型(True 代表有效,False 代表 Pad)。
- 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 致命陷阱)
-
半精度(FP16)Mask 数值溢出: 在 FP32 中,Mask 的无效值填入
-1e9是合理的;但在混合精度(AMP)或纯 FP16 中,FP16 的能表示的最小负数下界约是 -65504。如果强行写入 $-1e9$,张量会变成 -Inf 甚至 NaN,反向传播第一步直接全网 NaN。 解法:在代码中必须根据 dtype 动态获取:mask_val = torch.finfo(q.dtype).min。 -
忽略 Adam 优化器的
\epsilon精度灾难: 在混合精度训练中,如果分母(二阶矩均方根)加上极小值\epsilon时,\epsilon的默认值(如 $1e-8$)超出了 FP16 的精度范围(FP16的最小值阈值通常在1e-7量级)。这会导致除以零。 解法:在用 FP16 训练 Transformer 时,需手动将 Adam 的 eps 参数调整为 $1e-5$。 -
Weight Decay (L2 正则化) 误伤 LayerNorm 参数: 将所有模型参数无脑送入带 Weight Decay 的 Optimizer 中。LayerNorm 中的
\gamma和\beta参数,以及所有 Linear 层的 Bias 项,其本身维度极低且作用是保证仿射变换的平移和缩放能力。如果对其施加 Weight Decay,会限制归一化层的表达力,严重阻碍收敛。 解法:在构建优化器组时,必须使用脚本分离出所有的LayerNorm.weight、LayerNorm.bias、*.bias等一维张量,对它们应用weight_decay = 0.0。