Smarter Steps: Momentum, RMSProp and Adam

Plain gradient descent zigzags across steep valleys and crawls along flat ones. Three small changes to the update rule, remembering past gradients and scaling each knob's step, fix most of that, and one of them, Adam, is what most networks are trained with today.

49 min Prerequisites: Chapters 7 and 9 (the Tensor class saved as tensor.py)

Inspect Architecture: Multi-layer perceptron

1. The idea

Chapter 7 left us with a lopsided bowl. On the raw ride data the loss is 263 times steeper in one direction than the other, so the step size has to be small enough not to overshoot the steep direction, and then it crawls along the flat one. The best plain gradient descent could do from w = b = 0 was 450 steps. Standardizing the input fixed that particular bowl, but a real network has millions of knobs and no rescaling trick makes all their directions equally steep at once. So people changed the update rule instead.

The first change is momentum. Instead of stepping along today's gradient, keep a running velocity that adds each new gradient to a fraction of the old velocity, like a heavy ball rolling down the bowl. Along the steep direction the gradients keep flipping sign, so they cancel out in the velocity. Along the flat direction they all point the same way, so they add up. With β = 0.9 the ride model settles in 111 steps instead of 450. Theory gives the best possible momentum settings for a bowl like this, which settle in 89 steps, but they turn out to sit right at the edge of stability: round them slightly and it takes 170, lower β a little and the run blows up.

The second change, RMSProp, gives every knob its own step size. It keeps a running average of each knob's squared gradients and divides that knob's step by the square root. A knob whose gradients are huge gets small steps, a knob whose gradients are tiny gets big ones, so the lopsidedness largely disappears. Its weakness is that near the minimum its steps stay about η in size however small the gradient gets, so it never quite settles.

Adam combines the two: momentum's running average of gradients, divided by RMSProp's running scale, with a small correction for both averages starting at zero. Its first step moves every knob by almost exactly η, whatever the size of its gradient, and multiplying the whole loss by 1,000 doesn't change its path at all, while plain gradient descent would need a 1,000 times smaller learning rate. On the ride bowl Adam settles in 105 steps; on Chapter 9's activity network with mini-batches it reaches the lowest loss of the four. No optimizer removes the need to try a few learning rates, though: at η = 0.3 Adam does worse than plain gradient descent.

The same bowl as Chapter 7, four ways down. Momentum and Adam overshoot but arrive far sooner; plain gradient descent wastes its steps bouncing across the steep direction.

2. The math

The problem: one step size for a lopsided bowl

raw km, from (0, 0):  curvatures 120.04 and 0.456,  κ = 263
best plain step η = 0.016:  settles within 1% of L* after 450 steps
per step:  steep axis × (−0.92),   flat axis × 0.9927

These are Chapter 7's numbers. The steep direction overshoots back and forth while the flat direction shrinks by less than 1% per step.

Momentum: a running velocity

v ← β·v + g        θ ← θ − η·v
v = g_t + β·g_(t−1) + β²·g_(t−2) + …    constant g:  v → g / (1 − β)

β = 0.9 is the usual choice: the velocity remembers roughly the last 1 / (1 − β) = 10 gradients, and a steady gradient gets a 10× longer stride. Gradients that flip sign cancel instead.

Momentum on the ride bowl

β = 0.9, η = 0.01:   first within 1% at step 64,  settled at 111
stable only for  η < 2(1 + β) / λ_max = 3.8 / 120.04 = 0.0317
β = 0.9, η = 0.033:  blows up

Momentum allows a larger stable step than plain descent (0.0317 against 0.0167), but it overshoots on the way, which is why the loss oscillates before settling. The limit is derived in the first derivation.

The best momentum in theory, and why it's fragile

η* = 4 / (√λmax + √λmin)² = 0.0296     β* = ((√κ − 1) / (√κ + 1))² = 0.781
error per step:  × 0.884   (plain descent: × 0.9924)
exact η*, β*: 89 steps    rounded to 0.0296, 0.781: 170    β = 0.75: blows up

Momentum's rate depends on √κ instead of κ, which is a huge gain for badly shaped bowls. But η*·λmax = 3.549 sits just under the limit 2(1 + β*) = 3.562, so the optimum is a knife edge. In practice people use β = 0.9 and tune η.

RMSProp: a step size per knob

s ← ρ·s + (1 − ρ)·g²        θ ← θ − η · g / (√s + ε)
ρ = 0.9,  ε = 10⁻⁸ (only there to avoid dividing by zero)

