Large Batch Training: Scaling Rules, Limits, and LAMB

Training a neural network means processing data in batches: you take a chunk of examples, compute gradients, and update weights. The batch size you choose is one of the most consequential decisions in the entire training setup. It shapes how quickly you learn, how stable the training is, how well the model generalizes, and how efficiently you use your hardware.

For most of deep learning's history, small to medium batch sizes (32 to 256) were the default. Researchers discovered through painful experience that simply cranking up the batch size made models train faster per epoch but converge to worse solutions. The relationship between batch size and model quality seemed to be a fundamental constraint. Many practitioners accepted this as an immutable law of optimization, choosing their largest tolerable batch size and moving on.

That changed in the mid-2010s, driven partly by practical necessity. As datasets grew to billions of tokens and models reached hundreds of billions of parameters, training on small batches became prohibitively slow. Distributing training across thousands of GPUs required large aggregate batch sizes just to keep the hardware busy. A single forward pass on 4,096 GPUs processes one mini-batch per step; if your batch size is 256, most of your hardware sits idle waiting for gradient synchronization. Researchers had to figure out how to train effectively at scales that were previously unthinkable.

The result was a collection of techniques, with the most important insight being the linear scaling rule: when you multiply the batch size by a factor $k$, multiply the learning rate by the same factor $k$. This rule, backed by empirical and theoretical analysis, made it possible to train models on massive batches without degrading quality. Combined with the LAMB optimizer and careful attention to batch size limits, large batch training became standard practice.

In this chapter, you will understand why batch size affects learning in the way it does, how the linear scaling rule works and why it is justified, where the fundamental limits of batch scaling lie, and how the LAMB optimizer extends these ideas to the pre-training regimes of BERT and GPT.

Before examining the mechanics of scaling, it helps to understand what gradient descent computes and why batch size changes its character.

The true objective you want to minimize is the loss averaged over your entire dataset:

where:

• $L(\theta)$: the full-dataset loss as a function of model parameters $\theta$
• $N$: total number of training examples
• $\ell(x_i, y_i; \theta)$: the loss on the $i$-th example $(x_i, y_i)$

The true gradient of this loss with respect to the parameters is:

Computing this exactly would require a pass over all $N$ examples for each parameter update. For datasets with millions or billions of samples, that is computationally intractable.

Stochastic gradient descent avoids this by computing a noisy estimate using a mini-batch $\mathcal{B}$ of size $B$:

where:

• $g_B$: the mini-batch gradient estimate
• $B$: the batch size
• $\mathcal{B}$: the set of $B$ randomly sampled example indices

This estimate is unbiased: its expected value equals the true gradient. But it has variance that depends inversely on $B$:

where $\sigma^2$ is the per-sample gradient variance. Doubling the batch size halves the gradient variance. With a larger batch, your estimate of the true gradient direction becomes more accurate. This is simply the central limit theorem applied to gradient estimation: the sample mean of $B$ independent random vectors has variance $1/B$ times the variance of a single sample.

The noise in the gradient estimate plays a dual role, and understanding both sides is essential for making good batch size decisions.

On the positive side, gradient noise acts as a form of implicit regularization. Noisy updates prevent the optimizer from settling into sharp minima, those loss landscape features with high curvature in many directions. Sharp minima tend to correspond to solutions that overfit to training data. The parameters that minimize the training loss in a sharp minimum are very specific: small perturbations to the weights increase the training loss significantly. When you deploy the model, the test distribution differs slightly from training, and this difference looks like a perturbation to the weights. Sharp minima therefore generalize poorly.

Flat minima, where the loss is low and relatively insensitive to small perturbations, generalize better precisely because they tolerate the implicit shift between train and test. Gradient noise helps find these flat regions because the stochastic updates effectively prevent the optimizer from getting trapped in narrow, sharp basins. This connection between batch size, noise, and generalization was studied in depth by Keskar et al. (2017), who found that models trained with large batches consistently converged to sharper minima than those trained with small batches under the same budget.

