PaperlayerView exact revision on arXiv

Source-linked AI summary

Why Does Adaptive Batching Help LLM Pretraining? A Perspective from Unbounded Variance

Arda Fazla, Antesh Upadhyay, Ege C. Kaya, M. Berk Sahin, Abolfazl Hashemi

arXiv:2610.02355v1cs.LGmath.OC

TL;DR

The paper addresses the unclear theoretical basis for increasing batch size during LLM pretraining when gradient variance can grow with distance from initialization. It introduces BG-a theory and an adaptive batch scheduler; in OLMo2 pretraining, the scheduler lowers validation loss versus fixed-batch baselines under matched token budgets while using less than 10% of small-batch iterations.

  • Problem

    The theoretical justification for increasing batch size in LLM pretraining is unclear, while bounded-variance analyses do not capture distance-dependent stochastic-gradient noise.

    Source: §1 · p. 1

  • Method

    The paper introduces BG-a noise, derives growth-dependent complexity bounds, and uses estimated noise parameters to adapt batch size as iterates move from initialization.

    Source: §1 · pp. 2–3

  • Results

    The scheduler achieves lower validation loss than small- and large-batch training under matched token budgets, using less than 10% of small-batch training’s iterations.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The analysis links batch-size growth to distance-dependent noise and explains why batches can start small and increase as iterates move away from initialization.

    Source: §1 · pp. 2–3

    The upper bounds require a ≤ 2 because controlling the a-th displacement moment uses the second moment supplied by the variance assumption.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Increasing the batch size during training is a common practice in large language model (LLM) pretraining, yet the theoretical justification behind its success is not well understood. Analyses of stochastic optimization often assume uniformly bounded stochastic gradient variance, yet recent evidence suggests that this assumption fails in many practical nonconvex problems. The Blum--Gladyshev (BG-$0$) noise model relaxes this assumption by allowing the variance to grow quadratically with the distance from initialization, suggesting that batch size schedulers can help by controlling the variance growth during training. However, this growth can be overly conservative in practice. We empirically investigate variance growth in LLM pretraining and observe that a generalized BG model with a tunable growth exponent provides a tighter description of practical noise behavior. Motivated by this observation, we introduce the generalized BG-$a$ noise model, which interpolates between bounded variance ($a=0$) and BG-$0$ noise ($a=2$). Under $L$-smoothness, we derive an information-theoretic lower bound with growth-dependent oracle complexity $\Omega(\epsilon^{-(4+a)})$ and establish a matching upper bound in $\epsilon$-dependence by increasing the batch size as the iterates move away from initialization. Finally, we propose an adaptive batch scheduler that controls variance growth through dynamic batch size adjustments during training. In pretraining OLMo2 models of up to 1B parameters on C4, our scheduler achieves a lower validation loss than both small and large batch training under matched token budgets, while using less than 10\% of the iterations of small batch training.

1 INTRODUCTION

The paper asks why increasing batch size helps LLM pretraining, where distance-dependent gradient noise challenges standard bounded-variance analyses. It introduces BG-a theory and an adaptive scheduler, reporting lower validation loss than fixed-batch baselines under matched token budgets.

1 INTRODUCTION

  • Classical analyses commonly assume uniformly bounded stochastic-gradient variance, while distance-dependent variance can violate that assumption and incur greater oracle complexity.The BG-0 model allows variance to grow quadratically with distance from initialization.Source: §1 · p. 2
  • Empirical LLM noise growth is subquadratic, motivating BG-a, which interpolates between bounded variance at a = 0 and quadratic BG-0 noise at a = 2.Fitted exponents range from 0.98 to 1.49 and increase with model size.Source: §1 · p. 2, Figure 1 · p. 2
  • Under BG-a noise, lower and upper oracle-complexity bounds scale as ϵ^-(4+a), recovering the bounded-variance and BG-0 endpoint rates.The lower bound is information-theoretic; dynamic batching attains the matching upper-bound dependence.Source: §1 · p. 2
  • Dynamic batching tracks distance-dependent noise by increasing batch size as iterates move farther from initialization, while momentum can reduce per-iteration batch requirements at the cost of more iterations.The momentum analysis preserves the O(ϵ^-(4+a)) oracle complexity.Source: §1 · pp. 2–3
  • The adaptive scheduler achieves lower validation loss than both small- and large-batch training under matched token budgets, using less than 10% of small-batch training’s iterations.The reported pretraining experiments use OLMo2 models on C4.Source: §1 · p. 3

