PyTorch与TensorFlow交叉熵实现差异及数值稳定性疑问

Great question—this is such a common gotcha when digging into cross-entropy under the hood, since both frameworks hide clever numerical stability tricks behind their APIs. Let’s break this down step by step to clear up your confusion.

First: Why Naive Cross-Entropy Breaks #

Before we get to the frameworks, let’s recap why the "naive" cross-entropy implementation (calculate softmax first, then take -log(p) ) fails with extreme logits:

• When a logit z_k is very large, exp(z_k) will overflow to inf, making the softmax calculation inf/inf = nan (not a number).
• When all logits are very small, exp(z_k) will underflow to 0, making the softmax denominator 0—again leading to nan or -log(0) = inf.

These are the numerical issues both PyTorch and TensorFlow are solving.

PyTorch’s Solution: Merged Calculation + Log-Sum-Exp Trick #

Your confusion here comes from a common misunderstanding: PyTorch’s nn.CrossEntropyLoss does not compute a standalone softmax tensor before calculating cross-entropy. Instead, it merges the two steps into a single mathematical expression optimized for stability.

Here’s the breakdown:

1. The naive cross-entropy can be rewritten algebraically:

   CE = -log( exp(z_y) / sum(exp(z_k)) ) = log(sum(exp(z_k))) - z_y
   

2. To avoid overflow in sum(exp(z_k)), PyTorch uses the log-sum-exp trick:

   log(sum(exp(z_k))) = max_z + log(sum(exp(z_k - max_z)))
   

   Here, max_z is the largest value in your logit vector. Subtracting max_z from every logit ensures the largest exponent we compute is exp(0) = 1—all other exponents are exp(negative number) (smaller than 1). This makes the sum safe to compute without overflow.

By merging these steps, PyTorch never calculates the raw softmax values that would cause overflow. The entire computation stays within stable numerical bounds.

TensorFlow’s Solution: Same Core Trick, No Underflow Risk #

Your concern about TensorFlow’s implementation suffering underflow is unfounded—like PyTorch, it uses the merged cross-entropy + log-sum-exp approach to avoid both overflow and underflow.

Here’s why underflow isn’t an issue:

• We never compute exp(z_k) directly. Instead, we compute exp(z_k - max_z). Even if z_k is much smaller than max_z, modern floating-point systems can represent extremely small non-zero values (down to ~1e-308), so this won’t underflow to 0 for most practical cases.
• Even if some terms do underflow to 0, the sum sum(exp(z_k - max_z)) will always be at least 1 (from the exp(0) term corresponding to max_z). This means log(sum(...)) will never be log(0) = inf, keeping the entire calculation stable.
• TensorFlow doesn’t need to output the raw softmax probabilities, so it never encounters the p_k = 0 scenario that would break -log(p_k).

Key Takeaway #

Both frameworks solve numerical issues the same way: they skip computing standalone softmax values entirely, merging cross-entropy and softmax into a single calculation that uses the log-sum-exp trick to keep all operations within safe numerical ranges. Your initial confusion came from assuming they implemented the naive softmax-first approach, but the APIs are designed to hide these stability optimizations from users.

内容的提问来源于stack exchange，提问作者kusur