How Does Critical Batch Size Scale in Pre-training?
Authors: Hanlin Zhang
Depen Morwani
Nikhil Vyas
Jingfeng Wu
Difan Zou
Udaya Ghai
Dean Foster
Sham M. Kakade
Venue: ICLR 2025 Poster (2025)
Keywords: Pre-training
Language Models
Data Parallelism
Optimization
Primary area: foundation or frontier models, including LLMs

Abstract: Training large-scale models under given resources requires careful design of parallelism strategies. In particular, the efficiency notion of critical batch size (CBS), concerning the compromise between time and compute, marks the threshold beyond which greater data parallelism leads to diminishing returns. To operationalize it, we propose a measure of CBS and pre-train a series of auto-regressive language models, ranging from 85 million to 1.2 billion parameters, on the C4 dataset. Through extensive hyper-parameter sweeps and careful control of factors such as batch size, momentum, and learning rate along with its scheduling, we systematically investigate the impact of scale on CBS. Then we fit scaling laws with respect to model and data sizes to decouple their effects. Overall, our results demonstrate that CBS scales primarily with data size rather than model size, a finding we justify theoretically through the analysis of infinite-width limits of neural networks and infinite-dimensional least squares regression. Of independent interest, we highlight the importance of common hyper-parameter choices and strategies for studying large-scale pre-training beyond fixed training durations.

Published as a conference paper at ICLR 2025
How DOES CRITICAL BATCH SIZE SCALE IN PRE- TRAINING?
Hanlin Zhang' Depen Morwani 1.2
Nikhil Vyas'
Jingfeng Wu?
Difan Zou" Udaya Ghai
Dean Foster
Sham Kakade 1.2
'Harvard University University of California, Berkeley
Kemmper Institute, Harvard University
"The University of Hong Kong
AAaazon
ABSTRACT
Training large-scale models under given resources requires careful design of par- allelism strategies. In particular, the efficiency notion of critical batch size (CBS), concerning the compromise between time and compute, marks the threshold be- yond which greater data parallelism leads to diminishing returns. To operational- ize it, we propose measure of CBS and pre-train series of auto-regressive language models, ranging from 85 million to 1.2 billion parameters, on the C4 dataset. Through extensive hyper-paranee rrper-parameter sweeps and careful control of factors such as batch size, momentum, and learning rate along with its scheduling, we systematically investigate the impact of scale on CBS. Then we fit scaling laws with respect to model and data sizes to decouple their effects. Overall, our re- sults demonstrate that CBS scales primarily with data size rather than model size, a finding we justif'y theoretically through the analysis of wndthite-width limits of neural networks and dent interest, we highlight the importance of common typer-parameter choices and strategies for studying large-scale pre-training beyond fixed training dura- tions.!
least squares regression. Of indepen-
INTRODUCTION
Efficient optimization is critical in pre-training large models (LMs) at scale (McCandlish et al., 2018; Shoeybi et al., 2019; Kaplan et al., 2020). particular, large-batch training key to acceler- ating training, as it enables more efficient parallelism across hardware accelerators (You et al., 2017; Goyal et al., 2018). Specifically, understanding the scaling behavior of the critical batch size (CBS) is essential for optimizing data parallelism, it defines the point beyond which increasing the batch size may result computational efficiency degradation. Below the CBS, approximately linear scal- ing is achievable- doubling the batch size can proportionally reduce the number of optimization steps required to reach a target loss. However, beyond this threshold, further increases in batch size would lead diminishing returns, making it essential balance computational efficiency with model performance (Shallue et al., 2019; McCandlish et al., 2018). This trade-off presents a chal- lenge for studying pre-training given resource constraints as practitioners are compelled navigate difficult decisions in balancing compute, data, and training time.
We investigate the scaling laws governing CBS in the context of autoregressive transformer-based language modeling (Vaswani, 2017; Radford et al., 2018). Analyzing CBS in pre-train pre-training is chal- lenging due the absence of precise formalism relating to model and data sizes in the literature (Kaplan et al., 2020; McCandlish et al., 2018). Moreover, the interwined effects of scaling model and data sizes proportionally (Hoffmann al., 2022b) further complicate this analysis. Although previous works study the effects of batch size on optimization performance (DeepSeek-AA et al., 2024: Besiroglu et al., 2024; Porian et al., 2024), two crucial differences are (1) they do not decou- ple model size and data size; (2) they focus optimal batch size that reaches the minimum loss instead of critical batch size. We measure critical batch size as a metric, which represents the batch size that results in certain overhead compared to linear scaling: given certain target validation loss,
Code available https: : h//::////iii/ /ttpt/gttbb/.e////////////...t com/N /h1 10a
09/crits tical batch-h ze.

---

