AI/ML 详解：Loss 与 Optimization

前 9 期讲了「模型架构 + 推理」——一个训练好的 LLM 是怎么工作的。今天讲它怎么被训出来的四个齿轮：Cross-Entropy（评分函数：你错得有多离谱）、AdamW（步进引擎：每个参数自适应步长）、LR Schedule（调度器：什么时候快、什么时候慢）、Gradient Clipping（熔断器：防止单步把模型炸了）。这四件事组合起来，才是 GPT/Llama/Claude 万亿 token 训练能稳定不发散的真正原因。

Cross-entropy 就是训练时模型的「SLA 违约金」——预测概率分布和真实答案差多少，就罚多少。后端类比：错误率监控，但不是简单"对/错"二值，而是按"你给正确答案多少自信心"细粒度计费。给对的答案 99% 自信 = 几乎不罚；给对的答案 0.1% 自信 = 重罚——梯度告诉模型"下次给这个 token 多分点概率"。

LLM 训练任务的本质是下一个 token 预测：给前 N 个 token，模型输出一个长度等于词表（约 5 万-15 万）的概率分布，真实答案是某一个具体 token。需要一个 loss 把"分布预测 vs 单点真值"的差距量化成一个可微分的数字。

公式：H(p, q) = -Σx p(x) log q(x)。每个符号：p 是真实分布（在 LLM 里就是 one-hot——正确 token 位置为 1，其余为 0），q 是模型预测的分布（softmax 输出）。由于 p 是 one-hot，求和坍缩成一项：loss = -log q(正确 token)。

直觉：模型给正确 token 概率 0.99 → loss = -log(0.99) ≈ 0.01（几乎免罚）；给概率 0.001 → loss = -log(0.001) ≈ 6.9（重罚）；给概率 0 → loss = ∞（理论上罚到天上去，所以 softmax 永远不会输出真 0）。负对数这个形状的精妙在于："越自信越正确"奖励温和，"越自信越错"惩罚激烈——天然对齐了"模型该有多 calibrated"的需求。

为什么不用 MSE？分类任务上 MSE（(predicted - true)²）梯度小得多——softmax 输出已经在 [0,1]，平方再求导后梯度趋零，训练慢且容易卡在 bad local minima。Cross-entropy 配 softmax 时梯度形式漂亮得不可思议：∂loss/∂logit = (softmax_output - one_hot)——直接就是"预测 - 真值"，没有任何衰减项。这个数学巧合是为什么所有 LLM 训练都用 cross-entropy。

import torch
import torch.nn.functional as F

# 假设词表大小 50000，batch=2，模型对每个位置输出一个 logit 向量
logits = torch.randn(2, 50000)            # 模型的原始输出（未 softmax）
targets = torch.tensor([42, 7])           # 真实下一个 token 的 id

# PyTorch 的 cross_entropy 内置 log_softmax + NLL，数值稳定
loss = F.cross_entropy(logits, targets)
print(loss.item())  # 未训练时 ≈ ln(50000) ≈ 10.8（均匀猜测的 baseline）

# 手工等价实现，让公式具象化：
log_probs = F.log_softmax(logits, dim=-1)
manual_loss = -log_probs[range(2), targets].mean()
assert torch.isclose(loss, manual_loss)

# LLM 训练时 reduce='mean' 会把所有 token 位置的 loss 平均
# 训练良好的 LLM 在 web 文本上 loss 约 2.0-2.5（perplexity = exp(loss) ≈ 8-12）

AdamW 是「每个参数有自己自适应限流值的负载均衡器」——和你做分布式系统里的 per-key adaptive rate limiting 是同一个思想。某个参数最近梯度一直很大（说明它对 loss 影响大、是「热 key」），就自动给它小一点的步长；梯度一直小（「冷 key」）就放大步长。每个参数的 lr 都是被自己的历史梯度统计动态校准的。

朴素 SGD 全局一个 lr。但 Transformer 里不同层、不同参数的梯度量级差几个数量级——同一个 lr 要么把大梯度参数轰飞、要么把小梯度参数饿死。Adam（Kingma & Ba 2014）的洞察：用每个参数自己的梯度历史统计去归一化它的步长。AdamW（Loshchilov & Hutter 2017，ICLR 2019）在此基础上修正了一个 Adam 长期被忽视的 bug——把 weight decay 从 gradient 里抠出来单独应用，让正则化在自适应 lr 下也能正确工作。今天所有现代 LLM 训练都用 AdamW。

公式（每步）：

