How Does Critical Batch Size Scale in Pre-training?

## Abstract

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. 1 1 1 Code available at https://github.com/hlzhang109/critical-batch-size .

### 1 Introduction

Efficient optimization is critical in pre-training large models (LMs) at scale (McCandlish et al., ; Shoeybi et al., ; Kaplan et al., ).
In particular, large-batch training is key to accelerating training, as it enables more efficient parallelism across hardware accelerators (You et al., ; Goyal et al., ).
Specifically, understanding the scaling behavior of the critical batch size (CBS) is essential for optimizing data parallelism, as it defines the point beyond which increasing the batch size may result in computational efficiency degradation.
Below the CBS, approximately linear scaling 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 to diminishing returns, making it essential to balance computational efficiency with model performance (Shallue et al., ; McCandlish et al., ).
This trade-off presents a challenge for studying pre-training given resource constraints as practitioners are compelled to 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, ; Radford et al., ).
Analyzing CBS in pre-training is challenging due to the absence of a precise formalism relating it to model and data sizes in the literature (Kaplan et al., ; McCandlish et al., ). Moreover, the interwined effects of scaling model and data sizes proportionally (Hoffmann et al., ) further complicate this analysis.
Although previous works study the effects of batch size on optimization performance (DeepSeek-AI et al., ; Besiroglu et al., ; Porian et al., ),
two crucial differences are (1) they do not decouple model size and data size; (2) they focus on 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 a certain target validation loss, 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.t. model size and data size to systematically study the scaling of CBS.

Figure: 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-model size ratio $D/N=C_{\text{Chin}}$ 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).
Refer to caption: x1.png

#### 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 in tandem with model size, as suggested in the Chinchilla compute-optimal framework (Hoffmann et al., ). 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 can 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).

#### 1.2 Theoretical Implications

Theoretically, maximal update parameterization suggests that, beyond a 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, we 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 1 (Informal version of Theorem 2 ) .

In infinite width regimes (Yang and Hu, ), training dynamics and performance of the networks become effectively independent of the model size. Consequently, the critical batch size remains nearly 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)=\Theta(D^{c})$, where the exponent $c\geq 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 $c>0$, indicating CBS grows with data size.

###### Theorem 1 (Informal version of Theorem 2 ) .

###### Corollary 1 (Informal version of Corollary 2 ) .

### 2 Experimental Design and Empirical Findings

We first 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., ) using Adam (Kingma, ) with optimizer-specific hyper-parameters reported in Table 3. We adopt the tokenizer of Eleuther AI’s gpt-neox-20b that has a 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 $B_{\text{opt}}$ that is in the linear scaling regime that incurs no efficiency overhead as detailed in Appendix C and model architectures in Table 2.
We consider the validation loss of an optimal batch size $B_{\text{opt}}=256$ at step $t_{\textrm{Chin}}=C_{\textrm{Chin}}*N/(\textrm{ctx_len}*B)$ for each model size $N$ and batch size $B$ with context length ctx_len set to be 512.

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 hyper-parameters but also the unknown training dynamics of each model size and batch size.
Below, we outline several key aspects to approach the goal:

- 1.
As we focus on the number of training steps needed to achieve a target validation loss, learning rate decay strategies typically require predefining the total training duration (Loshchilov and Hutter, ; Hu et al., ; Hägele et al., ; Defazio et al., ). To address this, we propose using exponential weight averaging (EWA) (Polyak and Juditsky, ) to achieve the desired target validation loss, a 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).
- 2.
Training with proper hyper-parameters: ensuring proper sweeps of momentum and learning rate (Appendix A); adopting well-tuned values for the $\beta_{2}$ parameter and the exponential weight averaging decay rate $\tau$, tailored to each batch size (Appendix B).

#### 2.2 Training Beyond Fixed Durations for Reaching Target Validation Loss

Figure: Figure 2: Comparing and accounting for training dynamics. Throughout, we adopt Constant+EWA since it performs the best for large batch sizes and avoids setting a fixed training duration beforehand for reaching a target loss.
Refer to caption: x2.png

Benchmarking learning rate schedulers.
In practice, LMs are usually trained using fixed token budgets (Hoffmann et al., ), 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 a lower learning rate is kept at the end of training to enable better optimization.
However, our goal is to find the optimal performing run under various hyper-parameters and optimization conditions.
This implies a non-trivial decision regarding selecting the maximum training duration (Defazio et al., ; Hägele et al., ).
As training beyond fixed durations is particularly favorable in many large-scale pre-training scenarios, we benchmark recently proposed methods like schedule-free optimizer (Defazio et al., ), cosine, warmup-stable-decay schedule (WSD) (Hu et al., ) (or trapezoidal (Zhai et al., )) and our proposed constant+EWA strategy which maintains a running average of model weights $\xi_{t+1}=\tau\cdot\xi_{t}+(1-\tau)\cdot\theta_{t}$ to improve optimization and use $\xi$ for evaluations.
We show that our constant+EWA strategy can match the efficiency of cosine scheduling and WSD, especially for large batch sizes (Figure 2).

Prior research has shown the generalization (Izmailov et al., ) and optimization (Karras et al., ) 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 is useful in scenarios where a target loss must be achieved, but practitioners are uncertain of the exact maximum data size to set up the learning rate schedule for training.

#### 2.3 Ablation on Model Context Length

We adopt 512 as the context length of LMs for all of our experiments but it is unclear how it 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^{10},2^{11},2^{12}$ (Figure 3)
Overall all models in four different context lengths have very similar relative optimization efficiency across various batch sizes and thus justifies our use of 512 for all the experiments.

#### 2.4 Ablation on Model Width and Depth

Figure: Figure 3: Ablation results on context length using 151M models.
Refer to caption: x3.png

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 1 only involves a single way for scaling up models (Table 2), e.g. 604M model has $2\times$ 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 $4\times$ instead (detailed configurations in Table 5).

Figure: Figure 4: Scaling up width and depth are equivalent in efficiency gain for compute-optimal training.
Refer to caption: x4.png

Firstly, as shown in Figure 4 (left, middle), under Chinchilla settings where data and model size are scaled proportionally, increasing either model depth or width leads to a similar increase in the CBS.
Then through controlled comparison (Figure 4, right), we see that using two different ways to scale 151M models to 604M ones is equivalent in efficiency since both curves overlap.

Our findings may offer practical insights for scaling models under a fixed token budget that is allocated in proportion to model size. This is particularly relevant because scaling model width is often favored over increasing depth, as wider models tend to be more amenable to parallelization without incurring additional latency overhead (Shoeybi et al., ; Touvron et al., ; Erdil, ).

### 3 Critical Batch Size Scaling Law

Figure: Figure 5: Illustration of critical batch size, where $B^{*}=2^{11.87}$ and context length is 512 by default.
Refer to caption: x5.png

#### 3.1 Formal Definition of Critical Batch Size

Recall that CBS is the transition point where increasing the batch size by a factor of $k$, leads to a 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 a $20\%$ overhead compared to linear scaling. First of all, define $\mathcal{R}(N,D,B)$ as the best loss achievable for a model of size $N$ using a single pass on $D$ tokens with a 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 $\mathcal{R}_{\text{opt}}(N,D)=\min_{B}\\
\mathcal{R}(N,D,B),B_{\text{opt}}(N,D)=\argmin_{B}\mathcal{R}(N,D,B),$
as the minimal loss achieved optimizing over batch size and the optimal batch size respectively.
We define $f_{N,D}(B)$ to be the number of steps required to reach $\mathcal{R}_{\text{opt}}(N,D)$ as a function of batch size $B$. Clearly $f_{N,D}(B_{\text{opt}})=D/B_{\text{opt}}$. To define the Critical Batch Size, $B^{*}(N,D)$, we can define a linear scaling curve $f^{*}_{N,D}(B)=D/B$. $f^{*}$ matches $f$ at $B_{\text{opt}}$ and then scales down linearly as batch size goes up.
$B^{*}(N,D)$ is defined as the maximum batch size $B^{\prime}>B_{\text{opt}}(N,D)$ such that
$f_{N,D}(B^{\prime})\leq 1.2f^{*}_{N,D}(B^{\prime}).$

