Batch Size and Learning Rate
Batch size trades memory, throughput and gradient noise against each other, and it moves the right learning rate with it. Larger batches want larger rates — up to a point — and small batches regularise for free.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
A bigger GPU arrived and the batch size was raised to fill it. Throughput doubled and the model got worse. What did the batch size change besides speed?
An ads-ranking team moved training to accelerators with four times the memory and raised the batch size to match, expecting the same model faster. Examples per second went up as hoped; the held-out ranking metric went down and the retraining pipeline now produces a slightly worse model every week.
A larger batch gives a more accurate gradient and better hardware utilisation, so larger is strictly better as long as it fits in memory. Keep the learning rate and everything else the same.
With the same rate and the same number of epochs, the bigger batch made a quarter as many updates (Epoch, Batch, Step) and each update was no larger. The model is undertrained and the gap looks like a quality regression with no obvious cause.
- With the same rate and the same number of epochs, the bigger batch made a quarter as many updates (Epoch, Batch, Step) and each update was no larger. The model is undertrained and the gap looks like a quality regression with no obvious cause.
- When the team compensated by training longer, the model matched on training loss and still trailed on held-out ranking. The larger batch removed gradient noise that had been acting as a regulariser, and the run settled in a sharper minimum that generalised slightly worse.
- Scaling the learning rate up with the batch fixed most of the gap and then, at a further doubling, stopped helping: past a batch size the data-and-model combination sets, extra examples per step buy nothing, and the run simply costs more.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- Reach the same quality of minimum as before in less wall-clock time — the target of the migration was speed at equal quality, and a faster run to a worse model is not that.
- The surrounding system ranks ads by predicted click probability; a small quality loss is multiplied by a very large volume of decisions.
- Click logs, billions of rows, with a highly imbalanced label. One gradient estimate is the average over a batch of rows, and its variance falls with batch size.
- The data is not the constraint; the accelerator is. Batch size was chosen to maximise utilisation, which is a hardware decision made without reference to the optimisation it changes (GPU Fundamentals).
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- A mini-batch gradient is the true gradient plus noise whose variance scales as
1 / batch_size. Doubling the batch halves the noise variance and roughly doubles the compute per step; it does not double the useful information per step unless the noise was the thing limiting progress. - Under SGD, multiplying the batch by
kand the rate bykkeeps the expected progress per example roughly constant — the linear scaling rule — because the update isηtimes an average ofktimes as many gradients. The rule holds while the batch is small enough that the noise still dominates the step; beyond that critical batch size, the rate cannot be raised further without hitting the curvature limit (Gradient Descent), and larger batches give diminishing returns. - The noise removed by a large batch was doing work. Stochastic gradient noise pushes the iterate out of sharp, narrow minima towards flatter ones that are more robust to the shift between training and held-out data; small batches therefore regularise, and a large batch with a matched rate often needs explicit regularisation to compensate (Regularisation).
- From the hardware side, an accelerator spends fixed time per step on kernel launches and memory traffic and variable time on arithmetic; a larger batch amortises the fixed part and raises utilisation — until the arithmetic saturates and utilisation flattens (Inference Batching makes the same argument for serving).
Noise, throughput and the rate move together
Batch size is one knob that moves three quantities: the variance of each gradient estimate, the fixed-cost amortisation on the accelerator, and the learning rate that gives the same progress per example. Treating it as a memory setting changes the other two without anyone deciding to.
The pattern below shows the same run at four batch sizes with the rate adjusted by the linear rule and a critical batch size in the middle. Past it, the steps-to-target curve flattens: the extra examples per step are paid for and buy nothing.
| Batch | Rate (linear rule) | Gradient noise | Steps to target | Examples to target | Accelerator use |
|---|---|---|---|---|---|
| 64 | η | high | many | fewest | poor |
| 256 | 4η | moderate | about a quarter | about the same | good |
| 1024 | 16η | low | less than a quarter fewer | more | saturated |
| 4096 | 64η — exceeds stability | very low | barely fewer | four times more | saturated |
The migration that got faster and worse
The offline metric was honest and the migration still looked like a success for a fortnight, because the numbers on the dashboard were throughput numbers. The quality regression was small, consistent and attributed to the week's data until someone plotted held-out quality against the batch size in the run record.
Training throughput up several-fold; training loss at the end of the same epoch count slightly higher; held-out ranking metric slightly lower on every retraining run since the change.
A small, persistent decline in click-through on the ranked slots, visible only after aggregating over a week because the per-day noise is larger than the effect.
- 1The batch was raised without the rate or the training length, so each run makes a quarter of the updates at the same step size — an undertrained model every week.
- 2With training length restored, the remaining gap is the lost gradient noise: the larger batch settles in a sharper minimum that generalises slightly worse without added regularisation.
- 3Batch-norm layers in the tower now see different batch statistics, a second effect that changes the model independently of the optimisation.
Where the hardware argument stops
An accelerator wants large batches for the same reason a serving system does: fixed costs per call are amortised over more work (Inference Batching, GPU Fundamentals). The argument is correct and it has a limit that the hardware cannot see, because the limit is in the statistics of the gradient, not in the chip.
So the assumption to keep checking is that the batch has not crossed that limit. It is a property of the data and the model, and it moves when they do.
The production batch size is at or below the point where additional examples per step stop reducing the number of steps to a target loss, and the rate has been scaled for it.
holds when A steps-to-target sweep over batch sizes has been run for this model family and the chosen batch sits on the falling part of the curve, with the rate re-swept at that batch.
breaks when Hardware with more memory arrives and the batch is raised to fill it; a smaller model with a lower critical batch size replaces the current one; the data becomes cleaner and less noisy, lowering the point at which averaging stops helping.
respond Re-sweep the rate at the new batch, and if the batch is past the critical point, reduce it or use the extra memory for something else — a wider model, longer sequences — rather than a bigger batch.
How to build it
Most important first.
- When raising the batch size, raise the learning rate with it — linearly as the first guess, then re-sweep — and lengthen the warm-up, because a larger rate makes the early steps more dangerous.
- Find the critical batch size empirically: plot steps-to-target-loss against batch size on a log scale. Where the curve stops falling is where larger batches stop paying, and the batch should be at or below that point.
- Hold total updates or total examples seen constant deliberately when comparing batch sizes, and compare on held-out quality, not throughput.
- If a large batch is forced by hardware efficiency, add back regularisation — decay, dropout, label smoothing — and check that generalisation is recovered rather than assuming it.
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Held-out quality at a matched compute budget across batch sizes. This is the number the migration is about; examples per second is the number that motivated it, and it is not the same thing.
- Steps to reach a fixed training loss against batch size, on log axes, which locates the critical batch size and shows where scaling stops working.
- Accelerator utilisation as a separate number, so an efficiency win and a quality loss are reported side by side instead of netted into one impression.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The learning rate in the retraining config is the one tuned for the current batch size; a change to either without the other is a change to the optimisation.
- The current batch size is at or below the critical batch size for this data and model, so each added example per step still improves the gradient estimate.
- Any regularisation that was implicit in the smaller batch's gradient noise has been replaced explicitly if the batch was raised.
- Offline: run the batch-size sweep once per model family and record the critical batch size with the learning-rate sweep; assert that the production batch size sits at or below it.
- On migration: compare the new configuration against the old at a matched number of examples seen and a matched compute budget, on held-out quality with more than one seed, before switching the pipeline.
- Over time: track held-out quality per retraining run against the batch size and rate in the run record, so a hardware-driven change to either shows up as the cause of a quality step.
What can go wrong
- Linear scaling is applied past the critical batch size and the rate exceeds the stability limit; the run diverges, and the batch size is blamed rather than the rate.
- Batch normalisation statistics change meaning with the batch size, so a model with batch-norm layers behaves differently at a new batch size even at a correctly scaled rate (Normalisation Layers).
- Gradient accumulation is used to reach the large batch and the noise is reduced as intended, but the wall-clock gain is nil because the accelerator was already saturated at the smaller micro-batch.
- A batch-size sweep is several training runs and must be repeated when the model family changes — the cost of knowing where scaling stops.
- Small batches generalise well and use accelerators badly; large batches use accelerators well and need the rate, warm-up and regularisation retuned. Neither is free, and the balance moves with hardware.
- Holding quality constant under a larger batch usually means more total compute than the naive "same epochs" run, which erodes the speed-up that motivated the change.
- "A bigger batch gives a more accurate gradient, so training is better." It gives a less noisy gradient. Whether that helps depends on whether noise was the limit, and the noise was also regularising.
- "We doubled the batch and doubled the rate, so the run is equivalent." Linear scaling is an approximation that holds below a critical batch size and for SGD; under Adam the rate should scale less than linearly, and the warm-up must lengthen regardless.
- "Throughput doubled, so we are training twice as fast." Examples per second doubled. Progress per example fell. Whether the product went up is the sweep's job to answer.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- GENERALThat gradient noise scales inversely with batch size and that a critical batch size exists follows from the statistics of averaging, for any model trained by mini-batch descent.
- MODEL-SPECIFICThe linear scaling rule is derived for SGD; under adaptive optimisers the appropriate scaling is weaker, and models with batch normalisation change behaviour with batch size for a second, unrelated reason.
- SIMPLIFIEDThe "noise as regulariser" account is a useful simplification of a contested literature on flat and sharp minima; the practical claim that small batches often generalise slightly better at matched budgets is well supported, the mechanism less so.
Where the depth lives
This domain teaches the model and hands the rest off by name.