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

Piyush Sao

cs.LGMar 13, 2026

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

Authors: Piyush Sao

Organizations: Oak Ridge National Laboratory, Oak Ridge, TN, USA

Abstract

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$ to $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.

Figures & tables

Figure 1: We train six small architectures to convergence, then take one step at varying ratios

r=\tau/\rho_{a}

, where

\tau

is the step distance and

\rho_{a}

is our stability bound (defined in the text). Every network keeps its accuracy for

r<1

; once

r

crosses 1, collapse appears across all six architectures.

Figure 2: Taylor approximations

T_{n}(x)

f(x)=1/(x+a)

around

x_{0}=0

. Inside the convergence radius

R=a

(green), all orders approximate

f

well. Beyond

R

(red), higher-order approximations diverge faster, not slower.

Figure 3: Binary cross-entropy’s log-partition

\log(1+e^{x})

has convergence radius

\rho=\pi

set by the complex zero at

i\pi

(Euler:

1+e^{i\pi}=0

). Shown: derivative

\sigma(x)

on log scale. Green: Taylor converges (

|x|<\pi

). Red: diverges (

|x|>\pi

Figure 4: Half-plane geometry behind Lemma 4.5 . If the phases of

V_{k}=w_{k}e^{a_{k}t}

lie in an open arc of length

<\pi

, then all vectors remain in a common open half-plane, so no positive combination can cancel to zero.

Figure 5: Complex-

t

consequence of Theorem 4.6 . The half-plane obstruction excludes zeros from the strip

|\mathrm{Im}(t)|<\pi/\Delta_{a}

. Since the disk

|t|<\pi/\Delta_{a}

lies inside that strip, the minimal zero modulus satisfies

\rho^{*}\geq\pi/\Delta_{a}

Figure 6: Learning rate spike test (5 seeds, median with interquartile range (IQR) bands). Each row uses a different base LR. Left : training loss. Center : test accuracy. Right : step size

\tau

vs radius

\rho_{a}

(green = safe zone). Row 1: spike stays within

\rho_{a}

, Adam recovers. Row 2:

\tau

exceeds

\rho_{a}

{\sim}10\times

, accuracy drops. Row 3:

\tau

exceeds

\rho_{a}

{\sim}100\times

, learning is erased. The

\rho_{a}

-controller (gold) survives all spikes.

Figure 7: Phase transition across architectures. Six architectures trained to convergence, then one gradient step at varying

r

. (A) Test accuracy retained. Linear model transitions at

r\approx 1

. Other architectures retain confidence-margin slack. (B) Loss inflation (post-step loss / pre-step loss).

Figure 8: Random-direction sweeps. Both loss ratio and flip fraction transition near

r\approx 1

across all tested phases, supporting

r

as direction-independent in these settings.

Figure 9: Temperature-scaling fingerprint test. One-step loss inflation (post-step loss / pre-step loss) is measured across temperatures

T\in\{0.25,0.5,1,2,4,8,16,64\}

. A : plotted against raw step size

\tau

, collapse thresholds vary substantially across

T

(collapse std

=0.992

). B : plotted against the theory-normalized coordinate

r_{T}=\tau\Delta_{a}/(\pi T)=\tau/\rho_{a}(T)

, curves align much more tightly (collapse std

=0.164

Figure 10: Controller across architectures and spike severities (5 seeds, median with IQR bands). Each column increases the spike magnitude from

10\times

10\,000\times

. Red: plain Adam. Teal: Adam with gradient clipping (

\|g\|\leq 1

). Gold:

\rho_{a}

-controller. Gradient clipping helps partially at moderate spikes but still collapses at high severity. The

\rho_{a}

-controller survives all conditions.

Arch	Plain	Clip	$\rho_{a}$ -ctrl
Transformer	31.8%	38.7%	95.3%
MLP+LN	44.6%	51.5%	93.9%
CNN+BN	13.4%	10.3%	71.6%

Table 1: Median final test accuracy at

10\,000\times

spike. A seed is divergent if final loss exceeds 10.

Figure 11: Tiny transformer on Digits: adaptive vs fixed learning rates (5 seeds, median with IQR bands). Left: training loss. Center: test accuracy. Right: effective LR. The adaptive controllers (all-radii, output-only) vary their step size by orders of magnitude throughout training. Fixed LRs either converge slowly (1

\times

) or diverge (16

\times

). The most conservative strategy (all-radii, using

\min(\rho_{\text{out}},\rho_{\text{attn}},\rho_{\text{ffn}})

) reaches the lowest loss and highest accuracy.

Phase	FFN	Out	Attn	med. ffn	med. out	med. attn
Early	94/100	6/100	0/100	0.0164	0.0307	0.645
Middle	48/100	52/100	0/100	0.0111	0.0107	0.203
Late	8/105	97/105	0/105	0.0131	0.00793	0.168

Table 2: Which radius binds in the all-radii transformer run ( Figure 11 ; 5 seeds, 61 logged evaluations per seed). Counts report how often each term attained

\rho_{\text{net}}=\min(\rho_{\text{out}},\rho_{\text{attn}},\rho_{\text{ffn}})

within each training phase. Median radii are pooled over all logged points in that phase.

Figure 12: Natural instability detection: ResNet-18 on CIFAR-10 with SGD+momentum at four fixed learning rates and four

\rho_{a}

-controller targets (

r{=}0.5,1,2,4

; 5 seeds, median with IQR bands). Left : test loss. Center : test accuracy. Right : maximum

r=\tau/\rho_{a}

per epoch with the

r{=}1

boundary (dashed). Step size

\tau

includes the momentum buffer. The controller at

r{=}1

reaches the highest accuracy (

85.3\%

);

r{=}2

remains stable;

r{=}4

degrades with high variance. Even

\eta{=}0.005

occasionally exceeds

r{=}1

due to momentum.

Method	max $r$	$r{>}1$	Acc (med.)	IQR
$\eta{=}0.005$	1.11	1	80.4%	80.0–81.9%
$\eta{=}0.01$	1.23	13	82.6%	82.6–83.4%
$\eta{=}0.1$	5.87	58	78.5%	78.1–81.7%
$\eta{=}0.5$	29.34	70	57.5%	48.1–64.8%
$\rho_{a}$ -ctrl ( $r{=}0.5$ )	0.50	0	80.4%	80.2–80.5%
$\rho_{a}$ -ctrl ( $r{=}1$ )	1.00	0 †	85.3%	85.1–85.4%

Table 3: ResNet-18 CIFAR-10 natural instability summary (5 seeds, 10 epochs, momentum-corrected

\tau

). The controller at

r{=}1

is optimal;

r{=}2

retains accuracy from confidence-margin slack;

r{=}4

degrades. Even

\eta{=}0.005

exceeds

r{=}1

due to momentum.

Figure 13: Hessian curvature

\sigma(1-\sigma)\sim e^{-\delta}

vanishes exponentially while the radius bound

\pi/\sqrt{\delta^{2}+\pi^{2}}

decays only algebraically. At

\delta=10

, the mismatch exceeds

4000\times

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 205 papers

Oct 22, 2026

A Fix for Softmax and Cross-Entropy

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.

Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen

Google DeepMind

Neural Network Activation Functions

Ask PDF

May 21, 2026

A Boundary-Layer Mechanism for One-Third Scaling in Online Softmax Classification

Hard-label classification is usually trained with smooth surrogate losses such as softmax cross-entropy. We isolate a mechanism by which this mismatch between smooth surrogate and discrete labels produces power-law learning curves for a one-layer softmax student trained by online SGD on Gaussian inputs with argmax teacher labels. In the thermodynamic limit, the dynamics reduce to a growing student-teacher alignment and a residual student variance sustained by SGD noise. At late times, only shrinking layers of examples around teacher decision boundaries remain active. This yields an

α^{-1/3}

law in training time

α

not only for the test loss but also for the generalization error, one minus test accuracy. The error law is an online-noise effect. The loss law persists without gradient noise. Learning-rate annealing separates the two exponents, improving the error exponent from 1/3 toward 1/2 while slowing the loss. Simulations support the predictions. Correlated Gaussian inputs and whitened pretrained features change transients, not the late-time decay. Experiments with deep neural networks can show compatible error scaling. The mechanism is asymptotic and complementary to spectral explanations and data-statistics explanations of neural scaling laws.

Marcel Kühn, Yoon Thelge, Bernd Rosenow

Institute for Theoretical Physics, Leipzig University · ScaDS.AI Dresden/Leipzig

Neural Network Training Dynamics Multiclass Classification Online Learning Statistical Physics of Learning Neural Scaling Laws Teacher-Student Learning+1

Ask PDF

Jul 8, 2026

Avoiding unsafe sets when training with Langevin Dynamics

Training a model with noisy gradient descent can be idealized as overdamped Langevin dynamics, and a natural safety question is to bound the probability

ν_t(\mathcal{A}_H) = \mathbb{P}(Q_t \in \mathcal{A}_H)

that the trajectory lies in a designated failure region

\mathcal{A}_H

. We study this for a smooth, strongly convex loss in

d

dimensions, with

\mathcal{A}_H

separated from the minimizer by an energy gap. At the end of training, the equilibrium mass

π(\mathcal{A}_H)

is exponentially small in

d

, with a complementary energy-barrier rate when the noise is small. Along the trajectory, a shape-free bound

ν_t(\mathcal{A}_H) \le π(\mathcal{A}_H)(1 + \sqrt{χ_0^2/π(\mathcal{A}_H)}\,e^{-mt})

shows the in-set probability relaxes to (twice) the static value after a burn-in of order

d

, using only the global spectral gap

m

. A worked Ornstein-Uhlenbeck example shows this burn-in is necessary: an angular slice of the equilibrium shell can transiently swell by a factor exponential in

d

, though its equilibrium mass is tiny. To rule this out we introduce a local relaxation rate, defined through the spectral measure of the region's centered indicator rather than a Dirichlet-form Rayleigh quotient. For geometrically isolated regions this rate exceeds the global one, shrinking the burn-in, and with a maximum-principle ceiling it caps the trajectory probability uniformly in time. Strong convexity sets how fast training relaxes, but the shape of the unsafe set decides whether the trajectory bulges through it on the way to equilibrium.

Adam M. Oberman

Department of Mathematics and Statistics, McGill University, Montreal, QC, Canada; Mila, Quebec AI Institute; and LawZero

Langevin Dynamics Stochastic Optimization Convergence

Ask PDF

Abstract

Deep Learning Optimization

Explore similar work

A Fix for Softmax and Cross-Entropy

A Boundary-Layer Mechanism for One-Third Scaling in Online Softmax Classification

Avoiding unsafe sets when training with Langevin Dynamics