2 GRADIENT NOISE GROWS WITH DISTANCE IN LLM PRETRAINING

The section measures how stochastic-gradient variance changes with displacement from initialization in OLMo2 models. A tunable BG-a exponent fits subquadratic growth, with fitted parameters varying across model sizes.

2 GRADIENT NOISE GROWS WITH DISTANCE IN LLM PRETRAINING

  • The experiment evaluates randomly perturbed parameter states rather than following an optimizer-generated training trajectory.It samples 50 logarithmically spaced radii and 10 independent directions at each radius.Source: §2 · p. 3
  • Perturbation directions double the Gaussian scale for LM-head parameters before global normalization, reflecting their largest observed weight changes in pretraining runs.The resulting direction is normalized so displacement remains r.Source: §2 · p. 3
  • At each evaluation point, stochastic-gradient variance is estimated relative to a fixed large-batch reference gradient.The OLMo2 estimates use a 512-sequence reference batch and 32 minibatches of four sequences each.Source: §2 · p. 4
  • Fitted exponents are 0.98, 1.38, and 1.49 for OLMo2-100M, OLMo2-600M, and OLMo2-1B, respectively, and BG-a fits the observed variance increase better than fixed-exponent BG-0.The fitted parameters also show B decreasing while G and a increase with model size.Source: §2 · p. 4, Figure 2 · p. 4

3 THEORETICAL ANALYSIS UNDER THE BG-a ORACLE

The analysis studies smooth nonconvex optimization with an unbiased BG-a stochastic oracle and bounds the cost of finding an approximate stationary point. Its dynamic-batching results connect batch growth to displacement from initialization, with an upper-bound restriction of a ≤ 2.

3.1 PRELIMINARIES

  • The setup measures stationarity through expected squared gradient norm and counts stochastic first-order oracle calls, with minibatches sampled independently after batch size is selected.The analysis assumes a globally L-smooth objective bounded below.Source: §3.1 · p. 5, §3.2 · p. 5

3.2 ASSUMPTIONS

  • The stochastic oracle is assumed unbiased, with variance bounded by B²∥x − x0∥^a + G².The objective is globally L-smooth.Source: §3.2 · p. 5

3.3 CONVERGENCE RESULTS

  • For 0 ≤ a ≤ 2, dynamic-batching SGD and SGDM achieve O(ϵ^-(4+a)) expected oracle complexity, matching the growth-dependent lower-bound exponent.For B > 0, the matching exponents are ϵ^-4 at a = 0, ϵ^-(4+a) for 0 < a < 2, and ϵ^-6 at a = 2.Source: §3.3 · pp. 5–6

3.4 IMPLICATIONS FOR BATCH SIZE SCHEDULING

  • The upper bounds require a ≤ 2 because the variance assumption supplies a second moment for controlling the a-th displacement moment.When a = 0, the prescribed batch remains constant; for a > 0, it grows with displacement from initialization.Source: §3.3 · p. 6, §3.4 · p. 6
  • Dynamic batching starts with a smaller batch near initialization and increases it with displacement, avoiding the extra iterations of fixed small batches and early oversampling of fixed large batches.Momentum reduces the prescribed batch size by a factor of cβ at the cost of a 1/cβ larger iteration count, while preserving oracle complexity.Source: §3.4 · p. 6

4 ADAPTIVE BATCH SIZE SCHEDULER

The scheduler sets batch size from the iterate’s displacement from initialization using estimated BG-a noise parameters. It rounds batch multipliers upward to powers of two and adjusts learning rates according to the optimizer.