√s is a running root-mean-square of the knob's recent gradients, so g / √s is about ±1 for every knob. w's gradients are about 8 times b's at the start, but both take steps of about η.

RMSProp's catch: it never quite stops

near the minimum, g shrinks but so does √s:  the step stays ≈ η
η = 0.01:  within 1% at step 543, but ends at 0.6662, not 0.6644
η = 0.03:  within 1% at step 232, then bounces; ends at 0.681

Dividing by the gradient's own size throws away the information that the gradient is getting small. The usual cure is to shrink η over the course of training (Chapter 7's second exercise).

Adam: momentum divided by RMSProp's scale

m ← β₁·m + (1 − β₁)·g        v ← β₂·v + (1 − β₂)·g²
m̂ = m / (1 − β₁ᵗ)        v̂ = v / (1 − β₂ᵗ)
θ ← θ − η · m̂ / (√v̂ + ε)        defaults:  β₁ = 0.9,  β₂ = 0.999

m is a smoothed gradient (momentum), v a smoothed squared gradient (RMSProp). t counts steps. The hats undo the bias from starting both averages at zero, derived in the second derivation.

Adam's first step, and what it ignores

step 1:  m̂ = g,  v̂ = g²   ⇒   step = η · g / |g| = η · sign(g)
rides from (0, 0):  g = (−427.5, −51)  ⇒  (w, b) = (1.0, 1.0) after one step
loss × 1000:  Adam (η = 1.0) still settles at step 105;  SGD needs η ÷ 1000

Every knob moves by η at first, whatever the size of its gradient. Scaling the gradient doesn't change Adam's path, but rescaling a knob's units does (measure rides in metres and w must be about 0.003, which steps of size 1 overshoot).

The ride bowl, all four

settled within 1% of L* (from w = b = 0):
plain descent (η = 0.016)  450      momentum (β = 0.9, η = 0.01)  111
RMSProp (η = 0.01)  543, hovering      Adam (η = 1.0)  105

On a perfect quadratic bowl momentum is hard to beat. Adaptive methods earn their keep when knobs live on very different scales and gradients are noisy, which is the normal situation in a network.

Mini-batches: Chapter 9's activity network

16 hidden units, batches of 16, 30 epochs, mean of 3 random starts
best η for each:  SGD 0.279 (η = 1.0)    momentum 0.203 (η = 0.1)
                  RMSProp 0.268 (η = 0.1)    Adam 0.192 (η = 0.1)
too large:  SGD at η = 3.0 → 1.047    Adam at η = 0.3 → 0.468

Full-data loss after training. Adam and momentum win, but every optimizer has a learning rate that breaks it. A good default is Adam with η around 0.001 to 0.01 for big networks, then try a few values either side.

• Take one axis of the bowl with curvature λ, so the gradient is g = λ·θ (measuring θ from the minimum). Write the momentum step without v: since η·v_(t−1) = θ_(t−1) − θ_t, the update becomes θ_(t+1) = θ_t − η·λ·θ_t + β·(θ_t − θ_(t−1)).

• That's a recurrence with two terms of memory. Try θ_t = rᵗ: it works when r² − (1 + β − ηλ)·r + β = 0. The run converges exactly when both roots have |r| < 1.

• For a quadratic r² + a·r + c, both roots lie inside the unit circle when |c| < 1 and |a| < 1 + c. Here c = β, which is below 1, and the second condition reads |1 + β − ηλ| < 1 + β, that is 0 < ηλ < 2(1 + β).

• So with β = 0.9 the steep axis needs η < 3.8 / 120.04 = 0.0317, nearly double plain descent's limit (β = 0 gives back Chapter 7's 2/λ). The code confirms it: η = 0.033 blows up. ∎

• When the roots are complex, |r| = √β, whatever η is. The theoretical optimum chooses β so that both axes shrink at that same rate, which gives the per-step factor √β* = (√κ − 1)/(√κ + 1) = 0.884, and η* puts the steep axis right next to the limit.

• Unroll the first average from m₀ = 0: m_t = (1 − β₁)·(g_t + β₁·g_(t−1) + … + β₁^(t−1)·g₁). The weights are a geometric series.

• If the gradient were constant, g_k = g, the sum is (1 − β₁)·g·(1 + β₁ + … + β₁^(t−1)) = (1 − β₁ᵗ)·g. So m_t underestimates g by the factor (1 − β₁ᵗ), badly at the start: after one step m₁ = 0.1·g.

