Barron paper patches softmax and cross-entropy gradient bugs

Standard implementations of softmax, log-softmax and cross-entropy carry two numerical bugs that distort training behavior under weight decay, grokking and knowledge distillation, according to a new arXiv preprint by Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen.

The authors name them Softmax Collapse and Softmax Cancellation. The first, the paper says, "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 "occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal."

Their corrected implementations, the abstract reports, eliminate the previously observed "slingshot" phenomenon, speed up grokking with weight decay, prevent accuracy collapse without weight decay, and "stops student networks from drifting away from the teacher during distillation." The paper also claims that "after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's ε hyperparameter." The abstract carries no benchmark figures; the claims are stated as described behaviors. Two researchers we track flagged the preprint the day it posted.