Gradient descent visualiser

Summary

You build a 2D loss surface, derive its gradient, and write the gradient-descent update yourself. Then you run the optimiser from a fixed start and plot the path it takes over a contour map of the surface. You sweep the learning rate to watch a small rate crawl, a mid rate converge, and a large rate diverge, and you turn momentum on and off to see it cut a straighter path through a narrow valley. The surface is a few lines of numpy, the optimiser is your code, and the plot is where the behaviour becomes obvious.

Learning goals

• Feel that a gradient is the slope vector at a point, and descent is "step the opposite way, repeat".
• See the learning rate decide everything: too small crawls, too large overshoots or diverges, just right converges.
• Watch momentum smooth the zigzag through a narrow valley.
• See the gap between what the optimiser knows (the local slope at one point) and what you see (the whole surface in the plot).
• Recognise divergence as a numerical event (loss climbing, values going to inf or NaN), not a mystery.

Prerequisites

• L19 (gradients, descent, learning rate, minima and saddle points).
• B3 for numpy and matplotlib comfort. Depth-by-choice: skipping this does not block conceptual progress.

Estimated time

About 2 to 2.5 hours: roughly half an hour on the surface and its gradient, an hour on the update rule and the trajectory plot, and the rest on the learning-rate and momentum sweeps.

Deliverables

• descent.py: the surface, the hand-written gradient, the update rule, the trajectory loop, and the contour plot.
• One or two plot artefacts: a trajectory over the contours, and a learning-rate comparison.
• README.md: the update rule in your own words, the learning rate that converged, and what momentum changed.

Suggested file structure

builds/B4/
  descent.py        # surface + gradient + step + loop + contour/trajectory plot
  trajectory.png    # the optimiser path over the contours
  lr_sweep.png      # trajectories or loss curves for several learning rates
  README.md         # update rule, the lr that worked, what momentum did

One script is plenty. There is no data file: unlike B3, B4 is fully self-contained because the surface is a function.

Step-by-step instructions

1. Define a surface. Start with an elongated bowl, f(x, y) = 0.5 * (A*x**2 + B*y**2) with A and B far apart (say 1 and 20). Evaluate it on a grid with np.meshgrid and draw a contour plot so you can see the valley.
2. Write the gradient. Derive the slope by hand. For this bowl it is [A*x, B*y]. Sanity-check it points away from the minimum.
3. Write the update rule. step(p, v, lr, momentum): compute the gradient, update the velocity, step downhill, return the new point and velocity.
4. Run the loop. From a fixed start, record every point into a trajectory. Plot the trajectory over the contours.
5. Sweep the learning rate. Try 0.001, 0.01, 0.1, 0.5. Plot the trajectories or the loss-versus-step curves side by side. Find the one that crawls, the one that converges, and the one that diverges.
6. Toggle momentum. Run the same start with momentum 0 and 0.9 on the elongated bowl. Compare the paths.
7. Guard against divergence. Stop the loop if the loss climbs or goes non-finite, and report it rather than plotting garbage.
8. Record a table. Learning rate, momentum, converged yes or no, steps to a loss threshold, final loss.
9. Optional saddle. Swap in the saddle surface below, start just off-centre, and watch the optimiser slide away from the flat point. Keep this short.
10. Write the README.

Starter skeleton

Two functions carry the optimiser and are left for you to write. Everything else is scaffolding. Writing gradient and step yourself is the milestone.

import numpy as np
import matplotlib.pyplot as plt

A, B = 1.0, 20.0                       # an elongated bowl: a narrow valley along x

def f(p):                              # scaffolding: the surface
    x, y = p
    return 0.5 * (A * x**2 + B * y**2)

def gradient(p):
    # TODO (you write this): the slope vector at p.
    # for this bowl the gradient is [A*x, B*y]
    ...

def step(p, v, lr, momentum):
    # TODO (you write this): one optimiser update.
    #   velocity = momentum * v - lr * gradient(p)
    #   new_p    = p + velocity
    #   return new_p, velocity        (momentum = 0 gives plain gradient descent)
    ...

def run(start, lr, steps=200, momentum=0.0):   # scaffolding: the loop + recording
    p = np.array(start, dtype=float)
    v = np.zeros(2)
    traj = [p.copy()]
    for _ in range(steps):
        p, v = step(p, v, lr, momentum)
        traj.append(p.copy())
        if not np.isfinite(f(p)):              # divergence guard
            print("diverged at step", len(traj))
            break
    return np.array(traj)

def plot_contours(traj):                       # scaffolding: meshgrid + contour + path
    xs = np.linspace(-3, 3, 200)
    ys = np.linspace(-3, 3, 200)
    X, Y = np.meshgrid(xs, ys)
    Z = 0.5 * (A * X**2 + B * Y**2)
    plt.contour(X, Y, Z, levels=30)
    plt.plot(traj[:, 0], traj[:, 1], marker=".")
    plt.xlabel("x"); plt.ylabel("y"); plt.title("descent trajectory")
    plt.savefig("trajectory.png")

For the optional saddle demonstration, swap in a second surface and its gradient. Keep it brief; it shows a saddle point, it is not the main convergence exercise:

# a saddle: down in y, up in x. start just off-centre and watch it slide away.
def f_saddle(p):
    x, y = p
    return x**2 - y**2

def grad_saddle(p):
    x, y = p
    return np.array([2*x, -2*y])     # flat at the origin, but y still goes downhill

Expected output

Your exact numbers depend on the surface constants and the start point; these are illustrative, not targets:

• Trajectory. The path curves downhill into the minimum on the contour plot.
• Learning-rate sweep. A small rate barely moves, a mid rate settles cleanly, a large rate overshoots or blows up.
• Momentum. On the elongated bowl, vanilla descent zigzags across the valley while momentum cuts a straighter, faster path.
• Divergence. Above some learning rate the loss climbs to inf or NaN and the guard stops the loop.

A results table makes the sweep legible:

lr      momentum   converged   steps   final loss
0.001   0.0        no (crawl)   200      4.83
0.02    0.0        yes          138      0.001
0.02    0.9        yes           41      0.000
0.20    0.0        no (diverged) 7        inf

Validation criteria

Assess against the Build Track Validation Standard. The bar is understanding, not a converged number.

Common pitfalls

These are conceptual traps, distinct from code symptoms.

• Stepping the wrong way. The gradient points uphill. The step must subtract it. Add it and the optimiser climbs.
• Reading the plot as the optimiser's knowledge. You see the whole surface. The optimiser only ever feels the local slope at its current point.
• Expecting one learning rate to work everywhere. The right rate depends on the surface. The bowl's two axes even prefer different rates, which is the point of the valley.
• Thinking momentum is free. On a simple bowl too much momentum overshoots and oscillates.
• Confusing convergence with the global minimum. On a multi-minimum surface, the start point decides which basin you land in.

Troubleshooting

These are code symptoms and their likely causes, distinct from the conceptual pitfalls above.

Optional extensions

• A harder surface. Add Rosenbrock (a curved banana valley) or Himmelblau (four equal minima, where the start point decides the outcome). Their gradients are fiddly, so write them out carefully or use the numerical gradient below.
• Numerical gradient. Approximate the slope by finite differences (central difference) so any surface plugs in without hand-derivation. Pure numpy, no autograd.
• Other update rules. Add Nesterov momentum, RMSProp, or Adam and compare their paths against plain descent. Implement them by hand, not from a library.
• Animation. Animate the descent step by step instead of drawing the whole path at once.
• Basin map. Vary the start point across a grid and colour each by which minimum it reaches.
• Loss curves. Plot loss versus step for several learning rates on one axis.