4 ADAPTIVE BATCH SIZE SCHEDULER

  • At each iteration, the scheduler uses displacement from initialization and estimated BG-a parameters to choose batch size before drawing stochastic gradients.Its batch rule is designed to control the modeled variance growth.Source: §4 · p. 6
  • The scheduler rounds the displacement-dependent batch multiplier upward to a power of two, ensuring the selected batch meets the variance-control requirement.Power-of-two multiples also reduce potentially costly batch-size adjustments.Source: §4 · p. 7
  • The batch multiplier depends on the noise parameters through B²/G² and a, so common rescaling of B² and G² leaves it unchanged.The multiplier can be computed directly from the current iterate, initialization, and noise parameters.Source: §4 · p. 7
  • The scheduler adjusts learning rates with optimizer-dependent scaling: linear scaling for SGD-based optimizers and square-root scaling for Adam.The optimizer otherwise follows its usual state update.Source: §4 · pp. 7–8
  • In experiments, the initial batch matches the small-batch baseline because the target noise level is not known a priori.The chosen baseline batch is the smallest one found to train stably and achieve competitive final loss after learning-rate tuning.Source: §4 · p. 8

5 EXPERIMENTS

On C4, the BG-a scheduler achieves the lowest validation loss across all three OLMo2 model sizes under matched token budgets. It uses substantially fewer iterations than small-batch training, while the fixed-batch comparison varies with the fitted noise-growth exponent.

5 EXPERIMENTS

  • The experiments pretrained OLMo2-100M, 600M, and 1B on C4 with approximately 20 tokens per parameter and compared BG-a against fixed small and large batches.The corresponding budgets were 2B, 12B, and 20B tokens; training used AdamW with linear warmup followed by cosine decay.Source: §5 · p. 8
  • 8.76%, 8.10%, and 5.25% of small-batch iterations were used by BG-a for OLMo2-100M, 600M, and 1B, respectively, while BG-a achieved the lowest validation loss across all three sizes.The comparison uses matched token budgets; Figure 3 plots validation loss and the BG-a batch-size schedule.Source: §5 · p. 8, Figure 3 · p. 9
  • For OLMo2-100M (a = 0.98), small-batch training outperformed large-batch training, whereas small batches had the worst validation loss for OLMo2-600M and 1B, with larger fitted exponents.Small-batch learning rates were directly tuned, while large- and adaptive-batch rates used square-root scaling without additional tuning.Source: §5 · p. 8

6 CONCLUSION

The paper concludes that distance-dependent gradient-noise growth helps explain increasing batch sizes during training, and reports an adaptive scheduler that improves on fixed-batch training. Its related-work discussion positions BG-a as an extension of prior variance models and batch-scheduling analyses.

6 CONCLUSION

  • The conclusion reports that the adaptive scheduler improves on both small- and large-batch training in LLM pretraining, and that momentum trades smaller sufficient batches for additional iterations.The stated scope is the paper's LLM pretraining experiments and theoretical analysis.Source: §6 · p. 9

A RELATED WORK

  • Prior work studies batch-size efficiency, learning-rate relations, and gradient-noise-scale proxies for critical batch size, while this paper offers a complementary noise-growth perspective.The related-work discussion also distinguishes optimization in online LLM pretraining from offline, multi-epoch image classification.Source: A RELATED WORK · p. 14
  • The BG-a analysis extends prior BG-0 work by considering variance growth proportional to distance from initialization raised to a ∈ [0, 2], and relates the exponent to oracle complexity.The discussion cites Fazla et al. (2026) and Upadhyay et al. (2026) for BG-0 analyses.Source: A RELATED WORK · p. 15

B LOWER BOUND UNDER L-SMOOTHNESS

  • The lower-bound setup assumes an unbiased oracle whose variance is bounded by B^2∥x − x0∥^a + G^2, with a = 0 representing bounded variance and a = 2 representing BG-0.The objectives are lower bounded and globally L-smooth, with a bounded initial optimality gap.Source: B LOWER BOUND UNDER L-SMOOTHNESS · p. 15

B.1 DEFINITIONS AND AUXILIARY RESULTS

The auxiliary results formalize a chain construction and masked oracle used in the lower-bound proof. They bound the constructed function's range and show that oracle responses reveal progress along the chain only incrementally.

