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

## Abstract

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 italic_η and batch size B 𝐵 B italic_B . While techniques like μ 𝜇 \mu italic_μ P (Yang et al., 2022 ) provide scaling rules for optimal η 𝜂 \eta italic_η transfer in the infinite model size limit, the optimal scaling behavior in the infinite data size limit ( T → ∞ → 𝑇 T\to\infty italic_T → ∞ ) remains unknown. We fill in this gap by observing for the first time an interplay of three optimal η 𝜂 \eta italic_η scaling regimes: η ∝ T proportional-to 𝜂 𝑇 \eta\propto\sqrt{T} italic_η ∝ square-root start_ARG italic_T end_ARG , η ∝ 1 proportional-to 𝜂 1 \eta\propto 1 italic_η ∝ 1 , and η ∝ 1 / T proportional-to 𝜂 1 𝑇 \eta\propto 1/\sqrt{T} italic_η ∝ 1 / square-root start_ARG italic_T end_ARG with transitions controlled by B 𝐵 B italic_B and its relation to the time-evolving critical batch size B crit ∝ T proportional-to subscript 𝐵 crit 𝑇 B_{\mathrm{crit}}\propto T italic_B start_POSTSUBSCRIPT roman_crit end_POSTSUBSCRIPT ∝ italic_T . Furthermore, we show that the optimal batch size is positively correlated with B crit subscript 𝐵 crit B_{\mathrm{crit}} italic_B start_POSTSUBSCRIPT roman_crit end_POSTSUBSCRIPT : keeping it fixed becomes suboptimal over time even if learning rate is scaled optimally. Surprisingly, our results demonstrate that the observed optimal η 𝜂 \eta italic_η and B 𝐵 B italic_B dynamics are preserved with μ 𝜇 \mu italic_μ P model scaling, challenging the conventional view of B crit subscript 𝐵 crit B_{\mathrm{crit}} italic_B start_POSTSUBSCRIPT roman_crit end_POSTSUBSCRIPT 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 T → ∞ → 𝑇 T\to\infty italic_T → ∞ and to remain constant with μ 𝜇 \mu italic_μ P model scaling. We hope our results make the first step towards a unified picture of the joint optimal data and model scaling.

## 1 Introduction

Large Language Models (LLMs) have increasingly become a prominent area of study in the field of Natural Language Processing (NLP) and beyond. They have demonstrated significant improvement in performance across a wide range of tasks, such as language understanding, text generation, translation, and summarization, showing results comparable or outperforming those of an average domain expert (Dubey et al., ; OpenAI et al., ; Team et al., ). The primary advantage of LLMs is their ability to scale well with increased computational resources, which results in predictive improved performance (Kaplan et al., ; Hoffmann et al., ).

One of the main challenges in LLM scaling lies in the proportional scaling of computational resources required for hyperparameter tuning. To remedy this, $\mu$Transfer (Yang et al., ) technique was proposed as a way to transfer hyperparameters from a small (proxy) model to a large (target) one by introducing scaling rules for learning rate, weight multipliers and initialization scale, altogether referred to as Maximal Update Parametrization ($\mu$P). While significantly reducing the hyperparameter tuning cost coming with model scaling, its applicability is limited by requiring both target and proxy models to share the same batch size and number of training iterations. With current pretraining budgets surpassing trillions of tokens, it makes $\mu$Transfer computationally expensive to apply even with tuning a small proxy model.

One solution would be hyperparameter tuning performed both for the small proxy model and on the small dataset, followed by $\mu$Transfer to the larger model and larger dataset, under assumption of both datasets being sampled from the same underlying data distribution. This raises the question of $\mu$Transfer’s applicability in the infinite data limit, which can be formalized as an increase in the size of the training dataset, which in the LLM case is measured by the number of tokens. Understanding training dynamics in this limit would unlock hyperparameter transfer not only across model scales, but also across data horizons, thus removing the largest limitation of $\mu$Transfer.

The study of optimal hyperparameter evolution throughout the model training should also be complemented with a study of hyperparameter sensitivity, i.e. the measure of how the model performance is affected when the training is performed outside the optimal hyperparameter range. In practice, it is rarely possible to remain within the optimum due to statistical uncertainties in its estimation. It would be of large interest to find training regimes which have small hyperparameter sensitivity and penalize model performance the least if the optimal hyperparameters are missed by a small degree.

Expanding on this line of research, we consider a commonly used LLM pretraining setup and aim towards building a yet missing holistic picture of optimal learning rate and batch size dynamics as one scales up the model training – both in the data and model sizes. Our main contributions are summarized as follows:

- •
Optimal learning rate scaling: depending on batch size $B$ and point in time during the training, we observe three regimes of the optimal learning rate $\eta^{*}$ evolution in the limit of increasing pretraining token budget $T\to\infty$ (Fig. 1(a)):
–
$\eta^{*}\propto 1/\sqrt{T}$ for $B<B_{\mathrm{crit}}$,
–
$\eta^{*}\propto 1$ for $B\sim B_{\mathrm{crit}}$,
–
$\eta^{*}\propto\sqrt{T}$ for $B>B_{\mathrm{crit}}$,
where transitions between them are controlled by $B$ relative to the critical batch size $B_{\mathrm{crit}}$ (see Sec. 2.1 for definition and Appendix A.2 for general discussion). The region boundaries are dynamically evolving in time following the critical batch size evolution $B_{\mathrm{crit}}(T)$. Furthermore, we find these dynamics to be largely preserved within $\mu$P (Appendix A.9).
- •
Optimal batch size scaling: assuming $\eta$ is optimal for a given data horizon $T$, we observe an approximately $B^{*}\propto\sqrt{T}$ drift of the optimal batch size $B^{*}$ (Fig. 2(a)). The drift is correlated with the drift of the $B_{\mathrm{crit}}$ region, with $B^{*}$ falling within the region. Importantly, we show that naïve application of optimal $\eta$ scaling rules in the $T\to\infty$ limit with $B$ being indefinitely fixed becomes suboptimal over time: a joint $(\eta,B)$ scaling is required.
- •
Critical batch size: $B_{\mathrm{crit}}$ evolves in time with approximately $B_{\mathrm{crit}}\propto T$ linear dependency on the training token budget (Fig. 1(b)). This dynamics also affects optimal $\eta$ scaling regimes, driving the time transition between them. Surprisingly, we show evidence that $B_{\mathrm{crit}}$ is not exclusively defined by the value of the loss function (Eq. 7) as suggested by McCandlish et al. (): models within $\mu$P share the same $B_{\mathrm{crit}}$ region while having different performance in terms of loss.
- •
Learning rate sensitivity: the sensitivity is generally decreasing with an increase of the training token budget, which is interestingly more pronounced for the batch sizes in the critical batch size region (Fig. 3). We observe no significant change in the learning rate sensitivity with the change of the $\mu$P base model and within the $\mu$P width limit (Fig. 4).

## 2 Methodology

### 2.1 Terminology

