Optimizer Journey
Why did simple gradient descent fail on deep landscapes, and how did Adam solve it?
How this lesson unfolds.
Finding the lowest point on a high-dimensional loss surface is treacherous. Valleys have steep walls but gentle slopes, causing simple gradient descent to oscillate wildly or grind to a halt. Discover how adding momentum and adaptive per-parameter learning rates transformed optimization into the industry-standard Adam and AdamW algorithms.
Why did simple gradient descent fail on deep landscapes, and how did Adam solve it?
Picture rolling a heavy marble down a steep, narrow skateboard halfpipe that also slopes gently forward downhill. If you place a blind robot in the halfpipe that only takes steps proportional to the local slope (Vanilla SGD), the robot is thrown back and forth across the steep vertical walls, making almost zero forward progress. Now give the robot physical momentum: like a heavy bowling ball, its sideways oscillations cancel each other out over time, while its forward downhill speed accumulates momentum. Now add adaptive shock absorbers (RMSProp): if the terrain is violently steep in one direction, dial down step sizes along that axis; if the slope is nearly flat along another, dial up step sizes. Combining momentum and adaptive damping creates Adam: a smooth, confident roll straight to the minimum.
- Learning goalFormulate the update equations of Adam and AdamW, demonstrating how first moment m̂_t and second moment v̂_t resolve ill-conditioned Hessian curvature.
- Helpful foundationsMultivariate Gradients and Hessian Curvature · Adam: A Method for Stochastic Optimization (Kingma & Ba, 2014) · Decoupled Weight Decay Regularization / AdamW (Loshchilov & Hutter, 2017) · Optimization for Deep Learning (Ruder 2016 overview)
- Common trapConfusing L2 regularization with weight decay when using adaptive optimizers. In standard SGD, adding an L2 penalty 1/2λ \|θ\|² to the loss is mathematically identical to weight decay θ arrow (1 - η λ)θ. But in Adam, the L2 penalty gradient gets divided by √(v_t) + ε, meaning weights with large historical gradients receive almost no regularization, while weights with sparse gradients get penalized excessively! AdamW fixes this by decoupling weight decay directly from gradient updates.
- Next connectionHaving mastered fundamental neural architectures, backpropagation, vectorization, and optimization, we can now assemble full autoregressive models and explore how compute and data scale in Small Language Models.
More context & mechanics
- The Ill-Conditioned Ravine Challenge:
Let local loss curvature be approximated by Taylor expansion around minimum θ^*:
L(θ) ≈ L(θ^) + 1/2(θ - θ^)^T H (θ - θ^*)
where H = ∇² L is the Hessian matrix. If the condition number \kappa = (λ_(max)) / (λ_(min)) ≫ 1, vanilla SGD (θ_(t+1) = θ_t - η g_t) requires η < (2) / (λ_(max)) to prevent divergence, causing agonizingly slow progress along the flatter eigenvector corresponding to λ_(min).
- The Evolutionary Steps:
- Momentum (Polyak 1964):
m_t = β₁ m_(t-1) + (1-β₁) g_t, θ_(t+1) = θ_t - η m_t
Dampens high-frequency oscillations across steep ridges while accelerating along flat valleys.
- RMSProp (Hinton 2012):
v_t = β₂ v_(t-1) + (1-β₂) g_t², θ_(t+1) = θ_t - (η) / (√(v_t) + ε) \odot g_t
Normalizes step size by root mean square of recent gradients, providing per-coordinate learning rates.
- Adam (Kingma & Ba 2014):
Combines both moments with bias correction for initialization at m₀ = 0, v₀ = 0:
m̂_t = (m_t) / (1 - β₁^t), v̂_t = (v_t) / (1 - β₂^t)
θ_(t+1) = θ_t - (η) / (√(v̂_t) + ε) m̂_t
- AdamW (Loshchilov & Hutter 2017):
Decouples weight decay rate λ:
θ_(t+1) = θ_t - η λ θ_t - (η) / (√(v̂_t) + ε) m̂_t
AdamW with β₁ = 0.9, β₂ = 0.95 (or 0.98), ε = 10^(-8), and cosine learning rate schedules is the near-universal default optimizer for training modern foundation models including GPT-4, Llama 3, Gemini, and Claude.
Stochastic gradient descent was pioneered by Herbert Robbins and Sutton Monro in 1951. Boris Polyak introduced momentum in 1964. John Duchi et al. introduced AdaGrad in 2011. Diederik Kingma and Jimmy Ba presented Adam at ICLR 2015, which became one of the most cited computer science papers in history. Ilya Loshchilov and Frank Hutter introduced AdamW in 2017, solving generalization degradation in deep networks with weight decay.
- Golden Axiom: "Momentum remembers where you were going; RMSProp scales your speed by how bumpy the road is; AdamW decouples weight decay to keep weights lean."
- First moment m̂_t gives velocity direction; second moment v̂_t gives per-coordinate variance damping.
- Always use AdamW, not Adam, whenever using weight decay.