Mini-batches and SGD

Why a rough direction is good enough

Think of tasting soup. You don't drink the whole pot to check the salt. One spoonful from a well-stirred pot tells you almost the same thing. A random batch is that spoonful.

Each example in a batch gives its own gradient. Call the gradient from example 1 $g_1$, from example 2 $g_2$, and so on up to $g_B$. The batch's gradient is their average.

In plain words: add up the gradients from every example in the batch. Then divide by how many examples there are. That average is the direction of the step.

Real gradients have one number per weight. To keep it small, pretend there's only one weight. A batch of 4 examples gives gradients 3, 1, 2 and 6. They add up to 12, and 12 ÷ 4 = 3. So the batch's gradient is 3.

grads = [3, 1, 2, 6]               # one gradient per example in the batch
print(sum(grads) / len(grads))     # the batch's gradient: their average

This prints 3.0. Because the batch is picked at random, its average is right on average. Some batches point a bit too far one way, and some a bit too far the other. Over many steps, those errors mostly cancel out. The noise can even help: it can jostle the model out of a poor spot.

import random

random.seed(0)
grads = [random.gauss(2.0, 3.0) for _ in range(1000)]   # one made-up gradient per example
full = sum(grads) / len(grads)
print(f"full batch (all 1000): {full:.2f}")

for size in [1, 10, 100]:
    guesses = []
    for _ in range(5):                                   # five different random batches
        batch = random.sample(grads, size)
        guesses.append(round(sum(batch) / size, 2))
    print(f"batch size {size:>3}: {guesses}")

The DataLoader pattern

In PyTorch, DataLoader does the batching for you. It wraps a dataset and hands out (inputs, targets) pairs, one batch at a time. With shuffle=True, it reshuffles every epoch, so the batches differ on each pass.

That matters. Think of studying with flashcards. If the deck is always in the same order, you start learning the order instead of the cards. In the same way, batches that are identical every epoch can make training go round in circles.

Every PyTorch training loop you'll read has this shape:

loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
for epoch in range(num_epochs):
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()

The outer loop counts epochs. The inner loop takes one step per batch. By default the last, smaller batch is kept. Pass drop_last=True to skip it instead.

Choosing a batch size

Bigger and smaller batches each have strengths. A GPU is the graphics chip that most training runs on. It's fast when it has lots of work at once, but its memory is limited.

To generalize means to do well on new data the model hasn't seen.

In practice, pick the largest batch that fits comfortably in memory and still trains well. Then tune the learning rate for it. A common rule of thumb: when you scale the batch size up, scale the learning rate up by roughly the same factor. It's a starting point, not a law.

When the batch you want doesn't fit: gradient accumulation

Sometimes the batch you want is too big for the GPU's memory. There's a neat way around it.

In PyTorch, gradients pile up until you call zero_grad(). Each backward() adds to what's already there. You can use that on purpose. Run several small batches through backward(), and only then call step(). You've now taken one step with the gradient of all those batches combined. Each small batch is called a micro-batch.

One detail matters. Divide each loss by the number of micro-batches, so the result is an average and not a total:

accum = 4                                     # 4 micro-batches of 16 ≈ one batch of 64
optimizer.zero_grad()
for i, (xb, yb) in enumerate(loader):
    (loss_fn(model(xb), yb) / accum).backward()
    if (i + 1) % accum == 0:
        optimizer.step()
        optimizer.zero_grad()

The line if (i + 1) % accum == 0 is true on every fourth micro-batch. That's when the step happens and the gradients are cleared. Call the number of micro-batches $k$. Each one gives an average loss, and you add them up after dividing each by $k$.

In plain words: shrink each micro-batch's loss by the number of micro-batches, then add them up. The total is the average over all of them, the same as one big batch would give.

Say four micro-batches have losses 0.8, 0.6, 0.4 and 0.2. Divided by 4, they become 0.2, 0.15, 0.1 and 0.05. Those add up to 0.5, which is the plain average of the four losses. The gradients add up the same way, because each backward() adds its share.

losses = [0.8, 0.6, 0.4, 0.2]    # four micro-batch losses
accum = 4
total = 0.0
for loss in losses:
    total += loss / accum        # what accumulation adds up
print(round(total, 2))
print(round(sum(losses) / len(losses), 2))   # the plain average

Both lines print 0.5. Without the division, the total would be 2.0, four times too big. That's like quietly multiplying the learning rate by 4 (with plain SGD, as here).

LLM training batches are measured in millions of tokens. They're assembled exactly this way, across many GPUs and many accumulation steps.