Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training
Authors: William Merrill
Shane Arora
Dirk Groeneveld
Hannaneh Hajishirzi
Venue: NeurIPS 2025 spotlight ()
Keywords: large language models
pretraining
critical batch size
Primary area: deep_learning
TL;DR: We propose a simple way to measure the critical batch size for language model pretraining that alleviates issues with existing methods, and show that this can be used to train language models with fewer gradient steps in practice.

Abstract: The right batch size is important when training language models at scale: a large batch size is necessary for fast training, but a batch size that is *too large* will harm token efficiency. To navigate this tradeoff, McCandlish et al. (2018) suggest that a *critical batch size* (CBS), below which training will not substantially degrade loss, can be estimated based on the gradient noise scale during training. While their method has been adopted in practice, e.g., when training GPT-3, strong assumptions are required to justify gradient noise as a proxy for the CBS, which makes it unclear whether their approach should be trusted in practice, limiting its applicability. In this paper, we introduce a simple, empirical approach to *directly* measure the CBS and show how the CBS evolves over training. Applying our approach to the OLMo models, we find that CBS is near 0 at initialization, increases rapidly at first, and then plateaus as training progresses. Furthermore, we find that this trend holds across different model sizes (1B and 7B), suggesting CBS from small training runs can inform larger-scale training runs. Our findings about how the CBS changes over training motivate *batch size warmup* as a natural way to reliably train language models at large batch size: start the batch size small and increase it as the CBS grows. To validate this claim, we use batch size warmup to train OLMo 1B to slightly better loss than the original training run with 43% fewer gradient steps. This shows how our framework can be applied to reliably train language models at larger batch sizes, increasing data parallelism without compromising performance.

Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training
William Merrill
Shane Arora
Dirk Groeneveld Allen Institute for Al willm@allenai.c org
Hannaneh Hajishirzi
Abstract
The right batch size is important when training language models at scale: a large batch size necessary for fast training. but batch size that is too large will harm token efficiency. To navigate this tradeoff, McCandlish et al. (2018) suggest that a critical batch size (CBS), below which training will not substantially degrade loss, can be estimated based on the gradient noise scale during training. While their method has been adopted in practice, e.g., when training GPT-3, strong assumptions are required to justify gradient noise proxy for the CBS, which makes it unclear whether their approach should be trusted in practice, limiting its applicability. In this paper, we introduce simple, empirical approach to directly measure the CBS and show how the CBS evolves over training. Applying our approach to the OLMo models, we find that CBS is near 0 at initialization, increases rapidly first, and then plateaus as training progresses. Furthermore, we find that this trend holds across different model sizes (1B and 7B), suggesting CBS from small training runs can inform larger-scale training runs. Our findings about how the CBS changes over training motivate batch size warmup models at large batch size: start the batch size small and increase it as the CBS grows. To validate this claim, we use batch size warmup train OLMo IB slightly better loss than the original training run with 43% fewer gradient steps. This shows how our framework can applied to reliably train language models larger batch sizes, increasing data parallelism without compromising performance.
natural way reliably train language
Introduction
Increasing the throughput of training important for training large models. natural way to increase throughput is by increasing data parallelism, i.c., increasing the batch size used during training so that more data be processed once and the number of sequential gradient steps can decreased. However, naively picking very large batch size can degrade the performance achieved by fixed token budget, as larger batches can show diminishing returns in their ability to estimate the population gradient. Thus, in order to confidently train language models with higher token throughput, it is important to develop theoretical and empirical understanding of large-batch training.
One fundamental concept for large-batch training methodology is the following:
Critical Batch Size Hypothesis (McCandlish et al., 2018)
There is some critical batch size B* up to which increasing the batch size (and appropriately modifying the learning rate) approximately preserves the loss trajectory function of tokens trained, but, above which, the loss trajectory degrades.
such CBS 3 exists (and we can measure it), it represents a reasonable balance between efficiency and performance (i.e., loss) and is thus practically useful batch size at which to train. Working in
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

simplified theoretical setup, McCandlish et al. (2018) derive correspondence between the CBS and the gradient noise scale, i.e., the variance the per-example gradients from the training distribution. They suggest that an estimator for the gradient noise scale should be used in practice as proxy for the CBS, and this can in turn be used to set the batch size for large-scale pretraining runs. The noise scale also appears to have been adopted in practice as proxy for the CBS, having been mentioned explicitly in the GPT-3 technical report (Brown et al., 2020), and inspiring a flurry of methodological innovations for better noise scale estimation (Gray et al., 2023, 2024).
While appealing, the link between the gradient noise and the CBS requires several strong assumptions to justify: specifically, it assumes the SGD optimizer and that gradients are well-conditions (cf, Sec- tion 2). Thus, is unclear whether the noise scale should be a meaningful proxy for the CBS for language model pretraining in practice, where the Adam optimizer is often used and the optimization may not be well-conditioned. To this end, we aim to address the following practical questions that remain for effectively leveraging the CBS viewpoint to train language models at larger batch sizes:
1. How can measure the CBS cheaply with minimal assumptions before launching pretraining run?
2. How does the CBS change over the course of pretraining and as function of model size?
3. Having measured the CBS, how should we adapt the batch size, learning rate, and other parameters over the course of pretraining run?
This work documents our attempts to answer these questions in order to operationalize the CBS for large-batch training. We focus investigation the OLMo models (Groeneveld al., 2024; OLMo et al., 2025), due to their open weights and data, making the following contributions:
1. First, we introduc an empirical method to directly measure the CBS via branched training. Our method avoids strong assumptions needed to justify the noise scale method from prior work. This lets trust it more than the noise scale method, which we find unreliable.
2. We use our method to study how our CBS measurement changes over the course of training, finding it improves rapidly initially but than flattens off. Further, we find that the CBS not depend strongly on model size, in line with past findings using different methodology (Zhang et al., 2024).
3. Our knowledge of the local CBS across training checkpoints suggest a natural batch size warmup strategy for large batch training: begin training with a small batch size and double it whenever the CBS increases sufficiency. We use this strategy to train IB parameter models with 43% fewer gradient steps without degrading (and, in fact, slightly improving) final loss.
Overall, our empirical framework for measuring and leveraging the CBS provides simple and principled methodology for improving the efficiency of large-scale training runs and addressing other fundamental questions in the science of language model pretraining.
2 Background: Estimating CBS via Gradient Noise Scale
Past empirical work has aimed to measure the CBS by launching many training runs to the same target loss, which is expensive (Zhang et al.. 2019, 2024), or by using the gradient noise scale proxy (McCandlish et al., 2018; Gray et al., 2023, 2024). In contrast, we will introduce a new method that uses small amount of additional training to estimate the CBS, which is less expensive than launching many full training runs and does not make any strong assumptions like the noise scale. Before introducing
method, review the noise scale framework used to estimate the CBS in prior work and the underlying assumptions relies on.
McCandlish et al. (2018) suggest that the CBS can be measured in terms of the gradient noise scale, i.e., the variance of the gradients within a batch. Concretely, their recommendations to measure the CBS and adapt the learning rate are follows:

---

Existing Method: Noise Scale Proxy for CBS (McCandlish et al., 2018)
Let G be the full gradient and let Σ be the covariance matrix for the gradient across data examples. We first compute Biimle proxy for the CBS (using an efficient statistical estimator):
tr |G||2
B simple
We set the modified batch size Bsimple and linearly scale the learning rate 17"
B Baimle
B"
" B 7/-
(2)
This viewpoint is attractive due to its simplicity and tractability, and has inspired improved methods for estimating the noise scale (Gray et al.. 2023, 2024). However, the link between noise scale and CBS requires several strong assumptions to justify, which we will argue motivates revisiting other approaches for measuring the CBS. McCandlish et al. (2018) consider training a model on a loss surface where the loss landscape is well approximated by its second-order Taylor expansion. The first crucial assumption in their analysis is that optimizer used to decrease the loss is SGD:
Assumption 1: SGD Optimizer (MeCandlish al., 2018)
The step taken to reduce the loss is noisy estimate of the true gradient, computed via B samples.
This assumption may seem benign, but it is worth noting that it is not typically met in practice, as LMs are trained with the Adam optimizer (Kingma and Ba, 2017). Moreover, by analyzing training dynamics in terms of stochastic differential equations (Li ct al., 2021), Malladi et al. (2022) argue theoretically that Equation (2) is an appropriate scaling rule for SGD, but not for Adam, where a square-root scaling rule is more principled. The square-root scaling for Adam is also supported by the theoretical analysis of Li et al. (2024, Equation 4). Thus, when training with Adam, seems that the linear scaling rule assumed by McCandlish al. (2018) should not apply.
more fundamental issue that McCandlish et al. (2018) also require another strong assumption to derive the noise scale method for estimating the CBS. In general, their noise-scale-based estimate of the CBS has a more complex form than Baimple involving the Hessian H:
tr(EH) G HG
Booise
In order to justify that B" can be computed as Bsimples McCandlish et al. (2018) assume:
Assumption 2: Well-Conditioned Optimization (McCandlish et al., 2018)
The Hessian H is multiple of the identity matrix. It follows that Buoise
Baimple-
This a strong assumption required to justify using simple as proxy for the CBS because computing Hessians would be too expensive to be practice. McCandlish et al. (2018) suggest informally that, without this assumption, Baimlle may still correlated with Broise. but it is not obvious why this should be the case. Even if this is true, it still poses a real problem for the noise scale methodology. since the goal of the method is to produce an absolute measure of B". It unclear for practitioners what coefficient should be used to translate Bsimple to B²-and more fundamentally, whether it is even valid to assume that such a coefficient exists.
3 Our Method: Measuring the CBS via Local Branched Training
As discussed in Section 2, using the gradient noise scale Bample as a proxy to estimate the CBS relies on several strong assumptions. In light of this, it unclear whether we can trust the gradient noise scale as proxy for CBS We thus argue that we should instead aim to measure the CBS directly, without the need for any strong assumptions to justify an indirect proxy. After introducing measurement
3