Time $\boldsymbol{(T)}$: we often use the terms time, token budget, and data horizon interchangeably, both to specify the measure of the training data size in tokens, and to pinpoint the specific moment throughout the model training. From this perspective, an infinite data limit $T\to\infty$, as opposed to a fixed budget regime with $T=\mathrm{const}$, refers to an (infinite) increase of the number of tokens seen by the model during pretraining.

$\boldsymbol{\mu}$P: we refer to a model with width $d_{\mathrm{model}}^{\mathrm{base}}$ as a base model if $\mu$P scaling multipliers for learning rates, weight multipliers and initialization scale (Sec. 2.2) are computed relative to this width. This brings us to a broader view on $\mu$P where the base model “pinpoints” the training dynamics for all the other models obtained either by scaling up or down the base $d_{\mathrm{model}}^{\mathrm{base}}$ width. Together with the base model, we refer to this set of models as a $\mu$P model family or as a $\mu$P trajectory if the direction of scaling is implied. We also slightly distinguish between the base and proxy models, where the former is used to define a $\mu$P model family, while the latter is a model used to tune hyperparameters to be transferred with $\mu$Transfer to a target model.

Critical batch size $\boldsymbol{(B_{\mathrm{crit}})}$: similarly to the SGD study of Shallue et al. (), we define $B_{\mathrm{crit}}$ as the region where the $\eta\propto\sqrt{B}$ scaling rule for a fixed token budget breaks and, as we observe, directly switches to a $\eta\propto 1/\sqrt{B}$ scaling rule. This corresponds to the peak of the bell-shaped curve (Fig. 1(b)), which was shown by Li et al. () to equal the $B_{\mathrm{crit}}$ definition of McCandlish et al. () (see Appendix A.2 for extended discussion).

Sensitivity: as acknowledged by Wortsman et al. (), it is difficult to formalize this notion, also in the absence of a theory to be verified. We therefore define it in the most minimal way, namely as the variation of validation loss $\mathcal{L}_{\mathrm{val}}(\eta)-\mathcal{L}_{\mathrm{val}}(\eta^{*})$ for a given learning rate variation from its optimal value $\eta/\eta^{*}$. We refer to the corresponding loss vs. learning rate curve (both with and without $\mathcal{L}_{\mathrm{val}}(\eta^{*})$ normalization) as a loss profile.

### 2.2 Model configuration and datasets

For all our experiments we use a default MPT model architecture (MosaicML, ) as implemented in the llm-foundry codebase (MosaicML, ), with all the models sharing the same training configuration (Appendix A.3). We use the Decoupled AdamW optimizer (Loshchilov & Hutter, ) with $\beta_{1}=0.9$, $\beta_{2}=0.95$, $\epsilon=1e^{-8}$, weight decay $\lambda=0$ and gradient clipping by the $L_{2}$ norm value of 1.

$\mu$P is implemented according to Table 8 of Yang et al. (), so that when $d_{\mathrm{model}}$ is set to the base model width $d_{\mathrm{model}}^{\mathrm{base}}$, it replicates Standard Parametrization (SP). That makes our observations for the base models also applicable to setups that use SP rather than $\mu$P. Model weights are initialized from the normal distribution with the base model standard deviation $\sigma^{\mathrm{base}}=1/\sqrt{d_{\mathrm{model}}^{\mathrm{base}}}$. The models are scaled up/down only in width, with the head dimension $d_{\mathrm{head}}$ being always fixed and the number of heads being scaled proportionally to the width scaling.

The models are trained with the causal language modeling task on the train split of the Colossal Clean Crawled Corpus (C4) dataset (Raffel et al., ), tokenized with the GPT2 tokenizer (Radford et al., ) with a vocabulary size of 50257 and a context length of 1024 tokens. As a metric to evaluate model performance, we report the loss on the C4 validation split as $\mathcal{L}_{\mathrm{val}}$.

### 2.3 Hyperparameter grid

To investigate the interplay of learning rate and batch size in the infinite data limit $T\rightarrow\infty$, we define a 5D grid spanned by the following axes: $\eta$, $B$, $T$, $d_{\mathrm{model}}$, $d_{\mathrm{model}}^{\mathrm{base}}$ (see Appendix A.4 for exact definition). Fundamentally, we are interested in measuring how the loss profile $\mathcal{L}_{\mathrm{val}}(\eta)$ and its optimum value $\eta^{*}$ evolve in time $T$ depending on the choice of batch size $B$. As this measurement is moreover conditioned on the $\mu$P trajectory and a specific point therein, we firstly study this evolution for a trajectory pinpointed by one specific base model with $d_{\mathrm{model}}^{\mathrm{base}}$. We train a set of models within the defined $\mu$P trajectory with different widths $d_{\mathrm{model}}$, ranging in size from 32M up to 354M parameters, and measure for each of them the $\mathcal{L}_{\mathrm{val}}(\eta)$ profile at specific points in time $T$, ranging from 1B up to 275B tokens. Then, we repeat the same measurement for a new $\mu$P trajectory, pinpointed by a different value of $d_{\mathrm{model}}^{\mathrm{base}}$. This grid approach allows us to interpret results from multiple perspectives, as we detail in Sec. 3.

### 2.4 Learning rate schedule scaling

Since we study the training dynamics in the infinite data limit, it necessarily implies training models across different data horizons. This raises the question of how one should adjust the learning rate schedule in this limit. Motivated by recent work of Hu et al. (); Hägele et al. (), in all our experiments we use a warmup-stable (WS) version of the warmup-stable-decay (WSD) schedule consisting of a warmup phase with a linear increase of learning rate from 0 to $\eta_{\mathrm{max}}$ and a constant phase with learning rate fixed at $\eta_{\mathrm{max}}$, hereafter notated as $\eta$. Our version omits the decay phase to simplify experimentation as we observe that it does not affect the optimal $\eta$ position (Appendix A.6). The warmup duration is fixed across all horizons and across all experiments at an absolute value of $T_{\mathrm{warmup}}=2^{19}=524288$ tokens. Whenever batch size is varied, we adjust the number of gradient steps in the warmup phase accordingly so that the total amount of tokens seen by the model during warmup equals $2^{19}$. We also present additional experiments with different ways to scale the learning rate warmup and an added decay phase in Appendix A.6, with results largely confirming those of Hägele et al. (). The WS schedule allows us to reduce computational requirements by approximately a factor of two: contrary to retraining for each of the data horizons in the $T$ grid, we run indefinitely continued trainings and take evaluation snapshots on the way.

## 3 Results

### 3.1 Learning rate optimum drifts in time, with batch size interpolating between different scaling rules

Figure: (a)
Refer to caption: x1.png

