Mini Project: Gradient Descent Visualizer

A training run has been going for forty minutes. The loss drops from 2.4 to 0.9, wobbles, jumps to 6.1, then prints nan and every number after it is nan too. You look at the loss curve. It shows you exactly when the run died. It tells you nothing about why.

The quieter version: nothing crashes, the loss falls fast for two hundred steps, then flattens at 0.31 and stays there for fifty thousand more. Is that a local minimum? A learning rate too small to make progress? Or one so large the optimiser is bouncing between two walls of a narrow valley? The loss curve looks identical in all three cases: a line that goes flat.

The loss curve hides the reason because it throws almost everything away. At each step the optimiser holds a position in parameter space — possibly millions of numbers — and the curve collapses that whole position into one scalar. You are watching a shadow and trying to infer the shape of the thing casting it.

In two dimensions you do not have to infer anything. You can draw the whole loss surface as a contour map and draw the optimiser's actual path across it, step by step. Divergence looks like steps getting longer instead of shorter. A stall looks different from a slow creep. Zig-zagging looks like zig-zagging. This project builds that tool: a gradient descent visualiser for 2D functions, with sliders for learning rate and momentum, so the behaviours you will later have to diagnose blind become things you have actually watched.

The one line you are going to watch

Everything in this project revolves around a single update rule:

Read it piece by piece. $\theta_t$ (theta) is where you currently stand — in this project a pair of numbers $(x, y)$, in real training the full vector of model weights. $f$ is the function you are minimising: the loss. $\nabla f(\theta_t)$ (nabla, "the gradient") is the vector of partial derivatives of $f$ evaluated at your current position; each component says how fast $f$ changes if you nudge that one coordinate. $\eta$ (eta) is the learning rate: a single positive number controlling how far you move.

The minus sign is the whole idea. The gradient points in the direction of steepest increase — the way uphill. You want to go down. So you take the gradient, flip it, scale it by $\eta$, and step. If you wrote a plus sign there you would climb the surface with perfect efficiency, which is a fine algorithm for a completely different problem.

The gradient tells you which way is uphill and how steep. The learning rate decides how much you trust that information before you look again.

Here is the core function. Note what it returns.

1import numpy as np23def gradient_descent(f, grad, theta0, lr, n_steps):4    """f(x, y) -> scalar; grad(x, y) -> array([df/dx, df/dy]); lr is eta."""5    theta = np.array(theta0, dtype=float)6    path = [theta.copy()]7    values = [f(*theta)]89    for _ in range(n_steps):10        g = grad(*theta)11        theta = theta - lr * g          # the update rule, verbatim12        val = f(*theta)13        path.append(theta.copy())14        values.append(val)15        if not np.isfinite(val):        # blew up; stop before the array fills with nan16            break1718    return np.array(path), np.array(values)

Most optimiser code returns only the final answer, because in production that is all anyone wants. Here the final answer is the least interesting output. The path array — shape (n_steps+1, 2), one row per position visited — is the point of the project, and it is what you plot. Two runs can end at the same place by completely different journeys, and the journey is what tells you whether your learning rate is any good.

theta.copy() matters because NumPy arrays are mutable, and appending the same object repeatedly gives you a list of identical references. The isfinite check is not defensive padding: when a run diverges you want the last few real values kept so you can see the explosion, not a thousand rows of nan.

Four surfaces, four different ways to fail

A single test function teaches you almost nothing, because a well-behaved function makes every learning rate look reasonable. Use four, each engineered to break the optimiser in a different way.

You need the gradient of each, by hand, because you are going to check your code against it. Differentiate each term with respect to $x$ holding $y$ fixed, then the reverse:

1def bowl(x, y):        return x**2 + y**22def bowl_grad(x, y):   return np.array([2*x, 2*y])34def valley(x, y):      return x**2 + 10*y**25def valley_grad(x, y): return np.array([2*x, 20*y])67def rosen(x, y):       return (1 - x)**2 + 100*(y - x**2)**28def rosen_grad(x, y):9    return np.array([-2*(1 - x) - 400*x*(y - x**2), 200*(y - x**2)])1011def multi(x, y):       return np.sin(x)*np.cos(y) + 0.1*(x**2 + y**2)12def multi_grad(x, y):13    return np.array([np.cos(x)*np.cos(y) + 0.2*x, -np.sin(x)*np.sin(y) + 0.2*y])

The elongated bowl is the one whose lesson transfers, so it deserves a note now. Its curvature is $2$ along $x$ and $20$ along $y$: the surface is ten times steeper across the valley than along it. That ratio, $20/2 = 10$, is the condition number. A learning rate small enough to be safe in the steep direction is ten times smaller than the flat direction needs, so you creep along the valley floor while bouncing across it. Every adaptive optimiser you have heard of — momentum, RMSprop, Adam — exists to fix this.

The multi-modal function makes a different point. From $(-2,-2)$ it settles at $(-1.31, 0)$ with $f = -0.79$; from $(2,2)$ it settles at $(1.27, 2.57)$ with $f = +0.02$. Both loss curves fall smoothly and flatten, and both look like success. One found a far worse value and gave no signal that it had.

Momentum: giving the ball some weight

Plain gradient descent has no memory. Each step consults the gradient afresh, which is why it can spend a hundred steps flipping between two sides of a valley — the gradient genuinely does point across the valley each time. Momentum adds a velocity term that accumulates:

$v$ is the velocity, initialised to zero. $\beta$ (beta), typically $0.9$, is how much of the previous velocity survives, which makes $v$ roughly a running sum of the last ten gradients. Components that keep pointing the same way accumulate and the step grows; components that flip sign every step cancel against their own history and shrink. That is exactly the valley-floor versus zig-zag distinction, and momentum separates them for free.

1def gradient_descent(f, grad, theta0, lr, n_steps, beta=0.0):2    theta = np.array(theta0, dtype=float)3    v = np.zeros_like(theta)4    path, values = [theta.copy()], [f(*theta)]56    for _ in range(n_steps):7        v = beta * v + grad(*theta)     # beta=0 gives plain gradient descent back8        theta = theta - lr * v9        val = f(*theta)10        path.append(theta.copy()); values.append(val)11        if not np.isfinite(val):12            break1314    return np.array(path), np.array(values)

With beta=0 this reduces exactly to the original, so one function covers both cases and there are never two code paths to keep in sync.

Momentum is not free, and the visualiser shows the cost. On the round bowl $x^2+y^2$ from $(3,4)$ with $\eta = 0.1$, plain descent walks straight in and reaches $f < 10^{-6}$ in 39 steps. Add $\beta = 0.9$ and the path overshoots entirely — by step 6 it is at $(-2.13, -2.84)$, on the far side and nearly as far out as it started — then spirals back in over 119 steps. The inertia that rescues you in a valley carries you straight past a minimum that was easy to hit.

Momentum helps when consecutive gradients disagree and hurts when they already agree.

Drawing it: contours first, 3D second

Evaluate the function on a grid, draw filled contours, plot the path on top.

1import matplotlib.pyplot as plt23def plot_path(f, path, xlim=(-6, 6), ylim=(-6, 6), levels=30, minimum=(0, 0)):4    X, Y = np.meshgrid(np.linspace(*xlim, 400), np.linspace(*ylim, 400))5    Z = f(X, Y)                         # vectorised: one call, no Python loop67    fig, ax = plt.subplots(figsize=(6, 5))8    cs = ax.contourf(X, Y, Z, levels=levels, cmap='viridis', alpha=0.8)9    fig.colorbar(cs, ax=ax, label='f(x, y)')10    ax.plot(path[:, 0], path[:, 1], 'o-', color='red', ms=3, lw=1.2, label='path')11    ax.plot(*path[0], 'w*', ms=16, label='start')12    ax.plot(*minimum, 'kX', ms=12, label='minimum')13    ax.legend()14    return fig, ax

