Scaling Optimal LR Across Token Horizons

## Abstract

Abstract State-of-the-art LLMs are powered by scaling – scaling model size, dataset size, and cluster size. It is economically infeasible to extensively tune hyperparameters for the largest runs. Instead, approximately optimal hyperparameters must be inferred or transferred from smaller experiments. Hyperparameter transfer across model sizes has been studied in Yang et al. ( 2022 ) . However, hyperparameter transfer across dataset size – or token horizon – has not been studied yet. To remedy this we conduct a large-scale empirical study on how optimal learning rate (LR) depends on the token horizon in LLM training. We first demonstrate that the optimal LR changes significantly with token horizon – longer training necessitates smaller LR. Secondly, we demonstrate that the optimal LR follows a scaling law and that the optimal LR for longer horizons can be accurately estimated from shorter horizons via such scaling laws. We also provide a rule-of-thumb for transferring LR across token horizons with zero overhead over current practices. Lastly, we provide evidence that LLama-1 used too high LR, and argue that hyperparameter transfer across data size is an overlooked component of LLM training.

## 1 Introduction

Figure: Figure 1: Final validation loss of a 350 million parameter LLM for different learning rates (LR) and token horizons. The dashed lines indicate our fitted curve and the stars indicate the estimated optimal LR. The optimal LR decreases as the token horizon increases.
Refer to caption: x1.png

State-of-the-art LLMs are scaled in multiple dimensions. The models are becoming increasingly large, e.g. Grok-1.5 has 314 billion (B) parameters (xAI, ). The clusters used to train them are growing in size, e.g. the recently operational Memphis super-cluster contains over 100,000 H100 GPUs (Alcorn, ). Lastly, the training datasets are growing, e.g. LLama-3 was trained on 15 Trillion (T) tokens (Dubey et al., ). At these scales, it is infeasible to extensively tune hyperparameters. Practitioners must instead resort to hyperparameter transfer, a process where approximately optimal hyperparameters for large-scale experiments are inferred from experiments at a smaller scale. Perhaps the most famous work on hyperparameter transfer is muP (Yang et al., ) – a methodology for transferring optimal hyperparameter from a small model to a large model. While hyperparameter transfer across model size is a well-studied problem, transfer across dataset size remains understudied. This paper aims to remedy this shortcoming in the literature.

In this paper, we present a large-scale study on hyperparameter transfer across dataset sizes. We specifically focus on learning rate (LR), an important hyperparameter that influences training stability, generalization, and convergence speed. Our study is essentially a large ablation experiment where we vary LR and token horizon for a few different LLM models. We consider $>$250 training runs in total. To keep the scope and computational requirements manageable we focus on how LR depends on the token horizon, and only present some preliminary results on its interaction on model size. Our experiments are based on standard public resources – the Megatron codebase (Shoeybi et al., ) and the RefinedWeb dataset (Penedo et al., ). Our three main contributions are:

- •
We provide a large-scale study demonstrating that the optimal LR depends strongly on the token horizon, with longer horizons necessitating smaller LRs. This fact holds even when muP parametrization is used.
- •
We demonstrate that 1) the optimal LR for any given architecture follows a scaling law and 2) this allows hyperparameter transfer where the optimal LR for a long horizon is inferred from a shorter horizon. Furthermore, we provide a rule-of-thumb for transferring LR across token horizons with zero overhead over current practices.
- •
We provide a case study on the optimal LR for the LLama architecture. We provide evidence that LLama-1 used an LR that is too large, highlighting hyperparameter transfer across horizons as an overlooked component of LLM training.

## 2 Background

LLM Scaling Laws. LLMs are typically decoder-only transformers (Vaswani, ) trained via next token prediction on web-scale text (Radford et al., ). It has been empirically observed that the performance of LLMs scale well with model size – with the largest model showing emergent capabilities (Wei et al., ). Kaplan et al. () shows that LLM performance roughly scales as a power-law in the model size $N$. The performance here is measured by validation loss $L$, which is well known to correlate strongly with downstream metrics. Specifically they propose the following law (using constants $N_{c},\alpha_{N}$) for models trained on sufficiently large datasets: $L(N)=(N_{c}/N)^{\alpha_{N}}$. Hoffmann et al. () also showed that the performance scales well with dataset size and that the optimal performance for a given FLOPs budget is obtained by scaling the model and dataset size jointly. The current paradigm thus scales dataset size in addition to model size – e.g. LLama-3 was trained on $\textgreater$10x as many tokens as LLama-1 (Dubey et al., ). In this paper we use notation from Kaplan et al. (), denoting the dataset size (i.e. number of tokens) by $D$ and the model size (i.e. parameter count) by $N$. As is common in the literature, we will fit scaling laws to empirical observations. To measure goodness-of-fit we will use the $R^{2}$ measure – its value ranges from 0 to 1, where 1 indicates a perfect fit and 0 indicates no fit at all.

Hyperparameter transfer. Hyperparameters can strongly influence the performance of LLMs, and LR is a critical hyperparameter. The muP paper (Yang et al., ) first popularized hyperparameter transfer – the process of finding optimal hyperparameters from a small proxy experiment. Core to muP is the muP-parametrization – a way of parametrizing an LLM such that the optimal LR for a small model is also optimal for a larger model. The muP parametrization introduces some changes to the network, e.g. the attention scaling factor is changed from $\frac{1}{\sqrt{d}}$ to $\frac{1}{d}$. The original muP paper shows that LR and many other hyperparameters transfer from small to large models when muP parametrization is used, which is further corroborated in Lingle (). While Yang et al. () focuses on transfer between model sizes, we focus on transfer between token horizons.

Figure: Figure 2: Final validation loss as a function of learning rate (LR) and token horizon. The dashed lines indicate our fitted curve and the stars indicate optimal LR. The optimal LR decreases monotonically with longer horizons.
Refer to caption: x2.png

## 3 Experiments