---

OLMo 208
OLMo 4198
OLMO 6238
255
254
2.75
2.15
2:10
245
302
2005
13th
4000 1022
1999 300h
2763 V10 w00 790
3991
27.0
0091
2703 1.03 wiz 750 KIE
Batch
Balthh
526
Figure 1: Smoothed final loss after branched training at particular checkpoints, with B" shown as the dotted red line. Each point represents the loss achieved by a specific branched training run after 2B tokens. Our method detects the point which loss starts increase, heuristically tolerating noise within e = 0.01. These plots show how this plays out for three particular checkpoints; see Appendix A for loss curves for all checkpoints.
approach in this section, we will show in Section 4 how these measurements can be applied to train language models the same (or better) target loss with fewer gradients steps.
3.1 Method
We introduce simple branched training method that directly approximates the CBS by launched branched training runs from a checkpoint, which allows us to identify 8 as the largest batch size that does not degrade in loss relative smaller batch sizes as visualized in Figure 2. To make this tractable, we train only for a fixed token budger A, assuming that if 3 recovers in loss by A, its loss will continue to match smaller batch sizes onwards as well. This allows to estimate the CBS with only small amount of additional training (controlled by A). Further, as we will find later in Figure 2, the CBS trend remains consistent across model sizes, CBS measurements with small models could be used to inform large-scale training runs.
Our Method: Branched Training to Measure CBS
Given a training checkpoint with original batch size B and Icarning rate schedule 77. aim to measure the critical batch size B*. Let f(n) be the learning rate scaling rule: f(k) = & for SGD and vk for Adam. We create several training branches with modified batch size k B and learning rate f(k) 1I and train for small number of tokens A to get loss follollowin standard practice, we take La to be the smoothed loss. We then define k- as the maximum such that, for all l < k", Lg- T. e, where is tolerance parameter for "similar" losses. We then define the CBS B* and scaled learning rate 17" as
B" k- 17" = f(k))
Our method will empirically estimate the largest batch size B at which the local optimization trajectory recovers roughly to its original loss after A steps. In contrast to the strong assumptions needed to justify the gradient noise scale, the only (weak) assumption our method relies on is:
Assumption 3: Local Recovery
If the loss achieved by batch size B" recovers to match the loss with batch size B B" after training for A tokens, the loss trajectories will remain the same beyond A as well.
Implementation Details. Our method requires specifying two parameters: the window size and the loss tolerance We also apply smoothing to the loss to reduce noise. In more detail:
Window Size. Because the optimizer state must update when the batch size is changed, we expect an immediate spike in the loss after adjusting the batch size. The window size A represents the number of steps we are willing to wait for the loss recover from this bump. The CBS measurements could, in principle, depend on A, with larger values of A

---

CBS over OLMo 1B pretraining
CBS over OLMo 7B pretraining
8192
Empirical CBS
2048
1024
Empirical
1024
512
256
256
128
64
64
400
600
B0o
250
500
750
1000
1250 1500 (in billions).
1750 2000
billions)
Figure CBS over training for OLMo IB and 7B, measured in documents (4096 tokens per document). The qualitative trend is similar across both model sizes. The CBS starts near 0, grows rapidly but diminishingly, and plateaus around 4096.
potentially producing larger CBS estimates. We set A to 2B tokens, which we take small, conversative window size relative to our overall pretraining budget 600B tokens. 2. Loss Tolerance. Viewing loss as a function of batch size multiplier k Figure 1), we need a way to determine whether the loss at k" has increased relative to all k < kt. We operationalize this with tolerance parameter e, which we to 0.01 arbitrarily. In principle, tolerance could be set in a more principled way using a statistical test in future work.
3. Loss Smoothing. Pretraining loss is noisy the batch level, so, following standard practice, we apply exponentially moving average smoothing with parameter a set to 0.5.
It is worth comparing our method to other work that empirically measures the CBS. Our method of training for fixed token budget A and measuring the change in loss can be understood as the dual view measuring the number of steps required to achieve a target loss, which has been used in prior work (Zhang et al., 2019, 2024). Reformulating the measurement this way has the nice property that the training budget can be fixed in advance. Under the local recovery assumption, it also allows us to train for only A tokens, which means we do not have to launch full training runs for each batch apply our method at various checkpoints throughout
size. Finally, unlike Zhang et al. (2024), training, whereas they apply it only from initialization. This means we can estimate the local CBS at specific point in training rather than just the global CBS.
3.2 Experimental Setup
We aim to measure the CBS over the course of model training and the role of model size, Because method requires pretraining checkpoints and access the pretraining data, we the OLMo 1B and OLMo 7B models for our experiments (OLMo et al., 2025), whose pretraining data is openly available (Soldaini et al., 2024). For each model, take variety of checkpoints training, allowing us to assess how the CBS changes over the course of training; see Appendix A for more details. We also compute the noise scale across training checkpoints using the estimator proposed by McCandlish et al. (2018) to assess whether it is valid proxy for the CBS we measure.
the course of
We define the interval for the CBS each checkpoint by choosing B" (as defined in Section 3) lower bound, and by picking the least >> k- as an upper bound. The plotted point represents the geometric mean of these interval endpoints. We measure the CBS in documents, with pretraining sequence length of 4096 tokens per document.
3.3 Results: CBS Over Training and Across Model Sizes
Figure 2 shows the CBS B* measured via Section 3.1 across training checkpoints for OLMo IB and 7B. The CBS increases over training in a similar way for both model sizes: the CBS starts near 0, grows rapidly within the first 50k tokens, and then plateaus around 4096.
Impact of Model and Data Size. Prior work has suggested that the CBS largely independent of model size, scaling primarily with data size (Zhang et al., 2024; Bergsma et al., 2025). This is largely
5

---

Gradient nose scale over OLMo 18 pretraining
Gradient noise scale over OLMo 78 pretraining
8192
2048
Gradient noise scale
Gradient naise scale
200 400
600
U00
250
750
10000
1250
1500
1750
billions)
(in billions).
Figure 3: Gradient noise scale for OLMo IB and 7B computed via the estimator McCandlish et al. (2018) with 95% confidence intervals; details in Appendix B. The gradient noise scale underestimates the CBS (cf. Figure 2) and the qualitative trend does not clearly match, especially for OLMo 7B.
consistent with our findings, as the CBS curves like qualitatively similar at the 1B and 7B scales. In addition, the fact that the CBS grows over training suggests that the "aggregate" CBS should also increase train on more data since the average CBS over the course of training should increase. We elaborate on this in Appendix D: while CBS growth over training predicts that aggregate CBS should increase with data size, is unclear whether the CBS growth pattern we observe would predict the aggregate CBS & vD scaling law found in prior work (Zhang al., 2024; Bergsma al., 2025).
Comparison with Gradient Noise Scale. As discussed in Section 2. the gradient noise scale has been proposed as proxy to measure the CBS (McCandlish et al., 2018 8), though this connection relies on strong assumptions to justify. We thus empirically compare our mcasurement of the CBS to the gradient noise scale. Figure shows that, for both OLMo 1B and 7B, the gradient noise scale underestimates the CBS by several orders of magnitude. Furthermore, especially for OLMo 7B, the qualitative trend does not match the CBS. For OLMo 1B, the qualitative pattern is more similar. However, since this similarity is not found for both models, we conclude that, in general, the noise scale cannot be used reliably
proxy for the CBS.
Motivating Batch Size Warmup. A central takeaway from these results is that the CBS starts near 0, grows rapidly, and then plateaus. This suggests that batch size warmup, where the batch size is dynamically increased over the beginming of a training run, is a natural way to increase the effective batch size for most of training while avoiding training with batch size above the CBS. In the next section, we will discuss our implementation and validation of this idea.
Application: Batch Size Warmup for Larger Batch Training
A straightforward way to speed up training is to increase the batch size, and, as long as the new batch size less than the CBS, we can be confident that loss will be minimized about as effectively as before. Thus, we can leverage knowledge of the CBS to speed up training.
Furthermore, our local knowledge of the CBS can be leveraged to do this more effectively than if we only had a global sense of the CBS. If we simply increased a fixed batch size, this would mean that we are training with a batch size that too large for short period at the beginning of training, which could potentially destabilize training degrade final performance. To get around this, we can use our measurement of the local CBS to "warm up" the batch size. We will train the beginning of training when the CBS is small, and then switch to a larger batch size once the CBS grows large enough. More generally, we can aim to double the batch size whenever we determine the CBS has doubled. This should allow to benefit from training at a larger batch size for most of training without training at batch size that is too large at the beginning of training, which could degrade the final loss achieved by the training
smaller batch size for