• Dividing by (1 − β₁ᵗ) removes exactly that factor: m̂_t = g. The same argument with β₂ gives v̂_t = g². As t grows, βᵗ → 0 and the correction fades away.

• The correction matters most for v, because β₂ = 0.999 makes the bias last for thousands of steps. Without it the early √v would be far too small, and the first steps far too large. ∎

• At t = 1 the corrected values are exactly m̂ = g and v̂ = g², so the step is η·g/|g| = η·sign(g): the code shows w and b both moving from 0 to 1.0 with η = 1.

3. How it works

Keep state per knob

Each optimizer stores arrays of the same shape as each parameter: a velocity for momentum, a running square for RMSProp, both for Adam. They start at zero.

Same loop, new step

The training loop doesn't change: forward, zero the grads, backward, then opt.step(). Only the rule that turns a gradient into a change of the knobs is different.

Start with Adam

β₁ = 0.9 and β₂ = 0.999 rarely need changing. Try η = 0.001, 0.003 and 0.01 on a network, or larger values for tiny problems like this chapter's.

Watch for oscillation

Momentum and Adam overshoot before they settle, so a loss that wiggles early is normal. A loss that grows steadily means η is too large.

Shrink η at the end

Adaptive methods keep taking steps of about η near the minimum. Lowering η late in training, by hand or on a schedule, lets them settle.

4. The code (python)

core_ch10.py

# Chapter 10: better ways to step downhill. Momentum, RMSProp and Adam, as small classes.
# Uses the Tensor class from Chapter 6, saved as tensor.py (see Chapter 7).
import numpy as np
from tensor import Tensor


class SGD:
    """Chapter 7's step: θ ← θ − η·g."""
    def __init__(self, params, lr):
        self.params, self.lr = params, lr

    def step(self):
        for p in self.params:
            p.data = p.data - self.lr * p.grad


class Momentum:
    """Keep a running velocity: v ← β·v + g, then θ ← θ − η·v."""
    def __init__(self, params, lr, beta=0.9):
        self.params, self.lr, self.beta = params, lr, beta
        self.v = [np.zeros_like(p.data) for p in params]

    def step(self):
        for p, v in zip(self.params, self.v):
            v *= self.beta
            v += p.grad
            p.data = p.data - self.lr * v


class RMSProp:
    """Divide each knob's step by a running root-mean-square of its own gradients."""
    def __init__(self, params, lr, rho=0.9, eps=1e-8):
        self.params, self.lr, self.rho, self.eps = params, lr, rho, eps
        self.s = [np.zeros_like(p.data) for p in params]

    def step(self):
        for p, s in zip(self.params, self.s):
            s *= self.rho
            s += (1 - self.rho) * p.grad ** 2
            p.data = p.data - self.lr * p.grad / (np.sqrt(s) + self.eps)


class Adam:
    """Momentum and RMSProp together, with both averages corrected for starting at zero."""
    def __init__(self, params, lr, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params, self.lr, self.b1, self.b2, self.eps = params, lr, beta1, beta2, eps
        self.m = [np.zeros_like(p.data) for p in params]
        self.v = [np.zeros_like(p.data) for p in params]
        self.t = 0

    def step(self):
        self.t += 1
        for p, m, v in zip(self.params, self.m, self.v):
            m *= self.b1
            m += (1 - self.b1) * p.grad
            v *= self.b2
            v += (1 - self.b2) * p.grad ** 2
            m_hat = m / (1 - self.b1 ** self.t)          # bias correction
            v_hat = v / (1 - self.b2 ** self.t)
            p.data = p.data - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)


# === 1. Chapter 7's lopsided bowl: the four rides, raw km, from w = b = 0 ==========
km = np.array([2.0, 5.0, 8.0, 12.0])
minutes = np.array([11.0, 21.0, 28.0, 42.0])
BEST = 0.6643835616438343                          # the closed-form minimum from Chapter 7


def fit_rides(make_opt, steps=3000, loss_scale=1.0):
    w, b = Tensor(0.0), Tensor(0.0)
    opt = make_opt([w, b])
    history = []
    for _ in range(steps):
        loss = ((w * km + b - minutes) ** 2).mean() * loss_scale
        history.append(float(loss.data) / loss_scale)
        if not history[-1] < 1e12:
            return None, history, (w, b)                  # blew up
        for p in (w, b):
            p.grad = np.zeros_like(p.data)
        loss.backward()
        opt.step()
    return settled(history), history, (w, b)