First, we begin with setting $d_{\mathrm{model}}^{\mathrm{base}}=1024$ and scanning learning rate across different batch sizes and $d_{\mathrm{model}}$. We present results for the $\eta^{*}$ optimum evolution in time $T$ in Fig. 1(a) and full $\mathcal{L}_{\mathrm{val}}(\eta)$ profile scans in Appendix A.7. In order to reduce statistical uncertainties, in Fig. 1(a) each data point for budgets $T\leq 2^{35}$ is an average $\eta^{*}$ value across $\mu$P models for a given batch size and horizon length $(\overline{\eta^{*}},B,T)=\sum_{d_{\mathrm{model}}}(\eta^{*},B,T,d_{\mathrm{
model}})/3$, where $d_{\mathrm{model}}\in\{256,~{}512,~{}1024\}$ and all models share the same base model with $d_{\mathrm{model}}^{\mathrm{base}}=1024$(^1^11We believe this averaging approach is justified since all the three models share the same optimization trajectory in terms of the number of steps, batch size and data horizon length, therefore are theoretically guaranteed by $\mu$Transfer to share the same optimal learning rate. From the experimental side, we also observe no significant differences across the three models (Appendix A.9).).

Fig. 1(a) illustrates how batch size serves as a parameter interpolating between various learning rate scaling regimes. For the smallest probed batch sizes $B=\{2^{16},2^{18}\}$ we observe an initially perfect inverse square-root scaling of the optimal learning rate $\eta^{*}$ with increase of the data horizon $T$. The scaling, however, breaks down at some point and optimal learning rate plateaus. We speculate that this may be related to gradient moments accumulation and the fact that we apply no scaling rules for the $\beta_{1,2}$ AdamW parameters, in order to reduce analysis complexity. Such rules were derived by Hilton et al. (); Malladi et al. (), although only in the fixed data budget setup. Understanding the ways of extending this scaling regime further in time provides an interesting direction for future studies.

The $\eta^{*}$ drift in time is not per se a new phenomenon: a linear rule for SGD (Smith et al., ) and a square-root rule for Adam (Shen et al., ) are already known. We complement this picture with observing, for the first time, richer dynamics: namely, a continuous transition from inverse square-root to square-root scaling regimes. As the batch size increases, one first observes (inverse) constant-to-moderate scaling for $B=2^{20}(2^{22})$ which we also refer to as a metastable regime: to a certain degree, optimal learning rate transfers in time before breaking into one of the square-root regimes. Then, the metastable regime changes to square-root scaling for larger batch sizes $B=\{2^{24},~{}2^{26}\}$. The scaling does not appear immediately after the beginning of the training but rather after some time. We link this with the effect of “lost” epochs (Smith et al., ), which states that larger batch sizes require a larger number of iterations for the gradient accumulation to reach steady convergence dynamics. In our case, the transition to the square-root scaling happens only after iterating through $2^{34}\text{--}2^{35}\approx 17\text{--}34\mathrm{B}$ tokens and takes up to $137\mathrm{B}$ tokens to establish.

### 3.2 Optimally-tuned batch size increases in time

Figure: (a)
Refer to caption: x3.png

Second, we study how optimal hyperparameter values evolve in time to yield optimal loss values. For each batch size and horizon length, we select the best-performing run across the learning rate grid and plot model loss $\mathcal{L}_{\mathrm{val}}$ against batch size across time horizons for the configuration with $(d_{\mathrm{model}}=1024,~{}d_{\mathrm{model}}^{\mathrm{base}}=1024)$. Results are presented in Fig. 2, with a full set of plots across various combinations of $(d_{\mathrm{model}},~{}d_{\mathrm{model}}^{\mathrm{base}})$ in Appendix A.10.

We observe an approximate square-root scaling of the optimal batch size with increase of the pretraining token budget from $B^{*}|_{T=2^{30}}=2^{18}$ to $B^{*}|_{T=2^{35}}=2^{20}$ (Fig. 2(a)). Emergence of suboptimality is more pronounced when transposing the token budget and batch size axes (Fig. 2(b)), where the smallest $B=2^{16}$ batch size curve, with each point having learning rate scaled with the inverse scaling rule $\eta^{*}\propto 1/\sqrt{T}$, is being taken over in the $T\rightarrow\infty$ limit by the curves corresponding to larger batch sizes.

This result illustrates that, while naïve “pairwise” scaling rules for optimal learning rate, e.g. $\eta^{*}\propto 1/\sqrt{T}$, are convenient for predicting optimal values at scale, they do not necessarily result in the best model performance: taking batch size dynamics into account is required. In other words, the invariant induced solely by the $\eta^{*}\propto 1/\sqrt{T}$ scaling rule is not sufficient for the model performance to be optimal. We believe, similarly to Smith & Le (), that some broader notion of noise scale should serve as a more fundamental invariant to optimize for in the joint data and model size limit. We discuss this idea in more detail in Sec. 4.

### 3.3 Critical batch size region evolves in time, but is unchanged within μ 𝜇 \mu italic_μ P

In Fig. 1(b), we reinterpret Fig. 1(a) by transposing batch size and token budget axes and plotting the optimal learning rate and batch size scaling jointly per data horizon. This better emphasizes dynamics in the fixed token budget regime, rather than the trajectory of optimal learning rate evolution in time. As in Sec. 3.1, we perform an average across $\mu$P models sharing the same $d_{\mathrm{model}}^{\mathrm{base}}=1024$ and also plot the corresponding standard deviation. We include a similar plot for the other base model with $d_{\mathrm{model}}^{\mathrm{base}}=256$ in Appendix A.8 and individual plots for each of the $(d_{\mathrm{model}},~{}d_{\mathrm{model}}^{\mathrm{base}})$ configurations in Appendix A.9.

We observe that for a given time horizon, the $(\eta^{*},B)$ curve has a bell-like shape, as predicted by Li et al. (). The left-hand side of the peak represents a previously known $\eta\propto\sqrt{B}$ scaling rule (Malladi et al., ; Shen et al., ). However, with our experiments, we uncover a previously unseen right-hand side of the curve, also referred to as “surge” by Li et al. (), where the optimal learning rate for a fixed token budget scales inversely proportionally to the batch size scaling via the $\eta^{*}\propto 1/\sqrt{B}$ rule. Although statistical uncertainties are high for the experiment with $d_{\mathrm{model}}^{\mathrm{base}}=1024$, the square-root rules are more pronounced for the experiment with $d_{\mathrm{model}}^{\mathrm{base}}=256$ (Appendix A.8).

Furthermore, the peak position of the fixed token budget, which we refer to as the critical batch size $B_{\mathrm{crit}}$ (Sec. 2.1), is evolving in time via an approximately $B_{\mathrm{crit}}\simeq T$ scaling rule, showing the same trend as in McCandlish et al. (). Moreover, we note that the optimal learning rate corresponding to $B_{\mathrm{crit}}$ decreased by a factor of two with a horizon scaling from $2^{30}$ to $2^{37}$ tokens. This might be a statistical fluctuation since there is no significant decrease for the experiment with $d_{\mathrm{model}}^{\mathrm{base}}=256$ (Appendix A.8).

