SGD and Mini-Batch Training

Why noisy gradients help, batch size effects, and the implicit regularization of SGD.

You are training a 500,000-parameter network on 10 million examples. To take *one* honest gradient step, plain gradient descent must first run all 10 million examples through the network. On a GPU chewing through a thousand examples a second, that is nearly three hours — per step. One step, wait three hours, one step, wait three hours. Training would take decades. Before mini-batching, this was the actual wall people hit.

The fix is almost cheeky in hindsight: do not measure the slope from all 10 million examples. Grab a small *random handful* — say 32 — and estimate the slope from those. It is like working out which way is downhill by feeling the ground under a few random footsteps instead of surveying the whole mountain. Any single estimate is a little *noisy* — a different handful gives a slightly different direction — but on average it points the right way, and you can compute it in a blink. Trading one perfect step every three hours for a slightly-wobbly step every fraction of a second is a spectacular deal: with batches of 32, you go from *one* update per pass through the data to over 300,000. That is mini-batch stochastic gradient descent (SGD), and it is how essentially every neural network is trained.

The surprise: the noise actually helps

Here is the twist nobody saw coming. That gradient noise, which sounds like a necessary evil, turns out to make the final model *better*.

Picture the loss landscape as having two kinds of valleys: narrow, steep-walled sharp pits, and wide, gently-sloping flat basins. A perfectly noiseless optimizer slides straight into whichever valley is nearest and stops — sharp pit or not. But a *noisy* optimizer keeps getting jostled. In a narrow sharp pit, a jostle easily knocks it back out; in a wide flat basin, the same jostle is not enough to escape. So the noise acts like a filter: the optimizer keeps getting bounced out of sharp pits until it settles into a flat basin — and flat basins are exactly the ones that generalize better to new data. The messiness is doing quiet regularization work, for free.

Batch size is a dial, not just a memory setting

Most people pick batch size by "whatever fits in GPU memory." But because the batch size controls how much noise there is, it is really a *regularization dial*. Bigger batches mean *less* noise (each estimate averages more examples), so the optimizer behaves more like the noiseless one — faster and smoother, but drawn toward those sharp, worse-generalizing pits. Push the batch very large and you can match the training loss of a small-batch run while landing a couple of points *worse* on test accuracy, purely from the lost noise. So sensible defaults are modest batches (32 to 256); go bigger only when you truly need the speed, and expect to add explicit regularization (like weight decay) to make up for the noise you removed.

One rule you cannot skip: shuffle

Finally, a practical trap that quietly ruins training: shuffle your data before every pass. The whole "noisy but right on average" guarantee assumes each mini-batch is a random sample. If your data happens to be sorted by class — all the cats, then all the dogs — an unshuffled batch contains *only cats*, and its gradient screams "get better at cats" while forgetting dogs entirely. Consecutive batches then yank the model in wildly different directions. Shuffle before each epoch and every batch becomes a fair little snapshot of the whole dataset, which is what makes the estimate honest.

Epoch, step, iteration — pin down the vocabulary

These get muddled constantly, so be exact. One epoch is one full pass through the entire dataset. One step (also called one iteration) is a single mini-batch update. The relationship: updates per epoch = ceil(N / batch_size). So 1,000,000 examples at batch size 256 is ⌈1,000,000 / 256⌉ = 3,907 steps per epoch. When a paper says "trained for 100k steps" versus "100 epochs," these mean different amounts of compute unless you also know N and the batch size — always convert to one consistent unit before comparing runs.

"Small batch generalises better" is a tendency, not a law

The flat-minima story is real but *not universal*. Whether small batches actually generalise better depends on the learning rate, the dataset size, the architecture, whether you use normalisation, the schedule, and the training budget. With careful warmup, LR scaling, and enough steps, large-batch training can close much of the gap (this is how models train on thousands of GPUs). So don't state "small batch = better generalisation" as a rule — state it as a default tendency that a well-tuned large-batch setup can partly overcome.

Scaling the learning rate to the batch has limits

The linear scaling rule (multiply LR by k when you multiply batch size by k) works — but only up to a regime. Past a few thousand examples per batch it breaks down: the implied LR gets so large that early training destabilises, which is exactly why warmup (ramping the LR up over the first epochs) becomes essential at large batch. Some setups find square-root scaling (LR ∝ √k) safer than linear at scale. The takeaway: LR-vs-batch scaling is an approximation with a ceiling, not a guaranteed equivalence.

BatchNorm couples training to the batch size

If your network uses BatchNorm, batch size stops being just a noise/speed knob because BatchNorm computes its normalisation *statistics from the current batch*. Very small batches give noisy mean/variance estimates that destabilise training (a batch of 2 has almost meaningless statistics). Very large batches give near-exact statistics, which changes BatchNorm's own regularisation behaviour. This is a big reason transformers favour LayerNorm (which normalises per-example and doesn't depend on batch size). When you change batch size on a BatchNorm model, you're changing more than the gradient noise.

SGD versus Adam, fairly

"SGD generalises better than Adam" is too blunt. Adam/AdamW often *wins* — it's the standard for transformers and shines on sparse gradients and NLP. SGD-with-momentum can generalise better in some classic vision/CNN settings when well-tuned. Why the gap leans that way ties back to the noise story above: Adam's per-parameter adaptive steps damp gradient noise along some directions, so it explores less of the landscape and can settle into sharper minima — often training faster but finishing a touch worse than a well-tuned SGD run, whose undamped noise keeps bouncing it toward the flatter basins that generalize better. The honest summary: AdamW is the default for modern deep learning; SGD+momentum remains competitive or better in specific well-studied regimes. Pick based on the domain and tune both before declaring a winner.