Published as a conference paper at ICLR 2025
Empirical Takeaways:
1. In Chinchilla settings, CBS increases when model size N and data size D (or training duration thereafter, which we will use interchangeably) are jointly scaled up (Figure 1, left),
2. If we scale up training duration D while keeping N fixed (Figure 1, middle), the critical batch size increases to similar degree.
3. However, we find that CBS remains nearly invariant when scaling up while keeping fixed (Figure I, right), suggesting that CBS weakly depends model size N but more strongly depends data size D.
4. Our experiments small 151M proxy models provide insights into range of common hyper- parameters and optimization configurations, including transformer context length adjustments, and scaling strategies based on width versus depth, among others.
we measure the number of steps to reach it for different batch sizes; we derive the transition points that incur certain overhead when doubling the batch size; then we fit scaling laws w.r.f. model size and data size to systematically study the scaling of CBS.
logy (ngg(*sStepppsseeppp at 512)
Chinchilla Setting
Fixed Data 3
Fixed Model
(hipetions Lott)
Chegec Locd
Tlkker Spet
ISM (54)3.0)
ESH 371
27 20:
231M15.04:
333M 241
620M [3.32]
0:03
2.28 (27361
1.30 10.00
94 (1)
4
2
240
2
24
2
2
2.00
241
212
2
2
2.0
2
2
245
Batch Size (2")
Batch (2*)
Batch (2*)
Chinchilla Setting
Fixed Data Size 3.073
Fixed Model Size 302M
Dora Type Rited
Data Ypc Rited Porscasted
Data Type Fited -i+ Toracasted
Critial Batch Size
forecasted
2000
1000
-
500
1000
1500
2000
500
1000
1500
2000
40
Model
Model
Token (in B)
Figure 1: Optimization efficiency and scaling of critical batch size in Chinchilla (left) and controlled (middle, right) settings. To study the effect of CBS across different model sizes, we track the relative number of steps required to reach a certain target validation loss. In the Chinchilla setting (left), we keep the data-to-moded size ratio D/N Cenin constant and observe that CBS increases with scale. However, when controlling for either model size (middle) or data size (right), the growth in target losses becomes mostly dependent on data size rather than model size (Section 3).
1.1 EMPIRICAL TAKEAWAYS
Conceptually, we formalize the notion of critical batch size and examine the independent effects of both model and data size. We start by scaling up data size tandem with model size, as suggested in the Chinchilla compute-optimal framework (Hoffmann et al., 2022b). Through controlled studies, we propose scaling laws that decouple the growth of critical batch size from model and data size, leading to the following takeaways an aspect underexplored in previous research.
Overall, our empirical finding that CBS scales primarily with data size implies that when scaling up data, one reduce serial training time through greater data parallelism due to the increase of CBS, without a loss in computational efficiency that can be measured by floating point operations (FLOPs).
2

---

Published as a conference paper at ICLR 2025
1.2 THEORETICAL IMPLICATIONS
Theoretically, maximal update parameterization suggests that, beyond certain point, increasing the width of the neural network (while keeping data size fixed) does not further increase the critical batch size. In contrast, by analyzing a simple least-squares regression with mini-batch SGD, provide a theoretical basis for how the critical batch size continues to scale with increasing data size. Specifically, we introduce two informal theorems here and refer readers to Section 4 for more details.
Theorem (Informal version of Theorem 2). In infinite width regimes (Yang & Hu, 2. training dynamics and performance of the networks become effectively independent of the model size. Con- sequently, the critical batch size remains nearty invariant when scaling up the model size beyond this point, indicating that larger models do not require proportionally larger batch sizes to achieve optimal training efficiency.
Corollary 1 (Informal version of Corollary 2). Consider mini-batch SGD with D samples in the least square problems under power-law source and capacity conditions. The CBS, which enables mini-batch SGD to achieve the minimal expected excess risk while ensuring the fastest possible serial runtime, is given by B* (D) @(D)), where the exponent 2 0 is determined by the exponents of the source and capacity conditions. In the regime where the variance error tends to be dominant, we have 0- < 1/2, indicating CBS grows with data size.
EXPERIMENTAL DESIGN AND EMPIRICAL FINDINGS
We describe the experimental settings and refer readers to Appendix D for extra details. Throughout this paper, we use the abbreviations 'M' for million, 'B' for billion, and 'T' for trillion.
2.1 EXPERIMENTAL SETTINGS
Model and training details. We train a series of autoregressive LMs with a context length of 512 in different sizes ranging from 85M, 151M, 302M, 604M to 1.2B (Appendix D Table 2) on C4 (Raffel et al., 2020) using Adam (Kingma, 2014) with optimizer-speciffe haper-parameters reported in Table 3. We adopt the tokenizer of Eleuther Al's gpt-neox-20b that has vocabulary of size 50280. We use small 151M proxy models to analyze hyper-parameters in most ablation studies. We set the micro batch size smaller than the global one and use gradient accumulation to simulate the effects of large global batch sizes. We focus on fully synchronized distributed data parallelism scenarios where communication is frequent, which simplifies the evaluation and abstracts actual wall clock savings into a total number of optimization steps. More details on optimizer configurations and evaluation strategies are included in Appendix D.
Experimental design and outline. To study CBS in Chinchilla settings, we need to consider target loss on a holdout validation set and measure the number of optimization steps required to reach it. We consider the optimal batch size as Bopi in the linear scaling regime that incurs no efficiency overhead as detailed Appendix C. We consider the validation loss of an optimal batch size Byp = 256 at step &Chin Cchin N/(ctx_len x B) for each model size N and batch size B with context length cex_len set to be 512, where Cchin is the Chinchilla coefficient.
When scaling up model size jointly with data size, the above implies that each model size would have a different target loss. Achieving such a goal through the procedure above is challenging not only due to the combinatorially many each model size and batch size. Below, we outline several key aspects to approach the goal:
but also the unknown training dynamics of
1. As we focus on the number of training steps needed to achieve target validation loss, learning rate decay strategies typically require predefining the total training duration (Loshchilov & Hut- ter, 2022; Hu et al., 2024; Hägele et al., 2024; Defazio et al., 2024). To address this, propose using exponential weight averaging (EWA) (Polyak & Juditsky, 1992) to achieve the desired tar- get validation loss, simple approach that matches other popular choices (Figure 2). This enables training beyond fixed durations or data size, allowing to resume training from checkpoints until the target validation loss is achieved (Section 2.2).
3

---

Published as a conference paper at ICLR 2025
2. Training with proper hyper- paper-parameters: ensuring proper sweeps of momentum and learning rate (Appendix A); adopting well-tumed values for the 32 parameter and the exponential weight aver- aging decay rate т, tailored to each batch size (Appendix B).
2.2 TRAINING BEYOND FIXED DURATIONS FOR REACHING TARGET VALIDATION Loss
Benchmarking learning rate schedulers. In practice, LMs are usually trained using fixed token budgets (Hoffmann al., 2022b), which can determine the total number of iterations the training would undergo. This training process can be easily decomposed into learning rate warmup and decay phases so that lower learn- ing rate is kept at the end of training to enable better optimization. However, our goal is to find the optimal performing run under various myper-parameteer and optimization conditions. This implies non-trivial decision regarding selecting the maximum training duration (De- fazio et al., 2024; Hägele et al., 2024). As train- ing beyond fixed durations is particularly favor- able in many large-scale pre-training scenarios, we benchmark recently proposed methods like schedule-free optimizer (Defazio al., 2024), cosine, warmup-stuble-decan schedule (WSD) (Hu et al., 2024) (or trapezoidal (Zhai et al., 2022)) and our proposed constant+EWA strat- egy which maintains a running average of model weights ε111 = T optimization, where 0, is the actual model parameter at step / and we use & for evaluation. To ensure the baselines are close to optimal, we first get the optimal step counts from our constant+EWA runs, evaluate WSD scheduling by testing [0.1, 0.2, 0.3] multiples of those step counts. In parallel, for the cosine scheduler, we explore [0.9, 1.0, 1.1] times the same step counts. We show that our EOA strategy match the efficiency of cosine scheduling and WSD), especially for large batch sizes (Figure 2). The connections are explained in detail in (Morwani et al., 2025).
Target Validation Loss 3.24
17
Scheduler
16-
WSD
1697 (#Steps)
Cosine
15-
Schedule Free
Constant
14
111
11
2
210
2
2
Batch Size (2*)
Figure 2: Comparing and accounting for train- ing dynamics. Throughout, we adopt Con- stant+EWA since performs the best for large batch sizes and avoids setting fixed training du- ration beforehand for reaching a target loss.
(1 - r) O2 to improve Et +
Prior research has shown the generalization (Izmailov al., 2018) and optimization (Karras al., 2024) benefits of EWA, our findings further reveal that EWA can trade memory for optimization efficiency in LM pre-training, especially in large-batch regimes. This useful in scenarios where target loss must be achieved, but practitioners are uncertain of the exact maximum data size to set up the learning rate schedule for training.
Takeaway on learning rate scheduling: EWA consistently improves model training efficiency compared using constant learning rate without it. FWA proves to be effective approach pared to other baselines with decaying schemes, offering competitive performance while climinating the need to predefine training durations.
2.3 ABLATION ON MODEL CONTEXT LENGTH
We adopt 512 as the context length of LMs for all of our experiments but is unclear how would impact the efficiency of training and whether the scaling of CBS would vary when we enlarge the context length. So we sweep over several larger windows 2 211 212 (Section 2.4) Overall all models in four different context lengths have very similar relative optimization efficiency across various batch sizes and thus justifies our use 512 for all the experiments.
Takeaway on model context length: Different context lengths (2 ~ 222) have similar scaling w.r.t.
batch size.
4

---

Published as a conference paper at ICLR 2025
2.4 ABLATION ON MODEL WIDTH AND DEPTH
Model sizes can typically be scaled up in two main ways: by increasing the width, which involves enlarging the hidden size of the multilayer perceptron (MLP), or by increasing the depth, which entails adding more layers to the network. As the main result in Figure only involves a single way for scaling up models (Table 2), e.g. 604M model has 2x width than the 151M one. To explore alternative scaling strategies, we investigate how the model behavior changes when we scale the 604M model by increasing the depth by 4x instead (detailed configurations in Table 5).
steps at 512)
Width (Target Loss)
De (Target Loss)
Model Size (Target Loss)
20
4096 (3.24)
(3:24)
151M (3:24)
8192 (2.92)
24 (3:07)
604M-width (2.92)
2
10240 (2.82)
48 (2:94)
604Mi (2.94)
2-
2-
2
2'
2
2
2 2
2
2
200
2
2
2
2
2
2
22
Batch Size (2")
Batch Size (21)
Batch Size (2*)
Figure 3: Scaling up width and depth shares similar efficiency gain for compute-optimal training.
Firstly, as shown in Figure (left, middle), under Chin- chilla scaling- data and model size grow propor- tionally- increasing model depth or width has similar impact on CBS. Notably, according to previous results, the rise in CBS in compute-optimal scaling should be at-
loga (#$teps/##tepp at B5 2
Target Validation Loss 3.24
Context Length
512
tributed to the increased data size. Then through con-
2048
trolled comparison (Figure 3, right), we see that using
4096
two different ways to scale 15 models to 604M ones
-3
is equivalent in efficiency since both curves overlap. Our findings may offer practical insights for scaling models under a fixed token budget that allocated in proportion to model size. This is partieularly relevant because scal- ing model width is often favored over increasing depth, as wider models tend to be more amenable parallelization without incurring additional latency overhead (Shoeybi al., 2019; Touvron et al., 2023; Erdil, 2024; McLeish et al., 2025).
4-
2.5
2
2.5
2.5
2.
2
2
Batch Size (#Tokens)
Figure 4: Ablation results on context length using 151M models.
Takeaway on scaling transformer width and depth in compute-optimal regimes: Incrcasing width and depth has similar effects in critical batch size for cumpute-optimat pre-training.
CRITICAL BATCH SIZE SCALING LAW
3.1 FORMAL DEFINITION OF CRITICAL BATCH SIZE
Recall that CBS is the transition point where increasing the batch size by factor of k. leads to reduction in the required number of training steps by a factor that is less than k. We now define CBS as the batch size that leads to 20% overhead compared to linear scaling. First of all, define R(N, D, B) as the best loss achievable for a model of size N using single pass on D tokens with batch size B. This would be obtained by optimally tuning all other parameters of the optimizer, while keeping N. D, B fixed. Below is the formal definition of CBS:
Definition 1. Define Ropi(N, D) = ming arg min R(N, D. B), as the minimal loss achieved optimizing over batch size and the optimal batch size respectively. We define fN,D(B) to be the number of steps required reach Rope N, D) as a function of batch size B. Clearly JN.Dl (Bop)) D/ Begr. To
R(N, D. B), (N, (N,NN,) D)
5

---

Published as a conference paper at ICLR 2025
define the Critical Batch Size, B* D). define linear scaling curve fN.D(B) D/B. f matches f at Bog and then scales down linearly as batch size goes up. (N, is defined as the maximum batch size B' >> Bopi N, D) such that JN,D(B') S (.2)
As illustrated in Figure 5, (NV, D) is the batch size at which the number of steps is 20% higher than what is predicted by the linear extrapolation from
the optimal batch size. Note here that 20% can be replaced by any other suitable measure of increase from linear scaling.
Reaen Target Loss
2
3.2 SCALING LAWS W.E. MODEL
2
SIZE FOR CHINCHILLA-OPIMAN PRE-TRAINING
Relee no
6o (dd8)-al Dal gunear)
As observed in all the results above, doubling the batch size for larger models allows them to more ef- ficiently reduce the relative number of steps needed to reach the target loss. We ask whether those in- creased efficiencies can predicted via scaling law.
12 (20% Overhoas)
Ruy (Opdinal Baten Ston
Bept
(Cattor Datcs Stee)
2
Batth Size
Figure 5: Illustration of critical batch size,
where B = 2 and context length is 512 by default.
We begin our first step by fitting a power law of batch size (B) to the absolute number of steps (Y) to reach the target loss lYg = log((a a++ B- and then derive the critical batch size. Then we derive the CBS via B = 5g + 1.2B. ope. which is implied by a transition point where the total amount of data under this batch size would incur 20% overhead compared to
linear scaling: Dioaa (a- + b/ Bogr) ) * e.p = (a+ b/ B*) B, where a = 1, Bugi is set to be 256 chosen lic within the linear scaling regime as suggested in Appendix C. We report the parameters fitted to the power law relationship between the number of steps Y and batch size B in Appendix Table 6. We adopt the fixed a solution, as both strategies yield nearly identical forecasting results.
Secondly, we fit power law log = log(c+ ) with respect to the model size N (in million). The constant term e is set to be 0 by default (as B should be 0 at N = 0), which leads to B 93.20 NO.47, We visualize the curve fitted in Figure (left) and report more forecasts Table in the Appendix.
Overall, we observe an increase in CBS when scaling up compute-optimal training: In Figure (left), we have target losses selected according to chinchilla steps and we fit the power law of critical batch size with respect to model sizes N (in million) as B 93.20 * No.447 Our results suggest that critical batch size around 29 to 2 would be helpful to efficiently optimize models below IB on Chinchilla-optimal amount of tokens to study other empirical problems. However, it is common for the number of tokens trained to scale proportionally with the model's parameter count. So it is unclear whether the growth of CBS because of the increase in (1) model size or (2) the data size/training duration, a question we explore in the next subsection.
3.3 DECOUPLING CBS SCALING LAWS W.E. DATA SIZE AND MODEL SIZE
Controlled comparison with the same data size. Firstly, use the Chinchilla token size 3.072B of 151M models tchin to record the target validation loss for each model size and train all the 302M, 604M, 1.2B models with a smaller duration again to reach these target losses. To optimize for performance when training on fewer tokens, we also tune the warmup steps accordingly. Figure (top right) shows that all the curves behave similarly and we observe almost no increase in CBS when enlarging the model size. Moreover, we fit scaling law with respect to model size thereafter Figure (bottom right): keeping the data size fixed leads to a scaling law 8* 621.341 * N0.087 weakly dependent model size.
Controlled comparison with the same model size. Moreover, focusing on the 302M models, we conduct additional experiments by selecting target losses at 0.28x, 0.5x, 2x, and 4x the Chinchilla step for batch size 256 runs. This setup results in two under- under-training training and two over-training config- urations. To achieve optimal performance in the over-training scenarios, increase the warm-up

---

Published as a conference paper at ICLR 2025
ratio accordingly, while for the under-training cases, reduce the warm-up ratio proportionally. Results in Figure 1 (middle) show that as we enlarge the number of tokens being trained on, we see an increase of CBS, similar to what we have observed for training large models on chinchilla target loss. This can also be seen in the forecasted CBS curves shown in Figure 1 (middle) which shows increase of CBS, similar what
that enlarge the number of tokens being trained on, we have observed in the Chinchilla setting where model and data size are scaled up proportionally.
We also plot the results for scaling both N and D (Figure 1, left) and only scaling D (Figure 1, right) together in Figure 6. In the side-by-side comparison, we observe the following trends: (i) In the Chinchilla setting (indicated by the first column in the legend), models of various sizes (85M, 151M, 604M, 1.2B) trained on different token amounts exhibit an increase in critical batch size as scale grows. (ii) Additionally, each pair of curves with the same color overlaps significantly, indicating that models of different sizes trained on the same token quantity tend to have similar critical batch sizes. (iii) Finally, when model size is held constant and only the data size (second column in the legend) varies, we also observe increase in critical batch size with scale. Therefore, we can qualitatively understand that the increase of CBS is likely to be agnostic to model sizes but due to the increase in training duration.
logo StStpp//# Steps at BS 512)
0
-1-
-2-
Model (Target Loss]
85M (3.42)
302M (3.27)
-3
151M (3. 24)
302M (3.16)
302M (3.00)
1.28 (2.74)
302M (2.94)
29
2
2
212
23
Batch Size (2*)
Figure 6: Controlled comparison by training 302M models varying amounts of tokens and then comparing with other model sizes trained in similar amounts of tokens. Models with the same color or positioned in the same row of the legend represent this comparison. For fixed token count of 3.072B, we measure the target loss at that step for each model size.
Takeaway on scaling laws for critical batch sizes: Based on the scaling laws and controlled com- parisons, we conclude that the increase in CBS Chinchilla-optimal training is more strongly at- tributed extended data size or training durations rather than the increase of model size.
THEORY ON SCALING OF CRITICAL BATCH SIZE
Our experimental results show that CBS increases with larger data sizes but remains (nearly) invari- ant when scaling up the model size. We now formally investigate this observation using theoretical analysis for both scenarios.
4.1 FIXED DATA SIZE AND SCALING UP MODEL SIZE
Various previous works have established infinite width limits of neural networks (Yang Hu, 2021; Bordelon & Pehlevan, 2022). For initializations and architectures that obey these limits, we can theoretically claim, that for a fixed training duration and batch size, the performance of the neural networks asymptotes with increasing width. The formal statement is provided below:
Theorem 2.
For SGD with a given batch size B for gradient descent, i.e., - oc). training iterations t, an error tolerance e > 0, fixed learning rate schedule and data ordering, for any network and initialization satisfying Master Theorem (Theorem G.4) in Yang & Hu (202 I), there

---

Published as a conference paper at ICLR 2025
exists width such that for any two networks M1, M2 having widths W1, W2 w. [R(M1 t) R(M2, t) I e, where R(M.t) t) denotes the loss of network M at time t.
Proof of Theorem 2. The proof follows from the fact that the trajectory of the network approaches limit as width tends to oc, and thus, by definition of limits, there exists a width w, such that for any two networks with a width greater than w. their loss at time differs by at most c.
Note that the assumption of fixed learning rate schedule with increasing width might seem strong, but recent works (Yang et al., 2022) have shown, that one of these initializations, termed as Maxi- mal Update Parameterization (nPP), exhibits hyperparameter transfer with width. This initialization scheme has also recently gained popularity because of this property and has been used by many open-source implementations (Dey al., 2023a;b; Liu et al., 2023; Hu al., 2024). Moreover, works (Yang et al., 2022; Vaa et al., 2023) have empirically demonstrated that with μp, networks start exhibiting consistent loss curves practical widths.
Moreover, as the above theorem holds for fixed batch size B as well as B + oo, we expect that there exists a finite width w such that the above theorem holds for all batch sizes B. Thus, for fixed training tokens, we would expect that the critical batch size won't It scale with model width beyond a point. Although we have mostly talked about scaling model width, note that some recent results have also established such limits for infinite depth ResNets and transformers (Yang et al., 2024; Bordelon al., 2024b;c), and thus the arguments above also hold for these networks.
4.2 FIXED MODEL SIZE AND SCALING UP DATA SIZE
We now turn studying the impact of data size in mini-batch SGD for well-specified Gaussian linear regression problem. Let (x, y) be a pair of covariates and responses from population distri- bution. Let the population risk and the population distribution be
R(w)= = E(x v)² x~~ S~,,0,, H), y/x~N(x w", "²"))")
where w is the trainable parameter, the expectation is over the population distribution, and (H, w* oa)) specify the population distribution. Given D independent samples (x1, from the population distribution, we consider an estimate given by mini-batch SGD,
1 ((+1) Wo = Wi- = w. - 2
Σ (xJW: ,5)x,, t 0, O,...,n--1 jetB
where > 0 is a constant learning rate, B is the batch size, n == D/B is the number of steps, Wi is the initialization (without loss of generality), and the output is the average of the iterates, a E W,. Then the following theorem provides a tight bound on the excess risk achieved by the average of the mini-batch SGD iterates.
w : 1.2
We write ((D) S g(D) there is positive constant e such that f(D) S eg( for every D 2 1. We write J(D) = g(D) if f(D) & 9(D) & (D)) The proofs are all deferred to Appendix G.
Theorem 3. Let (Ac)so be the eigenvalues of H in nonincreasing order. Assume that l|
w" H S 0² Then every y S B/r(H/) H|[ we have
+ o2
h² + (Dn X?
ER(w) - 0² ~
W'Il Wo - + +l/ - w" Wii Dy
D
where k² == max (k: d 2 B/(D)) and the expectation is over the randomness of w.
The proof of Theorem 3 is motivated by Zou et al. (2023). We focus on well-specified Gaussian data distribution for simplicity, but this can be relaxed to misspecified cases under fourth-momen conditions following the results Zou et al. (2023). Theorem suggests that for a fixed data size D, the excess risk depends on the batch size B and the learning rate 7 only through their ratio 7/ B. Moreover, a large 7/ B tends to decrease the bias error (the terms depending on w") but increase the variance error (the terms depending on o²), and vice versa. This observation is exploited in the following corollary, where we compute the CBS that minimizes the sequential running time without sacrificing the rate of the attained excess risk.
8

---

Published as a conference paper at ICLR 2025
Corollary 2. Under the settings of Theorem additionally assume o² 1 and the following capacity and source conditions:
for a, > ARi (V, w/) =i,; EA,(v, w wg) 0 for i + j.
where (di, ve)szo are the eigenvalues and the corresponding eigenvectors of H, and the expectation is over prior of w". Then we have
1. When b S the optimal hyper-parameters (that minimize the expected excess risk up constant factors) are 7" ~ and B* =1
2. When b > the optimal hyper-parameters are " and B such that
0 0<>S1 1SB'SD, //B² D minph
Therefore, the CBS B D2-a/ min 2b, which (along with "- ~ 1) allows mini-batch SGD output w to attain the optimal rate of the expected excess risk (as data size D grows) with the smallest number of steps n.
The capacity and source conditions are from the nonparametric linear regression literature (Capon- netto & De Vito, 2007) and are recently used to study scaling laws theory (Bordelon et al., 2024a; Lin et al., 2024; Paquette et al., 2024). According to Corollary 2, when b I a, the bias error tends dominate the variance error, in this case, the CBS is B optimization steps. When b > a, the variance error tends to dominate the bias error, and the opti- mal choices of batch size and learning rate balance these two errors. While one can use 8 and a small Y to achieve the best excess risk rate, this leads to a suboptimal sequential runtime (n = D)B). this case, the CBS is B ~ D Dl-a/ (b, which achieves the optimal excess risk rate while minimizing the sequential runtime.
allow a maximum number of
Takeaway on the theory of batch size scaling when scaling up data and model size:
As we scale up model size while keeping the data size fixed, pP suggests that critical batch size does not scale with model width beyond point.
Fixing the model size, the critical batch size increases with the training duration. In the context high-dimensional linear regression, where the variance error dominates the bias error, is possible to choose large batch size (as function of the data size) for mini-batch SGD, allowing for reduced sequential runtime without compromising the rate at which excess risk is minimized.
5 RELATED WORK
Scaling laws. Scaling laws describe the parametric relationships among key factors involved in training neural networks: model size N, dataset size D, training cost C, and final training loss R. These laws enable the prediction of training loss R based available resources, making it possible to optimize resource allocation for efficient model training. For example, Hestness et al. (2017) found R D-a with [0.07, 0.35). Of the factors they varied, only tasks change the expo- nent a. Changing the architecture optimizers, regularizers, and loss functions, would only change the proportionality factor, not the exponent; Henighan et al. (2020) studied statistical relations be- tween N. D, C, R, over wide range of values and found similar scaling laws, over the range of N E [103, 109 C [1012 1021], and over multiple modalities (text, video, image, text to image, etc.). (Kaplan et al., 2020) states that N should be scaled faster than D. However, Chinchilla scaling (Hoffmann al., 2022a) found that models are under-trained, and then suggests that when given an increased budget (in FLOPs), achieve compute-optimal, model size N and data size D should scale in approximately equal proportions. Recent efforts (Pearce & Song, 2024; Besiroglu et al., 2024; Porian et al., 2024) have been made in reproducing the scaling laws from (Hoffmann et al., 2022a) and the (Kaplan et al., 2020). Different from our focus on measuring the efficiency notion of CBS, most of them focus deriving optimal hyper-parametent (Bi et al., 2024; Porian et al., 2024) including learning rate and batch size from small-scale training given a fixed compute budget FLOPs 6N without decoupling the effects of model size and data size.

---

Published as a conference paper at ICLR 2025
Optimization and critical batch size. Previous studies have shown that increasing batch sizes can be offset by a proportional adjustment to the learning rate in small-scale regimes (McCandlish et al,, 2018: Zhang et al., 2019; Kaplan et al., 2020; Li et al., 2021). McCandlish al. (2018) introduce the gradient noise scale, a measure that captures the variation gradients across different training examples, which helps predict the critical batch size (CBS). Their findings also suggest that small- batch training is more compute-efficient, while large-batch training requires fewer optimizer steps. Momentum-based methods extend scaling to larger batch sizes but converge the performance of standard SGD at smaller batch sizes (Shallue et al., 2019). Additionally, Zhang et al. (2019) analyze the impact of curvature on CBS using a noisy quadratic model, demonstrating that preconditioning techniques can increase the CBS. Golmant et al. (2018) show that the size of the dataset plays smaller role in determining training efficiency compared to factors like model architecture and data complexity. In contrast, Hilton et al. (2022) examine how performance can be maintained at smaller batch sizes. Mcanwhile, Smith et al. (2017); Smith & Le (2017) empirically investigated how the op- timal learning rate changes based on momentum and training set size. Thcoretical work has further sought to characterize CBS by analyzing SGD behavior in least-squares linear regression, especially in over-pane peee-paraettrizeed settings (Jain et al., 2018; Ma et al., 2018). Filatov et al. (2024) concurrently find that optimal batch size and CBS scale with data size. However, they do not explore how CBS scales with model size for models beyond 354M parameters, nor do they provide theoretical justifi- cations or address the challenge of selecting optimal runs across a broad range of hyperparamntter. Our work advances the optimization literature by formalizing CBS and quantifying its growth w.n.t data size and emphasizing the importance of common hyper-paun parameter choices. It also provides strategies for studying large-scale pre-training beyond fixed training durations.
6 CONCLUDING REMARKS
In conclusion, this study provides an extensive examination of the scaling laws for critical batch size in large-scale autoregressive language model pre-training. By systematically analyzing the relation- ship between model size, data size, and CBS, we found that while CBS increases with data size, it remains relatively invariant to model size. This finding suggests training on more data may enable greater data parallelism in pre-training. We further emphasize the role of key hyperparameters and exponential weight averaging, which can match the performance of cosine scheduling without re- quiring fixed training durations. These insights offer practical strategies for scaling models while maintaining efficiency, which is critical in resource-constrained scenarios.
ACKNOWLEDGEMENTS
We thank Dhruv Madeka. We thank Qirong Ho for discussions on distributed training; Jeremy Berstein for discussions on optimization; Dirk Groeneveld and Luca Soldaini for discussions on OLMo implementations. HZ is supported by an Eric and Susan Dunn Graduate Fellowship. SK and DM acknowledge the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence; SK acknowledges the support from the Office of Naval Research under award NO0014-22 11-2--1--2377 and the National Science Foundation Grant under award #IIS 2229881. NV and DM are supported by Simons Investigator Fellowship. NSF grant DMM-2134157, DARPA grant W91 INF2010021, and DOE grant DE-SC0022199
REFERENCES
Tamay Besiroglu, Ege Erdil, Matthew Barnett, and Josh You. Chinchilla scaling: A replication attempt. arXiv preprint 2rXiv:2404.10102, 2024.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954 2024.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. 2022.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In Forty-first International Conference Machine Learning, 2024a.
10

---

Published as a conference paper at ICLR 2025
Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics, 2024b.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise
hyperparameter transfer in residual networks: Dynamics and scaling limit. 2024c.
Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7:331-368, 2007.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao. Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier- Hellstern,, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
DeepSeek-AI, Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, Huazuo Gao, Kaige Gao, Wenjun Gao, Ruiqi Ge, Kang Guan, Daya Guo, Jianzhong Guo, Guangbo Hao, Zhewen Hao, Ying He. Wenjie Hu, Panpan Huang. Erhang Li, Guowei Li, Jiashi Li, Yao Li, Y. K. Li, Wenfeng Liang, Fangyun Lin, A. X. Liu, Bo Liu, Wen Liu, Xiaodong Liu, Xin Liu, Yiyuan Liu, Haoyu Lu, Shanghao Lu, Fuli Luo, Shirong Ma, Xiaotao Nie, Tian Pei, Yishi Piao, Junjie Qiu, Hui Qu, Tongzheng Ren, Zehui Ren, Chong Ruan, Zhangli Sha, Zhihong Shao, Junxiao Song, Xuecheng Su, Jingxiang Sun, Yaofeng Sun, Minghui Tang. Bingxuan Wang, Peiyi Wang, Shiyu Wang, Yaohui Wang. Yongji Wang, Tong Wu, Y. Wu, Xin Xie, Zhenda Xie, Ziwei Xie, Yiliang Xiong. Hanwei Xu, R. Xu, Yanhong Xu, Dejian Yang. Yuxiang You, Shuiping Yu, Xingkai Yu, B. Zhang, Haowei Zhang, Lecong Zhang, Liyue Zhang, Mingchuan Zhang, Minghua Zhang, Wentao Zhang, Yichao Zhang, Chenggang Zhao, Yao Zhao, Shangyan Zhou, Shunfeng Zhou, Qihao Zhu, and Yuheng Zou. Deepseek llm: Scaling open-source language models with longtermism, 2024.
Aaron Defazio, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, et al. The road less scheduled. arXiv preprint arXiv:2405. 2024.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference Machine Learning, pp. 7480-7512. PMLR, 2023.
Nolan Dey, Gurpreet Gosal, Zhiming, Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Cerebras-gpt: Open compute-optimal language models trained on the cerebras wafer-scale cluster, 2023a.
Nolan Dey, Daria Soboleva, Faisal Al-Khate,b, Bowen Yang, Ribhu Pathria, Hemant Khachane, Shaheer Muhammad, Zhiming, Chen, Robert Myers, Jacob Robert Steeves, Natalia Vassilieva, Marvin Tom, and Joel Hestness. Btlm-3b-8k: 7b parameter performance in a 3b parameter model, 2023b.
Ege Erdil. ing past
Data movement bottlenecks to large-scale model training: Scal-
URL https: /eppcchaippc.....e/ org/blog/ aetta-mvvemetteeee lentllesecksscssse www.ssaallggnsttainnee 1el 11-03.
le28 flop, 2024.
Accessed: 2024-
Oleg Filatov, Jan Ebert, Jiangtao Wang. and Stefan Kesselheim. Time transfer: On optimal learning rate and batch size in the infinite data limit. arXiv preprint arXiv:2410.0582 05.8888. 2024.
11

---

Published as a conference paper at ICLR 2025
Gabriel Goh. Why momentum really works. Distill, 2017. doi: 10.23915/dissill.000066 URL http:///isstill.pbb
Noah Golmant, Nikita Vemuri, Zhewei Yao, Vladimir Feinberg, Amir Gholami, Kai Rothauge, Michael W Mahoney, and Joseph Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint arXiv:1811.1299 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour, 2018.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00838, 2024.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010. 14701, 2020.
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang. and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv: 1712.0049, 2017.
Jacob Hilton, Karl Cobbe, and John Schulman. Batch size-invariance for policy optimization. Ad- vances in Neural Information Processing Systems, 35:17086-17099 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 2022a.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Infor- mation Processing Systems, 15:300066-00334 2022b.
Shengding Hu, Yuge Tu, Xu Han, Ganqu Cui, Chaoqun He, Weilin Zhao, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Xinrong Zhang, Zhen Leng Thai, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, da- hai li, Zhiyuan Liu, and Maosong Sun. MiniCPM: Unveiling the potential of small language models with scalable training strategies. 2024.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations, 2024.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv: 2018
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Paralleliz- ing stochastic gradient descent for least squares regression: mini-batching, averaging, and model misspecification. Journal of machine learning research, ((222)):1442 2018.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Tero Karras, Miik Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 14174-24184 2024.
Diederík Kingma. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
12

---

Published as a conference paper at ICLR 2025
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay bal- ances learning across neural networks. arXiv preprint WrXiv23355.12122 2023.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712- 12725, 2021.
Licong Lin, Jingfeng Wu, Sham M Kakade, Peter L Bartlett, and Jason D Lee. Scaling laws in linear regression: Compute, parameters, and data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Zhengzhong Liu, Aurick Qiao, Willie Neiswanger, Hongyi Wang, Bowen Tan, Tianhua Tao, Junbo Li, Yuqi Wang. Suqi Sun, Omkar Pangarkar, Richard Fan, Yi Gu, Victor Miller, Yonghao Zhuang, Guowei He, Haonan Li, Fajri Koto, Liping Tang, Nikhil Ranjan, Zhiqiang Shen, Xuguang Ren, Roberto Iriondo, Cun Mu, Zhiting Hu, Mark Schulze, Preslav Nakov, Tim Baldwin, and Eric P. Xing. Llm360: Towards fully transparent open-source Ilms, 2023.
I Loshchilov, Decoupled weight decay regularization. arXiv preprint arXiv:1711.05001,22 2017.
Ily Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In Interna- tional Conference on Learning Representations, 2022.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effec-
tiveness of sgd in modern over-parametrized learning. In International Conference on Machine PMLR, 2018.
Learning. pp.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprini 18XX::88/2206612 2018.
Sean McLeish, John Kirchenbauer, David Yu Miller, Siddharth Singh, Abhinav Bhatele, Micah Goldblum, Ashwinee Panda, and Tom Goldstein. Gemstones: A model suite for multi-faceted 2025.
scaling laws. arXiv preprint
Depen Morwani, Nikhil Vyas, Hanlin Zhang, and Sham Kakade. Connections between schedule- free optimizers, ademamix, and accelerated sgd variants. arXiv preprint arX1v:2502.02431, 2025.
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Noua- mane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. arXiv preprint 2rXhv:2305.16266 2023.
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+ 3 phases of compute- optimal neural scaling laws. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Tim Pearce and Jinyeop Song. Reconciling kaplan and chinchilla scaling laws, arXiv preprint arXiv:2406. 12907, 2024.
Boris Polyak and Anatoli Juditsky. Acceleration stochastic approximation by averaging. SIAM journal control and optimization, 30(4):838 055588888555 1992.
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. arXiv preprint arXiv:2406.19) 19146, 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. OpenAl Blog Post, 2018. URL htttps://oppntt www.ppanaaaaoaoooo..p..ommmoooot Index/Languane wttp://nnuuuguuuuuuuiiii..
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67 2020.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951.
13

---

Published as a conference paper at ICLR 2025
Christopher Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20( 12):1-49, 2019.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint vrXiv:/909.000
2019.
Samuel L Smith and Quoc V Le. A bayesian perspective on generalization and stochastic gradient descent. arXiv preprint 1rXiv:1710.06451, 2017.
Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv:1711.00489, 2017.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deep- speed and megatron train megatron-tunine -tturig nlg 530b, large-scale generative language model. arXiv preprint arXiv:2201.11999 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprini arXiv:2302. 13971, 2023.
Vaswani. Attention all you need. Advances in Neural Information Processing Systems, 2017.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cen- giz Pehlevan. Feature-learning networks are consistent across widths realistic scales. 2023.
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities, 2023.
Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural net- works. 2021.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2003.03466, 2022.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learning in infinite depth neural networks. 2024.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 011..77888.8888 2017.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M. Susskind. Stabilizing transformer training by preventing atten- tion entropy collapse. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 40770 40803. PMLR, 23-29 Jul 2023.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp- 12101, 2022.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva. George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- QuanQuan Luo. Adam can converge without any modification update rules. Neural Information Processing Systems, 2022.
14

---

Published as a conference paper at ICLR 2025
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign over- fitting of constant-stepsize sgd for linear regression. Journal of Machine Learning Research, 24 (326):1-58, 2023.
15

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Training Dynamics
B Additional Ablation Studies Studying Adam Optimizer
19
C Results Including Small Batch Sizes
20
D Additional Details on Experiments
21
Additional Details on Scaling Laws
F Reproducibility
24
G Complete Proofs in Section 4.2
24
16

---

Published as a conference paper at ICLR 2025
A TRAINING DYNAMICS
A simple strategy for setting warmup steps. To further prove that the critical batch size actually exists and the saturation of large batch sizes is not an artifact of training well with proper hyper-paragrams, we take into account the warmup fraction in training as well:
We sweep over warmup step ratios (how many fraction of training steps we need to linearly scale the learn- ing rate from zero) over 0.25 and 0.1 and find that 0.25 works best for 85M models. Therefore, we fix this num- ber of warmup steps to be 0.25 of the [Chin for future ex- periments. For 151M models, we sweep over the frac- tion of warmup steps in {0.15, 0.25, 0.35]. We show in Figure that using a warmup ratio of 0.25 can be a rea- sonable design choice it enjoys consistently better per- formance than 0.15 yet only slightly underperforms 0.35. After we find that setting the warmup steps according to this heuristic, we use the ratio proportionally for all the other model sizes. This strategy has also been shown to be effective in (Porian et al., 2024).
Target Validation Loss 3.24
Warmup Steps
12.0
0.15
log; (#Steps)
0.25
11.8
0.35
11.6
11.4
11.2
24
2 Batch Size (2*)
2
Figure 7: Ablation of warmup steps used in the linear LR warmup stage for large batch sizes.
Examining the last part of training. By closely exam- ining the final stages of the training process (Figure 8),
becomes apparent that applying Exponentially Weighted Averages (EWA) can help smooth out noise, allowing the optimization to converge to the target loss more efficiently. For example, very high EWA decay rate would be needed even for a 1.2B model with a moderate batch size of 1024. Moreover, observe that the optimization process is notably influenced by the final phase of train- ing. For instance, by step 10,000, most runs achieve a validation loss below 3.2 (Figure 8a), and similarly, a loss below 2.8 is reached by step 30,000 (Figure 8b). However, reach the target loss of 2.736, the difference between the best and second-best runs grows substantially, with the best run requiring over 5,000 fewer steps.
Evaluation Trajectories (Batch Size 1024)
Evaluation Trajectories (Batch Size 1024)
4.0
EWA
Adam
Validation Loss
0.9999
0.9995
0.998
0.9992
Validation Loss
3.6
3.8
0.95
0.99
0.999
0.9999
3.4-
2.75
3.2
2.70
40000
50000
60000
2.70
60000
45100
50n0c
55066
3-7.736
2.8-2.736
10000
20000
30000
40000
50000
60000
2.6
10000
20000 30000
40000
50000
60000
Training Steps
Training Steps
(a) EWA Ablation
(b) 32 Ablation
Figure 8: A large enough EWA decay rate and Adam B2 is important for long-duration training. We plot the evaluation curves of 1.2B models, as in Chinchilla settings, we scale up data size proportionally to model size. When increasing the number of training tokens, it is crucial to carefully set appropriate values for both B2 and r to effectively account for efficiency.
Scheduler comparison for other batch sizes. Figure 9, we include more comparisons on differ- ent schedulers that are reported in the main text (Figure 1). Overall, our Constant+EWA performs competitively with cosine scheduling and outperforms WSD scheduling, especially for large batch size regimes. Note that sweep over the decay steps as 0. 0.2, 0.3x total training steps for WSD scheduling. We tune cosine scheduling by conducting sweeps over various maximum optimization steps to identify the optimal value, and then rerun the training using this step count. This approach ensures that the model reaches the target loss near the end training, optimizing the performance of learning rate decay. For schedule-free optimizers, tune the 31 0.9, 0.95, 0.98. Under small
17

---

Published as a conference paper at ICLR 2025
batch sizes, the schedule-fre free optimizer (Defazio et al., 2024) is a competitive baseline but it is significantly worse for batch sizes larger than 1024.
Furd Coduarian Papettrrs Jarch Eoe 1251
Gencer
codle
Lega
NS tikel fer
12g
Exriig stps
Trading Steps
Farving steps
1a0ing Mept
(a) Batch size 64
(b) Batch size 128
(c) Batch size 256
(d) Batch size 512
4.2
Evalbuat Traectories (Betch Size 1226)
4.0
Evotuction Falecteries (Batch Sior 2043)
4.-
Evalushion Tri ecctriie (Batrh Sipe 4386)
Sched.ler
Schrdlin Fire
Cosing
wsh
Schellll Prse Condhot+
Codle
3.8
Scheate free
3.6
3.4
3.24
22
0
2002
1003
0002
8002
0
1050
2050
1000
0020
5300
32g
1003
2003
.000
4000
Triting Steps
Training Steps
Thaining Steps
(e) Batch size 1024
(f) Batch size 2048
(g) Batch size 4096
Figure 9: Scheduler comparison for different batch sizes. All are with model size 151M.
Longer training requires higher EWA decay rate т. Throughout the paper, we adopt a learn- ing rate of 0.00316 across most experiments, but it is unclear whether that would be sub-optimal, especially since training with longer duration may require a lower learning rate as suggested in (DeepSeck-AI et al., 2024). Therefore, we justify our design decision on tuning EWA decay rate for simulating learning rate decay on different training durations by conducting the follow- ing experiments on series of 151M models with (a) batch size 256, 0.5 Chinchilla tokens; (b) batch size 256, 20x Chinchilla tokens; (c) batch size 2048, 20x Chinchilla tokens, all with learn- ing rate swept over {0.00316, 0.00158, 0.01264, 0.00632, 0.00075} and EWA decay rate over (0.99, 0.9968, 0.999, 0.99968, 0.9999}. We set the number of warmup steps to be 0.25 of the total steps for (a) and 0.05 for (b) and (c). The results in Figure 10 denote the validation loss at the end of the training, which consistently show that within each group of experiments, learning rate of 0.00316 we use throughout the paper is consistently the best. Moreover, when enlarging the training data size from 0.5x Chinchilla tokens to 20x, the optimal EWA decay rate value would increase as well. This is also justified in the results presented in Figure 8 and Table which indicate that longer training durations may benefit from higher EWA decay rate to improve optimization performance.
3.3984 3.3787 237800 3.2905 3.3998
3.65
3.0122
3.0333
3.0052
3.0164
3.0157
3.03
2.9934
2.9794
2.9765 2.9791
2.9911
2.995
3.80
2.090
EIMA Deray Rata
0.300 3
3.3800
3.36 33558
3.3726
3.3922
3.0050
3.0245
3.0015
3.0102
3.0094
3.02
2.9867
2.9729
2.9702
2.9029
2.9549
3.55
3.3757
3.3535 23521
3,3651
3.3749
3.30
2.5972
3.0159
3.0032
3.01
2,0812
2.9674
2.0648
2.0677
2.9795
3.45
3.00
2.9515
3.40
3.5058 34624 36361 3.6515 3.6271
2. 2.9851
1.00m8
2.9842 2.9905 2. 2,9414
2.99
2.9961
2.9763 2.9574 2.9771 7.9913
-2.665
0.00075 505156 3,00316 00032 200254
0.20755 601156 0.00315
6.0002
1.01264
0.00075 0.00555 o 0.00as 6.10522 a.1264 Learning Rate
Learning Rate
Learnnng Rate
(a) Batch Size 256, Size N.5CoNNN
20Co
(c) Batch Size
Size 20Co..!
Figure 10: Impact of learning rate and EWA decay rate across various training durations. We report the validation loss at the end of training for each hyper-parameter combination. N denotes the model size. The best loss is marked by * symbol. Longer training durations, seen in (b) and (c), necessitate a higher EWA decay rate for a given learning rate.
18

---

Published as a conference paper at ICLR 2025
Effects of weight decay. Though weight decay (WD) does not provide generalization bene- fits for pre-training, previous works have shown that it might improve convergence in language model training (Hoffmann et al., 2022a; Kosson al., 2023). We adopt the default decoupled weight decay (Loshchilov, 2017) implementation in PyTorch and sweep over weight decay rate (0.01, 0.0316,0.1 for LR 0.01, (0.0316, 0.1, 0.316) for LR 0.00316. We show in Figure 11 that for constant learning rate with EWA, while weight decay offers a slight performance improvement, has little effect the critical batch size, which remains our primary focus. Consequently, we disable weight decay throughout the paper.
B ADDITIONAL ABLATION STUDIES ON STUDYING ADAM OPTIMIZER
We employ Adam as the default optimizer for large-scale model training throughout the paper. In this section, focus two key hyper-parameters that sig-
nificantly affect optimization efficiency and examine their im- pact in detail.
Target Validation Loss 3.24
14
loo3 ( #Steps)
The effect of momentum 81 of Adam on CBS. We sweep over several momentum 81 values in Adam for all learning rates and batch sizes: [0, 0.8, 0.9, 0.95, 0.975J. Overall, Figure 12 shows that language model pre-training may need a large momentum value to be efficient and 3 0.95 is slightly better (<0.02 gain on eval loss) than 0.9 for batch sizes. We observe that small batch size regimes like 26, the performance gap between opti- mizing with and without momentum 31 is small while the gap increases as we double the batch size (Shallue et al., 2019). Moreover, we show that momentum 0.9 and 0.975 have simi- lar effects the number of steps needed to reach target val- idation loss and critical batch sizes. On the other hand, small momentum, especially
12
WD
w/o WD
2
2
2 2
Batch Size (2*)
Figure 11: Comparison of effi ciency with and without weight decay
momentum, would hurt the optimiza- tion. This aligns well with the extensively studied acceleration of momentum in SGD with momentum (Goh, 2017).
Target Validation Loss 3.3
Momentum
16
logo (#Steps)
15
0.8
0.9
0.95
0.975
13-
12
11
27
2
2 10
21
2
2/3
Batch Size (2*)
Figure 12:
Ablation results momentum. All data points trained using 151M models and total number of optimization steps to reach a fixed target loss is reported.
The effeec of the second moment decay rate 3g in Adam. As reported in Appendix Table 4, we found that B2 in Adam, the exponential decay rate of the second momentum estimate of gra- dients to smooth out the model update, also has significant effects on training for small batch sizes. This might be because gradients in small-batch training are sparser. Specifically, we ab- late B2 E [0.95, 0.99, 0.999] for all model sizes and batch sizes in [64, 128, 256, 512). We find that the default value 0.95 in previous works that are set for millions of tokens batch size training might be sub-optimal (Smith et al., 2022; Wortsman et al., 2023; Groeneveld et al., 2024). For large batch sizes [1024, 2048, 4096, 8192]. finding that it is worse than the default 0.95 we choose. When training a larger model with a longer duration (e.g. Chinchilla settings in Appendix Figure 8b), high enough B3 is necessary.
experiment with small 8r = 0.9 with the model size 151M,
19

---

Published as a conference paper at ICLR 2025
Takeaway on Adam optimizer:
Momentum & is important in improving training efficiency: a value of 0.95 consistently performs well across various model sizes and batch sizes. However, setting it too high (0.975) or too low (0.8) leads sub-optimal results.
Smaller 32 0.95 is helpful for large-batch training over short durations, while a large $2 0.99, 0.999 0.9995 is helpful for long-duration training and substanfially improves small batch size training (<262k tokens).
C RESULTS INCLUDING SMALL BATCH SIZES
For completeness, demonstrate linear scaling behavior in small-batch regimes across all model sizes (Figure 13). This shows that all models exhibit linear scaling (with reasonable deviations) with batch size ranging from 2 2 where doubling the batch size roughly halves the number of steps needed to reach a target validation loss, as determined by the optimal run with batch size of 256 at the Chinchilla step.
Chinchilla Compute-optings Target Validation Loss
log2 #Steps/# # Steps at BS 64)
Model Size (Target Loss)
85M (3,42)
151M (3.24)
302M (3.07)
604M (2.92)
1.2B (2.736)
-2
-3
26
27
2
29
210
Batch Size (2*)
Figure 13: Linear scaling regimes: doubling the batch size can halve the optimization steps to reach the target loss.
Moreover, include all the results that contain the smallest several batch sizes (Figure 14). Note that the denominator is the number of steps to reach target loss batch size 64 instead 256 now. Now observe clear linear scaling of all the model sizes till around 2 for model sizes, while the largest three model sizes maintain linear scaling till almost 2 There are minor differences with the main plot in Figure because of the difficulty of optimizing with very small batch sizes like 64 but it does not affect the conclusions and takeaways we would like to convey. Since our focus primarily large batch sizes, consistently use 29 as the starting batch size throughout the main text. Furthermore, recall that we set Bopt 2 when selecting the target loss. Figure 13 and Figure confirm that 28 falls within the linear scaling regime, which justifies our design choice.
20

---

Published as a conference paper at ICLR 2025
Chinchilla Compute-optimal Target Validation Loss
Steps/#Stepp #Steps at BS 64)
0
Model Size (Target Loss)
85M (3.42)
-1
151M (3.24)
302M (3.07)
-2
604M (2.92)
-3
1.2B (2.736)
log2 #Steps/
-4
-5
-
-7
26
27
2
2 Batch Size (2*)
29
20
2
2
2
2 24
Figure 14: Full results for different models in Chinchilla settings. We include both a largest batch size 214 start the plot from several small batch sizes 26 27 2 Relative number of steps w.r.f. batch size 26 is reported.
D ADDITIONAL DETAILS ON EXPERIMENTS
Optimizer setup. For optimizers, we try both SGD (Robbins Monro, 1951) and Adam (Kingma, 2014) and find that SGD without momentum is significantly worse we Adam only for all the experiments. We disable weight decay in Adam as we observe that it does not significantly affect critical batch size (Figure 11). For generality, though the training set C4 might contain low- quality or duplicated documents that can potentially lead to training instability (Muennighoff) et al., 2023; Wortsman al., 2023), observed that these issues did not affect our primary target of interest- nanrest-namely... the final optimization efficiency. As result, we didn't 't explicitly adopt additional normalization like QK normalization (Dehghani al., 2023; Zhai al., 2023) or a z-loss (Chowd- hery et al., 2022) to mitigate loss spikes." 2 We set e to be 1e-8 by default, and refer to momentum as 31 in Adam by default throughout the paper.
Chinchilla steps for batch size 256 that determine the target losses. For each model size, we aim establish target validation loss by training with a global batch size of 256 on Chinchilla- optimal amount of tokens. Given the context length of 512 used throughout, we can determine the number of training steps required based on the following Table We use a token-to-mode size ratio Cchin of approximately 20.34 to study the halving effects of doubling the batch size and to observe its impact on the critical batch size.
Table 1: Chinchilla steps for determining the target loss for each model size.
Model Size
85M
151M
302M
604M
1.2B
Chinchilla Step
13193
23438
46875
93750
187500
Evaluation data size and frequency. To ensure frequent model evaluation holdout C4 validation set, it is important to maintain a balance between reliability and efficiency. larger evaluation size can provide more stable and reliable performance metries, but it must also be efficient to maintain practicality in each run. Using 151M models, we evaluated the variance across different token counts: 2. 7e-4 for 327,680 tokens, 4.53e-5 for 1,638,400 tokens, and 7.65c-6 for 3,276,800 tokens. Based on these results, we have set the default number of evaluation batches to
We ohserve irregular loss spikes despite adopting gradient clipping, but most runs can still be optimized
well in the end.
21

---

Published as a conference paper at ICLR 2025
100. is important to note that the total number of training steps varies across different batch sizes. To address this, we implement a hybrid evaluation protocol: the model is evaluated at intervals of 2' (where i Z every 1,000 steps, and 0.7n, 0.75m, 0.8n..... up during the last 30% of the total steps n. This approach ensures more frequent evaluations toward the end of training, allowing for more accurate assessment of the total training steps needed achieve target evaluation loss.
Hyper-parameter search details. Due to compute constraints, we cannot perform an exhaustive search all hyper-parameter configurations. Instead, as suggested by the ablation studies in the main text, we gain insights into hyper-parameters' by training smaller proxy models (15IM param- eters). We optimize the following parameertss in sequence: learning rate, momentum (B1), warmup steps, scheduler, and context length. Additionally, we tune 32 and T for each model size and batch size. Specifically, for large batch sizes (> 1024), smaller 33 and larger 7 tend to more effective, while the opposite holds true for smaller batch sizes, aligning with findings in (Porian al., 2024; Zhang al., 2022).
Below show the hyper-parameter choices (Table 3) and optimal ones (Table 4) report in our main plot for studying CBS with respect to model sizes. Additionally, Table presents various model size configurations and scaling methods, with models in bold indicating those used in our controlled experiments.
Table 2: Model architecture details.
Model Size
Thheuds
Plaayers
danodel
Hidden size of MLPs
85M
12
12
768
3072
151M
16
12
1024
4096
302M
16
24
1024
4096
604M
16
12
2048
8192
1.2B
32
24
2048
8192
Table Sweeping experiments settings. Default values after the hyper-parameter search are in Bold font. Bold font means the default extensive tuning. Not bolding implies full sweep for each model scale. The values in parentheses were not used for every sweep: for the 151M models, we tested learning rates of 3. 16.4 and le-2, but found that with EWA, these performed worse than 3.16.. The EWA decay rate of 0.99995 only used for long 1.2B runs.
that can closely reproduce our results without
Hlyper-parameter Model Size
Values
85M, 15TM, 302M, 604M, 1.2B
Batch size
26 214
Learning rate
(3.16e-4), 1e-3, 3.16e-3, (le-2)
Learning rate scheduler
constant+EWA, cosine, WSD, schedule free 0.15, 0.25, 0.35
Warmup fraction
Momentum
0, 0.8, 0.9, 0.95, 0.975
Adam Br
0.95, 0.99, 0.995, 0.999, 0.9995 0.95, 0.98, 0.99, 0.995, 0.998, 0.999, 0.9995, (0.99995)
EWA decay rate
Context Length
512, 1024, 2048, 4096 1.0
Grad clipping norm
22

---

Published as a conference paper at ICLR 2025
Table 4: Optimal hyyper-paaameee for different model sizes. The optimal means the number of steps to reach target validation loss. We refer B2 as the exponential decay rate for the second- moment estimates in Adam, 7 as the interpolation parameters in EWA (5++1 5. & (1 7) -8))
All the optimal runs are trained with momentum 31 = 0.95 and learning rate 3.16e-3.
Batch Size
52
Batch Size
32
85M
15IM
0.999
0.9995
64
0.99
0.998
128
0.999
0.9995
128
0.99
0.998
256
0.999
0.9995
256
0.99
0.998
512
0.999
0.998
512
0.99
0.998
1024
0.95
0.99
1024
0.95
0.95
2048
0.95
0.99
2048
0.99
0.99
4096
0.95
0.99
4096
0.99
0.99
8192
0,95
0.98
8192
0.95
0.95
16384
0.95
0.98
16384
0.99
0.99
Batch Size
t
Batch Size
B2
T
Batch Size
8g
302M
604M
1,2B
0.999
0.9995
64
0,9995
0.9995
0,999
0.9995
128
0.999
0.9995
128
0.9995
0.9995
128
0.999
0.9995
256
0.995
0.9995
256
0.9995
0,9995
256
0.995
0.9995
512
0.99
0.9995
512
0.9995
0.9995
512
0.99
0.9995
1024
0.99
0.999
1024
0,999
0.999
1024
0.99
0.999
2048
0.95
0.998
2048
0.998
0.998
2048
0.95
0.998
4096
0.95
0.995
4096
0.995
0.995
4096
0.95
0.995
8192
0.95
0.99
8192
0.99
0.99
8192
0.95
0.99
16384
0.99
0,99
16384
0.99
0.995
16384
0.95
0.99
Table Model architectures of the ablation study on the scaling of Depth and Width. Only models highlighted in bold are used, as they are more comparable in terms of model size.
Model Size
71haads
Ilaaers
dmodel
MLP hidden
151M
16
12
1024
4096
302.09M
16
24
1024
4096
604.18M
16
48
1024
4096
1.208B
16
96
1024
4096
Model Size
Tiadds
Tllayers
d modee
MLPP hidden
151M
16
12
1024
4096
339.81M
24
12
1536
6144
604.08M
32
12
2048
8192
943.84M
40
12
2560
10240
E ADDITIONAL. DETAILS ON SCALING LAWS
We first present the fitted power law relationship between the number of optimization steps required to reach the target loss and the batch size (Table 6). All the results are obtained by solving the equation Section 3.2 via scipy. opt imize. fsolve using default hyper-parameters.
We report forecasting results for various model and token sizes, extending beyond the plots presented the main text (Table 7). For each row, increasing either model size token size shows that the forecasting results remain comparable.
Note that our definition of CBS and its scaling law have similar interpretation with the one in (McCandlish et al., 2018; Kaplan et al., 2020)
Emin. :Smin denotes the minimum possible number Swin
23

---

Published as a conference paper at ICLR 2025
Table Fitted scaling law parameters for Chinchilla settings when fixing a 1: log(Y) log(a $ where Y is the number of steps reach Chinchilla target loss, B denotes the batch size, and the critical batch size is solved as B b++3ax x.2x 1.p:
Bupt 256.
(a) Fixed a (default)
(b) Fitted a
Model Size
logg(B")
Model Size
a
Iog,(B*)
1293.83
2834258.00
9,54
1348.1
3386537.23
1.03
9.34
151M
1752.42
5677478.78
9,90
151M
1943.53
8259867.95
1.07
9.51
302M
2095.35
11383269.88
10.44
302M
2281.48
13977184.09
1.04
10.20
604M
2459.93
19449688.59
10.88
604M
2733.81
23738850.26
1.04
10.62
1.2B
3897.31
43381130.22 130.22
11.31
1.2B
3388.53
36748556.71
0.97
11.62
Table 7: Additional forecasted CBS results for larger scale. Recall that fit B" 93.20 x NO.47.
B 22.91 x D where model size N in millions and data size D is in billions.
Model Size
Forecasted CBS
1og,(1*
Token Size
Forecasted CBS
loga (B*)
L.5B
2862.17
11.48
30B
2833.31
11.47
2B
3274.93
11.68
40B
3240.99
11.66
2.5B
3635.65
11.83
50B
3597.20
11.81
3B
3959.69
11.95
60B
3917.12
11.94
3.5B
4256.09
12.06
70B
4209.70
12.04
4B
4530.72
12.15
80B
4480.76
12.13
4.5B
4787.63
12.23
90B
4734.29
12.21
5B
5029.77
12.30
100B
4973.22
12.28
5,5B
5259.34
12.36
110B
5199.73
12.34
6B
5478.06
12.42
120B
5415.52
12.40
of steps taken reach target loss and Emin is the minimum possible number of training examples processed to reach target loss. In particular, recall that critical batch size can be analytically derived 1.p This relationship reflects the point where batch size scaling incurs 20%
as B" overhead when the batch size doubled while (McCandlish et al., 2018). Here, the parameter the specific overhead
plays a role analogous to Emin. while a corresponds Smin. depending chosen to characterize the diminishing returns from increasing the batch size. We also note that the diminishing return overhead can vary, leading the following observations: 10% B* 20.67 x D0.48 20% B" 22.91 x D0.47 50%% B = 30.50 x D..44
F REPRODUCIBILITY
In our training environment, we verify that, across multiple model sizes (2.4M, 9.4M, 19M, 42M, 85M, 151M, 302M), (Wortsman et al., 2023). We use nodes equipped with 8 A100 GPUs, each with 80GiB mem- ory, for model training. We built our training framework using the Olmo training suite (Groeneveld et al., 2024).
(approximately) reproduce the final evaluation loss of Figure 1 in
G COMPLETE PROOFS IN SECTION 4.2
Proof of Theorem 3. The work by Zou et (2023) studied SGD with batch size for linear regres- sion and established matching (up constant factor) upper and lower bounds on the excess risk. Our theorem generalizes theirs by further considering the effect of batch size. Our analysis uses their intermediate results through appropriate reductions. We first define set of operations on PSD matrices as follows:
(
(+Σxx))]
I I=10L MB E
M H H.
BB - & 1
MB - -
24

---

Published as a conference paper at ICLR 2025
where I is index set of B independent data. Note that
(MB M) A Cov (1 E x A
Cov ( A
B
For Gaussian data x EN(O0, H), we have
Cov(xx A ExEN(O.H) XX T Axx
HAH 2tr(HA)H.
Together, we obtain
(M - A - tr(HA)H. B
Now we compute the error propagation along the SGD steps. Let 77 wt w" be the enror vector. For convenience, G. - B Liez, x,X; be the empirical covariance of an independent batch. Then we can define the bias and variance iterates
bias -I- (G)
bius
t= 1=....., 1, bias
w
and
T = (I - (
variance
miaiaace
variance
+
Σ
Eixi, t= 1=1,.... 1, variance 0,
B
VEI.
where Es = Vi x, w' ~ N(0,0²) We then compute the covariance matrices of these two error
iterates
B E(m. bias
M bian]
C
El variance
variance]
n7
Using the operators, these covariance matrices take the following iterative updates:
B no 69 no
B F.G, [[ GG, (I- = (I -T B
72
CP 0, C EG. [I - GG (I- [G]]
Σ
EXXA
Σ E:Xi
132
iET
iET,
2002
= -I- f-1 +
H.
B
where the last equation is because
E
Σ EE
Σ Exx; KE T
E
x
= BH....
IET
W,. First, using Lemmas B.3 and C.I in (Zou al., 2023), get the following bias-variance decomposition (note that our setting is well-specified):
Recall that w
E[R(w)] min R(-) = bias + variance,
where
(I->H)
bias : 2 (H, E(y bias 2 n biast)
(1--H)* k- ( B,,
2n 2
(I- k- HH C
variance
H. variance
n varlance
(a- k- H. c
2n
1=0 k=t
25

---

Published as a conference paper at ICLR 2025
where
n-
R-1
blas
variance
variance
72
The remaining efforts are to characterize B, and C, for a batch size B. For the bias part, we have
B (I - B) B
(I - (I- B,, + (MB - -M) B
272
(I - 7 B +
ir (HB H. t= 1,
B
For the variance part, have
2002
C = (I -I- C C-1 H
1202 (I (I- o c + (MB - -M) o C.-1
B
2 B
7002
-I- c + tr (HC H
H.
t= 1..,,,--11
B
To obtain an upper bound on excess risk, we replace a in Assumption 2.2 of Zou et al. (2023) with
2/ B, the number of steps with n D/B, and the noise level o² with o² /B, then apply the proof of Theorem 2.1. Similarly, for lower bound on excess risk, we replace 3 in Assumption 2.4 with 2/B. the number of steps with n == T/B, B, and the noise level σ² with σ² o²/B, and apply the proof of Theorem 2.2. By doing the above, we obtain the following matching up to constant factors upper and lower bounds on the excess risk for mini-batch SGD:
1 ||wo w011 H D.r
ER(() (w) min R(-)
+l/wo- w 1
107
1/ B
[Wo w IP.... + #n/lwo w (' kt + (17)²²
(an))EEs
Ei>0 X
+
my
2
ha + (m7)²² Eisi X?
+
B
where A* max(k: (k: A > 1/(n7)), and sufficient stepsize condition (see Lemma 4.1, Theorems 2.1 and 2.2 in Zou et al. (2023)) is
B (r)
0< min (tr(H)
min =
H|
H
The assumption ||
11 S implies
|w 1' + ny|lwo
w"
11 He < ||
w <
my
which further simplifies the excess risk bounds to
ER(w) (w) min R(-)= (1 Iwo
σ² k2 + (ny)²
Eisk X?
w
+
l|wo
Finally, replacing D/B in the bounds completes our proof.
26

---

Published as a conference paper at ICLR 2025
Proof of Corollary 2. By X ~ i" can solve for k* to obtain k*
(Dy/B) We then
calculate the expected excess risk by Theorem using the capacity and source conditions:
B
2
ht +
(Dr/B))E X
BR(w)- - 0² R
Dy
H
D
B
D7
b+2a
Σ
k* +
20
Dy
D
B
B
max
[1k*) 1} + (1*)1~b
+
1
k² +
(Dy)
(1+))
Dry
D
Dr
(1-b)/a
D7
D7
1/a
max
+ D
B
B
B
We then discuss three cases.
1. When b a, we have
DDy B
(1-b)/a
D
1/a
Dr
(1-b3/a
ER(w)
+
D
B
B
where the last equality is because 7/B B S 1 so the first term dominates the second term. So the optimal hyper- asper-parameters are 7" and B" 1.
When a < < 2a + 1, we have
D7
(1-b//a
1
Dy
1/a
ER(w)= 0²
B
B
so the optimal hyper-parameters are
01 1SB SDD B/B Do/b-1
3. When b b> 2a + 1, we have
D2 B
-2
D
1/a
ER(w) - ~
so the optimal hyper-parametee are
0 0<>"$1, 1s //BB= ~ Da/2241)-1 Combining the second and third cases completes the proof.
27


Reviews:
[{"review_id": "9o8TV7sjg0", "reviewer": "Reviewer_FgiN", "summary": "The paper investigates the scaling behavior of critical batch size in the pre-training of autoregressive language models.\nThey first define the critical batch size (CBS) as the point where increasing the batch size no longer leads to significant gains in computational efficiency (>20% overhead when doubling the batch size vs. a linear scaling). They then perform experiments to determine the CBS of autoregressive Transformer-based language models of varying scales, finding that CBS scales primarily with the size of the training dataset rather than model size. They provide theoretical support for this finding by studying infinite-width limits of neural networks and infinite-dimensional least squares regression problems.", "questions": "- A central takeaway of the paper is that \"CBS remains invariant when scaling up N [model size]\" (line 93). However, isn't this partially contrasted by the results in Section 2.4, where you write \"Figure 4 shows that increasing depth and width [and thus model size] can both slightly increase the CBS\" (line 264)?\n- I found the plots to sometimes be a bit confusing or less accessible. Some suggestions:\n  - Could you add the linear scaling line as well as the 20% region?\n  - Could you highlight the critical batch sizes, e.g. by a star marker?\n  - Nit: The font sizes are inconsistent between subplots, e.g. the x-axis labels between Figure 1a and the subplots of Figure 1b.\n  - In the legend of Figure 1b (left), what is the number in parenthesis? Is it the (relative) number of training samples?\n  - The lines are sometimes hard to read. E.g. in Figure 2a, the Constant+EWA line is very bright. Similarly the lines in Figure 2b. I also had trouble distinguishing the shades of blue, e.g. in Figure 1. But it does look pretty :)\n  - Nit: The white spaces in Figure 4 are a bit weird. There is only a small space between the subcaption of subplot (a) and much more white space between the subcaption (b) and its subplot.\n- In line 110 you mention that \"traditional learning rate decay strategies typically require predefining the total training duration\" and you mention Defazio et al., (2024). However, isn't Schedule-Free a counter-example? Schedule-Free is also using a, at least to me, a relatively similar approach to your \"constant+EWA\" strategy, no (e.g. comparing the equation in line 158 to Eq. (5) by Defazio)? Could you elaborate on how your method differs from Schedule-Free?\n- For most of the paper, you sweep the batch size between $2^6$ and $2^{13}$ with the exception of Figure 3, where the batch size is between $2^{16}$ and $2^{22}$. Is this because in this figure, you count the number of tokens, so practically batch size * tokens per batch?\n- You provide a scaling law fit of the CBS. Did you try checking how close your thus predicted CBS tracks the true CBS for \"unseen data points\", i.e. test its predictive power?\n- In the concluding remarks you write \"Our findings contribute insights into [...] particularly highlighting the role of hyper-parameters such as learning rate scheduling and optimizer settings\" (line 530). What \"optimizer settings\" are you referring to?\n- Nits:\n  - Line 120: I believe the \"Refer\" is superfluous.\n  - Line 157: There is probably a closing parenthesis missing after the second citation.\n  - Line 185: Typo for \"maximum\".\n  - Figure 3a doesn't need to be a subplot as there is no other subplot.\n  - Line 260: This sentence seems grammatically weird to me \"As the main result in Figure 1 only involves a single way for scaling up models [...]\". I might also have misunderstood the sentence. For example, doesn't scaling from 151M to 302M double the depth, and therefore Figure 1 involves multiple ways of scaling up models?\n  - Line 271: \"Compute-optimal\" should probably be lowercase.\n  - Line 288: Typo with \"achieved\" and \"batch size\".\n  - Line 351: \"chinchilla\" should probably be capitalized.\n  - Line 366: \"They\" -> \"Then\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides an interesting finding that the critical batch size scales mostly with data set size, and is largely invariant to model size. This is a relevant and, to my knowledge, novel insight.\n- The paper considers models ranging from 85 million to 1.2 billion parameters and thus covers a reasonably large domain of models.\n- I really liked the highlighted practical takeaway blocks throughout the paper, which made it easy to understand, well-structured, and accessible.", "weaknesses": "Some of the takeaways seem to me a bit too bold or not backed by enough evidence for the given claim.\n\n- For example, in Section 2.2, they compare the efficiency of learning rate schedules across batch sizes by comparing the number of steps to achieve a given target validation loss. They conclude that \"EWA consistently improves model training efficiency. [...] while outperforming Cosine for large batch sizes [...] and even with appropriate learning rate decay, [Cosine] underperforms our constant+EWA strategy in large-batch settings.\" (line 188). However, looking at Figure 2b), we can see that the training duration of cosine was chosen inefficiently. It continues to decay well beyond the target validation loss, achieving the target at roughly 50% of its decay schedule. It is also clear to see for the WSD schedule, which hits the target loss shortly after starting its decay, i.e. it is still halfway in its decay phase when crossing the threshold loss. I believe making a statement like \"EWA consistently improves model training efficiency\" requires a more rigorous empirical analysis. For instance, Schedule-Free [1] suggests a similar running average strategy to improve efficiency, providing a much more comprehensive and rigorous analysis.\n- Similarly, Section 2.3 claims that different context lengths have similar CBS, based on a single (comparably small) model & dataset. It seems that the lines start to diverge a bit more for larger batch sizes, so is it possible that this is more pronounced for other settings (e.g. larger models, more training data, etc.)?\n- All claims are made for the definition of CBS using 20% overhead. The authors state that \"20% can be replaced by any other suitable measure of increase from linear scaling\". I wonder how the results and conclusions would change if one varies this parameter, e.g. to 10% or 50%. For example, what would a plot look like of the CBS per model or data size as a function of the overhead?\n\n[1] Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky; \"The Road Less Scheduled\"; arXiv 2024; <https://arxiv.org/abs/2405.15682>", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H3gk8kKOiL", "reviewer": "Reviewer_76qC", "summary": "The paper studies how the critical batch size (CBS) - the threshold beyond which increasing batch size causes diminishing return - scales with the model and data size for pre-training language models. The authors conduct careful experiments on models ranging from 82 million to 1.2 billion parameters and show that the CBS slightly scales with data size rather than the model size (remains invariant to model size). The authors also provide theoretical justifications for this phenomenon using infinite-dimensional least squares regression and the infinite-width limits of neural networks. These findings have important implications for designing efficient pre-training strategies.", "questions": "- (Minor) It would be nice to use $\\times$ instead of * (asterisk) in line 135.\n- (Minor) The color contrast in Figure 9 can be improved.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written with clear organization. This work would be a valuable contribution to the ICLR community. I especially appreciated the “key takeaways summary” after each section.\n- The experimental design is rigorous; for example, decoupling various hyperparameters makes the claims more convincing.\n- Detailed experimental procedures are provided in Appendix D.\n- The formalization of CBS (beyond [1] “An empirical model of large-batch training”) would be helpful in the literature. The key findings that the CBS scales slightly with data size (and stays invariant to model size) are interesting. These insights have important implications for efficient pre-training strategies.", "weaknesses": "- The experimental scope is limited to models up to 1.2B parameters trained on C4, which may not fully capture scaling behaviors at larger scales (e.g., models with over 50B parameters). On a similar note, key ablation studies are primarily conducted on smaller models (with C4). However, given the careful experimental design and clear theoretical analysis, I do not believe that these impact the validity of the findings.\n- It would be helpful to have a dedicated section discussing the limitations of both the theoretical analysis and empirical findings (e.g., Gaussian data distribution).\n- However, I believe that the work has substantial limitations that would prevent me from recommending acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qxeIXAyCFA", "reviewer": "Reviewer_ivw4", "summary": "The paper revisits the role of batch size in scaling language models during pre-training. The core finding is that the critical batch size needs to scale alongside the model. However, both the critical batch size (CBS) and model scale (specifically, model width) show diminishing returns when scaled under a given training budget. The authors demonstrate and conclude that CBS scaling is dependent on data scale but invariant to model scale. Various hyperparameter ablations provide additional insights into relative pre-training performance efficiency.", "questions": "1. Could the paper have bit more insight into why certain model scales were picked for certain experiments, since it is not always the smallest (151M) model that was selected?\n2. Any reason explaining the worsening _efficiency_ for a 1.2B model on larger batch sizes in Figure 1.b) (right)?\n3. For a study with even lesser bias or confounding factors, would it make sense to make conclusions on the critical batch size (CBS) or its scaling exponents without an EWA of model weights? \n    * In a similar vein, can we see in Figure 2 what only constant LR schedules look like?\n    * Is cosine here with or without warmups?\n4. Why does Figure 2 (right) have different step lengths for the same batch size assuming similar compute budgets overall, given same model size? \n    * Similar to above, why does Figure 9 all have different lengths?\n5. Why do we expect the rankings of different schedulers found on 151M scale would transfer to larger scales? Do we have a literature or empirical reference for it?\n6. Is the only notion of _efficiency_ in paper as denoted by the overhead ratio metric as defined by the measure on batch size 256/512?\n    * For every result seen as plot/table, does it mean there was an equivalent run made on a batch size 256/512?\n7. Would Figure 3 result hold for larger models with bigger context sizes? (I am not asking for an experiment at such scales but a cheaper experiment that could serve as proof-by-contradiction)\n    * How does the model size change when Figure 3 writes about a model of size 151M but with context size ranging from 512-4096?\n    * Are model sizes calculated without the embedding parameters?\n8. Could the 20% and the resulting 5 in Section 3.2 be more generalized in the formalism?\n    * Could the readers have a clearer reference/intuition for the 20% overhead?\n9. Could L318-319 have more support?\n10. Does Figure 5 have empirical backing regarding the predicted CBS and the formalism of staying within the 20% margin of steps for similar loss?\n    * Or if the decoupled prediction for batch sizes hold _reliably_ for different scales or ablations such as architecture, datasets, hyperparameters?\n11. How exactly are the tuning decisions undertaken especially in Section 3.3 (example, L371-372)?\n12. In Theorem 1, is it expected or meant to be $w_2 > w_1$? If so, then there is a typo and makes a significant difference to the interpretation.\n13. What is the significance of _t_ in Theorem 1? Should we compare losses under similar tokens for two vastly different model scales, especially following Chinchilla?\n14. What is the specification or how is $\\mathbf{H}$ defined in Section 4.2?\n15. Is there a more intuitive summary for the proofs for Theorem 2 and 3 especially for readers (like me), not familiar with Zou et al. (2023)?\n16. How exactly does the _hybrid evaluation_ (L968) provide _more accurate_ information?\n17. What or how are the hyperparameters (HPs) for parameter sweeps ranked or ordered (L977-979)? Does a different ordering yield potentially different results?\n    * What are the default values for other hyperparameters when tuning each of these in order?\n    * Assuming in this order that LR is considered to be the most important HP, why tune the least important HPs (in $\\beta_2$ and $\\lambda$) for each model scale and not LR (Table 4)?\n18. I may have missed it but could there be a reference for how the power law model fit is made for the values in the caption of Figure 7?\n\n\nOverall, the paper reads nice and does touch on an important and often under-studied aspect in scaling literature.\nHowever, the paper writing and presentation raises some eyebrows with really big plots, white spaces used, and a theory which feels like a sub-paper than something that brings the story together.\nGiven the expensive space of empirical experiments, the authors did well to consider multiple design choices and study them. \nUnfortunately, that also opens up more questions on how these designs were arrived at and that clarity would be my primary criticism.\nOne other thing would be the lack of _a_ clear takeaway in terms of _how_ a pre-training practitioner could use the insights from this paper when scaling (model or data).\n\n\nScores may be increased depending on suitable responses to most of the points raised above.\nThank you for the paper, it was a nice read overall.\n\n\nPS: I was unaware of Zou et al. 2023 and have not verified the proof on Pages 22-24.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and largely easy to follow.\n2. The related literature covers the important papers in the topic well.\n3. Formalizations and hypotheses are clearly outlined and help understand results better.\n4. Though personally I would like to reconsider its exact design and placement, the _Takeaway_ block was helpful while reading the paper first time.\n5. The model scales reported in experiments are adequate in applying the insights to large-scale pre-training.\n6. Formalizing the notion of critical batch size (CBS) through the 20% overhead assumption is novel and seemingly useful.\n7. Considering various hyperparameters ablated across selected model scales is important and welcome in a scaling-related paper.", "weaknesses": "1. Some lower scale experiment with repetitions over different seeds to show the robustness of the findings (laws, exponents) and insights (data dependence and model scale invariance).\n2. The work is mostly a benchmarking study with main contribution relying on the hypothesis constructed and how the experiment for it is setup, which therefore leaves more room for explaining some of the design choices, especially with model scale, hyperparameters (see, Questions below for examples).\n3. Section 3.3 mentions how different hyperparameters were adjusted _to achieve optimal performance_ (at 302M), however, the grid search (for ablations) were done on a different model scale (151M), and thus is unclear what were the different HP settings considered in fitting points for the power law.\n4. Theorem 1's proof by existence could do with more details and support, i.e., it is unclear how _training iterations t_ related to models of different width (thereby, scales), possibly different batch size scales and naturally different compute budgets under Chinchilla-optimality (see, Questions).\n5. Despite leveraging theoretical insights from muP, there is only a limited grid search over learning rate values over the smallest available scale while that is apparently used as-is for even the larger model trainings, i.e., hyperparameters are not suitably scaled (this is unclear overall).\n6. The outcome of the theoretical proofs and empirical results appear independent and do not knit the contributions as well.\n    * Is there a way to connect the toy regression experiment with the scaling law exponent fits?\n\n\nMinor issues/fixes:\n1. Discrepancy of batch size 256 or 512 as the batch size for calculating overhead ratios\n    * L52, Fig. 1,3,4,6,8 y-axes, L370, L888.\n2. Should likely be `b > a` in L1299.\n3. Typo on L366: perhaps should be \"They\" -> \"Then\"\n4. Discrepancy in Fig. 5 caption where $C_{chin}$ comes to be around 6 as per the equation given for $B^*$.\n5. Figure 4.a doesn't mention which model width or depth.\n6. Figure 10 mentions _context length_ but has nothing to show for it.\n\n\n\nNitpicks:\n1. Irregular and inconsistent plot sizes.\n2. Recommended to refer sections/tables/figures in the Contributions list in Section 1.\n3. Vertical lines in Table 4 separating different model sizes.\n4. Parantheses in L1119-1122 would enable easier parsing for readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BgCM3OdXfg", "reviewer": "Reviewer_SERg", "summary": "The authors reassess the notion of critical batch size (CBS) in the context of language model training, and investigate the impact of scale on CBS. They find that the CBS does not scale with model size in majority, but with data size, and highlight roles of optimizer choices, where they find exponential averaging with a constant LR to be performant, matching or outperforming cosine or WSD.", "questions": "I mention the main points in the section above, but repeat them here as questions (with some others):\n\n- Why did the authors choose to disable weight decay?\n- How did you use EWA in detail, or is there a specific implementation trick that was required?\n- Did you tune SFO with the same sweeps as Adam?\n- Why are the warmup steps chosen between 0.15, 0.25 and 0.35? These are quite large fractions, whereas in practice the warmup is often a very minor part of training (e.g. less than 5%).\n\nI would be very happy to engage in a discussion about these points. More broadly, I hope the authors see my comments as not only critical but encouraging.\n\nAnother minor comment: I think the choice of colors for the plots is not ideal, as the contrast is very low and it makes it hard to read and compare lines (especially for people with color vision deficiency).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "First of all, I thank the authors for looking into this problem — I believe the question of CBS is both extremely relevant and understudied. The paper is written in a clear manner, the experiments are very extensive, and the authors additionally provide theoretical studies. The work could be a starting point for further studies, e.g. going beyond Chinchilla optimal points.", "weaknesses": "While I very much appreciate the topic and investigation of the paper, I unfortunately have to be very critical of the experimental evaluation. \n\n- First, the authors note in the Appendix that they disable weight decay of AdamW, without justification; this means effectively only studying Adam and not AdamW. However, weight decay is not only a major part of modern large scale training (see e.g. Chinchilla Fig A7 https://arxiv.org/pdf/2203.15556 or Wortsman et al. (2024) https://openreview.net/pdf?id=d8w0pmvXbZ), but it also strongly changes the training trajectory, where disabling it leads to a practical decrease of the effective learning rate even for a constant schedule; see e.g. Kosson et al. https://arxiv.org/pdf/2305.17212. Therefore, it is unclear how the findings would generalize to actual practical settings (proper baselines), and assuming the same results is misleading.\n- The authors find a constant LR + EWA to be extremely performant. I have tried to replicate their results with the same hyperparameters (including disabled WD) in the same setting of LLM pretraining and fail to obtain the same results. While EWA can give a slight boost, it is far from the cooldown loss; even more, using e.g. a decay of 0.99 or higher results in a similar curve to the one of 0.9999 in Fig. 8 (increase in loss, but much earlier in training) and then EWA is *much worse* than the original model. Since the authors have not provided code, it is unclear to me where this discrepancy is coming from — I repeat this point as a question below.\n- For most experiments, the schedule free optimizer performs surprisingly bad; if well tuned, I would think it should be at least as good or below the stable phase of WSD from my experience.\n\nThe combination of these points makes me concerned of some of the main empirical takeaways of the paper. To be clear, I believe the paper has much merit, and I do not want to reduce it to these concerns. However, I think addressing them requires either rerunning most experiments or major rewrites.\n\n-- update:\n\nThe authors have engaged actively in the rebuttal and addressed the main concerns that I've raised. Expecting these points (weight-decay, fixed baselines, annealing to zero) to be well discussed in the camera-ready, and the findings of constant+EWA to be contextualized for the scope of the study, I have raised my score and look forward to the final version of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lrgh3LHQnw", "reviewer": "Reviewer_A37b", "summary": "The paper studies how a compute-optimal batch size, the critical batch\nsize (CBS) is influenced by scaling the amount of data or model\nparameters in a Transformer-based language modeling objective. Most\nimportantly, the authors show a correlation between the training\nduration and CBS, both empirically and theoretically. Importantly, the\nexperiments also show that model size does not have an effect on CBS.\nFinally, a secondary but interesting part of the experiments concerns\nallowing to change the training duration mid-training. It is shown\nthat a simple constant schedule combined with exponential weight\naveraging at the end of training outperforms both \"standard\" cosine\ndecay and the recently proposed warmup-stable-decay (WSD) schedule.\n\nThe empirical study also includes a controlled setting to study the\neffect of width- vs. depth-scaling on the CBS (the finding being that\nboth scaling methods influence the CBS similarly). Finally, the\nauthors additionally underline their empirical findings using\nNTK/tensor programs and Gaussian linear regression theory.", "questions": "### Please address\n\nIn Proof of Theorem 1 (page 8, line 418-419), it is stated that \"the\ntrajectory of the network approaches a limit as width tends to ∞\". I\nwould wish for what exactly is meant by \"trajectory\" to be defined\nmore clearly. E.g., is it the trajectory of the weights during\ntraining? Also, please cite a source for this statement.\n\nAppendix D, Evaluation data size and frequency (page 18, line 964ff):  \nWhich which model after what training have these evaluation variance\nnumbers been obtained?\n\n### Minor comments\n\nPage 3, line 158:  \nPlease describe the meaning of θ.\n\nPage 8, line 416:  \n$R(M, t)$ denotes the loss of network $N$ at time $t$.  \n-> $R(M, t)$ denotes the loss of network $M$ at time $t$.\n\nIn the experiments, ideally, the optimizer's ε hyperparameter would\nalso have been scaled with the model size\n(https://arxiv.org/abs/2309.14322).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "For their experiments, the authors used C4, a de-facto standard NLP\nresearch dataset, which makes interpretation and reproduction of the\npaper of the paper easier.\n\nA particularly strong section is that of relevant work, in which\nvarious important related papers are summarized well.\n\nI also really enjoy the \"takeaway\" boxes as section summaries.", "weaknesses": "The paper heavily relies on Chinchilla scaling theory and\nTransformer-decoder training with the language modeling objective. It\ndoes not account for settings which use fewer training steps than\nChinchilla-optimal training and does not study other data sets or\nmodels. This means it is difficult to extrapolate the results to the\ngeneral domain of deep learning.\n\nI would personally have preferred use of a more standard \"dynamic\nhorizon\" schedule such as WSD for the experiments in order to improve\ninterpretability/comparability with regard to other research. That\nsaid, the finding regarding constant+EWA quality is important.\n\nWeight decay was disabled for the experiments, which is hardly done in\npractice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)