The problem accumulation solves

Batch size is bounded by memory, not by wishes. A training step must hold at once the model weights, the optimizer state, and — the part that scales with batch size — the activations saved during the forward pass so the backward pass can reuse them. Activation memory grows roughly linearly with the samples in flight, so doubling the batch roughly doubles the footprint. On a small GPU, or a CPU box training an SLM, you hit the ceiling long before the batch size the optimizer actually wants.

Gradient accumulation breaks the coupling between the batch size the optimizer sees and the batch size the hardware processes at one time. You pick a micro-batch b that fits and a number of accumulation steps K, and you get an effective batch of B = K × b — paying in wall-clock time what you cannot pay in memory, while the gradient you feed the optimizer is exactly the one a batch of size B would produce.

Advertisement

The accumulate-then-step loop

The mechanism is a small change to the ordinary training loop. Normally every iteration does four things: forward, backward, step, zero-the-gradients. Accumulation decouples the last two from the first two — you backward K times before you step once.

optimizer.zero_grad()
for k, micro_batch in enumerate(accum_group):   # K micro-batches
    loss = criterion(model(micro_batch)) / K     # note the / K
    loss.backward()                              # ACCUMULATES into .grad
# after K backwards, .grad holds the summed gradient
optimizer.step()
optimizer.zero_grad()

The load-bearing detail is that backward() in most frameworks adds into the .grad buffers rather than overwriting them. That is usually treated as a nuisance you cancel with zero_grad() every iteration — here it is the whole feature. You deliberately skip zeroing for K−1 iterations, let the gradients pile up, and only after the K-th backward do you step and then clear the buffers for the next group.

Advertisement

Why it works: gradient linearity

The equivalence is not an approximation; it is the linearity of differentiation. Take a batch of B samples with a mean loss L = (1/B) Σ_i ℓ_i, where ℓ_i is the loss of sample i. The gradient with respect to parameters θ is

∇L = ∇[ (1/B) Σ_i ℓ_i ] = (1/B) Σ_i ∇ℓ_i

Now split the B samples into K groups of size b. Because the sum can be regrouped freely, the total gradient is the sum of the group sub-sums — and each group’s contribution is computed independently, with only its own activations resident in memory at that moment. That is the memory win: you never hold all B sets of activations simultaneously, only b at a time, yet the gradient you assemble is bit-for-bit the one the full batch would give (up to floating-point summation order). The optimizer, seeing only the final .grad, cannot tell whether it came from one pass or K.

The 1/K averaging factor

Everything hinges on making the accumulated gradient a mean over B samples, not a sum. Deep-learning losses almost always use reduction='mean', so a single micro-batch of size b yields ∇[ (1/b) Σ ℓ ] — already divided by b. If you simply add up K such micro-batch gradients you get

Σ_k (1/b) Σ_{i in k} ∇ℓ_i   =   (K/B) Σ_i ∇ℓ_i   =   K · ∇L

That is K times too large. To recover the true batch mean ∇L = (1/B) Σ_i ∇ℓ_i you divide each micro-batch loss by K before its backward pass — the / K in the loop. Scaling the loss scales its gradient by the same constant (linearity again), and K scaled micro-batch gradients sum to exactly ∇L. The factor belongs on the loss, not the optimizer, so gradient clipping and logging downstream see correctly-scaled numbers.

The normalization bug

The single most common accumulation mistake is dropping the / K. Nothing crashes. The shapes are right, the loss curve still descends, and the code looks correct in review — which is exactly why it is dangerous. What you have actually done is feed the optimizer a gradient that is K times too big, which is indistinguishable from multiplying your learning rate by K.

With K = 8, a carefully tuned learning rate becomes an effective 8× learning rate. Sometimes that merely trains faster and hides the error; often it pushes an already-aggressive rate into divergence, loss spikes, or NaNs you then misdiagnose as a data or model problem. The tell is a run that was stable at K = 1 and blows up the moment you raise K. Always sanity-check that the pre-step gradient norm is roughly invariant to K — if it scales with K, your normalization is wrong.

A worked numeric example

Suppose you want an effective batch of B = 32 but only b = 8 fits, so K = 4. Take a toy scalar parameter and say the four micro-batches produce mean-loss gradients of 2.0, 4.0, 1.0, and 5.0 respectively.

The true batch-of-32 gradient is the mean of the four micro-batch means (equal sizes), (2 + 4 + 1 + 5) / 4 = 3.0. In the loop, each micro-batch loss is divided by K = 4, so its gradient contribution becomes 0.5, 1.0, 0.25, 1.25; these accumulate to 0.5 + 1.0 + 0.25 + 1.25 = 3.0 — the correct batch gradient. Forget the / 4 and the buffer instead holds 2 + 4 + 1 + 5 = 12.0, four times too large. Same descent direction, quadrupled magnitude — the LR-inflation bug in one number.

Memory versus throughput

Accumulation is a pure trade of time for space. Peak activation memory is set by the micro-batch b, not the effective batch B — that is the entire benefit, and it is unbounded in principle: you reach any B by raising K. Throughput, however, does not improve; if anything it dips slightly, since you do the same total FLOPs as a real batch of B spread over K sequential passes plus a little per-pass overhead.

So accumulation does not make training faster — a true large batch on bigger hardware would finish the same step sooner because it parallelizes the micro-batches instead of serializing them. What accumulation buys is access: training at a batch size whose optimization behavior you want on hardware that could never hold it at once. For CPU SLM training, where memory is tight and you are time-rich anyway, that is usually the right trade.

Interaction with the learning rate

Because accumulation changes the effective batch size, it interacts directly with learning-rate tuning — and this is separate from the normalization bug. Larger batches give lower-variance gradient estimates, which generally tolerate (and often need) a larger learning rate to make comparable progress per step. The common heuristics are linear scaling (lr ∝ B) and square-root scaling (lr ∝ √B); which fits depends on the optimizer and regime.

The practical point: if you introduce accumulation to raise B from 32 to 256, you have changed the optimization problem, and the rate tuned for 32 is probably too small for 256. This is a legitimate reason to raise the LR, and must not be confused with the illegitimate K× inflation from a missing / K. Keep the two separate: divide correctly first so the gradient is the honest batch mean, then deliberately scale LR for the new batch size.