训练稳定性、复杂度与常见误区

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

在掌握了 Transformer 模型的核心组件——注意力机制、多头注意力、前馈网络和位置编码——之后，我们面临一个更实际的问题：如何稳定地训练它？Transformer 虽然结构优雅，但在训练初期，注意力分数的方差、梯度消失/爆炸等问题都可能让模型难以收敛。本节将深入剖析 Transformer 训练中的稳定性技巧（如缩放、残差连接、层归一化）、计算复杂度分析（为什么是 O(n²)），并澄清几个常见的理解误区。学完本节，你将能回答：为什么注意力要除以 $\sqrt{d_k}$？Transformer 的计算瓶颈在哪里？以及如何避免训练初期的数值不稳定。

先修：缩放点积注意力（本章 5.1 节）、多头注意力的形状与参数量（5.2 节）、残差与 LayerNorm（5.4 节）、从零拼出的最小 block（5.5 节）。后续：第 6 章《预训练范式》第 1 节《自监督与语言建模目标》——本节关于复杂度与训练稳定性的结论，是理解大规模预训练代价的前提。

💡 核心导读

1. 缩放因子 $\sqrt{d_k}$ 的数学根源：深入理解点积注意力中方差随维度 $d_k$ 增大的问题，以及缩放如何稳定 softmax 梯度。
2. Transformer 的 O(n²) 复杂度：明确计算复杂度来自注意力矩阵 QK^T，并分析其随序列长度 n 的平方增长关系。
3. 训练稳定性的三大支柱：残差连接（Residual Connection）、层归一化（LayerNorm）和权重初始化，它们如何协同工作以防止梯度问题。
4. 解码器的掩码（Mask）机制：区分训练时“前瞻未来”的掩码与推理时自回归生成的区别，理解其实现方式。
5. 常见误区与调试技巧：识别因形状错误、忘记缩放或掩码方向不对导致的典型错误现象及解决方法。

🧭 前置知识

1. 矩阵乘法：理解 (n, d_k) @ (d_k, m) -> (n, m) 的形状规则。
2. Softmax 函数：$Softmax(x_i) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$，其输出非负且和为 1。
3. 方差的基本性质：若随机变量 $X$ 和 $Y$ 独立且方差为 $\sigma^2$，则 $X+Y$ 的方差为 $2\sigma^2$，$X \cdot Y$ 的方差近似为 $\sigma^4$（在均值为 0 时）。
4. 残差连接：output = x + sublayer(x) 的基本概念，用于缓解梯度消失。

为什么注意力要除以 $\sqrt{d_k}$？—— 稳定 softmax 梯度

在 Attention(Q, K, V) = softmax(QK^T / √d_k) V 公式中，缩放因子 1/√d_k 并非随意添加。根据《Attention Is All You Need》（Vaswani et al., 2017，arXiv:1706.03762）第 3.2.1 节的解释，其根本目的是控制点积分数的方差，防止 softmax 函数进入梯度极小的饱和区。

理论推导（简化）：假设查询 q 和键 k 的每个维度都是均值为 0、方差为 1 的独立随机变量。那么它们的点积 q·k = Σ_{i=1}^{d_k} q_i k_i 的均值为 0，方差为 d_k。随着维度 d_k 增大，点积的绝对值倾向于变大。

问题：Softmax 函数对输入绝对值大的值非常敏感。假设输入向量为 [a, b, c]，若 a 远大于 b 和 c，则 softmax([a, b, c]) 会接近 [1, 0, 0]，此时梯度（关于 a、b、c 的导数）会变得非常小（因为输出几乎不变），这被称为梯度“消失”到饱和区，不利于训练。

解决方案：将点积结果除以 √d_k，相当于将其标准差重新缩放回 1 左右，从而使得 softmax 的输入保持在合理的范围内，梯度得以有效传播。

下面我们用代码来直观感受缩放前后的差异。

import torch

torch.manual_seed(0)

def mean_abs_dot(d_k, n_samples=2000):
    """抽样估计 |q·k| 的均值，以及除以 sqrt(d_k) 之后的均值"""
    q = torch.randn(n_samples, d_k)
    k = torch.randn(n_samples, d_k)
    dots = (q * k).sum(dim=-1)              # 每个样本一个点积
    return dots.abs().mean().item(), (dots / d_k ** 0.5).abs().mean().item()