Lastly, there is a difference of $B_{\mathrm{crit}}$ evolution between the $T$ and $\mu$P infinite width limits. Specifically, for a fixed token budget, we observe no significant change of the curves’ shapes and peak positions across $d_{\mathrm{model}}$ values within the same $\mu$P trajectory, and also with the change of the base model (Appendix A.9). At the same time, there is a noticeable drift of $B_{\mathrm{crit}}$ in the $T\to\infty$ limit with the model being fixed. As both limits are accompanied with a comparable change of the model performance(^2^22Back-of-the-envelope calculation from Fig. 2(a) and Appendix A.10: for $d_{\mathrm{model}}^{\mathrm{base}}=1024$, $B=2^{20}$, there is a loss change $\mathcal{L}_{\mathrm{val}}=3.4\to 2.8$ with a token budget increase $2^{31}\to 2^{37}$, resulting in $B_{\mathrm{crit}}$ drifting by $2^{4}$. For the same $(d_{\mathrm{model}}^{\mathrm{base}},B)$ configuration, there is no significant $B_{\mathrm{crit}}$ drift with a change of width by $2^{2}$ within $\mu$P, but the corresponding loss change is $\mathcal{L}_{\mathrm{val}}=3.5\to 2.9$.), this observation brings evidence that dependence of the critical batch size exclusively on the loss value suggested by Kaplan et al. () (Eq. 7) is not entirely complete. Or, contrary to experimental results in Li et al. (), the two definitions of the critical batch size region (Appendix A.2) are not the same and should be disentangled.

### 3.4 Learning rate sensitivity is reduced in time, and is unchanged within μ 𝜇 \mu italic_μ P

Figure: Figure 3: Learning rate sensitivity $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of the learning rate deviation from the optimal value $\eta/\eta_{\mathrm{optimal}}$, measured for batch sizes of $B=2^{18}$ (left column), $2^{20}$ (middle column), and $2^{22}$ (right column) tokens, separately for the $\mu$P base models with width $d_{\mathrm{model}}^{\mathrm{base}}=256$ (top row) and $1024$ (bottom row). The former model amounts to 32M and the latter to 354M trainable parameters. With an increase of the pretraining token budget (different marker styles) we observe a general decrease in the learning rate sensitivity, which is more pronounced for batch sizes $B\in\{2^{20},2^{22}\}$ in the critical region (Sec. 2.1) and for the 354M model. At the largest probed token budget $T=2^{35}$ tokens, the sensitivity equalizes across the models and batch sizes.
Refer to caption: x5.png

Figure: Figure 4: Learning rate sensitivity $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of learning rate $\eta$, measured for batch sizes of $B=2^{18}$ (leftmost column), $2^{20}$ (middle left column), $2^{22}$ (middle right column) and $2^{24}$ (rightmost column) tokens, separately for the $\mu$P base models with the width $d_{\mathrm{model}}^{\mathrm{base}}=256$ (top row) and $1024$ (bottom row). Different marker styles correspond to different models within the $\mu$P family, with all the models being evaluated at the data horizon of $T=2^{35}$ tokens. For the base model with $d_{\mathrm{model}}^{\mathrm{base}}=256$, we scale the width only downwards, while for the base model with $d_{\mathrm{model}}^{\mathrm{base}}=1024$, we scale it both upwards and downwards. We observe no significant difference in the sensitivity across all the $(d_{\mathrm{model}}^{\mathrm{base}},d_{\mathrm{model}})$ configurations. Note that for the configuration $(B=2^{24},~{}d_{\mathrm{model}}^{\mathrm{base}}=1024)$, the base and $d_{\mathrm{model}}=4\times d_{\mathrm{model}}^{\mathrm{base}}$ models share a different random seed compared to all the other models, to illustrate the loss penalty arising from the learning rate optimum variation.
Refer to caption: x6.png

After having studied the learning rate optimum dynamics, we turn our attention to a broader structure around the optimum from the sensitivity perspective. Specifically, we are interested in how the shape of the $\mathcal{L}_{\mathrm{val}}(\eta)$ curve changes in the time $T\to\infty$ and $\mu$P width limits. In Fig. 3, we present our observations for the two base models with $d_{\mathrm{model}}^{\mathrm{base}}=d_{\mathrm{model}}\in\{256,1024\}$, for token budgets $T\in\{2^{31},~{}2^{33},~{}2^{35}\}$. We note that since we implement $\mu$P in a way that the base model is also SP-parametrized, the results should be applicable to this parametrization as well.

We observe that there is a general decrease in the learning rate sensitivity by up to $2^{1}$ per each token budget increase by $2^{2}$ as measured by $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ value, where $\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}=\mathcal{L}_{\mathrm{val}}(\eta^{*})$ is the validation loss value in the learning rate optimum. This indicates that the model profits from longer training by having lower penalty for the misspecification of the optimal learning rate. Notably, the decrease is more pronounced for batch sizes in the critical region ($B=2^{20}$ and $2^{22}$), while for the region with the $\eta^{*}\propto 1/\sqrt{T}$ scaling rule ($B=2^{18}$), the effect is either reduced (base model $d_{\mathrm{model}}^{\mathrm{base}}=1024$) or shows asymmetric trends w.r.t. the learning rate optimum (base model $d_{\mathrm{model}}^{\mathrm{base}}=256$). However, within our measurement precision, the sensitivity evens out across batch sizes for the longest $2^{35}$ token horizon. Overall, our results motivate the choice of the training regime within the critical batch size region in order to minimize the risks of under- or overshooting the learning rate optimum. As we show in Appendix A.5, the learning rate optimum position can vary by a factor of two just depending on the random seed choice.

With respect to the $\mu$P width limit, we observe no significant deviation of the loss profile from the one of the base model, both for up- and down-scaled models within $\mu$P (Fig. 4 with and Fig. 18 without $\mathcal{L}_{\mathrm{val}}$ normalization). Evaluated for the data horizon of $T=2^{35}\approx 34$B tokens, this holds across the models with the number of trainable parameters ranging from 32M up to 5B. Likewise, changing the base model does not affect the profile shape, except for the optimum learning rate shift by $\times 2$, which is expected for the base models compared here due to our $d_{\mathrm{model}}^{\mathrm{base}}$-dependent weight initialization scheme (Sec. 2.2).

## 4 Discussion

While originally, we were aiming to find a golden recipe for hyperparameter transfer in the infinite data limit, we show that there is no simple and straight-forward answer. However, we do believe that there exists a deeper underlying perspective on the problem, as opposed to the one of simply tuning learning rate and batch size.

Fundamentally, the field of model parametrization research has originated from and is further converging towards preserving some notion of norm in some infinite (model width and/or depth) limit (Everett et al., ; Yang et al., ; Large et al., ). In fact, any parametrization itself is simply a set of scaling rules to be applied to hyperparameters in order to preserve these norms (e.g. of model weight matrices or weight updates). Expanding on this, one can argue that scaling rules follow from the requirement of keeping some underlying quantity invariant within the infinite limit. From this perspective, hyperparameter transfer is nothing but a consequence of such “conservation laws”.

