Adaptive Optimizers: Momentum, RMSProp, Adam & Decoupled AdamW

Standard Stochastic Gradient Descent (SGD) struggles on ill-conditioned loss surfaces:

1. The Ravine Pathology:

Loss landscapes often resemble narrow canyons: the surface curves sharply in one direction (high curvature, steep walls), but gently along another direction (low curvature, valley floor leading to the minimum).

1. Standard SGD Failure:

w_(t+1) = w_t - η ∇ L(w_t)

Because the gradient is dominated by the steep orthogonal walls, SGD takes huge oscillating steps bouncing back and forth across the ravine, making negligible progress along the gentle path to the minimum.

1. The Polyak Momentum Formulation (Heavy Ball Method):

Boris Polyak (1964) introduced physical inertia into optimization:

v_t = β v_(t-1) + η g_t

w_(t+1) = w_t - v_t

Where β ∈ [0, 1) is the momentum decay factor (typically β = 0.9).

• Oscillating Components Cancel: Along the steep walls, gradient signs alternate (+,-,+,-), causing the accumulated velocity components to cancel out: ∑ g_⊥ ≈ 0.
• Consistent Components Accelerate: Along the valley floor, gradient signs consistently point in the same direction, accumulating velocity like a rolling snowball.

Polyak momentum is the foundational velocity engine behind SGD-Momentum and Adam. In training computer vision models (ResNet) on ImageNet, SGD with momentum β = 0.9 consistently outperforms vanilla SGD in both training speed and generalization accuracy.

Boris Polyak published the "Heavy Ball Method" in 1964 in the USSR. In 1986, Rumelhart, Hinton, and Williams incorporated Polyak's momentum term into the backpropagation learning algorithm, making deep network optimization practical.

• Heavy ball rolls through noise: High inertia ignores sideways chatter and charges down the canyon.
• Velocity accumulates consistent signals; cancels alternating noise: The mathematical definition of momentum.

#### 1. Recursive Expansion of Velocity

The Polyak momentum recurrence relation is:

v_t = β v_(t-1) + η g_t

Assuming initial velocity v₀ = 0, we unroll the recurrence over t steps:

v₁ = η g₁

v₂ = β (η g₁) + η g₂ = η ( g₂ + β g₁ )

v₃ = β ( η (g₂ + β g₁) ) + η g₃ = η ( g₃ + β g₂ + β² g₁ )

In general:

v_t = η ∑_(k=0)^(t-1) β^k g_(t-k)

Past gradients are weighted by exponentially decaying powers of β.

Assume the optimizer traverses a constant slope where the gradient is invariant: g_t = g for all t.

v_t = η g ∑_(k=0)^(t-1) β^k = η g ( 1 + β + β² + ṡ + β^(t-1) )

Using the classical formula for the sum of a finite geometric series (S_t = (1 - β^t) / (1 - β)):

v_t = η g ( (1 - β^t) / (1 - β) )

Taking the limit as t → ∞ (since β ∈ [0, 1) ⟹ lim_(t → ∞) β^t = 0):

v_∞ = lim_(t → ∞) v_t = (η g) / (1 - β)

• The Speedup Factor: The effective step size accelerates by a factor of:

Speedup = (1) / (1 - β)

• Numerical Evaluation for β = 0.90:

Speedup = (1) / (1 - 0.90) = (1) / (0.10) = 10×

With β = 0.9, the optimizer moves 10 times faster down constant slopes than standard SGD! \blacksquare

Let learning rate η = 0.1, constant gradient g = 2.0, momentum β = 0.90.

• Step 0: v₀ = 0.0
• Step 1: v₁ = (0.9)(0.0) + (0.1)(2.0) = 0.0 + 0.20 = 0.20 (Step size: 0.20)
• Step 2: v₂ = (0.9)(0.20) + 0.20 = 0.18 + 0.20 = 0.38
• Step 3: v₃ = (0.9)(0.38) + 0.20 = 0.342 + 0.20 = 0.542
• Step 4: v₄ = (0.9)(0.542) + 0.20 = 0.4878 + 0.20 = 0.6878
• Step 10: v_(10) = 0.2 × (1 - 0.9^(10)) / (1 - 0.9) = 2.0 × (1 - 0.3487) = 1.3026
• Terminal Velocity:

