Property Tests Catch a Hidden Softmax Bug That Breaks PyTorch and JAX
A six-page follow-up uses Hypothesis to stress-test the recent DeepMind softmax fix, finds a new edge case, and ships a tougher implementation.
- Follow-up paper reproduces the DeepMind softmax bugs using property-based tests in Hypothesis.
- Finds Barron et al.'s fixed softmax returns NaN on an extreme float64 input pair.
- Proposes a more robust softmax that passes all tests and supports causal attention masking.
- Two invariants tested: Jacobian rows sum to zero, and cross-entropy gradient is zero when student equals teacher.
- Code, tests and counterexamples released at github.com/mmaaz-git/softmax-fix for PyTorch and JAX.
- Argues property-based testing should be standard practice for numerical ML library primitives.
A follow-up paper turns two softmax bugs into a compact case study for property-based testing. Using Hypothesis, the author reproduces the numerical failures described in the Barron et al. paper, finds an extreme float64 input that breaks its proposed fix, and supplies a rewrite that passes the published test suite and supports causal attention masks.
Softmax converts logits, the unnormalized scores produced by a model, into probabilities. Adding the same constant to every logit leaves those probabilities unchanged, so each row of the softmax Jacobian, the matrix of output derivatives with respect to inputs, should sum to zero.
Softmax Collapse violates that invariant when one logit greatly exceeds the others. The dominant probability rounds to one, its gradient rounds to zero, and smaller gradients remain nonzero. Their sum then drifts from zero, producing an internally inconsistent gradient.
Softmax Cancellation affects soft-target cross-entropy, including knowledge distillation. When student and teacher logits match, the gradient with respect to the student should be exactly zero while the teacher remains fixed. Current PyTorch and JAX implementations can instead return small nonzero gradients. Barron and colleagues report that Gemini Flash 3.7 originally identified both bugs.
Two Invariants Catch the Failures
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
