Gradient Accumulation

Train with a batch that does not fit, by not holding it all at once. The gradient is a sum, and sums can be built up.

Overview

Why it works

The loss over a batch is the mean of the per-example losses, and differentiation is linear. So the gradient of the batch is the mean of the per-example gradients.

That means the sum can be built in pieces. Run a forward and backward pass on 16 examples, keep the gradients, run another 16, add to them, and after 16 rounds you have exactly the gradient of a 256-example batch. Step the optimiser once.

The resulting update is identical — not similar, identical up to floating-point ordering — to what a single 256-example batch would have produced.

Gradient Accumulation

This module needs JavaScript: the numbers are computed in the page rather than recorded.

Worth knowing

The gradient of a mean loss over a batch is the mean of the per-example gradients. So it can be accumulated in pieces.
Run several small forward and backward passes, sum the gradients, and step once. The update is identical to the large batch.
Only activation memory scales with the micro-batch. Weights, gradients and optimiser state do not.
It costs wall-clock time, not correctness — except for batch normalisation, which sees only the micro-batch.

Gradient Accumulation

Getting the effect of a large batch out of a device that cannot hold one.

What actually consumes the memory

The bars separate the two cases at the memory budget you choose.

Training memory is roughly:

Fixed costs. Weights, their gradients, and the optimiser state. Adam keeps two extra tensors per parameter, so this is around four times the model size, and it does not depend on the batch at all.

Activations. Every intermediate value from the forward pass has to be kept for the backward pass. This scales linearly with the batch size, and for a deep model it dominates.

Accumulation attacks only the second, which is why it works. Drag the micro-batch down and watch the second bar fall while the fixed part stays put.

The dashed line is the budget. When the full batch crosses it and the micro-batch does not, that gap is the entire reason the technique exists.

What it costs

Time. Sixteen sequential passes take about as long as sixteen passes. Accumulation does not make anything faster; it makes something possible.

There is a small efficiency loss too, since very small micro-batches use the accelerator less well. A micro-batch of 1 is usually much less than a sixteenth as efficient as a micro-batch of 16.

The one thing that is not identical

Batch normalisation. Its statistics are computed over whatever is in the forward pass, which is the micro-batch, not the accumulated total.

So a model with batch norm trained with a micro-batch of 4 behaves like one trained with a batch of 4, however many accumulation steps follow. The normalisation is noisier, and at very small micro-batches it degrades badly.

This is a real reason to prefer [layer normalisation](layer_normalization.html) or group normalisation, both of which normalise per example and are indifferent to the batch. It is not a coincidence that transformers use layer norm and are the architectures trained with the most aggressive accumulation.

Getting it right

Two details cause most of the bugs.

Divide the loss by the number of accumulation steps before the backward pass, or the accumulated gradient is a sum where the optimiser expects a mean, and your effective learning rate is multiplied by the step count.

Zero the gradients at the right time — after the optimiser step, not after every backward pass. Zeroing every pass defeats the whole thing, and it fails silently: training still runs, just with the small batch you were trying to avoid.

Alongside

Activation checkpointing trades compute for memory in the other direction: discard activations during the forward pass and recompute them in the backward pass, for roughly 30% more time and a large memory saving. The two combine.

Mixed precision halves activation memory outright, and combines with both.

Model parallelism and ZeRO attack the fixed costs, which accumulation cannot touch. When the weights alone do not fit, accumulation is not the answer.

Where it goes wrong

Forgetting to divide the loss. Silently multiplies the learning rate.

Zeroing gradients every pass. Silently defeats the technique.

Using it with batch norm at a tiny micro-batch. Real quality loss.

Expecting a speed-up. It buys capability, not throughput.

Check yourself

0 of 3

Answer without scrolling back up.

  1. Why is the accumulated update identical to a large-batch update?

  2. Which part of training memory does accumulation reduce?

  3. Which layer type does gradient accumulation genuinely change?

Cheat sheet

Gradient Accumulation

The loss over a batch is the mean of the per-example losses, and differentiation is linear. So the gradient of the batch is the mean of the per-example gradients.

DEEP LEARNING · vizlearn.in/deep_learning/gradient_accumulation.html

About the author

Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.