With this perspective in mind, we draw a parallel between infinite model and data limits, and speculate that a similar notion of “norm” should exist and should be aimed to be preserved in the infinite data limit. In fact, there is already a good candidate for this, namely the noise scale (Eq. 4 and 8), which intriguingly also induces scaling rules for hyperparameters (see Appendix A.2 for in-depth discussion). However, the existing definition neither takes into account the adaptive nature of the optimizer, nor the scenario of jointly following the infinite data and model limits. From our experimental observations, the following definition might be an applicable extension:

$$ $B_{\mathrm{noise}}^{\mathrm{Adam}}\propto\eta_{\mathrm{eff}}\frac{\sqrt{T}}{ \sqrt{\frac{B_{\mathrm{crit}}(T,M)}{B}}+\sqrt{\frac{B}{B_{\mathrm{crit}}(T,M)} }},$ (1) $$

where $B_{\mathrm{crit}}(T,M)$ is the critical batch size as a function of time $T$ and model size contribution $M$ via the corresponding model parametrization, as motivated by Park et al. (); $\eta_{\mathrm{eff}}=\eta(\eta_{\mathrm{max}},\beta_{1},\beta_{2})$ is the effective learning rate, incorporating additional dependency on Adam’s $\beta_{1,2}$ parameters. In our experiments with $\mu$P, it appears that $B_{\mathrm{crit}}(T,M)=B_{\mathrm{crit}}(T)$ as we do not observe any dynamics changing across the $\mu$P model family (Sec. 3.3). Once the invariant is established, one can derive the corresponding scaling rules for $T,M,B,\eta$ in the joint $(M,T)\to\infty$, effectively resulting in the hyperparameter transfer.

The main limitation of our work is that we provide only experimental hints for Eq. 1. This formula is the fruit of empirical observation only, which makes it a scaling law and should not be taken as a theorem. Nonetheless, we hope that our results make the first step towards the unification of infinite data and model size limits via deriving such a joint scaling invariant, also inclusively across multiple model parametrizations. Finally, our insights into optimal scaling rules for learning rate and batch size might be valuable for practitioners who approach the problem of hyperparameter optimization in the infinite data and model size limit. We provide our summary and recommendations in Appendix A.1.

## 5 Related Work

$\boldsymbol{(\eta,B)}$ scaling rules    In efforts to accelerate model training, the $\eta\propto B$ rule for the SGD optimizer was found necessary to avoid performance loss due to increased batch size (Goyal et al., ), known as generalization gap (Keskar et al., ). Afterwards, additional usage of momentum (Smith et al., ) and model scaling (Park et al., ) was incorporated, and a $\eta\propto\sqrt{B}$ rule for Adam was observed (Hilton et al., ). From the theoretical side, experimentally observed rules were verified with the framework of stochastic differential equations (SDEs) (Smith & Le, ; Malladi et al., ), loss curvature analysis (Zhang et al., ; McCandlish et al., ; Li et al., ) and random matrix theory (Granziol et al., ). While most of the studies were performed in the fixed epoch budget, Shallue et al. () broadened the perspective to other target budget measures and studied the scope of the $\eta\propto B$ rule applicability across various datasets and model architectures. Looking beyond fixed budgets, Smith & Le () showed a linear relation between the optimal batch size and the dataset size (for fixed $\eta$), and Smith et al. () similarly presented hints for a linear relation between the optimal learning rate and the dataset size (for fixed $B$), with both works considering the SGD optimizer. In the modern LLM pretraining context, Hu et al. (); DeepSeek-AI et al. () approached this problem by deriving the joint $(\eta,B)$ scaling laws.

$\boldsymbol{\mu}$P    Originally developed within the Tensor Program series studying feature learning in the infinite width limit (Yang & Hu, ; Yang et al., ), $\mu$P has been gaining traction recently within the LLM community. It has been extensively tested and applied experimentally (Lingle, ; Blake et al., ; Gunter et al., ; Dey et al., ), as well as theoretically, with Yang et al. (); Bordelon et al. () extending it to the infinite depth limit, and Yang et al. (); Bernstein et al. () revisiting it from the spectral normalization perspective. Recently, Everett et al. () showed that other model parametrizations also induce hyperparameter transfer if taking weight alignment into account. Furthermore, they revealed that $\mu$Transfer does not work in the regime of Chinchilla-optimal scaling (Hoffmann et al., ). The most closely related work to ours, Shen et al. () expanded on this observation and proposed a learning rate scheduler combining $\mu$P and experimentally measured $(\eta,B)$ scaling rules to allow for the hyperparameter transfer in the $T\to\infty$ limit, however only limited to the $\eta^{*}\propto 1/\sqrt{T}$ scaling regime.

Sensitivity    The topic of loss sensitivity to suboptimal hyperparameter choice is less thoroughly studied, focusing exclusively on learning rate as the most affecting hyperparameter. Wortsman et al. () studied how various optimizer and model interventions, such as weight decay or $\mu$P usage, influence the learning rate sensitivity with the model size scaling. Hägele et al. () investigated the impact of various learning rate schedule choices, such as length and functional form of the decay phase.

## 6 Conclusion

In this work, we studied joint model and data scaling in the LLM context from the perspective of optimal learning rate and batch size dynamics. We observed an intricate interplay of three optimal learning rate scaling regimes in the infinite data limit, controlled by the batch size in its relation to the critical batch size as it evolves in time. This dynamic is preserved during model scaling with $\mu$P, as well as the loss sensitivity to the learning rate variation, highlighting the intriguing difference in how $\mu$P infinite width and time limits evolve the critical batch size. Overall, we hope our observations pave the way towards deeper understanding of the optimal scaling in the unified infinite data and model size limit.

## Acknowledgments

In alphabetical order, we thank Andrei Filatov, EleutherAI community, Ismail Khalfaoui Hassani, Lucas Dax Lingle, and Stepan Zakharov for helpful discussions and feedback on the manuscript. This research was supported by TrustLLM funded by Horizon Europe GA 101135671, by the Helmholtz Foundation Model Initiative as a part of the Synergy Unit, and by Helmholtz AI computing resources (HAICORE) of the
Helmholtz Association’s Initiative and Networking Fund through Helmholtz AI. We gratefully acknowledge the Gauss Centre for Supercomputing e.V. ([www.gauss-centre.eu](www.gauss-centre.eu)) for funding this project by providing computing time on the Supercomputers JUWELS and JURECA at Jülich Supercomputing Centre (JSC). Parts of computational resources were provided by the German AI service center WestAI.

## Appendix A Appendix

### A.1 Hyperparameter optimization in the infinite data and model size limit

We believe our observations provide useful hints on how to scale learning rate and batch size jointly in the infinite data and model size limits. We take the general $\mu$Transfer approach of tuning hyperparameters for a small proxy model and then transferring them either zero-shot or according to some scaling rules via extrapolation, across model sizes and data horizons.

