Softmax and cross-entropy bugs fixed improve model training stability
A Fix for Softmax and Cross-Entropy
Machine Learning
Summary
Softmax and cross-entropy are basic tools in AI that help computers learn from examples. The authors discovered two small but important bugs in how these tools are normally implemented, which cause errors in training. They provide simple fixes that prevent these errors, making AI training more reliable and efficient. Their improvements stop models from behaving strangely during training and keep models closer to their intended outputs.
What this means in practice
- •For machine learning engineers: Improve neural network training by fixing softmax and cross-entropy gradient errors to prevent accuracy drops and unstable learning.
- •For ai model distillation teams: Maintain better fidelity between teacher and student models during distillation by eliminating unintended gradient drifts caused by prior bugs.
Authors
Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
Abstract
Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously-observed "Slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's $ε$ hyperparameter.