On the negative side, noise means you sometimes step in the wrong direction. If the gradient noise is large relative to the true signal, you spend many updates wandering randomly rather than making consistent progress toward the minimum. With batch size 1, each update might point nearly opposite to the true gradient if you happen to sample an outlier example. The optimizer makes progress on average (the estimate is unbiased) but wastes many steps on unhelpful updates.

This tradeoff gives batch size its characteristic effect: small batches take many individually noisy steps, making slower but ultimately more stable progress, while large batches take fewer but more accurate steps, converging faster per epoch but risking poor generalization.

The appropriate learning rate depends on the curvature of the loss landscape. In regions of high curvature (sharp narrow valleys), a large learning rate causes divergent oscillations as each update overshoots the local minimum and bounces to the other side. In regions of low curvature (flat plateaus), a small learning rate makes progress unnecessarily slow.

When you increase the batch size, you reduce gradient noise. This changes the effective signal-to-noise ratio of each update. If you keep the learning rate fixed while increasing the batch size, the optimizer takes fewer but more accurate steps per epoch. The effective step size relative to the gradient signal has changed, and the training dynamics shift. Specifically, with a fixed learning rate, each more-accurate large-batch step covers less distance in gradient-normalized space than the noisy small-batch step would have.

This is the core insight behind learning rate scaling rules: batch size and learning rate are not independent settings. They are coupled through the statistics of the gradient estimate. Increasing batch size without adjusting learning rate is like making more accurate measurements and then using fewer of them without increasing your confidence accordingly.

Understanding where the linear scaling rule came from helps contextualize why it works when it does and fails when it does not. For most of the 2010s, practitioners treated batch size as primarily a hardware constraint. You used the largest batch that fit in GPU memory, typically 32 to 256, and accepted the associated learning rate.

The breakthrough came when distributed training made truly large batches possible. A cluster of 16 GPUs, each processing a batch of 32, produces an effective batch of 512 per step. A cluster of 512 GPUs produces an effective batch of 16,384. At these scales, the question of how to set the learning rate became urgent, because the naive answer (use the same learning rate as the single-GPU setting) clearly did not work.

The Facebook AI Research team's 2017 paper "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" provided the systematic answer. They trained ResNet-50 on ImageNet using 256 GPUs in one hour, matching the accuracy of a week-long single-GPU run. The key contributions were the linear scaling rule, the warmup strategy, and careful attention to batch normalization behavior at large scale. This work catalyzed a wave of research and engineering that ultimately enabled the training of GPT-3, PaLM, and similar models at scales previously considered impossible.

The most influential contribution to large batch training was the systematic analysis of how learning rate should scale with batch size. The rule itself is simple to state, but its derivation illuminates exactly why it works and where it breaks down.

Consider SGD with batch size $B$ and learning rate $\eta$. The derivation below assumes the loss is locally linear over the region swept by $k$ consecutive steps, which is the key approximation. After one step, parameters change by:

Now consider scaling the batch size by a factor $k$ to batch size $kB$. The new gradient estimate:

has variance $\frac{\sigma^2}{kB}$, which is $k$ times smaller than the variance of $g_B$.

The key intuition is this: with batch size $kB$, you need $k$ fewer steps to process the same amount of data. If you want the model after processing $n$ samples to be in the same state regardless of batch size, you need the cumulative weight update to be the same. One large-batch step with batch size $kB$ should replicate what $k$ small-batch steps with batch size $B$ would have achieved.

More formally, $k$ steps of SGD with learning rate $\eta$ and batch size $B$ produce:

One step with batch size $kB$ and learning rate $\eta' = k\eta$ produces:

When the loss landscape is locally linear (a first-order approximation), these two updates are equal in expectation, since $\mathbb{E}[g_{kB}] = \mathbb{E}[g_B] = \nabla L(\theta)$:

This is the formal justification for the linear scaling rule.

The derivation is mathematically clean, but the intuition is worth dwelling on. The argument assumes that $k$ consecutive small-batch updates do not change the gradient significantly, so the gradient computed on the first mini-batch is approximately the same as the gradient on the second, third, and so on. This is the local linearity assumption: the loss function behaves like a flat plane over the small region swept by $k$ SGD steps.

When this assumption holds, the $k$ small-batch steps are equivalent to a single step in the average gradient direction, scaled by $k$. A single large-batch step computes this average gradient more accurately (lower variance) and takes a proportionally larger step (with $k\eta$), arriving at the same expected location.

This framing also shows what can go wrong. If the model parameters change substantially over $k$ steps (high curvature, large gradients, or large learning rate), then the gradient estimated at step 1 points in a different direction from the true gradient at step $k$. The equivalence breaks down, and simply scaling the learning rate linearly overestimates how large a step you can safely take.

The linear scaling rule is exact only in the limit where the loss is locally linear (equivalently, where the gradient does not change significantly within a single step). In practice, the loss landscape has curvature, and the approximation breaks down at two extremes.

At small batch sizes, the gradient noise is so high that the local linearity assumption is violated. The gradient estimated on one mini-batch points in a very different direction than the gradient estimated on the next, so the simple equivalence argument does not hold.

At very large batch sizes, a different problem emerges. As you keep scaling the learning rate up proportionally to batch size, eventually the learning rate becomes so large that a single step overshoots the local minimum. The update magnitude exceeds the scale over which the linear approximation is valid. This is why there is an upper limit to how far the linear scaling rule can be pushed.

The practical range where the rule works reliably varies by architecture and dataset. For ImageNet training with ResNets, it holds up to batch sizes of around 8,192. For large language models trained on massive corpora, the effective range depends on model size, architecture, and stage of training. The rule is a powerful starting point, not an infallible prescription.

Before the linear scaling rule was fully established, researchers proposed an alternative: the square root scaling rule, which suggests using learning rate $\eta \cdot \sqrt{k}$ when scaling batch size by $k$. This rule comes from a different analysis of noise dynamics.

The argument runs as follows. The variance of the gradient estimate scales as $\sigma^2 / B$. The standard deviation (noise amplitude) scales as $\sigma / \sqrt{B}$. If you want to keep the signal-to-noise ratio of your update constant as you increase batch size, you should scale the learning rate by $\sqrt{k}$, because the noise has decreased by $\sqrt{k}$.

In practice, the square root rule tends to be too conservative for moderate batch size increases. The linear rule matches empirical results better for scaling up to 10x or 100x the original batch size. For very large scaling factors (beyond 1,000x), neither rule alone is sufficient, and more sophisticated techniques are needed. The linear rule has become the dominant standard, but knowing the square root rule exists helps you understand why some older papers use different scaling conventions.

At the start of training, the network parameters are far from their final values. The loss landscape in this early phase is highly non-linear, with large curvature in many directions. The linear scaling rule's underlying assumption of local linearity is worst at the beginning.

This is why large batch training almost always uses a warmup phase. As we discussed in the Learning Rate Warmup chapter, you start with a small learning rate and gradually increase it to the target value over some number of steps. For large batch training, warmup is often essential for stable convergence.

The Facebook AI Research paper found that jumping directly to the linearly scaled learning rate caused instability in the first epoch. Using a warmup of 5 epochs resolved this. The intuition is that during warmup, the network moves from a random initialization to a region of parameter space where the local linearity assumption begins to hold, making the scaled learning rate safe. Put differently, warmup gives the optimizer time to find a basin before it starts taking large steps inside that basin.

The optimal warmup duration scales roughly with the batch scaling factor $k$. If you are scaling batch size by 8x, a warmup of 5 epochs may suffice. If you are scaling by 1,000x (as in large LLM training), you may need to warm up for thousands of steps. The general guideline is to warm up for at least 1-3% of total training steps, but this varies significantly across tasks.