• mt = β₁·mt-1 + (1-β₁)·gt ← 一阶矩：最近梯度方向的指数移动平均（动量）
• vt = β₂·vt-1 + (1-β₂)·gt² ← 二阶矩：最近梯度幅度平方的 EMA
• θt = θt-1 - lr · m̂t / (√v̂t + ε) - lr · wd · θt-1 ← 更新（m̂、v̂ 是 bias-corrected 版本）

直觉：分子 m̂ 是"最近梯度的平滑方向"，让更新有惯性（穿过噪声）；分母 √v̂ 是"最近梯度的典型幅度"，让每个参数都被自己的尺度归一化（梯度大的参数除以大数 = 步小，梯度小的参数除以小数 = 步大）。最后一项 -lr·wd·θ 是解耦的权重衰减：每步把参数往零拉一点点，防止权重涨爆——这一项放在 gradient 外面是 AdamW 相对 Adam 的全部改动，看起来不起眼但训出来的模型泛化好很多。

显存代价：每个参数额外存 m 和 v，优化器状态 = 模型参数的 2 倍。70B 模型 FP32 参数 280 GB，加优化器 = 840 GB——这是为什么训大模型必须用 ZeRO/FSDP 把优化器状态分片到多卡。典型超参：lr=1e-4 到 3e-4（预训练）、1e-5 到 1e-4（微调）、β₁=0.9、β₂=0.95（LLM 用 0.95 比默认 0.999 更稳）、wd=0.1、ε=1e-8。

import torch
from torch.optim import AdamW

model = ...  # 你的 nn.Module

# 通常 weight decay 不应用到 bias 和 LayerNorm 的 scale
decay, no_decay = [], []
for n, p in model.named_parameters():
    if p.dim() >= 2: decay.append(p)        # 矩阵权重：加 wd
    else:               no_decay.append(p)     # bias / norm scale：不加

optimizer = AdamW(
    [{"params": decay,    "weight_decay": 0.1},
     {"params": no_decay, "weight_decay": 0.0}],
    lr=3e-4,
    betas=(0.9, 0.95),       # β₂=0.95 是 GPT/Llama 训练常用值
    eps=1e-8,
)

# 训练 loop 里：
for batch in data:
    optimizer.zero_grad()
    loss = model(**batch).loss
    loss.backward()           # 算梯度
    optimizer.step()           # 应用 AdamW 更新规则

和 TCP slow start + congestion avoidance 几乎一比一对应：先慢启动（warmup）摸清网络状况，再加速到稳态吞吐，最后看到拥塞信号缓步退避。也像数据库连接池的渐进式 ramp-up，或新员工入职的试用期→主力期→交接期。一句话：什么时候该快、什么时候该慢，不是常量，是时间的函数。

训练初期模型权重是随机初始化的，loss landscape 极陡，一开始就用 peak lr → 第一步 gradient 巨大 → 权重炸飞 → 训练发散（loss 变 NaN）。训练末期想精细微调到 loss 的局部极小值，大 lr 会让更新震荡跨过最优点。整段训练用同一个 lr 注定两头不讨好，所以需要一条随时间变化的曲线。

主流方案 = Warmup + Cosine Decay，GPT-3、Chinchilla、Llama 全用：

• Warmup（前 1-5% steps）：lr 从 0 线性升到 peak。lr(t) = lr_peak · t / T_warmup。让 Adam 的二阶矩 v 先攒够样本，避免训练开头几步用还不准的 √v 做归一化导致爆炸。
• Cosine Decay（剩余 95-99% steps）：lr 按余弦曲线从 peak 缓降到 peak 的 10%。lr(t) = lr_min + ½(lr_peak − lr_min)(1 + cos(π · progress))。余弦曲线两端慢、中段快——刚到 peak 时停留久些充分搜索，末段慢降避免越过极小值。

为什么是余弦？最早 SGDR 论文（Loshchilov & Hutter 2016）只是经验性发现 cos 比线性、阶梯衰减都好。后人猜测：cos 的两端导数为 0，给了"稳定起步"和"温柔收尾"两个好性质。社区也用线性 decay、梯形 schedule，差距通常 <1%，cos 因为先入为主成了默认。关键不是 cos vs linear，而是"必须有 warmup + 必须有 decay"——缺任何一头都会让 LLM 训练崩。

from transformers import get_cosine_schedule_with_warmup
from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))

total_steps = 100_000
warmup_steps = int(0.02 * total_steps)   # 2% warmup，GPT-3 风格

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps,
    # 默认衰减到 0，HuggingFace 想保 10% 末值可自行实现或用 num_cycles=0.5
)

