Scaling Optimal LR Across Token Horizons
Authors: Johan Bjorck
Alon Benhaim
Vishrav Chaudhary
Furu Wei
Xia Song
Venue: ICLR 2025 Poster (2025)
Keywords: LLMs
scaling laws
hyperparameters
mup
Primary area: foundation or frontier models, including LLMs
TL;DR: We show how the optimal LR for LLM pretraining depends on token horizon

Abstract: State-of-the-art LLMs are powered by scaling -- scaling model size, training tokens, 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 muP. However, hyperparameter transfer across training tokens -- 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 thus argue that hyperparameter transfer across data size is an overlooked component of LLM training.

Published as a conference paper at ICLR 2025
SCALING OPTIMAL LR ACROSS TOKEN HORIZONS
Johan Bjorck2. 2t.
Alon Benhaim 1.*
Vishrav Chaudhary
Furu Wei
Xia Song'
Microsoft
Nvidia
Meta
ABSTRACT
State-of-the-art LLMs are powered by scaling scaling model size, training to- kens, and cluster size. is economically infeasible extensively hyper- parameters for the largest runs. Instead, approximately optimal hyperparame- ters must inferred or transferred from smaller experiments. Hyperparameter transfer across model sizes has been studied in Yang al. (2022). However, hy- perparameter transfer across training tokens - or token horizon - has not been studied yet. To remedy this conduct large-scale empirical study how op- timal learning rate (LR) depends on the token horizon LLM training. We first demonstrate that the optimal LR changes significantly with token horizon - longer training necessitates smaller L.R. 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 rule-of- thumb for transferring LR across token horizons with zero overhead over current practices. Lastly, we provide evidence that LLama-1 I used too high LR, and thus argue that hyperparameter transfer across data size an overlooked component of LLM training.
INTRODUCTION
State-of-the-ar LLMs are scaled in multiple dimensions. The models are becoming increasingly large, e.g. Grok-1.5 has 314 billion (B) parameters (xAl 2024). 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 2024). Lastly, the training datasets growing, e.g. LLama-3 was trained on 15 Trillion (T) tokens (Dubey et al. 2024). At these scales, it is infeasible to extensively tune hyperparameters. Practitioners must instead resort to hyperparameter transfer, process where approximately optimal hyperparameters for large-scale experiments are inferred from experiments smaller scale. Perhaps the most famous work on hyperparameter transfer muP Yangg al., 2022 methodology for transferring optimal hyperparameter from small model to a large model.
Equal contribution. Work done while Microsoft.
25B tokens
Val Loss
3.15
3.05
400B tokens
Val Loss
2.88
2 21
2
2
Figure Final validation loss of 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.

---

Published as a conference paper at ICLR 2025
While hyperparameter transfer across model size is well-studied problem, transfer across token horizon remains understudied. This paper aims to remedy this shortcoming in the literature.
In this paper, present a large-scale study on hyperparameter transfer across token horizons, define the latter as the total number tokens seen during training. We specifically focus on learning rate (LR), an important hyperparameter that influences training stability, generalization, and con- vergence speed. Our study is essentially a large ablation experiment where we vary LR and token horizon for few different LLM models. We consider >50 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. 2019) and the Refined dataset (Penedo et al. 2023 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 long horizon is inferred from a shorter horizon. Furthermore, we provide mule-of-thumll 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 ev- idence that LI L 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 2017) trained via next token prediction on web-scale text (Radford et al. 2019). 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. 2022]. Kaplan al. (2020) shows that L.LM performance roughly scales as 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 Ne, (:x) for models trained sufficiently large datasets: L(N) (Ne/N)ac
Hoffmann et al. (2022) also showed that the performance scales well with token horizon and that the optimal performance for given FLOPs budget is obtained by scaling the model and token horizon jointly. The current paradigm thus scales token count in addition to model size e.g. LLama-3 was trained on > 10x many tokens LLama-1 (Dubey et al. 2024 In this paper we use notation from Kaplan et al. 2020). denoting the token horizon (i.c. number of tokens seen during training, some of which might be duplicate) 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. Following Kaplan et al. (2020), we will use power-laws of the form F(X) = AXx, where a and A are fitted and X the quantity of interest. To measure goodness-of-fi we will use the R²2 measure its value ranges from 0 to I, where indicates perfect fit and indicates no fit.
Hyperparameter transfer. Hyperparameters can strongly influence the performance of LLMs, and LR is a critical hyperparameteer The muP paper (Yang et al., 2022) first popularized hyperpurameter transfer the process of finding optimal hyperparameters from small proxy experiment. Core to muP is the small model is also optimal for larger model. The muP parametrization introduces some changes to the network, e.g. the attention scaling factor is changed from va to d' The original muP paper shows that LR and many other hyperparameters transfer from small to large models when muP parametrization is used, which further corroborated in Lingle (2024). While Yang al.) (2022) focuses transfer between model sizes, we focus transfer between token horizons.
a way of parametrizing an LLM such that the optimal LR for a
EXPERIMENTS
Experimental Setup. Our experimental setup closely follows the setup of the GPT-3 paper (Brown, 2020). We consider model sizes of 50 million (m), 125m, 350m, 760m, 1.3 billion (B), and 2.7B
2