Let us work through a concrete example to make the scaling rule tangible. Suppose you have determined through experimentation that batch size $B = 256$ trains well with learning rate $\eta = 0.1$ using SGD.

You now want to scale to a 64-GPU cluster where, to keep all GPUs busy, you need to use at least 16,384 examples per step.

The scaling factor is:

The linearly scaled learning rate is:

For warmup, a common practice is to linearly increase the learning rate from a small initial value to the target $\eta' = 6.4$ over the first few epochs, often 5 epochs for ImageNet-scale tasks.

At the end of training, the model has processed the same number of gradient updates relative to data seen. If the original run did 100 epochs of 1,000 steps with batch size 256, the scaled run does 100 epochs of roughly 16 steps with batch size 16,384. The same amount of data is processed overall, but in far fewer, more accurate steps.

In practice, this specific example is near the edge of where the rule holds. For batch sizes beyond roughly 8,192 on ImageNet, researchers have found that extra techniques (gradient noise injection, LARS/LAMB, more aggressive warmup) are needed to maintain quality. This example sits at 16,384, so you would likely need at least LAMB or careful warmup tuning to close the generalization gap.

The linear scaling rule describes how to adjust the learning rate when you scale up batch size. But there is a separate question: is there a limit to how large a batch you should use, regardless of learning rate scaling?

The answer is yes, and this limit is sometimes called the critical batch size.

The critical batch size is the point at which further increasing the batch size no longer improves the speed of convergence (measured in number of examples processed). Below this size, larger batches speed up convergence. Above it, you get no benefit from more parallelism.

Intuitively, this happens because beyond a certain batch size, the gradient estimate is already accurate enough. The bottleneck for learning is no longer gradient noise but the optimization geometry itself: the curvature of the loss landscape, saddle points, and movement across a complex high-dimensional surface. Adding more samples to a batch that is already providing a low-variance gradient estimate is simply redundant.

Think of it this way: if the true gradient points northeast, a batch that gives you an estimate pointing northeast-by-north is good enough. Doubling the batch size to get an estimate pointing northeast-by-northeast gives you a more accurate estimate but does not let you take a bigger step. You are still limited by the curvature of the loss surface, which determines the maximum safe step size.

Formally, the critical batch size is related to the ratio of gradient noise to gradient signal. McCandlish et al. (2018) formalized this as:

where:

• $B_{\text{crit}}$: the critical batch size
• $H$: the Hessian of the loss (capturing curvature)
• $\Sigma$: the per-sample gradient covariance matrix
• $g$: the true gradient vector
• $\text{tr}(\cdot)$: the matrix trace operator

This expression measures the ratio of gradient noise (numerator) to gradient signal (denominator) in a curvature-aware way. The critical batch size is larger when the gradient noise is high relative to the signal, and smaller when the signal is strong.

The numerator, $\text{tr}(H \Sigma)$, captures the total noise in the gradient weighted by curvature. Noise in directions of high curvature matters more because those directions are the ones where accurate gradient estimation most strongly affects convergence. The denominator, $g^T H g$, is the curvature-weighted gradient signal, measuring how much the gradient points in useful (high-curvature) directions. The ratio gives a curvature-sensitive signal-to-noise measure, and the critical batch size is the point at which increasing batch size no longer improves this ratio relative to the cost of the extra computation.

In practice, you cannot easily compute this quantity directly because you cannot compute the full Hessian for large models. But you can estimate the critical batch size empirically: train with progressively larger batches (adjusting the learning rate proportionally) and measure how many examples you need to process to reach a target validation loss. The batch size at which this number stops decreasing is approximately the critical batch size.

Research has found that the critical batch size varies widely across tasks and architectures:

• For image classification on CIFAR-10 with small models, critical batch sizes can be as low as a few hundred.
• For ImageNet with ResNet-50, the critical batch size is roughly 1,000 to 2,000 (much lower than the 8,192 that works with careful tuning, suggesting the extra scaling is not free from a convergence-per-example perspective).
• For language model pre-training (GPT-style), the critical batch size grows as training progresses. Early in training, when gradients are large and noisy, large batches help. Later in training, when the model is near convergence, smaller batches would suffice.
• For transformer-based language models, Kaplan et al. (2020) found that the critical batch size scales with the number of tokens processed, not just model size, suggesting a dynamic rather than fixed optimal batch size.

The implication for practitioners is important: there is an efficiency frontier. At batch sizes below the critical value, you are leaving GPU parallelism on the table. At batch sizes above it, you are wasting compute on redundant gradient averaging. Optimal large-scale training aims to hit this frontier.

However, hardware constraints often force you above the critical batch size anyway. When you have 10,000 GPUs, you need an enormous per-step batch just to utilize all the hardware. In these cases, the extra batch scaling beyond the critical size is a tax paid for hardware utilization. Engineers accept some inefficiency (in examples processed per useful gradient step) in exchange for faster wall-clock time.

Even at or below the critical batch size, large batch training can produce models that generalize slightly worse than small batch training with the same total compute budget. This is the "generalization gap" of large batch training, documented by Keskar et al. (2017).

The mechanism is the sharpness of minima. Large batch training, with its accurate gradient estimates, tends to converge to sharp minima where the loss is low but curvature is high. Small batch training, with its inherent noise, naturally explores the loss landscape more and finds flatter minima. Flat minima generalize better because small perturbations to the weights (representing test distribution shift) cause only small increases in the loss.

Several approaches address this gap:

• Ghost batch normalization: Run batch normalization with a smaller "ghost" batch size even when processing a larger batch, preserving the regularization effect of BN statistics computed on smaller groups. This decouples batch normalization from the training batch size.
• Linear warmup: Helps the optimizer find a good basin before scaling to the full learning rate, giving it time to identify flat regions.
• Sharpness-aware minimization (SAM): Explicitly penalizes sharp minima during training, regardless of batch size, by seeking parameters whose neighborhood has uniformly low loss.
• Data augmentation: Increases effective noise through input perturbations, partially compensating for the reduced gradient noise of large batches. More aggressive augmentation with large batches can close much of the generalization gap.
• Gradient noise injection: Deliberately adding noise to gradients to simulate the effect of smaller batches, while keeping the large batch for hardware efficiency.

In practice, for language model pre-training at the scales used in production (batch sizes of millions of tokens), the generalization gap is usually small and manageable with standard techniques. The model capacity and data volume dominate, and the batch size effect becomes a secondary concern.

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:

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:

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:

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.

We will implement and compare several training scenarios: standard SGD at small batch size, SGD at large batch size without scaling, SGD at large batch size with linear scaling, and LAMB. For tractability, we use a simple convolutional model on CIFAR-10.

We use CIFAR-10 with standard preprocessing. To make the experiments run quickly, we use a subset of the training data.

We implement a simplified LAMB to show the core mechanics. Production use would rely on the torch-optimizer or apex packages.

The compact experiment illustrates both the motivation for scaling and its limits. The small-batch baseline reaches the highest final accuracy. The unscaled large-batch run trails it, but applying the full 8x learning-rate increase is too aggressive here and performs worse. LAMB is the strongest of the large-batch configurations, although it still does not match the small-batch baseline. This is why the linear rule is a starting point rather than a guarantee: the scaled rate must remain within the stable range of the particular model, data, and training budget.

To make the statistical intuition concrete, we can directly compute how gradient variance changes with batch size on a fixed parameter snapshot.

The critical batch size can be approximated empirically by measuring how many training examples you need to process to reach a target loss level, at different batch sizes with linearly scaled learning rates.