As illustrated in Figure 5, $B^{*}(N,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.

###### Definition 1 .

#### 3.2 Scaling Laws w.r.t. Model Size for Chinchilla-optimal Pre-training

As observed in all the results above, doubling the batch size for larger models allows them to more efficiently reduce the relative number of steps needed to reach the target loss.
We ask whether those increased efficiencies can be predicted via a scaling law.

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 $\log(Y)=\log(a+\frac{b}{B^{\alpha}})$ and then derive the critical batch size. Then we derive the CBS via $B^{*}=\frac{b}{5a}+1.2B_{\text{opt}}$, 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: $D_{\text{total}}=(a+b/B_{\text{opt}}^{\alpha})*1.2B_{\text{opt}}=(a+b/B^{
\alpha})*B$
, where $\alpha=1$, $B_{\text{opt}}$ is set to be 256 chosen to lie 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 $\alpha=1$ solution, as both strategies yield nearly identical forecasting results.

Secondly, we fit a power law $\log(B^{*})=\log(c+\frac{d}{N^{\beta}})$ with respect to the model size $N$ (in million). The constant term $c$ is set to be 0 0 by default (as $B^{*}$ should be 0 0 at $N=0$), which leads to $B^{*}=93.20*N^{0.47}$.
We visualize the curve fitted in Figure 1 (left) and report more forecasts in Table 7 in the Appendix.

Overall, we observe an increase in CBS when scaling up in compute-optimal training: In Figure 1 (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*N^{0.47}$.
Our results suggest that a critical batch size around $2^{9}$ to $2^{11}$ would be helpful to efficiently optimize models below 1B 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 is 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.r.t. Data Size and Model Size

Controlled comparison with the same data size.
Firstly, we use the Chinchilla token size 3.072B of 151M models $t_{\textrm{Chin}}$ 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 1 (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 a scaling law with respect to model size thereafter Figure 1 (bottom right): keeping the data size fixed leads to a scaling law $B^{*}=621.341*N^{0.087}$ weakly dependent on 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.28\times$, $0.5\times$, $2\times$, and $4\times$ the Chinchilla step for batch size 256 runs. This setup results in two under-training and two over-training configurations. To achieve optimal performance in the over-training scenarios, we increase the warm-up ratio accordingly, while for the under-training cases, we 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 that as we enlarge the number of tokens being trained on, we see an increase of CBS, similar to what 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 an 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.

Figure: Figure 6: Controlled comparison by training 302M models on 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 a fixed token count of 3.072B, we measure the target loss at that step for each model size.
Refer to caption: x6.png

### 4 Theory on Scaling of Critical Batch Size

Our experimental results show that CBS increases with larger data sizes but remains (nearly) invariant 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 and Hu, ; Bordelon and Pehlevan, ). 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$ (or for gradient descent, i.e., $B\rightarrow\infty$), training iterations $t$, an error tolerance $\epsilon>0$, fixed learning rate schedule and data ordering, for any network and initialization satisfying Master Theorem (Theorem G.4) in Yang and Hu (), there exists a width $w$ such that for any two networks $M_{1},M_{2}$ having widths $w_{1},w_{2}>w$, $|\mathcal{R}(M_{1},t)-\mathcal{R}(M_{2},t)|\leq\epsilon$, where $\mathcal{R}(M,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 a limit as width tends to $\infty$, 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 $t$ differs by at most $\epsilon$.
∎

Note that the assumption of a fixed learning rate schedule with increasing width might seem strong, but recent works (Yang et al., ) have shown, that one of these initializations, termed as Maximal Update Parameterization ($\mu$P), 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 et al., , ; Liu et al., ; Hu et al., ). Moreover, works (Yang et al., ; Vyas et al., ) have empirically demonstrated that with $\mu$P, networks start exhibiting consistent loss curves at practical widths.

Moreover, as the above theorem holds for a fixed batch size $B$ as well as $B\to\infty$, 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 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., ; Bordelon et al., , ), and thus the arguments above also hold for these networks.

###### Theorem 2 .

###### Proof of Theorem 2 .

#### 4.2 Fixed Model Size and Scaling up Data Size

We now turn to studying the impact of data size in mini-batch SGD for a well-specified Gaussian linear regression problem.
Let $(\mathbf{x},y)$ be a pair of covariates and responses from a population distribution. Let the population risk and the population distribution be

$$ $\displaystyle\mathcal{R}(\mathbf{w}):=\mathbb{E}(\mathbf{x}^{\top}\mathbf{w}-y )^{2},\quad\mathbf{x}\sim\mathcal{N}(0,\mathbf{H}),\quad y|\mathbf{x}\sim \mathcal{N}(\mathbf{x}^{\top}\mathbf{w}^{*},\sigma^{2}),$ $$

where $\mathbf{w}$ is the trainable parameter, the expectation is over the population distribution, and $(\mathbf{H},\mathbf{w}^{*},\sigma^{2})$ specify the population distribution.
Given $D$ independent samples $(\mathbf{x}_{i},y_{i})_{i=1}^{D}$ from the population distribution, we consider an estimate given by mini-batch SGD,

$$ $\displaystyle\mathbf{w}_{0}=0,\quad\mathbf{w}_{t+1}=\mathbf{w}_{t}-\gamma\frac {1}{B}\sum_{j=tB}^{(t+1)B-1}(\mathbf{x}_{j}^{\top}\mathbf{w}_{t}-y_{j})\mathbf {x}_{j},\quad t=0,\dots,n-1,$ $$

where $\gamma>0$ is a constant learning rate, $B$ is the batch size, $n:=D/B$ is the number of steps, $\mathbf{w}_{0}=0$ is the initialization (without loss of generality), and the output is the average of the iterates, $\bar{\mathbf{w}}:=\frac{1}{n}\sum_{t=0}^{n-1}\mathbf{w}_{t}$.
Then the following theorem provides a tight bound on the excess risk achieved by the average of the mini-batch SGD iterates.

We write $f(D)\lesssim g(D)$ if there is a positive constant $c$ such that $f(D)\leq cg(D)$ for every $D\geq 1$. We write $f(D)\eqsim g(D)$ if $f(D)\lesssim g(D)\lesssim f(D)$.
The proofs are all deferred to Appendix G.

###### Theorem 3 .

Let $(\lambda_{i})_{i>0}$ be the eigenvalues of $\mathbf{H}$ in nonincreasing order.
Assume that $\|\mathbf{w}_{0}-\mathbf{w}^{*}\|_{\mathbf{H}}^{2}\lesssim\sigma^{2}$.
Then for every $\gamma\lesssim\min\{B/\mathrm{tr}(\mathbf{H}),1/\|\mathbf{H}\|_{2}\}$, we have

$$ $\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\sigma^{2}\eqsim\bigg{(} \frac{B}{D\gamma}\bigg{)}^{2}\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{H} ^{-1}_{0:k^{*}}}+\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{H}_{k^{*}: \infty}}+\sigma^{2}\frac{k^{*}+(D\gamma/B)^{2}\sum_{i>k^{*}}\lambda_{i}^{2}}{D},$ $$

where $k^{*}:=\max\{k:\lambda_{k}\geq B/(D\gamma)\}$ and the expectation is over the randomness of $\bar{\mathbf{w}}$.

The proof of Theorem 3 is motivated by Zou et al. ().
We focus on well-specified Gaussian data distribution for simplicity, but this can be relaxed to misspecified cases under fourth-moment conditions following the results in Zou et al. ().
Theorem 3 suggests that for a fixed data size $D$, the excess risk depends on the batch size $B$ and the learning rate $\gamma$ only through their ratio $\gamma/B$.
Moreover, a large $\gamma/B$ tends to decrease the bias error (the terms depending on $\mathbf{w}^{*}$) but increase the variance error (the terms depending on $\sigma^{2}$), 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.

###### Corollary 2 .

Under the settings of Theorem 3,
additionally assume $\sigma^{2}\eqsim 1$ and the following capacity and source conditions:

$$ $\displaystyle\text{for}\ a,b>1:\quad\lambda_{i}\eqsim i^{-a},\quad\mathbb{E} \lambda_{i}\langle\mathbf{v}_{i},\mathbf{w}^{*}_{i}\rangle^{2}\eqsim i^{-b}, \quad\mathbb{E}\lambda_{i}\langle\mathbf{v}_{i},\mathbf{w}^{*}_{i}\rangle \langle\mathbf{v}_{j},\mathbf{w}^{*}_{j}\rangle=0\ \text{for}\ i\neq j,$ $$

where $(\lambda_{i},\mathbf{v}_{i})_{i>0}$ are the eigenvalues and the corresponding eigenvectors of $\mathbf{H}$, and the expectation is over a prior of $\mathbf{w}^{*}$. Then we have

- 1.
When $b\leq a$, the optimal hyper-parameters (that minimize the expected excess risk up to constant factors) are $\gamma^{*}\eqsim 1$ and $B^{*}=1$.
- 2.
When $b>a$, the optimal hyper-parameters are $\gamma^{*}$ and $B^{*}$ such that
$\displaystyle 0<\gamma^{*}\lesssim 1,\quad 1\leq B^{*}\leq D,\quad\gamma^{*}/B
^{*}\eqsim D^{\frac{a}{\min\{b,2a+1\}}-1}.$
Therefore, the CBS is $B^{*}\eqsim D^{1-a/\min\{b,2a+1\}}$, which (along with $\gamma^{*}\eqsim 1$) allows mini-batch SGD output $\bar{\mathbf{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 (Caponnetto and De Vito, ) and are recently used to study scaling laws theory (Bordelon et al., ; Lin et al., ; Paquette et al., ).
According to Corollary 2,
when $b\leq a$, the bias error tends to dominate the variance error, in this case, the CBS is $B^{*}=1$ to allow a maximum number of optimization steps.
when $b<a$,
the variance error tends to dominate the bias error, and the optimal choices of batch size and learning rate balance these two errors.
While one can use $B^{*}=1$ and a small $\gamma^{*}$
to achieve the best excess risk rate, this leads to a suboptimal sequential runtime ($n=D/B$).
In this case, the CBS is $B^{*}\eqsim D^{1-a/\min\{b,2a+1\}}$, which achieves the optimal excess risk rate while minimizing the sequential runtime.

###### Theorem 3 .

###### Corollary 2 .

### 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 $\mathcal{R}$. These laws enable the prediction of training loss $\mathcal{R}$ based on available resources, making it possible to optimize resource allocation for efficient model training.
For example, Hestness et al. () found $\mathcal{R}\propto D^{-\alpha}$, with $\alpha\in[0.07,0.35]$. Of the factors they varied, only tasks can change the exponent $\alpha$. Changing the architecture optimizers, regularizers, and loss functions, would only change the proportionality factor, not the exponent; Henighan et al. () studied statistical relations between $N,D,C,\mathcal{R}$, over a wide range of values and found similar scaling laws, over the range of $N\in\left[10^{3},10^{9}\right],C\in\left[10^{12},10^{21}\right]$, and over multiple modalities (text, video, image, text to image, etc.).
(Kaplan et al., ) states that $N$ should be scaled faster than $D$. However, Chinchilla scaling (Hoffmann et al., ) found that models are under-trained, and then suggests that when given an increased budget (in FLOPs), to achieve compute-optimal, model size $N$ and data size $D$ should scale in approximately equal proportions.
Recent efforts (Pearce and Song, ; Besiroglu et al., ; Porian et al., ) have been made in reproducing the scaling laws from (Hoffmann et al., ) and the (Kaplan et al., ).
Different from our focus on measuring the efficiency notion of CBS, most of them focus on deriving optimal hyper-parameters (Bi et al., ; Porian et al., ) including learning rate and batch size from small-scale training given a fixed compute budget FLOPs $\approx 6ND$ without decoupling the effects of model size and data size.

##### 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., ; Zhang et al., ; Kaplan et al., ). McCandlish et al. () introduce the gradient noise scale, a measure that captures the variation in 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 to the performance of standard SGD at smaller batch sizes (Shallue et al., ). Additionally, Zhang et al. () analyze the impact of curvature on CBS using a noisy quadratic model, demonstrating that preconditioning techniques can increase the CBS.
Golmant et al. () show that the size of the dataset plays a smaller role in determining training efficiency compared to factors like model architecture and data complexity. In contrast, Hilton et al. () examine how performance can be maintained at smaller batch sizes. Meanwhile, Smith et al. (); Smith and Le () empirically investigated how the optimal learning rate changes based on momentum and training set size.
Theoretical work has further sought to characterize CBS by analyzing SGD behavior in least-squares linear regression, especially in over-parameterized settings (Jain et al., ; Ma et al., ).
Filatov et al. () 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 justifications or address the challenge of selecting optimal runs across a broad range of hyperparameters.
Our work advances the optimization literature by formalizing CBS and quantifying its growth w.r.t. data size and emphasizing the importance of common hyper-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 relationship 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 cosine scheduling without requiring 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 N00014-22-1-2377, and the National Science Foundation Grant under award #IIS 2229881. NV and DM are supported by a Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant W911NF2010021, and DOE grant DE-SC0022199.

## Appendix

### Appendix 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 not training well with proper hyper-paragrams, we take into account the warmup fraction in training as well:

Figure: Figure 7: Ablation of warmup steps used in the linear LR warmup stage for large batch sizes.
Refer to caption: x7.png

We sweep over warmup step ratios (how many fraction of training steps do we need to linearly scale the learning rate from zero) over 0.25 and 0.1 and find that 0.25 works best for 85M models. Therefore, we fix this number of warmup steps to be 0.25 of the $t_{\textrm{Chin}}$ for future experiments.
For 151M models, we sweep over the fraction of warmup steps in $\{0.15,0.25,0.35\}$. We show in Figure 7 that using a warmup ratio of 0.25 can be a reasonable design choice as it enjoys consistently better performance than 0.15 yet only slightly underperforms 0.35.
After we found that setting the warmup steps according to this heuristic, we use the ratio proportionally for all the other model sizes.

Examining the last part of training.
By closely examining the final stages of the training process (Figure 8), it 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, a very high EWA decay rate would be needed even for a 1.2B model with a moderate batch size of 1024.
Moreover, we observe that the optimization process is notably influenced by the final phase of training. For instance, by step 10,000, most runs achieve a validation loss below 3.2 (Figure 8(a)), and similarly, a loss below 2.8 is reached by step 30,000 (Figure 8(b)).
However, to 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.

Figure: (a) EWA Ablation
Refer to caption: x8.png

Scheduler comparison for other batch sizes.
In Figure 9, we include more comparisons on different 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 we sweep over the decay steps as $0.1,0.2,0.3\times$ 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 of training, optimizing the performance of learning rate decay.
Under small batch sizes, the schedule-free optimizer (Defazio et al., ) is a competitive baseline but it is significantly worse for batch sizes larger than 1024.

Figure: (a) Batch size 64
Refer to caption: x10.png

Longer training requires higher EWA decay rate $\tau$.
Throughout the paper, we adopt a learning 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 (DeepSeek-AI et al., ).
Therefore, we justify our design decision on tuning EWA decay rate for simulating learning rate decay on different training durations by conducting the following experiments on a series of 151M models with
(a) batch size 256, 0.5$\times$ Chinchilla tokens;
(b) batch size 256, 20$\times$ Chinchilla tokens;
(c) batch size 2048, 20$\times$ Chinchilla tokens, all with learning rate swept over $\{0.00316,0.00158,0.01264,0.00632,0.00075\}$ and EWA decay rate $\tau$ 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, a learning rate of $0.00316$ we use throughout the paper is consistently the best.
Moreover, when enlarging the training data size from 0.5$\times$ Chinchilla tokens to 20$\times$, the optimal EWA decay rate value $\tau$ would increase as well.
This is also justified in the results presented in Figure 8 and Table 4 which indicate that longer training durations may benefit from a higher EWA decay rate to improve optimization performance.

Figure: 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 a ★ symbol. Longer training durations, as seen in (b) and (c), necessitate a higher EWA decay rate for a given learning rate.
Refer to caption: x18.png

### Appendix 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, we focus on two key hyper-parameters that significantly affect optimization efficiency and examine their impact in detail.

The effect of momentum $\beta_{1}$ of Adam on CBS.
We sweep over several momentum $\beta_{1}$ values in Adam for all learning rates and batch sizes: [0, 0.8, 0.9, 0.95, 0.975].
Overall, Figure 11 shows that language model pre-training may need a large momentum value to be efficient and $\beta_{1}$ = 0.95 is slightly better ($<$0.02 gain on eval loss) than 0.9 for batch sizes.
We observe that in small batch size regimes like $2^{6}$, the performance gap between optimizing with and without momentum $\beta_{1}$ is small while the gap increases as we double the batch size (Shallue et al., ).
Moreover, we show that momentum 0.9 and 0.975 have similar effects on the number of steps needed to reach a target validation loss and critical batch sizes. On the other hand, small momentum, especially no momentum, would hurt the optimization.
This aligns well with the extensively studied acceleration of momentum in SGD with momentum (Goh, ).

Figure: Figure 11: Ablation results on momentum. All data points are trained using 151M models and a total number of optimization steps to reach a fixed target loss is reported.
Refer to caption: x19.png

The effect of the second moment decay rate $\beta_{2}$ in Adam. As reported in Appendix Table 4, we found that $\beta_{2}$ in Adam, the exponential decay rate of the second momentum estimate of gradients 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 ablate $\beta_{2}\in[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., ; Wortsman et al., ; Groeneveld et al., ).
For large batch sizes $[1024,2048,4096,8192]$, we experiment with a small $\beta_{2}=0.9$ with the model size 151M, 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 8(b)), a high enough $\beta_{2}$ is necessary.

### Appendix C Results Including Small Batch Sizes

For completeness, we demonstrate linear scaling behavior in small-batch regimes across all model sizes (Figure 12). This shows that all models exhibit linear scaling (with reasonable deviations) with a batch size ranging from $2^{6}$ to $2^{10}$, 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 a batch size of 256 at the Chinchilla step.

Figure: Figure 12: Linear scaling regimes: doubling the batch size can halve the optimization steps to reach the target loss.
Refer to caption: x20.png

Moreover, we include all the results that contain the smallest several batch sizes (Figure 13).
Note that the denominator is the number of steps to reach target loss at batch size 64 instead of 256 now.
Now we can observe clear linear scaling of all the model sizes till around $2^{10}$ for model sizes, while the largest three model sizes maintain linear scaling till almost $2^{11}$.
There are minor differences with the main plot in Figure 1 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 is primarily on large batch sizes, we consistently use $2^{9}$ as the starting batch size throughout the main text.

Furthermore, recall that we set $B_{\text{opt}}=2^{8}$ when selecting the target loss. Figure 12 and Figure 13 confirm that $2^{8}$ falls within the linear scaling regime, which justifies our design choice.

Figure: Figure 13: Full results for different models in Chinchilla settings. We include both a largest batch size $2^{14}$ start the plot from several small batch sizes $2^{6},2^{7},2^{8}$. Relative number of steps w.r.t. batch size $\mathbf{2^{6}}$ is reported.
Refer to caption: x21.png

### Appendix D Additional Details on Experiments

Optimizer setup.
For optimizers, we try both SGD (Robbins and Monro, ) and Adam (Kingma, ) and find that SGD without momentum is significantly worse so we use Adam only for all the experiments.
We disable weight decay in Adam.
For generality, though the training set C4 might contain low-quality or duplicated documents that can potentially lead to training instability (Muennighoff et al., ; Wortsman et al., ), we observed that these issues did not affect our primary target of interest—namely, the final optimization efficiency. As a result, we didn’t explicitly adopt additional normalization like QK normalization (Dehghani et al., ; Zhai et al., ) or a z-loss (Chowdhery et al., ) to mitigate loss spikes.(^2^22We observe irregular loss spikes despite adopting gradient clipping, but most runs can still be optimized well in the end.) We set $\epsilon=$1.0e-8 by default, and refer to momentum as $\beta_{1}$ in Adam by default throughout the paper.

Chinchilla steps for batch size 256 that determine the target losses.
For each model size, we aim to establish a target validation loss by training with a global batch size of 256 on a 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 1.
We use a token-to-model size ratio $C_{\textrm{Chin}}$ 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.
As we need to frequently evaluate the model on a holdout evaluation set, understanding how large evaluation set size can lead to a stable and reliable measure yet still pertain efficiency of each run can be important.
We evaluated the variance across different token counts: 2.17e-4 for 327,680 tokens, 4.53e-5 for 1,638,400 tokens, and 7.65e-6 for 3,276,800 tokens. Based on these results, we have set the default number of evaluation batches to 100.

It is important to note that the total number of training steps varies significantly across different batch sizes. To address this, we implement a hybrid evaluation protocol: the model is evaluated at intervals of $2^{i}$ (where $i\in\mathbb{Z}$ ), every 1,000 steps, and at $0.7n,0.75n,0.8n,\ldots$, up to $n$ during the last $30\%$ of the total steps $n$. This approach ensures more frequent evaluations toward the end of training, allowing for a more accurate assessment of the total training steps needed to achieve a target evaluation loss.

Hyper-parameter search details.
Due to compute constraints, we cannot perform an exhaustive search over 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 (151M parameters).
We optimize the following hyper-parameters in sequence: learning rate, momentum $\left(\beta_{1}\right)$, warmup steps, scheduler, and context length. Additionally, we tune $\beta_{2}$ and $\tau$ for each model size and batch size. Specifically, for large batch sizes ($>$1024), a smaller $\beta_{2}$ and larger $\tau$ tend to be more effective, while the opposite holds true for smaller batch sizes.

Below we show the hyper-parameter choices (Table 3) and optimal ones (Table 4) we report in our main plot for studying CBS with respect to model sizes.
Additionally, Table 5 presents various model size configurations and scaling methods, with models in bold indicating those used in our controlled experiments.

Table: Table 2: Model architecture details.

Table: Table 3: Sweeping experiments settings. Default values after the hyper-parameter search are in Bold font. Values in the parathesis mean used not for every sweep.

**Table 4: Optimal hyper-parameters for different model sizes. The optimal means the number of steps to reach a target validation loss. We refer $\beta_{2}$ as the exponential decay rate for the second-moment estimates in Adam, $\tau$ as the interpolation parameters in EWA ($\xi_{t+1}=\tau\cdot\xi_{t}+\left(1-\tau\right)\cdot\theta_{t}$). All the optimal runs are trained with momentum $\beta_{1}=0.95$ and learning rate 3.16e-3.**
| Batch Size | $\beta_{2}$ | $\tau$ |
| --- | --- | --- |
| 85M |  |  |
| 64 | 0.999 | 0.9995 |
| 128 | 0.999 | 0.9995 |
| 256 | 0.999 | 0.9995 |
| 512 | 0.999 | 0.998 |
| 1024 | 0.95 | 0.99 |
| 2048 | 0.95 | 0.99 |
| 4096 | 0.95 | 0.99 |
| 8192 | 0.95 | 0.98 |
| 16384 | 0.95 | 0.98 |

**Table 5: 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 | $n_{\text{heads}}$ | $\mathbf{n_{\textbf{layers}}}$ | $d_{\text{model}}$ | $\text{MLP}_{\text{hidden}}$ |
| --- | --- | --- | --- | --- |
| 151M | 16 | 12 | 1024 | 4096 |
| 302.09M | 16 | 24 | 1024 | 4096 |
| 604.18M | 16 | 48 | 1024 | 4096 |
| 1.208B | 16 | 96 | 1024 | 4096 |

### Appendix 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)

Table: Table 6: Fitted scaling law parameters for Chinchilla settings when fixing $\alpha=1$: $\log(Y)=\log(a+\frac{b}{B^{\alpha}})$, where $Y$ is the number of steps to reach Chinchilla target loss, $B$ denotes the batch size, and the critical batch size is solved as $B^{*}=(\frac{b+5a\times 1.2\times B_{\text{opt}}}{5a})^{\frac{1}{\alpha}}$, $B_{\text{opt}}=256$.

We report forecasting results for various model and token sizes, extending beyond the plots presented in the main text (Table 7).
For each row, increasing either model size or token size shows that the forecasting results remain comparable.

**Table 7: Additional forecasted CBS results for larger scale. Recall that we fit $B^{*}=93.20\times N^{0.47}$, $B^{*}=22.91\times D^{0.47}$ where model size $N$ is in millions and data size $D$ is in billions.**
| Model Size | Forecasted CBS | $\log_{2}(B^{*})$ |
| --- | --- | --- |
| 1.5B | 2862.17 | 11.48 |
| 2B | 3274.93 | 11.68 |
| 2.5B | 3635.65 | 11.83 |
| 3B | 3959.69 | 11.95 |
| 3.5B | 4256.09 | 12.06 |
| 4B | 4530.72 | 12.15 |
| 4.5B | 4787.63 | 12.23 |
| 5B | 5029.77 | 12.30 |
| 5.5B | 5259.34 | 12.36 |
| 6B | 5478.06 | 12.42 |

Note that our definition of CBS and its scaling law have a similar interpretation with the one in (McCandlish et al., ; Kaplan et al., ) as $\frac{E_{\min}}{S_{\min}}$, ${S_{\min}}$ denotes the minimum possible number of steps taken to reach target loss and $E_{\min}$ is the minimum possible number of training examples processed to reach target loss.
In particular, recall that critical batch size can be analytically derived as $B^{*}=\frac{b}{5a}+1.2B_{\mathrm{opt}}$.
This relationship reflects the point where batch size scaling incurs a 20% overhead when the batch size is doubled while (McCandlish et al., ).
Here, the parameter $b$ plays a role analogous to $E_{\min}$, while $a$ corresponds to $S_{\min}$, depending on the specific overhead chosen to characterize the diminishing returns from increasing the batch size.

### Appendix F Reproducibility

In our training environment, we verify that, across multiple model sizes (2.4M, 9.4M, 19M, 42M, 85M, 151M, 302M), we can (approximately) reproduce the final evaluation loss of Figure 1 in (Wortsman et al., ).
We use nodes equipped with 8 A100 GPUs, each with 80GiB of memory, for model training.
We built our training framework using the Olmo training suite (Groeneveld et al., ).

### Appendix G Complete Proofs in Section 4.2

###### Proof of Theorem 3 .

The work by Zou et al. () studied SGD with batch size $1$ for linear regression and established matching (up to a 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 a set of operations on PSD matrices as follows:

$$ $\displaystyle\mathcal{I}=\mathbf{I}\otimes\mathbf{I},\quad\mathcal{M}^{\mathrm {B}}=\mathbb{E}\bigg{[}\bigg{(}\frac{1}{B}\sum_{i\in\mathcal{I}}\mathbf{x}_{i} \mathbf{x}_{i}^{\top}\bigg{)}\otimes\bigg{(}\frac{1}{B}\sum_{i\in\mathcal{I}} \mathbf{x}_{i}\mathbf{x}_{i}^{\top}\bigg{)}\bigg{]},\quad\widetilde{\mathcal{M }}=\mathbf{H}\otimes\mathbf{H},$ $\displaystyle\mathcal{T}^{\mathrm{B}}=\mathbf{H}\otimes\mathbf{I}+\mathbf{I} \otimes\mathbf{H}-\gamma\mathcal{M}^{\mathrm{B}},\quad\widetilde{\mathcal{T}}= \mathbf{H}\otimes\mathbf{I}+\mathbf{I}\otimes\mathbf{H}-\gamma\mathbf{H} \otimes\mathbf{H},$ $$

where $\mathcal{I}$ is an index set of $B$ independent data.
Note that

$$ $\displaystyle\big{(}\mathcal{M}^{\mathrm{B}}-\widetilde{\mathcal{M}}\big{)} \circ\mathbf{A}=\mathrm{Cov}\bigg{(}\frac{1}{B}\sum_{i\in\mathcal{I}}\mathbf{x }_{i}\mathbf{x}_{i}^{\top}\mathbf{A}^{1/2}\bigg{)}=\frac{1}{B}\mathrm{Cov}( \mathbf{x}\mathbf{x}^{\top}\mathbf{A}^{1/2}).$ $$

For Gaussian data $\mathbf{x}\in\mathcal{N}(0,\mathbf{H})$, we have

$$ $\displaystyle\mathrm{Cov}(\mathbf{x}\mathbf{x}^{\top}\mathbf{A}^{1/2})=\mathbb {E}_{\mathbf{x}\in\mathcal{N}(0,\mathbf{H})}\big{[}\mathbf{x}\mathbf{x}^{\top} \mathbf{A}\mathbf{x}\mathbf{x}^{\top}\big{]}-\mathbf{H}\mathbf{A}\mathbf{H}=2 \mathrm{tr}(\mathbf{H}\mathbf{A})\mathbf{H}.$ $$

Together, we obtain

$$ $\displaystyle\big{(}\mathcal{M}^{\mathrm{B}}-\widetilde{\mathcal{M}}\big{)} \circ\mathbf{A}=\frac{2}{B}\mathrm{tr}(\mathbf{H}\mathbf{A})\mathbf{H}.$ $$

Now we compute the error propagation along the SGD steps. Let $\bm{\eta}_{t}=\mathbf{w}_{t}-\mathbf{w}^{*}$ be the error vector.
For convenience, let $\mathbf{G}_{t}=\frac{1}{B}\sum_{i\in\mathcal{I}_{t}}\mathbf{x}_{i}\mathbf{x}_{
i}^{\top}$ be the empirical covariance of an independent batch.
Then we can define the bias and variance iterates as

$$ $\displaystyle\bm{\eta}_{t}^{\mathrm{bias}}=\big{(}\mathbf{I}-\gamma\mathbf{G}_ {t}\big{)}\bm{\eta}_{t-1}^{\mathrm{bias}},\quad t=1,\dots,n-1,\quad\bm{\eta}_{ 0}^{\mathrm{bias}}=\mathbf{w}_{0}-\mathbf{w}^{*},$ $$

and

$$ $\displaystyle\bm{\eta}_{t}^{\mathrm{variance}}=\big{(}\mathbf{I}-\gamma\mathbf {G}_{t}\big{)}\bm{\eta}_{t-1}^{\mathrm{variance}}+\gamma\cdot\frac{1}{B}\sum_{ i\in\mathcal{I}_{t}}\xi_{i}\mathbf{x}_{i},\quad t=1,\dots,n-1,\quad\bm{\eta}_{ 0}^{\mathrm{variance}}=\bm{0},$ $$

where $\xi_{i}=y_{i}-\mathbf{x}^{\top}_{i}\mathbf{w}^{*}\sim\mathcal{N}(0,\sigma^{2})$.
We then compute the covariance matrices of these two error iterates

$$ $\mathbf{B}_{t}^{\mathrm{B}}:=\mathbb{E}[\bm{\eta}_{t}^{\mathrm{bias}}\otimes \bm{\eta}_{t}^{\mathrm{bias}}],\quad\mathbf{C}_{t}^{\mathrm{B}}:=\mathbb{E}[ \bm{\eta}_{t}^{\mathrm{variance}}\otimes\bm{\eta}_{t}^{\mathrm{variance}}].$ $$

Using the operators, these covariance matrices take the following iterative updates:

$$ $\displaystyle\mathbf{B}_{0}^{\mathrm{B}}=\bm{\eta}_{0}\otimes\bm{\eta}_{0}, \quad\mathbf{B}_{t}^{\mathrm{B}}$ $\displaystyle=\mathbb{E}_{\mathbf{G}_{t}}\big{[}(\mathbf{I}-\gamma\mathbf{G}_{ t})\mathbf{B}_{t-1}^{\mathrm{B}}(\mathbf{I}-\gamma\mathbf{G}_{t})\big{]}=\big{ (}\mathcal{I}-\gamma\mathcal{T}^{\mathrm{B}}\big{)}\circ\mathbf{B}_{t-1}^{ \mathrm{B}},$ $\displaystyle\mathbf{C}_{0}^{\mathrm{B}}=\bm{0},\quad\mathbf{C}_{t}^{\mathrm{B}}$ $\displaystyle=\mathbb{E}_{\mathbf{G}_{t}}\big{[}(\mathbf{I}-\gamma\mathbf{G}_{ t})\mathbf{C}_{t-1}^{\mathrm{B}}(\mathbf{I}-\gamma\mathbf{G}_{t})\big{]}+\frac {\gamma^{2}}{B^{2}}\mathbb{E}\bigg{[}\bigg{(}\sum_{i\in\mathcal{I}_{t}}\xi_{i} \mathbf{x}_{i}\bigg{)}\bigg{(}\sum_{i\in\mathcal{I}_{t}}\xi_{i}\mathbf{x}_{i} \bigg{)}^{\top}\bigg{]}$ $\displaystyle=\big{(}\mathcal{I}-\gamma\mathcal{T}^{\mathrm{B}}\big{)}\circ \mathbf{C}_{t-1}^{\mathrm{B}}+\frac{\gamma^{2}\sigma^{2}}{B}\mathbf{H},$ $$

where the last equation is because

$$ $\displaystyle\mathbb{E}\bigg{[}\bigg{(}\sum_{i\in\mathcal{I}_{t}}\xi_{i} \mathbf{x}_{i}\bigg{)}\bigg{(}\sum_{i\in\mathcal{I}_{t}}\xi_{i}\mathbf{x}_{i} \bigg{)}^{\top}\bigg{]}=\mathbb{E}\bigg{[}\sum_{i\in\mathcal{I}_{t}}\xi_{i}^{2 }\mathbf{x}_{i}\mathbf{x}_{i}^{\top}\bigg{]}=\sigma^{2}B\mathbf{H}.$ $$

Recall that $\bar{\mathbf{w}}=\frac{1}{n}\sum_{t=0}^{n-1}\mathbf{w}_{t}$.
First, using Lemmas B.3 and C.1 in (Zou et al., ), we get the following bias-variance decomposition (note that our setting is well-specified):

$$ $\displaystyle\mathbb{E}[\mathcal{R}(\bar{\mathbf{w}})]-\min\mathcal{R}(\cdot)= \mathrm{bias}+\mathrm{variance},$ $$

where

$$ $\displaystyle\mathrm{bias}:=\frac{1}{2}\langle\mathbf{H},\mathbb{E}[{\bar{\bm{ \eta}}}^{\mathrm{bias}}\otimes{\bar{\bm{\eta}}}^{\mathrm{bias}}]\rangle\begin{ dcases}\leq\frac{1}{n^{2}}\sum_{t=0}^{n-1}\sum_{k=t}^{n-1}\big{\langle}( \mathbf{I}-\gamma\mathbf{H})^{k-t}\mathbf{H},\mathbf{B}_{t}^{\mathrm{B}}\big{ \rangle},\\ \geq\frac{1}{2n^{2}}\sum_{t=0}^{n-1}\sum_{k=t}^{n-1}\big{\langle}(\mathbf{I}- \gamma\mathbf{H})^{k-t}\mathbf{H},\mathbf{B}_{t}^{\mathrm{B}}\big{\rangle}, \end{dcases}$ $\displaystyle\mathrm{variance}:=\frac{1}{2}\langle\mathbf{H},\mathbb{E}[{\bar{ \bm{\eta}}}^{\mathrm{variance}}\otimes{\bar{\bm{\eta}}}^{\mathrm{variance}}] \begin{dcases}\leq\frac{1}{n^{2}}\sum_{t=0}^{n-1}\sum_{k=t}^{n-1}\big{\langle} (\mathbf{I}-\gamma\mathbf{H})^{k-t}\mathbf{H},\mathbf{C}_{t}^{\mathrm{B}}\big{ \rangle},\\ \geq\frac{1}{2n^{2}}\sum_{t=0}^{n-1}\sum_{k=t}^{n-1}\big{\langle}(\mathbf{I}- \gamma\mathbf{H})^{k-t}\mathbf{H},\mathbf{C}_{t}^{\mathrm{B}}\big{\rangle}, \end{dcases}$ $$

where

$$ $\displaystyle{\bar{\bm{\eta}}}^{\mathrm{bias}}:=\frac{1}{n}\sum_{t=0}^{n-1}{ \bar{\bm{\eta}}}^{\mathrm{bias}}_{t},\quad{\bar{\bm{\eta}}}^{\mathrm{variance} }:=\frac{1}{n}\sum_{t=0}^{n-1}{\bar{\bm{\eta}}}^{\mathrm{variance}}_{t}.$ $$

The remaining efforts are to characterize $\mathbf{B}_{t}^{\mathrm{B}}$ and $\mathbf{C}_{t}^{\mathrm{B}}$ for a batch size $B$.
For the bias part, we have

$$ $\displaystyle\mathbf{B}_{t}^{\mathrm{B}}$ $\displaystyle=\big{(}\mathcal{I}-\gamma\mathcal{T}^{\mathrm{B}}\big{)}\circ \mathbf{B}_{t-1}^{\mathrm{B}}$ $\displaystyle=\big{(}\mathcal{I}-\gamma\tilde{\mathcal{T}}\big{)}\circ\mathbf{ B}_{t-1}^{\mathrm{B}}+\gamma^{2}\big{(}\mathcal{M}^{\mathrm{B}}-\tilde{ \mathcal{M}}\big{)}\circ\mathbf{B}_{t-1}^{\mathrm{B}}$ $\displaystyle=\big{(}\mathcal{I}-\gamma\tilde{\mathcal{T}}\big{)}\circ\mathbf{ B}_{t-1}^{\mathrm{B}}+\frac{2\gamma^{2}}{B}\mathrm{tr}\big{(}\mathbf{H}\mathbf {B}_{t-1}^{\mathrm{B}}\big{)}\mathbf{H},\quad t=1,\dots,n-1.$ $$

For the variance part, we have

$$ $\displaystyle\mathbf{C}_{t}^{\mathrm{B}}$ $\displaystyle=(\mathcal{I}-\gamma\mathcal{T}^{\mathrm{B}})\circ\mathbf{C}_{t-1 }^{\mathrm{B}}+\frac{\gamma^{2}\sigma^{2}}{B}\mathbf{H}$ $\displaystyle=(\mathcal{I}-\gamma\tilde{\mathcal{T}})\circ\mathbf{C}_{t-1}^{ \mathrm{B}}+\gamma^{2}\big{(}\mathcal{M}^{\mathrm{B}}-\tilde{\mathcal{M}}\big{ )}\circ\mathbf{C}_{t-1}+\frac{\gamma^{2}\sigma^{2}}{B}\mathbf{H}$ $\displaystyle=(\mathcal{I}-\gamma\tilde{\mathcal{T}})\circ\mathbf{C}_{t-1}^{ \mathrm{B}}+\frac{2\gamma^{2}}{B}\mathrm{tr}\big{(}\mathbf{H}\mathbf{C}_{t-1}^ {\mathrm{B}}\big{)}\mathbf{H}+\frac{\gamma^{2}\sigma^{2}}{B}\mathbf{H},\quad t =1,\dots,n-1.$ $$

To obtain an upper bound on excess risk, we replace $\alpha$ in Assumption 2.2 of Zou et al. () with $2/B$, the number of steps with $n:=D/B$, and the noise level $\sigma^{2}$ with $\sigma^{2}/B$, then apply the proof of Theorem 2.1.
Similarly, for a lower bound on excess risk, we replace $\beta$ in Assumption 2.4 with $2/B$, the number of steps with $n:=T/B$, and the noise level $\sigma^{2}$ with $\sigma^{2}/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:

$$ $\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\min\mathcal{R}(\cdot)$ $\displaystyle\eqsim\bigg{(}\frac{1}{n\gamma}\bigg{)}^{2}\|\mathbf{w}_{0}- \mathbf{w}^{*}\|^{2}_{\mathbf{H}^{-1}_{0:k^{*}}}+\|\mathbf{w}_{0}-\mathbf{w}^{ *}\|^{2}_{\mathbf{H}_{k^{*}:\infty}}$ $\displaystyle\quad+\frac{1/B\big{(}\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{ \mathbf{I}_{0:k^{*}}}+n\gamma\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{H} _{k^{*}:\infty}}\big{)}}{n\gamma}\cdot\frac{k^{*}+(n\gamma)^{2}\sum_{i>k^{*}} \lambda_{i}^{2}}{n}$ $\displaystyle\quad+\frac{\sigma^{2}}{B}\cdot\frac{k^{*}+(n\gamma)^{2}\sum_{i>k ^{*}}\lambda_{i}^{2}}{n},$ $$

where $k^{*}=\max\{k:\lambda_{k}\geq 1/(n\gamma)\}$, and a sufficient stepsize condition (see Lemma 4.1, Theorems 2.1 and 2.2 in Zou et al. ()) is

$$ $\displaystyle 0<\gamma\lesssim\min\bigg{\{}\frac{1}{\alpha\mathrm{tr}(\mathbf{ H})},\frac{1}{\|\mathbf{H}\|_{2}}\bigg{\}}\eqsim\min\bigg{\{}\frac{B}{\mathrm{ tr}(\mathbf{H})},\frac{1}{\|\mathbf{H}\|_{2}}\bigg{\}}.$ $$

The assumption $\|\mathbf{w}_{0}-\mathbf{w}^{*}\|_{\mathbf{H}}^{2}\lesssim\sigma^{2}$ implies

$$ $\frac{\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{I}_{0:k^{*}}}+n\gamma\| \mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{H}_{k^{*}:\infty}}}{n\gamma}\leq \|\mathbf{w}_{0}-\mathbf{w}^{*}\|_{\mathbf{H}}^{2}\lesssim\sigma^{2},$ $$

