Gradient Descent

Training a model means walking downhill on its loss surface. Drop three optimizers at the same point and watch how plain gradient descent, momentum and Adam find their way to a minimum, or fail to.

Intermediate

Loss surface

f(x, y) = 0.1x² + 2y²

Curvatures 0.2 (along x) and 4 (along y). GD is stable only for η < 2/4 = 0.5, a step that barely moves along x: watch it zig-zag across the valley.

Click or tap anywhere on the plot to start all three optimizers from that point. Darker colours mean lower loss; white crosses mark the minima.

Gradient descentstep 0, f = 5.2290
Momentum (β = 0.9)step 0, f = 5.2290
Adamstep 0, f = 5.2290

Loss versus step

Plotted as log₁₀(f − f*), where f* = 0.000 is the lowest minimum. Each gridline is a factor of 10. A curve that flattens above 0 is stuck in a worse local minimum; one that shoots off the top has diverged.

-8-6-4-202450 steps0

How to use this visualizer

  1. 1
    Choose a loss surface. Dark regions have low loss; white crosses mark the minima.
  2. 2
    Click or tap anywhere on the contour plot to start all three optimizers from that point.
  3. 3
    Press Run (or Step) and compare the red (GD), blue (momentum) and green (Adam) trails.
  4. 4
    Drag the learning-rate slider. On the bowl and the valley, push η beyond 2 divided by the curvature and watch GD diverge.
  5. 5
    Read the loss chart: it shows log₁₀(f − f*), so a straight downward line means the error shrinks by a constant factor every step.

What gradient descent does

The gradient ∇f points in the direction in which the loss increases fastest. Gradient descent takes a step of size proportional to the learning rate in the opposite direction, then repeats. On a smooth bowl this always reaches the minimum, provided the steps are small enough.

Real loss surfaces are rarely round bowls. They have long narrow valleys, where the gradient points mostly across the valley instead of along it, and several minima of different quality. Momentum and Adam are two widely used fixes: momentum remembers past gradients, and Adam rescales each parameter's step by how large its gradients have been.

Learning rate too small

Every step is safe but tiny. On the Rosenbrock surface plain GD is still far from (1, 1) after 2000 steps.

Learning rate too large

On the valley, raise η past 0.5 and GD overshoots further each step along y until it flies off the plot.

Momentum

Accumulates velocity along directions where the gradient keeps pointing the same way, and cancels zig-zags across a valley.

Adam

Divides each coordinate by a running RMS of its gradient, so steep and flat directions both take steps of about α.

Key formulas

Plain gradient descent

θt+1 = θt − η ∇f(θt)

θ
parameters (here the point (x, y))
η
learning rate

Momentum (heavy ball)

vt+1 = β vt − η ∇f(θt), θt+1 = θt + vt+1

β
momentum coefficient, 0.9 here

Adam

mt = β1mt−1 + (1 − β1)gt, st = β2st−1 + (1 − β2)gt2

θt+1 = θt − α m̂t / (√ŝt + ε)

β₁, β₂
0.9 and 0.999
m̂, ŝ
bias-corrected moments: m / (1 − β₁ᵗ), s / (1 − β₂ᵗ)
ε
10⁻⁸, avoids division by zero

When does gradient descent diverge?

On a quadratic with curvatures λ1, …, λn (the eigenvalues of the Hessian), each step multiplies the error along direction i by (1 − ηλi). GD converges only if |1 − ηλi| < 1 for every i, that is η < 2 / λmax. The steepest direction sets the limit, and the flattest direction then converges slowly: the ratio λmax / λmin (the condition number) measures how hard the problem is.

Why Adam has its own step size

Adam's update is roughly α in each coordinate regardless of how steep the surface is, while GD's step is η times the gradient. The same number means very different step lengths for the two, so the visualizer gives Adam a separate slider.

In code

Python
import numpy as np

def grad(p):                        # elongated valley f = 0.1x² + 2y²
    return np.array([0.2 * p[0], 4 * p[1]])

def gd(p, lr=0.4, steps=100):
    for _ in range(steps):
        p = p - lr * grad(p)
    return p

def momentum(p, lr=0.4, beta=0.9, steps=100):
    v = np.zeros_like(p)
    for _ in range(steps):
        v = beta * v - lr * grad(p)
        p = p + v
    return p

def adam(p, lr=0.1, b1=0.9, b2=0.999, eps=1e-8, steps=100):
    m, s = np.zeros_like(p), np.zeros_like(p)
    for t in range(1, steps + 1):
        g = grad(p)
        m = b1 * m + (1 - b1) * g
        s = b2 * s + (1 - b2) * g * g
        m_hat, s_hat = m / (1 - b1 ** t), s / (1 - b2 ** t)
        p = p - lr * m_hat / (np.sqrt(s_hat) + eps)
    return p

Practice problems

Problem 1: Two steps by hand

Minimise f(x, y) = x² + y² with plain gradient descent, starting at (2, −1) with η = 0.1. Where are you after two steps, and by what factor does f shrink per step?
Show solution

∇f = (2x, 2y), so θ ← θ − 0.1 × 2θ = 0.8 θ

Step 1: (2, −1) × 0.8 = (1.6, −0.8)

Step 2: (1.6, −0.8) × 0.8 = (1.28, −0.64)

f scales with the square of the distance: 0.8² = 0.64 per step

θ₂ = (1.28, −0.64); f drops from 5 to 2.048 (5 × 0.64²)

Problem 2: The largest stable learning rate

For f(x, y) = 0.1x² + 2y², what is the largest learning rate for which plain gradient descent converges? At that limit, roughly how many steps does it take to shrink the x-error by a factor of 10?
Show solution

Hessian = diag(0.2, 4), so λmax = 4 and λmin = 0.2

Stable when η < 2 / λmax = 2 / 4 = 0.5

Along x each step multiplies the error by 1 − ηλmin = 1 − 0.5 × 0.2 = 0.9

0.9n = 0.1 gives n = ln 0.1 / ln 0.9 ≈ 21.9

η < 0.5; about 22 steps per factor of 10 along x

Discussion

Questions, ideas or a different approach? Join the conversation.