Lesson 13 of 16

Mini-Batches: Trading Off Noise and Speed

Every train_step you've written computed a gradient from a handful of points and updated immediately, or averaged over the entire dataset first. Both are extremes of the same idea, called batch size: how many examples' gradients get averaged together before the weights move. Real datasets can have millions of examples, averaging all of them before a single update would make training glacially slow, but updating after every single example is noisy and can't take advantage of fast, parallel hardware doing many examples at once.

Batch size 1 (pure stochastic gradient descent): update after every single example. Cheap per step, and you get many updates per epoch, but each one is based on only one point's gradient, noisy, and can bounce the loss around rather than descending smoothly. Full-batch: average the gradient over the entire dataset before updating once. Smooth, low-variance steps, but only one update per epoch, however large the dataset. Mini-batch is the practical middle ground almost everyone actually uses: average over a modest chunk, say 32 or 256 examples, getting a much cleaner gradient estimate than batch size 1 while still updating many times per epoch.

Same network, same data, three different batch sizes

updates / epoch

total updates (25 epochs)

500

final loss

0.000

Batch = 1 updates the weights after every single example, 20 times per epoch, each update noisy (based on one point's gradient alone) but there are far more of them. Full-batch updates once per epoch using every point's averaged gradient, a smoother but much slower-ticking descent. Batch = 4 sits in between. This is the real tradeoff: smaller batches mean noisier, cheaper, more frequent updates; larger batches mean smoother, more expensive, less frequent ones.

All three batch sizes train the exact same network on the exact same data for the same number of epochs, only how the gradient is averaged before each update differs. Compare the "updates per epoch" and "total updates" numbers as you switch: batch size 1 gets 20 updates every epoch (one per point), full-batch gets exactly 1. More frequent updates, even noisy ones, often win the race to a low loss within a fixed number of epochs, which is a big part of why mini-batch training dominates in practice.

The only change from a full-batch training loop: slicing the dataset before each update instead of using all of it.

Python

def train_epoch(params, data, batch_size, lr, train_step_fn):
    for i in range(0, len(data), batch_size):
        batch = data[i:i + batch_size]
        train_step_fn(params, batch, lr)  # averages the gradient over just this batch

# batch_size=1: one update per example, len(data) updates per epoch
# batch_size=len(data): one update per epoch, using every example at once
# anything in between: mini-batch, the practical default

data = list(range(20))  # stand-in for 20 training examples
for batch_size in [1, 4, 20]:
    num_batches = len(range(0, len(data), batch_size))
    print(f"batch_size={batch_size:>2}: {num_batches} updates per epoch")

def train_epoch(params, data, batch_size, lr, train_step_fn):
    for i in range(0, len(data), batch_size):
        batch = data[i:i + batch_size]
        train_step_fn(params, batch, lr)  # averages the gradient over just this batch

# batch_size=1: one update per example, len(data) updates per epoch
# batch_size=len(data): one update per epoch, using every example at once
# anything in between: mini-batch, the practical default

data = list(range(20))  # stand-in for 20 training examples
for batch_size in [1, 4, 20]:
    num_batches = len(range(0, len(data), batch_size))
    print(f"batch_size={batch_size:>2}: {num_batches} updates per epoch")

Your turn

Implement num_batches(dataset_size, batch_size): how many batches (updates) one epoch produces, rounding up for a leftover partial batch.

Python

Batch size changes how often and how noisily weights update, but every batch size so far has been trained on a dataset the network was allowed to memorize freely. The next module asks a different question: does a network that fits its training data perfectly actually generalize to data it hasn't seen?

Why might batch size 1 reach a low loss faster (in terms of training epochs) than full-batch, despite each individual update being noisier?

← Training on a 2D Toy Dataset Overfitting and Regularization →