- 1.
If one can afford tuning a $\mu$P proxy model on the data horizon of the target model, then it is sufficient to simply perform a grid search over learning rate and batch size values to find the best combination, following $\mu$Transfer (Yang et al., ). As we describe in Sec. 5, $\mu$Transfer has been established to successfully transfer hyperparameters to O($10\mathrm{B}$) model sizes, albeit with potential limitations arising from very long range extrapolation in the infinite width limit (Blake et al., ; Gunter et al., ).
- 2.
Otherwise, a proxy model has to be tuned on a shorter data horizon than the target one. In that case, we suggest running a 2D grid search across learning rate and batch size values roughly around the optimal ones, where each training follows a WSD schedule (Sec. 2.4), for as long as compute budget allows. We suggest both the warmup and decay of the schedule to be fixed to the one of the target model in absolute number of tokens, which in turn should be about $10\text{--}20\%$ fraction of the target model horizon to be optimal (Kosson et al., ; Hägele et al., ). This is due to the observed drift of the learning rate optimum with the change of the number of steps (Appendix A.6). It is still not yet clear how scaling of warmup/decay length and Adam’s $\beta_{1,2}$ parameters (which we keep constant in our experiments) can be incorporated into the total horizon scaling. We leave this as an interesting direction for future work.
- 3.
After the grid search, one should be able to obtain a plot similar to Fig. 1(b) and Fig. 2(a). Provided long enough WSD horizon, a drift in time of the critical batch size region, associated to the peak of the fixed token budget curve in Fig. 1(b), should be visible. Likewise, there should be a drift of the optimally tuned (i.e. assuming optimal learning rate is used) batch size in time as in Fig. 2(a). Since we observe a strong correlation but still a mismatch between the optimally-tuned batch size and the critical batch size, we suggest the following approach for selecting optimal hyperparameter values:
(a)
Derive scaling rule by extrapolating the batch size optimum drift in time $T$ based on Fig. 2(a) (in our case, approximately $B^{*}\propto\sqrt{T}$). Estimate the expected optimal batch size value $B^{*}_{\mathrm{target}}$ for the target data horizon $T_{\mathrm{target}}$ under assumption of the optimally tuned learning rate.
(b)
Extrapolate the drift of the fixed budget curve in time based on Fig. 1(b) (in our case, approximately $B^{\mathrm{crit}}\propto T$) and derive the expected critical batch size for the target horizon $B^{\mathrm{crit}}_{\mathrm{target}}$.
(c)
Set optimal learning rate for the target horizon as:
$\eta^{*}_{\mathrm{target}}=\begin{cases}\vspace{2mm}\eta^{\mathrm{crit}}\cdot
\sqrt{B^{*}_{\mathrm{target}}/B^{\mathrm{crit}}_{\mathrm{target}}}&\text{if $B
^{*}_{\mathrm{target}}\leq B^{\mathrm{crit}}_{\mathrm{target}}$}\\
\vspace{2mm}\eta^{\mathrm{crit}}\cdot\sqrt{B^{\mathrm{crit}}_{\mathrm{target}}
/B^{*}_{\mathrm{target}}}&\text{if $B^{*}_{\mathrm{target}}>B^{\mathrm{crit}}_
{\mathrm{target}}$}\end{cases},$
(2)
where we correct the learning rate value for the corresponding scaling regime, assuming the optimum learning rate corresponding to the critical batch size remains constant in time $\eta^{\mathrm{crit}}=\eta^{*}(T)|_{B^{\mathrm{crit}}(T)}=\mathrm{const}$.
- 4.
Apply optimal values of learning rate $\eta^{*}_{\mathrm{target}}$ and batch size $B^{*}_{\mathrm{target}}$ to the target model, scaled up with $\mu$P, and to the target training horizon. As we show in this work, $\mu$P does not impact the dynamics of the critical batch size evolution in the infinite data limit, therefore we expect no interference between the two limits.

We suppose it is also possible to adjust the recipe above to the continual learning setting (Çağatay Yıldız et al., ; Ibrahim et al., ): under assumption of $\eta^{\mathrm{crit}}$ being constant in time and of the golden path hypothesis (Vyas et al., ), one could indefinitely run the model training with the same learning rate but dynamically adjust the batch size to follow the critical one (peak of the fixed budget curve in Fig. 1(b)), or, alternatively viewed, to remain on the pareto curve of Fig. 2(b) (inset plot).

### A.2 On critical batch size and noise scale

There are two perspectives on the critical batch size $B_{\mathrm{crit}}$. Firstly, McCandlish et al. () define it as a batch size which results in an optimal trade-off between data sample efficiency and gradient update step efficiency:

$$ $B_{\mathrm{crit}}\coloneq\frac{E_{\mathrm{min}}}{S_{\mathrm{min}}},$ (3) $$

where $E_{\mathrm{min}}$ ($S_{\mathrm{min}}$) are the minimum possible number of training examples (steps) to reach a specified level of performance. Additionally, they introduce a notion of a noise scale (for SGD-like optimizers):

$$ $B_{\mathrm{noise}}^{\mathrm{curv}}\coloneq\frac{tr(H\Sigma)}{G^{T}HG},$ (4) $$

where $G$ is the noiseless true gradient, $H$ is the true hessian of the loss function and $\Sigma$ is the minibatch covariance. For $B\ll B_{\mathrm{noise}}^{\mathrm{curv}}$ one obtains the linear learning rate scaling rule, while for $B\gg B_{\mathrm{noise}}^{\mathrm{curv}}$ increasing $B$ does not yield any loss improvement.

Under assumption of the Hessian being a multiple of the identity matrix, one obtains a simplified version:

$$ $B_{\mathrm{simple}}^{\mathrm{curv}}\coloneq\frac{tr(\Sigma)}{|G^{2}|},$ (5) $$

and McCandlish et al. () argue that

$$ $B_{\mathrm{crit}}\approx B_{\mathrm{noise}}^{\mathrm{curv}}\propto B_{\mathrm{ simple}}^{\mathrm{curv}},$ (6) $$

thus bridging together mathematical loss curvature and pragmatical compute resource utilization views. Approximation with $B_{\mathrm{simple}}^{\mathrm{curv}}$, being computationally less expensive to estimate, is shown to be to a good degree applicable across multiple tasks, datasets and model architectures. Both the critical batch size and the noise scale are shown to grow in time as one progresses in the training, with the only dependence on the loss value via a power law, with parameters $B_{0}$ and $\alpha_{B}$ to be determined empirically (Kaplan et al., ):

$$ $B_{\mathrm{crit}}=\frac{B_{0}}{L^{1/\alpha_{B}}}.$ (7) $$

Notably, Smith & Le () introduce from a different SDE perspective another definition of the noise scale:

$$ $B_{\mathrm{noise}}^{\mathrm{SDE}}\coloneq\eta(\frac{T}{B}-1)\approx\eta\frac{T }{B},$ (8) $$

where $T$ is the training set size. It is suggested that one should aim at finding the optimal noise scale in the first place, rather than optimal batch size and learning rate. Within the suggested Bayesian framework, Smith & Le () argue that the optimality arises from the trade-off between depth and breadth in the Bayesian evidence. In a follow-up work, Park et al. () take one step further and extend the noise scale to a model width limit and introduce a modified noise scale accounting for the change of the model width in the standard (SP) and Neural Tangent Kernel (NTK) parametrizations (Jacot et al., ):