which further simplifies the excess risk bounds to

$$ $\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\min\mathcal{R}(\cdot) \eqsim\bigg{(}\frac{1}{n\gamma}\bigg{)}^{2}\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{ 2}_{\mathbf{H}^{-1}_{0:k^{*}}}+\|\mathbf{w}_{0}-\mathbf{w}^{*}\|^{2}_{\mathbf{ H}_{k^{*}:\infty}}+\frac{\sigma^{2}}{B}\cdot\frac{k^{*}+(n\gamma)^{2}\sum_{i>k ^{*}}\lambda_{i}^{2}}{n}.$ $$

Finally, replacing $n=D/B$ in the bounds completes our proof.
∎

###### Proof of Corollary 2 .

By $\lambda_{i}\eqsim i^{-a}$, we can solve for $k^{*}$ to obtain
$k^{*}\eqsim({D\gamma}/{B})^{1/a}.$
We then calculate the expected excess risk by Theorem 3 using the capacity and source conditions:

$$ $\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\sigma^{2}$ $\displaystyle\eqsim\mathbb{E}\bigg{(}\bigg{(}\frac{B}{D\gamma}\bigg{)}^{2}\| \mathbf{w}^{*}\|^{2}_{\mathbf{H}^{-1}_{0:k^{*}}}+\|\mathbf{w}^{*}\|^{2}_{ \mathbf{H}_{k^{*}:\infty}}\bigg{)}+\frac{k^{*}+(D\gamma/B)^{2}\sum_{i>k^{*}} \lambda_{i}^{2}}{D}$ $\displaystyle\eqsim\bigg{(}\frac{B}{D\gamma}\bigg{)}^{2}\sum_{i\leq k^{*}}i^{- b+2a}+\sum_{i>k^{*}}i^{-b}+\frac{1}{D}\bigg{(}k^{*}+\bigg{(}\frac{D\gamma}{B} \bigg{)}^{2}\sum_{i>k^{*}}i^{-2a}\bigg{)}$ $\displaystyle\eqsim\bigg{(}\frac{B}{D\gamma}\bigg{)}^{2}\max\big{\{}(k^{*})^{1 -b+2a},1\}+(k^{*})^{1-b}+\frac{1}{D}\bigg{(}k^{*}+\bigg{(}\frac{D\gamma}{B} \bigg{)}^{2}(k^{*})^{1-2a}\bigg{)}$ $\displaystyle\eqsim\max\bigg{\{}\bigg{(}\frac{D\gamma}{B}\bigg{)}^{(1-b)/a},\, \bigg{(}\frac{D\gamma}{B}\bigg{)}^{-2}\bigg{\}}+\frac{1}{D}\bigg{(}\frac{D \gamma}{B}\bigg{)}^{1/a}.$ $$