Experimental Setup. Our experimental setup closely follows the setup of the GPT-3 paper (Brown, ). We consider model sizes of 50 million (m), 125m, 350m, 760m, 1.3 billion (B), and 2.7B parameters using the architectures of Table 2.1 in the GPT-3 paper. The table is replicated as Table 4 in Appendix A and also lists the LRs used in the GPT-3 paper. We use hyperparameters following GPT-3 – weight decay of 0.1, gradient clipping of 1.0, and cosine learning decay schedule. The full list of hyperparameters can be viewed in Table 3 in Appendix A. The LR will vary with the training step since we use an LR schedule, and we will use the term LR to refer to the largest LR of a run. We make three adjustments compared to the GPT-3, mainly to prevent divergence of the model for high LRs. 1) For warmup we use the maxima of 1% of the training steps following Muennighoff et al. () and 1000. Having too short a warmup stage is known to cause divergence (Wortsman et al., ) which we want to avoid for short training runs. 2) We use qk-norm following Dehghani et al. (), this is known to prevent divergence without hurting validation loss (Wortsman et al., ). 3) The original GPT-3 paper uses different batch sizes for different model sizes. To avoid confounders we simply use the same batch size of 0.5m tokens all model. We use the RefinedWeb dataset (Penedo et al., ), a common-crawl derived dataset of roughly 600B tokens which is known to be of high quality (Penedo et al., ). Experiments are run on the Megatron codebase (Shoeybi et al., ). Unless mentioned we will use the same seed for all runs. For curve fitting we use least-square fitting with either a first or second-degree polynomial using Numpy and Scipy (Harris et al., ).

Figure: Figure 3: Final validation loss as a function of max learning rate (LR) and token horizon for four models. The dashed lines indicate our fitted curve. The optimal LR, denoted by a black star, decreases monotonically with longer horizons for all models.
Refer to caption: x3.png

### 3.1 Ablations

In our first experiment, we consider the 350m LLM model from Table 4 in Appendix A. We perform an ablation study where we vary the LR and token horizon and measure the final validation loss. We consider $\{25,50,100,200,400\}$ billion tokens. We will start with the LR of $3\!\times\!10^{-4}$ from Table 4 which is what the GPT-3 paper used (Brown, ). We then multiply this base LR by factors $\{0.25,0.5,1,2,4\}$. The LRs we consider are thus $\{7.5\!\times\!10^{-5},1.5\!\times\!10^{-4},3\!\times\!10^{-4},6\!\times\!10^{
-4},1.2\!\times\!10^{-3}\}$. For each combination of LR and token horizon, we train a model and record its final validation loss. To make sure we have the best possible resolution around the minima we find the optimal learning rate $LR^{*}$, and then further train with the LRs halfway between $LR^{*}$ and its two nearest neighbors – a procedure we repeat twice. From these losses, we fit a curve that estimates how the final validation loss depends on the LR. For each token horizon we fit a second degree polynomial in $\log(LR)$ which provides an excellent fit. The $R^{2}$ of the fit are 0.995 or better, see Table 6 in the Appendix for exact values. The validation losses and fitted curve are shown in Figure 2. We also estimate the optimal LR by taking the minimizer of the fitted curve. From Figure 2 we make two observations: 1) the optimal LR decreases with longer token horizons, and 2) the quadratic fit provides excellent agreement with experimental data.

We now repeat these experiments for more model sizes – specifically the 50m, 125m, 760m, 1.3B, and 2.7B parameter models from Table 4. For computational reasons, we only consider shorter horizons for the larger models. The 50m and 125m models go up to 800B tokens, the 760m and 1.3B models go up to 200B tokens while the 2.7B model only goes up to 100B tokens. For the larger models, we don’t increase the sampling rate around the minimizer. For each model, we consider the base LR as the one used in the GPT-3 paper (also listed in Table 4) and then multiply it by e.g. $\{0.25,0.5,1,2,4\}$. The results of these experiments are shown in Figure 3. Two runs (model sizes 50m and 125m, 800B tokens, and highest LR) diverge, and we remove these. The results look similar to those of Figure 2 and we observe that the optimal LR decreases with longer token horizons across model sizes. Figure 10 in the Appendix show similar results for the 1.3B model. We thus conclude that this is a robust phenomenon.

Figure: Figure 4: Scaling laws for optimal LR versus token Horizon. We compare the empirically best LR (dots) to the smooth scaling law of Equation 1 with fitted constants. The $R^{2}$ of these fits are in the range 0.99 - 0.96. Across all model sizes, we see that the scaling law provides a good fit to the empirical data.
Refer to caption: x4.png

### 3.2 Scaling Laws

We have seen that longer token horizons require smaller LR. We now investigate if this insight allows us to derive scaling laws and do hyperaprameter-transfer – i.e. finding the optimal LR to a long token horizon from experiments on a shorter horizon. To do this we will fit scaling laws to our empirical results. Given some fixed model architecture and training recipe, let ${LR}^{*}(D)$ denote the optimal LR for some token horizon $D$. We will use the following functional form:

$$ $LR^{*}(D)=BD^{-\beta}$ (1) $$

Here $B$ and $\beta$ are two constants that might e.g. depend on the model architecture. Taking the logarithm of both sides of Equation 1 we get

$$ $\log\big{[}LR^{*}(D)\big{]}=\log B-\beta\log D$ (2) $$

We thus have a linear equation in the unknowns $\log B,\beta$, and can fit these to our experimental results with least squares. Specifically, we use the minimizer of the quadratic fits of Figure 3 as the data points $LR^{*}$. Fitting Equation 2 gives good fits – the curves for four models are shown in Figure 4. The $R^{2}$ of these fits are in the range 0.99-0.96 (see Table 5 in Appendix A for exact values), a rather high number. In Figure 11 in the Appendix we also show the fits for the 1.3B and 2.7B models, which also show a good fit.

To evaluate the scaling laws we cannot solely rely on $R^{2}$ as that is essentially a measure of training loss. We instead need to evaluate the fit on some held-out data. To simulate hyperparameter transfer we thus consider fitting the constants $\log B,\beta$ on token horizons $25B,50B$ and $100B$ – and then use these constants to predict the optimal LR at longer horizons. We will use the optimal LR we have empirically found with quadratic fits as the correct LR. The results are illustrated for the 50m model in Table 1. We see a reasonably good fit with a relative error of 10-15% – a clear improvement over not scaling LR at all which has a relative error of $>250\%$. We thus conclude that it is possible to perform hyperparameter transfer across token horizons with our scaling laws. Table 1 is repeated for more models with similar fits in Appendix B.

