A Fix for Softmax and Cross-Entropy

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.

• The paper identifies numerical pathologies in standard implementations of softmax and cross-entropy that can alter training dynamics, such as Softmax Collapse and Softmax Cancellation.
• The proposed fixed implementations improve computational stability and precision, maintaining accurate gradients and preventing loss instabilities in lower-precision arithmetic by using operations like softplus, expm1, and log1p.
• The solutions solve precision issues in grokking experiments, ensure reliable distillation, and meet margins dependant on the Adam optimizer's epsilon parameter, showing potential to improve handling of precision cliffs in various training scenarios.

The paper identifies two numerical pathologies in standard implementations of softmax, log-softmax, and cross-entropy and proposes algebraic fixes that preserve the exact gradient structure in low-precision arithmetic. The central claim is that these are not merely benign rounding effects: under adaptive optimization, they can qualitatively alter training dynamics, generate apparent loss instabilities, impair grokking, and induce parameter drift during distillation. The proposed implementations operate directly in log space and use softplus, expm1, and log1p to avoid numerically ill-conditioned additions and subtractions (2610.31642).

Numerical failure modes in standard cross-entropy

For logits $z \in \mathbb{R}^K$, standard implementations subtract the maximum logit before evaluating softmax or log-softmax. This prevents exponent overflow, but it introduces a precision problem when the largest logit is substantially larger than the others. Let $t$ denote the target class and define the target margin

$m = z_t - \log \sum_{j \ne t} e^{z_j}.$

The exact target probability is

$q_t = \frac{1}{1+e^{-m}},$

and the exact target gradient for hard-target cross-entropy is

$g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$

The standard computation forms $1+e^{-m}$ and then subtracts one from the resulting probability. Once $e^{-m}$ becomes smaller than the spacing of floating-point numbers near one, the denominator rounds to exactly $1$. The target probability therefore becomes exactly one, and the target gradient becomes exactly zero. The non-target gradients, however, are computed from exponentials that remain representable near zero and therefore remain positive. The resulting gradient no longer satisfies the exact conservation law

$\sum_k g_k = 0.$

The paper calls this failure Softmax Collapse. Its precision thresholds are approximately $m \approx 5.5$ in bf16, $t$0 in fp32, and $t$1 in fp64. These thresholds are determined by the precision available near one, not by exponent underflow. Consequently, increasing the floating-point exponent range does not address the problem; the relevant issue is cancellation and rounding in the computation of a quantity close to one.

Figure 1: Softmax Collapse zeroes the target gradient beyond precision-dependent logit margins, while Softmax Cancellation produces a non-zero gradient floor when student and teacher logits coincide.

The proposed hard-target formulation excludes the target logit from the log-sum-exp calculation. It computes the margin from the non-target logits and evaluates the loss as

$t$2

Because the target coordinate is represented by an exact zero rather than by a quantity obtained through subtraction and exponentiation, autodifferentiation produces the target gradient as the negative sum of the non-target gradients. This preserves the zero-sum property even when the denominator $t$3 rounds to one. The target gradient remains accurate down to exponent underflow, which occurs much later than the precision cliff.

This distinction is practically important. For logits $t$4 in fp32, standard JAX and Optax routines return a zero target log-probability, zero loss, and zero target gradient, while retaining positive non-target gradients of approximately $t$5. The corrected implementation returns a target loss of approximately $t$6 and a target gradient of $t$7, exactly balancing the two non-target gradients. The forward loss and the backward gradient therefore remain mutually consistent.

Softmax Cancellation in soft-target objectives

The second pathology occurs when cross-entropy or KL-based distillation compares student logits $t$8 with teacher logits $t$9. The conventional implementation computes student and teacher probabilities independently,

$m = z_t - \log \sum_{j \ne t} e^{z_j}.$0

and forms the gradient as $m = z_t - \log \sum_{j \ne t} e^{z_j}.$1. When $m = z_t - \log \sum_{j \ne t} e^{z_j}.$2 is close to $m = z_t - \log \sum_{j \ne t} e^{z_j}.$3, both probabilities are nearly identical. Subtracting them after they have been rounded discards the information carried by the small logit difference. In addition, the forward softmax and backward log-softmax may use different operation sequences, so the two probability vectors can differ numerically even when the logits are identical.

