How training actually works

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.

SGD is the plain version:

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.

Momentum and Adam are engineering fixes

Plain SGD can zig-zag in narrow valleys. Momentum remembers recent direction like a ball carrying velocity: the smoothing effect builds over several steps, not on the first.

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 (velocity starts at zero, so it equals the gradient). Only afterwards does the rolling velocity pull the momentum run ahead, here fast enough to overshoot the bottom and swing back. That is the whole point: momentum is a multi-step effect.

Adam goes further by keeping running averages of gradients and squared gradients for each parameter, then adapting the step size per parameter. In the notebook, you implement it as bookkeeping over tensors.

The optimizer loop

Here is the whole shape of training, with PyTorch doing gradient calculation:

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. This week names each moving part and asks you to implement optimizers yourself.

Hyperparameters and experiment tracking

A hyperparameter is a setting you choose rather than learn: learning rate, batch size, momentum, weight decay, schedule. The optimizer can only optimize weights; humans still choose the training recipe.

Weights and Biases is a dashboard for comparing runs. A sweep is a structured set of runs over different hyperparameters. The point is answering "which change actually helped?" without relying on memory.

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 Adam often reaches the basin quickly. If you do the W&B stretch, you should see multiple runs on your dashboard and be able to point to the learning rate or schedule that changed the curve.

Where to go next

Week 4 removes the last black box by making you write backward() itself. The ARENA notebook's recommended reading for this week is unusually good:

• 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.