def settled(history):
    """The first step after which the loss stays within 1% of the best possible."""
    good = [h < 1.01 * BEST for h in history]
    if not good[-1]:
        return None
    i = len(good) - 1
    while i > 0 and good[i - 1]:
        i -= 1
    return i


# The bowl's curvatures (Chapter 7), and the momentum settings theory calls best for them
X = np.column_stack([km, np.ones(4)])
lam_min, lam_max = np.linalg.eigvalsh(2 / len(km) * X.T @ X)
kappa = lam_max / lam_min
eta_hb = 4 / (np.sqrt(lam_max) + np.sqrt(lam_min)) ** 2
beta_hb = ((np.sqrt(kappa) - 1) / (np.sqrt(kappa) + 1)) ** 2
print(f"heavy ball: η* = 4/(√λmax + √λmin)² = {eta_hb:.4f},  β* = ((√κ − 1)/(√κ + 1))² = {beta_hb:.3f}")
print(f"  error shrinks per step by {(np.sqrt(kappa) - 1) / (np.sqrt(kappa) + 1):.3f} (plain descent: {(kappa - 1) / (kappa + 1):.4f})")
print(f"  steep direction: η*·λmax = {eta_hb * lam_max:.3f}, stability limit 2(1 + β*) = {2 * (1 + beta_hb):.3f}")
print(f"momentum at β = 0.9 is stable for η < 2(1 + β)/λmax = {2 * 1.9 / lam_max:.4f}\n")

runs = [("plain gradient descent, η = 0.016", lambda ps: SGD(ps, 0.016)),
        ("momentum β = 0.9, η = 0.01", lambda ps: Momentum(ps, 0.01)),
        ("momentum β = 0.9, η = 0.033", lambda ps: Momentum(ps, 0.033)),
        ("heavy ball, exact η*, β*", lambda ps: Momentum(ps, eta_hb, beta=beta_hb)),
        ("heavy ball rounded, η = 0.0296, β = 0.781", lambda ps: Momentum(ps, 0.0296, beta=0.781)),
        ("heavy ball with β = 0.75", lambda ps: Momentum(ps, eta_hb, beta=0.75)),
        ("RMSProp, η = 0.01", lambda ps: RMSProp(ps, 0.01)),
        ("RMSProp, η = 0.03", lambda ps: RMSProp(ps, 0.03)),
        ("Adam, η = 1.0", lambda ps: Adam(ps, 1.0))]
print("optimizer                                 first within 1%   settled   final loss")
for name, make in runs:
    when, hist, _ = fit_rides(make)
    first = next((i for i, h in enumerate(hist) if h < 1.01 * BEST), None)
    final = "blew up" if hist[-1] >= 1e12 else round(hist[-1], 4)
    print(f"{name:42} {str(first):>8}       {str(when):>8}    {final}")

# Adam's first step moves every knob by about η, whatever the size of its gradient
w, b = Tensor(0.0), Tensor(0.0)
loss = ((w * km + b - minutes) ** 2).mean()
loss.backward()
print(f"\nfirst gradients: dw = {float(w.grad)}, db = {float(b.grad)}")
opt = Adam([w, b], 1.0)
opt.step()
print(f"after one Adam step with η = 1: w = {float(w.data):.4f}, b = {float(b.data):.4f}")

# Scale the loss by 1,000: SGD needs a new learning rate, Adam doesn't
for name, make in [("SGD, η = 0.016", lambda ps: SGD(ps, 0.016)),
                   ("SGD, η = 0.000016", lambda ps: SGD(ps, 0.000016)),
                   ("Adam, η = 1.0", lambda ps: Adam(ps, 1.0))]:
    when, hist, _ = fit_rides(make, loss_scale=1000.0)
    print(f"loss × 1000, {name:18}: settled at {when}" + ("  (blew up)" if hist[-1] >= 1e12 else ""))

# === 2. Mini-batches: Chapter 9's activity classifier with 16 hidden units ==========
rng = np.random.default_rng(9)
X = np.round(rng.random((150, 3)), 2)
labels = np.where(X[:, 1] > 0.5, 2, np.where(X[:, 2] > 0.6, 1, 0))
noisy = rng.random(150) < 0.1
labels[noisy] = rng.integers(0, 3, noisy.sum())
Y = np.eye(3)[labels]


