Optimization

What it is and why it matters

You're on a hill in thick fog and want to reach the bottom of the valley. You can't see the valley. You can only feel the ground under your feet: which way slopes down, and how steeply. So you take a step downhill, feel again, take another step. That's gradient descent, and it's how nearly every model in this course is trained, from a two-number line fit to an LLM with billions of weights.

Three things decide whether the walk works:

• The loss function: the height of the hill, a single number that says how wrong the model is.
• The gradient: the slope under your feet (last lesson).
• The learning rate: how big a step you take.

How it works

The loss function turns "how wrong" into one number. For predicting numbers, the usual one is mean squared error, the average of (prediction − answer)². For picking classes you'll use cross-entropy (Part 2). Training means: find the weights that make this number small.

The update rule. Repeat:

w = w − learning_rate × gradient(w)

The minus sign walks against the slope, so downhill. Where the hill is steep, the gradient is big and the step is big. Near the bottom the gradient shrinks, so the steps shrink by themselves.

The learning rate is the one knob you set by hand:

• Too small: every step is tiny, and training takes forever.
• Right: the loss drops fast and settles.
• Too big: each step overshoots the bottom and lands higher on the other side. The loss zig-zags or explodes to infinity.

Local minima. A real loss surface isn't a clean bowl. It can have several valleys, and gradient descent stops in whichever one it rolls into first. Where you start (the initial weights) can change where you end. In huge neural networks this matters less than you'd fear: most valleys they find are about equally good. Saddle points and flat regions, where the slope is nearly zero but you're not at the bottom, slow training down more often.

Visual

Practical example

The same one-weight model from the last lesson, trained from w = 0 with three learning rates. Then a bumpy loss with two valleys, started from two different places.

# Gradient descent on loss(w) = (w*2 - 10)^2, the one-weight model from the last lesson.
# The best w is 5. Watch three learning rates.
def loss(w):
    return (2 * w - 10) ** 2

def grad(w):
    return 2 * (2 * w - 10) * 2

for lr in [0.01, 0.1, 0.3]:
    w = 0.0
    trail = []
    for step in range(1, 21):
        w = w - lr * grad(w)  # the whole algorithm is this line
        if step in (1, 2, 5, 10, 20):
            trail.append(f"{step}:{w:.2f}")
    print(f"lr={lr:<4} w after steps {'  '.join(trail)}  final loss {loss(w):.4f}")

# A bumpy loss with two valleys: where you start decides where you end.
def bumpy(w):
    return w**4 - 3 * w**2 + w

def bumpy_grad(w):
    return 4 * w**3 - 6 * w + 1

for start in [2.0, -2.0]:
    w = start
    for _ in range(200):
        w -= 0.01 * bumpy_grad(w)
    print(f"start {start:+.1f} -> w {w:+.3f}, loss {bumpy(w):+.3f}")

package main

import (
	"fmt"
	"strings"
)

// Gradient descent on loss(w) = (w*2 - 10)^2, the one-weight model from the last lesson.
func loss(w float64) float64 { e := 2*w - 10; return e * e }
func grad(w float64) float64 { return 2 * (2*w - 10) * 2 }

func bumpy(w float64) float64     { return w*w*w*w - 3*w*w + w }
func bumpyGrad(w float64) float64 { return 4*w*w*w - 6*w + 1 }

func main() {
	for _, lr := range []float64{0.01, 0.1, 0.3} {
		w := 0.0
		var trail []string
		for step := 1; step <= 20; step++ {
			w = w - lr*grad(w) // the whole algorithm is this line
			switch step {
			case 1, 2, 5, 10, 20:
				trail = append(trail, fmt.Sprintf("%d:%.2f", step, w))
			}
		}
		fmt.Printf("lr=%-4v w after steps %s  final loss %.4f\n", lr, strings.Join(trail, "  "), loss(w))
	}

	// A bumpy loss with two valleys: where you start decides where you end.
	for _, start := range []float64{2, -2} {
		w := start
		for i := 0; i < 200; i++ {
			w -= 0.01 * bumpyGrad(w)
		}
		fmt.Printf("start %+.1f -> w %+.3f, loss %+.3f\n", start, w, bumpy(w))
	}
}

Output:

lr=0.01 w after steps 1:0.40  2:0.77  5:1.70  10:2.83  20:4.06  final loss 3.5605
lr=0.1  w after steps 1:4.00  2:4.80  5:5.00  10:5.00  20:5.00  final loss 0.0000
lr=0.3  w after steps 1:12.00  2:-4.80  5:31.89  10:-139.63  20:-4178.41  final loss 70003769.6591
start +2.0 -> w +1.131, loss -1.070
start -2.0 -> w -1.301, loss -3.514

Reading it:

• lr = 0.01 is crawling. After 20 steps w is only 4.06 and the loss is still 3.56.
• lr = 0.1 reaches w = 4.00 in one step and 5.00 by step 5. The loss is 0.
• lr = 0.3 overshoots to 12, then −4.8, then 31.89, and by step 20 the loss is seventy million. Same model, same data, same code: only the step size changed.
• On the bumpy loss, starting at +2 settles at w = 1.131 (loss −1.07), and starting at −2 settles at w = −1.301 (loss −3.51). Both runs stopped where the slope is zero, but only one found the deeper valley.

In production

• The learning rate is the first thing to tune. If training misbehaves, try 3x and 10x smaller and larger before changing anything else. Most frameworks also use a schedule: a short warm-up, then a slowly shrinking rate.
• Use better optimisers. Plain gradient descent is the base. Momentum and Adam (Part 3) keep a running memory of past gradients, so they roll through flat spots and handle weights that need very different step sizes. Adam is the default for transformers.
• Watch the loss curve. Log the loss at every step. A curve that falls then rises, a nan, or a curve that never moves tells you more than any single number.
• Each step costs a full pass over the data's gradient. Computing it on all the data is too slow, so real training uses small random batches (mini-batch SGD), which also adds a bit of useful noise that helps escape bad spots.
• Guard against nan. One overflow poisons every weight. Clip gradients to a maximum size, and stop the run automatically if the loss becomes nan or infinite, so you don't burn GPU hours on garbage.

Key takeaway

Training is walking downhill on the loss: w = w − learning_rate × gradient, repeated. The learning rate decides everything: too small is slow, too big explodes, and where you start decides which valley you end in.

Practice

Run the example with learning rates 0.2 and 0.24. For each, does the loss go to 0, zig-zag, or explode? Then find the learning rate above which this model always explodes. Hint: one step multiplies the distance from w = 5 by (1 − 8 × lr).