**Table 1: Simulated hyperparameter transfer across token horizons for a 50m model. We use the empirically measured optimal LR at 25,50 and 100B tokens to estimate the optimal LR at longer horizons by fitting the constants of Equation 2. We find a reasonable fit when scaling up the horizon to 800B tokens, with a relative error of 10-15 %. The relative error of using the best LR at the 100B horizon for an 800B horizon is $>$ 250%.**
|  | 25B | 50B | 100B | 200B | 400B | 800B |
| --- | --- | --- | --- | --- | --- | --- |
| Optimal LR | $1.54\!\times\!10^{-3}$ | $9.79\!\times\!10^{-4}$ | $6.06\!\times\!10^{-4}$ | $3.33\!\times\!10^{-4}$ | $2.14\!\times\!10^{-4}$ | $1.71\!\times\!10^{-4}$ |
| Predicted LR |  |  |  | $3.81\!\times\!10^{-4}$ | $2.39\!\times\!10^{-4}$ | $1.50\!\times\!10^{-4}$ |
| Ratio |  |  |  | 0.873 | 0.894 | 1.14 |

### 3.3 muP parametrization

It is natural to ask if muP allows hyperparameter transfer across token horizons. To investigate this we consider the 50m model from Table 4, and use muP parameterization. Note, we will not perform hyperparameter transfer across model sizes, we just use the muP parameterization which slightly differs from the standard parameterization of transformers. We then run ablation experiments using token horizons $\{25,50,100\}$ billion tokens and LRs of $\{0.25,0.5,1,2,4\}$ times the base LR in Table 4. The results of this experiment are shown in Figure 5, and we can indeed see that the optimal LR decreases with a longer token horizon. We thus conclude that the optimal LR does not transfer across token horizons with muP.

Figure: Figure 5: Optimal LR vs token horizon for a 50m model using muP parameterization (Yang et al., ). We see that the optimal LR decreases with longer token horizons, demonstrating that LR does not transfer across horizons even with muP.
Refer to caption: x8.png

### 3.4 Quantifying Variance

To ensure that our results are reliable we here consider quantifying the variance of our experiments. We will use two different methodologies. Firstly we use bootstrapping, following Hoffmann et al. (). We consider the 350m model of Table 4, randomly remove 20% of the data points, and then fit the optimal LR at each token horizon and the scaling law of Equation 1. We repeat this procedure 1000 times and then measure the mean and std of the optimal LR and the constants of Equation 1. The results are given in Table 2. We see that the standard deviations are relatively small, which again suggests that the variance in our experimental results is low.

The second methodology we use is simply rerunning a small-scale experiment with multiple seeds. We use the 350m model from Table 4, a 100B token horizon, three learning rates, and two additional seeds. We then estimate the optimal LR via a quadratic fit as in Section 3.1. The results of this experiment are shown in Table 7 in Appendix B. We see that there are small differences in the losses, and hence small differences in the estimated optimal LR. The relative std $\sigma/\mu$ is $2.63\!\times\!10^{-2}$, so we can expect the relative error to be a few percent. When using more data points the variance could further decrease due to concentration of measure.

**Table 2: We estimate the mean $\mu$ and standard deviation $\sigma$ of the optimal learning rate $LR^{*}$ via bootstrapping. We sample 80% of the data from Figure 2 and then estimate the optimal LR with the procedure from Section 3.1. This bootstrapping procedure is repeated 1000 times. We see that the variance is small compared to the mean, and the relative error $\sigma/\mu$ is on the order of a few percent. This implies that the uncertainty in our estimates is small.**
|  | 25B | 50B | 100B | 200B | 400.0 |
| --- | --- | --- | --- | --- | --- |
| $\mu(LR^{*})$ | $7.04\!\times\!10^{-4}$ | $5.96\!\times\!10^{-4}$ | $5.07\!\times\!10^{-4}$ | $3.45\!\times\!10^{-4}$ | $2.44\!\times\!10^{-4}$ |
| $\sigma(LR^{*})$ | $1.28\!\times\!10^{-5}$ | $9.12\!\times\!10^{-6}$ | $1.96\!\times\!10^{-5}$ | $4.07\!\times\!10^{-6}$ | $4.31\!\times\!10^{-6}$ |
| $\sigma/\mu$ | $1.82\!\times\!10^{-2}$ | $1.53\!\times\!10^{-2}$ | $3.88\!\times\!10^{-2}$ | $1.18\!\times\!10^{-2}$ | $1.76\!\times\!10^{-2}$ |

### 3.5 Effect of Batch Size

It is well known that batch size $BS$ affects the optimal LR (Goyal, ). While we primarily focus on the setting of a fixed batch size, we here consider modifying the batch size for the 1.3B model of Table 4. Specifically, we double the batch size to 1m tokens and train for 25B and 50B tokens in total with different LRs. In Figure 9 in the Appendix we show the final validation loss as a function of the LR and estimate the optimal LR. In Figure 13 we show the optimal LR as a function of token horizon for the 1.3B model using a batch size of 0.5 or 1m tokens. We see that the optimal LR is higher with a larger batch size as expected. More importantly, we note that the linear fits are roughly parallel. This suggests that the optimal LR depends on the token horizon the same way irrespective of the batch size, i.e. that we can factorize Equation 1 as $LR(BS,D)=f(BS)D^{-\beta}$.

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

## 4 Scaling Law with respect to model size

So far we have considered Equation 1 with constants fitted independently for each model size. Fully determining the joint scaling properties of the model architecture and token horizon is outside of the scope of this paper – doing so would be computationally infeasible. Nonetheless, we will here provide some preliminary results and discussions regarding a scaling law for both model size $N$ and token horizon $D$. In Figure 6(b) we plot the optimal LR as a function of $N$, and in Figure 6(a) we plot it as a function of $D$. Both curves are straight lines when the axes are logarithmic. This observation motivates the following functional form, which is mathematically derived in Appendix C.

