Full Fine-tuning › Batch Size Effects

Batch size interacts with training in multiple ways, affecting memory usage, gradient quality, and convergence behavior. Understanding these interactions helps practitioners make informed choices when configuring their training runs. The batch size decision is rarely made in isolation: it must be considered alongside learning rate, gradient accumulation, and available hardware.

Think of batch size as the number of voters you poll before making a policy decision. Polling more voters gives you a more accurate estimate of public opinion, but it takes more time and resources. If you poll only 8 people (small batch), your estimate will be noisy and might lead you in the wrong direction on any given decision. If you poll 256 people (large batch), your estimate is reliable but each poll is expensive. The right number depends on how much you can afford per decision and how much noise you can tolerate. Importantly, with noisier estimates, you should be more conservative in how sharply you act on each one, which is why smaller batches pair with smaller learning rates.

The primary constraint on batch size is GPU memory. Each sample in a batch requires memory for input tensors (token IDs, attention masks), intermediate activations (stored for backward pass), gradients for all parameters, and optimizer states (momentum buffers for Adam). For a model like BERT-base with sequences of length 512, batch sizes of 8-16 are common on consumer GPUs with 12-16GB of memory.

The memory requirements break down into several components:

• Input tensors: the tokenized representation of each text sample, including token IDs and attention masks
• Intermediate activations: the outputs of each layer during the forward pass, cached for gradient computation during the backward pass
• Gradients: the partial derivatives of the loss with respect to each parameter, computed during backpropagation
• Optimizer states: additional buffers that adaptive optimizers like Adam maintain, specifically momentum and variance estimates for each parameter

The optimizer state memory deserves special attention. Adam and its variants maintain two additional state vectors per parameter: a running mean of gradients (first moment) and a running mean of squared gradients (second moment). For a model with $P$ parameters stored in 32-bit precision, the parameter storage alone requires $4P$ bytes. The optimizer states add another $8P$ bytes (4 bytes each for first and second moments), bringing the total memory for parameters plus optimizer states to $12P$ bytes. For BERT-base with 110 million parameters, this amounts to approximately 1.3 GB just for parameters and optimizer states, before counting activations and inputs.

Larger models or longer sequences require smaller batches or memory optimization techniques like gradient checkpointing, which trades computation for memory by recomputing activations during the backward pass rather than storing them. Gradient checkpointing can reduce memory requirements materially at the cost of approximately 30-40% more computation during the backward pass.

Smaller batches produce noisier gradient estimates. Each gradient is computed from fewer examples, which makes it a less accurate estimate of the true gradient over the entire dataset. This noise has both advantages and disadvantages that practitioners should understand.

Advantages of noise: Gradient noise can help escape local minima and may give implicit regularization, preventing overfitting. Some research suggests noisy gradients improve generalization by preventing the model from settling into sharp minima that may not transfer well to test data. For fine-tuning tasks with small datasets, the implicit regularization from noisy gradients can be beneficial, helping the model avoid memorizing the training set.

Disadvantages of noise: Excessive noise makes training unstable and can prevent convergence. Very small batches may require learning rate reductions that slow training sharply. The model may oscillate instead of making consistent progress toward the optimum. When the gradient noise is larger than the gradient signal, the optimizer is in effect making random steps, which is inefficient and potentially harmful.

The appropriate level of gradient noise depends on your task and dataset size. For very small datasets (a few hundred examples), using small batches gives more gradient noise, which can act as regularization. For larger datasets where the main concern is training efficiency rather than overfitting, larger batches and higher learning rates are preferable.

When memory constraints prevent using the desired batch size, gradient accumulation gives a workaround. Instead of updating weights after each batch, we accumulate gradients over multiple batches, then update. This technique effectively decouples the memory requirements (determined by the batch that fits in GPU memory) from the effective batch size (determined by how many batches we accumulate before updating).

If you want an effective batch size of 32 but can only fit 8 samples in memory, you accumulate gradients over 4 batches. Gradient accumulation allows us to simulate larger batch sizes by computing multiple forward and backward passes before performing a single parameter update. The relationship is expressed by the following formula:

$$B_{\text{effective}} = B_{\text{device}} \times N_{\text{accum}}$$

where:

• $B_{\text{effective}}$: the effective batch size used for parameter updates (32 in this example)
• $B_{\text{device}}$: the batch size that fits in device memory (8 in this example)
• $N_{\text{accum}}$: the number of gradient accumulation steps (4 in this example)
• The multiplication: combines multiple small batches into one large gradient update, trading time for memory

This formula shows that the effective batch size equals the device batch size multiplied by the number of accumulation steps. The key insight is that gradient accumulation computes $N_{\text{accum}}$ forward and backward passes, summing the gradients before performing a single parameter update. This produces mathematically equivalent results to processing all $B_{\text{effective}}$ samples simultaneously, but uses only $B_{\text{device}}$ worth of memory at any given time. The trade-off is computational: we perform more forward and backward passes, but each pass uses less memory.

Applying this formula to our example shows how the pieces fit together:

$$\begin{aligned} B_{\text{effective}} &= B_{\text{device}} \times N_{\text{accum}} && \text{(formula)} \\ &= 8 \times 4 && \text{(substitute values)} \\ &= 32 && \text{(compute)} \end{aligned}$$

Larger batches benefit from higher learning rates. The intuition is straightforward: with more samples, the gradient estimate is more reliable, so we can take larger steps confidently. When we compute a gradient from 32 samples instead of 16, the estimate is closer to the true gradient over the full dataset, and we can trust it more. The linear scaling rule suggests that learning rate should scale proportionally with batch size according to the following relationship:

$$\text{lr}_{\text{new}} = \text{lr}_{\text{base}} \times \frac{B_{\text{new}}}{B_{\text{base}}}$$

where:

• $\text{lr}_{\text{new}}$: the adjusted learning rate for the new batch size
• $\text{lr}_{\text{base}}$: the baseline learning rate that was validated with the baseline batch size
• $B_{\text{new}}$: the new batch size you want to use
• $B_{\text{base}}$: the baseline batch size used during validation
• The ratio $\frac{B_{\text{new}}}{B_{\text{base}}}$: the scaling factor by which the batch size changes

This formula adjusts the learning rate proportionally to changes in batch size. The linear scaling rule works because doubling the batch size doubles the number of samples adding to each gradient, making the gradient estimate twice as stable. This stability allows us to take steps that are twice as large (double the learning rate) while maintaining similar training dynamics. The rule preserves the expected magnitude of parameter updates per training sample, keeping the optimization trajectory similar despite different batch sizes.

However, this rule has limits. Very large batches (thousands of samples) often fail to benefit from proportionally scaled learning rates. The relationship becomes sublinear, and other techniques like LAMB (Layer-wise Adaptive Moments) become necessary for successful training. The breakdown occurs because extremely large batches can converge to sharp minima that generalize poorly, and at some point the additional samples give diminishing marginal benefit to gradient quality. Practitioners using batch sizes above 256 for fine-tuning should validate that the linear scaling rule still holds for their specific configuration.

For fine-tuning with typical batch sizes (16-64), the linear scaling rule gives reasonable guidance. If you've validated that a learning rate of 2e-5 works well with batch size 16, then batch size 32 might benefit from 4e-5. Let's verify this by applying the linear scaling rule:

$$\begin{aligned} \text{lr}_{\text{new}} &= \text{lr}_{\text{base}} \times \frac{B_{\text{new}}}{B_{\text{base}}} && \text{(linear scaling rule)} \\ &= 2 \times 10^{-5} \times \frac{32}{16} && \text{(substitute values)} \\ &= 2 \times 10^{-5} \times 2 && \text{(simplify fraction)} \\ &= 4 \times 10^{-5} && \text{(multiply)} \end{aligned}$$

This calculation shows that doubling the batch size from 16 to 32 suggests doubling the learning rate from 2e-5 to 4e-5. The linear scaling rule gives a principled starting point, though practitioners should still validate that the scaled learning rate works well for their specific task.