BP神经网络中Softmax与交叉熵损失函数的原理剖析与高效实现

1次阅读

没有评论

共计 2534 个字符，预计需要花费 7 分钟才能阅读完成。

在深度学习的分类任务中，Softmax 函数和交叉熵损失函数（Cross-Entropy Loss）的组合几乎是标配。然而，这一组合在实际应用中常常会遇到数值稳定性问题，导致训练过程中的梯度消失或爆炸，甚至出现 NaN 值。

Softmax 的数值溢出问题：当输入值过大或过小时，Softmax 函数的指数运算可能导致数值溢出（overflow 或 underflow），尤其是在处理极端值时，计算结果可能直接变为无穷大或零。
梯度计算特性：交叉熵损失函数与 Softmax 结合时，梯度计算会简化为一个非常简洁的形式，但如果不注意数值稳定性，这种简化反而会成为问题的源头。
典型错误案例：许多初学者在实现时直接使用原始 Softmax 公式，忽略了数值稳定性问题，导致训练过程中频繁出现 NaN 值，模型无法收敛。

Softmax 函数的定义为：

$$
\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^n e^{x_j}}
$$

其梯度可以通过链式法则推导出来。假设我们有一个损失函数 L，那么 Softmax 的梯度为：

$$
\frac{\partial L}{\partial x_i} = \text{Softmax}(x_i) \left(\frac{\partial L}{\partial \text{Softmax}(x_i)} – \sum_{j=1}^n \text{Softmax}(x_j) \frac{\partial L}{\partial \text{Softmax}(x_j)} \right)
$$

当交叉熵损失函数与 Softmax 结合时，梯度计算可以简化为：

$$
\frac{\partial L}{\partial x_i} = \text{Softmax}(x_i) – y_i
$$

其中，y_i 是真实标签的 one-hot 编码。这一简化大大减少了计算量，但也带来了数值稳定性问题。

为了避免数值溢出，我们可以使用 Log-Sum-Exp（LSE）技巧。具体来说，计算 Softmax 时，先减去输入的最大值：

$$
\text{Softmax}(x_i) = \frac{e^{x_i – \max(x)}}{\sum_{j=1}^n e^{x_j – \max(x)}}
$$

这样可以确保指数运算的输入值不会过大，从而避免数值溢出。

以下是一个使用 PyTorch 实现数值稳定的 Softmax 和交叉熵损失的代码示例：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 使用 torch.logsumexp 稳定计算
class StableSoftmaxCrossEntropy(nn.Module):
    def __init__(self):
        super(StableSoftmaxCrossEntropy, self).__init__()

    def forward(self, logits, labels):
        # 减除最大值，确保数值稳定性
        logits_max, _ = torch.max(logits, dim=1, keepdim=True)
        logits_stable = logits - logits_max
        # 计算 log_softmax
        log_softmax = logits_stable - torch.logsumexp(logits_stable, dim=1, keepdim=True)
        # 计算交叉熵损失
        loss = -torch.sum(labels * log_softmax, dim=1)
        return loss.mean()

# 示例：GPU 并行化处理
logits = torch.randn(128, 10).cuda()  # batch_size=128, num_classes=10
labels = torch.randint(0, 10, (128,)).cuda()
labels = F.one_hot(labels, num_classes=10).float()

criterion = StableSoftmaxCrossEntropy()
loss = criterion(logits, labels)
print(loss)