This is Softmax Cancellation. Its most direct violation is at $m = z_t - \log \sum_{j \ne t} e^{z_j}.$4: the exact gradient is zero, but the standard implementation can produce a non-zero residual. Under Adam, a residual gradient larger than the optimizer’s $m = z_t - \log \sum_{j \ne t} e^{z_j}.$5 parameter can be normalized into a finite update, converting an arithmetic error into parameter motion.

The paper avoids probability-space subtraction by expressing the student-to-teacher ratio in log space. It centers the logit difference under the teacher distribution and evaluates the resulting probability differences with expm1 and log1p. These functions retain relative accuracy for small arguments. A final projection subtracts the residual mean gradient so that the computed gradient sums to zero exactly in floating-point arithmetic.

The numerical example is consequential rather than cosmetic. For teacher logits $m = z_t - \log \sum_{j \ne t} e^{z_j}.$6 and student logits differing by only $m = z_t - \log \sum_{j \ne t} e^{z_j}.$7, the standard fp32 gradient is

$m = z_t - \log \sum_{j \ne t} e^{z_j}.$8

whereas the corrected result is

$m = z_t - \log \sum_{j \ne t} e^{z_j}.$9

matching the exact fp64 gradient. The standard computation shrinks one component by roughly a factor of three, reverses the sign of another, and fails to maintain a zero sum. The corrected computation restores the local geometry of the distillation objective near its optimum (2610.31642).

Why Softmax Collapse produces Slingshot dynamics

The paper provides a mechanistic explanation for the Slingshot phenomenon observed in grokking experiments. After memorization, target margins continue increasing even though training loss is already close to zero. With standard cross-entropy, examples crossing the precision cliff lose their target gradients while retaining non-target gradients. Once most examples have crossed the cliff, parameter gradients become very small and Adam’s second-moment estimates decay.

Weight decay and residual updates can subsequently reduce some margins below the cliff. The target gradients then reappear abruptly, but Adam’s second moments have already decayed. The restored gradients are therefore divided by a small adaptive denominator, generating a large parameter update and a sudden loss spike. The spike resets the second moment, after which the process repeats.

The fixed implementation removes the discontinuity. Target gradients decay smoothly as $q_t = \frac{1}{1+e^{-m}},$0, so Adam’s first- and second-moment estimates evolve continuously with the loss gradient. The result is the disappearance of recurrent Slingshot spikes.

The experiments distinguish two effects that are often conflated. With weight decay on $q_t = \frac{1}{1+e^{-m}},$1, standard cross-entropy reaches the fp32 cliff and obtains $q_t = \frac{1}{1+e^{-m}},$2 test accuracy, whereas fixed cross-entropy reaches $q_t = \frac{1}{1+e^{-m}},$3. Without weight decay on $q_t = \frac{1}{1+e^{-m}},$4, standard cross-entropy reaches approximately $q_t = \frac{1}{1+e^{-m}},$5 test accuracy before a spike drops it to $q_t = \frac{1}{1+e^{-m}},$6. Fixed cross-entropy continues smoothly to $q_t = \frac{1}{1+e^{-m}},$7. On $q_t = \frac{1}{1+e^{-m}},$8, the analogous drop is from approximately $q_t = \frac{1}{1+e^{-m}},$9 to $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$0.

Figure 2: Standard cross-entropy exhibits margin-triggered loss spikes and accuracy drops, whereas fixed cross-entropy crosses the floating-point precision cliff smoothly.

The paper also argues that the apparent benefits of Slingshot events in some prior grokking settings are contingent on architectural and optimization details. When LayerNorm precedes the output layer and no weight decay is used, increases in pre-LayerNorm weight norms can reduce the effective learning rate without changing the logits. Under those conditions, a Slingshot event may act as an accidental learning-rate decay. When weight decay already controls the norm, however, the same events slow optimization. When neither mechanism controls the norm, an individual spike can be destructive. This interpretation makes the phenomenon optimizer- and architecture-dependent rather than an intrinsic property of grokking.

Margin equilibria under AdamW

Removing the precision cliff does not imply unlimited margin growth when weight decay is present. The paper derives a separate equilibrium determined by Adam’s $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$1 hyperparameter and the decay coefficient.

