Standard implementations of softmax, log-softmax, and cross-entropy contain two numerical bugs that affect training dynamics. The first bug, Softmax Collapse, rounds the target logit's gradient to zero when one logit dominates the others, causing gradients that no longer sum to zero. The second bug, Softmax Cancellation, produces non-zero gradients when student and teacher logits are equal, which occurs when both are close together. These bugs appear in widely used deep learning frameworks and can silently degrade model training.
Origins of Softmax Collapse
Softmax Collapse happens when a single logit is much larger than the remaining logits. The gradient for that dominant logit gets rounded to zero, breaking the property that gradients should sum to zero. This means the model receives misleading signal during backpropagation, and learning signals for other logits become unreliable. The bug is especially likely in the early stages of training when logits have not yet been spread out.
Origins of Softmax Cancellation
Softmax Cancellation appears in knowledge distillation and similar student-teacher settings. When student logits are close to teacher logits, the cross-entropy gradient does not vanish even if the logits are exactly equal. This produces non-zero updates even when the student has already matched the teacher's predictions, causing the student to drift away from the teacher rather than staying aligned.
A corrected implementation
The authors present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. The key change involves modifying how the gradient flow is computed so that the target logit's gradient is neverzeroed out, and so that equal logits produce zero gradient as expected. The implementations avoid costly extra passes while preserving numerical stability. These fixes are short enough to replace standard library calls in most deep learning frameworks.
Effects of the fix
The corrected implementations eliminate the previously observed "Slingshot" phenomenon, in which parameters overshoot and then correct multiple times during training. The fix speeds up grokking with weight decay, which is the sudden emergence of test accuracy after a long period of slow training. It prevents accuracy collapse when weight decay is not used. The fix also stops student networks from drifting away from the teacher during distillation, leading to more stable transfer of knowledge. Additionally, the authors show that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's ε hyperparameter, giving a theoretical explanation for previously observed scaling behavior.
Practical implications for practitioners
For engineers training language models or other deep networks, the fix means that standard training loops can be made more reliable without changing architecture or hyperparameter search. The corrected softmax and cross-entropy can be dropped in as drop-in replacements. The results suggest that many previously observed training behaviors — grokking trajectories, accuracy plateaus, and distillation drift — may have been partially caused by these two long-standing numerical bugs. The fix is particularly relevant for research code, where training loops are frequently customized and may be more exposed to these edge cases.