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

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

cs.LGOct 1, 2026

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

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

Organizations: Elmore Family School of Electrical and Computer Engineering Purdue University West Lafayette, IN 47907, USA

Abstract

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 $Ω(ε^{-(4+a)})$ and establish a matching upper bound in $ε$ -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.

Figures & tables

Figure 1: Empirical stochastic-gradient noise of OLMo2 models on C4 as a function of displacement from initialization. For each radius

r

, we evaluate

x=x_{0}+ru

along ten independently sampled, normalized Gaussian directions and plot the corresponding minibatch variance estimates. The dashed red curve in each panel is the fitted BG-

a

model, (

B^{2}\left\lVert x-x_{0}\right\rVert^{a}+G^{2}

). The fitted parameters are: (a) OLMo2-100M:

B=0.379,G=5.41,a=0.98

; (b) OLMo2-600M:

B=0.132,G=8.64,a=1.38

; and (c) OLMo2-1B:

B=0.0916,G=9.77,a=1.49

. Across all three model sizes, the fitted curve captures both the initial noise floor governed by (

G^{2}

) and the distance-dependent growth governed by

B^{2}\left\lVert x-x_{0}\right\rVert^{a}

, with the fitted exponent

a

increasing with model size. The corresponding results for the vision models are provided in Appendix E.2.1 .

Figure 2: Estimated BG-

a

noise-model parameters across OLMo2 model sizes. We observe that as model size increases, (a)

B

tends to decrease, while (b)

G

and (c)

a

tend to increase.

Figure 3: Results for OLMo2-100M (top), OLMo2-600M (middle) and OLMo2-1B (bottom) pretrained on C4. (a) Validation loss for a fixed small batch (

b=32

for 100M and 600M,

b=64

for 1B), a fixed large batch (

b=512

for 100M and 600M,

b=2048

for 1B), and BG-

a

, with batch size across iterations shown in the inset. (b) Modeled variance ratio for BG-

a

, with parameters

B=0.379,G=5.41,a=0.98

(100M),

B=0.132,G=8.64,a=1.38

(600M) and

B=0.0916,G=9.77,a=1.49

(1B).

Model	Batch Size	Validation Loss	Iterations ( $T$ )	Iterations (%)
OLMo2-100M	Small ( $b=32$ )	$3.4609\pm 0.0037$	30,518	100.00%
	Large ( $b=512$ )	$3.5139\pm 0.0071$	1,908	6.25%
	BG- $a$	$3.4347\pm 0.0085$	2,673	8.76%
OLMo2-600M	Small ( $b=32$ )	$2.9188\pm 0.0062$	183,106	100.00%
	Large ( $b=512$ )	$2.8710\pm 0.0079$	11,444	6.25%
	BG- $a$	$2.8482\pm 0.0098$	14,824	8.10%

Table 1: Validation performance and optimization cost of OLMo2 models pretrained on C4. Validation loss is reported as mean

\pm

std across

3

seeds, with iteration percentages relative to small batch runs.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Batch Size	HellaSwag	CommonsenseQA
OLMo2-100M	Small ( $b=32$ )	$28.98\pm 0.08$	$29.03\pm 0.06$
	Large ( $b=512$ )	$27.97\pm 0.25$	$27.27\pm 0.01$
	BG- $a$	$29.15\pm 0.30$	$29.07\pm 0.46$
OLMo2-600M	Small ( $b=32$ )	$40.43\pm 0.11$	$32.43\pm 0.04$
	Large ( $b=512$ )	$43.05\pm 0.15$	$32.92\pm 0.11$
	BG- $a$	$43.11\pm 0.32$	$33.33\pm 0.34$

Appendix

Table 2: Downstream performance of OLMo2 models pretrained on C4. Entries are accuracy (%), reported as mean

\pm

std across

3

seeds.

Figure 4: Empirical stochastic-gradient noise of ResNet50 as a function of displacement from initialization. For each radius

r

, we evaluate

x=x_{0}+ru

along ten independently sampled, isotropic Gaussian directions and plot the corresponding minibatch variance estimates. The dashed red curve in each panel is the fitted BG-

a

model, (

B^{2}\left\lVert x-x_{0}\right\rVert^{a}+G^{2}

). The fitted parameters are: (a) CIFAR-10:

B=24.9

G=305

a=1.25

; (b) CIFAR-100:

B=29

G=337

a=1.19

; and (c) ImageNet-1K:

B=7.01

G=56.3

