Relationship Between Batch Size and Learning Rate

As we tune hyperparameters to optimize model training, two settings that often go hand-in-hand are the batch size and the learning rate. Changing one frequently necessitates adjusting the other for optimal performance. Let's examine why these two are so closely linked.

Recall that Mini-batch Gradient Descent computes the gradient estimate using a small subset (a mini-batch) of the training data in each step. The size of this mini-batch, $B$, directly impacts the quality and characteristics of the gradient estimate.

Gradient Noise and Batch Size

A fundamental aspect of Stochastic Gradient Descent (SGD) and its mini-batch variant is the noise introduced by using only a subset of the data to estimate the gradient.

• Small Batch Sizes: Using a small batch size (e.g., 16, 32, 64) results in a gradient estimate with high variance. Each batch can provide a significantly different gradient direction due to the small sample size. This noise can sometimes be beneficial, potentially helping the optimizer escape poor local minima or saddle points. However, it also makes the convergence path erratic.
• Large Batch Sizes: Conversely, using a large batch size (e.g., 256, 512, 1024 or more) leads to a gradient estimate with lower variance. The gradient calculated from a large batch is a better approximation of the true gradient over the entire dataset. This results in a smoother convergence path but potentially requires more computation per step and might converge to sharper minima, which sometimes generalize less effectively.

The variance of the mini-batch gradient estimate, $\hat{g}_B$, is roughly inversely proportional to the batch size $B$:

$$Var(\hat{g}_B) \approx \frac{1}{B} Var(g_i)$$

where $g_i$ represents the gradient for a single data sample. A larger $B$ reduces the variance (noise).

How Learning Rate Interacts with Gradient Noise

The learning rate, $\alpha$, determines the step size taken in the direction of the (estimated) negative gradient. The appropriate learning rate is closely tied to the reliability (noise level) of this gradient estimate.

• With Noisy Gradients (Small Batches): If the gradient direction fluctuates wildly from batch to batch, taking large steps (a high learning rate) can cause the optimizer to overshoot the minimum or even diverge. The updates might bounce around erratically. A smaller learning rate is often necessary to average out the noise and make steady progress towards a minimum.
• With Smoother Gradients (Large Batches): When the gradient estimate is more stable and accurately reflects the true gradient direction, larger steps can be taken more confidently. A higher learning rate can lead to faster convergence without the risk of divergence seen with noisy gradients.

The Linear Scaling Rule: A Common Heuristic

Based on the relationship between batch size, gradient noise, and appropriate step size, a popular heuristic has emerged, often called the Linear Scaling Rule:

If you multiply the batch size by a factor $k$, multiply the learning rate by the same factor $k$.

Intuition: If you use a batch size $k$ times larger, your gradient estimate is based on $k$ times more data and is therefore (roughly) $k$ times less noisy or, viewed differently, the sum of gradients in the larger batch is $k$ times larger in magnitude if individual gradients were similar. To maintain a similar update magnitude or to take advantage of the more reliable gradient direction, you can increase the learning rate proportionally. For example, if you switch from a batch size of 64 to 256 (a $4 \times$ increase), this rule suggests increasing the learning rate by $4 \times$.

Caveats: The Linear Scaling Rule is a heuristic, not a strict law.

• It often works best with SGD with Momentum. Adaptive optimizers like Adam might be less sensitive.
• It tends to break down for very large batch sizes, where generalization performance can degrade.
• It's often recommended to use this scaling in conjunction with a learning rate warmup period (discussed earlier), especially when starting with large learning rates, to prevent instability early in training.

Trade-offs and Practical Considerations

Choosing the batch size and learning rate involves several trade-offs:

1. Computational Cost: Larger batches can utilize parallel processing capabilities (like GPUs) more effectively, potentially leading to faster wall-clock time per epoch. However, they require more memory. Smaller batches require less memory but might lead to slower training overall if the hardware is underutilized or if more epochs are needed due to smaller learning rates.
2. Generalization Performance: This is an area of active research. Some findings suggest that smaller batch sizes introduce beneficial noise that helps the optimizer find flatter minima, which often generalize better than the sharper minima sometimes found by large-batch training. However, large-batch training can sometimes be tuned to achieve comparable results.
3. Convergence Speed: Larger batches paired with appropriately scaled larger learning rates can significantly speed up convergence in terms of the number of epochs required.

Illustration of loss curves for different batch size and learning rate combinations. Large batch/large learning rate often converges fastest initially, while small batch/large learning rate can become unstable.

Tuning Strategy

When tuning these hyperparameters:

• Start with a common batch size: Choose a batch size that fits your hardware memory (e.g., 32, 64, 128, 256) and find a good learning rate for it.
• Consider scaling: If you need to change the batch size (e.g., for memory reasons or to experiment with generalization), try adjusting the learning rate using the linear scaling rule as a starting point.
• Fine-tune: The linear scaling rule provides a starting point, not necessarily the final optimal value. You will likely need to perform some fine-tuning of the learning rate after changing the batch size.
• Monitor: Always monitor your training and validation loss curves. Instability (loss increasing or fluctuating wildly) often indicates the learning rate is too high for the chosen batch size. Very slow convergence might mean the learning rate is too low.

Understanding the relationship between batch size and learning rate is essential for efficient hyperparameter tuning. While heuristics like the linear scaling rule offer guidance, empirical testing and careful monitoring remain necessary to find the combination that works best for your specific model and dataset.