How Does Critical Batch Size Scale in Pre-training?

Thank you for your constructive feedback. We have addressed each of your points as follows:

The paper heavily relies on Chinchilla scaling theory

Please note that we do consider under-training and over-training compared to Chinchilla compute-optimal regimes (fig. 1 middle and fig. 6). This enables a more controlled comparison as we control for model size and data size separately, which further strengthens our claim.

"dynamic horizon" schedule

Thank you for the suggestion. We do compare extensively using Sec 2.1, Appendix A. Especially our results of fig. 10 show that tuning the EWA decay rate high is needed for training longer. This plays a similar role in tuning down the learning rate when training longer, as suggested in [1].

[1] Scaling Optimal LR Across Token Horizons Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, Xia Song.

Weight decay was disabled for the experiments.

We disabled weight decay as we found it did not change the critical batch size in preliminary experiments for our setup. If this is crucial for reviewers’ judgment of our work we would be happy to redo some of the experiments with 1e-4 decoupled weight decay (as suggested in Wortsman 2023).

We also note that our setup is based on the Olmo codebase. A previous work [2] (Figure. 5) also conducted experiments on the Olmo codebase and only observed very small difference for different weight decay values $0, 10^{-6}, 10^{-5}, 10^{-4}, 10^{-3}, 10^{-2}$ for autoregressive LM pre-training on C4.

[2] Deconstructing What Makes a Good Optimizer for Language Models. Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, Sham Kakade

"trajectory" to be defined more clearly. e.g., is it the trajectory of the weights during training? Also, please cite a source for this statement.

Thanks for the comment. Yes, we mean the trajectory of the weights. We will cite the Master Theorem of Yang and Hu (2021) in the proof as well.

Which model after what training have these evaluation variance numbers been obtained?

As suggested in Sec 2.1, we use 151M models for most ablation experiments to decide hyper-parameters and evaluation configurations. We use C4 consistently throughout the paper as our training and evaluation data.

Please note that the scaling of epsilon in Adam itself is an open research question and due to budget constraints we found that it has minor effects and didn’t extensively sweep over different other values. Note that we study models up to 1.2B, and [3] pointed out For standard parameterization models with up to a billion parameters, the typical default value of 1e-8 is likely acceptable

[3] Scaling Exponents Across Parameterizations and Optimizers. Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, Jeffrey Pennington.

Typos on R(M,t) and \theta

Thank you for pointing these out. We fixed them and updated the draft.

Does Figure 5 have empirical backing regarding the predicted CBS and the formalism of staying within the 20% margin of steps for similar loss?

Or if the decoupled prediction for batch sizes hold reliably for different scales or ablations such as architecture, datasets, hyperparameters?

Please note that in Sec 2.1 we have explained that we only take into account the best-performing run for each batch size. So the scaling laws hold if we have models well-tuned across hyper-parameters we consider in the paper. We only consider transformer-based auto-regressive language models on C4 as they are canonical settings that are practical and can be reproduced more easily.

How exactly are the tuning decisions undertaken especially in Section 3.3 (example, L371-372)?

Please note that we have mentioned that we scale the warmup tokens proportionally: in Fig. 1 we have compute-optimal setting 1x, and the 0.28x, 0.5x for under-training, 2x, 4x for over-training. If we set the warmup tokens for compute-optimal setting as $0.25C_{Chin}N$, then we have $0.07C_{Chin}N$, $0.125C_{Chin}N$ for over-training and $0.5C_{Chin}N$, $C_{Chin}N$ for over-training.

In Theorem 1, is it expected or meant to be w2>w1? If so, then there is a typo and makes a significant difference to the interpretation

No, only the two widths are greater than $w$, there is no relation between the two widths

What is the significance of t in Theorem 1? Should we compare losses under similar tokens for two vastly different model scales, especially following Chinchilla?

Yes, this theorem specifically only holds for models with different widths trained for the same number of tokens (and thus does not hold for Chinchilla setup where tokens increase with increasing model scale). So $t$ refers to the fixed number of tokens.

What is the specification or how is H defined in Section 4.2?

As shown in the beginning of Sec 4.2, H is the covariance matrix of the Gaussian data distribution.

Is there a more intuitive summary for the proofs for Theorem 2 and 3 especially for readers (like me), not familiar with Zou et al. (2023)?

Please refer to Sec 1.2, our newly added section, for an overview of the theory:

• Theoretically, maximal update parameterization suggests that, beyond a certain point, increasing the width of the neural network (while keeping data size fixed) does not further increase the critical batch size. In contrast, by analyzing a simple least-squares regression with mini-batch SGD, we provide a theoretical basis for how the critical batch size continues to scale with increasing data size.
• In infinite width regimes, training dynamics and performance of the networks become effectively independent of the model size. Consequently, the critical batch size remains nearly invariant when scaling up the model size beyond this point, indicating that larger models do not require proportionally larger batch sizes to achieve optimal training efficiency.

How exactly does the hybrid evaluation (L968) provide more accurate information?

The evaluation is for making sure the recorded #steps required to reach a goal loss would not be overly estimated because of infrequent evaluations. So we adopt the strategy to evaluate the runs more frequently when near the end of training.

What or how are the hyperparameters (HPs) for parameter sweeps ranked or ordered (L977-979)? Does a different ordering yield potentially different results? What are the default values for other hyperparameters when tuning each of these in order? Assuming in this order that LR is considered to be the most important HP, why tune the least important HPs (in β2 and λ) for each model scale and not LR (Table 4)?

Note that we have swept learning rate for various model scales as suggested by Tab.3. We include the search details for completeness but find that ordering does not result in too many differences. For example, our newly added experiments in Fig. 10 shows that our default learning rate 3.16e-3 is consistently the best among others with different EWA decay rates. We highlight the default values in Tab.3 using bold fonts.

I may have missed it but could there be a reference for how the power law model fit is made for the values in the caption of Figure 7?

The scaling law details are included in Sec 3.2 and Sec 3.3. Appendix E also includes additional details and interpretations.

How these designs were arrived at.

We include quite a few additional experiments on understanding design choices like learning rate in the Appendix. Please let us know which parts are unclear and we can clarify further.

Thank you for your constructive feedback. We have addressed each of your points as follows:

different seeds, design choices

Thanks for the feedback. Due to computation constraints, we did not sweep over multiple seeds but focused mainly on different hyper-parameters. But from our variance measure experiments in Appendix D, our runs are pretty consistent across seeds. Please note that we have included many results in the Appendix to study the effects of major hyper-parameters.

hyper-parameters are done with small-scale 151M models

Due to computation constraints, we use small-scale 151M models to study some hyper-parameters including context lengths, momentum $\beta_1$, scheduling and learning rate studies. However, for important ones that can lead to performance boost like EWA decay rate $\tau$ and $\beta_2$, we swept all the values as indicated in Tab.3.

A limited grid search over learning rate

Note that we have swept learning rate for various model scales as suggested by Tab.3. We found that under our constant+EWA strategy, learning rate does not affect the steps required to reach a target loss too much. In Appendix Fig. 10, we include more studies to understand the the impact of learning rate. We found that our default value 3.16e-3 is stable across the three regimes we considered. To improve model performance with a fixed learning rate, increasing the training duration may require a larger EWA decay rate for optimal results.

Connections between the outcome of the theoretical proofs and empirical results? Is there a way to connect the toy regression experiment with the scaling law exponent fits?

Please note that our theoretical proofs corresponds to two controlled settings we consider: (we included the following into Intro Sec 1.1, Sec 1.2 as well) Empirically:

1. If we scale up training duration D while keeping N fixed (Figure 1, middle), the critical batch size increases to a similar degree.

2. However, we find that CBS remains nearly invariant when scaling up N while keeping D fixed (Figure 1, right), suggesting that CBS weakly depends on model size N but more strongly depends on data size D.

Corresponding to the two cases above, theoretically:

1. In infinite width regimes [Yang et al, 2021], training dynamics and performance of the networks become effectively independent of the model size. Consequently, the critical batch size remains nearly invariant when scaling up the model size beyond this point, indicating that larger models do not require proportionally larger batch sizes to achieve optimal training efficiency.

2. Consider mini-batch SGD with $D$ samples in the least square problems under power-law source and capacity conditions. The CBS, which enables mini-batch SGD to achieve the minimal expected excess risk while ensuring the fastest possible serial runtime, is given by $B^*(D) = \Theta(D^c)$, where the exponent $c\ge 0$ is determined by the exponents of the source and capacity conditions. In the regime where the variance error tends to be dominant, we have $c>0$, indicating CBS grows with data size.

Discrepancy of batch size 256 or 512 as the batch size for calculating overhead ratios