a=0.96

Figure 5: Results for ResNet50 trained on CIFAR-10. (a) Training loss for fixed batches

b=64

b=512

, and BG-

a

, with batch size across iterations shown in the inset. (b) Modeled variance ratio for BG-

a

, with parameters

a=1.25,B=24.9,G=305

Figure 6: Results for ResNet50 trained on CIFAR-100. (a) Training loss for fixed batches

b=64

b=512

, and BG-

a

, with batch size across iterations shown in the inset. (b) Modeled variance ratio for BG-

a

, with parameters

a=1.19,B=29,G=337

Figure 7: Results for ResNet50 trained on ImageNet-1K. (a) Training loss for fixed batches

b=128

b=1024

, and BG-

a

, with batch size across iterations shown in the inset. (b) Modeled variance ratio for BG-

a

, with parameters

a=0.96,B=7.01,G=56.3

Dataset	Batch Size	Train Loss	Test Loss	Iterations ( $T$ )	Iterations (%)
CIFAR-10	Small ( $b=64$ )	$0.000262\pm 0.000020$	$0.2253\pm 0.0086$	156,200	100.00
	Large ( $b=512$ )	$0.000204\pm 0.000021$	$0.2508\pm 0.0028$	19,400	12.42
	BG- $a$	$0.000200\pm 0.000009$	$0.2412\pm 0.0028$	23,800	15.24
CIFAR-100	Small ( $b=64$ )	$0.002215\pm 0.000128$	$0.9522\pm 0.0052$	156,200	100.00
	Large ( $b=512$ )	$0.001496\pm 0.000009$	$0.9951\pm 0.0244$	19,400	12.42
	BG- $a$	$0.001746\pm 0.000193$	$0.9609\pm 0.0160$	23,440	15.00

Appendix

Table 3: Training and test performance of the ResNet50 model trained on CIFAR-10, CIFAR-100 and ImageNet-1K. Entries are loss, reported as mean

\pm

std across

3

seeds. Iterations are averaged across seeds and given as total iterations and percentage relative to the small-batch baseline within each dataset.

Dataset	Method	Epochs	Batch size	Base LR	Time (min)
CIFAR-10	Small batch	200	64	0.05	338
	Large batch	200	512	0.4	206
	BG- $a$	200	64 ∗	0.05 ∗	240
CIFAR-100	Small batch	200	64	0.05	337
	Large batch	200	512	0.4	206
	BG- $a$	200	64 ∗	0.05 ∗	240

Appendix

Table 4: Vision training configurations and mean wall-clock times, using four NVIDIA A100 GPUs. For BG-

a

, the batch size and base learning rate refer to the initial batch size and its associated base learning rate. (

*

) means that the given values are the initial starting values.

Model	Method	Tokens	Batch size	Base LR	Time
OLMo2-100M	Small batch	2B	32	0.004	28m 37s
	Large batch	2B	512	0.016	16m 13s
	BG- $a$	2B	32 ∗	0.004 ∗	17m 20s
OLMo2-600M	Small batch	12B	32	0.001	12h 58m
	Large batch	12B	512	0.004	9h 03m
	BG- $a$	12B	32 ∗	0.001 ∗	9h 17m

Appendix

Table 5: Language model pretraining configurations on C4 and mean wall-clock times, using four NVIDIA H100 GPUs. For BG-

a

, the batch size and base learning rate refer to the initial batch size and its associated base learning rate. (

*

) means that the given values are the initial starting values.

Model	Dataset	H100 GPUs	Time
ResNet50	CIFAR-10	1	2m 44s
ResNet50	CIFAR-100	1	2m 38s
ResNet50	ImageNet-1K	1	11m 49s
OLMo2-100M	C4	1	50m 08s
OLMo2-600M	C4	4	2h 34m 28s
OLMo2-1B	C4	4	2h 56m 39s

Appendix

Table 6: Computational resources and wall-clock times for the BG-

a

noise growth ablation studies.

	OLMo2-100M	OLMo2-600M	OLMo2-1B
Parameters	100M	600M	1B
Layers	12	16	16
Hidden dimension	512	1344	2048
Attention heads	8	16	16
Head dimension	64	84	128
FFN hidden dimension	2048	5376	8192

Appendix

Table 7: Architectural details of the OLMo2 models used in our experiments. All models use a decoder-only Transformer architecture with SwiGLU feed-forward layers, RMSNorm, QK-normalized attention, and rotary positional embeddings.

