Source-linked AI summary
Property-Based Testing for Robust Softmax and Cross-Entropy
Muhammad Maaz
TL;DR
Softmax and cross-entropy can exhibit numerical failures even with max-shifting. This paper uses property-based tests to rediscover those failures, finds an additional failure in a proposed fix, and presents a log1p-based alternative that passes the tests and supports causal attention masking.
Problem
Softmax collapse can zero the hard-target gradient, while softmax cancellation can leave the soft-target gradient nonzero as student logits approach teacher logits.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
Abstract
from arXiv · showhide
A recent paper of Barron et. al. (2026) identified two numerical issues with the softmax and cross-entropy functions in virtually all standard machine learning libraries (e.g., PyTorch and JAX). We show that these issues can be recovered through simple property-based tests in Hypothesis. However, our tests identify an additional edge case where their implementation fails. We then propose an even more robust softmax implementation which passes these tests and also supports use in causal attention masking. We compare our implementation with theirs and discuss some of the differences. Overall, this work shows that property-based testing can be a powerful tool for discovering and fixing numerical issues in machine learning libraries. Our implementations and property-based tests can be found at https://github.com/mmaaz-git/softmax-fix.
1 Introduction
The introduction frames softmax collapse and cancellation as numerical failures that persist despite max-shifting. It presents property-based testing as a way to rediscover them and proposes a more robust fix that supports causal attention masking.
1 Introduction
- Hard-target cross-entropy is the negative log-softmax at the target class, a one-hot-target special case of soft-target cross-entropy.The introduction also notes that max-shifting is widely used in practice, including in PyTorch and JAX.Source: §1 · p. 1
- Barron et al. (2026) identify softmax collapse and cancellation: the hard-target gradient can round to zero, while the soft-target gradient can plateau above zero.The paper notes that the collapse phenomenon was also studied in grokking by Prieto et al. (2025).Source: §1 · p. 1
- Property-based tests in Hypothesis can rediscover these failures by checking mathematical properties over generated inputs.The paper describes the approach as specifying a property and checking it against randomly generated inputs.Source: §1 · p. 2
- The paper reports that Barron et al. (2026)’s fix still fails on an additional edge case and proposes a more robust softmax that supports causal attention masking.Implementations and tests are released in the cited codebase.Source: §1 · p. 2
2 Rediscovering the Bugs with Property-Based Testing
This section tests mathematical properties of softmax and cross-entropy, explains how finite-precision arithmetic breaks them, and uses Hypothesis to find concrete failures in PyTorch.
2 Rediscovering the Bugs with Property-Based Testing
- When the target logit dominates, a sufficiently large margin makes 1 + e^-m round to 1, so the computed target probability is 1 while non-target probabilities remain positive.The standard backward formula then yields a zero target diagonal derivative but nonzero off-diagonal entries, violating the Jacobian row-sum identity.Source: §2 · p. 2
- Under round-to-nearest, collapse is predicted when margin m exceeds b ln 2: 5.545 for bf16, 16.636 for float32, and 36.737 for float64.These predicted thresholds align with Barron et al. (2026)’s numerical experiments.Source: §2 · p. 3
- For a fixed teacher, matching student logits should give a zero cross-entropy gradient, but different floating-point operation sequences can produce different rounding errors.The teacher probabilities use softmax(w), while the backward pass can compute student probabilities using exp(logsoftmax(z)).Source: §2 · p. 3
- The tests check that each softmax Jacobian row sums to zero and that identical student and fixed-teacher logits produce an exactly zero gradient.The Jacobian test uses a relative tolerance of 10^-12, scaled by the sum of absolute row entries.Source: §2 · p. 3
- Hypothesis generates up to 10,000 lists of at least two finite floats for each test, which are converted to tensors before the properties are checked.The implementation also allows users to choose which softmax or cross-entropy function to test.Source: §2 · pp. 3–4
- [0.0, 37.0] breaks the Jacobian row-sum property, while [0.0, 1.0] gives a nonzero gradient when student and teacher logits match.For the first vector, the target-logit margin of 37 exceeds the predicted float64 collapse threshold of 36.737.Source: §2 · p. 4
3 A Remaining Failure in the Proposed Fix
Barron et al.’s proposed softmax fix still fails for some finite inputs: max-shifting can produce an all-−∞ logsumexp whose backward pass is NaN. This differs from the built-in PyTorch failure, which returns zero backward values for the example discussed.
3 A Remaining Failure in the Proposed Fix
- Hypothesis finds that Property 1 fails for Barron et al.’s fix, producing a NaN Jacobian for an extreme finite input.Their implementation max-shifts logits, excludes the maximum from logsumexp, and applies softplus to the result.Source: §3 · p. 4
- Max-shifting can overflow a finite logit to −∞, leaving an all-−∞ logsumexp whose backward pass is NaN.The forward value is correctly −∞, but the backward value becomes NaN.Source: §3 · p. 4
- For [10308, −10308], both implementations return [1, 0], but standard PyTorch gives all-zero backward values while Barron et al.’s gives all NaNs.Source: §3 · p. 4
4 Our Fix
The proposed logsoftmax fix uses log1p to retain finite gradients in the extreme case where max-shifting leaves only −∞ non-maximum logits. It also handles causal attention masking, where that same pattern can occur.
4 Our Fix
- The fix max-shifts logits, sums exponentials of the non-maximum logits, applies log1p, then subtracts the log-normalizer.This computes logsoftmax and hence softmax.Source: §4 · p. 4
- At S = 0, log1p has finite derivative 1, avoiding logsumexp’s NaN backward on an all-−∞ vector; the implementation passes both tested properties.Source: §4 · p. 5
- For [10308, −10308], the implementation returns forward probabilities [1, 0] and an all-zero backward value.Source: Table 1 · p. 5, §4 · p. 5
- The fix reuses Barron et al.’s soft-target cross-entropy implementation, changing only logsoftmax and combining log1p with explicit maximum handling.Blanchard et al. (2021) also use log1p to compute logsumexp.Source: §4 · p. 5
- Causal attention masking can create a row with one finite logit and all other logits at −∞; the fix gives correct forward and backward values in this case.Source: §4 · p. 5
5 Conclusion
The conclusion presents property-based testing as a way to rediscover known numerical issues and catch failures in proposed fixes. It reports a log1p-based softmax fix that also handles causal attention masking and argues for wider use of these tests in machine-learning libraries.
5 Conclusion
- Property-based tests rediscover the softmax and cross-entropy issues identified by Barron et al. and expose NaNs from their proposed softmax on finite inputs.Source: §5 · p. 5
- The new log1p-based fix is more robust to overflow after max-shifting and adequately handles causal attention masking.Source: §5 · p. 5
- The authors hope property-based testing will be adopted more widely in machine-learning library test suites to discover numerical issues.Source: §5 · p. 5