Adaptive Step Sizes: AdaGrad, RMSProp, and Adam

A single scalar step size assumes that the objective has comparable curvature in every coordinate direction. That assumption fails constantly: sparse text features have gradients that fire in a handful of coordinates, embedding rows differ in update frequency by orders of magnitude, and a scaling difference between two input features changes the curvature of the corresponding directions without changing the problem. Adaptive methods address this by maintaining, for each coordinate, a running statistic of the gradients that have been seen in it, and dividing the step size by a function of that statistic. The result is a diagonal preconditioner computed from the trajectory itself rather than from any curvature estimate, which is why it is cheap and why its guarantees are subtle.

Where a fixed step size loses

Consider a quadratic whose Hessian is diagonal with entries spanning [μ, L], so κ = L/μ. Gradient descent at η = 1/L makes progress in the high-curvature directions immediately and crawls in the low-curvature ones. The optimal diagonal rescaling would set the step for each coordinate from its own curvature, giving an iteration that reaches the solution in a single step for a diagonal quadratic. An adaptive method tries to approximate that rescaling using only observed gradients. Its secondary benefit is scale invariance: if a feature is multiplied by a constant, its gradient is divided by that constant and the accumulator divides it out again, so the effective step in that coordinate is unchanged. Neither gradient descent nor the momentum methods described in this cluster has this property; their effective step size for a coordinate depends on its scale.

Accumulating squared gradients

AdaGrad keeps a per-coordinate sum of squared gradients from the start of optimization,

Gₜ = Gₜ₋₁ + gₜ², xₜ₊₁ = xₜ − η gₜ / (√Gₜ + ε),

with all operations applied coordinate-wise. The effective step for a coordinate is η/√Gₜ, so a coordinate that has received large gradients before gets a small step now. The normalisation is by accumulated magnitude, not by a bounded statistic, which means the step size is non-increasing and converges to zero at rate 1/√t for a coordinate with roughly constant gradient magnitude. That is the source of both the method’s safety and its weakness: safety because no step size can ever be catastrophically too large for a coordinate that has already produced a huge gradient, weakness because on a long run the steps become small enough that progress in frequently updated coordinates effectively stops. Constraints are orthogonal to the rescaling: a feasible set is handled by composing the step with a projection onto it, as in projected and proximal gradient methods.

The design motive was the ℓ₂-norm regularisation of online learning. AdaGrad’s published analysis bounds regret relative to the best fixed decision in hindsight by O(√T) for T rounds, matching the lower bound for general online convex problems up to constants, and the accumulation of squared gradients is exactly the potential function that makes the bound work. The mechanism therefore comes with its own reason to prefer it: it is the update rule whose natural bound is the optimal one for the online setting, not a heuristic that happens to work.

Exponential averaging and bias correction

RMSProp replaces the unbounded sum by an exponentially weighted average,

vₜ = β₂ vₜ₋₁ + (1 − β₂) gₜ², xₜ₊₁ = xₜ − η gₜ / (√vₜ + ε).

The statistic now reflects roughly the last 1/(1 − β₂) iterations rather than all of history, so the effective step stops decaying and the method is suitable for non-stationary objectives. At the conventional β₂ = 0.999 the window is about a thousand steps; at β₂ = 0.9 it is ten. Adam adds a momentum term over the rescaling, maintaining mₜ = β₁ mₜ₋₁ + (1 − β₁) gₜ with vₜ as above, and dividing the momentum by the square root of the second-moment estimate. With zeros for the initial state the estimates are biased toward zero, since the weighted average of nothing still has total weight 1 − β₁ᵗ rather than one; the correction m̂ₜ = mₜ/(1 − β₁ᵗ) and the analogous v̂ₜ remove exactly that factor, which matters most in the first few iterations when β₂ = 0.999 gives 1 − β₂ᵗ ≈ 0.001.

The typical defaults, β₁ = 0.9, β₂ = 0.999 and ε = 10⁻⁸, are published values from the Adam paper. The ε deserves more attention than it usually receives: it is added inside the square root’s denominator, so it only matters when √vₜ is comparable to it. For a coordinate whose gradients are persistently around 10⁻⁸ in magnitude, ε dominates and the effective step becomes η gₜ/ε, a scaled steepest-descent step with gain 10⁸. Coordinates of that size occur routinely in sparse embeddings and in parameters whose gradients are many orders of magnitude below the typical scale, which is why both larger and smaller ε values are used in practice and why the choice interacts with mixed-precision arithmetic, where the accumulator itself must stay in a wide format.