Please note that we use $B_{\text{opt}}=256$ to determine the target loss consistently, but report the relative number of steps with 512 batch size results as the denominator because we are only interested in large batch size regimes to build intuitions on CBS. We didn’t claim they are the same.

Should likely be b > a in L1299. Typo on L366: perhaps should be "They" -> "Then". Discrepancy in Fig. 5 caption $C_{chin}$. Figure 10 mentions context length. Plot sizes.

Thanks for pointing these out. We have greatly revised the draft to fix those issues.

Figure 4.a doesn't mention which model width or depth.

Please refer to Tab. 5 for the concrete setup, as mentioned in the main text.

Recommended to refer sections/tables/figures in the Contributions list in Section 1.

Thank you for the suggestion. We added Sec 2.1 and 2.2 accordingly for better introduction of our contributions.

Vertical lines in Table 4 separating different model sizes. Parantheses in L1119-1122 would enable easier parsing for readers.

Thank you for the suggestion and we would include the changes in the next draft.

I believe the trainings for all models >302M use no longer sub-epochs and tokens are repeated.

All the runs are one-epoch or sub-epoch without repeating data, a common practice in pertaining. C4 has 153.6 billion tokens and we are far from using all the data.

To clarify, do you suggest that instead of using LR cooldown towards the end of training budget, EWA is a better strategy?

We find that in our particular settings, constant+EWA is competitive. In contrast, [1] found that "We also experimented with an exponential moving average (EMA), which performed worse than SWA, and therefore we do not report EMA.". As indicated in the main text: as we focus on the number of training steps needed to achieve a target validation loss, learning rate decay strategies typically require predefining the total training duration. To address this, we propose using exponential weight averaging (EWA) to achieve the desired target validation loss, a simple approach that matches other popular choices (Figure 2). This enables training beyond fixed durations or data size, allowing to resume training from checkpoints until the target validation loss is achieved.

adding constant + cooldown.

WSD is exactly the Constant + Cooldown you are referring to, which was released earlier according to the arxiv timestamp https://arxiv.org/abs/2404.06395.

How does this work out for larger scales when the tuning is feasible only on the smallest scale available of 151M?

For important hyper-parameters, we already mentioned in Appendix that we swept for all Chinchilla settings for all model sizes using the ones in Tab. 3 and 4. The experiments where we only used 151M models including context length etc are not primary experiments. Those are already on a much larger scale than the papers we discussed.

Hyper-parameter captions and reproducibility

Learning rate (3.16e-4), (1e-2) are the ones we experimented for 151M models and found that with EWA it does under-performs 3.16e-3. EWA decay rate (0.99995) is only for long 1.2B runs. We would clarify further in the latest draft.

Bold font means the default hyper-parameters that can closely reproduce our results without extensive tuning. Not bolding implies a full sweep for each model scale. Please note that without such annotation makes it harder to reproduce as hyper-parameter search requires computational resources. For important ones like $\beta_2, \tau$ we listed separately in Tab.4. Please also refer to our actual code implementation where we keep yaml files separately for each experiment. We can also share wandb run reports upon request after peer review.

how long does it take for training batch size 64, model size 1.2B

It takes 22h 12m 37s on a node with 8 80GiB A100s.

Mention model sizes in Figure captions/titles. It is sometimes hard to find the corresponding model size from the related text.

We would incorporate the feedback into our latest draft.

Right now, the key points I take with me from your paper:

• A formalism for critical batch size (CBS).
• The existence of an increasing CBS under Chinchilla scaling.
• Decoupled influence of the number of tokens seen on CBS for a fixed-size model.
• Neat theoretical proof of a CBS limit with model-size scaling (along width) under a fixed size of tokens.

Thus, the paper is an interesting read with clean empirical results to support the claims. A key practical aspect that emerges: when scaling up model size, the critical batch size should not be scaled unless accompanied by an increase in dataset size.

I increase my score to 6. Thank you for the paper.

Thank you for your constructive feedback. We have addressed each of your points as follows:

Schedule free optimizer

Please note that From Fig. 2, schedule-free optimizer is competitive for small batch sizes, this aligns with what the original paper found. Notably the original paper only consider a batch size of around 2^9. However, it’s worse for large batch sizes even if we tuned the beta1 from 0.9, 0.95, 0.98.

Reproduction on constant LR + EWA.