print("d_k  | 无缩放 |q·k| | 缩放后 | 理论值 0.798*sqrt(d_k)")
print("-" * 56)
for d in (16, 64, 256, 1024):
    raw, scaled = mean_abs_dot(d)
    print(f"{d:5d} | {raw:11.2f} | {scaled:6.2f} | {0.798 * d ** 0.5:10.2f}")

d_k  | 无缩放 |q·k| | 缩放后 | 理论值 0.798*sqrt(d_k)
--------------------------------------------------------
   16 |        3.14 |   0.78 |       3.19
   64 |        6.42 |   0.80 |       6.38
  256 |       12.89 |   0.81 |      12.77
 1024 |       26.57 |   0.83 |      25.54

从输出可以看到两件事：不缩放时，$|q \cdot k|$ 的均值按 $\sqrt{d_k}$ 增长（16 维约 3.2，1024 维约 25.5，与理论值 0.798·$\sqrt{d_k}$ 基本吻合）；缩放后，不管 $d_k$ 多大，均值都被压回 0.8 左右——正好是标准正态绝对值的期望 $\sqrt{2/\pi} \approx 0.798$。这就是“把方差控制住”在数值上的样子：softmax 的输入不再随维度膨胀，梯度也就不会一上来就饱和。

Transformer 的计算复杂度：O(n²) 从何而来？

Transformer 被诟病的主要一点是注意力机制随序列长度平方增长。这里不做笼统的“O(n²)”，而是把每一项拆开算清楚。

先只看注意力子层，有四笔账（以下按“一次乘加算一次运算”计；若把乘、加各算一次 FLOPs，所有项同乘 2，比例不变）：

1. QKV 投影 + 输出投影：四个 $d \times d$ 矩阵乘 $n$ 个位置，约 $4n d^2$ 次运算。这一项随 $n$ 线性增长。
2. $QK^\top$：每个头算一个 $n \times n$ 的分数矩阵，$h$ 个头合起来约 $n^2 d$ 次乘加。
3. softmax：对 $n \times n$ 的每一行做归一化，约 $O(n^2 h)$。
4. $A V$：每个头把 $n \times n$ 权重乘上 $n \times d_v$ 的值矩阵，合计约 $n^2 d$ 次乘加。

所以注意力部分每层约 $4nd^2 + 2n^2d$：短序列时线性项占主导，长序列时平方项占主导。别忘了前馈网络——$d_{ff}=4d$ 时它还有 $8nd^2$，整层合计约 $12nd^2 + 2n^2d$，交叉点从 $n = 2d$（只看注意力）右移到 $n = 6d$（整层），但“短序列看线性、长序列看平方”的结论不变。显存的大头仍是那份 $n \times n$ 的注意力矩阵，每层要存 $h \cdot n^2$ 个数。

下面按公式算账（不跑计时——CPU 上毫秒级计时的抖动比信号还大，算 FLOPs 才是确定的）：

def layer_cost(n, d_model=512, h=8):
    proj = 4 * n * d_model ** 2      # QKV 投影 + 输出投影
    attn = 2 * n ** 2 * d_model      # QK^T 与 A@V 两项之和
    mem = h * n ** 2 * 4             # 每层要存的 n x n 注意力矩阵（float32）
    return proj, attn, mem

print("n | proj FLOPs | attn FLOPs | attn mem")
print("-" * 44)
for n in (256, 512, 1024, 4096):
    proj, attn, mem = layer_cost(n)
    print(f"{n:5d} | {proj / 1e6:8.1f}M | {attn / 1e6:8.1f}M | {mem / 2**20:6.1f}MB")

n | proj FLOPs | attn FLOPs | attn mem
--------------------------------------------
  256 |    268.4M |     67.1M |    2.0MB
  512 |    536.9M |    268.4M |    8.0MB
 1024 |   1073.7M |   1073.7M |   32.0MB
 4096 |   4295.0M |  17179.9M |  512.0MB

读法很直接（表里只算注意力子层：proj 是 QKV 与输出投影、attn 是 $QK^\top$ 与 $AV$ 之和、mem 是注意力矩阵显存）：$n$ 每翻一倍，投影项翻 2 倍，注意力项翻 4 倍。这两项在 $n = 2d$ 处相等（$d=512$ 即 $n \approx 1024$）；把前馈网络算进来后，整层的交叉点右移到 $n \approx 6d \approx 3072$。所以在几百 token 的上下文里，注意力还不是瓶颈；到 4096 时，单个样本单层的注意力矩阵就要 512 MB（$8 \times 4096^2 \times 4$ 字节），32k 上下文更是按 GB 计——FlashAttention 这类工作要解决的正是这块显存，而不是 FLOPs 本身。