---

4.1 Methodology for Batch Size Warmup
Given an existing training run (at a small batch size), we aim to adapt to achieve the the same final loss with fewer overall gradient steps (i.e., a larger batch size for most of training). Assuming an original fixed batch size B and base learning rate 7) our method for training with batch size warmup is as follows. First, use branched training CBS method (Section 3) to measure B the CBS after training for t tokens, using an existing checkpoint. When training a new model, we set the batch size and and learning rate follows:
Batch Size Warmup
1. At 0, initialize the batch size to Bo B and base learning rate to Y/o '7.
2. After training for tokens, if we determine that the CBS exceeds the current batch size (B; > 2B),, we double the current batch size and update the base learning rate following the square-root scaling rule (ef. Section 3):
Bi+1 2B,
Th+1 = v 21t-
This method will ensure that the batch size will increase over training safely, i.e., never exceeding the CBS. Since we found that the CBS increases quite rapidly at the beginning of training, this method will double the batch size twice early in training, reaching maximum batch size of 4096 by 503B tokens. This means that we will effectively train at a larger batch size for most of training compared to the original small-batch run. But, crucially, we can be confident that the batch size will never increase our batch size above the CBS. Thus, we expect the final loss will be comparable to that of the original training across thet training trajectory.
Implementation Details. The main implementation question is how to operationalize checking B > 2B, in order to double the batch size. In practice, we do this heuristically based the measurements in Figure 2: since we only double the batch size twice over training, this involves just picking two thresholds. In addition, this process could be automated for future training runs using online measurements of the CBS paired with some kind of curve fitting or statistical test.
Design Choices. As defined above, our method uses square-root scaling rule, which is well- motivated for Adam (Malladi et al., 2022), but in principle a linear scaling rule could also be used. As batch size warmup only modifies the base learning rate, it is compatible to overlay with existing learning rate schedules (in practice, the models we use follow a cosine schedule). Finally, we choose to only increase the batch size powers of two because it is not clear that having a more precise match to the CBS beyond the ander-of-magnitude level would be particularly useful and because the OLMo codebase requires the number of GPUs g divides the batch size. Thus, doubling is convenient because it easily guarantees the new batch size remains a multiple of g. However, in principle one could update more frequently, e.g., by setting updating the batch size the largest multiple of g S Bi-
Connection to Existing Methods. Our batch size warmup method is similar to preliminary experi- ments by McCandlish et al. (2018, Appendix D) using a dynamic batch size on the SVHN dataset. However, they use their noise scale method (Section 2) to set the batch size, which is potentially unreliable, and only consider small-scale classification. In the context of image classification, prior work explored replacing learning rate decay with increasing the batch size (Smith et al., 2018). This is conceptually related our batch size warmup approach in that it leverages scaling rules between batch size and learning rate to achieve larger batch training. However, our method can apply on top of existing learning rate schedules, can generalize to Adam (vs. SGD), and, most crucially, ensures that the batch size number exceeds the critical batch size.
4.2 Experimental Setup
We evaluate the viability of training with batch size warmup via the following training runs. For all models, we use the default OLMo IB pretraining hyperparameters unless stated otherwise.
Batch Size Warmup: We train OLMo IB with our batch size warmup method as detailed in Section 4.1. We initialize the batch size to 1024. Based on manual reading of the CBS
7

---

Training with different batch size schedules
4096
2.90
Small-Bale Control
Batch Size
2.85
Large-Batch Control
Batch Size Warmup
2048
Training Loss
2.75
2.70
2.60
Learning Rate
0,0010
2.65
2.58
2.60
2.56
0.0005
2.55
2.54
500
600
0.0000
2.50
200
200
400
600
Tokens (in billions)
Tokens billions)
Figure 4: Batch size schedule (left, top), learning rate schedule (left, bottom) and training loss (right) for the pretraining of an OLMo model with different batch size schedules. Training loss is smoothed by taking the moving average the past 10B tokens.
Method
PT Loss
1 MT Loss
Grad. Steps Saved
Batch Size Warmup (Ours)
2.5891
2.5433
43%
Small-Batch Control
2.6057
2.5486
0%
Large-Batch Control
2.5962
2.5506
75 %
Table Loss after pretraining (PT) and mid-training (MT) for each run, averaged over the past 10B tokens and the percentage of gradient steps saved the small-batch control (including annealing steps). Both before and after annealing, batch size warmup slightly outperforms both controls in loss.
measurements in Section 3, we determine that the CBS reaches 2048 by 168B tokens and 4096 by 503B tokens.! Our method thus doubles the batch size each of these points.
2. Small-Batch Control: We train OLMo IB with batch size B 1024 and base learning rate 17 = /2- 0.0004.2 We expect that our batch size warmup method should be able to achieve similar final loss to the small-batch control with fewer gradient steps.
3. Large-Batch Control: We train OLMo 1B with fixed large batch size of B 4096 and base learning rate 7) 2 0.0004. Since its batch size will exceed the critical batch size for the initial part of training, we expect that the large-batch control will show degraded loss compared to our batch size warmup method.
Evaluations. Language model training consists of two steps, pre-training and mid- training (OLMo et al.. 2025). We evaluate training loss as well as several measures of -itt-----dsributiin loss at the end of the pre-training stage as well as after the mid-training stage.
Loss after Pretraining. We evaluate the loss at the end of the pre-training stage for all models. The original training run for OLMo IB ran for 4T tokens, we do not have the
'These thresholds were determined from preliminary CBS measurements Our measurements changed slightly after point, but we deemed that was not worth it restart the expensive training runs because we do not think the results should he sensitive to a precise choice of threshold. Going forward. it could be useful to establish systematic methodology for choosing batch size warmup thresholds given the CBS measurements.
earlier stage of the project.
"The default hy hyperaarameters for OLMo IB use a batch size of 512 and learning rate 0.0004. Since we chose initial base size of 1024 for our experiments, we adapted the learning rate appropriately via the square-rool root scaling rule (Malladi et al., 2022).

---

Method
J Task BPB
C4
Pile
PT
MT
PT
MT
PT
MT
Batch Size Warmup (Ours)
1.0316
1.0076
2.8049
2.7597
2.1916
2.1521
Small-Batch Control
1.0112
0.9999
2.8196
2.7622
2.2073
2.1471
LLarge-Batch Control
1.0571
1.01927
2.8107
2.7658
2.1996
2.1586
Table 2:
According to the method from Bhagia et al. (2024), we evaluate downstream performance via BPB on downstream tasks, well as cross-entropy loss two held-out sets, C4 and the Pile, both after pretraining and after mid-training. Batch-size warmup generally performs comparably or better compared to the small-batch control, suggesting does not degrade downstream performance.
resources to replicate this full training run. Instead, we pre-train for 608B tokens using the original learning rate schedule for the longer run.
Loss after -ri-nrraniig. We report loss at the end of mid-training stage, which more closely reflects how these checkpoints
used in language model training. Specifically, starting from final pretraining checkpoint, we linearly anneal the remaining learning rate down to 0 for 50B tokens, keeping the batch size fixed at its final value (Figure 1). Because mid-training is a part of the standard language modeling pipeline, we take the loss after mid-training to represent the loss that would be achieved by these training runs in a practical context. Furthermore, OLMo et al. (2025) suggest that this kind of learning rate annealing can induce significant gains in loss for partial pretraining runs. We thus also take the loss after mrid-rrining loss as a proxy for how these runs would compare if we were to fully train them for the full learning rate schedule.
Out-of-Distribun Distributtinn Losses. To measure performance beyond the training loss, we also track three kinds of out-of-distribution losses. The first two are straightforward cross-entropy losses on common pretraining datasets, C4 (Dodge et al., 2021) and The Pile (Gao et al., 2020). The third -ut-of-distribution loss follows the method from Bhagia et al. (2024), which argues for computing the loss (in bits-per-byte; BPB) on the correct answers of multiple popular question-answ the full list of datasets used, see Appendix E.
datasets such as ARC-Easy, ARC-Challenge, MMLU, etc. For
4.3 Batch Size Warmup Results
We compare batch size warmup to the small and large-batch controls at two points: first, after the pretraining stage and, second, after the mid-training stage that anneals the remaining learning rate down O. After pretraining, batch size warmup reaches lower loss than both controls, as shown in Figure 2 and Table 1, outperforming the small-batch control in loss by a margin of 0,0166. After mid-training batch size warmup still slightly outperforms the small-batch control in loss, this time by smaller margin of 0.0053. The large-batch control now achieves the worst overall loss. Overall, batch size warmup slightly exceeds the small-batch control in loss while using 43% fewer gradient steps. In contrast, the large-batch control uses 75% fewer gradient steps, but its final loss degrades compared to the small-batch control. Thus, we find batch size warmup is reliable method to train with fewer gradient steps without degrading final loss.
Beyond the final loss, also evaluate whether batch size warmup leads to similar downstream performance as small-batch training using our downstream-task BPB evaluation as well as validation loss on C4 and the Pile. As shown in Table 2, we find that BPB on validation sets is broadly competitive with the small-batch control, performing best on three out of four conditions considered. With the caveat that precisely measuring the downstream impact of pretraining decisions is difficult, this suggests that, just as batch-size warmup does not degrade final loss, it should not degrade downstream measures of performance either.
Conclusion
In this work, we introduced simple empirical method for estimating the CBS throughout language model pretraining runs, which can be used to increase batch size (and thus effective token throughput) for large scale training runs without sacrificing performance. As we discussed, the existing noise

