简单softmax函数实现：在处理大输入值或小输入值时可能会遇到数值稳定性问题，比如溢出和下溢。

1
2
def softmax_naive(x):
		return torch.exp(x) / torch.exp(x).sum(dim=0)

数值稳定性问题（溢出与下溢）

“Softmax 实现的数值稳定性问题”主要源于计算机处理浮点数的精度和范围限制。

Softmax 的公式是：

$$\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$$

问题 A：溢出（Overflow）

  • 原因： 当输入 $x_i$ 是一个非常大的正数时。
  • 发生情况： $e^{x_i}$ 的值会迅速变得极大（呈指数增长）。如果这个值超过了计算机浮点数（如 32 位浮点数）所能表示的最大范围，就会导致 上溢（Overflow）。
  • 结果： 结果变为 $\text{Inf}$ (无穷大)，导致计算失败或返回 NaN。

问题 B：下溢（Underflow）

  • 原因： 当输入 $x_i$ 是一个非常小的负数时。
  • 发生情况： $e^{x_i}$ 的值会变得非常接近于零。如果这个值小于计算机浮点数所能表示的最小非零范围，就会导致 下溢（Underflow）。
  • 结果： 结果被强制舍入为 $\mathbf{0}$。虽然单个 $e^{x_i}$ 变为 0 似乎问题不大，但在求和 $\sum e^{x_j}$ 时，如果所有项都下溢为 0，分母就变成了 0，从而导致除零错误。

PyTorch 等框架的解决方案：稳定的 Softmax

为了解决上述数值问题，高性能的深度学习框架（如 PyTorch、TensorFlow）在实现 Softmax 时，通常会采用一个数学技巧来实现稳定的 Softmax。

核心技巧：最大值归一化（Max Normalization Trick）

这个技巧利用了一个数学特性：给输入向量 $\mathbf{x}$ 加上或减去一个常数 $C$，不改变 Softmax 的最终结果。

$$\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}} = \frac{e^{x_i - C} e^{C}}{\sum_j e^{x_j - C} e^{C}} = \frac{e^{x_i - C}}{\sum_j e^{x_j - C}}$$

在实践中，我们选择 $C$ 为输入向量 $\mathbf{x}$ 中的最大值：$C = \max(\mathbf{x})$。

为什么选择 $C = \max(\mathbf{x})$？

  1. 解决溢出问题： 将所有的 $x_i$ 都减去 $\max(\mathbf{x})$ 后，最大的那个指数项的输入就变成了 $\max(\mathbf{x}) - \max(\mathbf{x}) = 0$。
    • 这样，分子和分母中最大的指数项 $e^0 = 1$。
    • 所有其他的指数项 $e^{x_i - \max(\mathbf{x})}$ 的输入都是非正数（$\le 0$），因此它们的指数值都会在 $[0, 1]$ 之间，从而有效防止了上溢。
  2. 解决下溢问题： 虽然减去最大值后，一些项的输入会变得更小（负得更多），可能导致下溢，但关键在于：** Softmax 的结果并不会因为分母为零而失败**。只有当分子的项与分母的项的差值过大时，才会导致精度损失。通过这个归一化，我们确保了至少有一项（最大值对应项）的数值是稳定的（为 1），整个计算过程的数值稳定性大大提高。

知识串联：实践建议的理解

“在实践中建议使用 Softmax 的 PyTorch 实现，该实现经过了大量性能优化”

这句话的意思是：

  1. 数值稳定保障： PyTorch 的 torch.softmax()（以及 torch.nn.CrossEntropyLoss 内部实现的 LogSoftmax 技巧）已经内置了上述最大值归一化等数值稳定机制。
  2. 性能优化： PyTorch 的 Softmax 核心是在底层（通常是 C++ 或 CUDA）实现的，针对现代硬件（GPU）进行了高度并行和内存访问优化，运行速度比任何简单的 Python 实现都要快得多。

因此，应该理解这个数值稳定性的原理，但在编写代码时，应始终使用内置的、经过优化的库函数。