数值稳定的 Softmax 与 Log-Sum-Exp 实现

你好。上一课我们用离散分布计算了期望、方差与协方差：它们帮助我们描述 rollout 的平均奖励、波动程度，以及成功率与 token 成本等指标是否共同变化。

这一课进入这些概率量在模型计算中的实际载体：softmax、log-softmax 与 log-sum-exp。语言模型在每个位置都会输出一整个词表的 logits；这些数值必须被稳定地转换为概率或对数概率，才能用于采样、计算 token 对数概率，以及后续的 SFT、DPO、PPO 等训练目标。关键不只是会写公式，而是知道为何直接照公式实现会得到 inf、0 或 nan，并能用减去最大值的方式避免它们。

---

从 logits 到概率：softmax 做了什么

设一个模型在某个 token 位置对 $V$ 个候选 token 输出 logits：

$$\mathbf{z}=(z_1,z_2,\dots ,z_V)$$

logits 是任意实数分数，不是概率，也不要求为正，更不要求和为 $1$。softmax 将它们转换为概率分布：

$$p_i=\mathrm{softmax}(\mathbf{z}{)}_i=\frac{\exp (z_i)}{{\sum }_{j=1}^V\exp (z_j)}$$

因此每个 $p_i$ 满足：

$$0<p_i<1$$

并且：

$$\sum _{i=1}^V{p}_i=1$$

例如，若三个候选 token 的 logits 为：

$$\mathbf{z}=(1,2,0)$$

则：

$$\exp (\mathbf{z})=(\exp (1),\exp (2),\exp (0))$$ $$\approx (2.718,7.389,1)$$

归一化后：

$$\mathrm{softmax}(\mathbf{z})\approx (0.245,0.665,0.090)$$

第二个 token 的 logit 最大，因此得到最高概率；但 softmax 不会简单地“只选最大值”，而是保留其他候选的相对概率。

在语言模型中，若当前词表大小是 $V$，则一个位置的 logits 通常形状为：

$$[V]$$

一个 batch 的完整序列 logits 通常形状为：

$$[B,T,V]$$

其中 $B$ 是 batch size，$T$ 是序列长度。softmax 必须沿词表维 $V$ 计算；在 PyTorch 中这通常就是 dim=-1。如果误沿 batch 维或序列维归一化，程序也许不会报错，但得到的“概率”将没有正确的语言模型含义。

---

问题不在公式，而在有限精度计算

直接按定义实现 softmax，似乎只需要三步：

1. 对每个 logit 求指数；
2. 求指数之和；
3. 每项除以总和。

但浮点数能表达的数值范围有限。指数函数增长极快：当 logit 足够大时，$\exp (z_i)$ 可能无法表示，变成无穷大 inf。

考虑：

$$\mathbf{z}=(1000,1001,999)$$

数学上，这组 logits 完全合理：第二项仅比第一项大 $1$，第三项仅低 $2$。但在许多常用浮点格式中：

$$\exp (1000)=\infty$$

于是朴素计算会产生：

$$\frac{\infty }{\infty }$$

结果是未定义的 nan。

另一端也会出问题。若：

$$\mathbf{z}=(-1000,-1000,-1000)$$

直接计算时，每个指数项可能下溢为 $0$：

$$\exp (-1000)\approx 0$$

分母也变成 $0$，最终形成：

$$\frac{0}{0}$$

同样得到 nan。这两组 logits 虽然绝对值极大，但它们的相对差异并不奇怪；softmax 真正关心的也正是相对差异。

---

不改变 softmax 的关键：所有 logits 同时平移

softmax 有一个极其重要的性质：对所有 logit 加上或减去同一个常数，输出概率完全不变。

令 $c$ 为任意常数，则：

$$\mathrm{softmax}(\mathbf{z}{)}_i=\frac{\exp (z_i)}{{\sum }_j\exp (z_j)}$$

也等于：

$$\frac{\exp (z_i-c)}{{\sum }_j\exp (z_j-c)}$$

原因是分子和分母都带有同一个因子 $\exp (-c)$，它们会抵消：

$$\frac{\exp (z_i-c)}{{\sum }_j\exp (z_j-c)}=\frac{\exp (z_i)\exp (-c)}{{\sum }_j\exp (z_j)\exp (-c)}=\frac{\exp (z_i)}{{\sum }_j\exp (z_j)}$$

因此，最自然的选择是：

$$m=\underset{i}{\max }{z}_i$$

计算平移后的 logits：

$$\widetilde{z_i}=z_i-m$$

此时一定有：

$$\widetilde{z_i}\le 0$$

且至少有一个元素恰好为：

$$0$$

所以：

$$0<\exp (\widetilde{z_i})\le 1$$

指数阶段不再出现正向溢出。

回到刚才的例子：

$$\mathbf{z}=(1000,1001,999)$$

最大值为：

$$m=1001$$

平移后：

$$\tilde{\mathbf{z}}=(-1,0,-2)$$

