Training Deep Networks

Optimizers: momentum, RMSProp, and Adam

Smarter ways to step down the gradient

~12 minLesson 18 of 57Written by Claude Opus 4.8

Plain gradient descent is the engine; these are the transmission. The tricks that turn a marble that oscillates and stalls into one that rolls smoothly to the bottom.

In Lesson 11 you rolled a ball down a bowl with one rule: step against the gradient, scaled by a learning rate $η$ . It works - but on the lumpy, stretched-out loss surfaces of real networks that single rule oscillates, crawls, or blows up. Momentum, RMSProp, and Adam are three small edits to that update that make deep training actually converge. Every one is still gradient descent; each just decides how much to move, and in which directions, more cleverly.

The picture to hold in your head

Think of the loss as a landscape of hills and valleys and the parameters as a marble on it. Plain SGD is a marble with no mass: every step it teleports a little downhill, so in a narrow ravine it ping-pongs off the steep walls and barely creeps along the floor. The optimizers below give the marble momentum (so it coasts through small bumps) and a sense of scale (so it takes big steps where the ground is flat and tiny steps where it's steep). Same downhill instinct - better footwork.

One learning rate is never enough

The learning rate $η$ is the length of every step. Set it too large and the marble overshoots the valley and climbs the far wall higher each time - divergence. Set it too small and it inches along, needing millions of steps. But the deeper problem is that one number has to serve every direction at once. Real loss surfaces are ill-conditioned: steep along some axes, nearly flat along others. A step size safe for the steep walls is hopelessly timid on the flat floor - so plain SGD zig-zags across the ravine while barely advancing down it. You cannot fix this with a single $η$ ; you need the step to adapt.

The zig-zag, made precise

On a stretched quadratic bowl the two directions have different curvatures. Along a direction with curvature

c

, one SGD step multiplies the distance to the bottom by

(1 - η c)

. Pick

η

small enough that the steep direction (

c

large) stays stable and

(1 - η c)

is barely under 1 on the flat direction (

c

small) - so it crawls. Push

η

up to speed the flat direction and

∣1 - η c ∣

exceeds 1 on the steep one - so it explodes. There is no single

η

that is fast and stable. That gap is exactly what the next three ideas close.

Momentum: give the marble mass

Momentum keeps a running velocity $v$ - an exponentially decayed sum of past gradients - and moves the parameters along that velocity instead of along the raw gradient. Consistent gradients (the valley floor) accumulate into speed; alternating gradients (the ping-ponging walls) largely cancel. The marble coasts down the floor and stops sloshing side to side.

Momentum update

v_{t} = μ v_{t - 1} - η \nabla L (θ_{t - 1}), θ_{t} = θ_{t - 1} + v_{t} .

The momentum coefficient $μ$ (typically $0.9$ ) is how much of last step's velocity you keep. With $μ = 0$ this collapses back to plain SGD, $θ_{t} = θ_{t - 1} - η \nabla L$ . As $μ \to 1$ the marble grows heavier: in a steady downhill it approaches an effective step of $η / (1 - μ)$ - a $10 \times$ boost at $μ = 0.9$ - which is exactly how it powers across flat ground and rolls through small bumps.

RMSProp & Adam: a step size for every weight

Momentum fixes direction; it still uses one global $η$ . RMSProp fixes scale: it tracks a running average of each parameter's squared gradient $s$ , then divides the step by $s$ . Weights with large, persistent gradients (steep walls) get shrunk steps; weights with tiny gradients (flat floor) get amplified steps. Every parameter effectively gets its own, self-tuning learning rate.

Adam - the default optimizer almost everywhere - simply does both: momentum's velocity for direction (a 1st-moment estimate $m$ ) and RMSProp's per-parameter scaling (a 2nd-moment estimate $s$ ), plus a small bias correction so the very first steps aren't artificially tiny while the running averages warm up from zero.

Adam update (with bias correction)

m_{t} = β_{1} m_{t - 1} + (1 - β_{1}) g_{t}, s_{t} = β_{2} s_{t - 1} + (1 - β_{2}) g_{t}^{2}

\overset{m}{^}_{t} = \frac{m _{t}}{1 - β _{1}^{t}}, \overset{s}{^}_{t} = \frac{s _{t}}{1 - β _{2}^{t}}, θ_{t} = θ_{t - 1} - η \frac{m ^ _{t}}{s ^ _{t} + ϵ} .

Here $g_{t} = \nabla L$ , defaults $β_{1} = 0.9$ , $β_{2} = 0.999$ , $ϵ = 1 0^{- 8}$ . The $g_{t}^{2}$ is elementwise, so $\overset{s}{^}_{t}$ carries one number per weight. Drop the $m$ term and keep the divisor and you have RMSProp; drop the divisor and keep $m$ and you have momentum.

Roll a marble down the surface

Click the surface to drop the marble, then press Run to release it and trace its path. Switch the optimizer, drag the learning rate and momentum, and flip between a smooth bowl and a bumpy landscape with local dips. The start point is preserved as you change settings, so you can watch the same marble oscillate, diverge, get stuck, or glide home depending only on the algorithm.

Optimizer sandbox click to drop · Run · switch optimizer

trajectorymarbleglobal minimumdarker = higher loss

optimizersurface: convex bowl

learning rate η = 0.35momentum μ = 0.9

try a preset:

loss L(θ)-

gradient |∇L|-

steps0

Click the surface to drop the marble.

Schedules: warm up, then decay

Even a great optimizer wants a changing learning rate over training. Two moves dominate. Warmup: start

η

near zero and ramp it up over the first few hundred steps, so the moment estimates

\overset{m}{^}, \overset{s}{^}

stabilise before you take big steps - this alone tames the notoriously unstable opening of Transformer training. Decay: once warmed up, shrink

η

smoothly (cosine or linear) toward the end, so early steps explore boldly and late steps settle into the valley floor instead of rattling around it. In practice you tune the schedule as much as the optimizer.

# g = grad(loss, theta); lr, eps = 1e-8. states start at zero.

# --- SGD ---
theta = theta - lr * g

# --- SGD + momentum ---   (mu = 0.9)
v = mu * v - lr * g
theta = theta + v

# --- RMSProp ---          (rho = 0.9)
s = rho * s + (1 - rho) * g**2
theta = theta - lr * g / (sqrt(s) + eps)

# --- Adam ---             (b1 = 0.9, b2 = 0.999)
t = t + 1
m = b1 * m + (1 - b1) * g
s = b2 * s + (1 - b2) * g**2
m_hat = m / (1 - b1**t)
s_hat = s / (1 - b2**t)
theta = theta - lr * m_hat / (sqrt(s_hat) + eps)

One optimization step, four ways

Check yourself

Compared with plain SGD, what does momentum add to the update?

RMSProp divides each step by a running estimate of the...

Recall: what three ingredients does Adam blend, and what does each one buy you?

Try to state it, then check.

Lock it in

Momentum, RMSProp, and Adam are small edits to the SGD update that make deep training converge - each is still gradient descent, just deciding how much to move and in which directions more cleverly.
One global learning rate cannot serve an ill-conditioned surface (steep on some axes, flat on others), so plain SGD zig-zags: no single $η$ is fast and stable.
Momentum keeps a velocity (a decayed sum of gradients): consistent gradients accumulate into speed and oscillating ones cancel; at $μ = 0.9$ the effective step is roughly $10 \times$ .
RMSProp divides the step by $s$ from the running squared-gradient average, giving each weight its own step size; Adam does both plus a bias correction so early steps are not artificially tiny.
Even a great optimizer wants a schedule: warmup steadies the unstable opening, then decay (cosine or linear) settles into the valley - AdamW is what pretrains essentially every large language model.

Primary source

Sebastian Ruder's An overview of gradient descent optimization algorithms derives momentum, RMSProp, Adam, and the rest side by side with the exact update equations used above. For the geometric intuition of a marble descending a loss surface, 3Blue1Brown's Neural networks series is the clearest visual companion anywhere.

Ask your teacher

Here is the connection worth holding onto: AdamW - Adam with its weight decay decoupled from the gradient - is the optimizer that pretrains essentially every large language model you have heard of (the LLMs of Lesson 50). The reason is exactly the picture in this demo: an LLM's loss landscape is astronomically bumpy and ill-conditioned, and plain gradient descent (Lesson 11) would ping-pong and stall in it forever. Momentum rolls through those bumps; per-parameter scaling handles the wildly different curvatures; warmup steadies the opening. Ask me about why AdamW's decoupled decay beats classic L2 for Transformers, what a cosine schedule with warmup actually looks like, or when SGD-with-momentum still wins over Adam.