ML0009 Batch Size Selection

Written by

What are the best strategies for selecting the appropriate batch size?

Answer

Selecting an appropriate batch size is a crucial hyperparameter choice that trades training efficiency against optimization behavior. The practical strategy: start with a moderate value (16, 32, or 64) and adjust based on memory, gradient stability, and validation performance. Three factors drive the choice: memory constraints (larger batches need more GPU memory), dataset size (large datasets can sustain large batches; small datasets often benefit from the extra variability of small batches), and learning rate interaction (large batches usually allow (or require) a proportionally higher learning rate). Large batches train faster per epoch with stabler gradient estimates and higher memory cost, but can converge to sharp minima that generalize worse; small batches update more often with noisier gradients, which can explore flatter minima and generalize better at lower memory cost.

(1) Trade-Off Summary: Large batch = fast and stable but possibly sharp minima; small batch = slow and noisy but often better generalization.
(2) Practical Starting Points: Try 16/32/64 first, then scale with memory and dataset size.
(3) Interaction With LR: Scale the learning rate roughly proportionally when you change the batch size.

Optimization trajectories of small versus large batch size on a loss surface

Figure 1: Same loss surface, two batch sizes: the small-batch path is noisy but wanders into a wider basin, while the large-batch path is smooth but direct.

Mathematical Formulation:
$\hat{g} = \frac{1}{B}\sum_{i=1}^{B}\nabla_\theta \mathcal{L}_i$
$\mathrm{Var}[\hat{g}] \propto \frac{\sigma^2}{B}$

Where:

$\hat{g}$ is the gradient estimate computed from one mini-batch.
$B$ is the batch size, and $i\in\{1,\ldots,B\}$ indexes the samples in the batch.
$\nabla_\theta \mathcal{L}_i$ is the per-sample loss gradient.
$\sigma^2$ is the per-sample gradient variance; doubling $B$ halves the gradient noise, which is why large batches give stable estimates and small batches give noisy ones.