We then discuss three cases.

- 1.
When $b\leq a$, we have
$\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\sigma^{2}$
$\displaystyle\eqsim\bigg{(}\frac{D\gamma}{B}\bigg{)}^{(1-b)/a}+\frac{1}{D}
\bigg{(}\frac{D\gamma}{B}\bigg{)}^{1/a}\eqsim\bigg{(}\frac{D\gamma}{B}\bigg{)}
^{(1-b)/a},$
where the last equality is because $\gamma/B\lesssim 1$ so the first term dominates the second term. So the optimal hyper-parameters are $\gamma^{*}\eqsim 1$ and $B^{*}=1$.
- 2.
When $a<b<2a+1$, we have
$\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\sigma^{2}$
$\displaystyle\eqsim\bigg{(}\frac{D\gamma}{B}\bigg{)}^{(1-b)/a}+\frac{1}{D}
\bigg{(}\frac{D\gamma}{B}\bigg{)}^{1/a},$
so the optimal hyper-parameters are
$\displaystyle 0<\gamma^{*}\lesssim 1,\quad 1\leq B^{*}\leq D,\quad\gamma^{*}/B
^{*}\eqsim D^{a/b-1}.$
- 3.
When $b>2a+1$, we have
$\displaystyle\mathbb{E}\mathcal{R}(\bar{\mathbf{w}})-\sigma^{2}$
$\displaystyle\eqsim\bigg{(}\frac{D\gamma}{B}\bigg{)}^{-2}+\frac{1}{D}\bigg{(}
\frac{D\gamma}{B}\bigg{)}^{1/a},$
so the optimal hyper-parameters are
$\displaystyle 0<\gamma^{*}\lesssim 1,\quad 1\leq B^{*}\leq D,\quad\gamma^{*}/B
^{*}\eqsim D^{a/(2a+1)-1}.$