---

Published as a conference paper at ICLR 2025
parameters using the architectures of Table 2.1 in the GPT-3 paper. The table is replicated as Table 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 Appendix AA] The LR will vary with the training step since we use LR schedule, and will the term LR refer to the largest LR of run. We make three adjustments compared to the GPT-3, mainly to prevent divergence of the model for high
LRs. For warmup we use the maxima of 1% of the training steps following MMuennihoffi et al.
2024) and 1000. Having too short warmup stage is known to cause divergence (Wortsman 2023) which we want to avoid for short training runs. 2) We use qk-norm following Dehghani
al.
et al.
(202)). 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. 2023), a common Nommon-crawl derived dataset of roughly 600B tokens which is known to be of high quality (Penedo et al. 2024). Experiments are run on the Megatron codebase (Shoeybi et al. 2019). Unless mentioned we will use the same seed for all runs. For curve fitting least-square fitting with either first second-degree polynomial using Numpy and Scipy (Harris et al. 2020), and we always have more data points than free parameters in the fits.
2023).
3.1 ABLATIONS
In our first experiment, we consider the 350m LLM model from Table Appendix [ 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 10-4 from Table which is what the GPT-3 paper used (Brown) 2020). We then multiply this base LR by factors 00.25, 0.5, 1, 4]. The LRs we consider are thus {7.5 x 10-5, 1.5 x
x 10-4, 10-4 1.2 10 10-33. For each combination of LR and token horizon, we train 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), using quadratic polynomial as it provides an excellent fit and is the simplest polynomial with well-defined minimum. The R² of the fit are 0.995 or better, see Table
350m parameter model
3.2
25B tokens
50B tokens
100B tokens
200B tokens
3.1
400B tokens
Val Loss
3.0
2.9
2 2-13
2-12
2 2-11
2 10
LR
Figure 2: Final validation loss a function of learning rate (LR) and token horizon. The dashed lines indicate our fitted curve and the stars indicate optimal LRR. The optimal LR decreases monotonically with longer horizons.
3

---

