Gradient Accumulation

A batch of 256 won't fit in GPU memory, but 4 micro-batches of 64 will. Gradient accumulation sums gradients over micro-batches and steps the optimiser once. Done right, it gives exactly the same update as the big batch.

The model is linear regression with mean squared error, $\mathcal{L} = \frac{1}{N}\sum (Xw - y)^2$, so the gradient is $\frac{2}{N}X^\top(Xw - y)$.

Implement:

• mse_grad(X, y, w): the gradient for one batch.
• accumulated_grad(X, y, w, micro_batch): split the rows into consecutive micro-batches of size micro_batch (the last may be smaller), and return the gradient equal to the full-batch one. Weight each micro-batch's gradient by its share of the rows, since dividing by the number of micro-batches is wrong when the last one is smaller.
• naive_accumulated_grad(X, y, w, micro_batch): the common buggy version that averages the micro-batch gradients equally, for comparison.