Please note that a very large EWA decay rate would cause instability but may lead to better convergence as suggested in the plot you referred to. As our goal is to reach a certain target loss, we use different EWA values for better convergence during the last part of the training. We included our implementation in the supplementary materials. Note that the performance gap also depends greatly on the setting and especially on training durations ana batch size.

Weight decay was disabled for the experiments.

We disabled weight decay as we found it did not change the critical batch size in preliminary experiments for our setup. If this is crucial for reviewers’ judgment of our work we would be happy to redo some of the experiments with 1e-4 decoupled weight decay (as suggested in Wortsman 2023).

We also note that our setup is based on the Olmo codebase. A previous work [1] (Figure 5) also conducted experiments on the Olmo codebase and only observed very small difference for different weight decay values $0, 10^{-6}, 10^{-5}, 10^{-4}, 10^{-3}, 10^{-2}$ for autoregressive LM pre-training on C4.

Please note that the Wortsman et al. (2024) paper you mentioned shows that it does not affect the optimal LR for large enough models with size more than 85M (Figure E.10). To understand whether disabling weight decay would need a smaller learning rate, we add new experiments to show that our default value 3.16e-3 is consistently the best if the EWA decay rate is well tuned in Fig. 10.

[1] Deconstructing What Makes a Good Optimizer for Language Models. Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, Sham Kakade

How did you use EWA and reproduction

Please note that EWA is introduced in Sec 2.1 and we included all the possible values in Tab.3, as well as additional ablation studies in Fig. 10. Note that it’s common practice to tune EWA gradually from small values to larger ones. Especially we found that similar to the learning rate, the effectiveness of EWA+constant in different decay rates depends on the training duration. Please refer to our supplementary materials for reproduction.

Did you tune SFO with the same sweeps as Adam?

Can you please clarify what do you mean by SFO? We didn’t use this abbreviation in our paper.

Why are the warmup steps chosen between 0.15, 0.25 and 0.35?

Please note that choosing warmup proportional to data size is a common strategy that has been justified in previous work. In [2], they set the warmup period to be the minimum of the model size N and 20% of the total token budget. In the table 5 of [2], they also show that Durations of up to 4N achieve very similar and slightly better results. Therefore, we believe sweeping over 0.15, 0.25, 0.35 would be a reasonable design decision to achieve reasonable performance.

[2] Resolving Discrepancies in Compute-Optimal Scaling of Language Models Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, Yair Carmon.

Suggestion on colors.

Thank you and we choose colors to reflect the ablations on different scales like model size and data size. We use Fig. 6 for side-by-side comparison to further support the claim.

Thank you for your detailed feedback.

Hyper-parameter tuning.

Please note that as indicated in Appendix, the baselines are well-tuned as we swept over learning rate [1e-3, 3.16e-3, 1e-2] for all schedulers and decay ratio [0.1, 0.2, 0.3] for WSD. For cosine scheduling, we first use optimal steps we found the relaunch that number of optimization steps as the maximum steps to make sure it would reach the the target loss right near the end of the training.

The impact of weight decay.

We include the following experiments using decoupled weight decay for all the runs. We have hyper-parameter sweeps for 151M models using three schedulers - EWA, WSD, and ScheduleFree and our Constant+EWA. (due to time constraints we only swept over 4 batch sizes).

lr - [3.16e-3, 1e-2], decoupled weight decay rate - [3e-04, 1e-04], batch size 512~4096, WSD decay ratio [0.1, 0.2, 0.3]. For the rest of hyper-parameters, we chose the best one from previous runs. We also sweep over the total training steps for both WSD and cosine to make sure they are close to optimal performance:

The results are in an anonymous link at https://docs.google.com/document/d/1edDHtur9019FJDqLJ0hILIZiKuV4YXsEaMf4gzTlkbg/edit?usp=sharing.

EWA+constant remains to be competitive. We see in the top figure that weight decay does slightly improve the schedule-free optimizer but not for others. We do used the original implementation of Schedule-free Optimizer provided in their Github, which applies weight decay at y as you suggested. All the schedulers remain nearly the same compared to the bottom figure reported in our paper. As our focus is the critical batch size, we additionally emphasize that adding weight decay does not (qualitatively) change the critical batch size too much.

If the reviewers believe the claim should be revised, we are happy to do so. For instance, we propose that “EWA consistently performs as a competitive choice compared to the other three schedulers we evaluated, within the specific settings of our study.”

