Softmax and Cross-Entropy Fixes

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∈RKz \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 tt denote the target class and define the target margin m=zt−log⁡∑j≠tezj.m = z_t - \log \sum_{j \ne t} e^{z_j}. The exact target probability is qt=11+e−m,q_t = \frac{1}{1+e^{-m}}, and the exact target gradient for hard-target cross-entropy is gt=qt−1=−e−m1+e−m.g_t = q_t - 1 = -\frac{e^{-m}}{1+e^{-m}}. The standard computation forms 1+e−m1+e^{-m} and then subtracts one from the resulting probability. Once e−me^{-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 ∑kgk=0.\sum_k g_k = 0. The paper calls …