Published as a conference paper at ICLR 2025
50m parameter model
125m parameter model
3.625
3.3
Val Loss
3.575
3.525
3.2
3.475
258 tokens
50B tokens
2-1
2-
2
2-10
2
2
2-12
2-11
2-10
2-9
100B tokens
760m parameter model
2.7b parameter model
2008 tokens
4008 tokens
2.9
800B tokens
3.0
Val Loss
2.8
2.9
2.8
2.6
2 2
2-13
2 2-12
2
2 1
2
2-13
2
2-11
LR
LR
Figure 3: Final validation loss as function of max learning rate (LR) and token horizon for four models. The dashed lines indicate fitted curve. The optimal LR, denoted by black star, de- creases monotonically with longer horizons for models.
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, consider the base LR as the one used in the GPT-3 paper (also listed Table 4 and then multiply it by e.g. {0.25, 0.5, 1, 2, 4J. The results of these experiments are shown in Figure 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 and observe that the optimal LR decreases with longer token horizons across model sizes. Figure TO in the Appendix show similar results for the 1.3B model. We thus conclude that this is robust phenomenon.
don't 't increase the sampling rate around the minimizer. For each model,
3.2 SCALING LAWS
We have seen that longer token horizons require smaller LR. We now investigate this insight allows derive scaling laws and do hyperaprameteer-ransferr i.e. finding the optimal LR 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" LR) BD-
(1)
Here B and 3 are two constants independent of that might e.g. depend on the model architecture. get
Taking the logarithm of both sides of Equation
log [I.R* (D) log B B log 1)
We thus have a linear equation in the unknowns log B, 3. and can fit these to our experimental results with least squares. Specifically, the minimizer of the quadratic fits of Figure as the data
4

---

Published as a conference paper at ICLR 2025
50m model
125m model
2
Real
Real LR
Fitted LR
2-10
Fitted LR
Learning Rate
Learning Rate
2-10
2-11
2 2-11
12
2-12
2
2
237 Tokens
2
2 29
2
2
2 Tokens
2
2.39
350m model
760m model
Real LR
Real LR
Fitted LR
Ritted LR
Learning Rate
Learning Rate
2 2-11
2-11
2
236
237
2
2
2
237
Tokens
Tokens
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² of these fits are in the range 0.99 0.96. Across all model sizes, we see that the scaling law provides good fit to the empirical data.
points LR* Fitting Equation 2 gives good fits the curves for four models are shown in Figure The R of these fits are in the range 0.99-0.96 (see Table55jn Appendix ( for exact values), rather high number. In Figure [3 in the Appendix we also show the fits for the 1.3B and 2.7B models, which also show a good fit. In Figure 11 in the Appendix we repeat these experiments scale by using the Llama architecture on the smallest 50m model, and find that the 3 transfer well across architectures when using the same model size. In Figure 12 repeat these experiments by doing multiple epochs over a fixed dataset of 25B tokens, and observe similar scaling behavior.
small
To evaluate the predictive power of the scaling laws cannot solely rely on Rt 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, 8 on token horizons 25B, 50B and 100B and then use these constants to predict the optimal LR at longer horizons. We will the optimal LR we have empirically found with quadratic fits the correct LR. The results illustrated for the 50m model in T We see a reasonably good fit with a relative error of 10-15% clear improvement over not scaling LR all which has relative error of > 250%. We thus conclude that it is possible to perform hyperparameter transfer across token horizons with our scaling laws. Table Tis repeated for more models with similar fits in Appendix B
3.3 MUP PARAMETRIZATION
natural to ask if muP allows hyperparameter transfer across token horizons. To investigate this we consider the 50m model from Table 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, 2,44 times the base LR in Table 4 The results of this experiment
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.
shown in Figure 5 and
5

---

Published as a conference paper at ICLR 2025
25B
50B
100B
200B
400B
800B
Optimal LR
1.54 10
x x 10"
6.06
10
3.33 x 10
2.14x 10
1.71x 10
Predicted LR
x
2.39
10-4
1.50x10
Ratio
0.873
0.894
1.14
Table I: Simulated hyperparameter transfer across token horizons for 50m model. We use the empirically measured optimal LR at 25,50 and 100B tokens estimate the optimal LR at longer horizons by fitring the constants of Equation (2) We find reasonable fit when scaling up the horizon 800B tokens, with a relative emror 10-15 The relative error of using the best LR at the 100B horizon for an 800B horizon is > 250%. The scaling laws thus have predictive power.
50m parameter muP model
25B tokens
50B tokens
3.35
100B tokens
Val Loss
3.25
2-12
2-11
2 2-10
2-9
LR
Figure 5: Optimal LR vs token horizon for a 50m model using muP parameterization (Yang et al. 2022]. We see that the optimal LR decreases with longer token horizons, demonstrating that LR does not transfer across horizons even with muP.
3.4 QUANTIPYING VARIANCE
To ensure that results reliable we here consider quantifying the variance of our experiments. We will use two different methodologies. Firstly we use bootstrapping, following Hoffmann et al. (2022). We consider the 350m model of Table randomly remove 20% of the data points, and then fit the optimal LR at each token horizon and the scaling law of Equation (I) We repeat this procedure 1000 times and then measure the mean and of the optimal L.R and the constants of Equation ( 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 small-scale experiment with multiple seeds. We use the 350m model from Table a 100B token horizon, three learning rates, and two additional seeds. We then estimate the optimal LR via 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 o j is 2.63x 10-2, so we can expect the relative error to be few percent. When using more data points the variance could further decrease due concentration of measure.
3.5 EFFECT OF BATCH SIZE
is well known that batch size BS affects the optimal LR (Goyal 2017). While we primarily focus on the setting of fixed batch size, we here consider modifying the batch size for the 1.3B model of Table Specifically, we double the batch size to Im tokens and train for 25B and 50B tokens in total with different LRs. In Figure in the Appendix we show the final validation loss a function of the LR and estimate the optimal LR. In Figure 15 we show the optimal LR as a function of token horizon for the 1.3B model using batch size of 0.5 Im tokens. We that the optimal L.R is

---

Published as a conference paper at ICLR 2025
25B
50B
100B
200B
400.0
((RR))
7.04x 10
5.96 10
5.07 x 10
3.45 10
2.44 10 10 -4
(LR))
1.28 x 10-b
9.12x10
1.96 10 10-6
4.07 10
4.31x10-1
o/p
1.82x10-2 1.53x10-2
3 10-2
1.18x 10
1.76x10-
Table 2: We estimate the mean j and standard deviation of the optimal learning rate LR" 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 о/μ is on the order of a few percent. This implies that the uncertainty in our estimates is small.
N-350M OM
2-10
N=760M
2
508
N=1.30
1008
N-2.78
2008
N=
Optimal LR
Optimal
2
2-12
2-12
25B
50B
100B
2008
0.358
0.78 1.0B
3.5B
7.0B
Dataset size D
Model size N
Figure 6: (a) There is a linear relationship between log D and log LR for different model sizes N. (b) There linear relationship between log N and log LR" for different values of dataset size D.
higher with larger batch size expected. More importantly, we that the linear fits roughly parallel. This suggests that the optimal LR depends on the token horizon the same way irrespective of the batch size, i.e. that can factorize Equation I
LR(BS, D) F(BS)D
SCALING LAW WITH RESPECT TO MODEL SIZE
So far have considered Equation 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, provide some preliminary results and discussions regarding a scaling law for both model size N and token horizon D. In Figure (6b] we plot the optimal LR function of N, and Figure (6a] we plot it as 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
will here
LR' (N, D) CN-*D-
Here, the constant C' the ideal learning rate for 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 captures the fact that the optimal LR decreases with model size N, while the factor D captures that optimal LR also decreases with token horizon D. We now fit Equation 3 to data for the models of size 760m, 1.3B, and 2.7B, using the 7B model of Section 4.1 as held-out validation set. We use the huber loss with = x 10-3 and the BFGS algorithm, similar Hoffmann al. (2022). We observe good fit with RMSE= 2.2 10 and validation R² of 0.978. The numerical values we find are:
1.55 x 10 a ~ 0.23,
0.32
(4)
We depict the results in Figure 7 For smaller models (mainly 50m and 125m) larger is observed from the experiments. We thus consider Equation (3) to be valid in the regime of large models (2 760M parameters).

---

Published as a conference paper at ICLR 2025
0.0008
N=760M
0.0007
N=1.3B
N=2.7B
0.0006
N=7B validation
Optimal LR
0,0005
0.0004
0.0003
0.0002
0,0001
08
508
100B
1508
2008
2508
300B
Dataset size
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² of 0.978 on this validation data. Note, the 7B model uses the Llama architecture while the other data points use the GPT-3 architecture. This experiment thus demonstrates that the scaling laws have predictive power across model architectures.
2-11
2-12
Optimal LR
2-13
2-
Real LR
2 2-15
Fitted LR
Llama 1 used
Llama 1 optimal
2 2-16
10B
25B
508
100B
250B 500B 1000B
Dataset size
Figure 8: The optimal LR of the LLama-1 7B model as a function of token horizon D. We find that 1.15 10-4 4 optimal for 1T tokens whereas the Llama-1 paper used 10-4.
4.1 A CASE-STUDY ON LLAMA-1
We now consider evaluating if the LRs used by Llama-1 Touvron et al. (2023) are "correct" accord- ing scaling laws. To do this we adopt the LLama-] architecture (RMSnorm, Rope embed- dings, and so on) and run small-scale experiments with token horizons 25B, 50B, and 100B and dif- ferent LRs. Based upon these experiments, we find that values for Equation (1) are
8.29 10" and B = 0.3 (see Figure 14jin Appendix B). We then extrapolate to find the optimal LR at 1T tokens, which comes out as 1.15 10 A
LLama-1 used 3 10-4 an LR which too large by 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 Table g 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. AL L(D = 100B, LR = LR*) L(D 100B, LR = 2.6LR*) = 0.027.
Note, this is upper bound since the parabolas flatten with longer token horizons
8

---

Published as a conference paper at ICLR 2025
5 RELATED WORK
The study of scaling laws for LLMs originated in Kaplan al. (2020) and it is still an active area of research. Scaling laws are researched in the context post-training (Lin et al. 2024] Zhang et al.
2024at Gao et al. 2023), model architecture (Krajewski et al. 2024] Alabdulmohsin et al. 2024
Frantar et al. 2023] Wang et al. 2023). multi-modality (Aghajanyan et al., 2023] Cherti et al.
2023). inference (Sardana & Frankle 2023), data (Dohmatob et al. [024]] Fernandes et al., (202)
and other domains (Zhang et al. 2024b Neumann & Gros] 2022 [ro2|2002| Kadra et al. 2024). There are also
more theoretical studies (Michaud al. 2024 Caballero et al. (022)..
LR has long been known as an important hyperparameter, and hence there is ample work on how to select it (Goyal] 2017). Larger models are known to require smaller LR, and e.g. Kaplan et al. (2020) suggests the formula LR((N) Rs 0.003239 0.0001395 (00151gg(N)) for tuning LR as a function of model size. MuP (Yang et al. 2022) is a principled methodology for selecting LR when the model is scaled, and the method is actively researched (Lingle [024] Everett et al.] 2024] Noci al. 2024 Blake et al. 2024). Recently Evvreett et al.) (2024) showed that LRR transfer across model size is possible both with different parametrizations and optimizers. However, a limitation in both Yang et al. (2022) and Evveeett et al.] (2024) is the fixed training horizon assumption used in the theoretical derivations. Our work directly explores this limitation. One notable conclusion in [verett et al. (2024) 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, closer look at Table 11 in Everett et al. (2024) shows that for the Adam optimizer, the scaling exponents are similar for the two token horizons they considered. Another related paper is Wang & Aitchison (2024), which studies how set the weight decay of Adam when model and dataset is scaled. At last, we mention that Bi et al. (2024) 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.
DISCUSSION
Limitations. To limit the scope and computational requirements of our study we have intention ally focused 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. 2024).. 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. 2017), model width and depth, different atten- tion types (Ainslie al. 2023), and space models (Gu Dao 2023) could plausibly interact with both the IRR and token horizon. The number of repeated tokens can also play a role, and we only provide small-scale experiments on these theme in Figure [22 in the Appendix. Other dimensions to consider are additional hyperparameters such weight decay which interacts with LR (Bjorck) cet al., 2021). ). LR schedules, and multi- modaity (Huang et al. 2023). For computational reasons, we defer these topics to future work.
Our experiments show that the optimal LR decreases with the token hori- zon. 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 / 0.32 to generalize across architectures. To find the optimal LR T.R(Di) at some long horizon Dj practitioners can just find the optimal LR LR (D2) at a short horizon D2 and then estimate:
Advice for Practioners.
D2 D1
-0.3
LR (D2) Re LR (D:) (Di)
(5)
Since practitioners typically find the optimal LR for smaller horizons anyway during hyperparam- eter search, this methodology has no overhead over current practices. For practitioners with am- ple 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 ( can be found, and the optimal LR at longer horizons can be estimated. For practitioners using methods that provide

---

Published as a conference paper at ICLR 2025
zero-shot transfer of optimal LR across width Yang al. (2022) and Everett al. (024 )
recommend using an adjustment of Equation G3 (N x Rlayer model) which yields the formula LR (Mager, D) = Ca D-8 Then needs to be found using the usual muP sweep for LR with small width, depth and horizon. This allows hyperparameter transfer 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 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.
REFERENCES
Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for gen- erative mixed-modal language models. In International Conference on Machine Learning, pp- 265-279. PMLR, 2023.
Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305.1 13245, 2023.
Ibrahim M Alabdulmohsin, Xiaohua Zhai, Alexander Kolesnikov, and Lucas Beyer. Getting vit in shape: Scaling laws for compute-optimal model design. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Paul Alcom. Elon musk fires up the most powerful ai training cluster in the world, uses 100,000 nvidia h100 gpus on single fabric, 2024. Accessed: 2024-09-17.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXhv:2401.02954 2024.
Johan Bjorck, Kilian Weinberger, and Carla Gomes. Understanding decoupled and early weight decay. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 6777- 6785, 2021.
Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Y Prince, Bjöm Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr u-mup: The unit-scaled maximal update parametrization. arXiv preprint arXiv:2407.17465. 2024.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. arXiv preprint arXiv:2210.1489)) 2022.
Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gor- don, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. Reproducible scaling laws for contrastive danguage-imag learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2818-2829, 2023.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski. Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to billion parameters. In International Conference on Machine Learning, pp. 7480-7512. PMLR, 2023.
Elvis Dohmatob, Yunzhen Feng, Pu Yang, Francois Charton, and Julia Kempe. A tale of tails: Model collapse change of scaling laws. arXiv preprint urXiv:2402.0704. 2024.
10

