交叉熵（CE）：从一个预测到训练里的 loss

模型给正确答案 40% 的概率，和给它 1% 的概率，都可能被算成“预测错了”。但训练时，我们希望分清这两种错误：后者更需要改。交叉熵（cross-entropy，CE）就能表达这种差别。

先看一个三分类的小例子，再把它接到语言模型（language model, LM）。第一次读懂前 3 节就够了；要写训练代码，再看后面的 mask、平均方式和梯度。

正确类别拿到了多少概率？

假设任务是判断一句话在说“退款”“物流”还是“其他”。模型给出概率 [0.2, 0.3, 0.5]，这条数据的标签是“其他”。单条样本的 loss 是：

$$ \ell=-\log 0.5\approx0.6931. $$

这里用自然对数，单位是 nat。正确类别的概率变成 0.9 时，loss 约为 0.105；降到 0.01 时，loss 约为 4.605。正确答案拿到的概率越低，惩罚越大。

为什么不只数对错？当正确类别还排第 2 名时，它的概率从 0.1 涨到 0.4，准确率可能完全不变，CE 却能反映进步。反过来，CE 下降也不保证生成质量、校准（calibration）或每个群体的准确率都变好。

从 logits 直接算，为什么更稳？

令 $z_c$ 是类别分数，$y$ 是正确类别。把 Softmax 代入负对数：

$$ \ell=-\log p_y =\log\sum_c e^{z_c}-z_y. $$

直接算 exp(1000) 会溢出。我们先减去这一行最大的分数 $m$：

$$ \ell=\log\sum_c e^{z_c-m}-(z_y-m). $$

所有分数减去同一个数，不会改变 Softmax。最大的指数现在是 $e^0=1$，其他指数不超过 1。下面手写一个单 batch、类别在最后一维的版本，再和库函数核对：

import math
import torch
import torch.nn.functional as functional

def ce_from_logits(logits, targets):
    logits = logits.to(torch.promote_types(logits.dtype, torch.float32))
    shifted = logits - logits.max(dim=-1, keepdim=True).values
    log_partition = shifted.exp().sum(dim=-1).log()
    selected = shifted.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    return (log_partition - selected).mean()

logits = torch.tensor([[2., 3., 5.]], dtype=torch.float64).log()
targets = torch.tensor([2])
loss = ce_from_logits(logits, targets)
assert math.isclose(loss.item(), math.log(2))
torch.testing.assert_close(loss, functional.cross_entropy(logits, targets))
large_logits = torch.tensor([[1000., 999., -1000.]], dtype=torch.float64)
torch.testing.assert_close(
    ce_from_logits(large_logits, targets),
    functional.cross_entropy(large_logits, targets),
)

这里没有先生成概率再取 log。后者可能先把很小的概率舍入成 0，再得到无穷大的 loss。实际训练优先用现成的稳定实现，手写代码是为了看清它在算什么。

减去最大值只限制了每个指数项，还要留意求和用什么精度。例如 65,536 个 FP16 零 logits 对应的指数项全是 1；若求和结果仍存成 FP16，就会溢出，尽管正确 loss 只有 $\log 65536\approx11.09035$。因此代码先把 FP16/BF16 提升到 FP32 再计算，已有 FP64 则保留。这个转换不切断计算图：梯度仍回到原始 logits，并按原始 dtype 存储，不能因此免除低精度梯度的舍入或下溢。

直接调用库函数也不等于自动开启 autocast：在本文核对的 PyTorch 2.8 CPU 上，直接给 functional.cross_entropy 传入上述 FP16 输入也会得到 inf；显式转 FP32，或在 CPU autocast 中调用库 CE，结果才是有限的。CPU autocast 策略会把库 CE 放到 FP32（2026-10 对照 2.14 文档，cross_entropy_loss 仍在这张 FP32 列表里，但上面的 inf 结果没有在 2.14 上重跑），但不会自动把任意手写 exp → sum → log 组合都提升。不同设备的策略不同，这里不把 CPU 结果外推到 CUDA。

硬标签和软标签，差在哪里？

硬标签（hard target）只指定一个类别，相当于 one-hot 分布 [0, 0, 1]。软标签（soft target）可以是 [0.1, 0.2, 0.7]：目标不是把全部概率压到第 3 类，而是接近这一整个分布。

$$ \ell=-\sum_c q_c\log p_c. $$

沿用刚才的预测 [0.2, 0.3, 0.5]，软标签 CE 约为 0.8869。比硬标签的 0.6931 大，并不说明模型突然变差了：目标分布变了，不能直接比两个数字。

soft_targets = torch.tensor([[.1, .2, .7]], dtype=torch.float64)
log_probabilities = functional.log_softmax(logits, dim=-1)
soft_loss = -(soft_targets * log_probabilities).sum(-1).mean()
torch.testing.assert_close(
    soft_loss, functional.cross_entropy(logits, soft_targets)
)
assert math.isclose(soft_loss.item(), 0.8869413785, abs_tol=1e-9)

