The training loop: SGD, mini-batches, epochs

Real training never touches the whole dataset at once. It loops over shuffled mini-batches for many epochs, spending a little gradient noise to buy a lot of speed.

You already know how to take one gradient step: measure the loss, backprop to get the gradient, nudge every parameter downhill. The training loop is the boring, powerful machine that wraps that step and runs it thousands of times - but with a twist. Instead of computing the gradient over all your data every time (accurate but slow), it estimates it from a small random batch (noisy but fast). That single trade - a little noise for a lot of speed - is why modern models train at all.

The loop that wraps gradient descent

Here is the entire idea in code. Read it top to bottom: shuffle, slice into batches, and for each batch do one forward pass, one loss, one backward pass, one update. Then repeat for the next epoch. Every deep-learning framework is an elaboration of these lines.

for epoch in range(num_epochs):        # outer loop: full passes over the data
    shuffle(dataset)                    # fresh random order each epoch
    for batch in batches(dataset, B):  # inner loop: one mini-batch of size B
        preds = model(batch.x)         # forward pass   (Lesson 20)
        loss  = loss_fn(preds, batch.y) # how wrong, on THIS batch
        grads = loss.backward()        # backprop       (Lesson 22)
        for p in model.params:
            p -= lr * grads[p]         # gradient step  (Lesson 11)

Batch, stochastic, or mini-batch?

The only real design choice is how many examples go into each gradient estimate - the batch size $B$. That one number defines three classic variants, sitting on a spectrum from slow-and-exact to fast-and-noisy.

In practice mini-batch wins almost everywhere: batches of 32-512 give a decent gradient estimate, keep the GPU busy with parallel matrix work, and let you take hundreds of update steps in the time full-batch takes one. Confusingly, people still say "SGD" for the mini-batch case - the "stochastic" just means "gradient from a random sample," whatever the size.

A live training dashboard

Time to watch the loop breathe. Below, a tiny classifier is trained by mini-batch SGD on a two-blob dataset. Press Run to watch the loss descend in real time; the faint amber line is the raw per-batch loss (that jitter is the mini-batch noise), and the bold teal line is the per-epoch average trend. Shrink the batch size to make the amber line thrash; then crank the learning rate toward the top and watch the whole thing explode. Use Step epoch to advance one epoch at a time and see the decision boundary on the right snap into place.

Reading the loss curve

The loss curve is your instrument panel. A healthy run drops fast, then flattens into a noisy plateau - that flattening is your cue that a fixed rate has done its job. The next move is to decay the learning rate over time on a schedule: big steps early to cover ground, tiny steps late to settle precisely. The giant pretraining run of Lesson 50 does exactly this over billions of steps, and the next lesson makes momentum, Adam, and schedules precise. A curve that rises means the rate is too high; a curve that plateaus early and high means it is too low or the model is too small. Same loop, same reading, whether you train a toy classifier or a trillion-token language model.

Check yourself