On October 22, 2026, researchers Jonathan T. Barron, Dan Gnanapragasam, and Michael Tschannen from Google DeepMind released a preprint titled A Fix for Softmax and Cross-Entropy. The paper identifies two numerical defects, termed Softmax Collapse and Softmax Cancellation, occurring in standard machine learning libraries like PyTorch and JAX.
The authors demonstrate that when a model becomes highly confident, standard calculations in float32 precision can lead to the loss of small gradient components. Specifically, adding a very small value to one during softmax computation results in rounding errors, causing the gradient for the correct class to vanish prematurely while other components remain.
To address these issues, the researchers propose a mathematical reformulation based on algebraic simplification. By rewriting the fraction used in the gradient calculation, the small value is moved to the numerator, preventing it from being lost during the summation process and maintaining numerical stability.
The study highlights that while these errors do not stop the entire training process, they introduce distortions in the gradient signal. The proposed fix preserves the original mathematical intent while ensuring more accurate numerical results in deep learning frameworks.