Combining the second and third cases completes the proof.
∎

###### Proof of Theorem 3 .

###### Proof of Corollary 2 .

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
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)?
- I found the plots to sometimes be a bit confusing or less accessible. Some suggestions:
  - Could you add the linear scaling line as well as the 20% region?
  - Could you highlight the critical batch sizes, e.g. by a star marker?
  - Nit: The font sizes are inconsistent between subplots, e.g. the x-axis labels between Figure 1a and the subplots of Figure 1b.
  - In the legend of Figure 1b (left), what is the number in parenthesis? Is it the (relative) number of training samples?
  - 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 :)
  - 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.
- 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?
- 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?
- 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?
- 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?
- Nits:
  - Line 120: I believe the "Refer" is superfluous.
  - Line 157: There is probably a closing parenthesis missing after the second citation.
  - Line 185: Typo for "maximum".
  - Figure 3a doesn't need to be a subplot as there is no other subplot.
  - 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?
  - Line 271: "Compute-optimal" should probably be lowercase.
  - Line 288: Typo with "achieved" and "batch size".
  - Line 351: "chinchilla" should probably be capitalized.
  - Line 366: "They" -> "Then"?
soundness: 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.
- The paper considers models ranging from 85 million to 1.2 billion parameters and thus covers a reasonably large domain of models.
- I really liked the highlighted practical takeaway blocks throughout the paper, which made it easy to understand, well-structured, and accessible.
summary: The paper investigates the scaling behavior of critical batch size in the pre-training of autoregressive language models.
They 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.
weaknesses: Some of the takeaways seem to me a bit too bold or not backed by enough evidence for the given claim.

