DeepMind says standard Softmax and cross-entropy implementations contain numerical bugs, with Gemini credited for the initial fix

Google DeepMind researcher Jon Barron says standard implementations of Softmax and cross-entropy in PyTorch, JAX and even textbook formulations contain two subtle numerical bugs. According to Barron and a new arXiv paper, one issue can drive the correct gradient to zero when model confidence gets high enough, while another can leave a residual “ghost gradient” in distillation even when a student model exactly matches its teacher. The paper presents rewritten formulations intended to avoid adding extremely small values to 1 and to keep calculations in log space with functions such as expm1 and log1p. Barron also said in an X post that Gemini Flash 3.7 first found and fixed the bug, Gemini Flash 3.8 ran most of the experiments, and Gemini 4 Argon drafted the paper and figures. The reported results show test accuracy in a modular addition task rising from 93.3% to 100.0% in bfloat16 after the fix, while an independent engineer, Satyajit Ghana, said he reproduced the mechanism in PyTorch on CPU and confirmed the fp32 results described in the paper.