B.1 DEFINITIONS AND AUXILIARY RESULTS

  • The zero-respecting restriction requires queried coordinates to be supported by previous stochastic-gradient responses, and the lower-bound construction uses a deterministic chain whose gradient extends progress by at most one coordinate.The proof also imposes a further output restriction for its first stage.Source: §B.1 · pp. 15–16
  • The range lemma establishes that the constructed chain has finite supremum and infimum, with M_T ≤ 12T and inf_w f̄_T(w) ≥ −12T.The proof derives the bounds from the auxiliary scalar-function properties.Source: §B.1 · p. 16
  • The Bernoulli masked chain oracle is unbiased and has a probability-p zero-chain property: responses can extend progress by at most one coordinate, and extension requires ξ = 1.The stated support bounds restrict nonzero coordinates to the current progress index and at most the next coordinate.Source: §B.1 · pp. 16–17

B.2 LOWER-BOUND THEOREM AND PROOF

This section establishes ε-dependent oracle lower bounds for the generalized BG-a model, including an extension to unrestricted randomized algorithms, using an activated-chain construction in the proof. It also interprets how the bound scales with variance growth and relates to dynamically batched upper bounds.

B.2 LOWER-BOUND THEOREM AND PROOF

  • The activated-chain proof combines a travel coordinate, which places small-gradient points far from initialization, with an activation that delays chain noise until variance is sufficient.The chain length remains of order LΔ/ε^2, while the generalized variance budget changes the coordinate-revelation probability.Source: §B.2 · p. 17
  • The worst-case expected SFO complexity is Ω(ε^-(4+a)) under the generalized BG-a model.The theorem assumes fixed L, Δ > 0, B, G ≥ 0, and finite a ≥ 0.Source: §B.2 · pp. 17, 21

B.3 REDUCTION TO UNRESTRICTED RANDOMIZED ALGORITHMS

  • Random rotation and compression extend the lower-bound construction to unrestricted randomized algorithms, including unbounded queries, arbitrary outputs, and random stopping.A quadratic penalty prevents large-norm queries from exploiting compression to produce small gradients.Source: §B.2 · p. 21, §B.3 · p. 22
  • The proof concludes by establishing the expected SFO lower bound, completing the lemma.Source: §B.3 · p. 23

B.4 INTERPRETATION AND RELATION TO THE UPPER BOUND

  • For 0 ≤ a ≤ 2, the lower bound scales as Ω(ϵ^-(4+a)): Ω(ϵ^-4) at a = 0, Ω_a(ϵ^-(4+a)) for 0 < a < 2, and Ω(ϵ^-6) at a = 2.The corresponding variance bounds are B² + G² at a = 0, B²∥x − x0∥^a + G² for 0 < a < 2, and B²∥x − x0∥² + G² at a = 2.Source: §B.4 · pp. 23–24
  • For 0 ≤ a ≤ 2 and B > 0, the lower bound matches the ϵ-dependence of dynamically batched SGD and SGDM upper bounds under fixed problem parameters, target variance σ², and momentum coefficient β.The additional ϵ^-a factor comes from allowable variance at distance Θ(∆/ϵ) from initialization; when B = 0 and G > 0, the bound reduces to Ω(ϵ^-4), recovering Arjevani et al. (2023).Source: §B.4 · p. 24

C CONVERGENCE OF SGD

Under distance-dependent BG-a noise, dynamic batching yields an expected SFO complexity whose exponent grows with the variance-growth exponent a. The bounds recover bounded-variance SGD at a=0 and BG-0 at a=2.

C CONVERGENCE OF SGD

  • The convergence result assumes 0≤a≤2, a fixed target variance σ²>0, and target accuracy satisfying ϵ²≤min{2σ²,L∆}.The theorem uses SGD updates from fresh conditionally independent minibatch samples.Source: C CONVERGENCE OF SGD · p. 24
  • Expected SFO cost is O(ϵ−4) at a=0, O(ϵ−(4+a)) for 0<a<2, and O(ϵ−6) at a=2.For fixed σ², the bound is O(ϵ−(4+a)); the exponent interpolates between the two endpoints.Source: C CONVERGENCE OF SGD · pp. 24, 26
  • At a=0, the batch size is constant and the bound recovers O(ϵ−4); setting B=0 also gives a uniform variance bound G².For a=2, the variance bound is B²∥x−x0∥²+G², the BG-0 model considered in Fazla et al. (2026).Source: C CONVERGENCE OF SGD · p. 26
  • For 0<a<2, the exponent a sets the allowed variance growth with distance from initialization, while the iteration budget remains O(ϵ−4) for fixed σ².The proof's expected batch cost includes a displacement factor (∆/ϵ)^a.Source: C CONVERGENCE OF SGD · p. 26