$$ $LR^{*}(N,D)=CN^{-\alpha}D^{-\beta}$ (3) $$

Here, the constant $C$ is the ideal learning rate for a model of size $1B$ and token horizon $1B$, where the precise numerical value of $C$ may depend on the batch size (see section 3.5), vocabulary size and tokenization. The factor $N^{-\alpha}$ captures the fact that the optimal LR decreases with model size $N$, while the factor $D^{-\beta}$ captures that optimal LR also decreases with token horizon $D$. We now fit Equation 3 to our data for the models of size 760m, 1.3B, and 2.7B, using the 7B model of Section 4.1 as a held-out validation set. We use the huber loss with $\delta=1\times 10^{-3}$ and the BFGS algorithm, similar to Hoffmann et al. (). We observe a good fit with RMSE$=2.2\times 10^{-5}$ and a validation $R^{2}$ of $0.978$. The numerical values we find are:

$$ $C\sim 1.55\times 10^{-3},\quad\alpha\sim 0.23,\quad\beta\sim 0.32$ (4) $$

We depict the results in Figure 7. For smaller models (mainly 50m and 125m) a larger $\beta$ is observed from the experiments. We thus consider Equation 3 to be valid in the regime of large models ($\geq$ 760M parameters).

Figure: Figure 7: Fit of Equation 3 compared to the experimental results. The data points for the 7B model of Section 4.1 are excluded at the time of fitting and used as validation data. We have an $R^{2}$ of $0.978$ on this validation data.
Refer to caption: x11.png

### 4.1 A Case-study on Llama-1

We now consider evaluating if the LRs used by Llama-1 Touvron et al. () are ”correct” according to our scaling laws. To do this we adopt the LLama-1 architecture (RMSnorm, Rope embeddings, and so on) and run small-scale experiments with token horizons 25B, 50B, and 100B and different LRs. Based upon these experiments, we find that values for Equation 1 are $B=8.29\times 10^{-4}$ and $\beta=0.3$ (see Figure 12 in Appendix B). We then extrapolate to find the optimal LR at 1T tokens, which comes out as $1.15\!\times\!10^{-4}$. LLama-1 used $3\!\times\!10^{-4}$, an LR which is too large by a factor of $>2.5$ according to our results. This methodology is visualized in Figure 8. The numerical values used for our predictions are also tabulated in Table 8 in Appendix B, where we also compare to predicting the optimal LR for LLama-1 using our scaling law from Eq. Equation 3. If Llama used the wrong LR, what were the actual effects of this? We provide an estimate on the upper bound for the validation loss penalty between the optimal and used learning rate in Llama-1 using the parabola fit at $D=100B$. $\Delta L=L(D=100B,LR=LR^{*})-L(D=100B,LR=2.6LR^{*})=0.027$. Note, this is an upper bound since the parabolas flatten with longer token horizons

Figure: Figure 8: The optimal LR of the LLama-1 7B model as a function of dataset size. We find that $1.15\!\times\!10^{-4}$ is optimal for 1T tokens whereas the Llama-1 paper used $3\!\times\!10^{-4}$.
Refer to caption: x12.png

## 5 Related Work

The study of scaling laws for LLMs originated in Kaplan et al. () and it is still an active area of research. Scaling laws are researched in the context post-training (Lin et al., ; Zhang et al., ; Gao et al., ), model architecture (Krajewski et al., ; Alabdulmohsin et al., ; Frantar et al., ; Wang et al., ), multi-modality (Aghajanyan et al., ; Cherti et al., ), inference (Sardana & Frankle, ), data (Dohmatob et al., ; Fernandes et al., ) and other domains (Zhang et al., ; Neumann & Gros, ; Kadra et al., ). There are also more theoretical studies (Michaud et al., ; Caballero et al., ).

LR has long been known as an important hyperparameter, and hence there is ample work on how to select it (Goyal, ). Larger models are known to require smaller LR, and e.g. Kaplan et al. () suggests the formula $LR(N)\approx 0.003239-0.0001395\log(N)$ for tuning LR as a function of model size. MuP (Yang et al., ) is a principled methodology for selecting LR when the model is scaled, and the method is actively researched (Lingle, ; Everett et al., ; Noci et al., ; Blake et al., ). Recently Everett et al. () showed that LR transfer across model size is possible both with different parametrizations and optimizers. However, a limitation in both Yang et al. () and Everett et al. () is the fixed training horizon assumption used in the theoretical derivations. Our work directly explores this limitation. One notable conclusion in Everett et al. () is that extrapolating a LR to larger models may significantly overestimate optimal LR in the compute optimal setting. We reach a different conclusion, the exponent for model size is independent of the token horizon and vice versa. While their results might hold true for other optimizers, a closer look at Table 11 in Everett et al. () shows that for the Adam optimizer, the scaling exponents are similar for the two token horizons they considered. At last, we mention that Bi et al. () recently fits the optimal LR as a function of total compute (which will combine batch size, model size, and total training duration) for two small models.

## 6 Discussion

Limitations. To limit the scope and computational requirements of our study we have intentionally focused on a narrow area – scaling token horizons and changing LR with otherwise fixed LLM recipes. With this limited scope, there are naturally many limitations to our study. We have only extended the scaling laws to roughly 800B tokens, while many SOTA LLMs are trained significantly longer (Dubey et al., ). It is well-known that optimal LR depends on model size, and this study has only scratched the surface of this important topic. Beyond just model size, model architectural modifications like mixture of experts (Shazeer et al., ), different attention types (Ainslie et al., ), and state-space models (Gu & Dao, ) could plausibly interact with both the LR and token horizon. Other dimensions to consider are additional hyper parameters (Bjorck et al., ) and multi-modality (Huang et al., ). We leave these topics for future work.

Advice for Practioners. Our experiments show that the optimal LR decreases with the token horizon. This necessitates hyperparameter transfer across token horizons. For practitioners who are working on larger models (say $>=760m$) we recommend simply using Equation 3 where we have already found $\beta=0.32$ to generalize across architectures. To find the optimal LR $LR^{*}(D_{1})$ at some long horizon $D_{1}$ practitioners can just find the optimal LR $LR^{*}(D_{2})$ at a short horizon $D_{2}$ and then estimate:

$$ $LR^{*}(D_{2})\approx LR^{*}(D_{1})\bigg{(}\frac{D_{2}}{D_{1}}\bigg{)}^{-0.32}$ (5) $$

Since practitioners typically find the optimal LR for smaller horizons anyway during hyperparameter search, this methodology has no overhead over current practices. For practitioners with ample compute resources, we recommend finding the best LR at multiple short horizons using the quadratic fitting of Section 3.1. Thereafter the constants in Equation 1 can be found, and the optimal LR at longer horizons can be estimated. For practitioners using methods that provide zero-shot transfer of optimal LR across width (Yang et al. () and Everett et al. ()), we recommend using an adjustment of Equation 3 ($N\propto{n_{layer}}{d^{2}_{model}})$ which yields the formula $LR^{*}(n_{layer},D)=C{n^{-\alpha}_{layer}}D^{-\beta}$. Then $C$ needs to be found using the usual muP sweep for LR with small width, depth and horizon. This allows hyperparameter transfer of LR across depth.

Conclusions. We have investigated how LR and token horizon interact when training LLMs. First, we have shown that the optimal LR decreases as the token horizon gets longer. This finding is robust across model sizes. Secondly, we have shown that the optimal LR follows reliable scaling laws and that fitting these allows for hyperparameter transfer across token horizons. As a case study, we have applied our methods to the training of LLama, and have shown evidence that Llama-1 was trained with a LR which was significantly larger than optimal. We argue thus argue that hyperparameter transfer across token horizons is an understudied aspect of LLM training.

## Appendix A Hyperparameters

**Table 3: Hyperparameters. These follow Brown (). Hyperparameters not listed here follow defaults in the Megatron codebase Shoeybi et al. ().**
| Hyperparameter | Value |
| --- | --- |
| weight decay | 0.1 |
| grad clip norm | 1.0 |
| LR schedule | cosine |
| Adam $\beta_{1}$ | 0.9 |
| Adam $\beta_{2}$ | 0.95 |
| Context length | 2048 |
| Batch size (tokens) | 524288 |
| Warmup Steps | $\max(1000,0.01\times\textrm{train iters})$ |
| Min LR | $0.1\times$ Max LR |

**Table 4: Model architectures and base LRs we consider. These follow GPT-3 Brown ().**
| Model Name | params | layers | $d_{\mathrm{model}}$ | heads | Base LR |
| --- | --- | --- | --- | --- | --- |
| Tiny | 50M | 8 | 512 | 8 | $8\times 10^{-4}$ |
| Small | 125M | 12 | 768 | 12 | $6\times 10^{-4}$ |
| Medium | 350M | 24 | 1024 | 16 | $3\times 10^{-4}$ |
| Large | 760M | 24 | 1536 | 16 | $2.5\times 10^{-4}$ |
| 1.3B | 1.3B | 24 | 2048 | 16 | $2\times 10^{-4}$ |
| 2.7B | 2.7B | 32 | 2560 | 32 | $1.6\times 10^{-4}$ |
| 6.7B | 6.7B | 32 | 4096 | 32 | $1.2\times 10^{-4}$ |

## Appendix B Additional Experimental Data

**Table 5: The $R^{2}$ values and found $\beta$s of the fits in Figure 4. The $R^{2}$s are relatively close to 1, indicating a good fit. The $\beta$s are relatively similar for models $\geq$ 350m, but are larger for the smallest models.**
| Model size | $R^{2}$ | $\beta$ |
| --- | --- | --- |
| 2.7b | 0.9973 | 0.4184 |
| 1.3b | 0.9898 | 0.3171 |
| 760m | 0.9763 | 0.3155 |
| 350m | 0.9607 | 0.3799 |
| 125m | 0.9895 | 0.6531 |
| 50m | 0.9977 | 0.7029 |

**Table 6: The $R^{2}$ values of the fits in Figure 2. They are very close to 1, indicating a great fit.**
| Token Horizon | $R^{2}$ |
| --- | --- |
| 25B | 1- $6.68\!\times\!10^{-4}$ |
| 50B | 1 - $7.59\!\times\!10^{-4}$ |
| 100B | 1 - $4.68\!\times\!10^{-3}$ |
| 200B | 1 - $2.85\!\times\!10^{-3}$ |
| 400B | 1 - $1.76\!\times\!10^{-3}$ |

Figure: Figure 9: Learning rate and validation loss for different token horizons for the 1.3B model from Table 4 using 1m batch size.
Refer to caption: x13.png

Figure: Figure 10: Learning rate and validation loss for different token horizons for the 1.3B model of Table 4. We see that the optimal LR decreases as the token horizon increases.
Refer to caption: x14.png

Figure: Figure 11: Optimal LR as a function of token horizon for the 1.3B and 2.7B models. A straight-line provides a good fit, indicating that the power-law of Equation 1 works well. See Section 3.2 for further details.
Refer to caption: x15.png

**Table 7: We consider the loss $L$ as a function of the LR. We repeat the experiments of Section 3.1 with a 100B token horizon, a 350m model, and multiple seeds. With different seeds, we see slightly different final losses, and slightly different estimates of the optimal LR.**
| seed | $L($$1.5\!\times\!10^{-4}$$)$ | $L($$3\!\times\!10^{-4}$$)$ | $L($$6\!\times\!10^{-4}$$)$ | $\operatorname*{arg\,min}L$ |
| --- | --- | --- | --- | --- |
| 1 | 2.940372 | 2.919948 | 2.913585 | $5.81\!\times\!10^{-4}$ |
| 2 | 2.941199 | 2.919131 | 2.912387 | $5.76\!\times\!10^{-4}$ |
| 3 | 2.941648 | 2.920779 | 2.915190 | $5.47\!\times\!10^{-4}$ |

Figure: Figure 12: Optimal LR vs token horizon for a $7B$ LLama-1 model using batch size of $4M$.
Refer to caption: x17.png

