LARS, LAMB and Large-Batch Training

Training with a much bigger batch size — thousands of examples per update instead of a few hundred — lets you use far more parallel hardware and finish training faster, but naive scaling breaks down: if you just crank up the batch size with a standard optimizer, training becomes unstable or converges to a worse final result. The problem is that different layers of a deep network have very different scales of weights and gradients, and a single global learning rate that works for one layer can be far too aggressive or too timid for another once the batch (and thus the step size) grows large.

LARS (for plain SGD) and LAMB (for Adam-style updates) fix this per layer: each layer's update is rescaled by the ratio of the norm (overall size) of its current weights to the norm of its proposed update, so no single layer's weights get either blown up or left nearly frozen just because the batch size changed. This layer-wise trust-region-like scaling is what lets batch sizes scale up to tens of thousands without retraining from scratch to find new hyperparameters.

A worked example

Training a large transformer, going from a batch size of 512 to 32,768 with a standard optimizer typically requires re-tuning the learning rate from scratch and often still degrades final accuracy. Using LAMB, researchers have trained the same models with batch sizes over 32,000 while matching baseline accuracy and cutting wall-clock training time roughly linearly with the added parallelism — turning a multi-day training run into a few hours by using far more machines at once.