---

Published as a conference paper at ICLR 2025
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXhv:240..17783, 2024.
Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander Alemi, Roman Novak, Peter Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, et al. Scaling expo- nents across parameterizations and optimizers. arXiv preprint (txXiv:4077.0872,2 2024.
Patrick Fernandes, Behrooz Ghorbani, Xavier Garcia, Markus Freitag, and Orhan Firat. Scaling laws for multilingual neural machine translation. International Conference on Machine Learning, pp. 10053-10071. PMLR, 2023.
Elias Frantar, Carlos Riquelme, Neil Houlsby, Dan Alistarh, and Utku Evci. Scaling laws for sparsely-connected foundation models. arXiv preprint arXiv:2309.08520 2023.
Leo Gao. John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835-10866. PMLR, 2023.
P Goyal. Accurate, large minibatch sg d: training imagenet in hour. arXiv preprint
1XXiv:7006.06677, 2017.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Charles R Harris, K Jarrod Millman, Stéfan Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array program- ming with numpy. Nature, 585(7825) 2020.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203.15556 2022.
Shaohan Huang, Li Dong, Wenhui Wang, Yaru Hao, Saksham Singhal, Shuming Ma, Tengchao Lv, Lei Cui, Owais Khan Mohammed, Barun Patra, et al. Language is not all you need: Align- ing perception with language models. Advances in Neural Information Processing Systems, 36: 22096-72109 2023.
Arlind Kadra, Maciej Janowski, Martin Wistuba, and Josif Grabocka. Scaling laws for hyperparam- eter optimization. Advances in Neural Information Processing Systems, 36, 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXhv:2001.0836/ 2020.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóddz, Piotr Sankowski, et al. Scaling laws for fine-grained mixture experts. arXiv preprint arXiv:2402.07871 2024.
Haowei Lin, Baizhou Huang. Haotian Ye, Qinyu Chen, Zihao Wang, Sujian Li, Jianzhu Ma, Xiaojun Wan, James Zou, and Yitao Liang. Selecting large language model fine-tune via rectified scaling law. arXiv preprint arXiv: 2rXiv:2402.0233 02314, 2024.
Lucas Lingle. A large-scale exploration of mu-transfer. arXiv preprint arXiv:2404.05728, 2024.
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling-
Advances in Neural Information Processing Systems, 36, 2024.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling data-constrained language models. Advances in Neural Information Processing Systems, 36. 2024.
Oren Neumann and Claudius Gros. Scaling laws for a multi-agent reinforcement learning model. arXiv preprint 10.00992 2022.

---

Published as a conference paper at ICLR 2025
Lorenzo Noci, Alexandru Meterez, Thomas Hofmann,, and Antonio Orvieto. Why do learning rates transfer? reconciling optimization and scaling limits for deep learning. arXiv preprint arXiv:2402.17457, 2024.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon Ilm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116 2023.
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint w4Xv:2406.17557 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Nikhil Sardana and Jonathan Frankle. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws, arXiv preprint arXiv:2401.00444 2023.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 701.06538, 2017.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint 1eXiv:1909.08000 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Peihao Wang, Rameswar Panda, and Zhangyang Wang. Data efficient neural scaling law via model reusing. In International Conference on Machine Learning, pp. 36193-36204. PMLR, 2023.
Xi Wang and Laurence Aitchison. How to set adamw's 's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405. 13698, 2024.
Jason Wei, Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, al. Emergent abilities of large language models. arXiv preprint arXiv:2206.08 07682, 2022.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint 2rXivv2309.1133300 2023.
xAI. Open release of grok- March 2024. URL 2024-09-17.
grok-os Accessed:
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Biao Zhang, Zhongtao Liu, Colin Cherry, and Orhan Firat. When scaling meets I1m finetuning: The effect of data, model and finetuning method. arXiv preprint arXiv:2402.17199, 2024a.
Buyun Zhang, Liang Luo, Yuxin Chen, Jade Nie, Xi Liu, Daifeng Guo, Yanli Zhao, Shen Li, Yuchen Hao, Yantao Yao, et al. Wukong: Towards a scaling law for large-scale recommendation. arXiv preprint arXiv:2403.0 245.2450,, 2024b.
12

---

Published as a conference paper at ICLR 2025
A HYPERPARAMETERS
Hyperparameter
Value
weight decay
0.1
grad clip norm
1.0
LR schedule
cosine
Adam 81
0.9
Adam B2
0.95
Context length
2048
Batch size (tokens)
524288
Warmup Steps
max 1000, 0.01 x train iters)
Min LR
0.1x Max LR
Table 3: Hyperparameters. These follow Brown (2020).
Hyperparameters not listed here follow
defaults in the Megatron codebase Shoeybi et al. (219))
Model Name
params
layers
dmodel
heads
Base LR
Tiny
50M
8
512
8
8 10
Small
125M
12
768
12
10-4
Medium
350M
24
1024
16
10-4
Large
760M
24
1536
16
2.5 x 10-4
1.3B
1.3B
24
2048
16
2 10
2.7B
2.7B
32
2560
32
1.6 10-4
6.7B
6.7B
32
4096
32
1.2
10
Table 4: Model architectures and base LRs we consider.
These follow GPT-3 Brown 2020
13