EMA decay values.

We do not switch the value over the training but we show the values we swept in Tab. 3. We emphasize that as the learning rate is constant, we can more easily resume from checkpoints to reach certain target validation loss, which is well-suited for our case.

Thanks for the reply and the new experiments. Just to be clear: I did not imply that the hyperparameters that you have weren't properly swept or tuned, but you chose specific settings which -- under all current understanding in field -- lead to suboptimally trained models which you then compare to. I just wish the hyperparameter sweeps were done with a different setup in the first place.

Moreover, to be very rigorous: when you say you use 'decoupled' weight decay, did you fix the line of code that I mentioned before? The attached code does not perform the original decoupled (or independent) weight decay which therefore would require much higher WD values such as 0.1. Values in the range of 1e-04 would not make much difference.

Moreover, could you state the final learning rate for both WSD and cosine? As mentioned above, the code only decays to 10% per default and I do not see this config changed anywhere. This would lead to suboptimal final loss values as the LR is not fully annealed, especially for the large LR values.

I did not imply that the hyperparameters that you have weren't properly swept or tuned, but you chose specific settings which -- under all current understanding in field -- lead to suboptimally trained models which you then compare to. I just wish the hyperparameter sweeps were done with a different setup in the first place.

It’s possible that weight decay can slightly increase the performance in pre-training. But the performance gain is minor as and Wortsman et al 2023, Zhao et al 2024 shows. Especially our primary focus is on critical batch size, rather than the optimal performance metric commonly emphasized in previous works. Thanks for the careful read and we can make some more discussions on this.

Moreover, to be very rigorous: when you say you use 'decoupled' weight decay, did you fix the line of code that I mentioned before? The attached code does not perform the original decoupled (or independent) weight decay which therefore would require much higher WD values such as 0.1. Values in the range of 1e-04 would not make much difference. Final learning rate for both WSD and cosine? As mentioned above, the code only decays to 10% per default and I do not see this config changed anywhere. This would lead to suboptimal final loss values as the LR is not fully annealed, especially for the large LR values.

We meant the default weight decay implementation in torch. Following the setup you described and using the same settings above, we swept over

WD rate 1e-2, 3.16e-2, 1e-1 for LR 1e-2

WD rate 3.16e-2, 1e-1, 3.16e-1 for LR 3.16e-3

for all the schedulers. We decay cosine and WSD to zero as you suggested.

Results are included in the doc (top one): https://docs.google.com/document/d/1edDHtur9019FJDqLJ0hILIZiKuV4YXsEaMf4gzTlkbg/edit?usp=sharing.

We find that a large weight decay + decay to zero slightly improves cosine and schedule-free, but the performance between constan+EWA and cosine are close.

Please also note that decay to 0 may improve loss but it may not be usually adopted. For example, [2] observe that However, when evaluating on downstream benchmarks (see Section B.5), we see that annealing to zero can hurt metrics; we posit this comes from too early saturation.

[2] Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

Finally, we want to emphasize our focus and main contributions are on the empirical scaling behaviors of critical batch size followed by theoretical justifications. We did not claim in the main text that our constant+EWA approach offers a significantly improved convergence rate compared to other baselines. Instead, we highlighted its suitability for our study, treating it as a preparatory sub-section. The takeaway box is meant to be concise, where the claim “EWA consistently improves model training efficiency” is in comparison to pure “constant without EWA” based on our observations. We understand the takeaway as presented is a bit oversimplistic and would be happy to add more context. Thanks for noticing this and helping to properly focus our paper.

Thank you for your constructive feedback. We have addressed each of your points as follows:

“EWA consistently improves model training efficiency”

We included our code in Supplementary Materials. In the README.md or the assets folder, the experiments for cosine decay and WSD for various training lengths T. We can see that this is indeed the optimal training length for both of them. Note that these curves are at extremely high batch sizes, that’s why the decay indeed hurts as the noise induced by variance in the iterates is pretty small. We will still reword the statement to hold only in our setup, but we have done training length sweeps for cosine and WSD (for the 150m model), and we are comparing only after these sweeps. Note that we include the concrete sweeping settings in the Appendix. The decay phase is set to be last 0.1, 0.2, 0.3 of the training durations for all batch sizes. In the supplementary assets folder or README.md, we show that constant+EWA is competitive with Cosine and WSD with various decay ratios. fig. 9 (c) (d) (e) also shows that even if WSD reaches the target loss right at the end of training, it’s still comparable with constant+EWA.