现在再求指数：

$$\exp (\tilde{\mathbf{z}})\approx (0.368,1,0.135)$$

指数和为：

$$0.368+1+0.135\approx 1.503$$

因此稳定 softmax 为：

$$\mathrm{softmax}(\mathbf{z})\approx (0.245,0.665,0.090)$$

这与 logits 为 $(1,2,0)$ 时的结果相同，因为两组 logits 只相差同一个常数 $999$。softmax 表达的是相对偏好，而不是绝对 logit 原点的位置。

---

log-sum-exp：在对数域稳定地处理归一化常数

定义 log-sum-exp：

$$\mathrm{LSE}(\mathbf{z})=\log \left(\sum _{i=1}^V\exp (z_i)\right)$$

它看起来只是“指数、求和、取对数”的组合，却是深度学习中最常见的数值稳定操作之一。它正是 softmax 分母的对数。

利用最大值 $m$，稳定形式为：

$$\mathrm{LSE}(\mathbf{z})=m+\log \left(\sum _{i=1}^V\exp (z_i-m)\right)$$

注意：这里与 softmax 略有不同。

• softmax 在整体平移后保持不变；
• log-sum-exp 的值会随平移改变，因此需要在最后把 $m$ 加回来。

对前面的例子：

$$\mathrm{LSE}(1000,1001,999)$$

可稳定计算为：

$$1001+\log \left(\exp (-1)+\exp (0)+\exp (-2)\right)$$ $$=1001+\log (1.503)$$ $$\approx 1001.408$$

中间计算从未要求我们实际表示 $\exp (1001)$。

为什么说 LSE 是 max 的平滑近似

令：

$$m=\underset{i}{\max }{z}_i$$

由于平移后至少一项为 $0$，其指数为 $1$，并且所有指数项不超过 $1$，因此：

$$1\le \sum _i\exp (z_i-m)\le V$$

取对数并加回 $m$：

$$m\le \mathrm{LSE}(\mathbf{z})\le m+\log V$$

所以 LSE 至少等于最大 logit，最多只比它大 $\log V$。当一个 logit 比其余项大得多时，其余指数项几乎不贡献总和，LSE 会非常接近最大值。

这也解释了一个实践现象：若某个 token 的 logit 远高于其他 token，softmax 分布会非常尖锐；若多个 token logits 接近，概率质量则会在它们之间分散。

---

从 LSE 得到稳定的 log-softmax

后训练代码中，很多时候不需要直接保存 softmax 概率，而需要某个 token 的对数概率。例如，计算一段生成文本在当前策略下的逐 token log probability 时，使用的就是 log-softmax。

定义：

$$\log \mathrm{softmax}(\mathbf{z}{)}_i=\log \left(\frac{\exp (z_i)}{{\sum }_j\exp (z_j)}\right)$$

利用对数规则：

$$\log \mathrm{softmax}(\mathbf{z}{)}_i=z_i-\log \left(\sum _j\exp (z_j)\right)$$

也就是：

$$\log \mathrm{softmax}(\mathbf{z}{)}_i=z_i-\mathrm{LSE}(\mathbf{z})$$

将稳定 LSE 代入：

$$\log \mathrm{softmax}(\mathbf{z}{)}_i=(z_i-m)-\log \left(\sum _j\exp (z_j-m)\right)$$

这是应当直接实现和调用的形式。

一个常见但不稳定的写法是先算 softmax，再取对数：

log_probs = torch.log(torch.softmax(logits, dim=-1))

它在概率极小时可能先把该概率下溢为 $0$，随后计算：

torch.log(0)

产生负无穷 -inf。虽然某些极小概率的 log probability 本来就很负，但无意义的 -inf 会污染损失、KL、重要性采样比率和梯度计算。

应优先直接计算：

log_probs = torch.log_softmax(logits, dim=-1)

---

用 PyTorch 手写稳定实现

下面的版本假定沿指定维度的每个切片中至少存在一个有限 logit。这正是普通 softmax 的基本前提；若某一行全部被 attention mask 处理成负无穷，则该行不存在合法概率分布，必须在上游避免这种输入。

import torch


def stable_softmax(logits, dim=-1):
    max_logits = logits.max(dim=dim, keepdim=True).values
    shifted_logits = logits - max_logits

    exp_logits = torch.exp(shifted_logits)
    normalizer = exp_logits.sum(dim=dim, keepdim=True)

    return exp_logits / normalizer


def stable_logsumexp(logits, dim=-1, keepdim=False):
    max_logits = logits.max(dim=dim, keepdim=True).values
    shifted_logits = logits - max_logits

    lse = max_logits + torch.log(
        torch.exp(shifted_logits).sum(dim=dim, keepdim=True)
    )

    if keepdim:
        return lse

    return lse.squeeze(dim)


