RoBERTa › Larger Batches

BERT-base was trained with a batch size of 256 sequences. RoBERTa found that much larger batches improve both training stability and final performance. The key insight is that larger batches provide better gradient estimates, especially important for MLM where only 15% of tokens contribute to each update.

To understand why batch size matters so much for MLM specifically, consider what a gradient update represents. When you compute the gradient of the MLM loss over a batch of sequences, you are computing a noisy estimate of the true gradient over the entire dataset. Noise comes from two sources: variation in which sequences appear in the batch, and variation in which tokens within those sequences happen to be masked. With a small batch size, both sources of noise are large. The model takes many small, uncertain steps, each influenced heavily by the specific accidents of which sentences and which masks appeared in that particular batch.

Larger batches average out both forms of noise simultaneously. With a batch of 8192 sequences, you're averaging gradients over millions of masked token predictions, giving you a much more reliable signal about which direction to update the parameters. The model takes fewer, more confident steps. This is especially useful early in training when the parameters are far from their good values and you want each update to make steady progress rather than random-walking around the loss surface.

Why do larger batches help MLM specifically? Consider the gradient signal per update:

• With batch size 256 and 15% masking, each update aggregates gradients from roughly $256 \times 512 \times 0.15 \approx 19,700$ masked tokens
• With batch size 8192, this grows to $8192 \times 512 \times 0.15 \approx 629,000$ masked tokens

The larger gradient estimates are less noisy, allowing the optimizer to take more confident steps. This translates to faster convergence and better final performance.

RoBERTa used batch sizes up to 8192 sequences. To maintain the same number of parameter updates, larger batches require adjusting the learning rate. The linear scaling rule suggests multiplying the learning rate by the batch size increase factor, though warmup and careful tuning remain important.

The linear scaling rule has an intuitive justification. If you double the batch size, each gradient update now averages over twice as many examples. The variance of your gradient estimate halves, which means you can afford to take a larger step without risking instability. Empirically, multiplying the learning rate proportionally to the batch size increase tends to achieve similar convergence behavior while using fewer total gradient updates. The warmup schedule remains important because the linear scaling rule breaks down at initialization: when the model is very far from a good solution, large gradient steps in unpredictable directions can destabilize training regardless of how well-estimated those gradients are.

In practice, running experiments with batch sizes of 8192 requires either a large number of GPUs or gradient accumulation, where you compute gradients over several smaller mini-batches and accumulate them before applying an update. The latter approach is computationally equivalent to a large batch (the math is the same) and allows you to simulate RoBERTa-style training on hardware with limited GPU memory. Most modern training frameworks support gradient accumulation natively.