Gradient Accumulation › The Accumulation Procedure

Gradient accumulation works by splitting a logical large batch into a series of smaller micro-batches, running forward and backward passes on each one without applying the optimizer step, and accumulating the resulting gradients. Only after all micro-batches have been processed does the optimizer perform a single weight update using the accumulated gradients.

The procedure follows a precise sequence:

1. Zero the gradients at the start of each accumulation cycle. This is the same optimizer.zero_grad() call you know from ordinary training, but it now occurs once per logical batch rather than once per micro-batch.
2. Process each micro-batch: run the forward pass, compute the loss, scale the loss by $1/K$ where $K$ is the number of accumulation steps, and run the backward pass. The backward pass adds gradients to the gradient buffers without resetting them.
3. Step the optimizer once after all $K$ micro-batches have been processed. The gradient buffers now contain the accumulated sum, which after the $1/K$ scaling applied to each loss equals the mean gradient over the full logical batch.
4. Repeat for the next logical batch.

The memory benefit arises from step 2: at any point in time, only a single micro-batch occupies activation memory. The gradient buffers themselves must still hold the full model's gradients, but activation memory, which is proportional to batch size, is bounded by the micro-batch size.

The key structural difference from ordinary training is the placement of zero_grad() and step(). In ordinary training, both calls appear at every iteration. With gradient accumulation, zero_grad() appears at the start of every $K$-th iteration, and step() appears at the end of every $K$-th iteration. The $K-1$ iterations in between only call backward(), allowing gradients to accumulate.

The loss scaling in step 2 deserves careful attention. Standard training computes the mean loss over a batch and then backpropagates that mean. When you accumulate gradients across $K$ micro-batches, each micro-batch computes its own mean loss. If you simply sum these $K$ mean losses and backpropagate each one without scaling, your accumulated gradient is $K$ times larger than the gradient you would have computed from the full batch. This means the optimizer takes a step $K$ times too large, corrupting training.

To preserve mathematical equivalence with full-batch training, you must divide each micro-batch loss by $K$ before calling .backward(). This ensures that the accumulated gradients match the gradient from the combined batch exactly.

Let $\mathbf{g}_k$ denote the gradient vector computed from the $k$-th micro-batch with loss scaling applied. The scaled gradient for micro-batch $k$ is:

$$\mathbf{g}_k = \frac{1}{K} \cdot \nabla_\theta \mathcal{L}(\theta; B_k)$$

where:

• $K$: the total number of accumulation steps (micro-batches per logical batch)
• $\theta$: the model parameters
• $\mathcal{L}(\theta; B_k)$: the mean cross-entropy loss computed over micro-batch $B_k$
• $B_k$: the $k$-th micro-batch, a non-overlapping subset of the full logical batch $B$
• $\nabla_\theta$: the gradient operator with respect to $\theta$

After $K$ backward passes, PyTorch's autograd engine has accumulated the sum of all $K$ gradient contributions into the parameter gradient buffers:

$$\mathbf{g}_{\text{accum}} = \sum_{k=1}^{K} \mathbf{g}_k = \frac{1}{K} \sum_{k=1}^{K} \nabla_\theta \mathcal{L}(\theta; B_k)$$

where:

• $\mathbf{g}_{\text{accum}}$: the accumulated gradient stored in the parameter buffers after all $K$ backward passes
• $\mathbf{g}_k$: the gradient contribution from micro-batch $k$, already scaled by $1/K$

If the micro-batches partition $B$ evenly (each of size $|B|/K$), the accumulated gradient equals the gradient computed directly on the full batch $B$:

$$\mathbf{g}_{\text{accum}} = \nabla_\theta \mathcal{L}(\theta; B)$$

This equality holds because gradient computation is linear: the gradient of a mean is the mean of the gradients. Formally, if $B = B_1 \cup B_2 \cup \cdots \cup B_K$ is a partition of the full batch with $|B_k| = |B|/K$ for all $k$, then:

$$\begin{aligned} \nabla_\theta \mathcal{L}(\theta; B) &= \nabla_\theta \left[\frac{1}{|B|} \sum_{i \in B} \ell(\theta; x_i, y_i)\right] \\ &= \frac{1}{|B|} \sum_{i \in B} \nabla_\theta \ell(\theta; x_i, y_i) \\ &= \frac{1}{K} \sum_{k=1}^{K} \frac{1}{|B_k|} \sum_{i \in B_k} \nabla_\theta \ell(\theta; x_i, y_i) \\ &= \frac{1}{K} \sum_{k=1}^{K} \nabla_\theta \mathcal{L}(\theta; B_k) \\ &= \mathbf{g}_{\text{accum}} \end{aligned}$$

where:

• $\ell(\theta; x_i, y_i)$: the per-sample loss for sample $(x_i, y_i)$
• $|B|$: the full logical batch size
• $|B_k| = |B|/K$: the size of each micro-batch

The accumulated result is an unbiased estimate of the full-batch gradient, and the optimizer update proceeds exactly as if the large batch had been processed in one shot.

What happens if you forget the $1/K$ scaling? The accumulated gradient becomes $K$ times larger than intended, so the optimizer takes an update step $K$ times too large. This is equivalent to multiplying the learning rate by $K$, which typically causes training to diverge. This mistake is easy to make, especially when adapting a standard training loop. Always check that the loss passed to .backward() is divided by accumulation_steps before the call.