np.meshgrid turns two 1D coordinate arrays into two 2D grids, so X[i,j], Y[i,j] is the $(x,y)$ of pixel $(i,j)$. Because the functions are written with NumPy operations they evaluate on all 160,000 points in one call. This only works if you wrote np.sin rather than math.sin — math functions reject arrays, and that is the first error most people hit here.

Make contours the default view, not a 3D render, because a contour map is measurable. The gap between consecutive dots is the step length, so you can see at a glance whether steps are shrinking (converging) or growing (diverging). Whether the path crosses contour lines at right angles or slides along them tells you whether it is descending or stuck. Perspective and occlusion hide all of that in 3D.

The 3D version, honestly labelled

1f = valley                               # any of the four surfaces2path, values = gradient_descent(f, valley_grad, [5.0, 5.0], 0.05, 60)3X, Y = np.meshgrid(np.linspace(-6, 6, 400), np.linspace(-6, 6, 400))4Z = f(X, Y)                              # the later snippets reuse X and Y56fig = plt.figure(figsize=(7, 6))7ax = fig.add_subplot(111, projection='3d')8ax.plot_surface(X, Y, Z, cmap='viridis', alpha=0.6, linewidth=0)9zs = np.array([f(px, py) for px, py in path]) + 0.5   # lift so it is not buried10ax.plot(path[:, 0], path[:, 1], zs, 'r.-', lw=1.5)

The + 0.5 offset lifts the path above the surface; without it the line is buried in the mesh half the time. Build this view — it makes the word "surface" concrete and it looks good in a write-up. Just do not diagnose from it: depth is unreadable, hills hide the path behind them, and rotating to see one part hides another.

The loss curve, on a log axis

1fig, ax = plt.subplots(figsize=(6, 4))2ax.semilogy(values)                      # log scale on y3ax.set_xlabel('iteration'); ax.set_ylabel('f(theta)')4ax.grid(True, which='both', alpha=0.3)

Use semilogy, always. On the elongated bowl a good run goes from 275 down to about $10^{-8}$. On a linear axis the first three steps consume the whole height of the plot and the remaining 497 are a flat line pinned to zero — you cannot tell converged-to-$10^{-8}$ from stalled-at-$10^{-2}$. On a log axis, healthy convergence is a straight downward line, because each step multiplies the loss by a constant factor. A plateau is unmistakably horizontal. Divergence is a straight line going up.

The comparison grid

One run in isolation is hard to judge. Four on identical axes explain themselves.

1rates = [0.01, 0.05, 0.1, 0.12]2fig, axes = plt.subplots(1, 4, figsize=(20, 4.5), sharex=True, sharey=True)3for ax, lr in zip(axes, rates):4    path, values = gradient_descent(valley, valley_grad, [5.0, 5.0], lr, 60)5    ax.contourf(X, Y, valley(X, Y), levels=30, cmap='viridis', alpha=0.8)6    ax.plot(path[:, 0], path[:, 1], 'ro-', ms=3, lw=1)7    ax.set_xlim(-6, 6); ax.set_ylim(-6, 6)   # else the diverging run rescales every panel8    ax.set_title(f'lr = {lr}  |  final f = {values[-1]:.3g}')

This is the figure to keep. Left to right: a path that drops onto the valley floor and then crawls, one that walks in cleanly, one that ricochets down a corridor, and one that is thrown out of the frame on its very first step. Fixing the axis limits matters: the axes are shared, so without them the diverging run would stretch every panel out to billions and flatten the other three paths into dots.

Sliders, and one detail that matters

In a notebook, wrap the whole thing in ipywidgets:

1from ipywidgets import interact, FloatSlider, IntSlider23fig, ax = plt.subplots(figsize=(6, 5))           # created ONCE, outside45@interact(lr=FloatSlider(0.05, min=0.001, max=0.3, step=0.001),6          beta=FloatSlider(0.0, min=0.0, max=0.99, step=0.01),7          x0=FloatSlider(5.0, min=-6, max=6, step=0.1),8          y0=FloatSlider(5.0, min=-6, max=6, step=0.1),9          n_steps=IntSlider(60, min=5, max=500, step=5))10def explore(lr, beta, x0, y0, n_steps):11    path, values = gradient_descent(valley, valley_grad, [x0, y0], lr, n_steps, beta)12    ax.clear()                                   # reuse the existing axes13    ax.contourf(X, Y, valley(X, Y), levels=30, cmap='viridis', alpha=0.8)14    ax.plot(path[:, 0], path[:, 1], 'ro-', ms=3, lw=1)15    ax.set_title(f'final f = {values[-1]:.4g} after {len(path)-1} steps')

Creating the figure once and calling ax.clear() inside is the detail that matters. Call plt.subplots() on every slider event instead and you leak a figure per movement; dragging for ten seconds opens a hundred of them and the browser crawls. Outside a notebook, the same exploration works as a plain function with default arguments that you re-call by hand — the sliders are convenience, not substance.

The formula that predicts the divergence you just watched

Watching a run explode is useful. Predicting it in advance is better, and for these functions you can predict it exactly.

For a quadratic function, gradient descent converges if and only if