The lines start to diverge a bit more for larger batch sizes, so is it possible that this is more pronounced for other settings (e.g. larger models, more training data, etc.)?

Yes, this may be possible. However, based on Sec 2.3, we can expect that even if the curves diverge, that might happen beyond the critical batch size.

How the results and conclusions would change if one varies this parameter, e.g. to 10% or 50%

Please note that it’s a design choice for choosing the overhead level depending on concrete training setups. The overhead level can be generalized according to specific practical settings and we found that these do not affect the fitted scaling laws too much. For example, when having the following 10%, 50% overhead, we get data scaling laws for 302M models of similar form as

10%: B* = 20.67 * D^(0.48)

20%: B* = 22.91 * D^(0.47)

50%: B* = 30.50 * D^(0.44)

width vs depth scaling.

Please note that those settings are Chinchilla compute-optimal settings where the data size is scaled up proportionally with the model size. So large models are trained on more tokens. This means still increasing the scale via enlarging width or depth has similar CBS scaling behaviors, which we found in the next section to be mostly dependent on the data size.

In the legend of Figure 1b (left), what is the number in parenthesis? Is it the (relative) number of training samples?

It is the token size, denoted by how many times or proportions the Chinchilla tokens are used. We revised the plot to improve clarity.

Plot size, font size, white space.

Please refer to the current PDF. We have changed the formats of most plots substantially.

Linear scaling line as well as the 20% region and highlight the critical batch sizes, e.g. by a star marker?

Please refer to current Fig. 5 for an illustration. The corresponding definition is on the left. Also, Note that connecting diagonal lines of squares in each plot can lead to the linear scaling regime.

Fig. 2 plot formats.

Thank you for the reminder. We have updated the plot with new, more distinctive marker symbols to enhance visibility.

In line 110 you mention that "traditional learning rate decay strategies typically require predefining the total training duration" and you mention Defazio et al., (2024). However, isn't Schedule-Free a counter-example? Schedule-Free is also using a, at least to me, a relatively similar approach to your "constant+EWA" strategy, no (e.g. comparing the equation in line 158 to Eq. (5) by Defazio)? Could you elaborate on how your method differs from Schedule-Free?

Sorry for the confusion. What we mean is how to get rid of fixed training duration is an actively studied problem so we cite Defazio et al., 2024. We have revised the narrative for clarity. Mathematically they are equivalent, overall ScheduleFree SGD can be understood as a accelarated SGD followed by weight averaging. This can be shown by denoting $y_t =(1-\beta) z_t+\beta x_t, z_{t+1} =z_t-\gamma g\left(y_t\right), x_{t+1} =\left(1-c_{t+1}\right) x_t+c_{t+1} z_{t+1}$. Here, $y_t$ denotes the current weights of the model (note that the gradient is evaluated on $y_t$ ), while $x_t$ are the weights that will be used for evaluation. By recursively expanding $x_t$ to get that $x_t$ is an exponential average of $y_t$ for a constant $c_t$. However, ScheduleFree SGD couples the coefficients of momentum and weight averaging, which we decouple.

For most of the paper, you sweep the batch size between 2^6 and 2^13 with the exception of Figure 3, where the batch size is between 2^16 and 2^22. Is this because in this Figure you count the number of tokens, so practically batch size * tokens per batch?

