Two Major Numerical Bugs in Softmax and Cross-Entropy and Their Fixes

In deep learning models, softmax, log-softmax, and cross-entropy are components used in almost every classification network. However, standard implementations hide two numerical bugs that can lead to incorrect gradient computation, training instability, and even affect model generalization. A recent technical article analyzes these two issues in detail and proposes concise fixes.

Two Hidden Numerical Bugs

The article points out two flaws in standard implementations: Softmax Collapse and Softmax Cancellation.

Softmax Collapse: Gradient Vanishing and Non-Zero Sum

When the target logit is much larger than other logits, standard softmax rounds the gradient of that logit to zero. This causes the sum of gradients to no longer be zero, breaking the balance during optimization. This phenomenon is particularly evident when the model is highly confident and can trigger abnormal behavior in later training stages.

Gradient properties of Softmax: For softmax output (p_i = e^{z_i}/\sum_j e^{z_j}), the gradient of the loss with respect to logit (z_k) can be written as (\partial L/\partial z_k = p_k (\partial L/\partial p_k) - p_k \sum_j p_j (\partial L/\partial p_j)). When using cross-entropy with target (t), the gradient simplifies to (p_k - \delta_{kt}). Theoretically, the sum of all components is always zero, which is a necessary result of the probability simplex constraint. However, in floating-point arithmetic, when (p_t) is extremely close to 1, (p_t - 1) becomes 0 due to rounding error, and other (p_k) also underflow to 0, causing the gradient vector to be all zeros. This "gradient vanishing" does not come from activation function saturation but from insufficient numerical representation precision. It breaks gradient balance, preventing the optimizer from perceiving the relative relationship between the target class and other classes, thus causing training anomalies when the model is highly confident. The fix typically uses log-softmax and the log-sum-exp trick for stable computation, ensuring that gradients maintain correct signs and magnitudes even under extreme logit differences.

Softmax Cancellation: Spurious Gradients in Distillation

In knowledge distillation, when student logits approach teacher logits, the standard implementation produces non-zero gradients—even when the two are exactly equal. This means the student network is pushed in the wrong direction, causing drift. The article calls this phenomenon "Softmax Cancellation," which explains part of the reason for unstable student performance in distillation.

Gradient composition in knowledge distillation: The distillation loss typically consists of hard-label cross-entropy and soft-label KL divergence. The soft-label part is (L_{soft} = \tau^2 \mathrm{KL}(p_T^{\tau} | p_S^{\tau})), and its gradient with respect to the student's logits is (\tau (p_S^{\tau} - p_T^{\tau})). When the student and teacher logits are exactly equal, this gradient should theoretically be zero. But in the standard implementation, the numerator and denominator of softmax are computed separately, and subtracting two equal large numbers produces rounding error, so (p_S - p_T) is not strictly zero but leaves a tiny but non-zero "spurious gradient." This gradient direction is random and continuously pushes the student parameters to drift, especially in later training when the student has already fully fitted the teacher, the drift accumulates into performance degradation. The fix is to jointly compute log-softmax and directly subtract the logit differences, avoiding separate normalization and then subtraction, thereby ensuring that equal inputs produce strictly zero gradients.

Fixes and Empirical Results

The authors propose simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs simultaneously. These modifications do not increase computational complexity but bring several improvements in training dynamics.

Eliminating the "Slingshot" Phenomenon

The previously observed "Slingshot" phenomenon—where the training loss suddenly rebounds at some stage—disappears after the fix. This indicates that the numerical stability of gradient computation has a direct impact on the training trajectory.

Accelerating Grokking and Preventing Accuracy Collapse

With weight decay, the fixed implementation accelerates the emergence of grokking (the sudden transition from memorization to generalization); without weight decay, it prevents accuracy collapse. These two points show that numerical bugs interfere with regularization effects and generalization behavior.

Grokking and weight decay: Grokking refers to the phenomenon where, after the model has long overfitted the training set (training loss near zero), validation performance suddenly jumps from random level to good generalization after prolonged training. It usually occurs in settings with small datasets and strong weight decay, essentially because the optimizer gradually finds simpler solutions driven by weight decay. Numerical bugs interfere with this process: the gradient imbalance caused by Softmax Collapse may cause the model to lock into a memorized solution too early, delaying or preventing the transition to a generalized solution; after the fix, gradient balance is restored, and weight decay can more effectively guide parameters toward low-norm regions, thereby accelerating grokking. On the other hand, without weight decay, logits may grow unbounded, and numerical underflow in standard softmax causes gradients to vanish completely, preventing the model from further adjustment, manifesting as sudden accuracy collapse; the fixed implementation maintains effective gradients through stable log-softmax, preventing this collapse.

Stabilizing Knowledge Distillation

After the fix, the student network no longer drifts from the teacher network, and the distillation process becomes more stable. This is particularly important for tasks that rely on soft labels.

Revealing the Balancing Mechanism of Logit Growth

The article further finds that after the model memorizes the training data, logit growth is constrained by a balance between weight decay and the ε hyperparameter in the Adam optimizer. This finding provides a new perspective for understanding the dynamics of later training stages.

Adam's ε and logit norm: In Adam updates, the parameter step is approximately (\eta \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)). When gradients remain small, (\sqrt{\hat{v}_t}) is also small, and (\epsilon) in the denominator begins to dominate, compressing the actual step to (\eta \cdot \hat{m}_t / \epsilon), which is approximately proportional to the gradient magnitude rather than adaptively scaled. In later training, weight decay pulls parameters toward zero, while the classification loss pushes logits to increase to maintain correct classification, creating a tug-of-war. Without weight decay, logits continue to grow until numerical limits; with weight decay, the parameter norm is constrained, and after logits grow to a certain level, the weight decay term in the gradient balances with the classification loss term, stopping growth. (\epsilon) determines the effective step size near this equilibrium point, thereby affecting the final magnitude of logits. This mechanism explains why adjusting (\epsilon) changes the dynamics of later training and provides a new perspective for understanding the interaction between numerical implementation and optimizer hyperparameters.

Implications for Practitioners

For engineers engaged in model training and tuning, these findings remind us that numerical implementation details of fundamental components can have profound effects. Especially in knowledge distillation, large-scale classification tasks, and scenarios using weight decay, adopting the fixed softmax and cross-entropy implementations may lead to more stable training and better final performance.

The implementations provided in the article are simple and direct, and it is recommended to replace standard implementations in custom training loops to avoid these hidden numerical issues.

One-line Summary

Standard softmax and cross-entropy have two numerical bugs; fixing them stabilizes training, accelerates generalization, and reveals the logit growth mechanism.

Paragraph Summary

The article analyzes two hidden numerical bugs in standard softmax and cross-entropy implementations: Softmax Collapse, where gradients round to zero when the target logit is much larger than others, breaking gradient balance; and Softmax Cancellation, where spurious non-zero gradients arise in knowledge distillation when student and teacher logits are equal, causing student drift. The authors propose simple fixed implementations of log-softmax, softmax, and cross-entropy that do not increase computational complexity. Empirical results show that the fixes eliminate the Slingshot phenomenon of sudden training loss rebound, accelerate grokking, prevent accuracy collapse without weight decay, stabilize knowledge distillation, and further reveal that logit growth is constrained by a balance between weight decay and the ε hyperparameter in the Adam optimizer. The article recommends that engineers replace standard implementations in custom training loops to avoid these numerical issues.