How training actually works

The puzzle: 79,510 knobs and one number

Last week you trained an MNIST network and treated the training loop as ritual. That ritual set all 79,510 parameters in the little 784 → 100 → 10 network (784×100 weights + 100 biases, then 100×10 + 10) to values that recognize handwriting. Nobody searched each combination. A 79,510-dimensional space is too large to search directly. What information guides all those decisions at once?

The answer is one number, differentiated. The loss is a single scalar measuring how bad the model is on a batch. Calculus asks that loss, "If each parameter moved up a tiny bit, would you rise or fall, and how steeply?" The gradient answers for every parameter at once. One backward pass computes all 79,510 slopes for about the cost of one forward pass. The search now has a compass.

Loss is the terrain

The loss is one number measuring how bad the model is on a batch. Imagine that number as height and every model parameter as a coordinate. Real models have millions of coordinates, so the picture is impossible to draw, but the local idea still works: if the loss slopes upward in one direction, move the weights in the opposite direction.

The loss stays scalar, but its gradient is a collection of slopes with one entry per parameter. A weight matrix of shape (100, 784) receives a gradient of shape (100, 784); its bias vector receives a gradient of shape (100,). PyTorch preserves those shapes in each parameter's .grad field. The optimizer never extracts 79,510 decisions from the loss value alone. It uses the derivative of that value with respect to all 79,510 coordinates.

The figure is a 1-parameter cartoon of a 79,510-parameter reality. High-dimensional terrain contains features the cartoon cannot show. Most flat-looking points are saddle points, flat in some directions and sloped in others; true valley floors are rarer. The local move still holds: feel the slope, step against it, repeat.

SGD (stochastic gradient descent) is the plain version of that move. Each step uses the gradient from one random batch, giving "stochastic" its name and producing a noisy but cheap compass reading. Watch it work on the bowl f(x) = x²:

# 01: gradient descent on a bowl
def grad(x):
    return 2 * x              # slope of the bowl f(x) = x ** 2

x = 5.0
lr = 0.1
for step in range(6):
    x = x - lr * grad(x)      # step against the slope, toward lower loss
    print(round(x, 3))
# 4.0  3.2  2.56  2.048  1.638  1.311   -> sliding toward the minimum at x = 0

Each pass takes one step downhill, and the value creeps toward the minimum. The lr is the Learning rate : it scales every step. Too high and the optimizer jumps over the valley or diverges; too low and training crawls.

See both failure modes once because you will diagnose them later. Same bowl, three learning rates:

# 02: the learning rate has two failure modes
def grad(x):
    return 2 * x

for lr in [0.1, 0.6, 1.1]:
    x = 5.0
    for step in range(6):
        x = x - lr * grad(x)
    print(f"lr={lr}: x after 6 steps = {x:.3f}")

Expected output: lr=0.1 reaches 1.311 (safe but slow, still far from 0), lr=0.6 reaches 0.000 (large steps overshoot the minimum and bounce across it, but each bounce lands closer), and lr=1.1 reaches 14.930 (each bounce lands farther away and soon diverges to infinity). A "my loss went to NaN" bug is the third line at model scale.

Real training adds batch noise to this clean bowl. The bowl's gradient is exact, but SGD estimates a network's gradient from a mini-batch: a small sample of training examples processed together. Two batches contain different digits, so they give slightly different downhill directions. The wobble buys cheaper updates. A full-dataset step would be steadier, but it would cost one pass through every example before moving once. Mini-batches produce a stream of cheap, imperfect updates. They also average efficiently on GPU hardware: one example wastes parallel compute, while an enormous batch consumes memory and provides fewer updates per pass through the data.

A batch size of 64 means the loss averages the errors from 64 examples. The gradient asks how every parameter should change to improve that average. The next batch asks the same question of a different 64. Individual steps wobble while the trajectory heads downhill. A jagged loss chart may show healthy batch noise. Values that grow without bound point to a high learning rate or numerical instability.

An epoch completes one pass through the training set. With 60,000 examples and batches of 64, an epoch contains about 938 updates. This distinction matters when you compare runs: ten epochs with batch size 64 performs far more parameter updates than ten epochs with batch size 1,024. Track both epochs and optimizer steps when a learning curve moves at an unexpected speed.

Momentum and Adam are engineering fixes

Plain SGD struggles in long, narrow valleys with steep walls in one direction and a gentle floor sloping toward the goal in another. The gradient points mostly at the nearest wall, so SGD ricochets across the valley and advances toward the minimum in tiny sideways slivers.

Momentum is the fix, and the figure gives away why it works. Across the valley, successive gradients point in opposite directions (left wall, right wall, left wall...), so if you average recent gradients they cancel. Along the valley, successive gradients agree, so the average accumulates. Momentum remembers recent direction like a ball carrying velocity, and its smoothing effect appears after several steps.