---

Published as a conference paper at ICLR 2025
B ADDITIONAL EXPERIMENTAL DATA
Model size
R
B
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 5: The R values and found Bs of the fits in Figure 4
The R²s are relatively close 1,
indicating a good fit. The 3s are relatively similar for models 350m, but are larger for the smallest models.
Token Horizon
R²
25B
1- 1-6.88x 10
50B
1- 10
100B
1 10-3
200B
2.85 10-3 x.1.76x10-
400B
Table 6: The R2 values of the fits in Figure 2
They are very close to 1, indicating great fit.
1.3B parameter model
25B tokens
50B tokens
Val Loss
2.9
2.8
13
12
2 -11
LR
Figure 9: Learning rate and validation loss different token horizons for the 1.3B model from Table 4 Jusing Im batch size.
14

---

Published as a conference paper at ICLR 2025
1.3B parameter model
25B tokens
3.0
50B tokens
100B tokens
200B tokens
2.9
Loss
- 2.8
2.7
2-14
2-13
2-12
2-11
LR
Figure 10: Learning rate and validation loss for different token horizons for the 1.3B model of Table We see that the optimal LR decreases as the token horizon increases.
Liama model
Liama madel
258 tokens
50B tokens
1008 tokens
3.48
Learning Rate
Val
3.40
LA
2-12
2
Tokens
Figure 11: Experiments the 50m model using the Llama-1 architecture (RMSnorm, Rope em- beddings, and so on) of Touvron et al. (2023). Similar to the GPT-3 architecture, we see that optimal LR decreases with the token horizon. On the right plot the scaling law, using the same B as the GPT-3 model. We see an excellent fit, demonstrating that the transfers well across model archi- tectures.
seed
I(1.5x10-4)
L(3x10-4)
L(6x10-4)
arg min L
1
2.940372
2.919948
2.913585
5.81 10
2
2,941199
2.919131
2.912387
5.76 10
3
2.941648
2.920779
2.915190
x.47
Table 7: We consider the loss L as a function of the LR. We repeat the experiments of Section 3.1 with 100B token horizon, 350m model, and multiple seeds. With different seeds, see slightly different final losses, and slightly different estimates of the optimal LR.
15

