How Batch Size Changes Training

What is Batch Size?

Batch size is the number of data samples processed together in a single forward pass before calculating loss and updating weights.

1. Batch Size = 1 (Pure SGD): Updates weights after every single sample. High gradient noise, slow parallel hardware usage.
2. Mini Batch (Size 32 to 512): Standard sweet spot for single GPU training. Good balance between gradient noise and GPU parallel speed.
3. Large Batch (Size 4096 to 65536): Used in large scale distributed training across GPU clusters. Maximizes hardware throughput.

Small vs Large Batch Size Tradeoffs

┌──────────────────────────┬──────────────────────────┐
│ SMALL BATCH SIZE (32)    │ LARGE BATCH SIZE (4096)  │
├──────────────────────────┼──────────────────────────┤
│ High gradient noise.     │ Low gradient noise.      │
│ Acts as regularization.  │ Fast GPU tensor compute. │
│ Finds flat minima.       │ Risk of sharp minima.    │
│ Slower wall clock time.  │ High memory requirement. │
└──────────────────────────┴──────────────────────────┘

1. Small Batch Advantage: Implicit Regularization

Gradient updates from small mini batches are noisy because 32 samples are a rough approximation of the full dataset.

This gradient noise is actually helpful. It acts as implicit regularization, kicking the optimizer out of narrow, sharp local minima into broad, flat minima that generalize better to unseen test data.

2. Large Batch Advantage: GPU Parallel Efficiency

Modern GPUs excel at matrix operations on large data tensors.

Processing 4096 samples in a single batch saturates GPU tensor cores efficiently, resulting in much faster wall clock training time per epoch.

The Linear Scaling Rule

If you increase batch size by factor $K$ (for example, from 32 to 256, so $K = 8$), each step takes a cleaner average, but you perform fewer total updates per epoch.

To compensate, apply the Linear Scaling Rule (Goyal et al., 2017):

new_learning_rate = base_learning_rate * K

When scaling up batch size by $K$, multiply the learning rate by $K$ alongside a Linear Warmup schedule to maintain stable training.

Say this out loud

Batch size sets the number of samples processed per weight update. Small batch sizes introduce helpful gradient noise that acts as implicit regularization, helping models find flat minima that generalize better. Large batch sizes maximize GPU parallel hardware throughput for faster wall clock training, but require scaling the learning rate linearly to maintain stability.

Followups to expect

1. What is Gradient Accumulation? Simulating large batch sizes on a small GPU by running multiple forward and backward passes without clearing gradients, then stepping the optimizer once every N steps.
2. Why do flat minima generalize better than sharp minima? Flat minima are tolerant to small shifts between training data distribution and test data distribution. Sharp minima crash in performance under tiny input shifts.