This compact run does not show an initial efficiency gain from larger batches. Instead, examples required rise from batch size 32 through 128, and batches of 256 or more miss the target within the training budget. That outcome places the critical regime below 256 for this particular model, subset, learning-rate rule, and 15-epoch budget. In larger experiments the transition is often flatter, but the diagnostic principle is the same: once larger batches stop reducing the examples needed to reach a target, extra parallelism no longer improves statistical efficiency.

A key property of LAMB is that the trust ratio varies significantly across layers. Early layers (especially embeddings and the first convolutional layers) tend to have different weight-to-gradient-norm ratios compared to later layers. We can visualize this directly on our trained model.

Understanding each hyperparameter and its role helps you make principled decisions when configuring large batch training.

• batch_size: The number of training examples processed per gradient update. Should be chosen to balance hardware utilization and the critical batch size limit. The critical batch size varies by task; empirical measurement (as shown above) gives the best estimate.

• learning_rate ($\eta$): The global learning rate. Must be scaled linearly with batch size: $\eta' = k\eta$ when batch size is multiplied by $k$. The base learning rate $\eta$ should be validated at a small batch size before scaling.

• warmup_epochs: Number of epochs to linearly increase the learning rate from a small initial value to the target. Typically 5 epochs for ImageNet-scale tasks, often 1-3% of total training steps for LLM pre-training. A good rule of thumb is to warm up for at least as many steps as $k$ (the scaling factor).

• weight_decay ($\lambda$): L2 regularization coefficient. Incorporated directly into the LAMB update as a decoupled weight decay term. Decoupled weight decay (AdamW-style, not the original Adam L2 approach) is important for LAMB to behave correctly.

• beta1, beta2: Exponential decay rates for Adam's first and second moment estimates (LAMB). Standard values are 0.9 and 0.999. These generally do not need adjustment when scaling batch size.

• trust_ratio_clip: (Optional) Maximum value of the LARS/LAMB trust ratio to prevent instability for layers near zero weight norm. A common default is to clip trust ratios to the range $[0, 10]$.

Moving from theory to a working large-batch training setup involves several steps that are easy to get wrong. Here is a practical checklist for configuring large batch training in a real project.

First, establish your baseline at a small batch size. Train your model at batch size 32 or 64 with the default learning rate for your optimizer. Verify that training converges stably and achieves the expected validation performance. This baseline is the reference you will compare everything against.

Second, choose your target batch size. This should be driven by your hardware: how many GPUs you have and how large a batch each can process. For multi-GPU training, the total batch is the per-GPU batch times the number of GPUs. Aim to hit the critical batch size approximately; going much higher wastes compute efficiency without improving convergence.

Third, apply the linear scaling rule. Multiply your baseline learning rate by $k = \text{target batch size} / \text{baseline batch size}$. Add a warmup schedule starting from $\eta' / k$ or a fixed small value, ramping to $\eta'$ over the first 5% of training steps.

Fourth, choose your optimizer. Use Adam for small to medium batch sizes and standard transformer architectures. Use LAMB when you are training transformers with batch sizes above 8,192 tokens or when you observe training instability at the linearly scaled learning rate.

Fifth, monitor training stability carefully. Large batch training with a scaled learning rate can diverge quickly if something is misconfigured. Monitor gradient norms, weight norms, and per-layer update norms in the first few hundred steps. If any layer shows explosive gradient growth or the loss spikes, reduce the learning rate or extend the warmup.

Finally, validate generalization. Run your large-batch model on the held-out test set and compare to the small-batch baseline. If there is a meaningful accuracy gap, consider adding data augmentation, ghost batch normalization, or a longer training schedule to compensate.

The linear scaling rule is an approximation, and its applicability depends on assumptions that are often only partially satisfied. Understanding where it breaks down guides practical decisions.