蒸馏（distillation）可以用 teacher 的分布作目标；label smoothing 则把 one-hot 和一个平滑分布混合。这些目标都要检查非负、归一化及类别顺序。若目标 $q$ 固定，最小化 CE 与最小化 $D_{KL}(q\|p)$ 的梯度相同，因为两者只差与模型无关的 $H(q)$；数值不必相同。完整算例见蒸馏。

梯度为什么是“预测减目标”？

对单条、未加权、归一化目标的 Softmax CE：

$$ \frac{\partial\ell}{\partial z_c}=p_c-q_c. $$

硬标签例子的梯度是 [0.2, 0.3, −0.5]。梯度下降会降低前两类的分数，提高正确类别的分数。换成软标签，梯度变成 [0.1, 0.1, −0.2]，更新方向仍然合理，但没那么急着把第 3 类推到 100%。

二分类：别靠 clamp 补救 log(0)

二分类可以只输出一个 logit $z$，通过 sigmoid 得到正类概率。标签 $y\in\{0,1\}$ 时：

$$ \begin{aligned} \ell={}&-y\log\sigma(z)\\ &-(1-y)\log(1-\sigma(z)). \end{aligned} $$

它也能直接在 logits 上稳定计算。先算 sigmoid、再把概率夹到 [1e-7, 1−1e-7]，虽然可能避免 log(0)，却改了目标：落在截断区间外时，clamp 的梯度是 0。

例如正确标签为 1，模型却给出 logit −100。稳定 BCE 的 loss 约为 100，梯度约为 −1，应该明显提高分数。夹住概率后，loss 只有约 16.12，而且传回 logit 的梯度为 0，恰好不再纠正这个严重错误。

wrong_logit = torch.tensor([-100.], dtype=torch.float64, requires_grad=True)
binary_target = torch.ones_like(wrong_logit)
stable_loss = functional.binary_cross_entropy_with_logits(wrong_logit, binary_target)
stable_loss.backward()
assert math.isclose(stable_loss.item(), 100.)
assert math.isclose(wrong_logit.grad.item(), -1.)
wrong_logit.grad = None
clamped_loss = -wrong_logit.sigmoid().clamp(1e-7, 1 - 1e-7).log().mean()
clamped_loss.backward()
assert wrong_logit.grad.item() == 0.

BCEWithLogitsLoss把 sigmoid 与 BCE 合并计算，不需要你提前做 sigmoid。多标签任务通常给每个标签一个这样的 logit；“退款”和“物流”若可以同时成立，就不该强行要求它们的概率之和为 1。

到了语言模型：哪些位置算分？

语言模型每个位置输出词表（vocabulary）上的分类分数，常见 shape 是 [batch, sequence, vocabulary]。不过不能直接把整张张量平均：

• 先对齐目标：causal LM 的当前位置预测下一个 token；只 shift 一次。
• 再选位置：padding、未监督的用户文本等不计入 loss，但仍需正确的 attention mask。
• 最后定分母：这里按有效目标 token 数取平均，不按 padding 后的长度。

一个小 batch 的有效 loss 分别是 [2] 和 [1, 1, 1]。按 token 平均得到 $5/4=1.25$；先按样本平均再平均得到 $(2+1)/2=1.5$。前者每个 token 等权，后者每条样本等权。两者都能定义成目标，但别把它们当成同一件事。

同理，gradient accumulation 不能在有效 token 数不同时，未经调整就把每个 microbatch 的 mean 再平均。要先确定训练目标的分母；分布式时还要对齐跨 rank 的计数和梯度归约。按有效 token 总数归一化的做法见一次训练更新，完整对话例子见 SFT。

weighted_logits = torch.zeros(2, 2, dtype=torch.float64)
hard_targets = torch.tensor([0, 1])
class_weights = torch.tensor([1., 3.], dtype=torch.float64)
one_hot_targets = functional.one_hot(hard_targets, 2).double()
hard_mean = functional.cross_entropy(weighted_logits, hard_targets, weight=class_weights)
soft_mean = functional.cross_entropy(weighted_logits, one_hot_targets, weight=class_weights)
assert math.isclose(hard_mean.item(), math.log(2))
assert math.isclose(soft_mean.item(), 2 * math.log(2))

怎样确认自己真的写对了？

先拿小张量检查，不急着跑大模型。本文代码能按顺序在 CPU 执行；更完整的教学实现还检查输入、空 mask、soft target 和一次 shift，配套回归测试包含梯度和极端数值对照。

教学实现保留原有输入检查；检查通过后，软标签与 logits 的 CE 运算使用至少 FP32 的共同 dtype，任一输入为 FP64 就保留 FP64。不重新归一化或 detach 软标签，也不把 loss 转回低精度。提升精度不会恢复输入里已丢失的信息；忽略位置、有效 token 分母与一次 shift 的约定也不改变。

这些检查能发现实现错误，不等于证明模型效果好。跑通之后，还得回到数据覆盖、验证集和任务本身。