$$ $B_{\mathrm{noise}}^{\mathrm{norm}}\coloneq\dfrac{B_{\mathrm{noise}}^{\mathrm{ SDE}}}{|w|^{2}},$ (9) $$

where $|w|^{2}$ is model weight norm, normalizing $B_{\mathrm{noise}}^{\mathrm{SDE}}$ to have the unit $1/\text{loss}$.

The second perspective on $B_{\mathrm{crit}}$ is as a region where batch invariance breaks. Introduced by Hilton et al. (), batch invariance refers to a regime where the model performance remains invariant with the change of either learning rate or batch size within the corresponding scaling rule. As shown by Shallue et al. (), the breaking of batch invariance appears with an increase of batch size to sufficiently large values and looks like plateauing of the optimal learning rate. Zhang et al. () further investigated how the critical batch size is affected by using momentum, optimizer pre-conditioning and exponential moving average (EMA).

Intriguingly, Li et al. () expanded the approach of McCandlish et al. () and showed that in the case of Adam, the batch invariance does not break conventionally as in the SGD case. In fact, it is always preserved, with the only difference that the $\eta\propto\sqrt{B}$ scaling rule breaks at the peak value $B_{\mathrm{peak}}$ and transforms into a $\eta\propto 1/\sqrt{B}$ rule via:

$$ $\eta^{*}=\dfrac{\eta^{\mathrm{crit}}}{\frac{1}{2}(\sqrt{\frac{B_{\mathrm{peak} }}{B}}+\sqrt{\frac{B}{B_{\mathrm{peak}}}})}.$ (10) $$

They also show that $B_{\mathrm{peak}}\approx B_{\mathrm{crit}}$ in the definition of McCandlish et al. (), therefore bridging together the two $B_{\mathrm{crit}}$ perspectives outlined above.

### A.3 Model training configuration (cont.)

- •
24 layers, FFN expansion factor $f_{\mathrm{ffn}}=d_{\mathrm{ffn}}/d_{\mathrm{model}}=4$, multihead attention with the head dimension $d_{\mathrm{head}}=128$.
- •
GeLU activation function, Layer Normalization initialized with $1$ (Ba et al., ), RoPE with $\theta=10000$ (Su et al., ).
- •
Dropout is disabled and biases are included in all layers (initialized with 0), weights are shared between the input and output embedding layers.
- •
FSDP parallelization scheme (Zhao et al., ), bfloat16 precision, FlashAttention-2 (Dao, ).

### A.4 Hyperparameter grid (cont.)

The ($\eta$, $B$, $T$, $d_{\mathrm{model}}^{\mathrm{base}}$, $d_{\mathrm{model}}$) grid is defined with the following values:

- •
Learning rate $\eta$:
–
$\{2^{-12},~{}2^{-11.5},\dots,~{}2^{-7}\}$ for $d_{\mathrm{model}}^{\mathrm{base}}=1024$
–
$\{2^{-11},~{}2^{-10},\dots,~{}2^{-6}\}$ for $d_{\mathrm{model}}^{\mathrm{base}}=256$
- •
Batch size $B=\{2^{16},~{}2^{18},\dots,~{}2^{26}\}$ tokens
- •
Data horizon $T=\{2^{30},~{}2^{31}\dots,~{}2^{35}\}$ tokens
- •
Base model width $d_{\mathrm{model}}^{\mathrm{base}}=\{256,~{}1024\}$
- •
Model width $d_{\mathrm{model}}=\{256,~{}512,~{}1024\}$

For a configuration with $(B=2^{20}$, $d_{\mathrm{model}}^{\mathrm{base}}=1024)$, we perform longer runs with an extended set of horizons with $\{2^{36},~{}2^{37}\}$ token budgets, except for the smallest $B=2^{16}$ due to limited computational resources and low GPU utilization of this batch size on our hardware. A configuration for the largest batch size $(B=2^{26}$, $d_{\mathrm{model}}^{\mathrm{base}}=1024)$, we train until $T=2^{38}$ tokens to further establish the learning rate optimum drift (Sec. 3.1).

The total number of trainable parameters is 32M, 101M, 354M for the models with widths $d_{\mathrm{model}}=\{256,~{}512,~{}1024\}$, respectively. We also train 1.3B and 5B models up until $2^{35}\approx 34\mathrm{B}$ tokens with three selected learning rate values for a fixed batch size of $2^{20}$ and $2^{24}$ tokens, respectively, in order to study learning rate sensitivity change within $\mu$P (Sec. 3.4). The models share the same $\mu$P base model with $d_{\mathrm{model}}^{\mathrm{base}}=1024$ and have the corresponding width $d_{\mathrm{model}}=2048$ (1.3B) and $d_{\mathrm{model}}=4096$ (5B).

### A.5 Random seed variation

Figure: Figure 5: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for three different random seeds for the model configuration $(d_{\mathrm{model}}=d_{\mathrm{model}}^{\mathrm{base}}=1024)$.
Refer to caption: x7.png

### A.6 Learning rate schedule scaling (cont.)

Conventionally, the learning rate schedule consists of a warmup phase, followed by either a constant phase or a decay phase. When all of the three phases are enabled, one obtains a warmup-stable-decay (WSD) schedule (Hu et al., ):

$$ $\eta(t)=\begin{cases}\vspace{2mm}\dfrac{t}{T_{\mathrm{warmup}}}\cdot\eta_{ \mathrm{max}}&\text{if $t<T_{\mathrm{warmup}}$}\\ \vspace{2mm}\eta_{\mathrm{max}}&\text{if $T_{\mathrm{warmup}}\leq t<T-T_{ \mathrm{decay}}$}\\ \left(1-\dfrac{t-(T-T_{\mathrm{decay}})}{T_{\mathrm{decay}}}\right)\cdot\eta_{ \mathrm{max}}&\text{if $T-T_{\mathrm{decay}}\leq t<T$}\end{cases},$ (11) $$

where $T$ is the total length of the training horizon, $T_{\mathrm{warmup}}~{}(T_{\mathrm{decay}})$ is the length of the warmup (decay) phases, all measured in tokens.

As Hägele et al. () showed, there is no significant difference in terms of the final loss value and learning rate sensitivity between using cosine decay and WSD schedules. We run additional ablations in our setup and also arrive at the same conclusions: the structure of the learning rate optimum is marginally affected by the decay phase of the schedule and its type. Even though there appears to be a small increase in learning rate sensitivity if learning rate is decayed comparing to the schedule without decay, it does not affect the optimal $\eta^{*}$ location (Fig. 6).

Furthermore, we vary the warmup scaling strategy with an increase of the data horizon, specifically where all the horizons either share the same warmup length, or warmup is scaled together with the horizon length (with the fixed $f=T_{\mathrm{warmup}}/T=1/64$ fraction of the total horizon), or warmup is disabled. We observe that the addition of warmup decreases learning rate sensitivity and, interestingly, that scaling of the warmup proportionally with the horizon length leads to a drift of the learning rate optimum, as also indirectly observed earlier by Kosson et al. ().

