Why Adaptive Optimizers Underestimate Rare Tokens

Sangsidhya Kar

cs.LGSep 29, 2026

Why Adaptive Optimizers Underestimate Rare Tokens

Authors: Sangsidhya Kar

Organizations: Department of Statistics Presidency University, Kolkata

Abstract

In the softmax output layer, a rare token receives a small positive logit gradient on most steps and a much larger negative gradient on the few steps when it is the target. SGD simply adds these contributions. Coordinate-wise adaptive methods such as Adam, RMSProp, and sign descent instead divide each update by a running estimate of its magnitude, and that estimate is largest immediately after the token appears. This imbalance has two effects. At the level of the whole output layer, we characterize which optimizers preserve the mean output embedding: every method whose update is linear in past gradients does, as do Kronecker-factored and orthogonalized methods such as Shampoo and Muon. Adam, Adafactor, Lion, and sign descent do not, and for these methods we obtain an exact step-by-step expression for the change. At the level of an individual rare token, the same normalization shifts the training fixed point. In the unigram model, sign descent lowers the logit of every token that occurs in fewer than half of the minibatches at a constant expected rate. For RMSProp with periodic arrivals, we can solve the fixed point in closed form: if a token is absent for at least two consecutive minibatches, its equilibrium probability is strictly below its data frequency for every learning rate, and the ratio tends to $κ/(2(e^{κ/2}-1))$ . Here $κ$ is the mean number of steps between occurrences divided by the second-moment time constant $1/(1-β_2)$ . In the same model, SGD and AMSGrad retain the unbiased fixed point. We test these predictions both in a unigram model and in a small language model trained from a known generating distribution. With random arrivals, the bias is larger than the periodic formula predicts; in the language model, the optimizers with the biased fixed point also fit the generating distribution less well.

Figures & tables

Optimizer	Update	Conserved (Prop. 1)	$\\|\Delta\bar{\mathbf{w}}\\|_{2}$	$\|\Delta\bar{b}\|$
SGD	linear	yes	$5.4\times 10^{-17}$	$3.5\times 10^{-18}$
Heavy ball	linear	yes	$5.7\times 10^{-17}$	$3.5\times 10^{-18}$
Nesterov	linear	yes	$9.0\times 10^{-17}$	$3.5\times 10^{-18}$
Shampoo	Kronecker	yes	$1.8\times 10^{-16}$	$1.6\times 10^{-15}$
Muon (heavy ball on $b$ )	orthogonalized	yes	$7.5\times 10^{-17}$	$6.9\times 10^{-18}$
Coupled Adam	shared scaling	yes	$5.0\times 10^{-17}$	$0$

Table 1: Change of the mean output embedding and mean bias after training a softmax regression in float64. “Conserved” is the prediction of Proposition 1 ; Muon is applied to

W

, with heavy-ball momentum on

b

Figure 1: Fixed points for rare tokens. (a) Single token,

\beta_{2}=0.99

: measured

\log(p^{\star}/q)

for RMSProp with periodic and random arrivals and for Adam (

\beta_{1}=0.9

) with periodic arrivals; solid: exact fixed point of Theorem 4 (iii); dashed:

\log\rho(\kappa)

. (b) Unigram model,

V=5000

: time average of

\log(p_{i}/q_{i})

, binned by

\kappa_{i}=(1-\beta_{2})/\pi_{i}

, for three pairs

(\beta_{2},B)

