mini-batch SGD 為何每次只用少量訓練樣本？

Stochastic and mini-batch gradient descent: why every real training loop takes small, noisy steps.

Gradient descent, as we wrote it earlier in this series, has one quiet assumption. To take a single step, it looks at every training example. That is fine for four examples. It is not fine for a million.

Real training never does it that way. It takes a small random handful of examples, a mini-batch, and steps on that alone. This article shows why that works, what it costs, and how it looks in PyTorch. The numbers come from a tiny worked example and from the CS878 Deep Learning slides (lectures 07-08).

The gradient is an average

The loss of a model is an average over the training examples: L = (1/N) Σ ℓᵢ, where ℓᵢ is the loss on example i. A derivative of an average is the average of the derivatives, so the gradient is an average too: ∇L = (1/N) Σ ∇ℓᵢ.

Take a model with two weights, ŷ = w·x, a squared loss ℓᵢ = ½(ŷᵢ − yᵢ)², and four examples:

• x₁ = (4, −1), y₁ = 1
• x₂ = (0, 1), y₂ = 3
• x₃ = (1, 1), y₃ = 3
• x₄ = (1, −1), y₄ = 1

Start at w = (0, 0). There each example's gradient is −yᵢ·xᵢ, so the four arrows are (−4, 1), (0, −3), (−3, −3) and (−1, 1).

Put them tip to tail and they add up to (−8, −4). Divide by four and you get the true gradient, ∇L = (−2, −1). A step with learning rate ¼ moves the other way, to w = (½, ¼).

Four examples, four arrows, added tip to tail and divided by four: the true gradient (−2, −1).

Now make it a million examples. The slides use n = 2²⁰, about a million images. Every single step must touch all million. That is a lot of work for one small move.

A mini-batch is a cheap, honest guess

So take a shortcut. Shuffle the data, cut it into small batches, and average the gradients of one batch only:

g_B = (1/B) Σ over i in the batch of ∇ℓᵢ

With our four examples and B = 2, say the shuffle gives batches {1, 3} and {2, 4}.

• Batch one: ½[(−4, 1) + (−3, −3)] = (−7/2, −1). Too long.
• Batch two: ½[(0, −3) + (−1, 1)] = (−1/2, −1). Too short.

One-step guesses from them would land at (7/8, ¼) and (1/8, ¼), on either side of the true landing point (½, ¼). But their average is ½[(−7/2, −1) + (−1/2, −1)] = (−2, −1), exactly the true gradient.

That is the key property. Over the random choice of batch, the expected value of g_B is ∇L. The estimate is noisy, but it is unbiased: it does not lean in any direction.

Batch one's gradient is too long, batch two's too short; their average lands exactly on the true gradient.

How noisy is it?

If one example's gradient has a spread of σ around the true gradient, a batch of B examples has a spread of σ/√B (when the data set is much bigger than the batch).

That square root matters. Four times the batch costs four times the compute but only halves the noise. Going from 1 to 64 examples cuts the spread by a factor of 8. Past a point, a bigger batch buys very little extra accuracy for its cost. A small batch already points the right way.

The cloud of mini-batch gradients around the true one shrinks like σ/√B, while the cost grows like B.

A race on one budget

Give three methods the same budget: one pass over 1,024 examples, called one epoch, with the same learning rate.

• Full batch takes one careful step. After a whole epoch it is still far from the bottom.
• Pure SGD (B = 1) takes 1,024 shaky steps. It reaches the bottom but keeps jittering around it.
• Mini-batch (B = 64) takes 16 steady steps and ends near the bottom with only a small wobble.

There is a hardware reason too. A GPU computes many examples in parallel. A batch of 64 is one GPU call that keeps all its lanes busy, so the mini-batch epoch is 16 calls instead of 1,024.

If you plot the loss against iterations, as the lecture does, full batch is smooth, SGD is jagged, and mini-batch sits in between. That middle ground is why mini-batch SGD is the default.

One pass over 1,024 examples, all three at once: full batch takes 1 step, B = 64 takes 16, B = 1 takes 1,024; a B = 64 step lights all 64 GPU lanes together.

Try it in PyTorch

In PyTorch the batch size lives in the DataLoader. One batch is one iteration, one optimizer step. One full pass over the data is an epoch.

loader = DataLoader(data, batch_size=64, shuffle=True)
opt = optim.SGD(net.parameters(), lr=0.1)
for epoch in range(3):
    for xb, yb in loader:
        loss = loss_fn(net(xb), yb)
        opt.zero_grad()
        loss.backward()
        opt.step()

MNIST has 60,000 training images. In batches of 64 that is ⌈60,000 / 64⌉ = 938 iterations per epoch, and the last batch holds the 32 images left over. You can check it without any real data:

import torch.utils.data as D
import torch
X = torch.zeros(60000, 1)
dl = D.DataLoader(X, 64)
*_, last = dl
print(len(dl), len(last))

Output:

938 32

Change 64 to 1 and you get pure SGD: 60,000 steps per epoch. Change it to 60000 and you are back to full batch: one step per epoch. The optimizer is the same optim.SGD either way. The batches are what make it mini-batch SGD.

The PyTorch loop: one batch of 64 images flows through the network together, and 60,000 / 64 rounds up to 938 steps per epoch.

Noise can even help

At a flat saddle point the true gradient is almost zero, so full-batch descent barely moves. A mini-batch gradient keeps a little noise, and that is enough to shake the weights loose and let them roll downhill.

The noise does not fix everything. In a long, narrow valley the steps still zig-zag from wall to wall. Taming that zig-zag is the job of momentum, the next idea in this series.

Key takeaways

• The loss is an average over examples, so its gradient is an average of per-example gradients.
• A mini-batch gradient is a cheap estimate of the true gradient. It is noisy but unbiased: on average it is exactly right.
• Its noise shrinks like σ/√B, so four times the batch only halves the noise at four times the cost.
• In one epoch, mini-batches take many cheap steps that run in parallel on a GPU, between smooth full batch and jagged pure SGD.
• In PyTorch, batch_size in the DataLoader sets it: 60,000 images in batches of 64 is 938 steps per epoch.

This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.