C.1 DISCUSSION ON THE EXPECTED SFO BOUND WHEN a > 2

For a>2, the second-moment variance control does not suffice to bound the displacement moment needed for expected SFO cost. Under Assumptions 1 and 2 alone, the prescribed batch rule can have infinite expected cost.

C.1 DISCUSSION ON THE EXPECTED SFO BOUND WHEN a > 2

  • For a>2, the second-moment bound no longer controls the a-th displacement moment required for finite expected batch cost when B>0.The same control suffices for 0≤a≤2 by Jensen’s inequality.Source: §C.1 · pp. 26–27
  • For every fixed a>2 and target variance σ²>0, an oracle satisfying Assumptions 1 and 2 can give E[b1]=∞ and infinite expected total SFO cost for every T≥2.The construction can have a finite second moment but an infinite a-th moment after the first update.Source: §C.1 · p. 27
  • Recovering the same O(ϵ−(4+a)) expected SFO bound for a≥2 requires new assumptions controlling gradient-noise moments of order a, a direction left for future work.Finiteness alone would not establish the required dependence on ϵ.Source: §C.1 · pp. 27–28

D CONVERGENCE OF SGDM

Under BG-a noise, SGDM with dynamic batching achieves expected SFO complexity O(ϵ^-(4+a)) for fixed momentum and target variance, while momentum trades smaller batches for more iterations. The supplied experiments report downstream accuracy and show subquadratic noise growth in vision models.

D CONVERGENCE OF SGDM

  • O(ϵ^-(4+a)) expected SFO complexity is achieved by SGDM with dynamic batching for fixed β and σ² under the theorem’s conditions.The reported rates range from O(ϵ^-4) at a=0 to O(ϵ^-6) at a=2.Source: D CONVERGENCE OF SGDM · pp. 28, 31
  • Larger β prescribes smaller batches but requires a smaller learning rate and more iterations, without reducing total SFO cost as momentum increases.The batch comparison is at a common query point; different momentum choices generally produce different trajectories.Source: D CONVERGENCE OF SGDM · pp. 31–32

E.1 LLM EXPERIMENTS

  • BG-a attains the highest mean accuracy on HellaSwag and CommonsenseQA across the three OLMo2 model sizes, with clearest gains over small batches at larger sizes.The models were trained on C4 with 20 tokens per parameter; Table 2 reports accuracy as mean±std across three seeds.Source: §E.1 · p. 32, Table 2 · p. 32

E.2 VISION EXPERIMENTS

  • The vision experiments use ResNet50 on CIFAR-10, CIFAR-100, and ImageNet-1K to examine adaptive batching in multi-epoch training on fixed datasets.The comparison considers both training and test performance.Source: §E.2 · p. 32

E.2.1 GRADIENT NOISE GROWTH IN IMAGE CLASSIFICATION

  • Fitted exponents of 1.25, 1.19, and 0.96 show subquadratic variance growth on CIFAR-10, CIFAR-100, and ImageNet-1K, respectively, captured by BG-a’s noise-floor and growth terms.Figure 4 plots minibatch variance estimates across displacement radii and fits the BG-a model.Source: §E.2.1 · p. 32, Figure 4 · p. 33

E.2.2 EXPERIMENTAL RESULTS

Across three image-classification datasets, BG-a lowers training loss versus small-batch training using substantially fewer iterations, but its test-loss advantage depends on the dataset and baseline. The paper distinguishes these multi-epoch results from LLM pretraining, where validation loss reflects the optimized objective.