def stable_log_softmax(logits, dim=-1):
    max_logits = logits.max(dim=dim, keepdim=True).values
    shifted_logits = logits - max_logits

    log_normalizer = torch.log(
        torch.exp(shifted_logits).sum(dim=dim, keepdim=True)
    )

    return shifted_logits - log_normalizer

这里的 keepdim=True 很重要。假设输入是：

logits.shape == (batch_size, vocab_size)

沿最后一维取最大值后，max_logits 形状保持为：

(batch_size, 1)

于是它能自动广播到：

(batch_size, vocab_size)

对每一条样本分别减去各自的最大 logit。若不保留该维度，就需要额外调整形状；初学阶段保留维度通常更直观，也更不容易在 batch 计算中写错。

可以用极端输入验证手写实现：

logits = torch.tensor(
    [
        [1000.0, 1001.0, 999.0],
        [-1000.0, -1000.0, -1000.0],
    ]
)

manual_probs = stable_softmax(logits, dim=-1)
manual_log_probs = stable_log_softmax(logits, dim=-1)
manual_lse = stable_logsumexp(logits, dim=-1)

print(manual_probs)
print(manual_probs.sum(dim=-1))
print(manual_log_probs)
print(manual_lse)

应观察到：

• manual_probs 中没有 nan；
• 每一行概率之和接近 $1$；
• 第一行第二个 token 的概率最大；
• 第二行三个 logits 相同，因此三个概率都接近 $1/3$；
• manual_lse 是有限值，即使输入中含有 $1000$ 或 $-1000$。

再与 PyTorch 内置实现核对：

torch_probs = torch.softmax(logits, dim=-1)
torch_log_probs = torch.log_softmax(logits, dim=-1)
torch_lse = torch.logsumexp(logits, dim=-1)

print(torch.allclose(manual_probs, torch_probs))
print(torch.allclose(manual_log_probs, torch_log_probs))
print(torch.allclose(manual_lse, torch_lse))

对于上述有限输入，三个比较都应为 True。

---

手写用于理解，生产代码优先使用内置算子

掌握手写版本很重要：之后阅读训练框架源码时，你会不断看到“减最大值”“logsumexp”“log_softmax”“归一化常数”等操作；理解它们能帮助判断某段代码是在计算概率、对数概率还是某种稳定损失。

但在真实训练代码中，应优先使用 PyTorch 的内置算子：

probs = torch.softmax(logits, dim=-1)

log_probs = torch.log_softmax(logits, dim=-1)

lse = torch.logsumexp(logits, dim=-1)

原因有三点：

1. 数值稳定性。 内置实现专门处理了常见边界情况；
2. 性能。 底层通常使用融合 kernel，避免创建过多中间张量；
3. 梯度正确性。 自动微分路径经过了大量测试。

对语言模型的 logits：

logits.shape == (batch_size, sequence_length, vocab_size)

最常用的写法仍然是：

log_probs = torch.log_softmax(logits, dim=-1)

它代表：对每个 batch 样本、每个序列位置，在词表维度上形成一个对数概率分布。

一个实用的调试习惯是检查三个不变量：

probs = torch.softmax(logits, dim=-1)
log_probs = torch.log_softmax(logits, dim=-1)

assert torch.isfinite(probs).all()
assert torch.isfinite(log_probs).all()

assert torch.allclose(
    probs.sum(dim=-1),
    torch.ones_like(probs.sum(dim=-1)),
    atol=1e-6,
)

如果这些检查失败，优先排查：

• dim 是否指向词表维；
• 输入 logits 中是否已经有 nan 或 inf；
• 某一整行是否都被 mask 成了负无穷；
• 是否在其他地方先做了不稳定的 exp、除法或 log。

---

小结

本课的核心不是记住一段代码，而是建立一条可靠的计算原则：指数之前先平移，尽量在对数域中完成归一化。

关键公式如下。

softmax：

$$\mathrm{softmax}(\mathbf{z}{)}_i=\frac{\exp (z_i-m)}{{\sum }_j\exp (z_j-m)}$$

其中：

$$m=\underset{j}{\max }{z}_j$$

稳定 log-sum-exp：

$$\mathrm{LSE}(\mathbf{z})=m+\log \left(\sum _j\exp (z_j-m)\right)$$

稳定 log-softmax：

$$\log \mathrm{softmax}(\mathbf{z}{)}_i=(z_i-m)-\log \left(\sum _j\exp (z_j-m)\right)$$

在实践中：

• softmax 用于获得显式概率，如采样或概率检查；
• log-softmax 用于获得稳定的 token 对数概率；
• log-sum-exp 是稳定归一化常数的核心操作；
• 生产训练优先使用 torch.softmax、torch.log_softmax 与 torch.logsumexp；
• 手写版本则是阅读 loss、策略概率、KL 与采样源码的基础。

至此，第一模块的数学、概率与数值基础告一段落。下一模块将从“概率分布如何表示”继续走向“模型参数如何通过数据被估计”，首先学习似然函数与对数似然函数。