v_∞ = ((0.1)(2.0)) / (1 - 0.9) = (0.2) / (0.1) = 2.0000

The effective step size smoothly accelerates from 0.20 up to 2.00 (10× faster).

In PyTorch, momentum is configured in `torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)`. When tuning learning rates with momentum, engineers often decrease η by 10× to compensate for the 1/(1-β) terminal velocity amplification.

The geometric series sum formula dates back to Archimedes (circa 250 BC). Boris Polyak applied it in 1964 to prove linear convergence rates for smooth strongly convex functions.

• v_∞ = (η g) / (1 - β): The terminal velocity formula.
• β = 0.9 means 10×: A momentum of 0.9 accelerates you tenfold; β = 0.99 accelerates you a hundredfold.

Nesterov Accelerated Gradient (NAG) computes the gradient not at current position w_t, but at the anticipated lookahead position w_t - β v_(t-1):

v_t = β v_(t-1) + α ∇ L(w_t - β v_(t-1))

w_(t+1) = w_t - v_t

This lookahead provides predictive damping that brakes before overshooting sharp curves.

Implemented in PyTorch via `torch.optim.SGD(momentum=0.9, nesterov=True)` for faster convex convergence.

Invented by Yurii Nesterov (1983) in 'A method of solving a convex programming problem with convergence rate O(1/k²)', achieving theoretical optimal first-order convergence.

Look before you leap: Nesterov samples the slope where momentum will carry you, not where you stand.

AdaGrad accumulates the sum of squared historical gradients for each parameter i:

G_(t, i) = G_(t-1, i) + g_(t, i)²

w_(t+1, i) = w_(t, i) - (α) / (√(G_(t, i)) + ε) g_(t, i)

Parameters with frequent or large gradients experience monotonically shrinking learning rates.

The premier optimizer for sparse data domains (Google Ad Click Prediction, recommendation systems, GloVe word embeddings).

Published by John Duchi, Elad Hazan, and Yoram Singer (2011) in 'Adaptive Subgradient Methods for Online Learning and Stochastic Optimization'.

Rare features take large strides; frequent features take small steps: AdaGrad adapts to sparsity.

Anisotropic curvature poses a severe dilemma for standard optimization:

1. The Anisotropic Ravine Pathology:

Consider an elliptical loss surface:

L(w₁, w₂) = 100 w₁² + w₂²

• Axis w₁ has high curvature (100× steeper): Gradients (∂ L) / (∂ w₁) = 200 w₁ are massive.
• Axis w₂ has low curvature: Gradients (∂ L) / (∂ w₂) = 2 w₂ are tiny.
• If learning rate η is set high enough to make progress on w₂, updates along w₁ explode into instability. If η is set small enough to stabilize w₁, progress along w₂ grinds to a halt.

1. The RMSProp Solution (Geoffrey Hinton, 2012):

RMSProp maintains a running estimate of the magnitude (variance) of recent gradients independently for every coordinate:

s_t = β s_(t-1) + (1 - β) g_t²

Where g_t² is the element-wise square.

1. Adaptive Pre-Conditioning:

Parameters update by dividing the gradient by the square root of the running average:

w_(t+1) = w_t - (η) / (√(s_t + ε)) \odot g_t

• For steep coordinates (large g_t): √(s_t) is large ⟹ Step size is scaled down!
• For gentle coordinates (small g_t): √(s_t) is small ⟹ Step size is boosted!

The elliptical ravine is mathematically warped into an isotropic circular bowl where all coordinates progress at an equal pace.

RMSProp was the optimizer of choice for Deep Q-Networks (DQN, Mnih et al. 2015) in Atari gaming and recurrent neural networks. It forms the second-moment foundation of Adam and AdamW.

Geoffrey Hinton introduced RMSProp in Lecture 6e of his 2012 Coursera class "Neural Networks for Machine Learning". It was never formally published in a peer-reviewed paper, yet became one of the most cited algorithms in deep learning history.

• Divide by root of squared gradients: Tame the wild axes; amplify the quiet ones.
• Warp ellipses into circles: Equalize progress across all dimensions.

#### 1. Mathematical Formalism

Let gradient vector at step t be g_t = ∇ L(w_t).

1. Exponential Moving Average of Squared Gradients:

s_t = β s_(t-1) + (1 - β) g_t²