Please note that as we explain in Sec 2.1, the batch size should be multiplied by context length to get the number of tokens per batch. So the range is still calibrated and we have made the x-axis to be Batch Size (#Tokens) instead of just Batch Size.

Did you try checking how close your thus predicted CBS tracks the true CBS for "unseen data points", i.e. test its predictive power?

Due to computation constraints, we were not able verify the forecasted CBS.

In the concluding remarks you write "Our findings contribute insights into [...] particularly highlighting the role of hyper-parameters such as learning rate scheduling and optimizer settings" (line 530). What "optimizer settings" are you referring to?

We are referring to Adam optimizers including $\beta_1, \beta_2$.

Nits and typos.

Thank you for your careful read and we have fixed all the typos. We updated the draft significantly to address the formatting issues. Please double check the current versions.

Doesn't scaling from 151M to 302M double the depth, and therefore Figure 1 involves multiple ways of scaling up models?

Please note that by default we use the scaling configurations in Tab. 2. So all 302M models except for Fig. 4 are scaled from 151M ones through enlarging depth instead of width.

Thanks for the detailed reply. I will respond point-by-point.

EWA improves model training efficiency

Maybe I am misunderstanding something here, but in the mentioned README.md (you mean the traj_2048.png, right?) you are not sweeping training lengths. For WSD you are showing three different lengths of the decay phase. My point is that WSD trains for 5000 steps while reaching the target before that (~4500 depending on the exact decay ratio). And your claim that EWA is better is based on "time-to-reach" a loss of 3.24. If I wanted to reach 3.24 as fast as possible with the WSD schedule, I would, e.g., try a training duration of 4500 steps (with 0.1 of the duration for the decay). My issue is that you are setting up WSD (or cosine) to train for T steps, but then compare the time to a result that happens before the T steps. An alternative comparison would be "I trained every schedule for 5000 steps, which schedule gives me the best performance after 5000 steps?" In your attached plot (traj_2048.png), it looks like constant+EWA is actually the worst here, with (at least) WSD (0.1) providing better performance. I don't mean to be too critical of the constant+EWA schedule, however, I don't think that the current evidence is enough to support a statement like "EWA consistently improves model training efficiency". I do believe that it is an interesting idea and getting rid of having to define a training duration is very relevant (I wonder, though, if it currently just trades one hyperparameter (training duration $T$) for another (EWA decay rate $\tau$). However, I believe the current claims are too bold for the provided evidence.

Diverging lines for CBS vs. context length

This might depend on the exact definition of CBS, no? E.g. if CBS uses 50% overhead, context length could affect the CBS. But I understand your point, though I still wonder whether larger models might show a different enough behavior such that the lines would diverge before the CBS (when using 20%). Should the reference in Section 2.3 be for Figure 4 (instead of Section 2.4)?

Effect of CBS definition

Thanks for the additional scaling laws. This is indeed interesting.

Width vs depth scaling

Thanks for the clarification. This was my mistake.

Plots

Thanks for revising the plots. The new legend in Figure 1 does indeed help. While I also appreciate Fig 5., my suggestion was more directed at adding the 20% overhead and highlighting CBS in all other plots. E.g., in Fig 1, it is hard for me to eyeball what exactly the CBS is. Perhaps adding a light gray shaded background region to indicate the 20% overhead could be helpful. And/or marking the CBS with a star instead of a circle for each line. Btw. I still think that some of the font sizes are quite small, e.g., the x-ticks in the bottom row of Fig 1. And the colored lines in Fig 2 have very similar and very bright colors, but the markers help.

Figure 1 only involves a single way of scaling up models.

I think I am misunderstand something here. In Figure 1, you are showing models between 85M and 1.2B parameters. According to Table 2, going from 151M to 302M doubles the number of layers (thus scales depth). Going from 151M to 604M doubles $d_{\text{model}}$ and the hidden size of the MLPs, while having the same number of layers. Isn't this width-scaling? So doesn't Fig 1 involve both width-scaling and depth-scaling models? I feel like I am misunderstanding the entire sentence perhaps?

Thank you for addressing the other points as well. There, I don't have any follow-up questions. I am curious about the weight decay questions raised by Reviewer SERg and will be following the discussion.

Thank you for your constructive feedback. We have addressed each of your points as follows:

The experimental scope is limited to models up to 1.2B parameters trained on C4. However, given the careful experimental design and clear theoretical analysis, I do not believe that these impact the validity of the findings.

Thank you for the note. Due to computational constraints, we cannot extend our study to a larger scale. But we have carefully controlled for different factors in the settings we considered.

It would be helpful to have a dedicated section discussing the limitations of both the theoretical analysis and empirical findings (e.g., Gaussian data distribution).

Thanks for the suggestion. Please note that we do include the discussion of our setting is limited to C4 data in the Conclusion section. We would include the following: The theoretical analysis assumes Gaussian data distribution and infinite-width neural networks, which may not hold for real-world, non-Gaussian, high-dimensional data or finite models, potentially affecting the generalizability of CBS scaling laws. Empirical findings are restricted by limited hyperparameter sweeps and computational resources, which may prevent a full exploration of optimal configurations, particularly across varied tasks and larger architectures.

It would be nice to use instead of * (asterisk) in line 135. The color contrast in Figure 9 can be improved.

Thank you for the suggestions and we will include your feedback into the final version.