LAMB Optimizer

LAMB rescales each layer's update to match the norm of its weights, enabling stable training at very large batch sizes.

Optimizing for huge batches

Training large models faster often means increasing the batch size so more work runs in parallel. But naively scaling the learning rate with the batch size makes training unstable, because different layers tolerate different step magnitudes. LAMB (Layer-wise Adaptive Moments for Batch training) was designed to make very large-batch training reliable, most famously cutting the wall-clock time to pretrain BERT by allowing batch sizes in the tens of thousands.

Layer-wise trust ratio

LAMB starts from an Adam-style update, computing bias-corrected first and second moment estimates and forming a candidate step for each parameter. It then applies a layer-wise trust ratio: for each layer it scales the update so that the ratio of the update's norm to the weights' norm is controlled. Concretely, it multiplies the layer's update by the norm of its weights divided by the norm of its raw update. A layer with large weights gets a proportionally larger step, and a layer with small weights a smaller one, keeping every layer's relative change consistent.

step = m_hat / (v_hat.sqrt() + eps) + weight_decay * w
trust = w.norm() / (step.norm() + eps)      # per layer
w = w - lr * trust * step

• Normalizes update size relative to weight size, per layer
• Keeps training stable when batches are very large
• Builds on the Adam moment estimates
• Enables high data-parallel throughput without loss of accuracy

When it helps

LAMB matters when the goal is to shorten training by scaling batch size across many accelerators, since the trust ratio prevents any single layer from taking a destabilizing step. At ordinary batch sizes it offers little over AdamW and adds computation for the per-layer norms. The layer-wise trust idea generalizes an earlier method, LARS, which applied the same normalization to plain momentum SGD; LAMB combines it with Adam's adaptivity. See learning-rate schedules for the warmup that large-batch training also requires.