The paper "A Fix for Softmax and Cross-Entropy" identifies two fundamental numerical instabilities in standard machine learning implementations of softmax and cross-entropy—*Softmax Collapse* and *Softmax Cancellation*—and proposes robust, numerically stable alternatives.

1. Softmax Collapse

Softmax Collapse occurs during high-confidence training where the target logit $z_t$ is significantly larger than the non-target logits, resulting in a large logit margin $m = z_t - \log \sum_{j \neq t} e^{z_j}$. In standard implementations, $1 + e^{-m}$ is computed by adding $e^{-m}$ to $1.0$. Once $m$ exceeds a precision cliff ($c_d \approx 16.6$ in fp32), $e^{-m}$ is rounded to zero, causing the loss and the target gradient $\tilde{g}_t$ to become exactly zero. Crucially, non-target gradients remain non-zero, causing the total gradient to not sum to zero ($\sum_k \tilde{g}_k > 0$). This asymmetry drives the "Slingshot" phenomenon in Adam/AdamW optimizers, where weight decay and gradient updates compete until a sudden spike in loss occurs.

Methodology: The authors fix this by avoiding the direct addition of $1$ in the denominator. Instead, they define the margin $m$ excluding the target logit and utilize $\operatorname{softplus}(-m) = \log(1 + e^{-m})$. By computing $\mathcal{L}(\mathbf{z}, t) = \operatorname{softplus}(-m)$, the target gradient is derived as $g_t = -e^{-m} / (1 + e^{-m})$, which remains accurate even when $m > c_d$.

2. Softmax Cancellation

Softmax Cancellation occurs in knowledge distillation or KL regularization where a student model's logits $\mathbf{z}_S$ are forced to match teacher logits $\mathbf{z}_T$. Standard implementations compute $\mathbf{p} = \operatorname{softmax}(\mathbf{z}_T)$ and $\mathbf{q} = \operatorname{softmax}(\mathbf{z}_S)$ separately. As $\mathbf{z}_S \to \mathbf{z}_T$, subtracting these probabilities leads to precision loss, leaving a residual noise floor in the gradient even when the logits are equal.

Methodology: The authors reformulate the gradient calculation in log-space using $\operatorname{expm1}(x) = e^x - 1$ and $\operatorname{log1p}(x) = \log(1 + x)$. Given the logit difference $\boldsymbol{\delta} = \mathbf{z}_S - \mathbf{z}_T$, the difference $\mathbf{u} = \mathbf{q} - \mathbf{p}$ is computed as:
$$d_k = \delta_k - \sum_i p_i \delta_i$$
$$s_k = d_k - \operatorname{log1p}\left(\sum_i p_i \operatorname{expm1}(d_i)\right)$$
$$u_k = p_k \operatorname{expm1}(s_k)$$
This approach ensures the gradient $\mathbf{g} = \mathbf{u} - \mathbf{q} \sum_i u_i$ is identically zero when $\mathbf{z}_S = \mathbf{z}_T$, preventing the drift observed in standard implementations.

3. Margin Equilibria and Slingshot

The authors demonstrate that when using AdamW, the logit margin does not grow indefinitely but reaches an equilibrium $m^\star(\epsilon) \approx \log(1/\epsilon)$ where the normalized gradient step balances weight decay. In standard implementations, the precision cliff often forces an earlier, unstable termination of this growth. The proposed fix allows models to reach this natural equilibrium smoothly. Experimental results across bf16, fp32, and fp64 formats confirm that these fixes eliminate Slingshot spikes, accelerate grokking, and improve test accuracy by allowing the model to cross precision boundaries without destructive floating-point rounding errors.