Figure: Figure 13: Optimal LR as a function of token horizon for a 1.3B model using a batch size of 0.5m or 1m tokens. A larger batch size implies that a larger LR is optimal, as expected. The two lines are roughly parallel, suggesting that the dependence on token horizons is the same irrespective of the batch size.
Refer to caption: x18.png

**Table 8: Optimal LR for LLama-1 for different token horizon using the fit in Figure 12. Note that the real Llama-1 used an LR of $3\!\times\!10^{-4}$ which our results suggest is significantly too large.**
| Token Horizon | Optimal LR |
| --- | --- |
| 25B | $3.4\times 10^{-4}$ |
| 50B | $2.8\times 10^{-4}$ |
| 100B | $2.3\times 10^{-4}$ |
| 1T (Using Eq. 1 | $1.15\times 10^{-4}$ |
| 1T (Using Eq. 3) | $1.1\times 10^{-4}$ |

**Table 9: Simulated hyperparameter transfer across token horizons for a 125m model. The setup follows Table 9.**
|  | 25B | 50B | 100B | 200B | 400B | 800B |
| --- | --- | --- | --- | --- | --- | --- |
| Optimal LR | $1.34\!\times\!10^{-3}$ | $1.02\!\times\!10^{-3}$ | $6.60\!\times\!10^{-4}$ | $4.12\!\times\!10^{-4}$ | $2.51\!\times\!10^{-4}$ | $1.98\!\times\!10^{-4}$ |
| Predicted LR |  |  |  | $4.77\!\times\!10^{-4}$ | $3.35\!\times\!10^{-4}$ | $2.35\!\times\!10^{-4}$ |
| Ratio |  |  |  | 0.864 | 0.749 | 0.843 |

**Table 10: We show the mean $\mu$, standard deviation $\sigma$ and relative deviation $\sigma/|\mu|$ estimated via bootstrapping for two quantities : the $R^{2}$ and $\beta$ in Section 3.4. The relative deviation of $\beta$ is roughly 10%, demonstrating that the uncertainty is reasonably small in our scaling laws.**
|  | $R^{2}$ | $\beta$ |
| --- | --- | --- |
| $\mu$ | 0.961 | 0.384 |
| $\sigma$ | $8.85\!\times\!10^{-3}$ | $7.78\!\times\!10^{-3}$ |
| $\sigma/\mu$ | $9.21\!\times\!10^{-3}$ | $2.03\!\times\!10^{-2}$ |

## Appendix C Derivation

In section 3.2 we saw empirically that for a given $N$ and constants $B,\beta$ that may be dependent on $N$,

$$ $LR^{*}(N,D)=B(N)D^{-\beta(N)}$ (6) $$

As seen in Figure 6(b), we observe a linear relationship between $\log N$ and $LR^{*}$ for any given $D$. This suggests the relationship:

$$ $LR^{*}(N,D)=A(D)N^{-\alpha(D)}$ (7) $$

The constants $A,\alpha$ may depend on $D$.
The key question that arises is how the general notion of model size $N$ can be incorporated into the joint scaling law. Moreover, the scaling law formula from Eq. 7 for constant $D$ has to be representable by Eq. 6. It is anticipated to align with the latter, consisting of distinct power laws, each with specific parameters for different $N$ and $D$ values.
Consequently, the objective is to identify a function that fulfills these criteria

$$ $LR^{*}(N,D)=B(N)D^{-\beta(N)}=A(D)N^{-\alpha(D)}$ (8) $$

Dataset size $D$ for different Model Size $N$. As seen in Figure 6(a), since the lines are parallel for any given $N$, the slope $\beta$ (Eq. 6) is independent of the model size $N$. Therefore we can assume that $\beta(N)=\beta$.

Model size $N$ for different Dataset size $D$. As seen in Figure 6(b), since the lines are parallel for any given $D$, the slope $\alpha$ (Eq. 7) is independent of the dataset size $D$. Therefore we can assume that $\alpha(D)=\alpha$ is constant.

Using the fact that $\alpha,\beta$ are constant and multiplying Eq. 8 by $N^{\alpha}D^{\beta}$:

$$ $B(N)N^{\alpha}=A(D)D^{\beta}$ (9) $$

The LHS of Eq. 9 only depends on $N$, whereas the RHS only depends on $D$ so they should both equal some constant, $C$ (this step relies on our proof above that $\alpha,\beta$ are independent of $N,D$), resulting in the functional forms of $A(D),B(N)$

$$ $A(D)=CD^{-\beta},B(N)=CN^{-\alpha}$ (10) $$

Plugging the functional forms of $B(N)$ we finally get the final functional form for the joint scaling law as in Eq. 3.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: What do these results seem to suggest about a possible “$\mu$P extension” to a simultaneous scaling limit of the multiple dimensions (token, width)?

**Update**

I thank the authors for their additional answers.

I apologize if it was not very clear, but please let me stress the fact that my point about the semantics of what token horizon means underlies deep questions about the paper's claim and the extension of their validity. In my honest opinion, this paper presents exciting results about how the optimal learning rate evolves with more training time. However, it fails to define adequate boundaries for the (empirical) claims. In this respect, the fact that no specific meaning is attributed to token horizon stems from these missing experiments that would help to disentangle the various causes of the right shift of the optimal learning rate. Therefore I would not say that the new results are *"completely consistent with our previous experiments"*: however, they are consistent with the revised notion of token horizon that does **not** equal dataset size. And I do sincerely appreciate the authors for promising to update the paper, and carefully revising the notion of token horizon across the manuscript. 

The fact that $\beta=0.32$ transfers across architectures is itself a very counterintuitive claim, and in my opinion, it would require significantly more ablations to be verified. I would at least state that it could be that more investigation is needed for architectures that are not GPT or Llama-1, and $\beta=0.32$ might not be the right scaling exponent outside the tested framework (I thank the authors for pointing out Figure 11). At a fundamental level, no explanation is provided as to what and why the shift happens. Thus, I think these limitations have to be addressed (i.e. stated). 

I am referring to the papers that extend the lr transfer to the network's depth: https://openreview.net/forum?id=17pVDnpwwl and https://arxiv.org/abs/2309.16620.

Overall, I think that with the agreed revision of the storyline around the conceptualization of token horizon, and with the additional experiments, this paper deserves a slightly higher score. Thus, I am updating it to 6.
soundness: 3
strengths: The paper has significant strong points.

1. Experimental Scale. The authors had the capability to run large-scale experiments (billion scale). This makes the results very reliable in the specific experimental setting adopted here (i.e. architecture, model size, optimizer setting).

2.  The research question of how the learning rate should scale up is very important in practical settings and it is not covered in either theoretical or empirical research. Thus, the experimental findings are both novel and relevant.

3. The case study on LLAMA training is very interesting, for instance, by retrospectively questioning how hyperparameter tuning was performed in that setting.
summary: The paper studies how the optimal learning rate changes when varying the total number of tokens that are fed into GPT-3-like models (which, equivalently, means training for a longer time). Consistently, it is found that increasing the number of tokens corresponds to a decrease in the optimal learning rate. This inverse relationship is then studied more in detail by fitting a power law, and its extrapolation allows for a precise prescription of how the learning rate should scale with the number of tokens. Finally, this approach has discovered that LLAMA-1 has been potentially trained with a too-large learning rate.
weaknesses: The paper has some weaknesses, mainly due to the depth of the investigation that is performed. More concretely (in order of importance):

1. It is unclear whether the observed inverse relationship between the number of tokens used for pretraining and optimal learning rate is due to the fact the model is trained progressively for a longer number of time steps, or because the network has processed more data. This is quite a fundamental experiment, and it’s unclear what the view taken by the authors is, given that the meaning of the term “token horizon” is not entirely specified. This could be tested, for instance, by fixing the amount of data and training for multiple epochs, and by increasing the batch size while fixing the number of steps (this is already partially done in the current paper). Aggregating these experiments together with the ones already performed by the authors should elucidate the aforementioned question.

2. The experiment of Section 4 ignores the fact that as the scale gets larger, the value of $\beta$ is very different. Thus, with the provided evidence these results are hardly predictive of the optimal learning rate in the joint scaling (model size, token horizon), and more investigation is needed. In fact, the observation that $\beta$ changes with scales (which the authors make) should already advise against the approach of independently fitting the constants for model size and token horizon. However, the authors explicitly advertise this suggestion to practitioners. On the other hand, it would have been more appropriate to see reported how $\beta$ changes with scale, especially for $\mu$P. In particular, for $\mu$P we expect the learning rate to transfer across width, so the joint scaling properties should be more feasible to test.

3. The authors observe a different $\beta$ for different model size (Table 5). Fundamentally, the authors increase the model size by increasing both the width and at times the depth of the model. Thus, it is unclear whether the observed different $\beta$ at different model sizes is due to the width or the depth scaling.


Minor:

4. Another fundamental limitation of this work is that $\beta$ would potentially change as any architectural modification is made (e.g. QK norm, attention method, gating, etc..). Thus, the $\beta=0.32$ proposed to practitioners, on top of the problems stated above, is valid only in the GPT-3 setting studied. I think this should be emphasized more to not be misleading. 

5. I am not entirely clear on the purpose of the $\mu$P  experiment. $\mu$P is just not designed to exhibit hp transfer across a number of samples (in this case, the number of tokens, or equivalently training time). Thus it is not entirely clear what the expectation of this experiment was in the first place. However, in the related work section, the authors state that exploring this limitation was part of the objectives of the paper, hence I do not decrease my score for this reason.

6. The $\mu$P extension to depth scaling is derived in at least two existing works [1,2]. Thus, there is a parametrization that exhibits learning rate transfer across depth as well. However, I would partially still agree with the authors that the scaling with respect to depth still has to be fitted with power laws (as it is advertised to practitioners), due to the nature of Transformers that have >1 layers within a single residual block. Despite this, I would appreciate it if a more thorough/informed explanation of these suggestions were present in the paper.

[1] Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks

[2] Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit

I am willing to raise my score if my concerns are resolved.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: The warmup used by the authors is much smaller than that used in many recent works[1].


To confirm that the results are robust to values of other hyperparameters, could the authors report results on one of the setups but with 10% warmup steps and 0 weight decay? Assuming the results are robust, I would be happy to increase my score.

 
[1] SMALL-SCALE PROXIES FOR LARGE-SCALE TRANS-
FORMER TRAINING INSTABILITIES.
soundness: 3
strengths: As described, the problem studied by the paper is important and the paper's findings give an important starting point for predicting optimal LR.
summary: As LLMs are scaled up it is not possible to tune hyperparameters at the largest scale. This necessitates a way of predicting optimal hyperparameters. This paper focuses on predicting optimal LR as the training dataset size is scaled. Specifically, they fit a scaling law for optimal LR as a function of dataset size and find 1. Good fits to the scaling law and 2. The exponent is negative i.e. optimal LR decreases with increasing tokens. They also find that they can fit the loss vs log-LR curve by a quadratic (At least around the optimum value).
weaknesses: There are some other hyperparameters which strongly interact with learning rate such as weight decay and warmup. The paper does not explore the interaction between these factors and learning rate.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Can the authors provide any theoretical analysis or intuition for why the optimal learning rate scales with an exponent of approximately -1/3?

Could the authors use their fitted scaling law to propose a new LR schedule. For example, an LR schedule that decays with t^{-.32}. Does this outperform cosine decay? Is -.32 the best decay exponent for LR schedule? 

Could you strengthen your claims by testing the scaling law on other transformer architectures like Llama and Mistral, as well as non-transformer architectures like ResNets and MLPs? This would help demonstrate the generality of the observed scaling behaviour.
soundness: 3
strengths: The paper is clear and well written. It delivers what it promises.

This is an impactful research area. Efficient methods or scaling laws for finding best step-sizes helps to improve overall training efficiency and performance.
summary: The paper conducts an empirical study to characterize scaling law of optimal step-size with respect to the token used in the training of LLMs. They make some interesting observations: 1- best step-size shrinks approximately according to D^.32 where D is the training horizon and 2- this scaling law is almost consistent for different network sizes.
weaknesses: The paper provides no analysis or in-depth intuition on why optimal LR scales with exponent -.32 ~= -1/3. 

The scale of experiments are rather small for a fully empirical paper. Unless there is an analytical explanation or intuition provided for the observed patters, the trends might be unreliable for larger networks and longer horizons. In the same vein, would the scaling law remain the same for other transformers like Llama and Mistral? It would also be interesting to study how optimal LR scales with respect to training horizon in other architectures like ResNet and MLP.

In general, this is a solid paper, however in my view it falls marginally below the ICLR bar in terms of contributions. See below for some suggestions that I think would help in improving the contributions and the scores.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - If my understanding is correct, $\mu P$ describes a particular scaling scheme for how to adjust hypers as you scale width or depth (the latter being described in more recent depth-$\mu P$ extensions, like [1]). When you say you "scale in $\mu P$" but then sweep tokens at a *fixed* model size, I don't know what this means, or how it's different from standard parameterizations? Are you just referring to using a particular initialization scale, then? Can you explicitly tell me how you are scaling in the main (standard) vs $\mu P$ experiments and how they differ if you are not varying model size in the latter, for instance in Figure 5?
- It is known why we need to decrease LR with model size, $N$, and it has to do with reasons relating to how EoS/sharpness scale in $N$, for instance see [2]. Is there a similar theoretical reason why the optimal LR should be smaller on larger data budgets? Why should we expect a priori lower LR* when training for longer (especially since LR schedules usually decrease LR over training anyway)? 

[1] Bordelon, Blake, et al. "Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit." arXiv preprint arXiv:2309.16620 (2023).

[2] Noci, Lorenzo, et al. "Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning." arXiv preprint arXiv:2402.17457 (2024).
soundness: 3
strengths: I have not seen works touching on optimal LR with respect to data budget before, so this is neat. The message of the paper is clean and simple: optimal learning rate for language model pretraining decreases approximately as a power law in the data budget. In particular some strength are: 

- The message is short and simple, and the empirics are done with reasonably standard architectural decisions, and are -- most importantly -- at large scale. I imagine this was a reasonably expensive paper to write in terms of compute. 
- They check multiple parameterizations, including $\mu P$, which is important. 
- The validate their predictions and include a nice touch of showing an actionable consequence: that the LR of Llama was not tuned optimally. 
- The batch size factoring in the optimal learning rate is neat and intuitive and a good sanity check since the BS and LR are often optimized/tuned together in practice.
summary: The paper shows that the optimal learning rate is decreasing in dataset size across scaling parameterizations, in contrast to most existing work on scaling, which studies model size (depth/width) only, and not data budget. They find that the optimal learning rate decreases approximately as a power law in the data budget, and provide a heuristic to help practitioners select an optimal learning rate scaling over data budgets. Their experiments include language model pretraining runs up to 2.7B parameters and 800B tokens, on which they fit these scaling trends.
weaknesses: My gripes are mostly methodological. 

- In my mind, the main contribution is an empirical one: the dots you plot. I didn't put much weight into the actual fitted scaling laws because you fit a *separate set of constants to each curve* as you varied token horizons for a fixed model size. As I understand it, the main point of a scaling law is that you should be fitting a *fixed set of constants for all token horizons* and THAT is what you should be plotting. I suspect this is the reason for your unusually strong fit. The quote from von Neumann comes to mind: "With four parameters I can fit an elephant, and with five I can make him wiggle his trunk." Ideally, you should even be using one fitted set of constants ACROSS model sizes, but since your focus is on data and not model size scaling, I can forgive this. But one set of fitted constants per model size is a bare minimum: the current setup for how fitting is done is misleading if I am correctly understanding it. 

- It is unclear whether the main (scaling) experiments are done on the absolutely right type of architecture, ie. a modern "Transformer++," which is what they'd need to be to be most relevant to practice. Using architectural and hyperparameter choices from GPT3 is suboptimal because optimized Transformer++ architectures today use things like RMSnorm, RoPE embeddings, no linear biases, Adam $\beta$ values of $(0.9, 0.95)$, etc. Most people who work on pretraining know this, and the GPT3 architecture is certainly far from an optimized modern (2024) version by most people's standards. I see there are some preliminary experiments with the "Llama-1" architecture (though I'm not sure if it includes all the elements I outlined above). I believe the results, to be clear, and am not asking for new experiments, I just think there are some strange architectural choices even if mostly they are standard. Maybe including one ablation sweep with the architecture I described above (see perhaps the OLMo architecture for an example of a vanilla Transformer++) to check the same trends hold in the same way.  

- The functional forms you posit and fit are somewhat arbitrary. You choose a quadratic and power law, respectively, but do not explore other fits as I understand it, or even justify from any theoretical perspective why these should be the correct fits. Again, to be clear, I personally think these posited forms are fine, but they are indeed arbitrary and this requires justification.

- You definitely need to ablate LR schedule to make sure these are not an artifact of cosine LR schedule.

Meta-review: This paper studies the learning rate hyperparameter transfer problem across dataset sizes (i.e., token horizon) for large language models. The results show that the optimal learning rate becomes smaller as the token horizon length increases, and it follows a scaling law for any given architecture. This means it is possible to transfer the knowledge about optimal learning rate from shorter horizon to longer horizon. Interestingly the authors found that the learning rate used for LLamma-1 might be too large.

These are very interesting, novel and potentially impactful results. However, as the reviewers pointed out, the scope is very limited. Hyperparameters in the model training process are intertwined and can jointly affect the training results. Only focusing on learning rate can lead to biased conclusions. For example, it might be possible that the scaling law becomes different when other hyperparameters change values. 

During discussion, the reviewers pointed out that "the paper is slightly superficial, in the sense that it does not analyze whether the source of the shift is purely training time or more data" and "this paper does not do a lot beyond the simple singular claim it puts in the title". But the reviewers also praised that the paper is interesting and necessary since "the current theory behind learning rate scaling does not consider time and data as scaling quantities", and "its topic is important and crucially the experiments are done in a large scale and controlled settings with a clear trend that persists under ablations", and the new experiments resolved some concerns about the source of the shift.

Overall I believe the contributions are valuable and the ICLR community can benefit from it. Hence I recommend acceptance.

Decision: Accept (Poster)