PaperlayerView exact revision on arXiv

Source-linked AI summary

Why Adaptive Optimizers Underestimate Rare Tokens

Sangsidhya Kar

arXiv:2609.37535v1cs.LGstat.ML

TL;DR

Rare-token gradients meet coordinate-wise adaptive normalization in the softmax output layer, raising the question of whether optimizer updates distort token probabilities. The paper analyzes optimizer conservation and rare-token fixed points, then tests the predictions in unigram and small language models. It finds optimizer-dependent conservation and fixed-point behavior, with random arrivals producing more bias than the periodic formula predicts and biased-fixed-point optimizers fitting the generating distribution less well.

  • Problem

    Rare tokens receive small positive logit gradients on most steps and larger negative gradients when targeted, while adaptive methods scale updates by running gradient magnitudes, creating an imbalance that can affect their predicted probabilities.

    Source: §1 · pp. 1–2

  • Method

    The paper analyzes output-embedding conservation and rare-token fixed points, then tests the predictions in a unigram model and a small language model trained from a known generating distribution.

    Source: §1 · p. 2, §5.3 · p. 5

  • Results

    Random arrivals produced larger bias than the periodic formula predicts at κ ≥ 2, and in the small language model optimizers with biased fixed points had higher KL divergence from the generating distribution.

    Source: §5.2 · p. 5, §5.3 · p. 5

  • Takeaways & Limitations

    The common output-embedding shift leaves predictions unchanged for an untied output layer without weight decay, but rare-token fixed-point bias changes predicted probabilities.

    Source: §1 · p. 2, §3 · p. 3

    The theory relies on a single bias coordinate, periodic arrivals, β1 = 0, ϵ = 0, and assumption (A2), while experiments relax only the first three on small synthetic problems.

    Source: §7 · p. 7

Abstract

