Adam Optimizer › Why Per-Parameter Learning Rates

Michael BrenndoerferApril 24, 20255 min read
Language AI HandbookAll chapters
Part 10 of 67Neural Network Foundations
Chapter 68Adam Optimizer › Why Per-Parameter Learning Rates
Reading section 2 of 18 · Why Per-Parameter Learning Rates · 5 minSection 2 of 18 · 5 min

Book contents

Section 2 of 18 · 5 min read

Adam Optimizer › Why Per-Parameter Learning RatesOne section · Sections in this chapter2 / 18 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

To see why a single global learning rate is limiting, consider a sparse input problem like word embeddings. In language modeling, some words appear thousands of times per batch; others might appear once per epoch. When you compute gradients and average them over a minibatch, frequently occurring words generate large, consistent gradient signals. Rare words generate tiny gradients most of the time, with occasional large spikes. If you set the learning rate high enough for rare words to learn, frequent words overshoot. If you set it low enough for frequent words, rare words barely move.

The same tension appears in deeper networks. Parameters close to the output receive large, stable gradients; parameters in early layers receive small, noisy gradients that have been attenuated through many multiplications during backpropagation. We saw in the backpropagation chapter how the chain rule multiplies gradients across layers, which can leave early-layer parameters chronically under-updated or chronically over-updated depending on the global learning rate.

The Geometry of the Problem: Ill-Conditioned Loss Surfaces

The problem with a single learning rate becomes especially vivid when you think about the geometry of the loss surface. In a well-conditioned problem, the loss function is roughly spherical: the curvature is similar in every direction. SGD with a fixed learning rate moves at roughly the same speed in all directions. In an ill-conditioned problem, the loss surface looks more like a narrow canyon: very steep walls on one side, a very flat floor in another. The optimal path to the minimum runs along the canyon floor, but SGD with a fixed learning rate has to choose between two bad options. Use a large step size and you bounce back and forth off the steep walls, oscillating violently. Use a small step size and you crawl along the flat floor at a glacial pace.

This ill-conditioning is not a pathological edge case. It arises naturally in deep networks from weight initialization asymmetries, from the scale differences between input features, and from the structure of the weight matrices at different depths. A typical transformer training run will have gradient norms that vary by orders of magnitude across different parameter groups. The first embedding layer, the output projection, and the attention weight matrices all occupy different regions of the loss landscape with vastly different curvatures.

Adaptive learning rates address the ill-conditioning problem by measuring, for each parameter independently, how much curvature has been encountered. Parameters in steep dimensions accumulate large second-moment estimates and receive smaller effective learning rates. Parameters in flat dimensions accumulate small second-moment estimates and receive larger effective learning rates. The optimizer self-calibrates its step size to match the local geometry of each dimension simultaneously, without any manual tuning.

A Brief History: AdaGrad and RMSProp

The idea of parameter-specific step sizes had been explored before Adam. AdaGrad (Duchi, Hazan, and Singer, 2011) was the first major adaptive optimizer. It accumulated the sum of all squared gradients seen so far for each parameter, and used that sum in the denominator:

θt+1=θt−αGt+ϵgt\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{G_t} + \epsilon} g_t

where Gt=∑i=1tgi2G_t = \sum_{i=1}^{t} g_i^2 is the cumulative sum of squared gradients. AdaGrad was a significant step forward for sparse problems: parameters that receive occasional large gradients quickly build up a large GtG_t, causing their learning rate to shrink. Parameters that receive small gradients maintain a small GtG_t and continue to learn quickly. This worked well for training word vectors and other models with sparse inputs.

However, AdaGrad has a critical flaw for long training runs: the accumulator GtG_t only ever grows. It grows every time a gradient is non-zero, and it never forgets old gradient information. After enough training steps, GtG_t becomes so large that the effective learning rate α/Gt\alpha / \sqrt{G_t} shrinks to near zero for all parameters. Training effectively stalls. In practice, AdaGrad works well for shallow models trained for a small number of epochs, but fails to converge on deep networks trained for hundreds of epochs.

RMSProp, proposed by Geoffrey Hinton in an unpublished lecture (2012), addressed the monotonic accumulation problem by replacing the cumulative sum with an exponential moving average:

vt=β2vt−1+(1−β2)gt2,θt+1=θt−αvt+ϵgtv_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2, \quad \theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{v_t} + \epsilon} g_t

By using a decaying average rather than a cumulative sum, RMSProp's denominator reflects recent gradient magnitudes rather than the entire training history. Parameters that were frequently updated early in training but are now quiet will see their vtv_t decay, allowing their effective learning rate to recover. This makes RMSProp viable for long training runs where gradient statistics change substantially over time.

Adam combines RMSProp's adaptive second moment with momentum's first moment, wrapping both in a bias correction that makes the estimates accurate from the very first update. It is, in a sense, the synthesis of a decade of adaptive optimizer research into a single coherent algorithm.