At a stationary point, the Adam moments satisfy approximately $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$2 and $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$3. The AdamW update then has the stationary condition

$g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$4

where $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$5 is the weight-decay coefficient. The learning rate and Adam moment decay rates cancel from this local condition. The same stationary point therefore applies to RMSProp- and LaProp-like normalized updates under the assumptions used in the derivation.

For memorized examples, gradient magnitudes decrease approximately exponentially with margin. When $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$6, Adam normalizes the gradient to a nearly constant step. Once $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$7, the normalized step decreases proportionally to $g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$8. Balancing that decaying update against weight decay yields the approximate margin law

$g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}.$9

This law explains the observed plateau near $1+e^{-m}$0 for $1+e^{-m}$1, $1+e^{-m}$2 for $1+e^{-m}$3, and $1+e^{-m}$4 for $1+e^{-m}$5. The plateau is therefore not a floating-point artifact. It is an optimizer equilibrium whose location can nevertheless place training directly beyond the precision cliff of the selected datatype.

Figure 3: Fixed cross-entropy produces a smooth Adam-normalized step whose decay begins near $1+e^{-m}$6, while standard cross-entropy introduces abrupt precision-dependent discontinuities.

Experiments with fixed cross-entropy and LaProp support the analytical prediction over several orders of magnitude in $1+e^{-m}$7. With weight decay, margins level off near the predicted $1+e^{-m}$8 values. Without weight decay, the normalized gradient remains approximately unit magnitude as the true gradient shrinks, and margins continue growing until the gradients themselves underflow. This result identifies $1+e^{-m}$9 as a direct control on post-memorization logit growth, rather than merely a numerical safeguard.

Results across precision formats and benchmarks

The principal controlled experiment uses a width-64, one-layer ReLU Transformer trained on $e^{-m}$0 with a 50% training split, full-batch AdamW, and weight decay $e^{-m}$1. The comparison varies datatype, loss implementation, and $e^{-m}$2.

At $e^{-m}$3, standard cross-entropy stalls near the bf16 cliff at $e^{-m}$4 and repeatedly spikes near the fp32 cliff at $e^{-m}$5. Fixed cross-entropy crosses both thresholds and reaches $e^{-m}$6 test accuracy in bf16, fp32, and fp64. At $e^{-m}$7, the predicted equilibrium is $e^{-m}$8, beyond the fp64 cliff at approximately $e^{-m}$9. Standard fp64 cross-entropy consequently fails even in double precision, while fixed cross-entropy remains stable.

Precision	Loss	$1$0	$1$1	$1$2
fp32	Fixed CE	$1$3	$1$4	$1$5
fp32	Standard CE	$1$6	$1$7	$1$8
bf16	Fixed CE	$1$9	$\sum_k g_k = 0.$0	$\sum_k g_k = 0.$1
bf16	Standard CE	$\sum_k g_k = 0.$2	$\sum_k g_k = 0.$3	$\sum_k g_k = 0.$4
fp64	Fixed CE	$\sum_k g_k = 0.$5	$\sum_k g_k = 0.$6	$\sum_k g_k = 0.$7
fp64	Standard CE	$\sum_k g_k = 0.$8	$\sum_k g_k = 0.$9	$m \approx 5.5$0

The final margins further separate the methods. Fixed cross-entropy reaches approximately $m \approx 5.5$1, $m \approx 5.5$2, and $m \approx 5.5$3 for the three $m \approx 5.5$4 values, in agreement with the equilibrium law. Standard cross-entropy instead remains near the datatype-specific cliffs in bf16 and fp32, and shows degraded margins in fp64 at the smallest $m \approx 5.5$5.

Figure 4: Across bf16, fp32, and fp64, fixed cross-entropy reaches the optimizer-predicted margins, while standard cross-entropy fails when those margins cross the datatype’s precision cliff.

Two alternative interventions perform less consistently. Mean-gradient projection restores a zero-sum gradient after the target gradient has already been lost, but it assigns only approximately $m \approx 5.5$6 of the residual non-target gradient to the target coordinate. In the reported experiments, it stalls near margin $m \approx 5.5$7 and falls below $m \approx 5.5$8 test accuracy for $m \approx 5.5$9. StableMax avoids exponential softmax behavior by replacing the exponential with a piecewise rational function, but produces margins between roughly $t$00 and $t$01 and lower test accuracy than fixed cross-entropy. The comparison supports the paper’s preference for correcting the numerically unstable evaluation of the original objective rather than changing the objective or projecting an already corrupted gradient.