---

scale method (McCandlish et al., 2018) for estimating the CBS requires strong assumptions to justify, which our approach avoids. We used our method to study the evolution of the CBS during training for the OLMo models, finding that CBS increases monotonically but diminishing over the course of training, and that CBS does not seem to depend on model size, in line with prior work (Zhang et al., 2024). Guided by these findings, we showed that our measurements could be used to pick a batch size warmup schedule that enables larger batch training without harming final training loss. We take these results to demonstrate the validity and utility of our CBS measurement approach, and believe our framework could be useful for increasing the efficiency of future large-scale pretraining efforts.
There are several details and extensions of our method that would be interesting to explore going forward. First, would be interesting to carry out more systematic analysis of the impact of the hyperparamettr A on CBS measurements: how sensitive is the method to A and do different values for A potentially bias the measurement? It would also be interesting to further investigate the conditions under which noise scale might be meaningful CBS proxy and used for batch size warmup. When applying our CBS measurements to picking a batch size warmup schedule, we manually picked doubling threshold based the CBS. Going forward, it would be useful to establish systematic way to set threshold for increasing the batch size given CBS measurements. There are also other potential methodological improvements, such as removing the arbitrary power of 2 constraint and estimating the CBS in an online fashion. Overall, these improvements would allow batch size warmup to be applied more robustly and easily across different pretraining setups.
Acknowledgments
We thank Joel Hestness, Sadhika Malladi, and Ananya Harsh Jha for discussions.
References
J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, and C. Sutton. Program synthesis with large language models. arXiv preprint urXiv:2108.07732 2021.
S. Bergsma, N. Dey. G. Gosal, G. Gray, D. Soboleva, and J. Hestness. Power lines: Scaling laws for weight decay and batch size in Ilm pre-training, 2025. URL https://arxiv: orw...////........ 13738.
A. Bhagia, J. Liu, A. Wettig, D. Heineman, O. Tafjord, A. Jha, L. Soldaini, N. A. Smith, D. Groen- eveld, P. W. Koh, J. Dodge, and H. Hajishirzi. Establishing task scaling laws via compute-efficient model ladders. ArXiv, abs/2412.004003, 2024. URL https: //api. sewanticscholesttt... org/
CorpusID: 274514987.
Y. Bisk, R. Zellers, R. Le bras, J. Gao, and Y. Choi. PIQA: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):7432- 7439, Apt. 2020. doi: 10. 16099/aai.v34000 6.339, URL https: /ojs../j....../ /ta./aaaaa.aa.aaaaa/.aa...a..... awai.org///aaaaaaa.......a.....a index.ndxxx// php/ wataaiieeeiiiieeeeeeeeeeeeeeeeee
T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert- Voss. G. Krueger, Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, Sutskever, and Amodei. Language models are few-shot learners, 2020. URL https: / /arxiv.vvvv targ/abs/2000 14165.
M. Chen, J. Tworek, H. Jun, Q. Yuan, H. Ponde de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. urXiv preprint arXiv:2107.03374. 2021.
Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? Try ARC, the A12 reasoning challenge. CoRR, arXiv:1803.00 803405557, 2018.
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv: 2021.
10

---

J. Dodge, A. Marasovic, G. Ilharco, D. Groeneveld. M. Mitchell, M. Gardner, and W. Agnew. Documenting large webtext corpora: case study on the colossal clean crawled corpus. In Conference on Empirical Methods in Natural Language Processing, 2021. URL https: /tttpps:/app/ //ppi. semanticscholar org/ wwwgg/orousDD:... [CorpusID22375687722
L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The pile: An 800gb dataset of diverse text for lan- guage modeling, ArXiv, abs/2101.00022 2020. URL https: //api.
org/
CorpusID:230435733
G. Gray, A. Samar, and J. Hestness. Efficient and approximate per-example gradient norms for gradient noise scale. In Workshop Advancing Neural Network Training: Computational Efficciecy, Scalability, and Resource Optimization WWAN@NNerri 2023), 2023. URL https: //openreview nwt/feruunr wwwprrdddxnITMaPPPPPP
G. Gray, A. Tiwari, Bergsma, and J. Hestness. Normalization layer per-example gradients are sufficient predict gradient noise scale in transformers. In The ehirty- Anmual Conference on Neural Information Processing Systems, 2024. URL https: /tapps://pppt http://ppeeeeeeeeeeeeeee www.o/rmmteeeooometeemm ffrrmmrddd 87TH1pvH8i.
D. Groeneveld, Beltagy, E. Walsh, A. Bhagia, R. Kinney, O. Tafjord, A. Jha, Ivison, I. Mag- nusson, Y. Wang, S. Arora, D. Atkinson, R. Authur, K. Chandu, A. Cohan, J. Dumas, Y. Elazar, Y. Gu, J. Hessel, T. Khot, W. Merrill, Morrison, N. Muennighoff, A. Naik, C. Nam, M. Pe- ters, V. Pyatkin, A. Ravichander, D. Schwenk, S. Shah, W. Smith, E. Strubell, N. Subramani, M. Wortsman, P. Dasigi, N. Lambert, K. Richardson, L. Zettlemoyer, J. Dodge, K. Lo, L. Soldaini, N. Smith, and H. Hajishirzi. OLMo: Accelerating the science of language models. In L.-W. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15789-15809, Bangkok, Thailand, Aug. 2024. Association for Computational Linguistics. doi: 10. cettps:/acclatthlle
URL
org/2024 acl-long. 841/.
D. Hendrycks, C. Burns, Basart, Zou, M. Mazeika, Song, and Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICIR), 2021.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2017. URL https://arxiv. org/abs///et ht//a////111..999
A. Lewkowycz, A. Andreassen, D. Dohan, E. Dyer, H. Michalewski, Ramasesh, A. Slone, C. Anil, Schlag, T. Gutman-Solo, al. Solving quantitative reasoning problems with language models. arXiv preprint arXiv:2206. 10066/485.. /4858, 2022.
S. Li, P. Zhao, H. Zhang, S. Sun, H. Wu, D. Jiao, W. Wang, C. Liu, Z. Fang, J. Xue, Y. Tao, B. CUI, and Wang- Surge phenomenon in optimal learning rate and batch size scaling. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: /ppprr//eepeereieeeeee
Z. Li. S. Malladi, and S. Arora. On the validity of modeling sgd with stochastic differential equations (sdes), 2021. URL https: aaaxiviiiiv../..///eaa/eaaaa..aa /tt///aivv....aaaar/ee/reeeeraee abs//a//...22.//////// 12470.
I. Magnusson, N. Tai, B. Bogin, D. Heineman, J. D. Hwang, L. Soldaini, A. Bhagia, J. Liu, D. Groen- eveld, O. Tafjord, N. A. Smith, P. W. Koh, and J. Dodge. Datadecide: How predict best pretrain- ing data with small experiments. 2025. httpst //api. semanticschol www.cachhooa....ooooooooeeeeoeee Corpus 277787795.
S. Malladi, K. Lyu, A. Panigrahi, and S. Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL. https: http///ppppeeeeeeeeeeeeeeeeeeeee //ppprrrrevv..t//toommt........ F2mhzjHHkQP.
McCandlish, J. Kaplan, Amodei, and O. Team. An empirical model of large-batch training, 2018. URL https: arxiv. org/ babs///81..000 06162.

---