for step, batch in enumerate(loader):
    loss = model(**batch).loss
    loss.backward()
    optimizer.step()
    scheduler.step()                # ← 每步都调，lr 自动按曲线变
    optimizer.zero_grad()

    if step % 100 == 0:
        print(f"step="{step} lr={scheduler.get_last_lr()[0]:.2e} loss={loss.item():.3f}")
# 你会看到 lr 从 ~0 升到 3e-4，再缓降；loss 在 warmup 段下降快得惊人

Gradient clipping = 训练过程的熔断器 / API rate limiter。每步算完梯度，先看总长度有没有超阈值——超了就按比例缩到阈值内，方向不变。和你做高可用系统里"单次请求超时就拦截、防止打爆下游"完全同构。Pascanu et al. 2013 为 RNN 提出来防止 exploding gradient，今天所有 LLM 训练默认 clip_norm = 1.0。

训练几十万 step 的 LLM，loss 曲线绝大多数时候平滑下降，但偶尔会有 loss spike——某个 batch 里恰好含一段罕见组合（特殊符号、损坏文本、训练分布外的 pattern），让某层 activation 爆增 → backprop 出来的 gradient 比平时大几十几百倍 → 一步更新就把权重推到很远的地方 → 后续训练再也回不来 → 训练发散，几天算力白烧。Clipping 就是这种灾难的断路器。

算法（最常用的 global norm clipping）：

• ① 反向传播算出所有参数的 gradient g₁, g₂, ..., gn
• ② 算全局 L2 范数：‖g‖ = √(Σ‖gi‖²)（把所有参数的 gradient 拼成一个超长向量求长度）
• ③ 设阈值 c（通常 1.0）。若 ‖g‖ > c，所有 gi 同步乘以 c / ‖g‖；否则不变
• ④ 用裁剪后的 gradient 喂给 optimizer.step()

关键不变量：g' = g · min(1, c/‖g‖)。方向完全保留（所有分量按相同比例缩），只把过长的步长砍回安全长度。

诊断信号：训练中应该打日志记录 grad_norm。正常分布是窄峰（如 0.1-0.5），偶尔尖刺（10-100）说明 clip 救了你。如果 grad_norm 长期顶在阈值附近，说明 lr 太大或数据有问题，需要调；如果 grad_norm 突然爆到 1e6+ 而 clip 后 loss 还是 NaN，多半是 fp16 数值溢出，需要换 bf16 或开 loss scaling。GPT-3、Chinchilla、Llama 全部用 clip_norm = 1.0，已经成了不需要讨论的默认值。

import torch
from torch.nn.utils import clip_grad_norm_

CLIP_VALUE = 1.0   # GPT-3/Llama 默认值，几乎不用调

for step, batch in enumerate(loader):
    loss = model(**batch).loss
    loss.backward()                       # 算梯度

    # 关键一行：原地裁剪所有参数的 gradient，返回裁剪前的 norm
    grad_norm = clip_grad_norm_(
        model.parameters(),
        max_norm=CLIP_VALUE,
    )

    optimizer.step()                       # 用裁剪后的梯度更新
    scheduler.step()
    optimizer.zero_grad()

    # 必须监控 grad_norm 的分布，这是训练健康的最重要信号
    if step % 10 == 0:
        print(f"step="{step} grad_norm={grad_norm:.3f} loss={loss.item():.3f}")
        # 健康：0.1-0.5 区间，偶有尖刺到 10+ 被 clip 救下
        # 报警：长期顶在 1.0（lr 太大）；突然 NaN/Inf（数值溢出）

• Adam: A Method for Stochastic Optimization (Kingma & Ba 2014) — Adam 原论文，深度学习里被引最多的论文之一
• Decoupled Weight Decay Regularization (Loshchilov & Hutter 2017, ICLR 2019) — AdamW 原论文，讲清 Adam 的 weight decay bug 和修正
• SGDR: Stochastic Gradient Descent with Warm Restarts (Loshchilov & Hutter 2016) — Cosine annealing schedule 起源
• On the difficulty of training Recurrent Neural Networks (Pascanu et al. 2013) — Gradient clipping 起源，分析 exploding gradient 的几何根源
• GPT-3 Paper: Language Models are Few-Shot Learners (Brown et al. 2020) — 附录里的训练超参（warmup 375M tokens、cosine decay 到 10%、clip 1.0）是当今 LLM 训练事实模板
• Karpathy — nanoGPT — 300 行代码完整复现 GPT-2 训练，包含 AdamW + warmup + cosine + clip 全套，可以亲手跑
• Karpathy — A Recipe for Training Neural Networks — 神经网络训练实战经验大全
• Lilian Weng — Some Math behind Neural Tangent Kernel — 想深入理解优化器 vs 损失景观可参考