---

Published as a conference paper at ICLR 2025
25B unique takens
Fixed model
258 tokens
Rmal L. Tiited
100B tokens
3.60
Leatning Rate
Loss
3.56
3.52
2-42
2
2
Figure 12:
We sample a set of 25B unique tokens, and ablate LR when the training horizon is 25B, 50B and 100B tokens. This corresponds to one, two, or four epochs over the data. We use the 50m model. (Left) The optimal LR decreases when we increase number of unique tokens seen does not solely determine the optimal LR. (Right) The slope between the optima in this setting fits the scaling curve from training only only unique tokens. This suggests that it is the token horizon rather than the number of unique tokens which determines the scaling.
total token horizon, thus the total
1.3B model
2.7b model
2-11
Real LR
2-11
Fitted LR
Learning Rate
Learning Rate
Real
Fitted LR
235
23
237
2
2 36
Tokens
Tokens
Figure 13: 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 further details.
works well. See Section 3.2 for
Real
Fitted LR
Optimal LR
2-12
100B
Dataset size
Figure 14: Optimal LR vs token horizon for 7B LLama-1 model using batch size of 4M..
16

---

Published as a conference paper at ICLR 2025
0.5M 1M Batch Size
1M BS
0.5M
2 2-11
LR
235
2
237
Tokens
Figure 15: Optimal as a function token horizon for 1.3B model using a batch size of 0.5m or Im tokens. A larger batch size implies that 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.
Token Horizon
Optimal LR
25B
3.4 x 10
50B
2.8 x 10 44
100B
2.3 x 10-4
IT (Using Eq.
1.15 x 10-1
1T (Using Eq.
1.1 x 10
Table 8: Optimal LR for LLama-1 for different token horizon using the fit in Figure 14 Note that the real Llama-1 used LR of x 10-4 which our results suggest is significantly large.
25B
50B
100B
200B
400B
800B
Optimal LR Predicted LR Ratio
1.34 x 10 33 1.02 x 10
6.60 10 4
4.12 10
2.51 10 4
1.98
4.77 10
3.35
10-4
2.35 10-4
0.864
0.749
0.843
Table 9: Simulated hyperparameter transfer across token horizons for 125m model. The setup follows Table
R22
B
0.961
0.384
x8
7.78
10 3
o/p
9.21x10
2.03x100 -2
Table 10: We show the mean μ, standard deviation and relative deviation a /l/ estimated via bootstrapping for two quantities the R22 and 3 in Section 3.4 The relative deviation of / is roughly 10%, demonstrating that the uncertainty is reasonably small in our scaling laws.
17

---

Published as a conference paper at ICLR 2025
C
DERIVATION
In section 3.2 we saw empirically that for a given N and constants B, that may be dependent on
N,
LR (N, D) = B(N)D B(N)
(6)
As seen in Figure 6b we observe a linear relationship between log N and LR" for any given D.
This suggests the relationship:
(R D) A(D)N
The constants A, 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 align with the
from Eq. 7 for constant D has be representable by Eq. 6 is anticipated latter, consisting of distinct power laws, each with specific parameters for different N and D values. Consequently, the objective is to identify function that fulfills these criteria
LR" D) B(N)D P(N) A(D)N (D)
Dataset size D for different Model Size N. As seen in Figure 6a since the lines are parallel for any given N, the slope 3 (Eq. is independent of the model size N. Therefore we can assume that 8(N)= =
Model size N for different Dataset size D. As seen in Figure 6b since the lines are parallel for any given D, the slope a (Eq. 7 is independent of the dataset size D. Therefore can assume that
a(D)= = a is constant.
Using the fact that a, are constant and multiplying Eq. 8 by No
B(N) A(D)D
The LHS of Eq. only depends N, whereas the RHS only depends D they should both equal some constant, C (this step relies on our proof above that a, 3 are independent of N, D), resulting in the functional forms of A(D), B(N)
B(N) A(D) = CD B(N) = CN
(10)
Plugging the functional forms of B(N) we finally get the final functional form for the joint scaling law as in Eq. 3
18


Reviews:
[{"review_id": "o1JStWCd0M", "reviewer": "Reviewer_ikK6", "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.", "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)?\n\n**Update**\n\nI thank the authors for their additional answers.\n\nI 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. \n\nThe 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). \n\nI 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.\n\nOverall, 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper has significant strong points.\n\n1. 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).\n\n2.  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.\n\n3. The case study on LLAMA training is very interesting, for instance, by retrospectively questioning how hyperparameter tuning was performed in that setting.", "weaknesses": "The paper has some weaknesses, mainly due to the depth of the investigation that is performed. More concretely (in order of importance):\n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n\nMinor:\n\n4. 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. \n\n5. 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.\n\n6. 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.\n\n[1] Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks\n\n[2] Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit\n\nI am willing to raise my score if my concerns are resolved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GRhbxHFoFI", "reviewer": "Reviewer_8GV3", "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).", "questions": "The warmup used by the authors is much smaller than that used in many recent works[1].\n\n\nTo 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.\n\n \n[1] SMALL-SCALE PROXIES FOR LARGE-SCALE TRANS-\nFORMER TRAINING INSTABILITIES.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kpxd1sCrej", "reviewer": "Reviewer_yU6o", "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.", "questions": "Can the authors provide any theoretical analysis or intuition for why the optimal learning rate scales with an exponent of approximately -1/3?\n\nCould 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? \n\nCould 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is clear and well written. It delivers what it promises.\n\nThis is an impactful research area. Efficient methods or scaling laws for finding best step-sizes helps to improve overall training efficiency and performance.", "weaknesses": "The paper provides no analysis or in-depth intuition on why optimal LR scales with exponent -.32 ~= -1/3. \n\nThe 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.\n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O2DbmhlOAE", "reviewer": "Reviewer_VnBU", "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.", "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?\n- 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)? \n\n[1] Bordelon, Blake, et al. \"Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit.\" arXiv preprint arXiv:2309.16620 (2023).\n\n[2] Noci, Lorenzo, et al. \"Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning.\" arXiv preprint arXiv:2402.17457 (2024).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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: \n\n- 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. \n- They check multiple parameterizations, including $\\mu P$, which is important. \n- The validate their predictions and include a nice touch of showing an actionable consequence: that the LR of Llama was not tuned optimally. \n- 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.", "weaknesses": "My gripes are mostly methodological. \n\n- 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. \n\n- 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.  \n\n- 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.\n\n- You definitely need to ablate LR schedule to make sure these are not an artifact of cosine LR schedule.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)