Daily AI Paper

This paper is about a surprisingly practical but subtle problem: standard softmax, log-softmax, and cross-entropy are not just numerically imperfect in the usual overflow sense, but can produce qualitatively wrong gradients in common training regimes. That matters because these functions sit at the heart of classification, distillation, and many regularized training objectives. If they behave incorrectly at high confidence or when two distributions are nearly identical, then the model can get pushed in the wrong direction even when the math says it should be at or near equilibrium. To see why this is nontrivial, it helps to recall the standard story. Softmax turns logits into probabilities, and cross-entropy compares those probabilities to a target distribution. In exact arithmetic, the gradient of cross-entropy with respect to the logits is simply predicted probabilities minus target probabilities, and those gradients sum to zero. In practice, though, deep learning frameworks stabilize softmax by subtracting the maximum logit before exponentiating. That prevents overflow, but it does not solve everything. This paper identifies two remaining bugs. The first bug is called Softmax Collapse. It happens when the model is extremely confident in the correct class, so the target logit is much larger than all others. In that regime, the non-target exponentials become tiny. The issue is that floating-point numbers have much finer resolution near zero than near one, and standard implementations form the denominator by adding those tiny terms to 1. Once the margin is large enough, the tiny contribution gets rounded away and the denominator becomes exactly 1. That makes the target probability exactly 1, the target loss exactly 0, and the target gradient exactly 0, even though mathematically it should still be a small negative number. Meanwhile, the non-target gradients are still computed as small positive values. So the gradient vector no longer sums to zero. That asymmetry is not just a cosmetic error. When using adaptive optimizers like Adam or AdamW, which normalize updates based on gradient magnitudes, the mismatch can create a feedback loop. The paper connects this directly to the “Slingshot” behavior observed in grokking experiments: training loss sits near zero, then suddenly spikes, then falls again, often alongside jumps in test accuracy. The authors argue that this is a consequence of the collapsed target gradient interacting with the optimizer. The second bug is called Softmax Cancellation. This one shows up in knowledge distillation and KL-style regularization, where a student network is trained to match a teacher network. In the ideal case, if student logits equal teacher logits, the gradient should be exactly zero. But standard code usually computes the student probabilities and teacher probabilities separately, then subtracts them. When the two distributions are very close, that subtraction loses precision. Even worse, the forward path and backward path may use slightly different numerical formulas, so the two sides do not exactly match even when the logits are equal. The result is a spurious non-zero gradient that can keep nudging the student away from the teacher instead of letting it settle. The core idea of the paper is refreshingly direct: fix the implementations of softmax, log-softmax, and cross-entropy themselves, rather than patching around the symptoms with surrogate losses or higher precision everywhere. The authors build formulas that avoid the dangerous subtraction patterns. For the first bug, the key is to avoid adding 1 to a tiny exponential in a way that erases it. Instead of letting the target term vanish inside a denominator, they preserve the small exponential in the numerator and compute the expression so that the target gradient remains accurate all the way down to exponent underflow. In other words, the target entry no longer collapses to zero just because the probability is extremely close to one. For softmax and log-softmax, the paper makes a related change. The largest logit is still shifted to zero, but the implementation computes the log-sum-exp margin without including that zero term as an explicit 1 in the sum. That avoids the precision cliff. The same logic applies to hard-target cross-entropy: because the target slot is replaced by a constant zero rather than by a subtraction that can cancel, automatic differentiation can recover the correct target gradient as the negative sum of the other entries. The authors also provide a JAX implementation that preserves the exact zero-sum gradient structure. For the second bug, the solution is to compute the student-teacher difference directly in log space. Rather than forming two probability vectors and subtracting them, the loss function receives the teacher logits and uses stable primitives like expm1 and log1p to compute q minus p without throwing away the tiny differences. The paper emphasizes that this keeps the gradient accurate as the student approaches the teacher, and makes the gradient exactly zero when the logits are equal. They also explicitly remove any remaining floating-point drift so that the gradient sums to zero as it should. The experimental story in the paper is tied closely to these numerical pathologies. One set of experiments examines the Softmax Collapse phenomenon in settings where a model has already memorized the training data and the target probability is near one. The authors show that standard implementations can zero out the target gradient once the margin crosses the precision cliff, while their fixed implementation continues to produce the small but nonzero gradient expected from the math. This directly addresses the Slingshot behavior seen in grokking with AdamW, where loss spikes recur after memorization. The paper argues that the collapse in the target gradient is the mechanism behind those spikes. They also report that their fix speeds up grokking when weight decay is used, and prevents accuracy collapse when weight decay is absent. The main point here is not that the optimizer or the training setup is fundamentally broken, but that the numerical implementation of the loss was silently changing the effective dynamics once the model got very confident. On the distillation side, the paper shows that standard soft-target cross-entropy can produce gradients that are too large, slightly wrong in sign, or not exactly zero when student and teacher match. Their corrected version matches the exact fp64 gradient much more closely, even when the student logits are only a small fp32 perturbation of the teacher logits. That is important because distillation depends precisely on learning from small differences. If those small differences are numerically distorted, the student can drift instead of converging. The paper also draws an interesting conclusion about what happens after memorization. Once the model has driven the target probability close to one, logit growth does not continue without bound in the presence of weight decay. Instead, the authors observe that growth is bounded by an equilibrium between weight decay and Adam’s epsilon parameter. That is a useful reminder that in modern optimization, tiny numerical constants can shape the long-run behavior of training, especially when gradients become very small. As for limitations, the paper is focused on floating-point behavior in standard deep learning settings, so the fixes matter most when logits become large in magnitude or when two distributions are nearly identical. The paper does not claim that every training instability comes from these bugs. Rather, it isolates a specific and widespread source of error in the loss computations themselves. Another practical point is that the fixes are implementation-level changes, so their benefit depends on using these corrected versions rather than the default library routines. The broader lesson is that numerical stability is not only about avoiding overflow or NaNs. Even when a function returns finite values, the exact way it is computed can change the gradient qualitatively. Here, that difference shows up in two places that are central to modern training: very confident classification and teacher-student matching. The paper’s contribution is to make these operations behave like the mathematics says they should, even in the regimes where floating-point arithmetic is least forgiving.