- 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.
- 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.)?
- 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?

[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>

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - (Minor) It would be nice to use $\times$ instead of * (asterisk) in line 135.
- (Minor) The color contrast in Figure 9 can be improved.
soundness: 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.
- The experimental design is rigorous; for example, decoupling various hyperparameters makes the claims more convincing.
- Detailed experimental procedures are provided in Appendix D.
- 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.
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.
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.
- 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).
- However, I believe that the work has substantial limitations that would prevent me from recommending acceptance.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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?
2. Any reason explaining the worsening _efficiency_ for a 1.2B model on larger batch sizes in Figure 1.b) (right)?
3. 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? 
    * In a similar vein, can we see in Figure 2 what only constant LR schedules look like?
    * Is cosine here with or without warmups?
4. Why does Figure 2 (right) have different step lengths for the same batch size assuming similar compute budgets overall, given same model size? 
    * Similar to above, why does Figure 9 all have different lengths?
5. 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?
6. 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?
    * For every result seen as plot/table, does it mean there was an equivalent run made on a batch size 256/512?
7. 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)
    * How does the model size change when Figure 3 writes about a model of size 151M but with context size ranging from 512-4096?
    * Are model sizes calculated without the embedding parameters?
8. Could the 20% and the resulting 5 in Section 3.2 be more generalized in the formalism?
    * Could the readers have a clearer reference/intuition for the 20% overhead?
