Optimization at Scale: Large Batches, Noise, and Variance

Why bigger batches stop helping

A mini-batch gradient is a noisy estimate of the true gradient, and that noise helps — it lets SGD escape sharp regions and explore. When you scale distributed training across many devices, the natural move is to grow the batch so each device stays busy. But there is a ceiling: past a certain batch size, doubling it stops halving the number of steps to convergence, because the gradient estimate is already accurate and the extra examples are mostly redundant. You are paying twice the compute per step for a vanishing reduction in steps.

This transition is governed by the critical batch size: below it, more data per step buys proportionally faster convergence; above it, returns diminish sharply. Knowing roughly where it sits for your problem is the difference between a cluster that is well-utilized and one that is burning power on redundant examples.

Gradient noise scale: measuring the critical batch

Gradient noise scale gives you an actual number. The idea: estimate the variance of the per-example gradients and compare it to the squared norm of the mean gradient. When that signal-to-noise ratio is poor (gradients disagree a lot across examples), large batches help because averaging cancels the disagreement; when it is good, you are near the critical batch and further growth is wasteful. The noise scale is exactly that ratio, and it predicts the critical batch size remarkably well across tasks.

B_{\text{noise}}=\frac{\operatorname{tr}(\Sigma)}{\lvert G\rvert^{2}}

Gradient noise scale: per-example gradient variance over the squared mean-gradient norm sets the useful batch size.

# estimate by comparing gradient norms at two batch sizes b_small, b_big
# B_noise ~ tr(Sigma) / |g|^2  (variance of grads over squared mean-grad norm)
g_small = grad(loss, batch=b_small)
g_big   = grad(loss, batch=b_big)
# fit |g_est|^2 = |g_true|^2 + tr(Sigma)/B  across the two B values, solve for B_noise

The noise scale rises during training, so the useful batch size grows as the model learns.

LAMB: large-batch training that actually converges

Pushing batch size into the tens of thousands breaks naive optimizers: a single global learning rate that suits one layer overshoots another, because layers differ wildly in the ratio of their weight norm to their update norm. LAMB (Layerwise Adaptive Moments for Batch training) fixes this with per-layer trust ratios: it computes an Adam-style update, then rescales each layer's update so that its magnitude is proportional to that layer's weight norm. No layer is allowed to take a step that is huge relative to its own weights.

x_{t+1}^{(i)} = x_t^{(i)} - \eta_t\,\frac{\lVert x_t^{(i)}\rVert}{\lVert r_t^{(i)}\rVert}\, r_t^{(i)},\qquad r_t^{(i)}=\frac{m_t^{(i)}}{\sqrt{v_t^{(i)}}+\epsilon}+\lambda\,x_t^{(i)}

LAMB's layerwise trust ratio ‖x‖/‖r‖ rescales each layer's update so one global learning rate fits every layer.

This layerwise normalization is what let LAMB train BERT in large-batch regimes — tens of thousands of examples per step — without losing final accuracy, cutting wall-clock pretraining time dramatically when hardware is plentiful. The general lesson outlives the specific optimizer: at scale, the right unit of adaptation is often the layer, not the individual parameter, because layers are the natural blocks whose geometry differs.

Variance-reduced SGD: beautiful theory, awkward fit

If gradient noise is the obstacle, why not cancel it directly? Variance-reduced SGD — the SVRG and SAGA families — does exactly that. SVRG occasionally computes a full-dataset gradient as an anchor, then corrects each cheap stochastic gradient by the difference between the current and anchored gradients at a control point. On finite-sum convex problems this provably reaches linear convergence, dramatically better than plain SGD's sublinear rate.

g_t = \nabla f_{i}(x_t) - \nabla f_{i}(\tilde{x}) + \nabla F(\tilde{x}),\qquad \nabla F(\tilde{x})=\frac{1}{n}\sum_{j=1}^{n}\nabla f_{j}(\tilde{x})

SVRG's control variate subtracts the snapshot gradient and adds the full-batch anchor, cancelling noise without adding bias.

So why is it rare in deep learning? Three reasons. The anchor gradient over a giant dataset is expensive; the loss surface is nonconvex, so the guarantees weaken; and crucially, that gradient noise was helping generalization, so removing it can hurt the final model even as it speeds optimization. Variance reduction is a sharp reminder that faster optimization is not the same as better learning — a theme worth carrying into every optimizer choice.

Memory pressure and the lightweight turn

At scale the optimizer's state becomes a budget item. Adam stores two extra tensors per parameter — first and second moments — so the optimizer alone can double or triple the memory footprint of the weights. When a model barely fits, that overhead decides whether you can train at all. This pressure motivated Lion (EvoLved Sign Momentum), which keeps only a single momentum buffer and updates by the sign of an interpolated momentum, halving optimizer memory relative to Adam.

Because every Lion step is a sign — magnitude one in every coordinate — its effective step geometry is closer to a normalized update than to a curvature-scaled one, and it usually wants a smaller learning rate and a larger weight decay than Adam. Lion is the bridge to Guide 5, where we ask the bigger question: if a sign-based rule found by automated search can rival hand-designed Adam, how are good optimizers discovered at all?