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

A Fix for Softmax and Cross-Entropy
Oct 22, 2026Jonathan T. Barron, Dan Gnanapragasam, Michael TschannenNeural Network Activation Functions
Google DeepMind
A Boundary-Layer Mechanism for One-Third Scaling in Online Softmax Classification
May 21, 2026Marcel Kühn, Yoon Thelge, Bernd RosenowNeural Network Training Dynamics Multiclass Classification
Institute for Theoretical Physics, Leipzig University · ScaDS.AI Dresden/Leipzig
Avoiding unsafe sets when training with Langevin Dynamics
Jul 8, 2026Adam M. ObermanLangevin Dynamics Stochastic Optimization Convergence
Department of Mathematics and Statistics, McGill University, Montreal, QC, Canada; Mila, Quebec AI Institute; and LawZero

Oct 22, 2026J/K move · Enter open · S save

A Fix for Softmax and Cross-Entropy

Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen

Google DeepMind

Neural Network Activation Functions

Open paper AskPDF arXiv

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.

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