def adam_step(x, g, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    """One Adam update over coordinate lists; t counts updates and starts at 1."""
    for i in range(len(x)):
        m[i] = b1 * m[i] + (1 - b1) * g[i]
        v[i] = b2 * v[i] + (1 - b2) * g[i] * g[i]
        m_hat = m[i] / (1 - b1 ** t)          # bias correction
        v_hat = v[i] / (1 - b2 ** t)
        x[i] -= lr * m_hat / (v_hat ** 0.5 + eps)
    return x

The convergence counterexample

Adaptive methods with exponential averaging were introduced with an appeal to intuition rather than a proof for the general convex case, and the gap turned out to be real. Reddi and colleagues constructed a convex online problem on which Adam’s regret is not bounded by O(√T): the mechanism of the failure is that the second-moment estimate can decrease between rounds, so the effective step size can increase, and the algorithm can be made to oscillate around the optimum with a persistent gap. The repair is to enforce monotonicity of the denominator, keeping a running maximum v̂ₜ = max(v̂ₜ₋₁, vₜ); the resulting update, sometimes called AMSGrad, restores the O(√T) regret bound. Note precisely what the counterexample does and does not show. It shows that the popular update rule lacks the guarantee its predecessors had; it does not show that the rule fails on the problems people use it for, and the accumulating variant AdaGrad is unaffected because its denominator cannot decrease.

The related confusion is between ℓ₂ regularisation and weight decay. Adding λ‖x‖²/2 to the objective contributes λx to the gradient, and that term then passes through the adaptive rescaling, so coordinates with small gradient history get a disproportionately large shrinkage. Decoupled weight decay applies the shrinkage directly to the parameters instead, outside the adaptive step, which changes both the effective regularisation strength per coordinate and the interaction with the step size at convergence. The two are equivalent only for plain gradient descent.

What the state costs

The bookkeeping is not free. Plain stochastic gradient descent stores the parameters and the current gradient: two arrays. Adding momentum makes three, and Adam stores momentum plus second moment for four arrays, with the second moment and momentum both staying in a wide floating-point format in mixed-precision training. Dark-momentum accounting matters at scale: for a model with 10⁸ parameters, one array in 32-bit precision is 400 MB, so Adam’s optimiser state is roughly 800 MB on top of the parameters and gradients, and a correspondingly larger amount if master weights in 32-bit are kept alongside 16-bit copies for training. Methods that reduce the state, by storing the second moment in a narrower format or a quantised approximation of it, trade some numerical stability for a proportional memory saving.

The precision of the state is a second constraint. The second-moment estimate is a running average of squared gradients, so a single very small value stored in a narrow format can round to zero, turning the guarded division into a division by ε. Training in half precision therefore keeps the optimizer state in 32-bit even when the parameters and gradients are held in 16-bit, and loss scaling multiplies the gradients before the step to hold them inside the half-precision range. Because the accumulators persist across the whole run, their format cannot be changed midway without discarding the accumulated history, so the layout of the state is a decision made before the first step.

When the adaptivity is the problem

Three failure patterns are documented rather than hypothetical. Adaptive rescaling is essentially a diagonal approximation to the inverse Hessian built from gradient history, and on problems where the curvature is dominated by off-diagonal coupling it provides an inaccurate preconditioner while hiding the poor conditioning from the step-size selection, so the run appears stable while converging slowly. Weight-decay-through-the-gradients has a different effect on coordinates that are rarely updated, as noted above. And the interaction with momentum is not compositional: the temporal filter discussed in the momentum article and the per-coordinate scaling operate on the same gradient sequence, and neither analysis covers the combination, so the practical advice to tune β₁ and β₂ jointly rather than independently follows from the absence of a guarantee rather than from a theorem that one is worse. Where the objective is smooth, well-scaled and badly conditioned with coupling, the quasi-Newton methods described in quasi-Newton methods and the L-BFGS recursion use actual curvature information and typically converge in far fewer iterations, at a memory and per-iteration cost that is usually prohibitive for large models.