当完美公式遇上有限精度：Softmax数值稳定的第一课

发表于 2026-10-06 分类于深度学习 Waline：阅读次数：本文字数： 3.1k 阅读时长 ≈ 3 分钟

当完美公式遇上有限精度：Softmax数值稳定的第一课

在深度学习教材里，Softmax与交叉熵的组合总是优雅而简洁。给定logits向量 $z$ 和标签 $y$，联合损失通常写作：

$L = -\log\left(\frac{e^{z_y}}{\sum_j e^{z_j}}\right)$

这个公式在数学上无懈可击。于是许多初学者会自然地写出如下PyTorch代码：

# 危险示范：数学正确 ≠ 数值稳定
def naive_softmax_ce(logits, target):
    probs = torch.exp(logits) / torch.sum(torch.exp(logits))
    loss = -torch.log(probs[target])
    return loss

这段代码在小规模测试中运行完美，但在真实训练中可能是一颗定时炸弹。当某个logit值稍大（比如超过88），exp(89) 在float32下就会溢出为 inf；而当logit极小时，exp(-100) 会下溢为0，随后的 log(0) 直接产出 -inf 或 NaN，模型训练就此崩溃。

根本原因在于计算机的浮点数不是实数，其表示范围是有限的。

浮点数的安全区比你想象的窄

我们习惯认为数字可以无限大或无限小，但float32只有约7位有效十进制精度，其正数范围大约在 $1.2\times10^{-38}$ 到 $3.4\times10^{38}$ 之间。

想象一条数轴：右侧终点是 3.4e38，超过即上溢变为 inf；左侧靠近0处有个最小正数 1.2e-38，低于此值即下溢变为0；中间才是安全计算区。

而指数函数 $e^x$ 恰恰是最容易冲出安全区的操作：$e^{89} \approx 3.6\times10^{38}$ 已越界，$e^{-88} \approx 1.0\times10^{-39}$ 已下溢。这意味着logits只要偏离原点约90，朴素实现就会失效，而这在深度网络的训练过程中几乎必然发生。

数值稳定性要解决的，就是在有限的浮点表示内，等价地计算数学上定义良好的表达式。

推导安全的Softmax

回到朴素实现的两个致命伤：exp(z) 上溢、log(0) 下溢。我们能否在不改变数学结果的前提下，让 exp 的输入变小？

关键观察是：对Softmax分子分母同时乘以常数，结果不变。取对数后，这等价于给所有logits加上同一个常数 $c$：

$\log\frac{e^{z_y}}{\sum e^{z_j}} = \log\frac{e^{z_y + c}}{\sum e^{z_j + c}} = (z_y + c) - \log\sum e^{z_j + c}$

现在问题变成：选什么 $c$ 最安全？

若选 $c = -\text{mean}(z)$：平移后仍有正有负，最大值的 exp 仍可能上溢。
若选 $c = -\min(z)$：所有值变为非负，最大值反而更大，上溢风险加剧。
若选 $c = -\max(z)$：所有 $z_i + c \leq 0$，exp 的输出落在 $(0, 1]$ 区间，上溢彻底消除；而分母至少包含一项 $e^0 = 1$，求和后大于等于1，log 的参数永远为正，下溢也被阻断。

这就是LogSumExp技巧的核心：用一次减法，换整个计算的安全。

手写稳定版：理解原理

def stable_softmax_ce(logits, target):
    """数值稳定的Softmax+CE，仅用于学习原理"""
    # 关键：减去最大值，防止exp上溢
    max_z = torch.max(logits)
    shifted = logits - max_z

    # log(sum(exp(shifted))) 现在安全了
    log_sum_exp = torch.log(torch.sum(torch.exp(shifted)))

    # 最终损失 = -shifted[target] + log_sum_exp
    loss = -shifted[target] + log_sum_exp
    return loss

注意这里只做了一次 exp 和一次 log，相比朴素版本既稳定又高效。

生产环境：永远使用框架内置函数

手写版有助于理解，但绝不要在生产代码中使用。PyTorch等框架的内置实现不仅数值稳定，还经过底层优化，且自动处理边界情况：

# 正确做法：cross_entropy内部已融合log_softmax
# 实际使用时注意logits需包含batch维度，此处为演示省略
loss = torch.nn.functional.cross_entropy(logits, target)

cross_entropy 内部正是基于LogSumExp实现，且支持label smoothing、ignore_index等实用功能。数值稳定性是框架的责任，你的责任是正确调用它。

梯度与Loss的隐藏陷阱

使用了log_softmax，前向计算就高枕无忧了吗？未必。数值稳定性贯穿整个计算图，以下两个坑在新手实践中极为常见。

分离logits与labels导致的精度损失

有些教程会将Softmax和CrossEntropy拆成两步写：

#  不推荐：中间产物probs会丢失精度
probs = torch.softmax(logits, dim=-1)
loss = -torch.log(probs[target])

问题在于：softmax输出的概率值可能极小（如 $10^{-10}$），float32下虽未下溢为0，但小数点后很多位被四舍五入丢掉了；随后的log操作会将这个丢失的精度放大。即使前向没报NaN，梯度也可能因精度损失而失真。

正确做法始终是融合计算：cross_entropy或log_softmax + nll_loss在底层直接对logits操作，跳过概率中间态，从根源避免精度丢失。

自定义Loss时的log参数保护

当你设计自定义损失函数（如Focal Loss、对比学习中的InfoNCE）时，常需手动计算log(something)。若something由网络输出经sigmoid得到，理论上应在(0,1)区间，但浮点误差可能导致其恰好为0或略小于0。

防御性编程习惯是对所有手动log调用的参数加一个极小常数：

# 安全写法：假设prob来自sigmoid输出
prob = torch.sigmoid(raw_output)
eps = 1e-7  # 远大于float32最小正数，又远小于典型概率值
loss = -torch.log(torch.clamp(prob, min=eps))

注意clamp优于直接加eps（如prob + eps），因为后者在prob接近1时会引入不必要的偏移。clamp仅在危险区域生效，正常范围保持原值。当然，如果参数有数学保证严格为正（比如LogSumExp中的分母），这层防护可以省略。

写在最后

回顾全文，数值稳定性并非高深理论，而是几个可落地的习惯：

Softmax+CE永远用融合算子：cross_entropy或log_softmax + nll_loss是唯一正确写法，手写版仅用于理解原理。
手动log加保护是好习惯：自定义Loss中对不确定为正的参数用clamp(min=1e-7)防护，但有数学保证时可省略。
怀疑NaN先查数值：训练崩溃时，优先检查是否有未保护的指数、对数运算，而非盲目调学习率或换优化器。

框架已经封装了绝大多数稳定实现，你的核心任务是识别那些框架无法自动保护的边界场景。当你能预判哪里可能溢出、哪里会丢失精度，数值稳定性就从玄学调试变成了可控的工程实践。

下一步建议：在你的下一个项目中，刻意检查一遍Loss计算路径，确认是否全部使用了融合算子。这一个动作，就能避免未来无数次莫名其妙的训练失败。