The most significant limitation is the relationship between batch size and generalization. While careful engineering can close most of the accuracy gap between large and small batch training, the gap rarely disappears entirely. Training with very large batches tends to converge to sharper minima. Sharp minima are more sensitive to distribution shift between training and test sets, leading to worse generalization at deployment. In practice, large batch pre-trained models often require more aggressive fine-tuning data augmentation to match the downstream performance of models trained with smaller batches.

A second practical constraint is compute efficiency. The critical batch size is typically much smaller than what modern distributed training hardware can support at full utilization. For example, a 1,024-GPU cluster might require a batch size of 32,768 per step just to keep all GPUs busy, while the critical batch size for the task might be 2,048. The extra 16x scaling is necessary for hardware efficiency but does not improve convergence speed per example. In these regimes, engineers often accept some efficiency loss in exchange for better model quality, choosing a batch size that balances hardware utilization with convergence efficiency.

LAMB's layer-wise trust ratio addresses the adaptive optimizer scaling problem, but it introduces its own sensitivity. The trust ratio computation depends on the L2 norms of weights and Adam update directions. Early in training, when weights are small and gradients are large, the trust ratio can produce very small effective learning rates for some layers. This is usually mitigated by trust ratio clipping and careful initialization, but it requires tuning. LAMB's original validation was primarily on BERT and ResNets. For other architectures with unusual weight scale distributions across layers (especially models with heavily normalized activations), the trust ratio may behave differently and require additional tuning.

Batch normalization also interacts with batch size in complex ways. Batch normalization statistics (mean and variance computed over the batch) change significantly as batch size changes. With very large batches, the statistics become very accurate, but they also incorporate more diversity in the batch, potentially changing the effective normalization behavior. Ghost batch normalization sidesteps this by computing BN statistics on smaller sub-batches, but it introduces the hyperparameter of ghost batch size. For language models using layer normalization rather than batch normalization, this issue is less acute.

The field has also moved toward gradient accumulation as an alternative to true large batch training. Rather than physically processing a large batch on one step, you accumulate gradients from multiple smaller batches before applying the update. This achieves the same effective batch size without requiring hardware that can fit the large batch in memory. Gradient accumulation is straightforward for plain SGD and Adam, but interacts subtly with batch normalization statistics and some regularization techniques. The next chapter on gradient accumulation explores this approach in depth.

Despite these limitations, large batch training has become indispensable for modern AI development. Training GPT-3 with batch sizes of millions of tokens, or BERT with batch sizes of hundreds of thousands, was only possible because of the combination of linear scaling, LAMB, warmup, and careful monitoring of training stability. The techniques in this chapter represent the engineering foundation that made the scaling era of language models possible.

Large batch training changes the statistical properties of gradient estimation. Larger batches provide more accurate gradient estimates (lower variance), enabling the use of larger learning rates that make each step more consequential.

The linear scaling rule states that when you increase the batch size by a factor $k$, you should increase the learning rate by the same factor $k$. This rule emerges from the requirement that the expected weight update after processing $n$ training examples should be the same regardless of batch size. The rule holds well in regimes where the loss landscape is locally smooth but breaks at extremes: very small batches have too much noise, and very large learning rates cause overshooting.

The critical batch size marks the practical upper bound on useful batch scaling. Beyond this size, additional parallelism does not reduce the total number of examples needed for convergence. It represents the point at which gradient noise is no longer the bottleneck for learning.

LAMB extends the linear scaling idea to adaptive optimizers. It combines Adam's per-parameter moment adaptation with a LARS-style layer-wise trust ratio, preventing any layer from taking an unstable step when the global learning rate is large. LAMB enabled BERT pre-training at batch sizes that reduced wall-clock training time from days to minutes on large TPU pods.

Warmup remains essential across all large batch techniques, bridging the early phase where the loss landscape is non-linear and the later phase where the linear scaling assumption holds.

The next chapter on weight decay examines complementary regularization techniques that interact with batch size decisions, completing the toolkit for large-scale training optimization.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about large batch training, learning rate scaling, and the LAMB optimizer.