A Fix for Softmax and Cross-Entropy

Jonathan T. Barron; Dan Gnanapragasam; Michael Tschannen

cs.LGOct 22, 2026

A Fix for Softmax and Cross-Entropy

Authors: Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen

Organizations: Google DeepMind

Abstract

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, 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 bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously observed "slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's $ε$ hyperparameter.

Figures & tables

Panel of Figure 1 — Figure 1: The two bugs in standard softmax and cross-entropy. (a) Softmax Collapse: As a model becomes more confident in the correct class (thereby increasing the logit margin $m$ ), the true gradient on the target logit shrinks exponentially, but standard implementations incorrectly round that gradient to zero once $m$ crosses a precision cliff ( $\approx 5.5$ in bf16, $16.6$ in fp32, and $36.7$ in fp64). (b) Softmax Cancellation: In soft-target distillation, as student logits $\mathbf{z}_{\mathrm{S}}$ approach teacher logits $\mathbf{z}_{\mathrm{T}}$ , the true gradient shrinks toward zero, but standard implementations lose precision and plateau at a non-zero noise floor. Our fix eliminates both bugs and yields accurate gradients across all precisions.

Panel of Figure 2 — Figure 2: The Slingshot mechanism ( Thilak et al., 2022 ) with and without weight decay (5 random seeds, fp32). Top row ( $a+b\bmod 31$ , $\lambda=1$ , $\epsilon=10^{-8}$ , Varma et al. 2023 ): Standard cross-entropy hits the fp32 cliff ( $m=c_{32}\approx 16.6$ , where loss rounds to $0$ ) and repeatedly drops and rebounds, slowing and degrading grokking; fixed cross-entropy crosses $c_{32}$ smoothly, groks faster on all 5 seeds, and saturates just above $c_{32}$ at $m\approx 18.4$ (Section 5 ). Bottom row ( $a^{2}+b^{2}\bmod 97$ , $\lambda=0$ , $\epsilon=10^{-16}$ , Doshi et al. 2024 ): Without weight decay, hitting $c_{32}$ around step $1250$ triggers a spike that drops standard cross-entropy test accuracy from $61\%$ to $23\%$ , whereas fixed cross-entropy continues smoothly to $85\%$ .

Panel of Figure 3 — Figure 3: Normalized Adam step $|g_{t}|/(|g_{t}|+\epsilon)$ vs. logit margin $m$ at $\epsilon=10^{-8}$ . With our fixed cross-entropy, the step stays near $1$ and decays smoothly around $\log(\nicefrac{{1}}{{\epsilon}})\approx 18.4$ , whereas standard CE zeroes out $\tilde{g}_{t}$ at $c_{16}\approx 5.5$ and $c_{32}\approx 16.6$ .

Figure 4: Logit margins when training with our fixed cross-entropy and LaProp on $a+b\bmod 31$ in fp32 ( $\beta_{1}=0.9$ , $\beta_{2}=0.98$ ; 3 seeds per setting as solid lines, analytical predictions as dashed lines). (a) With weight decay ( $\lambda=1$ ), the margin for each $\epsilon>0$ grows along the $\epsilon=0$ trajectory and levels off at the predicted equilibrium $m^{\star}(\epsilon)=\log(\nicefrac{{1}}{{\epsilon}})$ (colored $y$ -ticks), whereas $\epsilon=0$ keeps growing until the fp32 gradients themselves underflow at $m\approx 81.2$ (gray). (b) Without weight decay ( $\lambda=0$ ), nothing balances the $g_{i}/\epsilon$ step, so the margins rise past $\log(\nicefrac{{1}}{{\epsilon}})$ and continue drifting upward.

Panel of Figure 5 — Figure 5: Standard cross-entropy fails at the precision cliff $c_{d}$ of its floating-point format; fixed cross-entropy removes the cliff across all precisions ( $a+b\bmod 31$ , $\lambda=1$ ; 2 of 5 seeds shown). Top: $\epsilon=10^{-8}$ . Bottom: $\epsilon=10^{-16}$ , where even fp64 standard cross-entropy fails at $c_{64}$ while fixed cross-entropy converges smoothly.

		Test Accuracy (%)									Final Margin $m$
	Variant $\epsilon=$	$10^{-8}$			$10^{-12}$			$10^{-16}$			$10^{-8}$	$10^{-12}$	$10^{-16}$
fp32	Fixed CE	100.0	$\pm$	0.0	99.6	$\pm$	0.7	99.0	$\pm$	2.3	17.9	27.4	36.8
fp32	Standard CE	94.3	$\pm$	5.4	96.4	$\pm$	5.1	96.5	$\pm$	4.3	13.0	15.1	12.7
bf16	Fixed CE	100.0	$\pm$	0.1	99.8	$\pm$	0.2	98.9	$\pm$	2.4	17.9	27.4	36.8
bf16	Standard CE	93.3	$\pm$	7.4	94.9	$\pm$	4.9	94.2	$\pm$	8.3	5.1	5.6	5.8
fp64	Fixed CE	100.0	$\pm$	0.0	99.5	$\pm$	1.1	99.0	$\pm$	2.0	17.9	27.4	36.8

Table 1: Comparison of loss functions on modular addition (

a+b\bmod 31

\lambda=1

5

seeds) across

\epsilon\in\{10^{-8},10^{-12},10^{-16}\}

. Fixed CE reaches the predicted equilibrium margin (

\log(\nicefrac{{1}}{{\epsilon}})\approx 18.4,27.6,36.8

) and

