Most engineering teams audit their ML stack from the top down: model architecture, data pipeline, hyperparameter choices, evaluation metrics. What they rarely audit is the numerical correctness of the primitives those choices rest on. A recent paper from researchers at Google DeepMind has identified two concrete bugs in the standard implementations of softmax and cross-entropy present across every major ML framework, and the failure modes they produce are precisely the kind that get misattributed to everything except their actual cause (Barron et al., arXiv 2026).
What the Bugs Actually Are
The first bug, which Barron et al. term Softmax Collapse, occurs when one logit is significantly larger than the others. In this regime, floating-point arithmetic causes the gradient of the dominant logit to round to zero. The result is that gradients no longer sum to zero, which is a property that correct softmax gradient computation must satisfy. The training signal becomes corrupted at exactly the point where the model is most confident, which is the moment when gradient precision matters most for fine-grained learning.
The second bug, Softmax Cancellation, is particularly relevant for distillation pipelines. It arises when student logits are close to teacher logits, producing non-zero gradients even when the two distributions are effectively identical. A student network that has successfully matched its teacher continues to receive a gradient signal pushing it away. The model drifts not because it has failed to learn, but because the loss function is numerically incapable of recognising convergence.
Both bugs are not edge cases in unusual configurations. They are triggered by the normal dynamics of training: confident predictions and near-converged distillation are routine states in any production training run.
How These Bugs Manifest in Production Training
The practical consequences documented by Barron et al. include the Slingshot phenomenon, a pattern of periodic loss spikes observed during training with weight decay that had previously been attributed to optimiser dynamics or learning rate scheduling. The fixed implementation eliminates it. That is a strong signal: a phenomenon that teams have been working around at the hyperparameter level turns out to be an artefact of numerical error at the primitive level.
The paper also shows that without weight decay, standard implementations allow accuracy to collapse entirely in certain training regimes, while the corrected implementation prevents this. When teams encounter accuracy collapse, the diagnostic path typically leads to data quality, augmentation strategy, or architecture changes. Numerical instability in the loss function is rarely on the checklist.
This is the structural risk. Silent numerical errors produce symptoms that are indistinguishable from the symptoms of other, more familiar failure modes. Engineering time gets spent on the wrong layer of the stack.
The Distillation Case Is the Most Commercially Significant
Knowledge distillation is now a standard technique for producing deployment-sized models from larger teacher models. It is central to how enterprises reduce inference cost without retraining from scratch. If the loss function used during distillation cannot correctly identify when a student has matched its teacher, the student will continue to update away from the teacher distribution after convergence.
Barron et al. demonstrate this directly: the corrected implementation stops student networks from drifting away from the teacher during distillation, while the standard implementation does not. For teams running distillation pipelines at scale, this means the compressed models they are deploying may have diverged from the intended teacher behaviour in ways that are not visible in aggregate accuracy metrics.
The divergence is also likely to be inconsistent across training runs, which makes it difficult to detect through standard evaluation. A model that behaves differently on two separate distillation runs from the same teacher is a reliability problem, not just a performance problem.
What a Numerical Audit Looks Like in Practice
The fix proposed by Barron et al. is implementable: corrected versions of log-softmax, softmax, and cross-entropy that address both bugs through careful reformulation of the computation. The paper provides the implementations directly. For most teams, the immediate action is to evaluate whether their current training and distillation pipelines are using the corrected primitives, and if not, to run controlled comparisons against the fixed versions on representative training jobs.
The broader lesson is methodological. Numerical correctness at the primitive level should be treated as an auditable property, not an assumption. This means:
- Maintaining a record of which framework versions and loss function implementations are used in each training run
- Including numerical sanity checks (gradient sum properties, loss behaviour at convergence) as part of training monitoring
- Treating unexplained training instabilities as potential numerical errors before exhausting architectural or data-side explanations
The cost of this audit is low relative to the cost of diagnosing training failures that have been misattributed for months.
The Broader Principle for Engineering Leadership
The softmax and cross-entropy bugs are a specific instance of a general problem: the further a bug sits from the model output, the longer it takes to detect and the more confidently it gets attributed to something else. Numerical errors in loss computation are maximally far from the output in diagnostic terms, because they affect every training step silently rather than producing a visible failure.
For CTOs and ML engineering heads, the practical implication is that the reliability of a production model is a function of the entire computational stack, including primitives that have been treated as solved for years. The fact that these bugs exist in every major framework is not an indictment of those frameworks. It is evidence that even well-maintained, heavily-used infrastructure can carry correctness assumptions that have never been formally verified against the full range of training dynamics.
Building in periodic numerical audits, particularly when adopting new training regimes or distillation strategies, is a reasonable engineering response to that reality.
Where Vector Labs Fits
We build production ML systems where numerical correctness is a design constraint, not an afterthought. In our cardiovascular certification work, class-weighted loss functions and rigorous validation against medical device software standards were central to achieving Class 2A certification on wearable ECG data. If you are running training or distillation pipelines where model reliability is a hard requirement, contact us at vector-labs.ai/contacts.
FAQs
Barron et al. identify the bugs in standard implementations across every major ML framework. The root cause is a shared numerical approach to softmax and cross-entropy computation rather than a framework-specific implementation choice, which is why the issue is present broadly. The corrected implementations the paper provides can be adopted independently of which framework your team is using.
The most direct diagnostic is to re-run a representative training job using the corrected loss implementation and compare training dynamics against the original run. If you have observed Slingshot-style loss spikes, accuracy collapse without an obvious data or architecture cause, or inconsistent distillation results across runs, these are the specific symptoms documented in the paper. Gradient sum checks during training can also surface the Softmax Collapse bug in real time.
The bugs affect gradient computation during training and therefore affect the weights the model converges to. At inference, softmax is applied without gradient computation, so the numerical error does not propagate in the same way. The risk is that models trained with the buggy implementations have learned subtly different weight configurations than they would have under correct gradient signals, and those differences are baked into the deployed model.
It is relevant for any training process that uses softmax-based cross-entropy loss, including fine-tuning. Fine-tuning typically involves fewer steps than training from scratch, which may reduce cumulative exposure to the bugs, but the Softmax Cancellation bug is specifically triggered by near-convergence conditions, which fine-tuning reaches quickly. Distillation-based fine-tuning is the highest-risk scenario given the documented student drift behaviour.
The corrected implementations are straightforward to adopt as custom loss functions without waiting for upstream framework patches. Given that the bugs affect training dynamics in ways that are difficult to detect after the fact, waiting introduces ongoing risk to any training runs executed in the interim. For teams running high-stakes training or distillation jobs, adopting the fix directly is the lower-risk path.