9. Could L318-319 have more support?
10. Does Figure 5 have empirical backing regarding the predicted CBS and the formalism of staying within the 20% margin of steps for similar loss?
    * Or if the decoupled prediction for batch sizes hold _reliably_ for different scales or ablations such as architecture, datasets, hyperparameters?
11. How exactly are the tuning decisions undertaken especially in Section 3.3 (example, L371-372)?
12. 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.
13. 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?
14. What is the specification or how is $\mathbf{H}$ defined in Section 4.2?
15. 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)?
16. How exactly does the _hybrid evaluation_ (L968) provide _more accurate_ information?
17. What or how are the hyperparameters (HPs) for parameter sweeps ranked or ordered (L977-979)? Does a different ordering yield potentially different results?
    * What are the default values for other hyperparameters when tuning each of these in order?
    * 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)?
18. 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?


Overall, the paper reads nice and does touch on an important and often under-studied aspect in scaling literature.
However, 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.
Given the expensive space of empirical experiments, the authors did well to consider multiple design choices and study them. 
Unfortunately, that also opens up more questions on how these designs were arrived at and that clarity would be my primary criticism.
One 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).


Scores may be increased depending on suitable responses to most of the points raised above.
Thank you for the paper, it was a nice read overall.


PS: I was unaware of Zou et al. 2023 and have not verified the proof on Pages 22-24.
soundness: 3
strengths: 1. The paper is well-written and largely easy to follow.
2. The related literature covers the important papers in the topic well.
3. Formalizations and hypotheses are clearly outlined and help understand results better.
4. Though personally I would like to reconsider its exact design and placement, the _Takeaway_ block was helpful while reading the paper first time.
5. The model scales reported in experiments are adequate in applying the insights to large-scale pre-training.
6. Formalizing the notion of critical batch size (CBS) through the 20% overhead assumption is novel and seemingly useful.
7. Considering various hyperparameters ablated across selected model scales is important and welcome in a scaling-related paper.
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.
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).
2. 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).
3. 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.
4. 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).
5. 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).
6. The outcome of the theoretical proofs and empirical results appear independent and do not knit the contributions as well.
    * Is there a way to connect the toy regression experiment with the scaling law exponent fits?


Minor issues/fixes:
1. Discrepancy of batch size 256 or 512 as the batch size for calculating overhead ratios
    * L52, Fig. 1,3,4,6,8 y-axes, L370, L888.
2. Should likely be `b > a` in L1299.
3. Typo on L366: perhaps should be "They" -> "Then"
4. Discrepancy in Fig. 5 caption where $C_{chin}$ comes to be around 6 as per the equation given for $B^*$.
5. Figure 4.a doesn't mention which model width or depth.
6. Figure 10 mentions _context length_ but has nothing to show for it.



Nitpicks:
1. Irregular and inconsistent plot sizes.
2. Recommended to refer sections/tables/figures in the Contributions list in Section 1.
3. Vertical lines in Table 4 separating different model sizes.
4. Parantheses in L1119-1122 would enable easier parsing for readers.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: I mention the main points in the section above, but repeat them here as questions (with some others):

- Why did the authors choose to disable weight decay?
- How did you use EWA in detail, or is there a specific implementation trick that was required?
- Did you tune SFO with the same sweeps as Adam?
- 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%).

I 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.

Another 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).
soundness: 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.
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.
weaknesses: While I very much appreciate the topic and investigation of the paper, I unfortunately have to be very critical of the experimental evaluation. 

- 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.
- 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.
- 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.

The 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.

-- update:

The 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: ### Please address

In Proof of Theorem 1 (page 8, line 418-419), it is stated that "the
trajectory of the network approaches a limit as width tends to ∞". I
would wish for what exactly is meant by "trajectory" to be defined
more clearly. E.g., is it the trajectory of the weights during
training? Also, please cite a source for this statement.

Appendix D, Evaluation data size and frequency (page 18, line 964ff):  
Which which model after what training have these evaluation variance
numbers been obtained?

### Minor comments

Page 3, line 158:  
Please describe the meaning of θ.

Page 8, line 416:  
$R(M, t)$ denotes the loss of network $N$ at time $t$.  
-> $R(M, t)$ denotes the loss of network $M$ at time $t$.

In the experiments, ideally, the optimizer's ε hyperparameter would
also have been scaled with the model size
(https://arxiv.org/abs/2309.14322).
soundness: 4
strengths: For their experiments, the authors used C4, a de-facto standard NLP
research dataset, which makes interpretation and reproduction of the
paper of the paper easier.

A particularly strong section is that of relevant work, in which
various important related papers are summarized well.

I also really enjoy the "takeaway" boxes as section summaries.
summary: The paper studies how a compute-optimal batch size, the critical batch
size (CBS) is influenced by scaling the amount of data or model
parameters in a Transformer-based language modeling objective. Most
importantly, the authors show a correlation between the training
duration and CBS, both empirically and theoretically. Importantly, the
experiments also show that model size does not have an effect on CBS.
Finally, a secondary but interesting part of the experiments concerns
allowing to change the training duration mid-training. It is shown
that a simple constant schedule combined with exponential weight
averaging at the end of training outperforms both "standard" cosine
decay and the recently proposed warmup-stable-decay (WSD) schedule.

The empirical study also includes a controlled setting to study the
effect of width- vs. depth-scaling on the CBS (the finding being that
both scaling methods influence the CBS similarly). Finally, the
authors additionally underline their empirical findings using
NTK/tensor programs and Gaussian linear regression theory.
weaknesses: The paper heavily relies on Chinchilla scaling theory and
Transformer-decoder training with the language modeling objective. It
does not account for settings which use fewer training steps than
Chinchilla-optimal training and does not study other data sets or
models. This means it is difficult to extrapolate the results to the
general domain of deep learning.

I would personally have preferred use of a more standard "dynamic
horizon" schedule such as WSD for the experiments in order to improve
interpretability/comparability with regard to other research. That
said, the finding regarding constant+EWA quality is important.

Weight decay was disabled for the experiments, which is hardly done in
practice.

Meta-review: The paper explores the scaling behavior of critical batch size (CBS) in autoregressive language model pretraining. CBS is shown to scale primarily with dataset size rather than model size, a finding supported by theoretical analysis. Empirical studies span models from 85M to 1.2B parameters, with controlled evaluations of depth and width scaling. The paper also assesses learning rate schedules, proposing a constant schedule with exponential weight averaging (EWA) as a competitive alternative.
Strengths highlighted by reviewers include the novelty of CBS scaling insights, clear experimental design, and the clarity of presentation through practical takeaways. Criticisms include the limited generalizability due to a limitation to the C4 dataset and relatively small model sizes. Incomplete hyperparameter tuning at larger scales also raised concerns about the validity of some conclusions. Despite these limitations, the paper offers valuable contributions to understanding CBS scaling.

Decision: Accept (Poster)