# 03: momentum is a multi-step effect
def grad(x):
    return 2 * x                     # slope of the bowl f(x) = x ** 2

def descend(momentum, lr=0.1, steps=6):
    x, v = 5.0, 0.0
    history = []
    for step in range(steps):
        v = momentum * v + grad(x)   # velocity accumulates past gradients
        x = x - lr * v
        history.append(round(x, 2))
    return history

print("plain SGD", descend(momentum=0.0))
print("momentum ", descend(momentum=0.9))
# plain SGD [4.0, 3.2, 2.56, 2.05, 1.64, 1.31]
# momentum  [4.0, 2.3, 0.31, -1.54, -2.9, -3.54]   <- velocity carried it past the minimum

The first step is identical for both because velocity starts at zero and therefore equals the gradient. The rolling velocity then pulls the momentum run ahead, fast enough here to overshoot the bottom and swing back. Momentum is a multi-step effect. On this symmetric bowl the overshoot looks like a bug; in a narrow valley, the same accumulated velocity carries you along the floor while damping wall-to-wall bounces.

Adam keeps running averages of gradients and squared gradients for each parameter, then adapts each parameter's step size. Parameters with consistently large gradients need small steps, while parameters with tiny gradients need larger steps; one global learning rate cannot suit both. Dividing a parameter's step by its running gradient scale evens out progress. Two more terms appear in the exercises. Weight decay shrinks every weight slightly toward zero on each step, pressing against large weights. Bias correction compensates for running averages that start at zero and begin too small. In the notebook, you implement Adam as bookkeeping over tensors.

The optimizer loop

The full training loop lets PyTorch calculate gradients:

# no-run
# 04: the five-line liturgy, named
import torch

model = torch.nn.Linear(2, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.randn(8, 2)
y = torch.randn(8, 1)

pred = model(x)
loss = ((pred - y) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
print(loss.item())

Week 2 introduced this loop as ritual. You now name each moving part and implement the optimizers yourself. opt.step() becomes fifteen lines you wrote, with a momentum buffer you can print. Next week opens loss.backward().

Each line passes data to the next. pred = model(x) creates predictions using the current parameters. The loss compresses all eight prediction errors into one scalar. zero_grad() clears gradient buffers left by the previous batch. PyTorch adds new gradients to those buffers, which supports gradient accumulation but surprises readers who expect assignment. backward() fills each parameter's .grad buffer. step() then reads those gradients and mutates the parameters. The printed loss is a positive scalar whose exact value varies because the model and data are random. It describes the batch before this update, so another forward pass must measure the updated model.

Training-loop bugs break that order. Forget zero_grad() and gradients from unrelated batches pile up. Call step() before backward() and the optimizer has no fresh direction. Reuse pred after step() and it still represents the old forward pass. Inspect one parameter three times: print its value before backward(), confirm its .grad becomes non-None afterwards, then confirm the value changes after step(). That trace separates "the model produced no gradient" from "the optimizer ignored a valid gradient" without guessing.

Hyperparameters and experiment tracking

A hyperparameter is a setting humans choose: learning rate, batch size, momentum, weight decay, or schedule. A schedule changes the learning rate during training. A typical schedule warms it up briefly from zero, then decays it gradually. The optimizer updates weights; humans choose the training recipe.

Repeated runs create a mundane engineering problem. By run seven, you will forget whether run three used lr=1e-3 or 1e-4. Weights and Biases logs losses, accuracies, and configs to the cloud, then plots runs side by side. A sweep runs a structured set of hyperparameter choices. The dashboard answers "which change helped?" without relying on memory. ARENA keeps using this workflow in later chapters, so the setup effort pays rent beyond this week.

Pair-session guide

Core work is ARENA 0.3 section 1: implement SGD (momentum and weight decay are built into that one exercise), then RMSprop, Adam, and AdamW, and race them across pathological loss surfaces. Stretch work is section 2 (W&B logging and sweeps on a ResNet finetuned on CIFAR10) and section 3 (distributed training). Pair rule: when an optimizer test fails, compare update equations line by line before changing code.

What you should see

You should see optimizer trajectory plots where SGD zig-zags, momentum smooths the path, and the plotted Adam run reaches the basin quickly. If you do the W&B stretch, your dashboard should show multiple runs and the learning rate or schedule that changed each curve.

The trajectory plots bring this week's diagram to life: your optimizer implementations draw the same narrow-valley ricochet. An SGD path that zig-zags and a momentum path that glides reproduce the reason these algorithms were invented.

Where to go next

Week 4 has you write backward() itself. These readings extend the optimizer story:

• A Visual Explanation of Gradient Descent Methods, ARENA's "if you only read one thing, make it this".
• Andrew Ng's short videos on momentum, RMSProp, and Adam (about 7-9 minutes each).
• Why Momentum Really Works on Distill, if you want the mathematical story with interactive plots.