Gradient Accumulation › A Worked Example

To make the mechanics concrete, consider a simple linear regression scenario that illustrates the mathematical equivalence between full-batch training and gradient accumulation.

Suppose you have a dataset of 8 samples $\{(x_i, y_i)\}_{i=1}^{8}$ and a model parameterized by a single weight $w$, with mean squared error loss:

$$\mathcal{L}(w; B) = \frac{1}{|B|} \sum_{i \in B} (w x_i - y_i)^2$$

where:

• $w$: the scalar model parameter
• $B$: the batch of samples
• $x_i, y_i$: the $i$-th input-output pair
• $|B|$: the batch size

The full-batch gradient over all 8 samples is:

$$g_{\text{full}} = \frac{\partial \mathcal{L}}{\partial w} = \frac{2}{8} \sum_{i=1}^{8} x_i(wx_i - y_i)$$

Now split the 8 samples into $K = 4$ micro-batches of size 2 each. For micro-batch $k$ (samples $2k-1$ and $2k$), the mean loss is:

$$\mathcal{L}_k = \frac{1}{2}\left[(wx_{2k-1} - y_{2k-1})^2 + (wx_{2k} - y_{2k})^2\right]$$

The scaled gradient contribution from micro-batch $k$ (after dividing by $K = 4$) is:

$$\mathbf{g}_k = \frac{1}{K} \cdot \frac{\partial \mathcal{L}_k}{\partial w} = \frac{1}{4} \cdot \frac{2}{2}\left[x_{2k-1}(wx_{2k-1} - y_{2k-1}) + x_{2k}(wx_{2k} - y_{2k})\right]$$

Summing across all 4 micro-batches to obtain the accumulated gradient:

$$\begin{aligned} g_{\text{accum}} &= \sum_{k=1}^{4} \mathbf{g}_k \\ &= \frac{1}{4} \sum_{k=1}^{4} \left[x_{2k-1}(wx_{2k-1} - y_{2k-1}) + x_{2k}(wx_{2k} - y_{2k})\right] \\ &= \frac{1}{4} \sum_{i=1}^{8} x_i(wx_i - y_i) \\ &= \frac{2}{8} \sum_{i=1}^{8} x_i(wx_i - y_i) \\ &= g_{\text{full}} \end{aligned}$$

The accumulated gradient equals the full-batch gradient exactly. This confirms that the $1/K$ scaling is necessary and sufficient to achieve mathematical equivalence.

The same derivation applies to any differentiable loss and any model architecture. The linearity of gradient computation means that you can always decompose a full-batch gradient into a sum of micro-batch gradients. The only requirement is that the micro-batches partition the full batch without overlap, and that each micro-batch uses the same model parameters.

In some settings, the last micro-batch of an epoch may be smaller than the others, because the dataset size is not perfectly divisible by the full logical batch size. For example, if you have 1000 samples, a logical batch of 64, and micro-batch size 16, then 15 logical batches will have all four micro-batches of size 16, but the final logical batch may have only 1000 mod 64 = 40 samples, requiring some micro-batches of size 10 or fewer.

In this case, the correct approach is to weight each micro-batch by its actual size rather than using uniform $1/K$ scaling. The properly weighted contribution of micro-batch $k$ with $m_k$ samples is:

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

where $M = \sum_{k=1}^{K} m_k$ is the total number of samples in the logical batch. This produces an unbiased estimate of the full-batch gradient even when micro-batches have different sizes.

Most training frameworks simplify this by dropping the last incomplete batch entirely, which avoids the unequal weighting problem. Whether to drop or include it is a hyperparameter choice: dropping it is simpler and slightly wastes data, while including it with proper weighting is more data-efficient but requires additional bookkeeping.