DeepMind paper reveals floating-point rounding bug in Softmax and cross-entropy; Gemini Flash 3.7 was the first to discover and fix it
AI教科书全错！底层公式惊天Bug曝光，Gemini率先自我修复
Jon Barron's team at Google DeepMind submitted the paper "Fixing Softmax and Cross-Entropy" to arXiv, pointing out that the standard Softmax and cross-entropy implementations in PyTorch, JAX and textbooks have a floating-point rounding bug. The standard implementation first subtracts the maximum score to prevent overflow; when the model is sufficiently confident, the extremely small numbers in the denominator are rounded to zero, and the correct gradient disappears accordingly. The paper calls this Softmax collapse, and the slingshot effect previously regarded as a grokking feature is actually caused by this bug. The fix is to use log-sum-exp to compute the margin for cross-entropy and return softplus, and in distillation scenarios to use expm1 and log1p in log space to eliminate ghost gradients. In the modular addition (a + b mod 31) task, under bfloat16 precision the standard code's test accuracy is only 93.3% and highly unstable, while after the fix it reaches 100.0%.
The article explains the mechanism by which the more confident the model is, the more easily the correct gradient is erased by floating-point rounding, and gives the corresponding Softmax and distillation fix formulations.