Where β ∈ (0, 1) is the decay rate (typically β = 0.90 or 0.99), and g_t² = g_t \odot g_t is the element-wise square.

1. Parameter Update Formula:

w_(t+1) = w_t - (η) / (√(s_t + ε)) \odot g_t

Where ε > 0 (typically 10^(-8)) is a numerical stabilizer preventing division by zero.

Unrolling the recurrence over past steps:

s_t = (1 - β) ∑_(k=0)^(t-1) β^k g_(t-k)²

The weight assigned to a gradient squared from k steps ago is (1 - β) β^k.

• Half-Life Definition: The number of steps k_(1/2) required for past gradient influence to drop by 50\% (factor of 0.5):

β^k_(1/2) = 0.5

Taking natural logarithms:

k_(1/2) ln(β) = ln(0.5) = -ln(2) ⟹ k_(1/2) = (-ln(2)) / (ln(β))

• Using the linear approximation ln(β) = ln(1 - (1-β)) ≈ -(1 - β) for β ≈ 1:

k_(1/2) ≈ (ln(2)) / (1 - β) = (0.693) / (1 - β)

• Effective Memory Horizon: The total effective window size is T_(eff) = (1) / (1 - β).
• For β = 0.90: T_(eff) = (1) / (0.10) = 10 steps (k_(1/2) ≈ 6.6 steps).
• For β = 0.99: T_(eff) = (1) / (0.01) = 100 steps (k_(1/2) ≈ 69 steps).

Unlike AdaGrad, which accumulates all past gradients forever (s_t = ∑ g² → ∞), RMSProp's leaky memory ensures the denominator never freezes.

Let η = 0.01, β = 0.90, ε = 1e-8, s₀ = 0.0.

Suppose coordinate 1 has massive oscillating gradients g = [+10.0, -10.0, +10.0], and coordinate 2 has gentle gradients g = [+0.5, +0.5, +0.5].

• Step 1:
• Coord 1: s₁^((1)) = 0.9(0) + 0.1(10²) = 10.0. Scale: √(10) ≈ 3.162.

Step₁ = (0.01) / (3.162) × 10.0 = 0.0316

• Coord 2: s₁^((2)) = 0.9(0) + 0.1(0.5²) = 0.025. Scale: √(0.025) ≈ 0.1581.

Step₂ = (0.01) / (0.1581) × 0.5 = 0.0316

Notice that despite the massive 20× difference in raw gradient (10.0 vs 0.5), RMSProp adjusted both coordinates to take the exact same step size of 0.0316!

In PyTorch, `torch.optim.RMSprop(model.parameters(), lr=0.01, alpha=0.99, eps=1e-8)` implements this exact algorithm (where `alpha` corresponds to β).

John Duchi, Elad Hazan, and Yoram Singer invented AdaGrad in 2011. While AdaGrad proved theoretically sound for convex optimization, its monotonic sum of squared gradients caused learning rates to decay to zero prematurely in deep networks. Geoffrey Hinton designed RMSProp's EMA leaky memory specifically to fix AdaGrad's premature freezing.

• s_t = β s_(t-1) + (1-β) g²: The leaky exponential moving average of gradient energy.
• Memory window is (1) / (1-β): Forget the ancient past; adapt to the terrain of the last few dozen steps.

Adam (Adaptive Moment Estimation) unites the strengths of Polyak Momentum and RMSProp into a unified first-order optimization algorithm:

1. Branch 1: The First Moment Vector (Polyak Momentum Engine):

Tracks the directional momentum (mean) of gradients via an exponentially weighted moving average:

m_t = β₁ m_(t-1) + (1 - β₁) g_t

Typically β₁ = 0.9. This accumulates velocity along consistent descent directions and cancels out high-frequency oscillating noise across opposing steps.

1. Branch 2: The Second Raw Moment Vector (RMSProp Adaptive Metric):

Tracks the uncentered variance (energy) of gradients per coordinate:

v_t = β₂ v_(t-1) + (1 - β₂) g_t²

Typically β₂ = 0.999. Coordinates with chronically huge gradients accumulate large v_t, whereas coordinates with sparse or tiny gradients yield small v_t.

1. The Parameter Update Equation:

By taking the element-wise quotient of the momentum velocity and the square root of coordinate variance:

θ_(t+1) = θ_t - (α) / (√(v̂_t) + ε) m̂_t

Adam normalizes parameter updates so that each coordinate moves at approximately the scale of the learning rate α, regardless of gradient magnitude scale differences.

Adam is the universal default baseline optimizer across deep learning, powering initial training for computer vision architectures (ResNet, ViT), speech synthesis, and reinforcement learning policy networks before specialized cosine-decay tuning.

Diederik Kingma and Jimmy Ba introduced Adam in their landmark 2014 paper, Adam: A Method for Stochastic Optimization. It quickly became the most cited optimization paper in computer science history, solving the tedious hyperparameter grid-search required by vanilla SGD.

• First moment steers the compass; second moment tunes the throttle: Momentum selects the direction m̂_t, while RMSProp scales the coordinate step size 1/√(v̂_t).
• Inertia over volatility: Adam drives fast through calm prairies and steps gently over rocky cliffs.

#### 1. The Cold-Start Defect

When initializing Adam, the moment vectors are anchored at zero: m₀ = vec(0) and v₀ = vec(0).

Unrolling the recursive formula for m_t:

m₁ = (1 - β₁) g₁

m₂ = β₁(1 - β₁) g₁ + (1 - β₁) g₂

m_t = (1 - β₁) ∑_(i=1)^t β₁^(t-i) gᵢ

Taking mathematical expectations of both sides, assuming gradients gᵢ come from a stationary distribution with true mean E[g]:

E[m_t] = E[(1 - β₁) ∑_(i=1)^t β₁^(t-i) gᵢ] = E[g] · (1 - β₁) ∑_(i=1)^t β₁^(t-i)

The summation is a finite geometric series:

∑_(i=1)^t β₁^(t-i) = β₁^(t-1) + β₁^(t-2) + ṡ + 1 = (1 - β₁^t) / (1 - β₁)

Substituting this back into the expectation formula:

E[m_t] = E[g] · (1 - β₁) · (1 - β₁^t) / (1 - β₁) = E[g] · (1 - β₁^t)

Because β₁ ∈ (0, 1), (1 - β₁^t) < 1. Thus, E[m_t] systematically underestimates E[g].

Dividing by the analytical discrepancy yields the exact unbiased moment:

m̂_t = (m_t) / (1 - β₁^t) ⟹ E[m̂_t] = (E[m_t]) / (1 - β₁^t) = E[g]

Identically for the second moment:

v̂_t = (v_t) / (1 - β₂^t) ⟹ E[v̂_t] = E[g²]

At step t = 1 with β₂ = 0.999:

1 - β₂¹ = 1 - 0.999 = 0.001

Without bias correction, v₁ = 0.001 g₁², which would be 1,000 times too small, causing the denominator √(v₁) to artificially inflate the first step by √(1000) ≈ 31.6×! Dividing by 0.001 restores exact scale. As t → ∞, β^t → 0, so (1 - β^t) → 1.

Bias correction is strictly implemented in PyTorch (`torch.optim.Adam`) and JAX (`optax.adam`). Without it, early training steps either stall completely or blow up weights due to mismatched early scaling between m_t and v_t.

In 2014, Kingma and Ba proved this exact algebraic cancellation to show that their estimators are asymptotically unbiased minimum-variance trackers, outperforming previous heuristics that simply warmed up learning rates arbitrarily.

• Divide by (1 - β^t) to undo the zero drag: At step 1, you divide by 1 - β; as time steps march to infinity, you divide by 1.
• Finite geometric series cancellation: The (1 - β) terms cancel cleanly, leaving (1 - β^t) as the sole correction factor.

AdamW decouples weight decay from the gradient computation entirely:

w_(t+1) = (1 - α λ) w_t - (α) / (√(v̂_t) + ε) m̂_t

where λ is the true decoupled weight decay coefficient. This ensures that every parameter shrinks exponentially at rate (1 - α λ) regardless of gradient variance.

The mandatory default optimizer for training all foundation models: GPT-4, LLaMA-3, Mistral, and Stable Diffusion.

Introduced by Ilya Loshchilov and Frank Hutter (ICLR 2019) in 'Decoupled Weight Decay Regularization', rescuing transformer generalization.

Decouple decay from gradients: true AdamW applies clean exponential shrinkage to keep parameters regularized.