where $L$ is the largest curvature of the surface — formally the largest eigenvalue of the matrix of second derivatives. For $f = x^2 + 10y^2$ the second derivatives are $\partial^2 f/\partial x^2 = 2$ and $\partial^2 f/\partial y^2 = 20$, so $L = 20$ and the threshold is $2/20 = 0.1$.
See why by doing the $y$ coordinate by hand. The update is $y_{t+1} = y_t - \eta \cdot 20 y_t = (1 - 20\eta)\, y_t$. Each step multiplies $y$ by the fixed factor $(1-20\eta)$. Starting from $y_0 = 5$:
$\eta = 0.01$: factor $= 1 - 0.2 = 0.8$. Shrinks. $5 \to 4 \to 3.2 \to 2.56$.
$\eta = 0.05$: factor $= 1 - 1 = 0$. Lands exactly on $y=0$ in one step.
$\eta = 0.1$: factor $= 1 - 2 = -1$. Flips sign, never shrinks. $5 \to -5 \to 5 \to -5$, forever.
$\eta = 0.12$: factor $= 1 - 2.4 = -1.4$. Flips and grows by 40% each step. After 500 steps $f$ is around $10^{148}$; left running, it overflows to inf at step 1,047.
Now run all four in the visualiser from $(5,5)$ and confirm it. This is the table your project should produce.
$\eta$
$\beta$
Steps to $f < 10^{-6}$
Final $f$ (500 steps)
What the path looks like
0.01
0
422
$4.2\times10^{-8}$
Drops onto the valley floor, then crawls towards the centre
0.05
0
81
$\approx 0$
One clean step to $y=0$, then a straight run in
0.10
0
never
250
Locked in a perfect zig-zag; $x$ shrinks, $y$ never does
0.12
0
never
$3\times10^{148}$
Oscillations widen by 40% a step; inf at step 1,047
0.01
0.9
153
$1.4\times10^{-21}$
Same route, roughly 3x faster along the flat direction
0.10
0.9
161
$1.4\times10^{-21}$
A few wobbles, then the zig-zag cancels itself out
0.18
0.9
170
$4.2\times10^{-20}$
Still stable at a rate that would have exploded without momentum
0.20
0.9
never
$2\times10^{179}$
Past the momentum stability bound too; inf at step 867
The last three rows are the payoff. Momentum does not merely speed things up; it moves the stability threshold to
$$\eta < \frac{2(1+\beta)}{L}$$
which for $\beta = 0.9$ and $L = 20$ gives $0.19$. And $0.18$ is stable while $0.20$ blows up, passing $10^{179}$ by step 500. A formula on paper, confirmed by a picture on your screen.
Rosenbrock obeys the same law and shows how brutal it gets. Its curvatures at the minimum $(1,1)$ work out to roughly $1002$ and $0.4$, so $L \approx 1002$ and the bound is $\eta < 0.002$. Test it: at $\eta = 0.002$ from $(-1.2, 1)$ the run reaches $f < 10^{-3}$ after about 3,200 steps; at $\eta = 0.003$ it is still stranded at $(0.74, 0.60)$ after fifty thousand. The ratio $1002 / 0.4 \approx 2500$ is the condition number, and it is why even the stable run takes thousands of steps to crawl along a valley it entered almost immediately.
Divergence is not bad luck. It is a learning rate above $2/L$, and $L$ is a property of the surface you can compute before you start.
What will go wrong, and how to recognise it
Everything becomes nan. Once one value overflows to inf, the next gradient is inf, and inf - inf is nan, which then poisons every subsequent number. The isfinite guard already in the loop is what keeps the last few real values so you can see the ramp-up. Reduce $\eta$ and it goes away.
Your gradient is wrong. This is the single most common bug and it is silent: the optimiser will happily descend some function, just not the one you plotted. Check every gradient against a central finite difference, $\frac{f(\theta + h) - f(\theta - h)}{2h}$, which approximates the derivative to within about $h^2$:
Python
WrapCopy
1def check_grad(f, grad, theta, h=1e-5):2    theta = np.array(theta, dtype=float)3    numeric = np.zeros_like(theta)4    for i in range(len(theta)):5        step = np.zeros_like(theta); step[i] = h6        numeric[i] = (f(*(theta + step)) - f(*(theta - step))) / (2 * h)7    analytic = grad(*theta)8    print(f"analytic {analytic}  numeric {numeric}  max diff {np.abs(analytic - numeric).max():.2e}")910check_grad(rosen, rosen_grad, [-1.2, 1.0])     # expect max diff below 1e-6
Run it once per function at a point where nothing is zero. Checking at the minimum tells you nothing, since both answers are zero whether or not your algebra is right. Anything above about $10^{-4}$ means a sign or a coefficient is wrong.
The contour plot looks empty. The path left the window on step two and every dot is off-frame. Set the axis limits from the path itself rather than hard-coding them.
Rosenbrock's contours are one solid block. Its values run from $0$ at the minimum to over $2000$ nearby, so 30 evenly spaced levels put 29 of them in the far corner and none anywhere near the valley. Space them logarithmically: levels=np.logspace(-1, 3.5, 30). The banana shape appears immediately.
Once the four functions work, extend. Add RMSprop and Adam behind the same interface and draw all three paths on one contour map; how they separate on the elongated bowl is the clearest explanation of adaptive methods anyone has drawn. Animate a run with FuncAnimation. Overlay the gradient field with plt.quiver to see the arrows the optimiser is following. Add a decaying learning rate schedule and watch the zig-zag shrink on its own. Or add Gaussian noise to each gradient to mimic mini-batch stochasticity, and watch the path rattle around a shallow minimum instead of settling — which is what your real training runs actually do.
What carries over to a real training run
The zig-zag you produced on $x^2 + 10y^2$ is not a toy phenomenon. It is what an ill-conditioned curvature matrix does, and in a network with ten million parameters the ratio between the sharpest and flattest directions is routinely in the thousands or worse. The mechanism is identical. The difference is that you cannot draw a ten-million-dimensional contour plot, so you never see it happening.
So you diagnose from symptoms instead, and the value of having built this tool is that you now know which symptom means what. A loss that explodes to nan within a few dozen steps is a learning rate above the stability bound; halve it and it will very likely just work. A loss that drops sharply then decays far too slowly, with gradient norms that stay large, is the crawl along the valley floor — fix it with momentum, an adaptive optimiser, or by normalising your inputs so the curvatures sit closer together. A loss that oscillates around a level without trending anywhere is the ricochet.
And $\eta < 2/L$ is why learning rate warm-up works. Early in training the curvature is often larger than it will be later, so a rate that is stable at step 10,000 is above the threshold at step 10. Ramping it up over the first few hundred steps keeps you under a moving bound. That is a real technique in real systems, and it is the same inequality you just watched decide whether a red dot spiralled into the centre of a picture or shot off the edge of it.