Time Transfer: On Optimal Learning Rate and Batch Size In The Infinite Data Limit

Authors: Oleg Filatov, Jan Ebert, Jiangtao Wang, Stefan Kesselheim
Primary area: foundation or frontier models, including LLMs
Keywords: optimal scaling, muP, hyperparameter transfer, learning rate sensitivity, LLM

Abstract:
One of the main challenges in optimal scaling of large language models (LLMs) is the prohibitive cost of hyperparameter tuning, particularly learning rate $\eta$ and batch size $B$. While techniques like $\mu$P (Yang et al., 2022) provide scaling rules for optimal $\eta$ transfer in the infinite model size limit, the optimal scaling behavior in the infinite data size limit remains unknown. We fill in this gap by observing for the first time an intricate dependence of optimal $\eta$ scaling on the pretraining token budget $T$, $B$ and its relation to the critical batch size $B_\mathrm{crit}$, which we measure to evolve as $B_\mathrm{crit} \propto T$. Furthermore, we show that the optimal batch size is positively correlated with $B_\mathrm{crit}$: keeping it fixed becomes suboptimal over time even if learning rate is scaled optimally. Surprisingly, our results demonstrate that the observed optimal $\eta$ and $B$ dynamics are preserved with $\mu$P model scaling, challenging the conventional view of $B_\mathrm{crit}$ dependence solely on loss value. Complementing optimality, we examine the sensitivity of loss to changes in learning rate, where we find the sensitivity to decrease with increase of $T$ and to remain constant with $\mu$P model scaling. We hope our results make the first step towards a unified picture of the joint optimal data and model scaling.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730739902444

thread_type:
official_review

<thread_object>
creator: reviewer_c7lk

<content>
summary:
This paper investigates the optimal learning rate $\eta$ and batch size $B$ as a function of the data size $T$ in language model (LM) training, using the maximal update parametrization (muP). While muP focuses on hyperparameter transfer as model width increases (with other parameters fixed), this work centers on how hyperparameters evolve as $T$ grows, assuming a large but fixed model size (with results averaged across three model sizes). The authors show that the optimal $(\eta, B)$ can vary in complex ways with $T$; for example, there are regions where $\eta \propto 1/\sqrt{T}$ and others where $\eta \propto \sqrt{T}$. Overall, this work offers valuable insights into the optimal scaling of learning rate and batch size as data size increases.

strengths:
In modern LM training, scaling both model size and data size is essential. Understanding how optimal hyperparameters scale with these factors is crucial. While muP provides insights into how hyperparameters scale with model size, less attention has been given to how they scale with data size. This work contributes meaningfully in this direction. Through a set of experiments, the paper offers intriguing insights into how optimal learning rate and batch size scale with increasing data size.

weaknesses:
While this paper addresses an important problem and provides interesting insights, certain clarifications would help strengthen the presentation. Below are some suggestions for improvement:

- **Line 107:** It is mentioned that the "crucial batch size" is defined as the region where $\eta \propto \sqrt{B}$ breaks, and that this comes from Shallue et al. (2019). Could you please point to the specific section in Shallue et al. (2019) where this is discussed?
- **Line 122:** The notation $\epsilon = 1e^{-8}$ seems more like Python syntax. It would be clearer to revise this to $10^{-8}$.
- **Lines 188 and 294:** These lines refer to the "critical batch size" (as a function of data size) as the batch size that leads to the best validation error for a given data size, if I understand it correctly. This definition appears different from the one given in Line 107. Could you please clarify this discrepancy?
- **Lines 265-266:** The paper claims that $B^* \propto \sqrt{T}$. However, in Figure 2a, $B^*$ is either $2^{18}$ or $2^{20}$. How is this square root rule derived with only two data points?
- **Lines 294-296:** The claim that $B_{\text{crit}} \propto T$ is also unclear. Can you explain how this scaling is derived?


For all the claimed scaling laws, if possible, providing the observed data points along with using a scientific fitting method to derive the scaling laws would make the claims more convincing. 

I think the paper could be more clear if the speculated laws were better explained and justified.

questions:
See above please.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730791958721

thread_type:
official_review

<thread_object>
creator: reviewer_hncm

<content>
summary:
This paper studies the optimal scaling of LLMs in terms of the training hyperparameters, including learning rate $\eta$ and batch size $B$, and their optimal choices with respect to different data sizes $T$. In particular, they show that for different batch size $B$, the optimal $\eta$ has entirely different scaling with $T$. Furthermore, this paper also shows that the optimal batch size is positively correlated with $T$. Then, this paper also shows that the sensitivity with respect to $\eta$ is decreasing when using larger data size $T$. The results can help build a unified picture of the joint hyperparameter tunning of learning rate and batch size.

strengths:
This paper develops extensive experiments on training LLMs with different choices of learning rates, batch sizes, model sizes, and data sizes. Then based on the experimental results, this paper identifies some patterns regarding the optimal hyperparameter configurations. In particular, this paper shows that