The same pattern appears on four grokking benchmarks: two one-layer Transformer tasks, $t$02 and $t$03, and two two-layer MLP tasks, $t$04 and $t$05. At $t$06, fp32 standard cross-entropy repeatedly breaks near the fp32 cliff, while fp32 fixed cross-entropy tracks fp64 standard cross-entropy and reaches $t$07 test accuracy. This correspondence suggests that the mechanism is not specific to one architecture or modular task, although the experiments remain concentrated on small synthetic grokking benchmarks.

Distillation stability

The Softmax Cancellation experiments initialize a student with the weights of a teacher that achieves $t$08 test accuracy on $t$09. The student is trained for 2,000 Adam steps using temperature $t$10, with no weight decay and $t$11. The temperature keeps scaled margins below the Softmax Collapse regime, isolating cancellation near the teacher solution.

At initialization, student and teacher logits are identical, so the exact distillation gradient is zero. Under standard soft-target cross-entropy, the residual numerical gradient is approximately $t$12 in fp64, $t$13 in fp32, and $t$14 in bf16. The resulting behavior follows the relationship between the residual and Adam’s $t$15: fp64 remains effectively stationary, fp32 exhibits a brief KL spike, and bf16 drifts away from the teacher and loses accuracy. The corrected implementation returns exactly zero at equal logits, maintaining $t$16 and $t$17 test accuracy across all three precisions.

Figure 5: In distillation initialized at the teacher solution, standard soft-target cross-entropy amplifies precision-dependent residual gradients, whereas the corrected objective remains stationary.

This result is particularly relevant for self-distillation, teacher-student initialization, and KL regularization, where the desired optimum can involve extremely small student-teacher logit differences. The paper does not claim that cancellation dominates ordinary supervised training; rather, it isolates a regime in which the objective’s gradient is expected to vanish and demonstrates that standard implementations can prevent that invariance from being realized numerically.

Limitations and open questions

The empirical evaluation is concentrated on synthetic modular-arithmetic grokking tasks, small Transformers and MLPs, and carefully selected high-confidence or near-identical-logit regimes. These settings are appropriate for exposing the failure modes, but they do not establish the frequency or optimization impact of the bugs in large-scale language, vision, or multimodal training. The paper explicitly notes that the extreme margins and very small student-teacher gaps used in the experiments may be uncommon in standard training.

The implementation also introduces additional operations, masking, and in the soft-target case a custom backward expression. The paper reports no meaningful slowdown, but the evidence provided is not a broad systems benchmark across accelerator architectures, kernel fusions, batch shapes, and compiler stacks. The effects of nondeterministic reductions and mixed-precision accumulation are likewise not characterized.

The equilibrium law $t$18 depends on assumptions about stationary Adam moments, exponential gradient decay, and a stable balance between normalized gradients and weight decay. It accurately describes the reported experiments, but its applicability to momentum transients, minibatch noise, parameter-wise heterogeneous gradients, or optimizers with materially different preconditioning remains open. Finally, the paper establishes that the corrected gradient is numerically faithful; it does not determine whether exact preservation of tiny gradients is always desirable when those gradients interact with stochastic regularization, gradient clipping, or noisy labels.

Conclusion

The paper shows that conventional numerically stable softmax implementations remain unstable in two distinct regimes. Softmax Collapse destroys the dominant-class gradient after a precision-dependent margin threshold, while Softmax Cancellation destroys information when student and teacher probabilities are nearly equal. Both failures violate exact gradient properties and can be amplified by adaptive optimizers.

The proposed log-space formulations preserve zero-sum gradients, maintain non-zero target gradients beyond the precision cliff, and compute student-teacher differences accurately near equality. Across grokking and distillation experiments, they eliminate Slingshot spikes, improve low-precision robustness, prevent teacher-student drift, and expose a margin equilibrium controlled by Adam’s $t$19 and weight decay. The remaining question is how often these numerically extreme regimes occur in large-scale workloads and whether the same corrections produce measurable benefits outside the controlled settings studied here (2610.31642).