T. OLMo, P. Walsh, L. Soldaini, D. Groeneveld, K. Lo, S. Arora, A. Bhagia, Y. Gu, S. Huang, M. Jordan, N. Lambert, D. Schwenk. O. Tafjord, Anderson, D. Atkinson, F. Brahman, C. Clark, P. Dasigi, N. Dziri, M. Guerquin, H. Ivison, P. W. Koh, J. Liu, S. Malik, W. Merrill, J. V. Miranda, Morrison, T. Murray, C. Nam, V. Pyatkin, A. Rangapur, M. Schmitz, S. Skjonsherg. D. Wadden, C. Wilhelm, M. Wilson, L. Zettlemoyer, A. Farhadi, N. A. Smith, and H. Hajishirzi. 2 olmo 2 furious, 2025. URL https:
wwtp:///....bbbbbbbbb////// 00656
K. Sakaguchi, R. Le Bras, C. Bhagavatula, and Y. Choi. WinoGrande: An adversarial winograd schema challenge scale. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05): 8732-8740, Apr. 2020. doi: 10. 10.1609/aaai. vaimaii.34005.6399 URL /tatpp://ojs.. aaai. www.orrginino /igdinddx.... php/
M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi. Social IQa: Commonsense reasoning about social interactions. In K. Inui, J. Jiang, V. Ng, and X. Wan, editors. Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4463-4473, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: 10. 10.18653/11/D19 1454. URL https: acl aa//a/loconn org/ Drgg/19--a 1454.
S. L. Smith, P.-. Kindermans, and Q. V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018. URL https / //ppn/eeieeeeeeeeeeeeeeemeeeeeee net/ forum?id-B1yy 1BCCZ.
L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, Dumas, Y. Elazar, V. Hofmann, A. H. Jha, S. Kumar, L.. Lucy, X. Lyu, N. Lambert, I. Magnusson, J. Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richardson, Z. Shen, E. Strubell, N. Subramani, O. Tafjord, P. Walsh, L. Zettlemoyer, N. A. Smith, H. Hajishirzi, Beltagy, D. Groeneveld, Dodge, and K. Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research, 2024. URL https: //arxiv.//////////a org/ abs/ 222.2.00001...
A. Talmor, J. Herzig, N. Lourie, and J. Berant. CommonsenseQA: A question answering challenge targeting commonsense knowledge. In J. Burstein, C. Doran, and T. Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4149- 4 4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10.11853/v1/1119144 1421, URL https: walaaaaanntttggy.oooooo
S. Wiegreffe, O. Tafjord, Y. Belinkov, H. Hajishirzi, and A. Sabharwal. Answer, assemble, ace: Understanding how Ims answer multiple choice questions. In International Conference Learning Representations, 2024. URL https: //ap:::://app.. semanticscholar org/ :orppuus:D:22776003222 276903925.
R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can a machine really finish your sentence? In A. Korhonen, D. Traum, and L. Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791 Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.19655/1//119-14722 URL https: // //ac//accceee anthology org/P19- PP///1991101 1472.
G. Zhang, L. Li, Z. Nado, J. Martens, S. Sachdeva, G. Dahl, C. Shallue, and R. B. Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, ed- itors, Advances in Neurul Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https:// / neurips. wtt.cccpppaaa wwtpcccpppppp.eeeeeee sateee/fflleepappe/eaaaae 60aadd983775342277111111aa832223 Papa4r.appppppe pdf.
H. Zhang, D. Morwani, N. Vyas, J. Wu, D. Zou, U. Ghai, D. Foster, and S. Kakade, How does critical batch size scale in pre- trainiann, 2024. URL https: orw.vvorrggaeee http://rxwi..rr/a////////// 2241/ 21676.
12

---

A CBS Measurement Details
Our empirical method for measuring the CBS Section 3 is, in principle, sensitive to the choice of checkpoints and batch size multipliers. We therefore document the checkpoints and multipliers used here.
OLMo IB. When measuring the OLMo 1B CBS with branching, we set the base batch size to 1024 and the base learning rate to 0.0004 v2, reflecting the default batch size of 512 and learning rate of 0.0004 in the OLMo codebase under square-root scaling rule (Malladi et al., 2022). We then chose the following checkpoints and multipliers k:
Step O: & ranging over 0.0625, 0.125, 0.25.
2. Steps 10K, 20K, ..... 50K: ranging over 0.5, I.....
3. Steps 100K, 150K, .... 450K: k ranging over I, 2, ...,
Figure shows loss batch size plots for all checkpoints of OLMo 1B.
OLMo 7B. We set the base batch size to 1024 and the base learning rate to 0.0003, as specified in the OLMo codebase. We then chose the following checkpoints and multipliers k:
1. Step O: k ranging over 0.0625, 0.125, 0.25.
2. Steps IK, 2K, 3K: ranging over 0.25, 0.5, 1, 2, 3, 4.
3.
Steps 1OK, 20K, 30K: k ranging over 1, 2, 3, 4, 5.
4. Steps 72K, 150K, 200K, 239K, 300K, 350K, 400K, 477K: k ranging over 1,2,3,4,5,, 7, 8.
Appendix A shows loss vs. batch size plots for all checkpoints of OLMo 7B. These checkpoints were chosen manually as we developed this project. Over the course of our experimentation, we launched many additional runs beyond the ones discussed above. Since the choice of k can influence the conclusions of our method, we filtered down the included runs to make the choice of & systematic.
IB runs were launched on single node of H100 GPUs, and 7B runs were launched on 8 nodes.
B Noise Scale Measurement Details
We use the gradient noise scale estimator proposed by McCandlish et al. (2018, Appendix A) to estimate the gradient noise scale. The method estimates the gradient noise scale using gradient norms at two different batch sizes Buig and Bamall according to:
S
Bsimple ~
where
|
S |Gsmall 1| ||
1/Bamall 1/ Boig
| Buig [ 1| Bsmall Gsmall |
Brig Bamall
We large batch size Bbig 64 and small batch size Bamall = 1.
It holds that E [S] tr(E) and E ||] || We thus average S and | over 4096 batches reduce variance and then return their ratio as our estimate of the noise scale Biimlle. using offline (i.e., unseen) data in each batch.
We estimate a confidence interval for Baimple in two steps. First, we estimate 95% confidence intervals for S and |||| 22 assuming the data are exponentially* and normally distributed, respectfully, based
Forr the exponential distribution, we use approximate contidence interval under "Confidence Intervals" here: https://en. wikipedia, orgooooi/ae....o.........oo..... wogg/wiki//gppaee viki:/Eponenttall http://pooontttaaaaaaaarrrrtoooa distribution.
13

---

