Adam's rare-token discount survives a smaller learning rate

For one softmax token, cross-entropy gives the logit gradient p − y. On most minibatches a rare token is absent, so the gradient is a tiny positive p and gradient descent nudges its logit down. When the token finally appears, the gradient is large and negative, so the logit should jump back up.

SGD adds those signed contributions directly. Adam and RMSProp first divide by a decaying root-mean-square estimate. The rare, large arrival makes that denominator largest exactly when the upward correction is needed. During the long absence that follows, the denominator shrinks, amplifying each small downward step.

The paper compresses rarity and optimizer memory into κ = (1 − β₂) / π, where π is the minibatch arrival probability. Roughly, κ is the mean gap between appearances divided by the second-moment memory 1 / (1 − β₂). When κ exceeds one, the token is usually absent for longer than RMSProp remembers its last large gradient.

Under periodic arrivals and a large period, the fixed-point probability relative to the data frequency has a closed form. The site runs the TypeScript below; Python and C++ are line-for-line translations.

import math

def probability_ratio(kappa):
    if kappa == 0:
        return 1.0
    return kappa / (2.0 * math.expm1(kappa / 2.0))

The closed form is a limit. The core module also rebuilds the paper's finite periodic second-moment orbit, then solves its unique balance point by bisection. With β₂ = 0.99, the exact solution approaches the asymptotic curve as the period grows.

Sign descent exposes the same asymmetry more brutally. If a token appears in only 10% of minibatches, its theorem-level expected logit drift bound at learning rate 0.001 is -0.0008 per step. The force does not fade as the predicted probability approaches zero.

Here is the buried boundary that matters most. The derivation assumes an appearance exactly every N steps. Real tokens arrive randomly. Near κ = 5, the periodic formula predicts log(p/q) = −1.50, while the paper's random-arrival RMSProp run reaches −3.55. Occasional very long gaps let the second moment decay far beyond what the mean gap suggests.

The authors train a 4,096-entry language model from a known Markov distribution. For rare observed tokens, Adam with β₂ = 0.95 assigns about 6.6% of their empirical frequency on average after converting the reported log ratio back to a probability ratio. Raising β₂ to 0.999 nearly removes that local bias, although its overall KL still trails SGD in this experiment.

The mechanism suggests concrete output-layer interventions: increase β₂ or batch size, keep the second moment's running maximum with AMSGrad, share scaling across vocabulary rows with Coupled Adam, or use SGD for the output bias. A positive ε also caps amplification once the decayed root second moment becomes tiny.

Inspect softmax outputs on the interactive Transformer page and follow parameter updates on the neural-network page. The next decisive experiment should measure token-frequency calibration, rare-word generation, and downstream utility in a realistic fine-tune while changing only the output-layer optimizer.