Deep Learning · Efficient Training

Gradient Accumulation: What Micro-Batching Preserves and What It Does Not

Splitting a batch into micro-batches keeps the arithmetic of a large batch but changes normalisation, batch-dependent layers, and the timing of every state update.

A batch of 256 samples that does not fit in device memory can be processed as eight micro-batches of 32, with the gradients summed before one optimiser step. The resulting parameter update is, up to floating-point summation order, the same as the update a single batch of 256 would have produced. That equivalence is exact for a narrow class of models and approximate for almost every real one, and the differences are not numerical noise.

Knowing precisely where the equivalence holds determines whether accumulation is a safe substitute for a larger batch or a change to the training procedure that needs its own validation.

The gradient is a sum, so accumulation is exact for that sum

For a loss that averages over samples,

$$\mathcal{L} = \frac{1}{B}\sum_{i=1}^{B} \ell(x_i)$$

the gradient with respect to a parameter is $\frac{1}{B}\sum_i \nabla \ell(x_i)$. Split the batch into $k$ micro-batches of size $b = B/k$. The sum over the full batch is the sum of the per-micro-batch sums. If each micro-batch’s loss is divided by $B$ rather than by $b$, the accumulated gradient equals the full-batch gradient term by term.

That division is the detail most implementations get wrong. The natural code scales each micro-batch loss by $1/b$, which is correct for a standalone step and off by a factor of $k$ when summed. The result is an effective learning rate multiplied by $k$, which usually looks like a run that diverges for no visible reason.

accum_steps = 8
optimizer.zero_grad(set_to_none=True)
for i, micro in enumerate(loader):
    loss = criterion(model(micro.inputs), micro.targets)
    (loss / accum_steps).backward()   # normalise by the full batch, not the micro-batch
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

When the last accumulation window is short because the loader ran out, the divisor must match the number of micro-batches actually summed; otherwise the final step of each epoch is scaled differently from the rest.

What breaks the equivalence

Three classes of operation depend on the batch beyond a simple sum.

Statistics computed across the batch. Batch normalisation computes a mean and variance over the samples in the tensor it sees. In a micro-batch of 32, those statistics come from 32 samples; in a batch of 256, from 256. The normalised activations therefore differ, and so do the gradients. The error shrinks as the micro-batch grows but never vanishes. Layer normalisation, which reduces over the feature dimension within each sample, is unaffected, which is part of why transformer architectures tolerate accumulation so easily. The architectural consequence is that accumulation is nearly free for a transformer and carries a measurable accuracy question for a convolutional network built on batch statistics.

Losses that are not a per-sample average. Contrastive losses with in-batch negatives compute similarity against every other sample in the batch. Splitting the batch removes negatives: a micro-batch of 32 offers 31 negatives where the full batch of 256 offered 255. The loss value changes and the gradient it produces is for a different objective. The same applies to any loss using batch-level ranking or hard-negative mining. Accumulation cannot repair this, because the missing samples were never in the same forward pass.

Layers whose state updates on every forward call. Running-statistics buffers in normalisation layers update once per forward pass. Under accumulation they update $k$ times per optimiser step, with each update computed from a smaller sample. With a momentum term near one, the running estimate becomes an average over micro-batches rather than over the batch, and the weighting is not uniform because the micro-batches arrive in a fixed order. The mitigation is to update such buffers only on the step boundary, or to accept the drift and monitor it.

A fourth effect is not a correctness issue but a scheduling one: dropout masks are drawn per forward pass. Eight micro-batches draw eight independent masks instead of one mask over the whole batch. For most purposes this is a slightly different regularisation pattern, not an error, but it means the two procedures are not bit-identical and a comparison between them has a noise floor.

The memory saving is real, the compute saving is not

Accumulation reduces the peak activation memory by a factor of $k$, because only one micro-batch’s activations are live at a time. It does not reduce the number of FLOPs: the same $6ND$ work is performed, as described in the step accounting. It also does not reduce the parameter, gradient, or optimiser-state memory, which is independent of batch size, and which the memory budget breakdown shows to be the dominant term for large models. A run that is out of memory because of optimiser state gains nothing from accumulation; the remedy there is sharding, not micro-batching.

There is a throughput cost as well. Each micro-batch runs the full sequence of kernel launches, and each launch has overhead. More importantly, the arithmetic intensity of each matmul depends on the micro-batch size: with fewer rows in the activation matrix, the weight matrix is reused fewer times per load, and the step becomes more bandwidth-bound. A micro-batch of 1 is dramatically less efficient than a micro-batch of 32 on the same hardware, and no amount of accumulation recovers that, because the small matmul is the unit of work.

The practical consequence is that the micro-batch should be as large as memory allows, and the accumulation factor chosen to reach the desired effective batch. The opposite arrangement, a small micro-batch with a large factor, gives the same gradient and worse throughput.

The effective batch is a hyperparameter, not a constant

Accumulation is usually introduced to reach a target effective batch size, and the target itself interacts with the learning rate. The linear scaling rule, established for large-batch SGD, holds that multiplying the batch size by $k$ permits multiplying the learning rate by $k$ while preserving the number of epochs to convergence, up to a limit beyond which the approximation fails and warmup becomes necessary.

The limit is where the gradient noise across the batch becomes small relative to the curvature of the loss. Beyond it, larger batches produce diminishing reductions in gradient variance and the step size that the linear rule prescribes overshoots. This is why large-batch training recipes pair the scaling with a warmup period: the early steps use a small rate while the running estimates in adaptive optimisers fill in.

Accumulation makes the effective batch a free parameter, which invites setting it very large. The compute cost per optimiser step grows linearly with it, and the benefit in gradient quality grows sublinearly. A defensible procedure is to fix the micro-batch at the memory limit, then choose the accumulation factor so the effective batch matches the largest value the learning-rate schedule was tuned for, rather than inventing a new effective batch and a new schedule simultaneously.

Interaction with the rest of the step

Accumulation composes with mixed precision but changes where the scaler sits. If loss scaling is used, the scale must be applied to each micro-batch loss and the check for non-finite gradients must happen after all micro-batches have been summed, because a single overflowing micro-batch corrupts the whole accumulated gradient. Framework scalers that expect one backward per step need the accumulation expressed in a way they can see, or the check fires at the wrong time.

Accumulation also interacts with gradient clipping: the norm must be computed on the summed gradient, not per micro-batch, or the clipping threshold applies to a partial sum and the effective clip is $k$ times tighter than intended. Both of these are ordering constraints rather than new mechanisms, and both are silent when wrong: the run continues, the gradients are simply wrong by a constant factor. That is the general character of accumulation errors, which is why the arithmetic should be checked once against a reference step at small scale before the factor is raised.

References

  1. Don't Decay the Learning Rate, Increase the Batch Size — Smith et al., ICLR 2018
  2. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour — Goyal et al., 2017
  3. PyTorch documentation — Automatic Mixed Precision examples — PyTorch project