Stochastic Optimization

Stochastic optimization studies how to find good decisions or model parameters when objectives, data, or feedback are random or noisy.

See all 209 papers

May 18, 2026

Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates

It is widely believed that stochastic gradient descent (SGD) performs significantly worse than adaptive optimizers such as Adam in pre-training Large Language Models (LLMs). Yet the underlying reason for this gap remains unclear. In this work, we attribute a large part of the discrepancy to SGD's inability to sustain learning rates comparable to Adam's much larger effective learning rates. Through empirical and theoretical analysis of LLM pre-training dynamics, we identify that training is characterized by small gradient norms and large weight-to-gradient ratios, an effect that becomes more pronounced with larger batch sizes typical in pre-training, necessitating such large effective learning rates. However, we find that output-layer gradient magnitudes become highly uneven across token classes, and that large gradient spikes frequently occur during training. Together, these effects severely restrict the admissible learning rate of SGD. Guided by this understanding, we show that simple clipping mechanisms that stabilize SGD at large learning rates enable it to recover most of Adam's performance. In our large-scale experiments, the validation loss gap between large-learning-rate SGD and Adam shrinks from more than 50% to only about 3.5% when pre-training a 1B-parameter LLaMA model with a 1M-token batch size.

Athanasios Glentis, Dawei Li, Chung-Yiu Yau +1

University of Minnesota

Language Model Pretraining Gradient Clipping Adaptive Gradient Methods Stochastic Gradient Descent

Ask PDF

Oct 6, 2026

The Best Optimizer Depends on Batch Size

A plethora of new adaptive optimizers are designed to efficiently estimate and use minibatch gradient statistics to shape parameter updates, but they are typically benchmarked at a single batch size. Hyperparameter scaling rules promise to preserve performance as batch size and gradient noise change, suggesting that the best optimizer at one batch size should remain the best at another. We challenge this approach to developing and evaluating optimizers by showing: (1) no principled scaling rule for Muon works consistently across training settings, and (2) the best optimizer for language model pretraining changes with batch size even after extensive hyperparameter tuning. To understand these failures, we analyze a noisy quadratic model and show that the batch-size scaling rule that preserves performance is direction-dependent, shifting from square-root to linear as a direction's curvature-to-noise ratio rises. We further prove that gradient noise drives a bias-variance tradeoff in preconditioning that can reverse optimizer rankings across batch sizes, even under optimal learning rate and momentum tuning. Consistent with this, when we increase the batch size

16\times

partway through language model pretraining, keeping small-batch updates in only the sharpest

0.001\%

of directions removes up to

59.5\%

of the resulting loss gap over the next

8\%

of training, while a random subspace of the same size barely helps. Taken together, our results show that principled optimizer design must account for batch size and empirical benchmarking should be performed across batch sizes.

Xingyu Dang, Kaiyue Wen, Sadhika Malladi

Princeton University · Stanford University · University of California, San Diego

Deep Learning Optimization Hyperparameter Transfer Muon Optimizer

Ask PDF

Sep 30, 2026

How Does Local Landscape Geometry Evolve in Language Model Pre-Training?

The scale and expense of pre-training language models make efficient hyperparameter tuning essential, yet a principled guidance is still missing. In this work, we analyze language model pre-training dynamics from a local landscape geometry perspective. Our study reveals two distinct phases. In Phase I, sharpness of the local landscape is initially high, leading to instability and loss plateaus under large learning rates (LRs). The landscape shifts from sharp to flatter regions early in training. This dynamic explains the necessity of LR warmup and further suggests that larger peak LRs require proportionally longer warmup periods. In Phase II, the local landscape is governed by the gradient noise scale. Our theory identifies a depth flatness trade-off: high noise from smaller batches widens the loss basin, whereas reduced noise from larger batches deepens it. This theory motivates a dynamic batch-size (BS) scheduler that begins with a small BS and increases it late in training. Together, we provide a unified view of loss landscape evolution, which translates into actionable tuning strategies for large-scale pre-training.

Zhanpeng Zhou, Yuhan Sun, Bingrui Li +4

Shanghai Jiao Tong University · Tsinghua University · Peking University

Language Model Pretraining Learning Rate Scheduling Neural Network Training Dynamics Loss Landscape Geometry

Ask PDF

Abstract

Stochastic Optimization

Explore similar work

Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates

The Best Optimizer Depends on Batch Size

How Does Local Landscape Geometry Evolve in Language Model Pre-Training?