A recent paper published on arXiv by researchers Jonathan T. Barron, Dan Gnanapragasam, and Michael Tschannen details the discovery and resolution of two significant numerical bugs affecting widely used machine learning functions: softmax, log-softmax, and cross-entropy. These bugs, termed "Softmax Collapse" and "Softmax Cancellation," can lead to unstable model training and unexpected behavior, particularly in scenarios involving high confidence predictions or knowledge distillation. The researchers have developed simple, algebraically equivalent implementations that address these issues without introducing additional computational overhead.
Softmax Collapse occurs when a model becomes highly confident in its prediction. In such cases, the gradient for the target logit can be rounded to zero due to floating-point precision limitations. This results in gradients that no longer sum to zero, violating a fundamental property of these functions and potentially causing training instability, such as the "Slingshot" phenomenon where loss spikes occur. The paper explains that this bug arises because the calculation involves adding a very small number (resulting from the exponentiation of a large negative logit difference) to 1.0. Once this small number is below the precision threshold of the floating-point format, it is lost to rounding, causing the denominator to evaluate to exactly 1.0 and effectively zeroing out the target gradient. The researchers observed that this precision cliff varies depending on the floating-point format, occurring at margins of approximately 5.5 for bf16, 16.6 for fp32, and 36.7 for fp64.
The second identified bug, Softmax Cancellation, impacts knowledge distillation and KL regularization. This issue arises when the logits of a student model are very close to those of a teacher model. In these situations, standard implementations can produce non-zero gradients even when the student and teacher logits are identical. This occurs because the probabilities are computed separately and then subtracted, leading to a loss of precision when dealing with nearly equal values. The research demonstrates that this can cause student networks to drift away from their teachers during training.
The proposed solution involves rewriting the softmax and cross-entropy functions in log-space using functions like `softplus`, `expm1`, and `log1p`. This approach avoids the problematic addition of small numbers to 1.0 and directly computes the difference between student and teacher logits in log-space, thereby maintaining gradient accuracy across different floating-point precisions. The authors state that these fixes eliminate the previously observed "Slingshot" phenomenon, speed up model "grokking" (the point at which a model generalizes beyond its training data) when weight decay is applied, and prevent accuracy collapse in models trained without weight decay. Furthermore, the fixes ensure that student networks do not drift away from their teachers during distillation. The paper also notes that these revised implementations can be easily integrated into existing deep learning frameworks without significant performance degradation.