* the optimal learning rate has a different evolution in the limit of the data sizes when choosing different batch sizes. 
* the optimal batch size is approximately proportion to $\sqrt{T}$ when choosing optimal $\eta$.
* the critical batch size is approximately proportion to $T$ that will affect the optimal choice of learning rate.
* The learning sensitivity can be mitigated when using larger $T$.

Overall these findings provide unified results for guiding the hyperparameter choices when having different token budgets.

weaknesses:
The major weakness of this paper is that it seems to be an experimental report that summarizes the key patterns observed from the experiments, while no further explanation and deeper analysis are conducted. Although there are extensive experimental results and the findings seem to be reliable, it remains unclear whether the findings can be generalized to other settings. I suggest the authors include some understanding of the empirical results, especially about why the learning rate can be much larger when  $T$ is large and we are using a large batch size.

Besides, I am not sure whether only considering $\mu P$ models is sufficient to understand the model with different sizes. It would be good to have more experiments and discussions on the models beyond the $\mu P$ ones.

The definition of $B_{crit}$ is a bit confusing. This may be a kind of clarity issue. The authors begin to mention $B_{crit}$ in the abstract and many places in the introduction. However, the reader cannot quickly get what's the formal definition of $B_{crit}$.

More details about the definition of the effective learning rate for Adam (line 428) should be provided.

questions:
See the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730977445204

thread_type:
official_review

<thread_object>
creator: reviewer_fvhf

<content>
summary:
The authors investigate optimal learning rate, optimal batch size and critical batch size as the number of training training tokens increases.  The authors also investigate the effects of changing mode size within $\mu$P and find that optimal learning rate and batch size do not vary with model size.  The authors also investigate the sensitivity of the loss to learning rate as data and mode size change finding decreased sensitivity as data increases, but similar relationships within $\mu$P scaling.

strengths:
Understanding how optimal learning rate and batch size scale as data scales is a very important and potentially impactful direction.  Understanding the diminishing returns of batch size significantly impact practitioner decisions regarding data parallelism.  Implications here can help go beyond compute optimal scaling paradigm that couples data and model size. In addition, strong understanding of relationships between hyperparameters can help reduce the need for expensive hyperparameter sweeps.

Aligning with existing theoretical models and providing the conjecture of eq. 1 can be useful modes of thinking to eventually provide a deeper understanding of these phenomena.

The scaling of batch size in data presented in Fig 2b and the sensitivity analysis of Fig 3 are both intriguing and useful for the community empirically.

weaknesses:
Overall, the writing needs work. The empirics of this paper are not comprehensive enough to make such a strong claim on training regimes as 3 different growth regimes are being implied from 6 batch sizes and 8 token budgets.  We essentially have 6 trajectories from 8 points and the scaling does not even start immediately because of the "lost" epochs (mainly looking at Fig 1a here).  Given this, the origin of these explicit scaling rules need to be justified by existing work or theory in my opinion.  Content from A.2 should be much more prominent in the main paper.  In addition, the intro feels misleading; the contributions section makes this sound like a theoretical contribution and the specific growth proposed isn't justified by empirics alone.

Fig 2: a 2^2 times budget increase in budget increase of 2^5 is oberved from the sequence with only one step increase?  Extrapolating a general rule from one bump in $B_{\text{crit}}$ is not convincing. In my opinion, it is valuable to look into this, but it is a detrimental to quantitative claims from such little evidence.

questions:
Fig 1a pattern in batch size is not monotone, what's going on here? 

Overall, my opinion of the paper is mixed but I lean towards acceptance.  I think much of the empirics are valuable to the community, but the paper is dragged down by overselling some weak empirics.  While I think it is completely reasonable that we don't have comprehensive answers to these questions (presumably from an onerous compute requirement), I'd like to see the claims tempered to mach the strength of the empirics.  I would recommend suggesting further work to investigate what looks like it could be a trend and provide some intuition on why the $T^{-0.5}$, constant and $ T^{0.5}$ regimes are a good conjecture, but hold off on a more forceful claim.

The content of A.2 and I think would be useful to integrate more directly in the main text.  

A more rigorous definition of critical batch size (line 107) would be useful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731064904437

thread_type:
official_review

<thread_object>
creator: reviewer_jan7

<content>
summary:
This paper studies how to set the learning rate $\eta$ and batch size $B$ when the total number of steps $T$ increases. A series of conclusions are drawn from the experiments on MPT of scale 32M, 101M and 354M:
1. A scaling law of the optimal learning rate given $B$. It depends on whether $B$ is smaller than a threshold $B_{crit}$.
2. A scaling law of the optimal batch size, assuming the learning rate is tuned to the optimal.
3. $B_{crit} \propto T$.

strengths:
1. This paper studies an important problem. Hyperparameter tuning is very costly in LLM pretraining. The optimal choices of these hyperparameters are particularly unclear when the number of steps $T$ is a variable.
2. Experiment results are well-organized.
3. A series of scaling laws have been derived. These laws could be useful for tuning the learning rate and batch size.

