Gradient Accumulation

~9 min · custom-train

Level 0Keras Apprentice

0 XP0/97 lessons0/20 achievements

0/120 XP to next level120 XP to go0% complete

The problem: batch size you can't afford

Some models train best at a large batch size — transformers especially — but a large batch may simply not fit in GPU memory. Gradient accumulation decouples the statistical batch from the physical one: you run several small micro-batches, sum their gradients, and only call the optimizer once the sum represents the batch you wanted. The weight update is (almost) identical to having run one big batch.

How the math stays honest

Two details make it correct. First, divide each micro-batch's loss by accumulation_steps so the summed gradient is an average, not a sum — otherwise your effective learning rate scales with the accumulation count. Second, apply and then reset the accumulator on the same boundary, so each optimizer step consumes exactly N micro-batches. The example below shows the manual-loop form with TensorFlow's tape; folding it into a train_step() with the accumulator as a model attribute is the production version.

Code

Gradient accumulation (TensorFlow backend)·python

import tensorflow as tf

accumulation_steps = 4  # Effective batch = batch_size x 4
accumulated = None

for step, (x, y) in enumerate(dataset):
    with tf.GradientTape() as tape:
        y_pred = model(x, training=True)
        # Divide so the summed gradient is an AVERAGE
        loss = loss_fn(y, y_pred) / accumulation_steps

    grads = tape.gradient(loss, model.trainable_variables)

    # Accumulate
    if accumulated is None:
        accumulated = grads
    else:
        accumulated = [a + g for a, g in zip(accumulated, grads)]

    # Apply every N micro-batches, then reset
    if (step + 1) % accumulation_steps == 0:
        optimizer.apply_gradients(
            zip(accumulated, model.trainable_variables)
        )
        accumulated = None

External links

Exercise

Take a working CIFAR-10 model. Implement gradient accumulation in a custom train_step (accumulate over 4 micro-batches before applying gradients). Verify training behavior matches a 4× larger batch run.