训练稳定性的三根支柱

残差、LayerNorm、初始化——机制细节在 5.4 节拆过，这里只补两条和“稳定性”直接相关、5.4 没细说的结论。

• 残差：x + Sublayer(x) 让每层雅可比变成 $W^\top + I$，那个 $I$ 保证至少有一条不被缩小的梯度通路（5.4 节 4 层堆叠实测差 40 多倍）。
• LayerNorm：$\gamma \odot \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}} + \beta$，$\mu$、$\sigma^2$ 沿特征维逐位置计算，统计量与 batch 无关，所以训练和推理行为一致。
• 初始化：方差要小（标准差 0.02 量级很常见），否则一上来 $QK^\top$ 就很大、softmax 直接饱和。GPT-2 之后还按残差层数 $N$ 把残差分支权重缩小 $1/\sqrt{N}$，抵消深层残差路径的方差累积（Radford et al., 2019）。

三者是配套的：调模型时单独改其中一个（比如只调初始化、不动 LN），往往看不到预期效果。

解码器掩码：防止“偷看”未来

在训练序列到序列（如机器翻译）任务时，解码器在预测第 t 个输出 token 时，只能依赖于已知的前 t-1 个 token 以及完整的编码器输出。为了防止模型在训练时“偷看”未来的答案，需要在自注意力计算中应用一个因果掩码（Causal Mask）。

掩码是什么？ 一个形状为 (n, n) 的布尔矩阵：下三角（含对角线）为 False（允许看），上三角为 True（屏蔽）。用 masked_fill 把 True 的位置填成 -inf，再经过 softmax，这些位置的权重就变成 0。

import torch
import torch.nn.functional as F

torch.manual_seed(0)

def generate_causal_mask(seq_len):
    """生成因果掩码：上三角（不含对角线）为 True，表示这些位置要屏蔽。"""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
    # 等价写法：mask = torch.arange(seq_len).view(-1, 1) < torch.arange(seq_len).view(1, -1)
    return mask  # shape: (seq_len, seq_len)

seq_len = 6
causal_mask = generate_causal_mask(seq_len)
print("因果掩码矩阵（True 表示需要屏蔽的未来位置）:")
print(causal_mask)

def show(mat):
    for row in mat:
        print("  " + " ".join(f"{v:7.3f}" for v in row))

# 模拟注意力分数
scores = torch.randn(1, seq_len, seq_len)  # (batch, n, n)
print("\n原始注意力分数矩阵（6x6，保留3位小数）:")
show(scores[0].tolist())

# 应用掩码：将未来位置设为非常大的负数
masked_scores = scores.masked_fill(causal_mask, -1e9)
print("\n被屏蔽位置的最大分数:", masked_scores[0][causal_mask].max().item())

# 计算注意力权重
attn_weights = F.softmax(masked_scores, dim=-1)
print("\n最终的注意力权重矩阵（每行和为1）:")
show(attn_weights[0].tolist())
row_sums = attn_weights.sum(dim=-1)[0].tolist()
print("各行和:", [round(s, 4) for s in row_sums])
print("被屏蔽位置的最大权重:", attn_weights[0][causal_mask].max().item())

因果掩码矩阵（True 表示需要屏蔽的未来位置）:
tensor([[False,  True,  True,  True,  True,  True],
        [False, False,  True,  True,  True,  True],
        [False, False, False,  True,  True,  True],
        [False, False, False, False,  True,  True],
        [False, False, False, False, False,  True],
        [False, False, False, False, False, False]])

原始注意力分数矩阵（6x6，保留3位小数）:
   -1.126  -1.152  -0.251  -0.434   0.849   0.692
   -0.316  -2.115   0.322  -1.263   0.350   0.308
    0.120   1.238   1.117  -0.247  -1.353  -1.696
    0.567   0.794   0.440   0.112   0.641   0.441
   -0.216  -0.743   0.563   0.260   0.523   2.302
   -1.469  -1.587   1.203   0.085  -1.200  -0.005

被屏蔽位置的最大分数: -1000000000.0

最终的注意力权重矩阵（每行和为1）:
    1.000   0.000   0.000   0.000   0.000   0.000
    0.858   0.142   0.000   0.000   0.000   0.000
    0.148   0.452   0.400   0.000   0.000   0.000
    0.265   0.333   0.234   0.168   0.000   0.000
    0.134   0.079   0.292   0.215   0.280   0.000
    0.037   0.033   0.542   0.177   0.049   0.162