E.2.2 EXPERIMENTAL RESULTS

  • BG-a achieves lower training loss than small-batch training across all three datasets while using only 14–16% of its iterations.Against large batches, BG-a has lower training loss on CIFAR-10 but slightly higher training loss on CIFAR-100 and ImageNet-1K, with similar iteration counts.Source: §E.2.2 · p. 33
  • On CIFAR-10 and CIFAR-100, small-batch training has the lowest test loss, while BG-a improves test loss over large-batch training on both datasets.For ImageNet-1K, test-loss ordering closely follows training-loss ordering.Source: §E.2.2 · p. 33
  • Repeatedly revisiting samples in multi-epoch image training means lower training loss does not necessarily translate into lower test loss, unlike online LLM pretraining.The paper notes that image-classification gradients cease to be independent fresh draws after the first epoch.Source: §E.2.2 · p. 33

F EXPERIMENTAL DETAILS

The experimental details describe ResNet50 vision runs and OLMo2 language-model pretraining, along with baseline selection, scheduler adjustments, and noise-growth ablations. BG-a uses measured noise parameters and adjusts batch size and effective learning rate during training.

F EXPERIMENTAL DETAILS

  • Figures 5–7 pair training-loss curves for fixed and BG-a batch schedules with modeled variance ratios, using dataset-specific fixed batch sizes.The figures cover CIFAR-10, CIFAR-100, and ImageNet-1K, respectively.Source: Figure 5 · p. 34, Figure 6 · p. 34, Figure 7 · p. 35
  • Vision experiments train ResNet50 for 200 epochs on CIFAR-10 and CIFAR-100 and 90 epochs on ImageNet-1K, using SGDM with momentum 0.9.CIFAR models replace the initial 7×7 convolution with a 3×3 convolution; ImageNet uses the standard PyTorch implementation.Source: F EXPERIMENTAL DETAILS · pp. 33–35
  • OLMo2-100M, 600M, and 1B are pretrained on C4 for approximately 20 tokens per parameter, using AdamW and sequence length 2048.The corresponding training budgets are 2B, 12B, and 20B tokens.Source: F EXPERIMENTAL DETAILS · p. 35
  • The small-batch baseline and BG-a share the selected stable initial batch size and tuned base learning rate; the large-batch baseline uses that rate with scaling but no additional tuning.For BG-a, batch-dependent scaling jointly adjusts the learning rate during training: linearly with batch multiplier for SGDM and by its square root for AdamW.Source: F EXPERIMENTAL DETAILS · p. 36
  • BG-a noise-growth ablations measure variance at 50 logarithmically spaced radii with 10 independently sampled directions per radius, without training.Language-model directions scale LM-head coordinates by two before normalization; vision directions use unscaled Gaussian samples.Source: F EXPERIMENTAL DETAILS · p. 36

G DISCUSSION

The discussion frames adaptive batching as especially suited to online training, while distinguishing optimization gains from generalization in multi-epoch image classification. It also identifies scale, method, and evaluation boundaries for the current findings.

G DISCUSSION

  • Adaptive batching lowers validation loss than both small- and large-batch training in online LLM pretraining, while in multi-epoch image classification it lowers training loss but small batches may yield lower test loss.The paper attributes the distinction to sample reuse and small batches’ implicit regularization in image classification.Source: G DISCUSSION · p. 37
  • The observed exponent rises from 0.98 for OLMo2-100M to 1.49 for OLMo2-1B; testing whether this trend continues at larger scales and relates to critical batch size is left for future work.The paper suggests larger models eventually require larger training batches.Source: G DISCUSSION · p. 37
  • The scheduler scales learning rate with batch size and uses AdamW for language models, whereas the analysis covers fixed-step-size SGD and SGDM.The initial batch size is selected empirically because the target noise level is not known a priori.Source: G DISCUSSION · p. 37
  • Experiments cover OLMo2 models up to 1B parameters on C4 with Chinchilla-optimal budgets; broader scales and corpora remain unevaluated.The theory’s upper bounds cover a≤2; extending them above 2 requires controlling higher moments of displacement.Source: G DISCUSSION · p. 37

Related papers

Loading 2610.02355v1…