def softmax_cross_entropy(Z, Yb):
    E = (Z + (-Z.data.max(axis=1, keepdims=True))).exp()
    P = E / (E @ Tensor(np.ones((3, 1))))
    return -(Tensor(Yb) * P.log()).sum() / Z.shape[0]


def train_mlp(make_opt, seed, epochs=30, batch=16):
    r = np.random.default_rng(seed)
    W1, b1 = Tensor(r.normal(size=(16, 3))), Tensor(np.zeros(16))
    W2, b2 = Tensor(r.normal(size=(3, 16))), Tensor(np.zeros(3))
    params = [W1, b1, W2, b2]
    opt = make_opt(params)
    shuffle = np.random.default_rng(1)
    for _ in range(epochs):
        order = shuffle.permutation(150)
        for start in range(0, 150, batch):
            rows = order[start:start + batch]
            loss = softmax_cross_entropy((Tensor(X[rows]) @ W1.T + b1).relu() @ W2.T + b2, Y[rows])
            for p in params:
                p.grad = np.zeros_like(p.data)
            loss.backward()
            opt.step()
    return float(softmax_cross_entropy((Tensor(X) @ W1.T + b1).relu() @ W2.T + b2, Y).data)


print("\n150 days, batches of 16, 30 epochs: full-data loss, mean of 3 random starts")
grid = [("SGD", SGD, [0.1, 0.3, 1.0, 3.0]), ("momentum", Momentum, [0.01, 0.03, 0.1, 0.3]),
        ("RMSProp", RMSProp, [0.003, 0.01, 0.03, 0.1]), ("Adam", Adam, [0.003, 0.01, 0.03, 0.1, 0.3])]
for name, cls, lrs in grid:
    cells = []
    for lr in lrs:
        mean = np.mean([train_mlp(lambda ps: cls(ps, lr), seed) for seed in range(3)])
        cells.append(f"η={lr}: {mean:.3f}")
    print(f"  {name:9}" + "   ".join(cells))

5. Practice

Work these out on paper (or in Python) and type the number. Answers are checked with a small tolerance for rounding.

P1 With momentum β = 0.9 and a gradient that stays constant at g, what does the velocity v approach, as a multiple of g?

P2 Momentum starts with v = 0. The first two gradients are 4 and then −4. With β = 0.9, what is v after the second step?

P3 For momentum with β = 0.9 on an axis with curvature λ = 50, what is the largest stable learning rate?

P4 RMSProp with ρ = 0.9 starts with s = 0. The first gradient is g = 3. What is s after the update?

P5 Adam with β₁ = 0.9 starts with m = 0. After one step with gradient g = 5, what is m (before the correction)?

P6 Adam's first step from (0, 0) on the rides, with η = 0.5. The gradients are −427.5 for w and −51 for b. What is b after the step?

P7 A curvature ratio κ = 100. What factor does plain descent's error shrink by per step at best, (κ − 1)/(κ + 1)? (4 decimal places.)

P8 With β₂ = 0.999, what is the bias-correction factor 1 − β₂ᵗ after t = 1,000 steps? (3 decimal places.)

6. Go further

Add Nesterov momentum: compute the gradient at the look-ahead point θ − η·β·v instead of at θ, then update v and θ as before. (With the engine, set the knobs to the look-ahead point, run forward and backward, then put them back.) Does it settle faster than plain momentum on the ride bowl at η = 0.01, and does its stability limit change?
Give RMSProp a decaying learning rate, η_t = 0.03 / (1 + t/100), and rerun it on the ride bowl. Does it now settle at the minimum? Try the same schedule on Adam at η = 3.0, which reaches the 1% band at step 35 but takes thousands of steps to stay there.
AdamW changes one line: after the Adam step, also shrink every weight a little, θ ← θ − η·λ·θ with λ = 0.01. Add it to the Adam class as an option and train the activity network. What happens to the size of the weights and to the loss? Chapter 11 explains why shrinking weights can help.

7. Check yourself

Answer all 5 questions correctly to complete the chapter · 0 / 5 done

Q1/5 Why does momentum speed up gradient descent on a long, narrow valley?

Q2/5 What does RMSProp divide each knob's gradient by?

Q3/5 Why does Adam divide m by (1 − β₁ᵗ)?

Q4/5 You multiply the loss by 1,000. What happens to plain gradient descent and to Adam with the same learning rates as before?

Q5/5 The theoretically best momentum settings for the ride bowl settle in 89 steps. Why don't people simply use them?

In the catalog

Finished this chapter?

Learn machine learning from scratch