; circles: RMSProp, squares: Adam. (c) The same for SGD, AMSGrad and Coupled Adam (

\kappa_{i}

computed with

\beta_{2}=0.99

Figure 2: Small language model. (a)

r_{y}=\log(\bar{p}_{y}/f_{y})

, binned by the expected number of occurrences in a minibatch; dashed and dotted:

\log\rho(\kappa)

from the unigram theory for

\beta_{2}=0.95

and

0.999

. (b) Mean logit over the vocabulary on held-out contexts.

Optimizer	Test CE	KL	$\bar{r}$ , $e_{y}{<}0.05$	$\bar{r}$ , $e_{y}{\geq}1$	Mean logit	Unused $\log p$
SGD + momentum	$3.467$	$0.0940$	$-0.22$	$-0.06$	$0.00$	$-13.5$
Adam, $\beta_{2}{=}0.95$	$3.603$	$0.2278$	$-2.72$	$0.04$	$-46.83$	$-27.4$
Adam, $\beta_{2}{=}0.999$	$3.537$	$0.1586$	$-0.06$	$-0.01$	$-8.45$	$-21.7$
RMSProp, $\beta_{2}{=}0.95$	$3.594$	$0.2190$	$-2.65$	$0.03$	$-53.42$	$-26.6$
AMSGrad, $\beta_{2}{=}0.95$	$3.500$	$0.1241$	$-0.14$	$0.01$	$-5.25$	$-16.4$
Coupled Adam, $\beta_{2}{=}0.95$	$3.499$	$0.1238$	$-0.18$	$-0.02$	$-0.00$	$-13.2$

Table 2: Small language model after training. KL: divergence from the generating distribution to the model, averaged over contexts.

\bar{r}

: mean of

r_{y}

over tokens in the stated range of

e_{y}

. Mean logit: over the vocabulary and 2048 held-out contexts. Unused

\log p

: average log-probability of the entries that never occur.

Adaptive Gradient Methods

Adaptive Gradient Methods adjust updates using gradient history, curvature, or noise estimates to improve optimization across diverse problems.

See all 157 papers

Nov 30, 2025

Provable Benefit of SignGD: A Minimal Model Under Heavy-Tailed Class Imbalance

Adaptive and non-Euclidean optimizers often outperform Euclidean methods such as stochastic gradient descent (SGD) in language modeling by a large margin. Existing theory usually explains this gap by assuming favorable smoothness geometry or noise structure tailored to the specific optimizer. We instead ask whether such geometry can be induced from a concrete learning setting. Starting from an optimizer gap that persists across realistic language-modeling experiments, we progressively remove sequence dependence, architectural complexity, and stochasticity. We find that the gap exists in a minimal setting: the softmax unigram model with heavy-tailed data. This model exposes a simple deterministic mechanism under heavy-tailed class imbalance. We prove that GD learns rare tokens slowly because the corresponding logits receive only tiny updates, while SignGD removes this magnitude dependence and moves rare and common coordinates on a more comparable scale. We make this precise with upper and lower bounds for the convergence rate of GD and upper bounds for the convergence of SignGD. Our stochastic bounds contain additional noise-dependent terms that can obscure this advantage in the convergence guarantees and can be reduced by increasing the batch size

Robin Yadav, Shuo Xie, Tianhao Wang +1

Toyota Technological Institute at Chicago · University of British Columbia · University of California, San Diego

Sign-Based Optimization Class-Imbalanced Learning Gradient Descent Language Modeling Stochastic Optimization Convergence

Ask PDF

May 20, 2026

Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers

Preconditioned optimizers are central to language model training, but their stochastic update rules are usually treated as direct approximations to population preconditioned descent. We show that this view misses two finite-sample biases. First, the gradient and preconditioner are typically estimated from the same minibatch, introducing gradient--preconditioner coupling bias. Second, even when the preconditioner estimate is unbiased, its inverse or inverse-root is generally biased because inversion is nonlinear. We propose a single-batch bias-correction framework that addresses both effects: cross-fitted preconditioning estimates the numerator and preconditioner from independent microbatch groups, while variance-corrected inversion uses microbatch variability to subtract the leading delta-method bias term. The framework applies to diagonal moment, diagonal curvature, and matrix preconditioning methods, instantiated in AdamW, Sophia, and Shampoo. Bias correction reduces held-out pretraining loss on Qwen2.5-0.5B by

0.15

0.07

, and

0.11

nats, respectively; the effects on mixed-quality pretraining and downstream instruction tuning are consistently neutral-to-positive. Together, these results establish bias correction as a practical mechanism for reducing finite-sample update bias and improving the performance of preconditioned optimizers.

Nikhil Nayak, Julia White, Urchade Zaratiana +7

Fastino Labs

Efficient Language Model Training Adaptive Gradient Methods

Ask PDF

May 9, 2026

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Training large language models requires optimization algorithms that are not only statistically effective, but also computationally and memory efficient at extreme scale. Although Adam remains the dominant optimizer for large-scale language-model pretraining and fine-tuning, recent work has revisited nearly every component of the optimization stack: adaptive moment estimation, decoupled weight decay, memory footprint, curvature approximation, sign-based updates, large-batch stability, low-rank gradient structure, and matrix-wise orthogonalized updates. This survey reviews optimizer design for large language models through a systems-and-optimization lens. We organize the literature into classical first-order optimizers, adaptive optimizers, memory-efficient variants, second-order and curvature-aware methods, sign-based and discovered optimizers, low-rank and projection-based methods, and matrix-based optimizers such as Muon. We also discuss benchmarking methodology, including hyperparameter fairness, scale dependence, wall-clock efficiency, token efficiency, memory overhead, and downstream evaluation. We argue that optimizer research for LLMs is entering a new phase: moving from single-algorithm speedup claims toward rigorous, scale-aware comparisons that jointly evaluate convergence, stability, memory, and implementation complexity.

Aditya Ranganath

Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

Deep Learning Optimization LLM Input Optimization Memory-Efficient Optimization Matrix Optimization LLM Training

Ask PDF

Abstract

Adaptive Gradient Methods

Explore similar work

Provable Benefit of SignGD: A Minimal Model Under Heavy-Tailed Class Imbalance

Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers