Neural networks, deep learning, NLP · insight-1406
Measure steps to target at several batch sizes before paying for bigger batches
- Evidence type
- empirical finding
- Difficulty
- hard
- Roles
- ML engineer
- Industries
- cross-industry
Summary: situation, action, reason, boundary
Situation
A team wants to move a training job from 4 to 32 accelerators and assumes the run will finish about 8 times sooner.
Action
For three or four batch sizes, tune the learning rate separately and record steps to reach the same validation target. Plot steps against batch size and buy hardware only along the part of the curve that still falls.
Reason
Shallue and colleagues measured perfect scaling, then diminishing returns, then a maximum useful batch size, and found it depends strongly on the model and data.
Boundary
Planning distributed training, or choosing between gradient accumulation and more devices.
Symptom or risk
Beyond some batch size the number of steps needed stops falling, so extra hardware buys little, and the learning rate chosen for one batch size is wrong for the next.
Exceptions
Very small models that fit on one device rarely need the experiment.
Trade-offs
The sweep itself costs several short runs; use a truncated target to keep it cheap.
Pitfalls
Fixing the learning rate across batch sizes measures your tuning, not the batch size.
Lessons and concepts
Claims and their sources
Shallue et al. report that the effect of batch size on steps to a target error varies greatly between workloads and that conflicting prior findings often came from differences in tuning and compute budgets.