Figure: Figure 6: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for schedules with the following phases: warmup and constant (top left); warmup, constant and linear decay to 0 (top middle); warmup and linear decay to 0 (top right); warmup and cosine decay to 10% of the maximum $\eta$ (bottom left); warmup scaled as $1/64$ fraction of the total horizon and constant (bottom middle); constant (bottom right). Warmup duration is always set to $T_{\mathrm{warmup}}=2^{19}=524288$ tokens, except for the case with warmup phase scaling. The model configuration is $(d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=1024,~{}B=2^{20})$.
Refer to caption: x10.png

### A.7 Loss profiles per ( d model base , d model ) superscript subscript 𝑑 model base subscript 𝑑 model (d_{\mathrm{model}}^{\mathrm{base}},d_{\mathrm{model}}) ( italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT ) configuration

#### A.7.1 d model base = 1024 , d model = 1024 formulae-sequence superscript subscript 𝑑 model base 1024 subscript 𝑑 model 1024 d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=1024 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 1024 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 1024

Figure: Figure 7: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=1024)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x16.png

#### A.7.2 d model base = 1024 , d model = 512 formulae-sequence superscript subscript 𝑑 model base 1024 subscript 𝑑 model 512 d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=512 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 1024 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 512

Figure: Figure 8: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=512)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x22.png

#### A.7.3 d model base = 1024 , d model = 256 formulae-sequence superscript subscript 𝑑 model base 1024 subscript 𝑑 model 256 d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=256 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 1024 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 256

Figure: Figure 9: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=1024,~{}d_{\mathrm{model}}=256)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x28.png

#### A.7.4 d model base = 256 , d model = 256 formulae-sequence superscript subscript 𝑑 model base 256 subscript 𝑑 model 256 d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=256 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 256

Figure: Figure 10: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=256)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x34.png

#### A.7.5 d model base = 256 , d model = 512 formulae-sequence superscript subscript 𝑑 model base 256 subscript 𝑑 model 512 d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=512 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 512

Figure: Figure 11: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=512)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x40.png

#### A.7.6 d model base = 256 , d model = 1024 formulae-sequence superscript subscript 𝑑 model base 256 subscript 𝑑 model 1024 d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=1024 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256 , italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT = 1024

Figure: Figure 12: Loss profile $\mathcal{L}_{\mathrm{val}}-\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$ as a function of maximum learning rate $\eta$ for $(d_{\mathrm{model}}^{\mathrm{base}}=256,~{}d_{\mathrm{model}}=1024)$ for batch size $B=2^{16}$ (top left), $B=2^{18}$ (top middle), $B=2^{20}$ (top right), $B=2^{22}$ (bottom left), $B=2^{24}$ (bottom middle), $B=2^{26}$ (bottom right) across various token budgets.
Refer to caption: x46.png

### A.8 μ 𝜇 \mu italic_μ P-averaged optimal learning rate and batch size joint scaling for d model base = 256 superscript subscript 𝑑 model base 256 d_{\mathrm{model}}^{\mathrm{base}}=256 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256

Figure: Figure 13: (left) Evolution of the optimal learning rate with an increase of the pretraining token budget $\eta^{*}_{\mathrm{norm}}(T)$ , normalized to $\eta^{*}|_{T=2^{30}}$, for a set of batch sizes (in tokens). Each point is obtained by averaging optimal learning rate values across $\mu$P model family, as described in Sec. 3.1. Dashed lines correspond to a square-root $\eta^{*}\propto\sqrt{T^{-1}}$ scaling rule. (right) Transposition of (left): optimal learning rate $\eta^{*}$ per batch size, against a range of pretraining token budgets. Each point is $\mu$P-averaged as in (left), with color bands visualizing the corresponding standard deviation. We note that experiments were performed with a coarser learning rate resolution of $2^{1}$ compared to a $2^{0.5}$ step in experiments with $d_{\mathrm{model}}^{\mathrm{base}}=1024$.
Refer to caption: x52.png

### A.9 Per-model optimal learning rate and batch size joint scaling

#### A.9.1 d model base = 256 superscript subscript 𝑑 model base 256 d_{\mathrm{model}}^{\mathrm{base}}=256 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256

Figure: Figure 14: Individual curves contributing to Fig. 13 for models with $d_{\mathrm{model}}=256$ (left column), $512$ (middle column), $1024$ (right column) showing evolution of the normalized to $T=2^{30}$ tokens optimal learning rate $\eta^{*}_{\mathrm{norm}}$ in time per batch size (top row), and joint optimal $(\eta,B)$ curves per token budget (bottom row), for $d_{\mathrm{model}}^{\mathrm{base}}=256$.
Refer to caption: x54.png

#### A.9.2 d model base = 1024 superscript subscript 𝑑 model base 1024 d_{\mathrm{model}}^{\mathrm{base}}=1024 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 1024

Figure: Figure 15: Individual curves contributing to Fig. 1 for models with $d_{\mathrm{model}}=256$ (left column), $512$ (middle column), $1024$ (right column) showing evolution of the normalized to $T=2^{30}$ tokens optimal learning rate $\eta^{*}_{\mathrm{norm}}$ in time per batch size (top row), and joint optimal $(\eta,B)$ curves per token budget (bottom row), for $d_{\mathrm{model}}^{\mathrm{base}}=1024$.
Refer to caption: x60.png

### A.10 Per-model validation loss evolution in time depending on batch size with optimally-tuned learning rate

#### A.10.1 d model base = 256 superscript subscript 𝑑 model base 256 d_{\mathrm{model}}^{\mathrm{base}}=256 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 256

Figure: Figure 16: Analogue of Fig. 2(b) (top row) and Fig. 2(a) (bottom row) for models with widths $d_{\mathrm{model}}=256$ (left column), $512$ (middle column), $1024$ (right column) and the base model width $d_{\mathrm{model}}^{\mathrm{base}}=256$.
Refer to caption: x66.png

#### A.10.2 d model base = 1024 superscript subscript 𝑑 model base 1024 d_{\mathrm{model}}^{\mathrm{base}}=1024 italic_d start_POSTSUBSCRIPT roman_model end_POSTSUBSCRIPT start_POSTSUPERSCRIPT roman_base end_POSTSUPERSCRIPT = 1024

Figure: Figure 17: Analogue of Fig. 2(b) (top row) and Fig. 2(a) (bottom row) for models with widths $d_{\mathrm{model}}=256$ (left column), $512$ (middle column), $1024$ (right column) and the base model width $d_{\mathrm{model}}^{\mathrm{base}}=1024$.
Refer to caption: x72.png

### A.11 Learning rate sensitivity in the μ 𝜇 \mu italic_μ P width limit

Figure: Figure 18: Same as Fig. 4 but without y-axis normalization with $\mathcal{L}_{\mathrm{val}}^{\mathrm{min}}$.
Refer to caption: x78.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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$.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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.
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: See the weakness section.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: See above please.
soundness: 2
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.
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.
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.

Meta-review: 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.

Decision: Reject