99

–

100\%

test accuracy across all precisions, whereas Standard CE stalls near its precision cliff

c_{d}

(

c_{16}\approx 5.5

in bf16,

c_{32}\approx 16.6

in fp32, and

c_{64}\approx 36.7

in fp64).

Panel of Figure 6 — Figure 6: Results across four grokking tasks at $\epsilon=10^{-12}$ (2 of 3–5 seeds shown): 1-layer Transformers (a, b) and 2-layer MLPs (c, d). Fixed cross-entropy in fp32 (blue) matches fp64 standard cross-entropy (red), while standard cross-entropy in fp32 (orange) breaks at $c_{32}$ .

Panel of Figure 7 — Figure 7: Soft-target distillation when the student is initialized with the teacher’s weights ( $\boldsymbol{\theta}_{\mathrm{S}}^{(0)}=\boldsymbol{\theta}_{\mathrm{T}}$ , $T=4$ , $\epsilon=10^{-8}$ , 2 of 3 seeds shown). Fixed soft-target cross-entropy stays at $\mathrm{KL}=0$ and $100\%$ accuracy across all precisions; standard soft-target cross-entropy amplifies rounding error into parameter drift in bf16 and a brief $\mathrm{KL}$ spike in fp32.

Neural Network Activation Functions

Neural network activation functions shape how models represent nonlinear patterns, learn, remain stable, and map efficiently to hardware.

See all 78 papers

Mar 13, 2026

Ghosts of Softmax: Complex Singularities That Limit Safe Step Sizes in Cross-Entropy

Optimization analyses for cross-entropy training rely on local Taylor models of the loss to predict whether a proposed step will decrease the objective. These surrogates are reliable only inside the Taylor convergence radius of the true loss along the update direction. That radius is set not by real-line curvature alone but by the nearest complex singularity. For cross-entropy, the softmax partition function

F=\sum_j \exp(z_j)

has complex zeros -- ``ghosts of softmax'' -- that induce logarithmic singularities in the loss and cap this radius. To make this geometry usable, we derive closed-form expressions under logit linearization along the proposed update direction. In the binary case, the exact radius is

ρ^*=\sqrt{δ^2+ π^2}/Δ_a

. In the multiclass case, we obtain the lower bound

ρ_a=π/Δ_a

, where

Δ_a=\max_k a_k-\min_k a_k

is the spread of directional logit derivatives

a_k=\nabla z_k\cdot v

. This bound costs one Jacobian-vector product and reveals what makes a step fragile: samples that are both near a decision flip and highly sensitive to the proposed direction tighten the radius. The normalized step size

r=τ/ρ_a

separates safe from dangerous updates. Across six tested architectures and multiple step directions, no model fails for

r<1

, yet collapse appears once

r\ge 1

. Temperature scaling confirms the mechanism: normalizing by

ρ_a

shrinks the onset-threshold spread from standard deviation

0.992

0.164

. A controller that enforces

τ\leρ_a

survives learning-rate spikes up to

10{,} 000\times

in our tests, where gradient clipping still collapses. Together, these results identify a geometric constraint on cross-entropy optimization that operates through Taylor convergence rather than Hessian curvature.

Piyush Sao

Oak Ridge National Laboratory, Oak Ridge, TN, USA

Deep Learning Optimization Convergence Analysis

Ask PDF

Oct 27, 2026

Property-Based Testing for Robust Softmax and Cross-Entropy

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.

Muhammad Maaz

Robust Loss Functions

Ask PDF

Oct 4, 2026

E $^2$ -OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation

On-policy self-distillation (OPSD) provides dense token-level supervision without a second model: one network acts as teacher with the reference solution and as student with only the problem. We identify a specific failure mode of this recipe. During training, student token entropy rises past the teacher's and remains elevated, a pattern we call entropy overshoot. We trace it to both sides of distillation. The reference-conditioned teacher is confident along its answer-directed reasoning path, but this confidence transfers poorly to student-generated prefixes, making its supervision overly tied to answer-specific cues rather than reusable reasoning patterns; meanwhile, the forward KL used by OPSD continually diffuses the student's predictive distribution without pulling it back. We introduce E

^2

-OPSD to address both causes. Exemplar-guided teaching replaces the current answer with a retrieved solved neighboring problem, providing transferable reasoning guidance without revealing the destination and better matching student-reachable states. Entropy-aware distillation uses the student-teacher entropy gap to determine the direction and strength of each token's correction. E

^2

-OPSD improves math reasoning by up to 4.3 points in mean@16 over OPSD, while out-of-domain evaluations show gains over the corresponding base models of up to 4.9 points in mean@16 and 5.5 points in pass@8. Despite these gains, E

^2

-OPSD remains simple, requiring no additional forward passes or networks. Code repo: https://github.com/orionfei/E2-OPSD.

Yifei Liu, Minghao Fang, Xinyu Gu +6

The Chinese University of Hong Kong · Zhejiang University · Shanghai Artificial Intelligence Laboratory +5

On-Policy Self-Distillation Mathematical Reasoning Uncertainty-Aware Knowledge Distillation Knowledge Distillation

Ask PDF

Abstract

Neural Network Activation Functions

Explore similar work

Ghosts of Softmax: Complex Singularities That Limit Safe Step Sizes in Cross-Entropy

Property-Based Testing for Robust Softmax and Cross-Entropy

E2^22-OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation

E $^2$ -OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation