Large Batch Training › The LAMB Optimizer

The linear scaling rule works well for SGD and its momentum variants, but modern large-scale models are trained with adaptive optimizers like Adam. Scaling Adam to large batch sizes introduces a different challenge, and the LAMB optimizer was developed specifically to address it.

Adam maintains per-parameter learning rates, scaling each parameter's update by the square root of its historical gradient variance. This adaptive rescaling lets Adam work reliably across diverse architectures and tasks. It handles the varying gradient magnitudes across parameters automatically, requiring less hyperparameter tuning than SGD.

When you scale the batch size with Adam and apply the linear scaling rule to the overall learning rate, something subtle goes wrong. Adam's adaptive update for parameter $\theta_j$ is roughly:

$$\theta_j \leftarrow \theta_j - \frac{\eta \cdot m_j}{\sqrt{v_j} + \epsilon}$$

where:

• $m_j$: the first moment (exponential moving average of gradients)
• $v_j$: the second moment (exponential moving average of squared gradients)
• $\epsilon$: a small constant for numerical stability

The problem is that the ratio $m_j / (\sqrt{v_j} + \epsilon)$ is essentially normalized: it tells you which direction to move but not how far. The global learning rate $\eta$ controls the step magnitude, but it applies uniformly to all parameters.

Different layers in a deep network have very different gradient magnitudes. Embedding layers may have gradients many orders of magnitude smaller than the final output layer. When you apply a single global learning rate, you are either learning too fast in the final layers or too slowly in the embeddings. This tension becomes acute at large learning rates, which is what large batch training requires.

With small batches and modest learning rates, this imbalance is tolerable because the absolute update magnitude stays small enough that no layer diverges. But when you multiply the learning rate by a factor of 64 or 256, a layer that was marginally stable becomes explosively unstable. The layer-agnostic nature of the global learning rate, which Adam handles gracefully at moderate scales, becomes a critical failure mode at large learning rates.

There is also a subtler issue: the second moment estimate $v_j$ in Adam accumulates gradient variance over time. At the start of training with a large batch, the gradient estimates are already low-variance (that is the point of the large batch). This means the second moment estimates are small, making the $1 / (\sqrt{v_j} + \epsilon)$ factor large, amplifying the effective per-parameter learning rate at exactly the moment when stability is most critical.

Before LAMB, there was LARS (Layer-wise Adaptive Rate Scaling), developed for large batch training of convolutional networks. LARS scales the learning rate for each layer based on the ratio of the layer's weight norm to its gradient norm:

$$\eta_l = \eta \cdot \frac{\|\theta_l\|}{\|g_l\|}$$

where:

• $\eta_l$: the effective learning rate for layer $l$
• $\eta$: the global learning rate
• $\|\theta_l\|$: the L2 norm of the parameters in layer $l$
• $\|g_l\|$: the L2 norm of the gradients for layer $l$

The intuition is that the update should be proportional to the current weight scale of the layer. A layer with large weights can tolerate larger absolute updates without destabilizing. Conversely, a layer with small gradients relative to its weights is already learning slowly; the trust ratio gives it permission to take a larger effective step.

LARS also provides a natural safety check: if the gradient norm becomes very large (indicating an unstable update direction), the trust ratio shrinks automatically, reducing the effective learning rate for that layer. This self-limiting behavior prevents the cascading divergence that can occur when one layer's unstable gradients propagate through the network.

LARS worked well for training ResNets on ImageNet with very large batch sizes (up to 32,768 or more) without loss of accuracy. However, it was designed for SGD-style updates and did not incorporate the adaptive learning rate logic of Adam. For transformer architectures with weight tying, mixed precision, and complex gradient flows, pure SGD is rarely competitive with Adam. A solution that combined LARS's layer-level safety with Adam's per-parameter adaptivity was needed.

LAMB (Layer-wise Adaptive Moments optimizer for Batch training), introduced by You et al. (2019) for BERT pre-training, combines Adam's adaptive per-parameter scaling with LARS's layer-wise trust ratio.

The LAMB update rule is:

$$\begin{aligned} m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \\ \hat{m}_t &= \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \\ r_t &= \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \\ \theta_t &= \theta_{t-1} - \eta \cdot \frac{\|\theta_{t-1}\|}{\|r_t + \lambda \theta_{t-1}\|} \cdot (r_t + \lambda \theta_{t-1}) \end{aligned}$$

where:

• $g_t$: the gradient at step $t$
• $m_t$: the first moment (bias-corrected: $\hat{m}_t$)
• $v_t$: the second moment (bias-corrected: $\hat{v}_t$)
• $\beta_1, \beta_2$: exponential decay rates (typically 0.9 and 0.999)
• $r_t$: the Adam update direction (before layer scaling)
• $\lambda$: the weight decay coefficient
• $\|\theta_{t-1}\| / \|r_t + \lambda \theta_{t-1}\|$: the trust ratio (the LARS-style layer scaling)

The trust ratio $\phi_t = \|\theta_{t-1}\| / \|r_t + \lambda \theta_{t-1}\|$ scales each layer's update. When the Adam update direction $r_t$ has small norm relative to the weight norm, the trust ratio is large, allowing the layer to take bigger steps. When the update is large relative to the weights, the trust ratio shrinks, giving a safety net against destabilizing steps. This self-regulating property is what makes LAMB stable at the large learning rates required by large batch training.

The key insight behind LAMB's effectiveness is that the trust ratio decouples the update direction from the update magnitude. Adam computes an excellent direction for each parameter (using the first and second moment estimates). LAMB then scales the magnitude of the resulting layer-wise update to be appropriate for that layer's current weight scale.

This two-level structure mirrors the intuition behind the linear scaling rule but operates at finer granularity. The global learning rate $\eta$ is scaled up proportionally with batch size. But each layer's effective step is additionally modulated by its own trust ratio, preventing any single layer from taking catastrophically large steps when $\eta$ is large.

The trust ratio also adapts over training. Early in training, when weights are small (near initialization) and gradients are large (the model is far from convergence), trust ratios tend to be small, limiting the step size. As training progresses and the model approaches a good solution, weights grow and gradients shrink, and the trust ratio naturally allows larger steps. This dynamic adaptation is not an explicit design choice but an emergent property of the weight-norm-to-gradient-norm ratio.

In the original BERT paper and follow-up experiments, LAMB enabled training BERT-base in 76 minutes on 1,024 TPU chips using a batch size of 65,536, compared to roughly 3 days on 16 TPU chips with batch size 256. The key was that LAMB maintained model quality (BERT's downstream task performance on GLUE benchmarks) across this dramatic batch size increase. Without LAMB, Adam at the equivalent linearly scaled learning rate was unstable, causing training divergence or significant accuracy degradation.

Understanding when to use each optimizer guides practical decisions:

• Adam is the default for most tasks at moderate batch sizes (up to a few thousand). It handles heterogeneous gradient magnitudes across parameters naturally and is well-understood. Use Adam unless you have specific large-batch scaling requirements.
• LARS is appropriate for convolutional networks with SGD-style training at large batch sizes. It is simpler than LAMB and well-validated on image classification tasks. If your architecture does not benefit from Adam's adaptivity (ResNets trained with SGD and cosine decay, for instance), LARS is a solid choice.
• LAMB is the right tool for transformer models trained with Adam at very large batch sizes. The combination of Adam's per-parameter adaptivity and LARS's layer-wise safety net matches the needs of architectures like BERT and GPT where both are required for stable large-batch training.

The trust ratio in LAMB is computed per parameter tensor (one trust ratio per weight matrix or bias vector), not per individual parameter. This is intentional: computing a trust ratio per parameter would approach Adam's behavior (the second moment already provides per-parameter scaling). The layer-level granularity gives LAMB its distinct character.