Big batches live across many GPUs: distributed training

At scale the batch is split across devices, which adds vocabulary. Global batch size is the total across all GPUs; per-device batch size is what each one processes. Gradient accumulation simulates a large batch on limited memory by summing gradients over several forward/backward passes before one update. Data parallelism replicates the model on each GPU and all-reduces the gradients — which costs communication bandwidth, often the real bottleneck at scale. So "batch size 8,192" usually means a global batch spread over many devices, not one machine's memory.

Rare classes can get squeezed out of batches

Random mini-batching assumes every batch is a fair sample — but under heavy imbalance a batch of 32 from a 0.1%-positive dataset frequently contains *zero* positives, so many steps carry no signal about the rare class. Fixes: stratified batches (force a minimum number of rare-class examples per batch), weighted sampling (oversample the rare class into batches), and hard-example mining (bias batches toward the examples the model currently gets wrong). Under imbalance, how you *build* the batch matters as much as its size.

Key points

• Default to modest batch sizes (32–256); go bigger only when speed forces it, and pay for it with extra regularization. Small batches supply the noise that finds flat, well-generalizing minima. If your GPU is sitting mostly idle, the batch is too small and you are wasting throughput; if test accuracy is a point or more below published numbers for the same architecture, the batch may be too large. When you do scale the batch up by k, scale the learning rate up by about k too (the "linear scaling rule") — but know that this only fixes step size, not the lost noise, so generalization still slips once batches get very large.
• The trap: cranking the batch size for speed without realising you switched off the free regularization. The symptom is sneaky: training loss matches the benchmarks, but test accuracy sits 1–3% low, so you blame the data or the architecture. The real culprit is a batch of 4,096 where the benchmark used 256. Fix it by shrinking the batch, or by adding explicit regularization (weight decay, dropout) to replace the noise you removed — and always report batch size alongside the generalization gap.
• The diagnostic: watch training and validation loss, and how the loss moves per step versus per epoch. Validation loss much higher than training loss points to too large a batch (a sharp minimum, overfitting). A loss that is jittery step-to-step but steadily falling epoch-to-epoch is healthy mini-batch behaviour. A loss that is perfectly smooth every single step means you are effectively doing full-batch descent — worth asking whether you actually want that, and what it is costing you in generalization.
• Pin the vocabulary and qualify the small-batch claim. One epoch = one full pass; one step/iteration = one mini-batch update; updates per epoch = ceil(N / batch_size), so "100k steps" and "100 epochs" differ unless you know N and batch size. "Small batch generalises better" is a tendency, not a law — it depends on LR, dataset size, architecture, normalisation, schedule, and budget, and well-tuned large-batch training (with warmup) closes much of the gap. The linear LR-scaling rule holds only up to a regime; beyond a few thousand it needs warmup and sometimes √-scaling is safer.
• Watch BatchNorm coupling, judge SGD-vs-Adam fairly, and handle scale and imbalance. With BatchNorm, batch size changes the normalisation statistics themselves — very small batches destabilise, very large ones alter its regularisation (a reason transformers use batch-independent LayerNorm). Don't blanket-claim SGD beats Adam: AdamW is the modern default and wins for transformers/sparse gradients, while SGD+momentum can win in some tuned vision settings. At scale, distinguish global vs per-device batch, use gradient accumulation for limited memory, and mind all-reduce communication cost. Under heavy imbalance, random batches can contain zero rare-class examples — use stratified batches, weighted sampling, or hard-example mining.

Recap

• Full-batch is infeasible at scale: to take *one* honest step it runs all 10M examples through the network first — ~3 hours per step on a GPU doing 1k examples/sec. One step, wait three hours: training would take decades. This was a real wall before mini-batching.
• Mini-batch SGD: estimate the slope from a small random handful (say 32) instead of the whole dataset — any one estimate is a little noisy but on average points the right way, and you go from 1 update per pass to 300,000+. This is how essentially every neural net is trained.
• The noise turns out to help: a noiseless optimizer slides into whatever valley is nearest, sharp or not; a noisy one keeps getting jostled out of narrow sharp pits but stays in wide flat basins — and flat basins generalize better. Free regularization, for nothing.
• Batch size is a regularization dial, not just a memory setting: bigger batch = each estimate averages more examples = *less* noise = the optimizer drifts toward sharper, worse-generalizing minima. Push it very large and you can match training loss but land a couple of points worse on test accuracy.
• Shuffle before every epoch: the "noisy but right on average" guarantee assumes each batch is a random sample. Data sorted by class gives all-cats-then-all-dogs batches whose gradients scream "get better at cats" and yank the model around — shuffling makes every batch a fair snapshot.
• Pin the vocabulary: one epoch = one full pass over the data; one step / iteration = one mini-batch update; updates per epoch = ceil(N / batch_size). "100k steps" and "100 epochs" mean different compute unless you also know N and the batch size — always convert to one unit before comparing runs.
• Know the caveats so you don't overstate them: "small batch generalizes better" is a *tendency* (depends on LR, data size, architecture, norm, schedule, budget), the linear LR-scaling rule breaks past a few thousand and needs warmup, BatchNorm couples training to batch size, and heavy class imbalance can leave batches with zero rare-class examples (use stratified/weighted sampling).

Check your understanding