← Back to all news
arXiv Machine Learning By Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen

A Fix for Softmax and Cross-Entropy

Read the original on arXiv Machine Learning →

The paper identifies two numerical bugs in standard softmax, log‑softmax, and cross‑entropy implementations: Softmax Collapse, where a dominant logit causes its gradient to be rounded to zero, and Softmax Cancellation, where student and teacher logits that are equal still produce non‑zero gradients. The authors provide straightforward implementations that correct both issues. Their fix removes the previously observed "Slingshot" phenomenon, accelerates grokking with weight decay, prevents accuracy collapse without weight decay, keeps student networks aligned with teachers during distillation, and demonstrates that post‑memorization logit growth is bounded by an equilibrium between weight decay and Adam’s epsilon hyperparameter.

Machine-generated by The Flow from the publisher's headline and feed description — not written or checked by a human. The full article lives at arXiv Machine Learning.

Related stories

arXiv Machine Learning
Aug 27

A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

The paper investigates how the sampled-token reverse-KL loss in on‑policy distillation distributes updates across tokens. By analyzing the gradient of the per‑token K2 estimator, the authors find that tokens with low student probability and large teacher‑student gaps receive disproportionately large gradient norms. They propose Surprise‑aware Reweighting (SuRe), a lightweight weighting rule that further amplifies this allocation, and demonstrate that SuRe improves math metrics on Qwen3 student models without harming out‑of‑domain performance.

By Bing Shao, Jiazheng Zhang, Long Ma, Yujiong Shen, Senjie Jin, Xin Guo, Yuming Yang, Mingxu Chai, Zhiheng Xi, Tao Gui, Qi Zhang, Xuanjing Huang
efficiencybenchmarks
More like this →
arXiv Machine Learning
Oct 2

How Accurate Is Accurate Enough?

The paper investigates how the magnitude of numerical errors in a learning system should be evaluated, arguing that error impact depends on the current learning state, the class affected, and the class weight. By analyzing softmax cross‑entropy, the authors derive exact bounds on loss changes for fixed error sets and use these to provide finite‑error guarantees for losses, probabilities, predictions, and gradients. They then invert these guarantees to certify the maximum allowable primitive error for a given learning‑level requirement, demonstrating through von Mises‑Fisher learning that identical errors can have vastly different effects across states.

By Ningkang Peng, Qianfeng Yu, Jingyang Mao, Xiaoqian Peng, Yanhui Gu
More like this →
Hugging Face Trending Papers
Sep 30

How Accurate Is Accurate Enough?

How accurate must a numerical approximation be within a learning system? Primitive error alone cannot answer this question: errors of the same magnitude can have very different consequences for losses...

More like this →
arXiv Machine Learning
Oct 9

Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes

Deep neural networks can exhibit periodic loss spikes known as the "Slingshot Mechanism" during long‑term training. This paper shows that such spikes arise from finite‑precision floating‑point arithmetic: when the correct‑class logit becomes much larger than the others, the gradient for that class is rounded to zero while others remain nonzero, breaking the zero‑sum constraint and creating a positive feedback loop that inflates feature and classifier norms. The authors term this Numerical Feature Inflation (NFI) and demonstrate that it explains the rapid norm growth, reappearance of gradients, and loss spikes, while also showing that partial absorption can still cause hidden norm growth without visible spikes.

By Hanqing Liu, Jianjun Cao, Yuanze Li, Zijian Zhou
More like this →