各行和: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
被屏蔽位置的最大权重: 0.0

可以看到，第 0 行的注意力权重全部集中在第 0 个 token 自身（未来位置全被屏蔽）；第 1 行可以分给前两个 token，依此类推。这就是“只能看到过去和当前”的因果约束。

两点容易踩的细节：

• 这里用的是 masked_fill 语义：布尔矩阵里 True = 屏蔽。而 F.scaled_dot_product_attention 的布尔 attn_mask 是 True = 参与注意力（5.1 节提过），语义正好相反，混用会得到完全相反的结果。
• 填充值 -1e9 只适合 fp32。半精度下 -1e9 可能溢出成 -inf 甚至 NaN，正统写法是直接填 float("-inf")，softmax 内部会处理成 0 权重。

📝 动手练一练

问题：把 generate_causal_mask 里的 diagonal=1 改成 diagonal=0，也就是连“自己”一起屏蔽。运行一下，看看第 0 行的注意力权重会变成什么。想一想：为什么真实实现必须保证每一行至少留下一个可看的位置？

import torch
import torch.nn.functional as F

torch.manual_seed(0)

n = 4
scores = torch.randn(1, n, n)

# 错误做法：连对角线一起屏蔽
bad_mask = torch.triu(torch.ones(n, n), diagonal=0).bool()
w_bad = F.softmax(scores.masked_fill(bad_mask, float("-inf")), dim=-1)
print("diagonal=0 时第 0 行权重:", w_bad[0, 0].tolist())
print("是否出现 NaN:", bool(torch.isnan(w_bad).any()))

# 正确做法：保留对角线
good_mask = torch.triu(torch.ones(n, n), diagonal=1).bool()
w_good = F.softmax(scores.masked_fill(good_mask, float("-inf")), dim=-1)
print("\ndiagonal=1 时第 0 行权重:", [round(v, 4) for v in w_good[0, 0].tolist()])
print("是否出现 NaN:", bool(torch.isnan(w_good).any()))

diagonal=0 时第 0 行权重: [nan, nan, nan, nan]
是否出现 NaN: True

diagonal=1 时第 0 行权重: [1.0, 0.0, 0.0, 0.0]
是否出现 NaN: False

常见错误

本章小结

1. 缩放因子 √d_k 是稳定训练的关键，它通过控制点积分数的方差，防止 softmax 梯度饱和。
2. 注意力子层每层约 $4nd^2 + 2n^2d$，加上前馈网络的 $8nd^2$ 后整层约 $12nd^2 + 2n^2d$：短序列看线性项，长序列看平方项；显存大头是每层 $h \cdot n^2$ 的注意力矩阵，这是长上下文最贵的部分。
3. 残差连接、层归一化和合适的初始化共同构成了训练稳定性的支柱，三者是配套的，单独改一个往往效果打折。
4. 解码器中的因果掩码通过在训练时屏蔽未来信息，强制模型学习自回归生成，是序列生成任务正确工作的前提。
5. 注意力权重的计算必须对最后一个维度（dim=-1）应用 softmax，以确保每个查询对应的权重分布和为 1。

📋 行动清单

• 修改一段注意力计算代码，故意去掉缩放因子 √d_k，观察 softmax 输出的极端值（如出现 0.9999 和 0.0000），并与缩放后的输出对比。
• 编写一个函数，输入序列长度 n，生成对应的因果掩码矩阵，并手动验证其下三角属性。
• 用本节 layer_cost() 的公式，算一下你关心的模型（比如 n=2048、d=768、h=12）单层注意力矩阵要吃多少显存，再乘上层数，感受长上下文的代价。

延伸阅读

本段 PyTorch 官方文档/教程条目的框架出处：PyTorch: An Imperative Style, High-Performance Deep Learning Library（Paszke et al., 2019，arXiv:1912.01703）；页面内容以 PyTorch 2.14.1 为准。

1. Attention Is All You Need（Vaswani et al., 2017，arXiv:1706.03762）—— 本节所有核心机制（缩放点积注意力、多头注意力、位置编码、模型架构）的原始论文。
2. PyTorch 官方文档：torch.nn.functional.softmax（本节代码在本地 PyTorch 2.14.1 实测；stable 链接随版本更新）—— dim 参数的语义与隐式维度的弃用说明 - https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html

—— 小象教研组