Gradient accumulation: count every example

The maths

A small forward and backward pass

$$h=w_1x,\quad y=w_2h,\quad L=\tfrac12(y-t)^2,\quad \frac{\partial L}{\partial w_1}=(y-t)w_2x,\quad \frac{\partial L}{\partial w_2}=(y-t)h$$

Both derivatives use the same old weights. The film traces every factor on the first example.

Weight examples equally

$$\bar g=\frac{\sum_b\sum_{i\in B_b}g_i}{\sum_b |B_b|}=\frac{\sum_b |B_b|\bar g_b}{N}$$

If each microbatch already reports a mean, multiply that mean by its example count before the final division by N.

Why a mean of means can be wrong

$$\tilde g=\frac1M\sum_{b=1}^{M}\bar g_b,\qquad \text{example weight in batch }b=\frac{1}{M|B_b|}$$

With microbatch sizes one and two, the weights become one half, one quarter and one quarter instead of one third each.

One optimizer step

$$w^+=w-\eta\bar g,\qquad (2,-1)-0.1(5/3,-10/3)=(11/6,-2/3)$$

The toy uses plain SGD. Keep the parameter snapshot fixed while accumulating this effective batch, then update once.

Transcript

The narration, chapter by chapter. A timestamp opens the film at that moment.

0:00Every example counts

Three examples can arrive as one microbatch and then two. If every example should count equally, the split must not decide how far we move the weights. Let us test that.

0:11A tiny forward pass

Our tiny network multiplies an input by two, then by minus one. With input one and target zero, its prediction is minus two, and half the squared error is two.

0:22Trace the gradient

Backpropagation carries the error through both multiplies. This example contributes gradients two and minus four. Compute the other examples at those same weights, so their gradients can be added before any update.

0:35Add microbatch sums

The first microbatch has one example. The second has two, whose gradient sum is three and minus six. Across all three examples, the accumulated sum is five and minus ten.

0:47The mean-of-means trap

Here is the trap: averaging the two microbatch means gives the single example half the influence. It should receive one third. Changing the grouping can now change the update, even though the data is identical.

1:01Divide by examples

Instead, divide the total gradient sum by three examples. The correct mean is five thirds and minus ten thirds. With learning rate zero point one, make exactly one optimizer step after that division.

1:16One update, same result

The new weights are about one point eight three three and minus zero point six six seven. Try another split below: the correct result stays fixed. This scalar example uses plain gradient descent and excludes stateful batch layers.

Also available as captions (WebVTT), the video file and a Markdown copy of this page.

Sources and model assumptions

Follow the original mechanism behind this explainer. The interactive examples identify their toy data and simplifying assumptions above.

• PyTorch: gradient accumulation and effective-batch optimizer steps
• PyTorch: automatic differentiation and the chain rule

Related terms

• TrainingTraining is the process that produces a model: showing it enormous amounts of text and adjusting its parameters until it gets good at predicting what comes next. It happens once, before you ever use the model.
• ParametersParameters are the learned numbers (weights) inside a model that hold everything it appears to know. The count of them is what people mean by model size, and they are fixed once training ends.
• ModelA model is the trained artifact at the centre of every AI coding tool: a large file of numbers (parameters) that, given some text, produces the most likely continuation. When people say "which model are you using," this is the thing they mean.