weaknesses:
My major concern is whether the experiment results can indeed support their scaling laws.
1. It is hard to see the claimed scaling of optimal LR from Figure 1. The curves in Figure 1 are very noisy, and the scaling formulas match the experiment results only for a few choices of $T$. It is thus unclear whether the scaling formulas here describe a real trend or just some noise.
2. It is even harder to see the claimed scaling of optimal batch size from Figure 2. In fact, the optimal batch size only changes once in the figure, which makes me wonder if any scaling concluded from just a single change of optimal batch size could be an overclaim.

Minor weaknesses:
1. It would be good to do some sanity checks on downstream tasks: Does the optimal LR for reducing the pretraining loss also lead to good performance on downstream tasks?
2. Experiments are conducted on only one dataset, C4.

questions:
I wonder if the authors could provide stronger empirical evidence to support the scaling of optimal LR and batch size proposed in the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732638999650

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General comment

comment:
Dear Reviewers,

We thank you for valuable and critical reviews, and apologise for the additional time required to approach the raised comments in a systematic way. In the general response, we would like to address the points which are common to all of the reviewers, and specific questions are answered separately for each of the reviewers. 

The major raised weakness is the gap between empirics and the claimed scaling rules. Also for us, the situation was unsatisfactory. Motivated by the constructive reviews we addressed it in a twofold way: (a) We added further experiments, mainly aiming at reducing the statistical uncertainty of the measured data points. (b) We added thorough analysis based on theory, and now have made major steps towards a comprehensive understanding of the intricate interplay of the different factors. Reflecting this, we have revised our presentation, rewriting the contributions to more accurately reflect our findings. 

Specifically, the gap between empirics and claims concerns the (inverse) square root scaling of the optimal learning rate $\eta^*$, the linear increase of the critical batch size $B_\mathrm{crit}$, and the square root increase of the optimal batch size $B^*$ (all with an increase of the token budget $T$). The latter (concluded from old Fig. 2a, now Fig. 3a) is indeed an overclaim and we toned it down to stating objectively that we only observe it increasing with $T$. For the former two, we approach the problem from a deeper perspective:

* We extend our experimental data to reduce statistical noise: Figure 1 now contains 4 additional runs per point for the four highest most sensitive to statistical fluctuations batch sizes. We added variance shading to old Figure 1a (now 1b, we swapped Fig. 1a and 1b in the updated version of the paper) to display this uncertainty.
* We now perform curve fits per token budgets T to data points of old Fig. 1b (now 1a) with the theoretical model of Li et al. (2024) (https://arxiv.org/abs/2405.14578) (Eq. 3.1) and update Fig. 1 with the model fit accordingly. 
* We empirically extend the fit model to incorporate the model parameters’ $\eta_\mathrm{crit}$, $B_\mathrm{crit}$ dependence on $T$. By fitting the power law to the measured $\eta_\mathrm{crit}$ and $B_\mathrm{crit}$ data points, we derive their dependence on the token budget $T$ and plot it on the additional Fig. 2. In the formulation of Li et al. (2024),  $B_\mathrm{crit}$ is the critical batch size, and from our fits we observe  $B_\mathrm{crit} \propto T^a$, where $a = 1.0 \pm 0.2$. This is compatible with the claim we made earlier, however now it is more grounded in the statistical method and underlying theory.
* We interpret the fitted model in various regimes of $B$ in relation to  $B_\mathrm{crit}$ and add the corresponding discussion to Sec. 3.2. Our initial claim of square-root / inverse square-root scaling is actually supported by fitted model, but only in the limit of constant $\eta_\mathrm{crit}$, and linear scaling of $b_\mathrm{crit}$ with $T$. We therefore reduce our previous claims in the contribution section to root them into a more objective interpretation of our fit model results. 
* We added limitations and future work paragraph in the Discussion section to reflect points raised in the reviews and to provide perspective and directions on the potential improvements.

The last common comment was regarding the critical batch size definition, which we agree was not sufficiently precise before. Since we now ground our interpretation of the data into the theoretical framework of Li et al. (2024), the definition of critical batch size becomes unambiguous, as we update accordingly in Sec. 2.1.    

We hope that moderating our claims, inclusion of additional supporting evidence, and deeper analysis of observation suffices to address the gap between empirics and the stated results. In comparison to before the improvement of our analysis based on the theory of Li et al., our paper in general was more descriptive than concise, but with the inclusion of the theory, we judge our own findings now as consistent and very encouraging.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper investigates the scaling behavior of optimal learning rate and batch size in the infinite data limit, analyzing their dependence on the pretraining token budget and critical batch size. While the reviewers appreciated the relevance of the topic and the authors' effort to address practical challenges in scaling large language models, significant weaknesses were noted. Key concerns included the lack of theoretical rigor and unclear explanations of empirical findings. Despite clarifications in the rebuttal, these issues remained unresolved. I recommend rejection at this time.