from arXiv · show

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 $\kappa/(2(e^{\kappa/2}-1))$. Here $\kappa$ is the mean number of steps between occurrences divided by the second-moment time constant $1/(1-\beta_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.

1 Introduction

Rare-token gradients interact asymmetrically with coordinate-wise normalization, shifting individual tokens’ training fixed points. The paper distinguishes this prediction-changing bias from a common output-embedding shift and characterizes both effects.

1 Introduction

  • A rare token’s small positive gradients between target appearances are normalized by a decaying second-moment estimate, while its larger negative target gradient arrives when that estimate is highest.This suppresses the upward update after an observation relative to later downward updates.Source: §1 · p. 1
  • The common embedding shift breaks mean-embedding conservation, but for an untied output layer without weight decay it leaves predictions unchanged.The separate rare-token fixed-point shift changes predicted probabilities.Source: §1 · pp. 1–2
  • The paper characterizes mean-embedding conservation across optimizer classes and derives an exact per-step change for coordinate-wise methods; sign descent and RMSProp also have distinct rare-token fixed-point results.Sign descent lowers logits for tokens appearing in fewer than half the minibatches; periodic-arrival RMSProp has a closed-form biased fixed point, while SGD and AMSGrad remain unbiased.Source: §1 · p. 2
  • At κ > 10, the predicted ratio ρ(κ) is approximately 0.03; with β2 = 0.95 and batch size 4096, this applies to tokens below frequency 1.2×10−6.Tokens below frequency 1.2×10−5 have κ > 1 in that setting.Source: §1 · p. 2

2 Setting

The setting defines the softmax output layer and its loss, then uses gradient identities and an output-layer symmetry to frame optimizer updates.

2 Setting

  • The output layer maps a network representation to logits through z = Wh + b, with W and b denoting the output weights and bias.The loss is log-sum-exp of the logits minus the target-weighted logits, with predictions p = softmax(z).Source: §2 · p. 2
  • For each minibatch, the logit gradient is p − t and its entries sum to zero.This is the stated basis for the output-layer gradient property used in the analysis.Source: §2 · p. 2
  • For an untied output layer, adding the same vector to every output embedding and the same scalar to every bias leaves the loss and gradients unchanged.The transformation leaves p unchanged, so the representation gradient is unchanged as well.Source: §2 · p. 2
  • The optimizer definitions cover Adam, RMSProp, AMSGrad, Coupled Adam, sign descent, Lion, and Adafactor.They differ in momentum, second-moment scaling, or sign-based updates.Source: §2 · p. 2

3 The mean output embedding

The section characterizes when output-layer updates preserve the mean embedding and gives an exact expression for changes under coordinate-wise scaling. It also separates this common shift from prediction-changing rare-token bias.

3 The mean output embedding

  • SGD and momentum updates, specified Kronecker-factored updates, orthogonalized updates such as Muon, and vocabulary-independent scaling preserve the mean output embedding.With decoupled weight decay, the mean instead decays geometrically.Source: §3 · pp. 2–3
  • For coordinate-wise methods, the paper gives an exact per-step identity for the change in the mean output embedding.The identity applies to arbitrary coordinate-wise scaling of momentum-based updates.Source: §3 · p. 3
  • For sign-based updates, the mean-embedding change is expressed through counts of positive and negative entries in each update column.The source specifies those counts but the supplied equation text does not preserve the full expression.Source: §3 · p. 3
  • For a bias coordinate, stale positive momentum combined with a small second moment can make the mean bias decrease under coordinate-wise scaling.This is the mechanism associated with the positive covariance between momentum and scaling.Source: §3 · p. 3
  • For an untied output layer without weight decay, the common shift leaves predictions unchanged in exact arithmetic, whereas the separate rare-token effect changes predictions.Finite precision, weight decay, tied embeddings, and auxiliary objectives can make the common shift matter.Source: §3 · p. 3

4 Biased fixed points for rare tokens

The unigram model isolates how token arrivals affect adaptive-optimizer fixed points. Sign descent keeps pushing down logits for sufficiently infrequent tokens, while periodic-arrival RMSProp has a stable equilibrium below the token’s data frequency.

4.1 A unigram model

  • In the unigram model, linearly combining gradients yields the usual fixed point p=q, the population-loss minimizer.A token’s minibatch arrival probability is π_i=1−(1−q_i)^B; when q_iB is small, its mean gap between appearances is about 1/π_i.Source: §4.1 · p. 3

4.2 Sign descent

  • Sign descent lowers a token’s logit at expected rate at least η(1−2π_i) when π_i<1/2, regardless of how small its predicted probability becomes.For a token that never occurs, the logit decreases exactly linearly; the theorem concerns logits, and part of the decrease is a common shift.Source: §4.2 · p. 3

4.3 RMSProp: the fixed point for periodic arrivals

  • For periodic arrivals every N≥3 steps, RMSProp’s stable fixed point has probability p⋆<q_i, and its ratio to q_i does not depend on the learning rate.The ratio depends only on N and β2; at κ=1, the large-N ratio is 0.771, with κ comparing the inter-arrival interval to the second-moment time constant.Source: §4.3 · p. 4
  • Under the same periodic-arrival assumptions, SGD and AMSGrad retain the fixed point p=q_i, while sign descent has no fixed point for N≥3.For SGD, the fixed point is stable; for AMSGrad with β1=0, the running maximum converges and the period change has the same zero.Source: §4.3 · p. 4
  • The RMSProp fixed-point theorem does not cover β1>0, random arrivals, or output-weight rows whose gradients depend on token contexts.The paper tests these effects separately.Source: §4.3 · p. 4

5 Experiments

The experiments test optimizer conservation and rare-token fixed-point predictions in synthetic models and a small language model. Random arrivals produce larger bias than the periodic formula predicts, and biased optimizers fit the generating distribution less well in the language-model experiment.

5 Experiments

  • Every experiment in the section ran on a CPU in under two hours, with full configurations provided in Appendix B.Source: §5 · p. 4

5.1 Conservation across optimizers

  • Optimizers covered by Proposition 1 preserved the mean output embedding and mean bias up to a largest change of 1.6 × 10−15, while remaining optimizers changed them by at least 3.6 × 10−1.The experiment used softmax regression with 512 classes, including 128 that never occurred.Source: §5.1 · p. 4, Table 1 · p. 5

5.2 A single token and the unigram model

  • Periodic-arrival fixed points agreed with Theorem 4(iii) within 0.009 in log p, while random arrivals at κ ≥ 2 averaged −5.88 below the periodic prediction.For κ < 1, random-arrival results remained close to that prediction, with mean deviation −0.078; Adam with β1 = 0.9 differed from the periodic RMSProp fixed point by at most 0.06.Source: §5.2 · p. 5, Figure 1 · p. 6
  • For random arrivals, tokens with κ > 1 received less probability than their frequency under RMSProp and Adam, and the bias grew with κ.At κ near 5, RMSProp with β2 = 0.99 gave −3.55 against the periodic prediction log ρ(5) = −1.50; at larger κ, the three (β2, B) settings separated.Source: §5.2 · p. 5, Figure 1 · p. 6

5.3 A small language model

  • The small language model used 4096 vocabulary entries, half never occurring, and measured predictions against its known generating distribution over contexts.The model was trained on 5.2 × 10^5 pairs; the plotted rare-token statistic was ry = log(average predicted probability / empirical frequency), binned by expected occurrences per minibatch.Source: §5.3 · p. 5, Figure 2 · p. 6
  • With β2 = 0.95, rare tokens (ey < 0.05) had mean ry of −2.72 for Adam, −2.65 for RMSProp, and −2.25 for AdamW; raising β2 to 0.999 reduced the bias to −0.06.The biased methods with β2 = 0.95 also had the highest KL and test cross-entropy; SGD had the lowest KL, followed by AMSGrad and Coupled Adam.Source: §5.3 · p. 6, Figure 2 · p. 6

6 Related work

The related work connects the paper to optimizer conservation laws, output-embedding shifts, rare-class adaptation, Adam convergence, and long-tailed calibration. The paper distinguishes its focus on token-level fixed points from prior work on common shifts and convergence behavior.

  • Prior symmetry work studied conserved quantities and symmetry-breaking effects; this paper places Shampoo and Muon in a conserving class and gives exact per-step changes for coordinate-wise methods that break the conservation law.The cited methods include Kronecker-factored and orthogonalized optimizers (Gupta et al., 2018; Bernstein & Newhouse, 2024).Source: §6 · p. 6
  • Earlier work connected Adam’s second-moment normalization to common embedding shifts and output-logit divergence, whereas this analysis studies token fixed points that centering leaves unchanged.The cited work includes Gao et al. (2019), Bi´s et al. (2021), Wortsman et al. (2024), and Stollenwerk et al. (2026).Source: §6 · p. 7
  • Prior rare-class work found slower gradient-descent progress under heavy-tailed imbalance; this paper adds that adaptive normalization’s faster progress can accompany a shifted fixed point.Kunstner et al. (2024) reported the progress comparison; the connection between Adam and sign descent is discussed by Balles & Hennig (2018) and Kunstner et al. (2023).Source: §6 · p. 7
  • The rare-token gradient pattern matches Reddi et al.’s Adam non-convergence construction, while Theorem 4 gives the corresponding fixed point in closed form.AMSGrad was introduced as a remedy in that construction; Adagrad appears here as the κ → 0 limit.Source: §6 · p. 7
  • Logit adjustment explicitly corrects rare-class logits, while this analysis identifies an additional optimizer-dependent mechanism that can shift rare-class probabilities.The cited calibration work includes Menon et al. (2021) and Guo et al. (2017).Source: §6 · p. 7

7 Discussion and limitations

The theory relies on simplifying assumptions, and experiments relax only some of them on small synthetic problems. Whether the observed language-model effect extends to large models or downstream rare-word generation remains open.

  • The theory assumes a single bias coordinate, periodic arrivals, β1 = 0, ϵ = 0, and assumption (A2), while experiments relax the first three only on small synthetic problems.The small language model showed higher KL divergence for optimizers with biased fixed points, but whether this is measurable in large language models and affects rare-word generation remains open.Source: §7 · p. 7
  • During pretraining with millions of tokens per batch, κ > 1 occurs only for extremely rare tokens, whereas fine-tuning and other small-batch settings with β2 = 0.95 can place a substantial fraction of vocabulary above κ = 1.For output layers, proposed bias-reduction options include higher β2, larger batches, AMSGrad, Coupled Adam, SGD for the output bias, or positive ϵ when √v falls below it.Source: §7 · p. 7
  • With tied embeddings, the common shift affects loss because it also shifts input embeddings.Source: §7 · p. 7

A Proofs

The appendices prove conservation and fixed-point claims by tracking optimizer updates and the periodic second-moment recursion. They also establish the sign-descent drift bound and the limiting rare-token probability ratio.

  • The proof of Proposition 1 uses the zero-sum gradient property to show that specified optimizer updates preserve the mean output embedding.For Shampoo, the proof uses the all-ones vector as an eigenvector of the relevant symmetric matrix, then shows the transformed update retains the required form.Source: A Proofs · p. 9
  • For sign descent, absent tokens have positive gradients and decrease their logits by η, while target appearances change them by at most η; a martingale bound yields the long-run result.The proof conditions on the token’s appearance probability and applies Azuma–Hoeffding with Borel–Cantelli.Source: A Proofs · p. 9
  • Under periodic arrivals, the second-moment recursion is an affine map over each period, and summing normalized gradients gives the per-period change Δ(x).The proof establishes a unique fixed point x⋆ > N and hence p⋆ < qi for N ≥ 3; stability follows from the sign of Δ around that point.Source: A Proofs · p. 9
  • With β2 = 1 − κ/N, the periodic fixed-point ratio converges to ρ(κ) = κ/(2(e^(κ/2) − 1)).Source: A Proofs · pp. 9–10
  • The AMSGrad running maximum converges to a positive limit, so the per-period change becomes proportional to the sum of gradients; sign updates have one negative arrival step and N − 1 positive absent-token steps.Source: A Proofs · p. 10

B Experimental details

The experiments cover optimizer conservation, single-token and unigram dynamics, and language-model fitting, with settings specified separately for each setup.

B Experimental details

  • The conservation experiment uses a 512-class softmax regression with Zipf-weighted labels, random-mean Gaussian features, and 64-example minibatches for 300 steps.It reports changes rather than losses; the learning rates set the scale of those changes. Optimizer rates are SGD 0.5; heavy ball and Nesterov 0.05; Shampoo and Muon 0.02; Adam, RMSProp, AMSGrad, Coupled Adam, and Adafactor 3 × 10^-3; Lion 3 × 10^-4; and sign descent 10^-3.Source: B Experimental details · p. 10
  • The single-token run uses β2 = 0.99, a = 10^-3, η = 3 × 10^-3, and 4 × 10^5 steps, averaging the logit over the final third.It uses ϵ = 0, no bias correction, and initializes the logit to log(a/N).Source: B Experimental details · p. 10
  • The unigram experiment trains on 5,000 tokens with qi ∝ i^-1.2 for 3 × 10^5 steps and records log p every 20 steps in the second half.All runs with the same batch size use the same minibatches; learning rates are 4 × 10^-3 for RMSProp, Adam, AMSGrad, and Coupled Adam, 1 for SGD, and 10^-3 for sign descent.Source: B Experimental details · p. 10
  • The language-model experiment trains a width-64 residual MLP on 2 × 10^4 minibatches of 256 pairs from a 2,048-token Markov chain.Training uses 8,192 length-64 chains after a 50-step burn-in; evaluation uses 1,024 length-32 chains. The model has LayerNorm, a 64 → 256 → 64 GELU MLP, and an untied linear output layer; test cross-entropy uses held-out pairs, while KL uses the exact transition matrix weighted by the empirical context distribution.Source: B Experimental details · p. 10

Related papers

Loading 2609.37535v1…