Optimization by gradient descent

Almost every machine-learning model is trained by minimizing a function — the "loss" that measures how wrong the model is. The workhorse algorithm for doing that is gradient descent, and the idea is beautifully simple:

Stand on a hillside in the fog. Feel which way is downhill (the gradient). Take a small step that way. Repeat.

The "small step" size is the learning rate. Too small and you crawl; too big and you leap right over the valley and bounce out. Below you can set the starting point and the learning rate and watch the ball roll down a bumpy landscape, all live in WebAssembly.

begin
    using PlutoUI, WasmMakie
end

A bumpy landscape

We'll minimize f(x) = 0.2 x⁴ − x² + 0.3 x, a curve with two valleys. Gradient descent only ever feels the slope right under its feet, so where it ends up depends on where it starts — a first taste of local versus global minima.

start x = -18 ÷10

learning rate = 12 ÷1000

steps = 30

let
    # derive the (float) start and learning rate from the integer sliders, then run
    # gradient descent, recording every point we visit so we can draw the path
    x = Float64(x0i) / 10.0
    lr = Float64(lri) / 1000.0
    px = Vector{Float64}(undef, nsteps + 1)
    py = Vector{Float64}(undef, nsteps + 1)
    px[1] = x
    py[1] = f_land(x)
    for i in 1:nsteps
        x = x - lr * slope(x)
        px[i + 1] = x
        py[i + 1] = f_land(x)
    end

    # sample the landscape curve over a fixed window with an integer loop
    cx = Vector{Float64}(undef, 241)
    cy = Vector{Float64}(undef, 241)
    for k in 0:240
        t = -2.6 + 5.2 * Float64(k) / 240.0
        cx[k + 1] = t
        cy[k + 1] = f_land(t)
    end

    fig = Figure(size = (600, 360))
    ax = Axis(fig[1, 1])
    lines!(ax, cx, cy)        # the landscape
    lines!(ax, px, py)        # the path the ball took as it rolled downhill
    fig
end

opt_stats = let
    # run the descent in a SEPARATE bond-dependent cell so the markdown below
    # interpolates the result live (values inside a markdown cell's own `let` bake to
    # the slider defaults).
    x = Float64(x0i) / 10.0
    lr = Float64(lri) / 1000.0
    for i in 1:nsteps
        x = x - lr * slope(x)
    end
    (floor(x * 1000.0) / 1000.0, floor(f_land(x) * 1000.0) / 1000.0)
end;

After 30 steps the ball is at x = -1.677, where f = -1.734. Start on the left versus the right and it settles in a different valley. Crank the learning rate up past about 0.05 and watch it overshoot and fly apart – the step size matters as much as the direction.

Why this scales

In one dimension you could just look at the graph. The power of gradient descent is that the exact same loop — evaluate the slope, step downhill — works in millions of dimensions, where you could never plot anything. That is how neural networks with billions of parameters are trained: follow the gradient, one small step at a time.

The two dials you played with here are the two that matter everywhere: where you start (initialization) and how big a step you take (the learning rate).

Appendix

The MIT lecture reaches for ForwardDiff.jl / Optim.jl / JuMP. WebAssembly can't run those in the browser, so we compute the slope with a one-line finite difference and draw with WasmMakie. The algorithm — and the lessons about learning rate and local minima — are exactly the same.

slope

Its slope, by a finite difference (we pretend we don't know the formula).

begin
    "The landscape we are minimizing."
    f_land(x::Float64) = 0.2 * (x * x * x * x) - x * x + 0.3 * x

    "Its slope, by a finite difference (we pretend we don't know the formula)."
    function slope(x::Float64)
        h = 0.0001
        return (f_land(x + h) - f_land(x - h)) / (2.0 * h)
    end
end