425
3.4
Lo05
13
Batch Size
Datchh Size
Batch Size
tokenns
tokens
1235
237
314
2.4
Batthh Size
Batch Size
Batthh Size
18,
18, tokens
2.50
3.4
2.5
2.00
2.5
238
2.54
235
1,05
2.5
2.4
3.75
3.1
0.2
245
312
245
Batch Size
(Batch Size
Batoh Size
18,
tokens
2.5
2.00
2.75
243
2.00
Batthh Size
Batthh Size
Batch Size
18,
2.73
2.5
245
Batch Size
Batch S1ze
Figure 5: All loss vs. batch size plots for OLMo IB. Overall, the red line moves to the right over time, showing that the CBS increases.
14

---

78,
4.5
L205
37-
396
2900 vhe
Batthh Size
Batthh Size
Batch Size
tokenns
4.7
332
2.99
2.1
Lo055
2.14
235
216
240
190
voc 4891
Batch Size
Batch Size
Batthh Size
78,
78, tokens
2.52
2.45
232
2.48
2.5
244
237
242
2.4
Batthh Size
(attch Size
Batton Size
2.44
2.55
224
Batch Size
Batch Size
Batch Size
78,
78, tokens
2350
2.8
2399
2309
2.91
Batch Size
Batch Size
Batch Size
15

---

150
160
140
120
100
100
40
-2.04
-0.02
0.50
0.02
Variance (OLMp 18, step 0)
Grad
(OLMo 18, step
Figure 6: Representative histograms for S and || showing data from the 1st to 99th percentiles. The distribution for S is positive, leading us to use exponential distribution, while the fact that some samples of |||| are negative motivates normal distribution.
on manual inspection of their distributions (ef. Figure 6). We denote these intervals (as, bs] and [d1191,, b gra), respectively. We then define the confidence interval for Bsimple as follows:
as
bs
bigl!
our estimates for or |/||² (or their lower or upper bounds) come out negative, we consider them to be O.
The checkpoints considered for OLMo 1B are steps 10K, 20K, 40K, .... 100K, 200K, ... 400K.
For OLMo 7B, we use checkpoints at steps O, 10K, ... 40K, 60K, 70K, 100K, 200K, 400K. The noise scale experiment for each checkpoint (for both the IB and 7B models) was launched on single GPU.
C License Information
The OLMo models (Groeneveld et al., 2024; OLMo al., 2025) and pretraining code, which we use, are released under Apache-2.0 license. C4 (Dodge et al., 2021) is released under ODC-BY license. The Pile (Gao al., 2020) released under MIT license.
D Deriving CBS Scaling Laws: An Attempt
In this section, we explore whether our empirical fits for the critical batch size over training can be used derive scaling laws for aggregate critical batch size that have been derived in prior work. These scaling laws assume we want to use fixed batch size B over training, and then train many different models to the same target loss. They then measure the critical B* up to which increases in batch size do not diminish token efficiency. The standard finding from such work is that CBS grows VT, where T is the total training budget in tokens. This is consistent with our finding that CBS increases over the course of moemiingmmrroverrr we now seek to analyze whether this scaling law can be derived from our empirical measurements of local CBS. If so, this would provide converging evidence and simpler method for fitting CBS scaling laws that only requires training single model.
To begin, we assume that the goal of picking a fixed batch size B is to minimize the L2 distance to the local CBS over the course of training. is not obvious that minimizing L2 distance the right way to pick the fixed CBS: for instance, we might want to weight training at batch size above the local CBS more negatively than training below Regardless, we will proceed for now under the assumption that this is the right perspective. We also make the weaker assumption that f(t) 0, in line with our empirical findings (Section 3). follows that the best batch size to train at (i.e., fixed CBS) is simply the average local CBS over training:
16

---

Proposition 1. Let ft) be integrable with f(0) = and define
R₂ S (B - ((1)² at.
Then R2 is minimized by B" 7 oo sttddt.
Proof. We can first simplify the expression for (R2)²:
T (12)) = (B-) dt
{' (B² 2Bf((t f(t)²) ((t)))dt dt
/
B I'"
- ( dt.
Now, taking the derivative with respect to B, we
-2 S[ f(t)dt.
(R2) =B3T dB
Note that the second derivative 27' is positive. Thus, setting the derivative to 0 and solving for B, we conclude that the following value of B minimizes Rg:
1 f f(t)dt. B T
Thus, under the assumptions we have made, if we are trying to pick fixed batch size that best approximates the local CBS throughout training, we can simply pick the average CBS over training. We can use Proposition to derive scaling law for the fixed B" as function of the final CBS or, equivalently, the total steps We now consider various reasonable functional forms f(t) for the CBS.
D.1
Power Law CBS Scaling
We first consider the prediction for the fixed CBS scaling law the local CBS evolves as a power law. Proposition 2 (B* for power-law CBS), Let f(t) = t for > 0. Then the fixed CBS is
B"
Tec
C
Proof. Plug in and solve the integral:
1 T
B
dd
T J
1
T
c c
Тe
c+1
In the case where e = 1/2 (square root), B* proposed by prior work (Zhang et al., 2024). 3
2 Br.
VT,, which derives the /T scaling law
17

---

task
split
# shots
reference
ARC-Challenge
Test
5
(Clark et al., 2018)
ARC-Easy
Test
(Clark et al., 2018)
CommonsenseQA
Val
5
(Talmor et al., 2019)
HellaSwag
Val
5
(Zellers et al., 2019)
MMLU
Val and Test
5
(Henddycks et al., 2021)
PIQA
Val
5
(Bisk et al., 2020)
Social IQa
Val
(Sap et al., 2019)
WinoGrande
Val
5
(Sakaguchi et al., 2020)
GSM8K
Gold
5
(Cobbe et al., 2021)
Minerva
Gold
0
(Lewkowycz al., 2022)
Humaneval
Gold
O
(Chen al., 2021)
MBPP
Gold
(Austin et al., 2021)
Copycolors 10-way
0
(Wiegreffe et al., 2024)
D.2 Logarithmic CBS Scaling
Proposition 3 (B" for log CBS), Let f(1) = log(t + 1). Then the fixed CBS is
T
B
log(T" 1.
T+1
Proof. Plug in and solve the integral:
I log(! + 1)dl
B
T
T
1
T ((t (1++1) 1)lg(t + +1) 1111-t -)
T
log(T
T
Thus, for large T, the fixed CBS will scale as B log
D.3 Discussion
These results show that, if we are choosing the fixed batch size to minimize average distance to the CBS as it evolves
training, we should pick it, more less, simple function that slightly discounts the final CBS. Specifically, if we believe that the local CBS grows VT during training, then this derives the VT scaling law for B* proposed in prior work.
One limitation of this view that the L.2 residuals may not the right way to measure closeness to the CBS. In particular, it may be worse to overestimate the CBS compared to underestimate, as training above the CBS (with a scaled up learning rate) can unstable. We thus not read to much into this analysis. but view it as potentially useful starting point for future empirical and theoretical that derives CBS scaling laws from the development of the local CBS over training.
E
BPB Evaluation on Downstream Tasks
This section lists the datasets used compute BPB measures for downstream tasks. For multiple- choice tasks, we use the Cloze/Completion Formulation (CF), and compute the BPB metric on the gold answer. For completion tasks, we simply compute BPB over the correct answer. This approach was inspired by Bhagia et al. (2024). The selection of tasks follows the guidelines from Magnusson al. (2025).
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
CBS measurements in Section 3 and batch size warmup tested in Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Assumptions and limitations our CBS method are discussed in Section 3.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
19

---

Justification:
Our (minor) theoretical results are fully justified in Appendix D.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Hyperparameter information and other experimental choices are documented in Appendices A and B.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
20

---

Answer: [No]
Justification: accepted, we will release code with the camera-ready version. Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/ httpc//ggides//ooeeeeeee
ceeioonoliyyy olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips.. ctt.cc/ppuu. wtcp//ublcggguaiieeeeeeeeecopeee
Sttp://CooSSSuiiiseeee wwtuuuimmssiooni cesssinoolicooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Provided in Appendices A and B. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Figure shows an interval representing lower and upper bounds on the CBS, as documented in Section 3. Figure reports confidence interval for the noise scale whose details are documented in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
21

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Provided in Appendices A and B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https https://neett hetp://nneuipppe cttptppp.pppppppppe httpr//p..../pblltc///ccccccieee
Answer: [Yes]
Justification:
No explanation of special circumstances necessary.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is foundational research on the science of pretraining language models with no immediate impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No pretrained models or datasets released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The licenses for OLMo, C4, and the Pile acknowledged in Appendix C.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com con/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects involved.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No human subjects involved.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "Gorosn96PX", "reviewer": "Reviewer_AKeu", "summary": "This paper presents a practical method to empirically determine the **critical batch size** in large-scale language model pretraining—the largest batch size at which performance does **not** degrade compared to training with smaller batches and more steps. Beyond this threshold, further increasing batch size while reducing the number of steps begins to hurt final model quality.\n\nThe authors introduce a simple on-the-fly test for identifying this threshold: take a checkpoint from an ongoing training run, resume training with a larger batch size and appropriately scaled learning rate (e.g., following square-root scaling), and observe the smoothed training loss after a short additional training segment. If the loss remains within a fixed tolerance (e.g., within 0.01 of the loss without changing the batch size), the new batch size is deemed acceptable.\n\nUsing this method, they analyze the OLMo-1B and OLMo-7B models and find that the critical batch size increases sharply in the early stages of training (for around 200B) but plateaus later. The analysis also shows that **Gradient Noise Scale**—a previously proposed heuristic for estimating the largest usable batch size underestimates the critical batch size by roughly a factor of two, making it overly conservative.\n\nBased on these findings, the authors propose a batch size warm-up strategy: begin training with a smaller batch size and increase it gradually during training, adjusting the learning rate accordingly. They decide the steps to increase the batch size by referring to the complete OlMo 1B run. This strategy improves training efficiency by reducing the number of steps needed to reach target loss while maintaining or slightly improving downstream performance compared to a small batch size baseline.", "questions": "1. Is the batch size set to a constant in pretraining and mid-training in these experiments? Is there a possibility that it should be different at the end of pretraining versus at the beginning of mid-training?\n\n2. In terms of wall time, how much saving can the batch size warmup produce?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. The proposed method is simple, easy to understand, and has minimal assumptions, making it a convincing method to determine CBS in practice.\n\n2. The paper provides strong and thorough empirical evidence to support their findings on the critical batch size, showing that their method effectively improves training efficiency.\n\n3. Overall, the paper is clearly written and easy to follow.\n\n**Weaknesses:**\n\n1. The authors didn't discuss a related work on optimal learning rate and batch size scaling for Adam [1], which generalizes the gradient noise scale method. Including a discussion or comparison with this related work would have strengthened their argument.\n\n2. There seems to be some minor arbitrariness in how the thresholds for increasing the batch size are set. In their experiments, the critical batch size appeared to be around 4096 as early as 200 billion tokens, but the authors only set the threshold to increase the batch size at 600 billion tokens, which is much later in the training process.\n\n[1]  Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling", "was_revised": true, "final_justification": "This paper offers a clean way to specify the critical batch size in pretraining and I don't see any significant issues about the submission.", "rebuttal": "Thanks for your review! We are glad to hear that you find our paper well-written and that its findings are convincing.\n\nWe’ll make sure to incorporate some comparison with [1] as you suggest.\n\nRegarding the minor arbitrariness in setting the batch sizes thresholds for the warmup experiment, we set the threshold manually early on the research project based on a preliminary version of Figure 2 and then launched the batch size warmup pretraining runs. We agree that, all else equal, it would be nice to eliminate this arbitrariness. However, the pretraining runs for the warmup experiments were quite expensive, and we judged that the substantial cost for rerunning those experiments was not worth it to eliminate minor arbitrariness, especially since we believe the results are not especially sensitive to a precise choice of threshold. We will add some brief discussion of this limitation in the paper and highlight that making the choice of threshold more systematic would be appealing in future work."}, {"review_id": "O6h0dEwl17", "reviewer": "Reviewer_6mhC", "summary": "This paper studies \"critical batch size\" (CBS), the largest batch size a training run can use without hurting the performance (assuming hyperparameters like learning rates are adjusted correctly). Prior methods use gradient noise to estimate CBS, but it is based on the assumption of using SGD and well-conditioned optimization, which is unrealistic. \n\nThe three main contributions of this paper is very clear:\n\n(1) This paper proposes a new empirical way to estimate CBS by (a) starting from a checkpoint in the middle of training and (b) training for a short period of time and checking if the performance matches the baseline. Compared to prior method, this only needs to train the model for a little bit. It also allows for discovering the local CBS, while prior method can only find a global CBS.\n\n(2) By using the aforementioned method, the authors discovered that CBS is small at the beginning of training, then increases, and then plateaus. The final CBS does not differ much across scales (the authors used 1b and 7b).\n\n(3) To accommodate the finding mentioned before, the authors proposed a batch size warmup regime (double the batch size every time CBS increases). The authors showed that this leads to slightly better performance than using either the initial (smallest) or the critical (largest) batch size throughout the training.", "questions": "From Figure 2, it seems that the 1B model should reach the batch size of 4K at 100B tokens? However, the authors say that this is reached at 503B tokens in section 4. Not sure where I get it wrong?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strength**\n\nThe new way of empirically finding CBS is novel and allows us to examine the local CBS at any given point of training. The finding that CBS changes throughout training is super interesting and potentially reveals how gradient noise changes throughout training as well. The batch size warmup regime, though not significantly better than using one CBS throughout training, is something.\n\n**Weakness**\n\nMy main complain is on the experiments with the batch size warmup regime. First, there should be more evaluation (for example, few-shot downstream tasks evaluation). The authors used the loss on these tasks, which is not as straightforward to interpret. Second, if I understand correctly, the schedule is based on the results from Section 3. It would be more complete of a method if the authors propose a way to online check the CBS inflection point.", "was_revised": true, "final_justification": "I'm very happy with the author response and will maintain my positive score.", "rebuttal": "Thanks for your review!\n\nRegarding the downstream evaluations, we have found the loss-based evaluations as represented in https://arxiv.org/abs/2412.04403 to be some of the most reliable for evaluating pretraining interventions. Of course, this kind of evaluation is an unsolved problem with a lack of scientific consensus. If you have some specific suggestions for alternate types of evaluations we could use here, we would be curious for suggestions.\n\nIndeed, the batch size warmup schedule is based on the results in Section 3. Because we only needed to determine two thresholds for increasing the batch size, we determined them manually based on Figure 2 before launching the batch size warmup pretraining runs. We did not have the resources to relaunch the pretraining runs using a more online approach, but we agree this would be a useful refinement of the approach and revise the paper to clarify this.\n\n> From Figure 2, it seems that the 1B model should reach the batch size of 4K at 100B tokens? However, the authors say that this is reached at 503B tokens in section 4. Not sure where I get it wrong?\n\nWe chose the thresholds manually and conservatively in Figure 4, with the goal of having the warmup batch size upper bounded by the curve in Figure 2. Based on a noisier initial version of Figure 2, we preferred a later threshold for increasing to batch size 4K, and once we produced the newer Figure 2 plot, we did not have the resources to relaunch the pretraining runs. We will clarify this manual choice in the paper and emphasize that there is an opportunity for future work to determine batch size thresholds more systematically."}, {"review_id": "g6oLE0a8nZ", "reviewer": "Reviewer_Hdu9", "summary": "This paper proposes and analyzes a method for estimating \"critical batch size\" (CBS). The paper measures the CBS by performing \"branched training\" for relatively small token windows, across multiple batch sizes, defining the CBS as the largest batch size at which the branched training does not reduce loss (with some tolerance/smoothing), as in Figure 1. The paper specifically compares this method to the gradient noise scale method of estimating CBS proposed by [McCandlish et al., 2018]. They point out that the gradient noise scale has specific, stringent underlying assumptions, and show empirically that it can underestimate CBS.\n\nThe paper then uses their branched training estimate of CBS to define a batch size warmup schedule. The authors show that when training a 1B parameter model, this slightly outperforms a small-batch size baseline in terms of loss and downstream evaluation metrics, while using significantly fewer gradient estimates. On the other hand, it does significantly better than a large-batch size baseline in terms of loss, while requiring somewhat more gradient steps.", "questions": "1. What happens when you use the CBS estimate from [McCandlish et al., 2018] to do batch size warmup, keeping everything else the same? That is, how much does the underestimate of the CBS from the gradient noise estimate translate into reduced performance on loss and/or downstream evaluation metrics?\n\n2. How does the branched training estimate vary with the number of tokens $\\Delta$ used for branched training? For example, would training with 1B tokens in Figure 1 give similar CBS estimates? What about 4B? You mention in L131-133 that larger $\\Delta$ may produce larger CBS estimates - I don't immediately see the reason. Viewing CBS as a random function of $\\Delta$, should we expect that its expected value changes with $\\Delta$?\n\n3. More generally, do you have any heuristic for selecting $\\Delta$, potentially as a function of some model-scale-dependent training horizon (e.g. as some fraction of the Chinchilla-optimal number of training tokens)? \n\n4. I found Appendix D to be quite interesting - is the reason it is not in the main body just due to lack of space? In the context of fitting such a power law, what would happen if you tried to fit your local CBS estimates to a power law? Do they obey this behavior, or do they exhibit qualitatively different behavior?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\nI will be forthright: I think this is a well-written, well-executed, and novel paper, and should be accepted as such. Its clarity is perhaps its greatest virtue. It does a terrific job of being absolutely explicit, while still making the paper easy (even enjoyable) to read. I will also note that the paper bucks the trend of overstating contributions. The paper does a great job in its abstract and introduction in detailing exactly what they do. Their experimental setup is extremely clear, with explicit discussion of every experimental choice. I also think the paper does a good job of avoiding confusing, extraneous, or only tangentially related experimentation. Every figure and table has a clear purpose in contributing to the paper's overall message, and I would encourage the authors to continue to write papers like this.\n\nThe paper is sound, novel, and clearly relevant to a large body of work on pre-training dynamics. I will note that the paper's approach to measuring CBS is straightforward. I wish to fend off potential criticism on this front preemptively. This is a strength of the paper, not a weakness. The fact of the matter is that CBS has been an opaque phenomenon for years, and after reading this work I felt I understood it better. Not fully of course (the paper makes no mechanistic explanations for it, nor does it need to). But we can't study what we can't measure, and I'm particularly appreciative that the question of how to measure is approached so carefully. The approach leads to plenty of interesting open questions (some of which I detail below). Moreover, the empirics are well-done (with one caveat, as I discuss below) and make good use of datasets, models, and techniques training methodology that are widely used, increasing the paper's relevance.\n\n## Weaknesses\n\nGenerally speaking, most my questions about the paper are in the vein of wanting to know more about the behavior of the branched training estimate. This is honestly a good critique to have - I have some questions that I think the authors could address, not because it's necessary for the paper to be accepted, but because I think the paper is quite compelling. That being said, there are a few issues that I want to emphasize.\n\n### Ground-truth CBS\n\nFigure 3 and the discussion around it (L173-180) is useful, but has an implicit assumption that I'm not sure I understand. The authors write that the gradient noise estimate of the CBS \"underestimates the CBS by several orders of magnitude\". What is the ground truth CBS that the authors are using in such a statement? Figure 3 would suggest that the answer is 4096 (at least, towards the end of training) but I'm not sure why this is treated as the actual CBS. This number is closer to what their own methodology suggests is the correct CBS, but that doesn't make it the ground truth. For example, is it possible that the method in Section 3.1 *overestimates* the CBS?\n\nThis lack of ground truth could be dealt with another way, via downstream evaluation, which unfortunately brings me to my next point.\n\n### Why not train with a gradient noise scale-based batch size schedule?\n\nI think that Figure 3 is somewhat, but not wholly convincing in terms of the superiority of the method in 3.1 to the gradient noise scale method of estimating CBS. To that end, I'm surprised that the authors didn't train a model using a batch size schedule derived from the gradient noise scale CBS. The warmup could be the same as in 4.1, just change the CBS estimator. If (and I suspect it would) the gradient noise scale yields an inferior model to the batch size based on branched training, that'd be a great point of evidence in favor of branched training! Regardless of being accepted, this is something that I hope the paper includes at some point - I think it really would make the comparison point much clearer.\n\n### How does the CBS estimate depend on $\\Delta$?\n\nThis is a more minor issue, but one that was on my mind while reading. The branched training depends on the window size $\\Delta$. L128-L133 discusses the role of $\\Delta$ more briefly than I wanted. For example, the statement that larger $\\Delta$ may yield larger CBS estimates makes some sense, but it'd be nice to see some kind of exploration of this fact. For example, would a miscalibrated $\\Delta$ yield better or worse CBS estimates than gradient noise scale estimates? Moreover, how would the CBS estimate change between settings where $\\Delta$ is too small versus settings where $\\Delta$ is too large?\n\nI will note that no paper can provide comprehensive answers to everything. I think it's reasonable to treat this specific \"critique\" as evidence for accepting the work - the questions are natural ones that could (and should) be answered by follow-up work.\n\n## Other comments\n\nI have left other comments on the work, some overlapping with the above, some not, in the form of questions below.", "was_revised": true, "final_justification": "I have read all other reviewer responses and author feedback. Nothing about the strengths of the paper have changed in my opinion, and the other reviewers basically seem to have a positive impression as well.", "rebuttal": "Thank you for your review and your appreciation of our work! We are very happy to hear that you appreciated the contributions and clarity of the paper.\n\nRegarding your questions about the branched training method, we have some thoughts.\n> The authors write that the gradient noise estimate of the CBS \"underestimates the CBS by several orders of magnitude\". What is the ground truth CBS that the authors are using in such a statement? Figure 3 would suggest that the answer is 4096 (at least, towards the end of training) but I'm not sure why this is treated as the actual CBS. This number is closer to what their own methodology suggests is the correct CBS, but that doesn't make it the ground truth. For example, is it possible that the method in Section 3.1 overestimates the CBS?\n\nThis is a good question, and we will try to make our reasoning more explicit. Indeed, the comparison between Figures 2 and 3 suggests that the gradient noise scale underestimates the CBS relative to our method. A priori, it is possible that our method is an overestimate and that the gradient noise scale is correct. However, our pretraining results in Figure 4 provide evidence against this (i.e., they show that our method is not an overestimate), because we can train with the larger batch size predicted by our method and obtain no degradation in final loss. By definition of the CBS, this means the batch size produced by our method does not exceed the CBS (and thus the smaller gradient noise estimate is an underestimate).\n\n> Why not train with a gradient noise scale-based batch size schedule?\n\nWe think this could have been a nice, though not entirely central, comparison to have. The prediction based on the previous sections would be that we could still obtain decent pretraining loss (since we’re training below the CBS) but that the number of gradient steps would be significantly larger than training with the branched training batch size. The latter claim is true by construction and thus not worth testing, but it could have been interesting to measure the loss attained. On the other hand, it takes a while to launch the pretraining runs behind the results in Figure 3, and we were not able to launch this pretraining run during the short rebuttal period, especially given that the results would not be essential to our main findings.\n\n> How does the CBS estimate depend on $\\Delta$?\n\nUnder the local recovery assumption (Assumption 3), increasing $\\Delta$ up to the total pretraining token budget will only increase the estimate of the CBS, so CBS estimates with smaller $\\Delta$ are still a lower bound. The choice of $\\Delta$ can thus be understood as striking a tradeoff between finding the largest possible CBS lower bound and the compute budget one is willing to spend (since larger values of $\\Delta$ require more training). We did not explore this tradeoff in further empirical detail since we chose to allocate our attention and compute towards the large scale batch size warmup pretraining runs instead. However, we’ll try to clarify the role of $Delta$ in the paper and add a note that understanding the role of $\\Delta$ is an open question for future work. \n\n> I found Appendix D to be quite interesting … In the context of fitting such a power law, what would happen if you tried to fit your local CBS estimates to a power law?\n\nWe spent some time trying to fit a parametric form to our CBS plots, and we did not find any form that we were confident fit the shape of the function. This could be because there is a phase change between the first and second parts of the trend. Because we did not have any conclusive insights here, we did not make any claims about the functional form of the CBS trend in the paper, though we thought it would still be interesting to include the theoretical investigation related to this in the appendix since it might be relevant for future research."}, {"review_id": "z7ch8P74FF", "reviewer": "Reviewer_xcpd", "summary": "This paper aims to answer the question that what’s the suitable batch size during training language models at scale. One existing work utilizes the gradient noise scale during training to estimate the critical batch size (CBS), below which training will not substantially degrade loss. However, this method requires two strong assumptions: SGD Optimizer and Well-Conditioned Optimization. In practice, these requirements are not typically met or be too expensive to realize, which limiting its applicability. To this end, this work introduces a simple, empirical approach to directly measure the CBS through branched training. The authors find that the CBS evolves over training, that is, CBS is near 0 at initialization, increases rapidly at first, and then plateaus as training progresses. Then, they present a natural way to reliably train language models at large batch size, i.e., batch size warmup method. Experiments with 1B- and 7B-size models demonstrate the effectiveness of the proposed approach.", "questions": "Line 58: but than -> but then?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. They introduce an empirical method to directly measure the CBS via branched training.\n2. They discover the finding that the CBS changes over the course of training, i.e., it improves rapidly initially but then flattens off. And, the CBS not depend strongly on model size.\n3. They propose a batch size warmup strategy for large batch training: begin training with a small batch size and double it whenever the CBS increases sufficiency. \n\nWeaknesses:\n1. In table-1, compared to Large-Batch Control method, the proposed approach: Batch Size Warmup, gains less margin. For example, it only surpasses 0.7% in terms of PT Loss. It seems to be a performance fluctuate. Besides, it cannot save more gradient steps than Large-Batch Control method. In another word, it is better to use a larger batch size, e.g., 4096, to train a LM, rather than the proposed one.\n2. The work claims that the proposed method utilizes 43% fewer gradient steps than the original training run. But in Lines 71-73, they also stated that their method is more expensive than launching many full training runs. These claims seem conflicting.\n3. Does the data scale impact the CBS? In the experiments, 2B data is utilized. I wonder whether the CBS changes with another data size. Further, do different model architectures also impact the CBS? It lacks some discussion about that.\n4. In the experiments, it lacks the comparison with the gradient noise scale method.\n5. In the proposed batch size warmup strategy (Line 207), why double the current batch size? How to specify the hyper-parameter: 2? If we train another LM model with different model architecture (e.g., MoE),data scale, model scale, and optimizer, how to set it?", "was_revised": true, "final_justification": "I have a positive attitude toward the acceptance of this paper.", "rebuttal": "Thank you for your review!\n\nRegarding weakness (1), our claim is that batch size warmup attains comparable (in some cases slightly better) loss compared to the small-batch run while using fewer gradient steps. On the other hand, the large-batch run degrades loss, as shown in Figure 4 and Table 2. The takeaway we want to highlight is that the warmup method is a way to achieve comparable loss with fewer gradient steps; while the large-batch method saves more gradient steps, we find it degrades loss.\n\nRegarding Lines 71-73, this was a typo about the branched training method. Thanks for catching this! In fact, branched training is *less expensive* compared to training many full runs, and batch size warmup does save 43% gradient steps relative to the small-batch control. We will update lines 71-73 to correctly say “less expensive”.\n\nRegarding data scale, we agree this is an interesting question, but outside the scope of this project, where we focus on the evolution of CBS over the course of training given a fixed data distribution. As mentioned in Section 2, studies [1] have suggested that CBS scales with data scale.\n\nRegarding the comparison of gradient noise scale with the branched training approach, the contrast between Figures 2 and 3 does give a comparison of these approaches for estimating the CBS. We do agree that it could have been interesting to consider batch size warmup with a gradient noise schedule, but we did not have time to launch this experiment during the rebuttal period, and, while it would have been interesting to have, this experiment is not central to the main claims in the paper (see further discussion in our response to Reviewer Hdu9).\n\nRegarding doubling the batch size, we chose this approach because of the practical constraint that we want the batch size to remain a multiple of the number of GPUs whenever we increase it. Because our number of GPUs during pretraining is a power of 2, doubling whenever we reach a new power of two was an easy way to maintain this invariant. In general, one could imagine more complex approaches, such as explicitly selecting the largest multiple of the number of GPUs less than the CBS. We will clarify our motivation and the potential alternatives around this hyperparameter choice.\n\n[1] https://arxiv.org/abs/2410.21676"}]

Author rebuttal:

Meta-review: This paper proposes a new method for estimating the critical batch size in LLMs that avoids restrictive assumptions and shows how it will changes over the course of training. The generality of the method to MoE’s or beyond LLMs is not studied and there were some concerns on the warmup scheduled. Overall the work is clearly written, easy to follow, and backed by experiments on realistic model scales. The reviewers were unanimous about accept.
Decision: Accept (spotlight)