A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules
Authors: Kairong Luo
Haodong Wen
Shengding Hu
Zhenbo Sun
Zhiyuan Liu
Maosong Sun
Kaifeng Lyu
Wenguang Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Large language model
Learning rate scheduler
Scaling Law
Hyperparameter optimization
Primary area: foundation or frontier models, including LLMs
TL;DR: Loss curve prediction and optimized learning rate schedule

Abstract: Training large models is both resource-intensive and time-consuming, making it crucial to understand the quantitative relationship between model performance and hyperparameters. In this paper, we derive an empirical law that predicts pretraining loss for large language models for every intermediate training step across various learning rate schedules, including constant, cosine, and step decay schedules. Our proposed law takes a multi-power form, combining a power law based on the sum of learning rates and additional power laws to account for a loss reduction effect as learning rate decays. We validate this law extensively on Llama-2 models of varying sizes and demonstrate that, after fitting on a few learning rate schedules, it accurately predicts the loss curves for unseen schedules of different shapes and horizons. Moreover, by minimizing the predicted final pretraining loss across learning rate schedules, we are able to find a schedule that outperforms the widely-used cosine learning rate schedule. Interestingly, this automatically discovered schedule bears some resemblance to the recently proposed Warmup-Stable-Decay (WSD) schedule (Hu et al, 2024) but achieves a slightly lower final loss. We believe these results could offer valuable insights for understanding the dynamics of pretraining and for designing learning rate schedules to improve efficiency.

Published as a conference paper at ICLR 2025
A MULTI-POWER LAW FOR Loss CURVE PREDICTION ACROSS LEARNING RATE SCHEDULES
Kairong Luo
Haodong Wen Shengding Hul
Zhenbo Sun'
Zhiyuan Liu
Maosong Sun Kaifeng Lyu
Wenguang Chen 144
l Deppartment of Computer Science and Technology, Tsinghua University
2 Qian Xuesen College, Xi'an Jiaotong University
Simons Institute, University of California, Berkeley *Peng Cheng Laboratory
(luckr24, sunzb20) émaii i 1s t.s wnggua.eduuu.. fherrywenh, weceennniingggglln (eaaallm@mmmm..m com
kaifenglyu@berke edu
{liuzy, sms, cwg stgighuung edu. cn
ABSTRACT
Training large models is both resource-intensing and time-consuming, making it crucial to understand the quantitative relationship between model performance and hyperparameters. In this paper, present empirical law that describes how the pretraining loss of large language models evolves under different learning rate schedules, such as constant, cosine, and step decay schedules. Our proposed law takes a multi- multi-power power form, combining a power law based on the sum of learn- ing rates and additional power laws to account for a loss reduction effect induced by learning rate decay. We extensively validate this law on various model sizes and architectures. and demonstrate that after fitting on a few learning rate schedules. the law accurately predicts the loss curves for unseen schedules of different shapes and horizons. Moreover, by minimizing the predicted final pretraining loss across learning rate schedules, we are able to find a schedule that outperforms the widely used cosine learning rate schedule. Interestingly, this automatically discovered schedule bears some resemblance to the recently proposed Warmup-Stable-Deew (WSD) schedule (Hu et al., 2024) but achieves a slightly lower final loss. We be- lieve these results could offer valuable insights for understanding the dynamics of pretraining and designing learning rate schedules to improve efficiency.
INTRODUCTION
Large Language Models (LLMs) can achieve strong performance if pretrained with appropriate configuration of hypeparrameters, such as model width, depth, number of training steps, and learning rate. However, tuning these hyperparameters at scale is extremely costly since one pretraining run can take weeks or even months.
To reduce the cost of hyperparameter tuning, various scaling laws have been proposed to predict pretraining loss or downstream performance by capturing empirical relationships between key hy- perparameters and model performance. A notable example is the Chinchilla scaling law (Hoffmann et al., 2022), which approximates the final pretraining loss as a simple function of the model size N and total training steps T' (or total training tokens), L(N, 7) Lo + A 7 + B - 9 By fitting parameters Lo, A, B, o, 3 from few training runs with varying N and T, one can use the formula infer the optimal choice of N and T given fixed compute budget C NT.
A key challenge that existing scaling laws have not addressed is how to set the Learning Rate (LR) optimally over time. LR is arguably the most critical hyperparameter in optimization, as it can significantly affect the training speed and stability. A large I.R can quickly reduce the training loss, but in the long term, may cause overshooting and oscillation along sharp directions on the loss
CCoreepondding authors.
Code Implementation: htttps://github.ccm
1uo/Mu 1 Powerlaa

---

Published as a conference paper at ICLR 2025
3,6
3.0
2.5
2.0
3.2
Leaiing Rate
1.5
1246 oueer JSM
1.0
WDD
3,0
Cosizee
WSD.
22000
23003
W (Tuned)
WSD (Tumed)
0.5
SOLD (Timed)
WSDLD (Tume))
Opt
Ope
0.0
10000
15000
20000
23000
5000
15000
20000
25000
Stop
Step
(a) Comparison of Optimized LR Schedules
(b) Loss Curves for Different LR Schedules
Figure 1: Optimizing the LR schedule induces a schedule (Opt) better than cosine and WSD schedules. We conduct evaluation experiments on 400M Llama-2 (Touvron al., 2023) model trained over 12B tokens. Zoom-in regions facilitate the readers who are interested in the local details. (a) Our optimized schedule comprises constant and decay stages post-warmup, aligning with WSD (Hu et al., 2024). (b) Loss curves demonstrate that our optimized schedule outperforms cosine schedules and two major variants of WSD with tuned hyperparameters (WSD with exponential decay and WSDLD with linear decay).
landscape. In contrast, a small LR ensures a more stable training process but also slows down the convergence. Practitioners often balance these trade-offss by starting training with large LR and then gradually reducing it over time, following a Learning Rate schedule (LR schedule) (Bengio, 2012). These LR schedules sometimes include a warmup phase at the beginning, where the LR linearly increases from zero to large value over few thousand steps, and only after this warmup phase does the L.R start to decay. The most commonly used L,R schedule in L.M pretraining the cosine schedule (Loshchilov & Hutter, 2017), which decays the LR following cosine curve. Other schedules include the cyclic (Smith, 2017), Noam (Vaswani et al., 2017), and Warmup- "aarmup-Stable-Dae Stable-Decay (WSD) schedules (Hu et al., 2024), but there is no consensus on the optimal choice.
Existing scaling laws sidestep the complexity of LR schedules by fitting parameters on fixed fam- ily of LR schedules. For instance, Hoffmann et al. (2022) fitted the parameters in the Chinchilla scaling law for training runs that have gone through the entire cosine LR schedule. As a result, does not generalize well to other LR schedules, or even to the same schedule with early stopping. Moreover, existing scaling laws lack a term to account for LR schedules, limiting their ability to provide practical guidance on setting the LRs. This issue can become even more pronounced when scaling up training to trillions of tokens (Dubey et al., 2024; Liu et al., 2024), where the extreme cost of training makes it impractical to experiment with multiple LR schedules.
In this paper, we aim to quantify how LR schedules influence the evolution of training loss in LLM pretraining through empirical analysis. More specifically, we study the following problem, which we call the schedule aware loss curve prediction problem: Can we use a simple formula to accurately predict the training loss curve C(t) (1 S T') given a LR schedule E := {m ....... Mrr for T steps of training? To align with standard practices in L.M pretraining and to enable more precise analysis tailored to this setting, we impose the following reasonable restrictions on the problem. First, we take fresh samples from a data stream at each training step, so there is generalization gap between the training and test loss. Second, we focus on LR schedules that decay the LR over Finally, most LR schedules used in practice start with
time, i.e., ? 1/2 1Js warmup phase before the LR decays, we make a minor modification to the problem and include a fixed warmup phase before the decay phase the peak LR Mmax of the warmup phase have been carefully picked, potentially through a series of only interested in understanding how different LR decay schedules
interested in. We assume that the shape and
short training runs, and after warmup affect the training loss curve. For convenience, we shift the time index so that = 1 corresponds to the first step after the warmup phase.
In contrast to most existing scaling laws that rely on only two or three hyperparameters (Kaplan al., 2020; Hoffmann al., 2022; Muennighoff al., 2023; Goyal et al., 2024), solving the above problem poses unique challenges, as it requires predicting the loss curve based on the entire LR schedule, which inherently nigh-dimensional. This complexity necessitates more sophisticated approach to understand and quantify the relationship between the LR schedule and the loss curve.
Our Contribution: Multi-Power Law. In this paper, we propose the following empirical law (1) for schedule-aware loss curve prediction:
((t) = Lo + A (1) + -ww - LD((),, where S((1) =
T
(1)
2

---

Published as a conference paper at ICLR 2025
LR Lass Prediction
Training Set
Cosine
Conss
Tast Set
WSDLD
3.00
3.25
4.00
2.90
2.50
3.20
3.75
2.85
2.00
3.50
25M
3.15
2.80
1.50
3.25
Learning Rate (x10-+)
1.00
3.10
3.00
2.75
Q00M
0.50
3.05
2.75
2,70
8000
24000
Loss
8000
12000
8000
16000
24000
20000
24000
3.00
3.25
4.00
3.40
2.50
3.75
3.35
2.00
3.50
25M
3.15
1.50
3.25
100M
1.00
3.00
0.50
3.05
2:75
3.25
8000
8000
12000
16000
24000
20000
Step
Step
Law (MPL) with parameters fitted on cosine, constant, and two-stage schedules can accurately predict the loss curves of unseen schedules, including WSDLD, WSD, and two-stage schedules with a different LR in the second stage. See Table for evaluation metrics.
Figure 2: The
Here, Sw denotes the sum of learning rates used in the warmup phase. The first two terms La + A (S)(1) - S can be viewed as extension of the Chinchilla scaling law by replacing the number of steps T with the cumulative sum of learning rates up to step t, while neglecting the dependence the model size. While this alone provides crude approximation of the loss curve by linearizing the contribution of the LR at each step (see Section 3.1 for further discussion), it does not account for the specific shape of the LR decay. The additional term Lt(tt) serves as correction term, which captures the effect of LR decay in further reducing the loss:
LD((): = B Σ (1k-1- ma) 7
(x) = Σ Mr
(2)
More specificallly, LD(t) is linear with a cumulative sum of the LR reductions 7]- - Th over time, scaled by nonlinear factor (t)))) This factor gradually saturates to a constant as the training progresses, which follows a power law in a scaled sum of learning rates V)2 (t).
We call this law of L(t) the Multi-Power Scaling Law (MPL) as it consists of multiple power-law forms. Lo, A, B, C, ar B,y are the parameters of the law and can be fitted by running very few pretraining experiments with different LR schedules. Our main contributions are as follows:
1. We propose the Multi-Power Law (1) for schedule-aware loss curve prediction, and empirically validate that after fitting the parameters of the law at most 3 pretraining runs, it can predict the loss curve for unseen LR schedules with remarkable accuracy (see Figure 2). Unlike the Chinchilla scaling law, which relies solely the final loss of each training run to fit its param- eters, our approach utilizes the entire loss curve of each training run to fit the parameters, thus significantly reducing the number of training runs and compute resources needed for accurate predictions (Figure 5). Extensive experiments are presented for various model architectures, sizes, and training horizons (Section 4).
2. Our Multi-Power Law is accurate enough to be used to search for better LR schedules. We show that by minimizing the predicted final loss according to the law, we can obtain an optimized LR schedule that outperforms the standard cosine schedule. Interestingly, the optimized schedule has a similar shape as the recently proposed WSD schedule (Hu et al., 2024), but its shape is optimized so well that it outperforms WSD with grid-searched hyperparameters (Section 5).
3. We use novel "bottom-up" approach to empirically derive the Multi-Power Law. Starting from wtagstage schedules, we conduct series of ablation studies on LR schedules with increasing complexity, which has helped us to gain strong insights into the empirical relationship between the LR schedule and the loss curve (Section 3).
4. We provide theoretical analysis for quadratic loss functions and demonstrate that the Multi- Power Law emerges when the Hessian and noise covariance matrices exhibit certain types of power-law structures (Appendix B).
3

---

Published as a conference paper at ICLR 2025
2 PRELIMINARY
Learning Rate Schedule. A learning rate (LR) schedule is sequence E {m ......., m that specifies the LR at each step of the training process. For language model pretraining, the cosine LR schedule (Loshchilov & Hutter, 2017) the most popular schedule, which can be expressed as xt Th 2 Mmmax + - Mmax Cos( T Here, Mmax is the peak LR and a is usually set to 0.1. The 1=at
Harmuu--Saale--Deca (WSD) schedule (Hu et al., 2024) is a recently proposed L. schedule. This schedule first goes through a warmup phase, then maintains at stable LR Tmnax with Tatale steps, and finally decays in the form of f(t Tt'abee Imax for stable I S Tooaa. Here f(x) E (0, 1) can be chosen linear or exponential decay functions. We visualize these two LR schedules in Figure ((a).
Warmup Phase. Many LR schedules, such as WSD, include a warmup phase in which the LR gradually increases from 0 to the peak LR 7mmax over a few thousand steps. We denote the number of warmup steps as W. By default, the LR increases linearly, so the total LR sum during warmup is given by Sw 2 Mmax W. Our analysis focuses on the training process after the warmup, where the LR is decaying in almost all LR schedules. We count training steps starting from the end of warmup and set f schedule, and the LR at the last warmup step roo = 7max is the peak LR of the entire schedule.
the first step after warmup. Accordingly, (n,-... "r) represents the wost-warmup
Power Law of Data Scaling Prior studies (Hoffmann et al., 2022; Kaplan et al., 2020) demonstrate that, for fixed model size, the final loss follows a power law of the data size or, equivalently, the total training step number in a constant-batch-see
setting. This relationship is expressed as: L(T) Re L(T) == Lo + A T
where Lu, A, a are parameters to fit. This law is typically fitted over the final losses of a set of training curves generated from specific LR schedule family, such as a cosine schedule with given peak LR (7max), ending LR (00)max) and warmup steps (W). However, applying (3) directly to intermediate steps (t T) introduces bias, as the LR schedule up to t bears insufficient decay compared to the full schedule over T, resulting in different loss trajectories. This discrepancy is confirmed in Figure 5(b). We refer to (3) the Chinchilla Data Scaling Law (abbreviated as CDSL) throughout the paper since it is simplified from the Chinchilla scaling law (Hoffmann et al., 2022) to highlight the data dimension.
EMPIRICAL DERIVATION OF THE MULTI-POWER LAW
In this section, we present the empirical derivation of the Multi-Power Law (MPL) for schedule- aware loss curve prediction. Our key insights are summarized as follows:
1. If two training runs share the same sum of learning rates, Et= The then their final losses tend to be similar, though a non-negligible discrepancy remains (Section 3.
2. In particular, for a training run with a given LR schedule, the final loss L(T) is similar to that of another training run using a constant learning rate schedule with the same total LR sum. This motivates us to decompose L(T) into two components: (1) the final loss of the corresponding constant LR run; and (2) residual term that captures the effect of LR decay, defined as the difference between the final loss of the target run and the constant LR run. (Section 3.1)
3. Empirically, observe that training runs with constant learning rates exhibit Chinchilla-like pawer-law behavior in the loss curve and can thus be well approximated by a simple power law. (Section 3.2.
4. To approximate the residual term, instead of analyzing it directly, we imagine a sequence of training runs with schedules that gradually transition from a constant LR the target schedule, all while maintaining the same total LR sum. Using novel "botttomupp approach, we derive an approximation formula for the loss difference introduced by each incremental change in the LR schedule, first by analyzing simple two-stage schedules and then extending the results to more complex schedules. (Sections 3.2.2 and 3.3)
Finally, we sum up all the approximation terms above, leading to our MPL. Below, we elaborate on our approach in detail.
4

---

Published as a conference paper at ICLR 2025
3.54
3.0
Equal LR Sum
1.0
Learning Rate (x10-4)
2.5
9/
3.52
25
Ar
2.0
Leariin Rate (x10
2.0
3.50
1.5
Loss
1.5
99.1
3.48
1.0
LD(8720)
0.8
1.0
9/
3.46
6.5
0.5
6
3.44
Equal LR Sum
8000
8180
8360
8540
8720
8000
8180
8360
8540
8720
Step
Step
Figure 3: A multi-stage schedule (Appendix A.2) example illustrate the learning rate (L.R) sum matching (Section 3.1) and fine-grained loss reduction decomposition (Section 3.2.2). The steps with equal LR the final step T3 = 8720 are marked and linked with the dash-point line. Each stage spans 90 steps. T₃ = 8000, T22 = 8090, (()) Zrg (1b). ((2) Zrg (1a). See Appendix F.3 for experiment details. Left: The actual auxiliary processes. LR gap between adjacent points denotes the reduction Ary") = (1( 79) Right: Corresponding training curves for the multi-stage schedule and the auxiliary processes. The total loss reduction is (Tb) and be decomposed as the intermediate loss reduction sum. The loss gap between adjacent points denotes the stage-wise loss reduction LD( (1(9).
multi stage schedule and schedules
3.1 OUR APPROACH: LEARNING RATE SUM MATCHING
Auxiliary Training Process. As introduced above, we construct a series of auxiliary training runs with LR schedules gradually changing from constant LR schedule to the target schedule (Th...., }. Our construction is detailed as follows. We define the k-th auxiliary process shares the first steps of learning rates, {M,,,.... 7). with the actual training process with LR schedule E, and continues with the constant LR м) afterwards. The corresponding loss curve for the k- th auxiliary process is denoted as (t). In particular, the 0-th auxiliary process shares only the warmup phase with the actual training process and uses constant LR no 7mma after warmup. We especially call it the constant process and use Looust (t) to represent its loss curve. The T- auxiliary process coincides with the actual training with the target LR schedule, so Lr (t) L(t).
Learning Rate Sum Matching Decomposition The Multi-Power Law (MPL) approximates the loss curve C(t) of the actual training process through the following decomposition. We define Z(t) as the equivalent step in a constant LR process that shares the same cumulative LR sum as the actual s(()) and S(t) Σt 1hr represents the sum of wost-warmup
process up to step i, where (t) = Tr
LRs. The loss at step / is then decomposed as:
C(t) = Lconst (Z()) - (Ccoust (Z(t)) - C(t)),
(4)
::I((t)
where Lconst (t)) interpolates the loss for non-integer Z(1) the constant LR process. We first approximate C(t) using the training loss Looust (Z()) at step Zt),, and then write the residual term LD(t) representing the approximation error. We call Lt(t) the Loss reDuction term, as it quantifies the loss reduction due LR decay. We will approximate these two terms by parts in Section 3.2, with Lconse( (t))) detailed in Section 3.2.1 and LD(t) in Section 3.2.2.
Motivation: Continuous Approximations of the Training Dynamics. The rationale behind this approach is that two training runs with the same LR sum should result in similar training losses, thus making natural decompose the loss curve into major term corresponding the loss of a run with the same LR sum and a small residual term. To see this, we use SGD as an example. If the learning rates ....... TIT are small, then SGD can be seen as first-order approximation of its continuous counterpart, gradient flow, under mild conditions (Li et al., 2017; Cheng et al., 2020; Elkabetz & Cohen, 2021). Here gradient flow describes a continuous-time process in which the dØ(r)
parameters (T) evolve according to the differential equation dr = - -V)(Or)) where V V(()) the gradient at 0, and denotes the continuous time. In this approximation, the 1-th step of SGD
5

---

Published as a conference paper at ICLR 2025
corresponds to evolving ((7) over small time interval of length 7/-- When the learning rates sufficiently small, the parameters after t steps of SGD close to (T) at time 7 Σi=1 T]-- This connection naturally motivates us to compare the losses of two training runs with the same LR sum. While we use SGD for illustration, other optimization methods such as Adam can be similarly approximated by their continuous counterparts (Ma et al., 2022).
3.2
APPROXIMATION BY PARTS
3.2.1 CONSTANT PROCESS Loss APPROXIMATION
Motivated by the continuous approximation of the training dynamics, we hypothesize that losses of constant LR processes with identical LR sums are closely aligned. This insight inspires us to function of Stt) + Sw, where S(t) + Sw represents the cumulative
represent Leonst (Z))) LR sum up to step i, including the warmup phase part Sw. Analogous to (3), we propose that Coonst (t)) follows power law over the LR sum:
Ecuns (Z (Z( Lo + A (S(t) + Sw)"
where is parameter counterpart of Å. We perform extensive empirical validation and ablation studies across different model sizes, training horizons, and learning rates to confirm the robustness of (5), detailed in Appendix F.1 and illustrated Figure 11.
3.2.2 LOsS REDUCTION APPROXIMATION
Now turn the loss reduction term LD(t). We start by proposing simple yet effective linear approximation as a warmup, then we further break down the term with a finer-grained LR sum matching approach.
Warmup: Crude Linear Approximation. We first generate training loss curves across var- ious LR schedule types, including cosine and WSD schedules, alongside the loss curves of their corresponding constant processes, Then we can compute the loss reduction LD(t) for different LR schedules and analyze their dependency. As demonstrated in Figure 10, (t( is approximately pro- portional to the LR reduction, Anp = Yoo Th across different schedules. This leads to the following approximation:
(6) where B is constant. This finding highlights strong correlation between the loss gap and the LR gap at equivalent LR sum points on the loss landscape. However, while the linear approximation offers insights into the shape of (D),)), deviations from the actual loss reduction remain. Notably, when the LR decreases abruptly (e.g., in step-wise schedules), predicts an instant loss drop at the stage switch, whereas the true loss decline remains smoother during the training process. Sec Appendix C for further discussion.
LD(()= Re B(1o- 1)),
-rainGd LR Sum Matching Decomposition. In practice, the loss reduction term LD(() can have a more complex dependency on the LR schedule. To provide a more accurate approxima- tion than the linear approximation above, employ LR sum matching between adjacent auxiliary processes and decompose the loss reduction LD(()) into a telescoping sum of intermediate loss re- ductions between adjacent auxiliary processes.
More specifically, consider the step the actual training process. Similar to Z(t), we define t4 = (t) as the equal-LR-sum step in the k-tth auxiliary process, which is given by
t = (t) := k 1 Sx(t).
T
where (z) =t= Mr- Then, for the k-th and (k + 1)-th processes, we define the intermediate loss reduction as:
(8) Intuitively, this term compares the loss at step (e+1 in the 1k- process with the loss at the equal- LR-sum step in process that stops decaying the LR after the first k steps, i.e., the k-th process. We then decompose the loss reduction term as telescoping sum of intermediate loss reductions:
LD == Cx Cx+1(tu+1).
LD(t) = Leomu (Z(t)) - L(t) = (z) (z))(t)) Σ LD&(te+2..
(9)

---

Published as a conference paper at ICLR 2025
3.0
3.0
3.5
Launing Rale (x10
3,5
Lauiing Mato 0-10
2.0
3.4
3.4
1.0
Loming
3,3
Laming
Stollppooc
Malli poice
3.2
Ome-goner
0.00
0.5
40H00
50000
60000
40000
GUINO
Stop
Stop
(a) Cyclic Schedule
(b) Random-Polyline Schedule
Figure The examples of long-horizon non con-oonooonic schedules. The one-power line represents the constant process prediction.
(a) The cyclic schedule with 72000 steps, where cach half-cycle spans 8000 steps, and the first decay begins after 16000 steps. (b) The random-polylin schedule, consisting of piccewise linear interpolation between randomly selected intermediate learning rates the range of 3 10 to 10-4, with LR milestones occurring at intervals 8000 steps.
By leveraging this fine- grained decomposition, good estimation of (tx((t++) can lead to more accurate approximation of LD(t). Where the context is clear, we simplify notation by omitting subscripts and denoting intermediate loss reduction LD&
3.3 BOTTOM- UP DERIVATION: TWO-STAGE, MULTI- SALLEE AND GENERAL SCHEDULES
The challenges in approximating the intermediate loss reduction (t) twofold. First, for commonly used schedules, the learning rate (LR) reduction at intermediate steps is often too small to induce measurable loss reduction. Second, LD& may depend intricately on all previous learning rates ( 7 which we refer to as the LR prefix in this section.
To address these issues, derive the form of Lt&(t) using "bottom-up" approach regarding schedule structures. First, we propose its form through schedules comprising two constant LR stages, leveraging significant LR reductions. Next, we examine its dependency on the LR prefix using schedules of multiple stages. Our main finding is that LDx(t) depends weakly on the LR prefix and can be approximated by the following form:
(tx(t) (tx(t)= = B(T Mh+1) (1- (I-(On)
+1) +1)
(10)
with LR-prefix independent constants B, C. 7 and 3. Due to space constraints, we refer readers to Appendices A.1 and A.2 for detailed derivations of (10).
For general LR schedules, we extrapolate this findings and propose to approximate the total loss reduction term as: 1-1
(-1 B(T= mx+1) (1 11-(CiiTT(t4 k)++ 11)+)
(
(t) Σ LDD
k=0
By the definition of tx+1 (7), have tu+1 - -k k= Sk+1ll) Therefore, conclude
Me+1
LD(t) ~ LD((t)
Σ B -17) Th) "Cn,"" (t)(t) - +))
(11)
A=
where we also change the subscript indices from ++ 1 to k. Combining the above ansatz for the loss reduction term with the power-law ansatz for the auxiliary loss in (5) leads Multi-Power Law:
((t) z To A (1) Sw) Σ B(n-1
(1 (C (t) 1)
(12)
k=1
See Appendix for the ablation studies on different components of the Multi-Power Law.
EMPIRICAL VALIDATION OF THE MULTI-POWER LAW
The Multi-Power Law (MPL) comes from speculations based experiments with special types LR schedules. Now we present extensive experiments to validate the law for common LR schedules used in practice. Our experiments demonstrate that MPL. requires only two three LR

---

Published as a conference paper at ICLR 2025
Table 1: Evaluation metrics for the Momentum Law and Multi-Power Law on predicting the loss curves of 25M, 100M, and 400M models with unseen schedules. Rº. MAE. RMSE. PredE, and WorstE are the coefficient of determination, Mean Absolute Error, Root Mean Square Error, Prediction Error, and Worst-case Error, respectively.
Model Size
Method
R 1
MAE 4
RMSE +
PredE +
WorstE
Momentum Law
0.9904
0.0047
0,0060
0.0014
0.0047
25M
Multi-Power Law (Ours)
0.9975
0.0039
0.0046
0.0012
0.0040
Momentum Law
0.9959
0.0068
0.0095
0.0022
0,0094
100M
Multi- Law (Ours)
0.9982
0.0038
0.0051
0.0013
0.0058
Momentum Law
0,9962
0.0071
0.0094
0.0025
0,0100
400M
Multi- Law (Ours) 0.9971
0.0053
0.0070
0.0019
0.0070
schedules and their corresponding loss curves in the training set to fit the law. The fitted MPL can then predict loss curves for test schedules with different shapes and extended horizons.
4.1 RESULTS
Generalization to Unseen LR Schedules. MPL accurately predicts loss curves for LR schedules outside the training set. As illustrated in Figure 2 and Table 1, despite the absence of WSD schedules in the training set and the variety of decay functions, MPL successfully predicts their loss curves with high accuracy. Furthermore, MPL generalizes to two-stage schedules with different y/n values from the training set, effectively extrapolating curves for both continuous and discontinuous cases.
Generalization to Longer Horizons. MPL. demonstrates the ability to extrapolate loss curves for horizons exceeding three times the training set length. In our runs, the training set contains approximately 22000 post-warmup steps, while the test set includes curves with up to 70000 post- warmup steps. These results validate MPL's capability to generalize to longer horizons. Notably, the data-to-model ratio for a 25M -erammeeeterr model trained over 72000 steps (36B tokens) is comparable to Llama2 pretraining (70B model, 2T tokens), consistent with trends favoring higher data volumes for fixed model sizes (Dubey et al., 2024).
Generalization to Non-monotonic Schedules. MPL. extends effectively to complex non- monotonic schedules, although derived for monotonic decay schedules. We test the fitted MPL over challenging cases such as cyclic schedules and the random- aando-ppoylinee schedule, where LR values are randomly selected at every 8000 steps and connected by a polyline. These experiments, conducted on a 25M -5M--mraeeter model over 72000 steps, also represent a demanding long- morghorizon scenario. As shown in Figure 4, MPL accurately predicts these long- morgooorizon non-monotonic schedules.
4.2 COMPARISON WITH BASELINES
Comparison with Chinchilla Law. While Chinchilla-stlle data scaling laws, which we abbreviate as CDSLs, are widely adopted (Muennighoff et al., 2023; Hoffmann et al., 2022), MPL offers several distinct advantages: (1) MPL incorporates LR dependency, unlike CDSLs, and (2) MPL predicts the entire loss curve, whereas CDSLs are limited to estimating only the final loss. These advantages enable MPL to achieve higher sample efficiency than CDSLS. Notably, we demonstrate that single constant and cosine schedule curve suffices to fit MPL with strong generalization. As illustrated in Figure 5a)), MPL reduces final loss prediction to less than 1/3 that of CDSLs while requiring about 1/5 compute budget. Furthermore, MPL excels in fitting the open-source 7B OLMo (Groeneveld al., 2024), as shown in Figure 5(b). Additional details of the comparison with Chinchilla Law are provided in Appendix G.2.
Comparison with Momentum Law. The MPL outperforms the recently proposed Momentum Law(MTL) (Tissue et al., 2024) in both accuracy and applicability to discontinuous learning rate schedules. While MTL incorporates LR annealing effects by modeling loss reduction through the momentum of LR decay, it indicates an exponential loss reduction for two-stage LR schedules, in- consistent with our observations (see Appendix A. Across the diverse schedules in the test set, MPL consistently outperforms MTL in both average and worst-case prediction accuracy, sum- marized in Table Additionally, for WSD schedules with linear LR decay, MPL more accurately captures the loss reduction trend during the decay stage, as highlighted in Figure 14(b), compared to MTL. Further details on MTL and its relationship to MPL can be found in Appendix C, with fitting specifes provided in Appendix G.2.
8

---

Published as conference paper at ICLR 2025
3.2
2.575
X
Loss Curves( C)
Loss Curvest M)
2.5
Loss
power
3.0
2.530
Loss Ends(C)
2.4
Chinchilla
2.5
Learning Rate (110
Lass Ends(M) Pred(C)
Loss
12.50
12.75 13.00
Pred(M)
Loss Curve(Test)
1.5
Target: Loss
1.0
2.6
0
12
Step 6x10%)
Step (x10%)
(a) Fitting Sample Efficiency Comparison
(b) Whole Curve Fitting Comparison
Figure 5: (a) Target loss predictions 128000-step for cosine schedule using MPL and CDSL fitting, with 400M model. CDSL fitting requires six cosine losses (Loss Curve(C)) from 14,960 steps 72000 steps but relies solely on their final losses (Loss Ends(C)). In contrast, MPL leverages the entire 24000-step constant and cosine loss curves (Loss Curves(M)). Final loss predictions are denoted as Pred(C) for CDSL and Pred(M) for MPL respectively. (b) Comparison of MPL and CDSL fittings the whole loss curve of the open-source 7B OLMo model, trained with a linear schedule.
Table 2: Downstream performance comparison for the cosine and our optimized schedules. Percent- age changes (f L) indicate relative improvements regressions compared the cosine schedule.
Schedule
LAMBADA
HellaSwag
PIQA
ARC-E
C
RTE
Cosine
46.54
37.12
65.13
43.56
48.44
52.71
48.71
37.74
65.07
44.09
50.30
53.79
Optimized
(t 2.17%)
(1 0.62%)
(L 0.06%)
(1 0.53%)
(f 1.86%)
(t 1.08%)
THE MULTI-POWER LAW INDUCES BETTER LR SCHEDULES
Due to the high cost of each pretraining run and the curse of dimensionality for LR schedules, is generally impractical to tune the LR for every training step. To address this, propose leveraging the Multi-Power Law (MPL) to predict the final loss as surrogate function to optimize the entire LR schedule, achieving a lower final loss and outperforming the cosine schedule and WSD variants.
5.1 METHOD
The Multi-Power Law (MPL) provides an accurate loss estimation, enabling its final loss prediction to serve as surrogate for evaluating schedules. We represent the learning rate (LR) schedule T-dimensional vector E = parameters. Our goal is to find the optimal LR schedule E- arg min L(E) Using MPL, parameterize the predicted final loss as (e(E) with parameters θ = {L,, A, B, C, 3. esti- mated outlined in Section 4. We approximate E" by optimizing the surrogate loss Le subject to monotonicity constraints:
with the final loss denoted as L(E) under given hyper-
min Le(E) S.L. 0 S nh-1 A1 ViStST.
(13)
This optimization induces an "optimal" schedule under the MPL approximation. In practice, the peak LR To = 3 x 10-4 We view E as a high-dimensional vector and optimize using the Adam optimizer. Further details are provided in Appendix H. Results for 400M model are shown in Figure 1, with additional experiments for 25M and 100M models in Figure 18.
5.2 RESULTS
Optimized LR Schedule Exhibits Stable-Decay Pattern. Our optimized LR schedule follows Staale---bbabay (WSD) pattern, comprising two main post-warmup phases: a stable phase with constant peak LR, and decay phase ending with lower LR, illustrated in Figures and 18. By contrast, the momentum law (Tissue et al., 2024) theoretically yields a collapsed learning rate schedule, which will prove in Appendix I.
Optimized LR Schedule Outperforms Cosine Schedules. Across comparison experiments of different model sizes and training steps, optimized schedules consistently outperform the co-

---

Published as a conference paper at ICLR 2025
Cosine
2.0
4.0
Const loss
Const pred
3.2
Opt(Ours) WSDSC(OOus)
LR
Cosine loss
1.5
Leaming Rate (x10
3.5
25M
Cosine pred
3,0
Loss
2,425
100M
2.8
2,400
1.0
3.0
400M
2.375
2.6
0.5
1B
2.5
2.4
0.0
0
2
Step (x10*)
Step (x10*)
(a) Long- MorizHorrzzon Prediction of MPL
(b) Loss Curves Comparison for IB Models
Figure 6: (a) Long-hooe horizoo loss predictions using MPL for cosine and constant schedules, with model sizes ranging from 25M to IB (top to bottom). (b) Loss curve comparison for 1B models across the optimized sched- ule (Opt), cosine schedule (Cosine), and simplified optimized schedule (WSDSC, see Section 5.2), featuring WSD schedule with sqrt-cube decay.
sine schedules, achieving a margin exceeding 0.02. Notably, no WSD-like schedule is present in the training set, highlighting MPL's extrapolation capability, Figure 19 extends this comparison to longer training horizons and Figure 6(b) validates the superiority for 1B model. we further validate the effectiveness of our optimized schedules by evaluating the downstream task performance. As shown in Table 2, our optimized schedule leads to overall improvements in downstream tasks against the cosine schedules. showing practical gains from loss improvements. Ablation details for longer horizons and larger models are in Appendix H.
Optimized LR Schedule Outperforms Tuned WSD Variants. Our optimized schedules lead to smaller final loss than the WSD and WSDLD schedules proposed in Hu et al. (2024). For a 400M model, we find that the decay step of a 24000-step optimized schedule (Figure 1) is close to the optimally tuned step (~6000) for these WSD schedules, determined via grid search over {3000, 4000, 5000, 6000, 70003. However, even when the decay ratios of WSD and WSDLD schedules are optimally tuned, between our optimized schedule and these two WSD schedules as follows. The optimized schedule decays to below 1 1/20 of the peak LR, even approaching to zero, while WSD schedules decay linearly or exponentially to 1/10 of the peak LR. However, simply adjusting the ending LR to near-zero (Appendix H) does not close the gap. Another key difference is the decay function: we find through symbolic regression that in the decay phase, the optimized schedule roughly follows power decay function rather than a linear or exponential decay: π) = 7maa (1 -) 1.5, where is the step number in the decay phase, normalized to [0, 1]. Motivated by this, we propose WSD variant with squt-cube decay (WSDSC), which decays the LR exactly as 1/+ = T/max - (1-r)1.5 1.5 WSDSC is effective across various model sizes and architectures and outperforms the WSD schedule, as shown in Figures 6(b) and (a), though it still falls behind our optimized schedule. See Appendix H for details.
optimized schedule still outperforms them. Further, analyze key differences
6 CONCLUSIONS
This paper proposes the Multi-Power Law (MPL) to capture the relationship between loss and LR schedule. The fitted MPL accurately predicts the entire loss curve while requiring much fewer training runs compared to existing scaling laws. Furthermore, our MPL is accurate enough to be used for optimizing schedules, and we extensively validate the superiority of our optimized schedules over commonly used ones. However, we do observe slight deviations between our predictions and actual training curves, especially for long-horizon and high peak LR cases like in Figures and 16. likely due to several simplifications in our derivation: (I) the coefficient 3 remains constant across different LR scales; (2) the intermediate loss reduction does not depend on the LR prefix; (3) variations in LR during the warm-up phase are ignored.
In future work, we aim to (1) further explore the theoretical foundation of our MPL to uncover its underlying mechanisms; (2) investigate empirical laws for schedule-aware loss curve prediction with varying peak LRs and other hyperparameters; and (3) refine our MPL to further enhance its prediction accuracy and generalizability.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We would like to thank Kaiyue Wen, Huanqi Cao, and all anonymous reviewers, for their insightful comments and feedback. We also thank Hongzhi Zang for improving figure readability. This work is supported by the National Natural Science Foundation of China under Grant Number U20B2044.
REFERENCES
Steven Adriaensen, Herilalaina Rakotoarison, Samuel Müller, and Frank Hutter. Efficient bayesian learning curve extrapolation using prior-data fitted networks. In A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 19858- 19886. Curran Associates, Inc., 2023.
Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. In S. Koyejo, S. Mohamed, Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems. volume 35, pp. 22300-22312.
Curran Associates, Inc., 2022.
Alexander Atanasov, Jacob A Zavatone-Vett and Cengiz Pehlevan, Scaling and renormalization in high-dimensional regression. arXiv preprint arXiv:2405.00592 2024.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):e231187882) 2024. doi: 10.1003/pnas.231 2311188 231187112..
Yoshua Bengio. Practical Recommendations for Gradient- Based Training of Deep Architectures, pp. 437-478. Springer Berlin Heidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-642-35288 doi:
26.
Shane Bergsma, Nolan Simran Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to zero: Why linearly decaying the learning rate zero works best for LLMs. In The Thirteenth International Conference on Learning Representations, 2025.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, et al. Deepseek LLM: Scaling open-source language models with longtermism. arXiv preprint S.K1::2001.029544 2024.
Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. PIQA: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):7432-7439 Apr. 2020. doi: 10. 1609/aaaaii v34105.623333 6.339.
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1034-1034. PMLR, 13-18 Jul 2020.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. dynamical model of neural scaling laws. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings Machine Learning Research, pp. 4345-4382. PMLR, 21-27 2024.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning improve neural scaling laws. In The Thirteenth International Conference on Learning Representations, 2025.
David Brandfonbrener, Nikhil Anand, Nikhil Vyas, Eran Malach, and Sham M. Kakade. Loss-to- loss prediction: Scaling laws for datasets. Transactions on Machine Learning Research, 2025.
ISSN 2835-8856. Featured Certification.
Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. arXiv preprint arviv:2100. 14891, 2022.
11

---

Published as a conference paper at ICLR 2025
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and align- ment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1):29 2021.
Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and Langevin pro- cesses. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1810-1819. PMLR, 13-18 Jul 2020.
Daeyoung Choi, Hyunghun Cho, and Wonjong Rhee. On the difficulty of dnn hyperparameter opti- mization using learning curve prediction. In TENCON 2018 2018 IEEE Region 10 Conference, pp. 0651-0656, 2018. doi: 10.1 1011/EENCON.2011 8655007.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 00011103.00000 000.055577 2018.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs the edge of stability. In International Conference on Learning Representations, 2021.
Jeremy Cohen, Alex Damian, Ameet Talwalkar, J Zico Kolter, and Jason D. Lee. Understanding optimization in deep learning with central flows. The Thirteenth International Conference on Learning Representations, 2025.
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborova. Generalization error rates in kernel regression: The crossover from the noiseless noisy regime. In Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gra- dient descent at the edge of stability. In The Eleventh International Conference on Learning Representations, 2023.
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. Optimal linear decay learning rate schedules and further refinements. arXiv preprint arXiv:2310.07833 2023.
Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, Pa- quet, J. Tomezak, and C. Zhang (eds.). Advances in Neural Information Processing Systems, volume 37, pp. 9974-10007. Curran Associates, Inc., 2024.
Tobias Domhan, Jost Tobias Springenberg, and Frank Hutter. Speeding up automatic hyperparam- eter optimization of deep neural networks by extrapolation of learning curves. In Proceedings of the 24th International Conference on Artificial Intelligence, IJCAI' 15, pp. 3460-3468. AAAI Press, 2015. ISBN 9781577357384.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, et al. The llama 3 herd models, arXiv preprint arXiv:2407.21783 2024.
Omer Elkabetz and Nadav Cohen. Continuous vs. discrete optimization deep neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 2021.
Curran Associates, Inc.,
Elias Frantar, Carlos Riquelme, Neil Houlsby, Dan Alistarh, and Utku Evei. Scaling laws for sparsely-comnected foundation models. arXiv preprint arXiv:2309.088:5 2023.
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Bimix: Bivariate data mixing law for language model pretraining. arXiv preprint arXiv:2405. 14908, 2024.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski. Aapo Kyrola, An- drew Tulloch, et al. Accurate large minibatch SGD: Training imagenet in hour. arXiv preprint 1rXiv:1706.0000 02677, 2017.
12

---

Published as a conference paper at ICLR 2025
Sachin Goyal, Pratyush Maini, Zachary C. Lipton, Aditi Raghunathan, and J. Zico Kolter. Scal- ing laws for data filtering- data curation cannot compute agnostic. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 22702-22711, June 2024.
Dirk Groeneveld, Iz Beltagy, Evan Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Jha, et al. OLMo: Accelcrating the science of language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), pp- 15789- 11899, Bangkok, Thailand, aug 2024. Association for Computational Linguistics. doi:
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. In A. Glober- son, L. Mackey. D. Belgrave, Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 76232-76264. Curran Associates, Inc., 2024.
Tatsunori Hashimoto. Model performance scaling with multiple data sources. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4107 4116. PMLR, 18-24 Jul 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701, 2020.
Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer. arXiv preprint
2021.
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, et al. Deep learning scaling predictable, empirically. arXiv preprint arXiv: 171..0099,, 2017.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer. generalize better: closing the general- ization gap in large batch training of neural networks. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Shengding Hu, Yuge Tu, Xu Han, Ganqu Cui, Chaoqun He, Weilin Zhao Xiang Long, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. In First Conference on Language Modeling, 2024.
Peter Huber. Robust estimation of location parameter. Breakthroughs in statisties: Methodology and distribution, pp. 492-518. Springer, 1992.
Mareus Hutter. Learning curve theory. arXiv preprint arXiv:2102.04074 2021.
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats Leon Richter, Quentin Gregory Anthony, Eugene Belilovsky, Timothée Lesort, et al. Simple and scalable strategies to continually pre-train large language models. Transactions Machine Learning Research, 2024. ISSN 2835-8856.
13

---

Published as a conference paper at ICLR 2025
Ayush Jain, Andrea Montanari, and Eren Sasoglu. Scaling laws for learning with real and surrogate data. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 110246-110289. Curran Associates, Inc., 2024.
Arlind Kadra, Maciej Janowski, Martin Wistuba, and Josif Grabocka. Scaling laws for hyperpa- rameter optimization. In A. Oh, I. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 47527-47553 Cur- ran Associates, Inc., 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, et al. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
Aaron Klein, Stefan Falkner, Jost Tobias Springenberg, and Frank Hutter. Learning curve prediction with bayesian neural networks. In International Conference on Learning Representations, 2017.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradi- ent algorithms. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 Proceedings Machine Learning Research, pp- 2101-2110. 10. PMLR, 06-11 Aug 2017.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In International Conference on Learning Representations, 2020.
Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter L. Bartlett, and Jason D. Lee. Scaling laws in linear regression: Compute, parameters, and data. A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances Neural Information Processing Systems, volume 37, pp. 60556-60606. Curran Associates, Inc., 2024.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437 2024.
Xiaoran Liu, Hang Yan, Shuo Zhang, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05209, 2023.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference Learning Representations, 2017.
Jan Ludziejewski, Jakub Krajewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciehiera, et al. Scaling laws for fine-grained mixture of experts. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. $3270-33288 PMLR, 21-27 Jul 2024.
Bochen Lyu, Di Wang, and Zhanxing Zhu. A solvable attention for neural scaling laws, In The Thirteenth International Conference on Learning Representations, 2025.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of nor- malization layers: Sharpness reduction. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 34788 Curran Associates, Inc., 2022.
Chao Ma, Lei Wu, and Weinan. A qualitative study of the dynamic behavior for adaptive gradient algorithms. In Mathematical and Scientific Machine Learning, pp. 671-692 PMI.,, 2022.
14

---

Published as a conference paper at ICLR 2025
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp- 7697-7711. Curran Associates, Inc., 2022.
Alexander Maloney, Daniel A Roberts, and James Sully. A solvable model of neural scaling laws. arXiv preprint arXiv:2210.16859, 2022.
Ben Mann, N Ryder, M Subbiah, J Kaplan, P Dhariwal, A Neelakantan, P Shyam, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 1:3, 2020.
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. In A. Oh, Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 28699-28722 Curran Associates, Inc., 2023.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, al. Scaling data-constrained language models. In Oh, Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 50358-50376. Curran Associates, Inc., 2023.
Yoonsoo Nam, Nayara Fonseca, Seok Hyeong Lee, Chris Mingard, and Ard A. Louis. An ex- actly solvable model for emergence and scaling laws in the multitask sparse parity problem. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 39632-39693, Curran As- sociates, Inc., 2024.
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, et al. 2 olmo 2 furious. arXiv preprint arXiv:2501.0000 2024.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, et al. The LAMBADA dataset: Word prediction requiring broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1525- 1534, Berlin, Germany, aug 2016. Association for Computational Linguistics. doi: 18.13vv// P16- 1144.
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute- optimal neural scaling laws. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine I Learning Research, 21(140): 1-67, 2020.
Fabian Schaipp, Alexander Hägele. Adrien Taylor, Umut Simsekli, and Francis Bach. The surpris- ing agreement between convex optimization theory and learning-rate scheduling for large model training. arXiv preprint arXiv:2501.18965 2025.
Utkarsh Sharma and Jared Kaplan. Scaling laws from the data manifold dimension. Journal of Machine Learning Research, 23(9):1-34, 2022.
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adri- ana Meza Soria, al. Power scheduler: A batch size and token number agnostic learning rate scheduler. arXiv preprint arXivv2408. 13359, 2024.
Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pp. 464-472 IEEF, 2017.
Stefano Spigler, Mario Geiger, and Matthieu Wyart. Asymptotic learning curves of kernel methods: empirical data versus student paradigm. Journal of Statistical Mechanics: Theory and Experiment, 2020{12):224011, 2020.
15

---

Published as a conference paper at ICLR 2025
Kai Sun, Dian Yu, Dong Yu, and Claire Cardie. Investigating prior knowledge for challenging
Chinese machine reading comprehension. Transactions of the Association for Computational Linguistics, 8:141-155, 2020. doi: 10.1 00305.
Yunfei Teng, Jing Wang, and Anna Choromanska. Autodrop: Training deep learning models with automatic learning rate drop. arXiv preprint arXiv:2111.1533 2021.
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing. arXiv preprint arXiv:2408. 10029, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint Stt.:2307.092285 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, et al. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, et al. Superglue: A stickier benchmark for general-purpose language understanding sys- tems. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 23549-23588 PMLR, 17-23 Jul 2022.
Kaiyue Wen, Zhiyuan Li, Jason S. Wang, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-deen -eacee learning rates: A river valley loss landscape view. In The Thirteenth International Conference Learning Representations, 2025.
Xingyu Xie, Kuangyu Ding, Shuicheng Yan, Kim-Chuan Toh, and Tianwen Wei. Optimization hyper-parameter laws for large language models. arXiv preprint arXiv:2409.04777, 2024.
Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule. arXiv preprint 1rXiv:1909.000 199909979922.. 2019.
Jaaheng Ye, Peiju Liu, Tianxiang Sun, Yunhua Zhou, Jun Zhan, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint
2rXiv:2403. 16.522 2024.
Kaichao You, Mingsheng Long, Jianmin Wang, and Michael I Jordan. How does learning rate decay help modern neural networks? arXiv preprint arXiv:1908.01878. 2019.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, jul 2019, Association for Computational Linguistics. doi: 10.18653/ 11/211--472.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12104-12113, June 2022.
16

---

Published as a conference paper at ICLR 2025
LRA LKB
Arra A
M Mark
Arre B
LacA
Abrk-A
0.08
Loss
Mark
2-523
0.05
Emi-1.2eed5
Lanming Rare (UR
3.475
0.04
Ponii
LOTT, ex
LOIT,
lisponentia
3425
Model
7003
5030
5002
10000
11000
12000
3.400
10000
2000 Sky t
0.00
1021
Los Rodacion
300
3000
4003
Supp
Sty
(a) LR Step &
(h) Loss Step f
(c) Loss Reduction Step
Figure 7: Loss reduction (L.D) of two-stage schedule exhibits a power law. Example setting: te 11000,
œя 3000, T/a = 9 10 77A 3 10 TA = 8000. (a) and B have the equal LR sums: 2 = 900,
LA = 8900. (b) Loss reduction at B: LD(TA + zR) = (A(AA) - (a(Bs). (c) Fitting loss reduction LD(T' + Æe) with power form results in 0.13(1 - (1 + 0.21x) -0.15) Fitting with exponential form results in 0.0790(1 - e -0.0fz). ). The shape of loss reduction is closer to power form than exponential.
A BOTTOM-UP DERIVATION: TWO-STAGE, MULTI-STAGE (SECTION 3.3)
A.1 CASE 1: STAGE LEARNING RATE SCHEDULE
The two-stage schedule keeps learning rates at 7JA for T' steps, directly drops to yjn and continues for T3 steps. Then the LR reduction 7/ - 7/B could be significant enough to induce LDt, (t). which is also the loss reduction Lt(t) for step f on Stage 2. See Appendix F.2 for experiment details.
Loss Reduction Term Follows a Power Law. As shown in Figure 7, the number of steps æ == t- TA in Stage 2 increases, LD(TA + .)) monotonically rises from
around 0.09 and eventually saturates. This motivates to approximate LD(TTA + 22) in the form B- (1 U (nn=))) where B is parameter and (s) is a function that decreases from to 0 as 8 = T increases from 0 to infinity. The reason we choose T instead of as the argument of U will be clear in the general case.
But at what rate should (s) decrease? After trying different forms of (s) to fit LD(I'A + x), we find that the power-law form (s) = (C s + 1) for some C, 8 > 0 fits most properly as shown in Figure 7, which leads the following power-law form for the loss reduction term:
LD(TA + x)) = LD(TA + x) == B(1 (C 7/22 + 1)-
(14)
Appendix A.I shows that this power law aligns well with the actual loss reduction term LD(T)
In contrast, the exponential form (s)) = e- Bs (so LD(TA + x) s A(1 c Bma))) struggles to
match the slow and steadily increase of LDD(A x) when x: is large.
Parameter Pattern of Power Law. We further investigate how estimate the parameters B, Č, in the power law. Based on our preliminary experiments, we set 3 = 0.4, a constant that works well. Then we conduct experiments to understand how the best parameters B, Ò to fit LD(t) depend on MA, 1B, TA, where set default values 77 x 10-4 1/B x 10- T'A 8000 and change one variable at time. The details of ablation experiments can refer to Appendix F.2. The observations are summarized as follows.
(1) B Linear LR Reduction. As shown the first row of Figure B linearly decreases with fa and approximately increases linearly with 7A. especially when 7A is not too large. Moreover, the slope of B over 7ı and 7/8 are approximately opposite to each other. This motivates to hypothesize that B x AA -/B and reparameterize B as B = B( 7/s, where B is constant.
(2) Ć Follows a Power Law of r/B As shown in the second row of Figure 8, is very sensitive to Th but much less dependent on 7JA- We hypothesize that Č follows a power law Č' x 1h5 and reparameterize Ö as Ö Cmp²" where C >> O and O are constants.
(3) LR Reduction Term Depends Less on Tx We also find that В and Ö are less sensitive to Ta relatively stable as TA varies, as shown in the last column in Figure 8. This suggests that the loss reduction has weak dependency of loss reduction on LR prefix length.
17

---

Published as a conference paper at ICLR 2025
0.15
vs
B vs Ta
Rt
0.10
400.52 1/s
0.05
1 Fit
6 403 270r + 0.126
Cv T.
č vs Rt
C vs TA
1000
5.707g
500
2.5
5.0 (x10-1)
7.5
(x10-1)
Ta (x10f)
Figure 8: The dependency patterns B, C over T)a, TB and Ta the two-stage cases. B is approximately proportional AA VJ,, and Ò manifests power-law pattern over fj8- The dependency of 7a over Ö and the impacts of TA on B, C' are unpredictable or negligible, which are approximately ignored in our discussion.
Approximation Form. Putting all the above observations together. we have the final approxima- tion form for the loss reduction term in the two-stage schedule:
LD(TA 2) LD(TA + B(Π 7B) (1- 72 + +))-))
(15)
A.2 CASE 2: MULTI-STAGE LEARNING RATE SCHEDULE
In the two-stage case, the LR prefix is constant at 7a. leaving uncertainty about whether the interme- diate loss reduction conforms the power form when the LR prefixes vary. To investigate this, we analyze the multi-stage step decay schedule. Consider an n-stage LR schedule E = {m..... ......... 1/1 ) where the i-th stage spans from step T, + to T:+1 and uses the LR 17 (0 S Ti < ... < Tn+1 = T,
with = (0) ( > ... n("), 1 S i < n.). An example is illustrated in Figure 3.
Stage- Loss Reduction. In the multi-stage schedule, given stage index S S n, the stage-wise loss reduction defined (D) LDr, (t) The LR reduction between stages, (n) = (fi-1) ()) is also measurable. Using this, we estimate the shape of (D for different and define
t Tr As shown in Figure 9a), T.D (T: + æ) approximately conforms similar power law as (14) for the two-stage case: () 2)
stages. Regard T; TA in the two-stage
LD (T; (T; ~ LD (8) (T++
B()) (1 (1- (1-
(((0)))())))))) 1/ x
(16)
where B(0) and (1)) are constants dependent on the LR prefix {1)...... nr. for stage
Intermediate Loss Reduction Weakly Depends on the LR Prefix Shape. For stage i, the LR prefix is { }. which varies length and scale across stages. To evaluate the effect of the LR prefix on the intermediate loss reduction form, we examine its impact on B(s) and (t()) Inter- estingly, as shown in Figure 9(b), we observe that B3()) B( 17(8) and (()) As (*)
which align closely with the two-stage results. Here, B, C. and are constants largely independent of the stage index. This suggests that intermediate loss reductions are relatively insensitive to the LR prefix compared the LR reductions (n)() and the stage LR r Moreover, this weak depen- dence on the LR prefix may extend to general schedules, indicating a broader applicability of the power-law form for intermediate loss reduction.
B How MIGHT THE MULTI-POWER LAW ARISE?
In this section, we present preliminary theoretical analysis understand how the Multi-Power Law might arise. More specifically, we consider simple setting where SGD optimizes a quadratic loss function with noisy gradients, and show that the Multi-Power Law naturally emerges when the Hessian and noise covariance matrices exhibit certain types of power-law structures. While this analysis does not fully capture the complexity deep learning, we believe it offers insight into how the Multi-Power Law relates to underlying spectral properties in the optimization landscape.
Note that T LDD (f) for each T, + 1 S {(t) Tet1, as these auxiliary processes for specific stage coincide.
18

---

Published as a conference paper at ICLR 2025
0,030
503
Lowe Robietion (x30 3)
12.5
0.025
10.0
0.030
500
1-16-15
4-4-5
p-1e-55
500
403
2.5
8-347.757Un K-
300
C= 4,951(("")
0.0
0.005
20
40
50
Stnge Step
7 - 71th 0c10-
7 (x10-))
(a) Power Fitting of (L)
(b) LD Parameter Patterns
Figure 9: The intermediate loss reductions of multi-stage schedule (Figure 3) and their shape patterns. (a) The loss reduction (t) between the adjacent stages of the multi-stage schedules still follows the power form. (b) B x ff(h11 - ()) x (1(*))=² ), The paramcter patterns in the two-stage setting hold the stati-stage setting approximatcly. The shape of patterns is similar to the patterns in the two-stage experiments, shown in Figure 8.
B. SETUP
where A E RRd represents the trainable parameters, 0, is the ground truth and H E Raxd is the Hessian matrix. Linear regression special case of this formulation. More generally, any sufficiently smooth loss function can be locally approximated by such a quadratic form near minimizer.
We consider a quadratic loss function = = (0-
{m...... .r to optimize the loss function, where the 1-th iteration is given by 0, 0 7.9r. with I. being the stochastic gradient at step /. We assume that the stochastic gradient gi equals the true gradient (L(00-1) = HO,-1 plus Gaussian noise (0, Σ), where Σ E gdxd is the covariance matrix. That gt (HO,, E).
We use SGD with LR schedule
From spectra to scaling law. The scaling behavior of the loss during training typically deter- mined by the spectra of the Hessian matrix H and the noise covariance matrix (Canatar al., 2021; Spigler et al., 2020; Maloney et al., 2022; Cui et al., 2021; Brandfonbrener et al., 2025). By carefully analyzing the training dynamics, matrices can lead to a scaling behavior similar to our empirical Multi-Power Law. In the following, use A; denote the i-th eigenvalue of H and use Σn denote the i-th diagonal entry of Σ in the eigenbasis of H (i.e., v. Σv, where v, is the i-th eigenvector of H). We initialize the parameters at 0n and use A, to denote the i-th corrdinate of 0 0, in the eigenbasis of H (i.e., v. (6p 0.))- We consider scenario where Hessian, noise covariance, and initial point drawn from certain distributions before training, and make the following assumptions:
show that certain structures of Hessian and noise covariance
Assumption 1. For all 1 <i S d, the marginal distribution of (Ai, Eii, Ai) is fixed distribution p(x, E, A) with following properties:
À is supported on A) for some A > O, and p()) & X for some exponent v E [0, 1). That is, p(A) = (XE(O. x for some normalization constant 7 0.
Ep1E A] x P expp(rX) for some p < 1 and 0.
EAAA² A1 D X for some E [0, v) and D > 0
B.2 MAIN THEOREM
Consider an SGD training process with an arbitrary LR schedule E = (n...., .r)) Define
Sc(t) Ei= 7/ for S k < T. Fix any no 0. We define the following function
Z(t) as an estimate of the expected loss at step t: LD(t),
((t) Lo + A A.
(17)
LD(() B Σ (m--1 mm) ())((((()),, Ō(x) = 7(8, (2r +r)A)
(Crr + 1)-8, (18)
r))
19

---

Published as a conference paper at ICLR 2025
where the constants Lu. at, B, above are given by d
d d-T(a) D², a := t0:=2-100 2- B=
d - := 1----pp C == (19)
To = 4 7 (20 A = 20+12 Z
Here, [(2) := f0 e denotes the gamma function, and (s, == -o ts l 't denotes
the lower incomplete gamma function.
The theorem below establishes E(I) as precise estimate of the expected loss curve of SGD. See Appendix for the detailed proof.
Theorem 1. Under Assumption 1, for all t T, if Mmax == maxogier [n,} is sufficiently small and S (t) is sufficiently large, then we have the following estimate E[
[[[((+)] [(t] = (Tmmxx
min(a+1.9) + mmaax)
Here, Equation (17) is the same as Equation (1), and we can see that the exponent a is determined by the rate of eigenvalue decay of the Hessian and the rate of variance decay of the initial parameter. The coefficients Lo x To A x 132 depend on the learning rate, variance of gradient noise, and the initial distance to the optimal parameter.
Equation (18) gives the loss reduction term (D),,, which is similar to the loss reduction term in Equation (2). A change in the learning rate at step k induces a loss reduction ( GSSx (t)), where B x [E][] depends on the variance of gradient noise. Similar to (2). this loss reduction saturates as (1k=1 (1 (S when Sx becomes large, where 3 is determined by the rate of eigenvalue decay of the Hessian and the rate of variance decay of the gradient noise. But a slight misalignment between theory and practice here is that the form of (()) is not the same as "S, (t)) in (2). The main discrepancy here is that G(T Sg has explicit dependence on the current learning rate T/e while ((Sx(t) does not. We suspect that this is due to that in practice, changing the learning rate also change the local loss landscape, such as the well-known Edge of Stability (EoS) phenomenon (Cohen et al., 2021; Damian et al., 2023; Lyu et al., 2022; Cohen al., 2025), but the quadratic loss function we consider here is not complex enough to exhibit such a behavior. In addition, the function G(x) is in a simple power form, while G(x) also involves a lower incomplete gamma function 7(3 (2x + r)) and only approximately follows the power form. We argue that this a minor difference since these functions are approximately the same especially for large F, as we plot in Figure 21. Further, if A - +00, i.e., the maximum eigenvalue of the Hessian is very large, it easy to that G(x) - G(z) with the same C and B.
Extending the analysis of our Multi- Hulti-Power Law beyond quadratic cases is a key direction for fu- ture work. A deeper understanding of loss landscape properties in deep learning is crucial for this generalization. For example, a recent paper (Wen et al., 2025) conjectures that the loss landscape in LLM pretraining exhibits river valley landscape, which is similar to deep valley with river at its bottom. Based on this conjecture, they further explained the success of WSD schedules. For future work, it would be interesting to extend our analysis to this river valley landscape or other frameworks that better capture the complex structure of the loss function in practical deep learning scenarios.
C FORMULA COMPONENT ABLATION
To understand and evaluate the role of each component in our Multi-Power Law (MPL; See (1)), systematically simplify the MPL formula various levels and explore alternative formulations. Table 3 summarizes the fitting performance of these simplified versions and variants of the MPL. The fitting experiments are conducted on 25M models, using the same experimental setup described in Appendix E.
The necessity of the loss reduction term LD(t) can be assessed by fitting One-Power Law (OPL), a simplified MPL where (t)(t) = 0 or equivalently B 0:
No Loss Reduction.
(o) = Lo + A (S) +ww) " Si(1) = Σ Tr
(20)
T=1
20

---

Published as a conference paper at ICLR 2025
This formulation approximates the loss curve by matching the LR without correction term, as discussed in Section 3.1. The fitted results (first row of Table 3) exhibit significant degradation compared to the full MPL, demonstrating the critical role of (t)
Linear Approximation of Loss Reduction. Based on the observation in Section 3.2.2, the loss reduction term (t) (defined in Equation (2)) can be simplified by treating the sealing function ((r)) as a constant:
LD(L)
B(m 7h) B(10 7t))
(21)
Despite its simplicity, we observe a near-linear relationship between LD(t) and the LR reduction (%) m.), regardless of the LR schedule type, shown in Figure 10. This motivates the Linear Loss reDuction Law (LLDL):
LLDDL (t) = Lo + A- (ti + Sw) + B(% - -m).
(22)
As shown in Table 3, LLDL achieves significantly better accuracy than OPL., although it underper- forms the full MPL. However, this formulation is unsuitable for optimizing schedules, as its results collapse to trivial solution:
7/ % when &o I T and ris when k = T.
Loss Reduction Without 7/- Next, we simplify G(x) by setting 0, yielding the No-y Law:
NoO- Lo A (1) Sw) -M B (1k-1 - mk) G(S)(1))
(23)
Results (third of Table 3) show slight performance drop, confirming that y enhances fitting accuracy with minimal additional computational cost. Thus, we retain 7 in the final MPL.
Step-Based Approximation. An alternative is to replace G(m with a step-based formu- lation, G(t 1 + 1). This yields the Step Power Law (SPL):
LsPl = Lo + A (S) Sw) B
Σ (m--1-1) -16 mR) G(t - + 1).
(24)
A=1
While simpler, this approximation reduces prediction accuracy and contradicts empirical results, because it implies loss reduction continues increase when LR reaches zero.
Exponential Approximation. Substituting G(x)) with an exponential function G(z)) = C gives the Multi-Exponential Law (MEL):
LVEL. Lo + A (1) + Sw) B
Σ (m ma) G(S)(()).
(25)
k=1
Results (fifth row of Table 3) show a performance drop compared to the power-based MPL., consis- tent with observations in Appendix A.1 that Ütt m7s takes power form rather than an exponential form.
Relation to Momentum Law. of
The concurrently proposed Momen'Tum Law (MTL.) is in the form
(MIL(T) Lo + A (S1 Sw) B. S2, where S
Σ Ts. S2
ΣΣ M
A=1
where is hyper-parameter of MTL and 1. is indeed variant of MPL since
1 (
S2 (mh-1 ===
(Mk-1
7pk
Σ (
im k=1
ビ=
k=
Thus, MTL is a variant of MPL with an exponential step- approximation:
LMIL (t) = Lo + A (Si + Sw) + +B' Here, B' = B second row Table 3) even lags behind MEL, highlighting the limitations of step-based approxi- mations.
= B'. G(t - k + 1), G(z)=1-ed 1-c)=1-c-o0
ア C log A. MTL incorporates step-based decay and its performance (last
21

---

Published as a conference paper at ICLR 2025
3.0
Const
WSDLD
4.0
2.5
WSD-Cosine
Leamihg Rate D010-t
0.00
WSD
Cosice
3.8
2.0
1055
3.6-
0.04
0.02
3.4
loss
linear regression
(LR reduction, Loss reduction)
25000
0.5
1.5 LR Reduction (x10-*)
1.0
2.0
2.5
Figure 10: Lincar regression of loss reduction versus LR reduction across different schedules for 25M model at 4000 for WSD and its variants, among which WSD-Cosine specifically
over 24000 steps. Decay steps denotes the WSD schedule with cosine decay function. Left: Visualization learning rate schedules and their corresponding loss curves. Right: Scatter plot of loss reductions against LR reductions, accompanied by a linear regression fit (mean R² 0.9980), demonstrating strong linear relationship between the two variables.
Table 3: Summary of fitting results for simplified versions and variants of the MPL. Metrics include Ra. MAE, RMSE, PredE, and WorstE, where higher H² values and lower values of other metrics indicate better fitting performance. See Table for metric definitions.
Formula
Features
R2 1
MAE
RMSE
PredE!
WorstE
OPL.
LD(t) = (B = =0)
0.8309
0.0378
0.0412
0.0111
0.0241
LLDL
G(r)= =1
0.9797
0.0077
0.0101
0.0023
0.0108
No-7
7=0
0,9961
0.0046
0.0053
0,0014
0,0041
SPL
T I k
0.9921
0.0066
0.0075
0.0020
0.0069
MEL
G(x) 1
Cr 7
0.9934
0.0044
0,0057
0.0013
0.0047
MTL
G(x) = 1 e
C2
0.9904
0.0047
0.0060
0.0014
0.0047
MPL
G( = 1
(Cx + +1 1)-8
0.9975
0.0039
0.0046
0.0012
0.0040
(Ours)
x 1J =
D RELATED WORK
Scaling Laws and Loss Curve Prediction. Scaling laws reveal empirical power-law relationships between training losses and various factors such model size, dataset size, and computational re- sources. Hestness et al. (2017) initially observed that generalization errors in deep learning decrease predictably with larger model and dataset scales, following power-law trend. Subsequently, Ka- plan et al. (2020) investigated these scaling laws in the context of training Transformers (Vaswani al., 2017), demonstrating a persistent power-law decay in loss and contributing to the rise of LLMs (Mann et al., 2020; Bi et al., 2024; Touvron et al., 2023; Dubey et al., 2024). Later stud- ies refined these scaling laws by improving fitting and evaluation pipelines, refining parameter size metrics, ensuring consistent hyperparameter configurations across scales, and extending their appli- cability to broader training scenarios (Hoffmann et al., 2022; Henighan al., 2020; Bi al., 2024; Caballero et al., 2022; A Alabdulmohsin et al., 2022).
Scaling laws have been used to predict the model performance in various settings, hyperparameter optimization (Kadra et al., 2023), multi-epoch training (Muennighoff et al., 2023), training sparsely- connected models (Frantar et al., 2023; Ludziejewski et al., 2024), length extrapolation (Liu et al., 2023), transfer learning (Hermandez et al., 202 1), and data mixing (Hashimoto, 2021; Ge et al., 2024; Ye et al., 2024; Jain et al., 2024). However, most existing work neglects the impact of the learning rate, making their predictions unreliable for assessing model performance throughout training (Hoff- mann et al., 2022). As a result, deriving scaling laws for final losses under specific schedules (c.g., the cosine schedule (Hoffmann al., 2022; Muennighoff al., 2023)) requires more than 10 full
22

---

Published as a conference paper at ICLR 2025
training runs. Another line of works extrapolate loss curves with Bayesian methods (Klein al., 2017: Domhan et al., 2015; Choi et al., 2018; Adriaensen et al., 2023) for hyperparameter opti- mization, but they also need large number of training runs. In comparison, our work introduces multi-power law that extends the existing power-law form of specific LR schedules by incorporating terms to capture L.R schedule effects. Unlike prior work, our approach is L--deependent and can predict full loss curves across various schedules using fewer than three training curves.
Concurrent to our work, Tissue al. (2024) proposed a momentum law that also incorporates LR into scaling laws. However, our work outperforms theirs by identifying a power-law decay rather than an exponential decay in loss reduction relative to LR reduction (detailed in Section 4.2). This results in more accurate formula capable of inducing optimized schedules, whereas the momentum law provably yields suboptimal collapsed schedules (detailed in Appendix 1). In another concurrent work, Xie et al. (2024) derived a LR -.Reeepedeent scaling law based on theoretical insights drawn from SDE-based continuous-tin
approximation of training dynamics. However, their scaling law relies on manually defined training phases and is limited to predicting the final loss value. In contrast, our proposed scaling law is arguably more general, as it can predict the entire loss curves for various LR schedules, even including discontinuous and non-monotonic ones. This also allows us to optimize the LR schedule by minimizing the predicted loss over all possible LR decay schedules, which may not be easily achievable using their method.
Theoretical Insights on Scaling Laws. Many works have explored theoretical explanations for the observed power-law behavior. Sharma & Kaplan (2022) attribute the power-law behavior to the intrinsic dimension of data manifold. Hutter (2021); Michaud et al. (2023) draw insights from a toy case where an infinite amount of distinct knowledge pieces need to be memorized, and the power law of the loss curve can arise when the frequency of knowledge exhibits a power-law distribution. Many other works analyze the scaling law of linear models (Spigler et al., 2020; Bordelon et al., 2020; Maloney et al., 2022; Bahri et al., 2024; Wei et al., 2022; Bordelon et al., 2024; Lin et al., 2024; Atanasov al., 2024; Paquette et al., 2024), assuming certain power-law properties in the input data or ground truth functions. A few others examine how power law behaviors arise in simple neural networks (Nam et al., 2024; Bordclon et al., 2025; Lyu et al., 2025). Similar to these works, we also provide a theoretical explanation for our multi-power law assuming certain power-law properties in the optimization landscape, but our analysis is accurate enough to capture the effects of learning rate schedules on the loss curve.
Optimal Learning Rate Schedule. Setting a proper schedule for the learning rate is crucial for training deep neural networks. He al. (2016) introduced the warmup strategy, which is now stan- dard in modern schedules. Early work by Smith (2017) proposed cyclical LR schedules, featuring periodic linear decay with warmup restarts, later extended to cosine decay with warmup restarts by Loshchilov & Hutter (2017). Some works explored adaptive approaches, such as Bayesian or reinforcement learning-based methods (Xu al., 2019; Teng et al., 2021), but they are compu- tationally expensive for LLMs. Li & Arora (2020) demonstrated that training with one schedule accompanied by weight decay is equivalent to training with the same network with an exponentially increasing LR schedule without weight decay. Goyal et al. (2017); Hoffer et al. (2017); Malladi et al. (2022) studied how to scale the LR when increasing the batch size. Li et al. (2019); You et al. (2019) showed that IR decay can benefit generalization by suppressing the memorization of noisy data early in training and learning complex patterns late in training.
In the context of large model training, Hu et al. (2024) introduced the Warmup-Stable-Dees (WSD) schedule, which starts with a warmup phase, continues a main stable phase, and ends with a rapid decay phase. This schedule has shown strong performance in LLM pretraining and efficient con- tinual training. Similar patterns have been adopted in other works (Zhai et al., 2022; Hägele et al., 2024). Ibrahim et al. (2024); Zhai et al. (2022); Raffel et al. (2020) adopt a reciprocal (inverse)-sqrt LR schedule in full process or as a component. Wen et al. (2025) analyze the benefits of WSD schedules by conjecturing that the loss landscape exhibits a river valley structure, and propose alter- natives for WSD in continual training. Inspired by these works, recent open-source models ad advocate schedules with a slow-decay or stable phase followed by rapid decay (Liu et al., 2024; OLMo al., 2024).
Our resulting scaling law can induce optimized schedules share a similar pattern with the WSD schedule, even though we do not fit the law on WSD schedules. Concurrent to our work, other schedules have claimed optimality under specific conditions. Defazio et al. (2023) proposed linear
23

---

Published as a conference paper at ICLR 2025
decay schedules as optimal based worst-case analysis. Shen et al. (2024) introduced power schedule for continual training, which outperforms WSD schedules. Schaipp et al. (2025) drew parallels between convex optimization and LR scheduling for LLMs, using simulation results to guide continual training strategies and peak LR selection. Bergsma et al. (2025) argued for linear- to-zero LR schedule optimal, ablating peak L., data size, and model size. Defazio et al. (2024) proposed a schedule-free approach using weight averaging techniques, but it underperforms WSD schedules (Hägele et al., 2024). Existing methods often optimize schedules under specific constraints, such as ending LR (Bergsma et al., 2025), decay ratio (Schaipp et al.. 2025), continual training (Shen et al., 2024), or worst-case convergence bounds (Defazio al., 2023). In contrast, our approach integrates LR schedules into scaling laws, which enables gradient-basee -ased optimization over possible schedules.
E EXPERIMENT SETTING
Unless otherwise specified, the model training in the Section 3, 4 and 5 follows the following set- tings.
Codename
Embedding Dimension
#Heads
#Layers
#Non-embeddings
#Params
25M
640
5
5
25
89
100M
1024
8
8
101
205
400M
1536
12
12
340
493
IB
2048
32
16
822
1026
GPT-2
768
12
12
85
162
Table 4: The model series run in all the experiments. Hoffmann et al. (2022) utilizes the number of non-embedding parameters (#Non-embeddings) to count model sizes, while Kaplan et al. (2020) counts the total number of parameters (#Params). The unit of the Parameter is M in this table.
Default Hyperparameter
Value
Sequence Batch Size
128
Sequence Length
4096
Optimizer Type
AdamW
$1
0.9
0.95
1 x 10
Weight Decay
0.1
Gradient Clipping
1.0
Peak Learning Rate
3 10-4
Final Learning Rate
3 10-5
Warmup Steps
2160
Table 5: Hyperparameters related to model training.
Our validation contains two steps: (1) fitting schedule-curve pairs from the training set and (2) pre- dicting the loss curves for schedules in the test set. The training set contains only single 24, St0-s constant and cosine schedule pair, alongside a 16,000-step two-stage schedule of 7/B 0.31 The test set has one 72,000-step constant and cosine schedule, 24,000-step unseen WSD and WSDLD schedules, and 16,000-step two-stage schedules with 7/ = 0.17A and 7/B = 0.67A The details are provided in Table 6. We train Llama2 (Touvron al., 2023) models of 25M, 100M, and 400M, and collect their loss curves, with model parameter details in Table 4. Training employs the Adam optimizer, with a weight decay 0.1, gradient clipping at 1.0, 3 = 0.90, and 32 = 0.95, consistent with the Llama2 training setup. Default hyperparameters include peak LR of 3 10-4, warmup period of 2160 steps, and a batch size of 0.5M. Additional hyperparameters are detailed in Table 5. In ablation studies, simplify the experiment fit short constant and cosine schedules and predict the loss for long-horize morgoorizon cosine schedule. The MPL fitting employs Huber loss (Huber, 1992) as the objection function, aligning with prior work (Hoffmann al., 2022; Muennighoff al., 2023),
24

---

where -= (B, 3). and we set 8 = 1 10-2, For each experiment, we use the Adam optimizer with a learning rate at 10-4 and total steps of 20000. Here we do not conform to the L-BFGS algorithm like Hoffmann et al. (2022) due to its sensitivity to the initialization. In fitting, the parameters initialized based on the loss reduction curve shape: B corresponds to the estimation step Published as a conference paper at ICLR 2025
peed
4.0
Powd:4oMM
1osse400M)
Powd/ 100M)
Lewmiing Rate (110r1)
LoseG25MI
3.7
3.2
3.6
3.0
3.5
3.4
2000
4000
6000
8000
10000
12000
14000
10000 20000
30000
40000
50000
60000
70000
Step
Step
(a) Peak LR Ablation for Constant Schedules
(b) MMddel Size Ablation for Long- Mori-ooizzo Training
Figure 11: Loss curves for constant LR schedules. pred denotes the fitted law prediction and loss represents the ground-truth loss curve. See Appendix F.1 for details.
and uses the Adam optimizer for optimization. Unless otherwise specified, we report validation loss. For fitting approaches and additional details see Appendix G.
F DISCUSSIONS OF MULTI-POWER LAW DERIVATION (SECTION 3)
F.1 CONSTANT PROCESS Loss APPROXIMATION (SECTION 3.1)
The constant process employs constant LR schedule with the same warmup phase and peak LR as the actual process schedule. We validate (5), the LR sum power law of the loss curves for constant schedules, through two series of experiments. First, we conduct ablation over the peak LR, ranging from 3.0 x 10- to 3.6 x 10-3 over 14,400 steps, achieving an MSE of 1.55 x 10- and R² of 0.9976 (Figure (a)). Second, we validate the power form over long-horizon curves (72000 steps) for model sizes of 25M, 100M, and 400M, with a peak LR of 3.0 x 10-4, yielding an average MSE of 8.04 10-5 and R² of 0.9947 (Figure 11(b)).
F.2 TWO-STAGE EXPERIMENTS (SECTION A.1)
In this section, we provide details on the investigation of the variation of coefficients in the power law for two-stage LR schedules.
Experiment Setting and Law Fitting. The experiment setting aligns with Appendix E. Default configuration uses 7A = 3 x 10-4, 1/B 3 10-5 TA 8000. In the ablation experiments, 7a ranges from x 10 5 to x 10-3, ranges from 4 x 10 5 to 2.9 x 10 and T'A ranges from 4000 28000. The second stage lengths spanning 1000 over 6000 steps. Validation loss is sampled every 2 steps due to the rapid loss changes after the stage switch. Following Hoffmann et al. (2022), we fit the law utilizing Huber loss as the objection function (Huber, 1992),
min Huber, LD (T' + x) log LD(TA +
(26)
of asymptotic values of loss reduction and Ō can be estimated according the slope (Equation (14)).
Fixed 3 Experiments for Parameter Patterns. We fit the power-law form in Equation (14) across ablation experiments to identify the loss curve shape and power-law parameter patterns. For the sake of further derivation, we fix the exponent, 33 I. L.R-independent parameter 0.4 based on the warmup
25

---

Published as a conference paper at ICLR 2025
experiments. Then re-fit the loss curves fixing 3 0.4 to confirm the validity of the power form. Figure 12 includes the re-fitted curves and ground truths for the ablation experiments over 7a and 7m, showing feasible error margins for further derivation despite fixed B. We further investigate the dependency of different parameters on the 7/A, J/B, and TA, with pair-wise relations presented in Figure and summarized in Appendix A.1.
0.00
0.12
Loss Rcductinn
0.08
Gx ba10 ?
0.06
0.06
2.02
0.02
Average Eman: 2.75e:06
reduction
reduction
1000
2000
3don
4000
5000
6000
Step
(a) A Ablation
(b) 1/B Ablation
Figure 12: Power-law fitting of loss reductions versus steps two-stage LR schedules.
F.3 MULTI-STAGE EXPERIMENTS (SECTION A.2)
We analyze intermediate loss reduction dependency on the LR prefix, through experiments multi-stage schedule and its auxiliary (intermediate) processes. As shown in Figure 3, our multi- stage schedule consists of 9 stages, with first stage of 8000 steps at 3 x 10 followed by eight 90-step stages. The validation interval is also set 2 steps. For adjacent stages i - and i (z S 90), We compute (.D (T, + æ²), defined in Appendix A.2 and fit it going through the same law fitting process as Equation (26). The fitting of loss reductions for different stages is presented in the left panel of Figure 9. Moreover, parameter trends, analogous to the two-stage findings, reveal B(O) changing with (66--1) (8) and (()) changing with (7()), shown in the right two sub-figures of Figure 9.
G DETAILS OF VALIDATION EXPERIMENTS (SECTION 4)
G.1 TRAINING SET AND TEST SET
Set
Schedule Type
Total Lengths
1B/1A
Constant
24000
Training
Cosine
24000
Two-stage
16000
0.3
WSD
24000
WSDLD
24000
Two-stage
16000
0.1
Test
Two-stage 16000
0.6
Constant
72000
Cosine 72000
Table 6: Summary of training and test sets.
Our validation frames the Multi-Power Law (MPL) fitting machine learning task, training schedule-lose curve pairs from the training set and predicting loss curves for the test set. The training set contains 24000-step constant and cosine schedule pair, and 16000-step two-stage schedule with 7h 0.37A- The test set includes a 72000-step constant and cosine schedule, a 24000-step unseen WSD and WSDLD schedule, and 16000-step two-stage schedules with 7/B 0.1 and
26

---

Published as a conference paper at ICLR 2025
30
33
Cosine
36
WSD
4.50
2.0
Loss
4.25
3,
26
Learaing Rate
1.5
8.75
0.5
3.25
52-
0.5
5000
10000
15000
20000
25000
10500
Supp
(a) Loss Curve Comparison for GPT-2
(b) Long Horizon Prediction for GPT-2
Figure 13: Loss curves of GPT-2 models with Multi-Power Law fitted over 24000-step constant and cosine schedule losses. (a) Comparison between the cosine, WSD, and WSDSC schedules (see Section 5.2); (b)
Prediction for a 72000-step cosine schedule loss curve.
Seed 45018
WIDLD 3.8
4.0
Seed 337
Seed 1660
Learning Rate (x10-1)
2,5
3,18
3.6
3,3000
Loss
3.8
53375
33100
3.4
3.6
1.5
2.85
3.2
22.0
26.5
23.0
23.5
24.0
Loss
3.4
0.5
One powee
3.0
Momentum
5
10
15
20
25
5000
10000
15000
20000
25000
Step(x
Step
(a) Random Seed Ablation
(b) Comparison with Momentum Law
Figure 14: (a) Experiments with a 25M model over 24000 steps across different seeds, showing a final loss standard deviation of 0.0007 and maximum gap of 0.0014. (b) Comparison between Multi-Power Powtr Law (MPL) and Momentum Law (MTL). In the decay stage, MPL achieves higher fitting accuracy and matches the curva- ture of the loss curve, whereas MTL fits the stable stage but predicts counterfactual concave curve during the decay stage.
7B = 0.67A The peak learning rate is 3 10 10-4, and the ending learning rate is x 10-5 for the co- sine, WSD, and WSDLD schedules. For all two-stage schedules, TA 8000. All schedules include warmup phase of 2, 16 steps. Detailed descriptions of the training and test sets are summarized in Table
G.2 FITTING THE LAW
Similar to the two-stage fitting, fit the parametric law using the Huber loss the objective (Hu- ber, 1992):
min Σ Huber, Lo (X,) log Cgr (X.))
(27)
where (x,) denotes the ground truth of validation loss, Le (X₂) is the predicted loss, is hyperparameter for the Huber loss, and θ denotes parameters to fit. The total fitting loss sums up the Huber loss over the validation steps. In practice, we compute the area under the linearly interpolated polyline of the learning rate at validation steps surrogate for the LR sum. This approach reduces the computational cost since requiring only step numbers, learning rates, and losses at validation steps.
[A, 3, Lo). and X. (m..... Mc We use the Adam optimizer fit the MPL due to its flexibility, with learning rate of 5 x 10 for the index parameters (a, B, and 7) and x 10-2 for the coefficient or constant parameters (A, B, C. and Lo). We also perform second optimization with a learning rate of 1 x 10 and 1 x 10 -0 initialized with parameters from the first optimization. Each optimization runs for 5 x 1 steps, selecting the lowest training loss result. Fitted parameters listed in Table 7.
Multi- Powter Law. For the Multi-Power Law (MPL),
27

---

Published as a conference paper at ICLR 2025
Table 7:
Parameter values for optimized schedules across different model sizes, rounded to two
decimal places.
Model Size
A
B
C
B
Y
Lo
400M
0.66
614.30
0.16
0.42
0.88
0.56
2.52
100M
0.59
521.40
0.24
0.46
0.60
0.65
2.79
25M
0.51
446.40
2.07
0.53
0.41
0.52
3.17
In the discussion of Appendix C, we also fit simplified MPL or MPL variants this manner, except for the momentum law (Appendix G.2). In Figure 15, present the fitting and prediction results for a subset of experiments, with a zoom-in window highlighting predictions near the end of training. In long-horizon experiments, the zoomed-in view reveals slight discrepancies between the MPL predictions and the actual training curves, targeted for future refinement.
Momentum Law. For the momentum law (MTL; Appendix C), θ {A, B, a, Lo), with as a tunable hyperparameter. The input X, for MTL is the same as MPL's input. Fol- lowing Tissue et al. (2024), we L-BFGS to minimize Equation (27), grid-searching À (0.95, 0.99, 0.995, 0.999, 0.9995} and selecting the best fit based on training accuracy. Predictions are evaluated across the test set (Table 6). with comparisons to MPL in Table and Figure 14. In Figure 14, we compare them specifically over the WSDLD schedule. In the decay stage, MPL not only achieves higher fitting accuracy but also aligns with the curvature of the loss curve. In contrast, MTL fits the stable stage well but predicts a counterfactual concave curve during the decay stage.
Chinchilla Data Scaling Law. The Chinchilla Data Scaling Law (CDSL) is similar the one- power law mentioned in Appendix C, but uses the power of steps instead of the LR sum, with θ = {A,a ax La), and X, = (final steps only) for Equation (27). The fitting of CDSL. follows Hoffmann et al. (2022) and uses the L-BFGS algorithm to minimize the Huber loss. With regard to sample efficiency (Figure 5(a)), CDSL uses cosine curves at 14960, 20080, 27760, 40560, 53360, and 72000 steps, requiring 4.8 times more compute than MPL (two 24000-step curves), with prediction errors of 0.007 (MPL) versus 0.024 (CDSL). MPL achieves less than one-third the prediction error of CDSL. In Figure 5(b), CDSL fits all intermediate steps, ignoring the effect of LR schedule and loss reductions for the comparison with MPL.
Discussion the Optimization Method. We also explored the of the L-BFGS algorithm for fitting MPL. but found highly sensitive to parameter initialization. For instance, under certain initializations, the fitted parameters may include a high value and near-zero C. Note that (1 + Cx)-9= = 1 exp(-8 log + Cx)) ss 1 exp(- in this case, making MPL resemble multi-exponnntial form. In practice. this issue can be mitigated by constraining parameters such as B and to the interval (0, 1). Additionally, we can initialize C, B, and y through grid search to obtain more feasible results. However, using the Adam optimizer is not without limitations, as lacks theoretical convergence guarantees. Future work will focus on enhancing the fitting process to achieve greater robustness and stability.
28

---

Published as a conference paper at ICLR 2025
3.50
3.05
3,50
Learstng Mnta 10-10-m
27
Learutg 2oth 0:10
3.6
2108 20
5.
2.95
20
3.40
Laameng
15
2.90
Leareng
3.35
Loos
Loss
Molt-poeer
2.85
Molt-pouer
3.30
One-posee
7500
10000
12500
15000
17500 20000
22500
6000
2000
10000
12000
14000
16000
Stop
Stop
3.0
3.0
Learning Bater 00 0CE0-1
2.5
Leatning Rate 0220 4
2.9
3.450
3.425
Loss
Leomung Bate
2.7
Learning Rate Lina
1.0
3.400-
Mulli
Oue
Que Que-poons
Mult-pooese
0.5
10000
11000 12000 13000
14000
10000
30000
30000
60000
70000
Stap
Stap
3300
3.25
3.25
3.20
Lourning 12C0
310
3.15
-20
300
10.00 500
M107 c2n Balo Pid wide
3.10
-1.5
3.05
3.05
2.85
Lasming
3.00
3.00
2.0
Mattiggoue
Oac-ponerr
2.95
paw
7500
19000
12000
12000 17500 Step
20000
22500
10000
20000
30000
40000
50000 Stop
G0000
70000
Figure 15: Fitting and Prediction Details. Subfigures depict loss curve fitting (training set) and prediction (test set) across various configurations, labeled (X, Y) for row X, column Y. The columns the accom- panying table indicate: F/P Fitting (F) or Prediction (P), Model Size, Step Length, and Learning Rate Schedule. Subfigure details follow:
(X,Y)
F/P
Model Size
Step Length
LR Schedule
()
25M
24000
Cosine
(1,2)
400M
16000
2-stage (3 x 10-4 + + 10
12
25M
16000
2-stage (3 10-4 1.8 10 10-4)
(2,2)
P
400M
72000
Cosine
(3.1)
PP
100M
24000
WSD
(3.2)
100M
72000
Constant
29

---

Published as a conference paper at ICLR 2025
3.2
Learzing 183
Credi
3.
Cosine
TwoSSuge
1.21833 x35
33
5003
10000
15000
20300
25000
10203
20009
Step
Stop
317
Two Stuge
Two Stage
WSD13
50.0
10000
15000
20000
25000
5002
10203
15000
20200
25000
Step
Snp
0.04
Joss
prod
Lewering Hasia (130
3.8-
Lace
3.6
20000
40000
50000
10000
2000
30300
4,0000
56000
000.0
70000
3.8
Lawming
3.6
1.8xx
3.0
LANTSSnn (210
3.0
Exx 10-1
Two Stuge
3.6
3.4
WSD
WSDLD
20203
Lornnn (220-1
3.6
10000
20000
70000
40500
50000
60000
70000
10000
20000
30000
40000
50000
G0000
70000
Step
Stop
Figure 16: Ablation study on peak learning rates. Left: Learning rate schedules: Right: Corresponding loss curves. Layout: The first three rows show the results for peak LR of 4 x 10 while the last three rows are for the peak LR of 6 10 Within each of the three the first shows the fitting the training set, second row displays the prediction over unseen schedules the third row demonstrates extrapolation capability on a long horizon loss curve.
30

---

Published as a conference paper at ICLR 2025
3.0
5.0
Launing Satt
2.0
2.0
Roto
Learingg
Tha--Sape
15
Oa 3x:10
Ocent
1.0
Gremte
10000
10000
20000
25303
20000
60000
Jos
Jo66
4.0
peed
peed
3.9
4.0
3.8
3.8
3.7
3.5
35.
355
3.6
20000
03032
lcsa
1039
4.0
3.8
pood
3.8
prod
3.7
3.4
3.2
Stop
Stup
Figure 17: Ablation study on batch sizes, with R values of 0.9977 (hatch size 64) and 0.9973 (batch size
256). The subfigure layout is as follows. Rows: (1) Learning rate schedules, (2) Loss curves for batch size
64, (3) Loss curves for batch size 256. Columns: (1) Training set results, (2) Test set results (same horizon as training), (3) Test set results (extended horizon).
31

---

Published as a conference paper at ICLR 2025
Cusine
2.5
Oyt
Learaine Bato (xd0
333-
1.5
3.00
relo 10ime 2110 zoin 2llo
0.5
WSDLD
DosLoM
WED
5000
10000
15000
20000
25000
5000
10000
15000
20000
25000
Step
Step
3.0
Cusine
WSDLD
WSD
2.5
3.6
Opt
1005
Leemingg
0.5
WSD
3.0
Opt
5000
Step
Step
Figure 18: Comparison of our optimized LR schedules and their loss curves with cosine, WSD, and WSDLD schedules over 24000 steps. The decay step for WSD and WSDLD is set to 4000. Upper: 25M model; Lower: 100M model. Left: Learning rates over steps. Right: Losses over steps.
H DETAILS OF LR SCHEDULE OPTIMIZATION (SECTION 5)
Optimizing the Surrogate Objective. To enhance optimization stability. we redefine the learning
rate schedule F. {m ( MF) using A = (Ai, A2,..., Ar), where A, = `/- - 1/+ and roo denotes the initial peak LR. Thus, Tte = 7/0 -Et= Åg. establishing one-to-one mapping between
E' and A. We transform the objective Le(E) from (13) into (a(A) and alternatively optimize:
min A
(A
s.t.
A
> 0,
Vt == =,
Σ A S 170-
In practice, we enforce these constraints through clipping: after each optimization step, we restrict A, into [0, mu] and set A, when 1): S c, with e 10-10 to ensure numerical stability. Applied to the MPI. fitted from Appendix G. this reformulation empirically stabilizes optimization by aligning learning rate reductions with zero initialization. For optimization, we use the Adam optimizer with constant learning rate, grid searched from 2 x 10 to 1 x 10-9, over 50,000 to 200,000 for better convergence.
Optimized Schedule of Longer Horizons and Different Model Sizes. Beyond Figure I and Fig- ure 18, we validate the optimized schedules for extended horizons and different model sizes. For models ranging from 25M 400M, optimize LR schedules for 72000-step training based on the MPL over the training set. As shown Figure 19, the resulting schedules exhibit a WSD- like shape, consisting of stable phase and decay phase, outperforming cosine schedules across
32

---

Published as a conference paper at ICLR 2025
3.0
3.2
Learnny Bate (x10-1)
2.0
2.9
1.0
Cusine
0g
Qp11100M
Opt:4P0M)
2.7
Learning Rate
2.6-
10000 20000 30000 10000
50000
60000
70000
50000
55000
60000
65000
70000
Figure 19: Left: Optimized and cosine LR schedules over 72000 steps for models ranging from 25M to 400M. Right: Corresponding loss curves for optimized and cosine schedules.
WSDLD
4.0
Leinis Ratt (X300"
2.0
3.8
3.2-
3
034
1.0
WIILD
DM DW
WSDLD
WSD
WSDNZE1
5300
10000
15000
20000
25000
5000
10000
15000
20000
25000
Siep
Step
Figure 20: Comparison between the optimized schedules and WSD variants with near-zero ending LR. WSD (ZE) and WSDLD (ZE) denote WSD and WSDLD schedules with an ending learning rate of 3 10-7. Learning rate comparison. Right: Loss comparison.
Left:
sizes. For the IB model, derive 72000-step schedule based on the MPL fitted from 24000- step constant and cosine schedule curves, with results in Figure 6(b) confirming superiority over the cosine schedule. Additionally, for the 1B model, MPL-induced schedule against the cosine schedule across several tasks, including LAMBADA (Pa- perno et al., 2016), HellaSwag (Zellers et al., 2019). 9), PIQA (Bisk al., 2020), ARC-easy (Gu & Dao, 2023; Clark et al., 2018), C (Sun et al., 2020), and RTE (Wang et al., 2019). The MPL- induced schedule achieves an average score improvement of 1.03 compared the cosine schedule, as shown in Table 2. This highlights the effectiveness of the MPL-induced schedule in enhancing model performance across diverse downstream tasks.
evaluate the downstream performance of the
Zero-Ending Learning Rate Experiments. Optimized schedules consistently outperform WSD variants with near-zen zero learning rates (3 10 approximately 1/100 of the default setting). To higher ending LR (e.g.. 1/10 peak LR) degrades baseline performance, we compare the optimized schedules against WSD(LD) variants with near-zero ending learning rates and the original ones. As shown in Figure 20, the optimized schedule still outperforms these WSD variants. In addition, lower ending learning rates do not consistently improve the final loss (c.g., zero-ending WSD exceeds baseline loss), suggesting complex interaction between the ending learning rate and the decay function. This highlights the advantage of the optimized schedule reducing the need for extensive hyperparameter tuning.
WSD with Sqrt-Cube Decay (WSDSC). We derive the decay function for optimized schedules by analyzing the decay phase across Ilama2 models ranging from 25M 400M. We compute
33

---

Published as a conference paper at ICLR 2025
normalized steps and learning rates (LRs) within the decay phase for schedules of varying step counts and model sizes. After averaging the normalized LRs, we perform symbolic regression
-t 1.5
against the normalized steps and approximate the decay function as (table)
Tura Tants
Validation experiments on 1B Llama2 and GPT models confirm its efficacy: Figure 6(b) shows that WSDSC outperforms the cosine schedule for the 1B model, though it falls short of the MPL- optimized schedule. Figure 13(a) demonstrates WSDSC's superiority over both the standard WSI) and cosine schedules for GPT.
34

---

Published as a conference paper at ICLR 2025
1 OPTIMAL LEARNING RATE SCHEDULE FOR MOMENTUM LAW
In this section, we derive the optimal learning rate schedules for the Momentum Law (Tissue et al., 2024):
L(T) Lo + A S2, where S1 Σ Be and S2 = Eー ranges from 0.99 to 0.999, and Lo, A, C are parameters.
- 7):) X hyperparameter typically
Similar to Section 5, here we could also optimize this law to get a learning rate schedule achieving lowest final loss by solving
min 01:02
Le
71))
(A)
S.L. 0 The T-- L1 SIST,
where = (Lo, A, C, } represents the hyperparameters and parameters in L(T). For simplicity of derivation, we introduce 7/ in front of E' as the maximal LR. Compared with our Multi-Power Law (MPL), this optimization problem obviously convex, so we can characterize its minimizer more easily in math. In result, the Momentum Law yields optimal schedules that go through stable phase at peak LR and then directly drop to zero LR in no more than two steps. However, these kinds of schedules clearly far from the optimal schedules in practice. As comparison, in Section 5, MPL can induce a WSD-like schedule, which is empirically effective.
Next, formalize the above arguments mathematically.
Theorem 2. For any LR schedule F. (Mi..... yir that minimizes the optimization problem (A), there exists E {0, 1, 1,..., T} such that the following holds for all t. E (1,..., TT):
1. 1ft < k then vi = rJo:
2. Ift 2 k 2, then ni = 0.
For convenience, we first prove the following lemma.
Lemma 1. For a function f(x) = x M(1 X²) with M > 0 and 0 < X < 1, we have the following properties:
1. f(x) is strictly convex and has unique minimizer over E (0,
2. If fw) 2 0 for some y E [0, oo), then f(x) 2 fly) all (c).
Proof: First, it is easy to check
ʃ(0) = 0, de (z) - 1 MX log X, d&f
df
du (x) = MX (log 2)2 > 0.
Therefore, f(x) is strictly convex. Then we can discuss the property of f(x) over x E oo) by discussing (0).
(1) When (0) 2 0, di (x') > of (0) 2 0 for all x' E (0,z). Thus, f(x) is monotonically
da
de
increasing over (0, x) and f(x) > f(0) 0. So x is the unique minimizer over a E (0)
and f(x 2 fly) when 2 0.
(2) When =()) < O, lim # (x) = 1. Then there exists x" E (0, (0,00) oo) such that o (x*) = 0. de
Thus, f(x) monotonically decreases over (0, and monotonically increases over (x", , oo).
Hence x² is the unique minimizer over E oo). Moreover, f(x*) < f(0) 0 and
limg- f(x) 00, so there exists ă (a", 20), such that fxx over (0,5) and f(x)
over (ž, 2o). Clearly, (x) monotonically increases over x E [ oo). Therefore, if f(w) 2
for some y E [0, oo), then y 2 ă and f(x) 2 f(y for all x E (c).).
This completes the proof.
Next, we prove Theorem 2.
35

---

Published as a conference paper at ICLR 2025
Proof for Theorem 2. First, reparameterize 1/t 1i no -Ek=1 Ak. then the optimization
problem (A) becomes
min Ar
İ= ,,...,, ,,,
s.t.
4 200
LA vI
T
Σ
A
70:
==1
where (r(A1, A2..... Ar) is given by
T
a
T
A2 Ar) Lo A
Tru-
ΣΣ
A
-
Σ
A
[=1 km
1=1
Define the Lagrangian by
(A).A,))= L=(A1. Ar) Σ A + p
Σ 1=1
A
1=1
where A1,..., Xr and are the Lagrange multipliers associated with the constraints A, 2 0 and E A S fj. respectively. By Karush-Kuu Kuhn-Tucker (KKT) conditions, there exist Y1..... Ar 2
and μ 2 such that the following conditions hold:
Complementary Slackness: A, 0 for all t ,,..., and u ΣT= A. 10 =0
Stationary: AA,
aig (A1, Ar) - X + p 0 for all / = 1,..., T.
Here, we have
di: 0A.
= aAd a- -ar-a-a.-------- (C---11----0 +... +..--
= aAΦ aa-1 (T-++1) C 1- T-
1-X
Kf(T - ++ - 1),
where Thu
Σl_ Ei A.. K aAΦ a--1 > 0, and f(x) = x MM1 X²) with
M [1- 1jA >
Note that does not depend on t. We can rewrite the stationary condition as
X Kf[T - + 1) μ.
By Lemma 1, fxx) strictly convex and has a unique minimizer over x E [0, oc). Let x² be this unique minimizer. Let fiin mineg(1....7T)) {T - t 1)) be the minimum value of f(T 11,, and S be the set of indices that minimize f(T t + 1). Then [S| < 2 and S C {[T x² +1]. [T- x² 11].
Now discuss the following two cases by the value of K fmin + p.
Case 1. If Kfmin + l > 0, then > > for 1,..., T. By the complementary slackness condition, A, = 0 for all I,...,T T. This implies that E is constant schedule, 7ju n/ nÿ-
7/
Case 2. If Kfmmn + μ 0. then A 0 for all t S and X 0 for all t S. By the complementary slackness condition, the latter implies that A, = 0 for all t $ S. Then E* falls into one of the two categories:
1. If 5 = {s) for some 8, then 7/o = = = "15 = ..... 1/=-1 and 11, = 1/,11 = .... 17;
2. If S {s 1,8) for some 8, then f/p 1/i n
1/-2 and 11/ 1//11 n/
36

---

Published as a conference paper at ICLR 2025
ceweplemennttineee slackness condition. Moreover, 8 < implies T & S, and then we have 0 < Xr =
We claim that 7/ = if < T. If not, then we have μ due to the complementary
+ -
By Lemma 1, f()) 2 f(1) 0 for all 2 1, which implies that X = Kf - KJ(T- - t ++11 1) > 0 for all L = 1,...., 7, which contradicts the fact that X, = 0 for all & e S.
Putting all these together, conclude that E* must exhibit the pattern described in the theorem.
37

---

Published as a conference paper at ICLR 2025
J PROOF OF THEOREM
The proof of Theorem consists two main parts. First, we explicitly derive the formula for any quadratic loss function after steps, without making Assumption 1. Then, we take the expectation over Σ, H, and the initialization 0g to prove the theorem.
WLOG, assume that H diag(>1, Ad), and set 0, O. Otherwise, due to the rotational and translational invariance of SGD, we can always transform the coordinate system so that the eigenbasis of H is the standard basis, and the optimal solution is the origin. In this case, En is just the i-th diagonal entry of Σ in the standard basis.
J.1 GENERAL THEOREM FOR ALL QUADRATIC LOss FUNCTIONS
In the first part, we establish the following theorem that characterizes the expected loss. We use @(0, E') to denote the distribution of the T-th iteration Or of SGD with initialization 6g and LR schedule E := (m1 17......77'
Theorem 3. For Or ~ @(Gu. E) and any fixed > O, we have the following estimate of ([(())))
1
1 1 exp(-
M(O,, E) = 2
A (xp((22A,51) + no
2
1=1
1
1 - exp(-xx2,555
2
7h-
7pk
Σ
Ea
k=1
i=1
where S₁ := =TE The estimation error is bounded as
15
[[C(Or)] - M E)] 51/max
S exp(ep2
+ mmax
2
Endi
2
where Maax == maxoskeT T---
To prove the theorem, first introduce some notations and auxiliary expectations. We define
U) = (0) exxxxx((222 + E/E - 11-ex(-22X,5)
2
We decompose the expected loss (0,,,)))00,, into a telescoping sum of T + 1 auxiliary expectations Ag, .1....,A4
10) Ao Σ (AN Ak-1), k=1
(28)
A B0.-0(00.5 (0 [T Me, Sk41)l,
where we define ST+1 0, Ar -or- (0.. (0....E((C00 [[(OT)].
The above theorem needs the following two lemmas.
Lemma 2. If
[0, 1]. then
361 100 0]
(1 - (1=x)² -exppxxpxx (1 {1#?), 36 10, 0] (1 2x)) exp( exp(-2x)) 2xr) (1 E2x²)
s.t.
Proof. The above inequalities hold for x = 0. For the case of x E (0. 11. we must have
f = (1 ) cxp(2x)
1 1- (1 2x) cxp(2x)
x?2
x2
So it suffices to show that both 1- exp(2n) and
1- 2xr (2x) 2²
lie in [0, 10).
Noting that
zz
1 - 2x < (1 x²2 < exp( - 2x), we obtain the following lower bounds:
1 - (1- - 2x) exp(22 2xx) 1 (1 - )² exp(2x)
1 - -exp(
2x exp(2x) 0.
x²2
2²
z²
Also note that
1-(1- exp an increasing function of x. So we have
2x) exp(2x)
1
1) cxp(2)
< 10.
2
12
38

---

Published as a conference paper at ICLR 2025
Putting these two inequalities together, we have
10 2 (1 2x.) exp 2x))
1 1 x))
exp((22 expp(
2 0,
x²
x²
which completes the proof.
Lemma 3. Tmax S Amax
then for all [T].
k-1
1
Σ
T/e exp(-2X; S exp(-
2A;
4 - exp( (xp(-22,54).
Tr mexp(-22/54+1)
X;
Proof. The first inequality follows from the fact that lower Darboux sum is smaller than the cor- responding Darboux integral
k-1
1-1
Σ
(S.
exp()x
St
exp(-22;5)d5
(exp(-2X;5k) exp(-22;55)
< exp(-2x,
24:
For the second inequality,
n exp((2X,St+1)
1) expppp (xx0-2x))))xp(22m)
17: -2,,-)) exp(2)
exp(2) exp(-22)Skk
2\;
< X
exp( -22;)))
which completes the proof.
The following lemma characterizes the difference between two consecutive auxiliary expectations Ag and AA-1- Lemma 4. Tmax S 1, then for all E [T%
Amaxy
exp(-22,Sk)
A AR- 1a-1=-500 (1IK (1
7Tk
Ex + Ck.
where the error term is bounded by
d
lex| <5
10 + >
me
3
1=1
i=1
Proof. By the definition of A and Ap-1, we have
A A = (00,0000,EE5)) (U
12, Mh, Sk+1)] - 10.0
=Eo Esa- AU (0x=1)))
where
AÚ 10,00 (0p-1) = Eg.~N(H0, x)) I
(0,-1 (0
Thgk, Ak. Sk+1) -0-1--U(B)) (U -U Sx).
39

---

Published as a conference paper at ICLR 2025
Expanding (0k-1)) based the definition of U gives
Ü (0 Bgg~N(HH)) 1.22)
10--1 m.gK.) 22 exp(-exppppp -2+1)+ OR-1
=:01 1)
d
1 - exp( (22,
2
1=1
=002(0-11
For 01(0x-1) evaluating the expectation gives
01(0x-1)=
A exp 2x;(-22;Sx+1)) ((1 (1-m) (1 + -x),
1=1
Then, we split (0x-)) into two parts:
1
(1(0x=1)= =
A - 21,
1,
Σ
A, expl -2,, S
2
1m1
i=1
=:01((F-11)
=122(02-1)
Let (0 = 012(0x-1) + (2(0x=1). Then ((04-1) = Ú11(0k-1) + (0(-)) We can
rewrite (1:0)=1)
1
Dc 1- e (xxp(-2);8x)CC - 21px
2
i=1
Since A; 10, for all i, by Lemma 2, we can find E1,i, E2, E 10, 0] such that )x))
+ (n (1 - 21),,) = exp( 2m) + x2)
Then we can rewrite Un (0: -1)=)aa as
(11(0=1))
(1 E1 )2 1xp(-22;Sk)0pp
1=1
1 Σ
Σ A exp( 1opp-22,5B)6pp +
E, 2
exp((2A,
Similarly, we can rewrite (3(0x=1)
1 (1 (1+ ( exp(-2X:Sa)
(3(0x-1)
Mh Eijt
d
1- 1- 1-eep((-2Ap3
T/ Eii
3
exp(
2
2
1=1
1=1
Therefore, we can rewrite
U(Oa-1) as
d
exp( exp(-2X;5t)
U(0x-1)
\ 1exp(-22,55)88-111 + E
+ (
2
2
1=1
where
d
d
(0c = 2
Σ
! exp( 20
Σ
En exp(-2X;,44
2
1=1
1=1
Subtracting U(0,-10 11, from the above expression, we can obtain the following formula for
(0 == (0 - (@0 x - Sk),
AÜ 1
(1k-1 7h Σ
1
exp(-2X,08) En + (0x-1))
1=1
40

---

Published as a conference paper at ICLR 2025
Taking the expectation of AÙ (0x-0))) over 8x-1 ( Esk-1), we have
A Ax-1 100.--00(0,.EEa-.. (0k-1)
1 - (9k -1=- 77k)
exp(-22:55s)
En + R 1eLL
Letting 6y: = Bol-i 1) (0x=((0)11) completes the proof.
The following lemma gives an upper bound for the term Eg, 100.00 1)102-1,1 that appears Lemma
Lemma 5. Mmax S 1 then for all l E [T] and E [d].
Amna'
100 11,2 - 0%, exp( (2 - Sk]) + A, Thuax Dis.
Proof. By the update rule, for all < 1 k - 1, have
E = (1 - n/ 2 + +11 E...
Since (1 - 712;;² I exp( 2m, ):) and S 7muxx we have the following bound:
E S A + MmTmmax Dig.
Expanding the recursion, have
1, exp(-22/121) Sz)) Σ Mmaa En exp(-22 -sx)
k-1
06 exp(<=2,,(55) St)) exp(2X, M
Mmax Dig
Σ
1/t expl 1app--2hp -21,
t=1
where the first line uses the identity II. II==144 (20(-2n;1)) = exp( 122; Sk)).
Further, by Lemma 3, we have Σk=1 7)t exp( -x; 22;5+++1 S $ expppep220 -21,5k)) Thus, we have
4
E 10021,11 < 0g exp(exp2222 22:(5 - -k)) + exp(2).5x 7max Eis
exp(-22x,8)
X;
exp( (x, Sa)) X, Immxx Dii,
which completes the proof.
Lemma 6. In the setting of Lemma 4, we can bound the sum of the error terms as
15
<
51Jmax
Si exp(-2
+ 2 Tmax
Egdi
1=1
1=1
Proof. By the upper bound of |EN|e
T
T
T
Ek
<
1=1
1=1
1
1=1
mh En expl
E2.
For E1,i, we apply Lemma 5 and have
T
E1
xx
(xxpp-22(SS------1 -S)) + X, Mmax Eis
exp( 1xp(-22xSS)d6
2
41/max X ME
th exp(-e22 Sw),
1
1=1
k=1
41

---

Published as a conference paper at ICLR 2025
For the first term, we have E AE I Mmax Σ
kk=1 7/
M/aax S1. For the second term, by Lemma 3,
we have
T
T
2 (xpp-22X,5) - Maax
cxp((22;5k))
-
Mmux
exp(-22;8r+1) = 7max
2X
2
Jmn
Rum
Putting these bounds together, we have
E1,o 7max exp(- (5p(-22,81)06, + 21mmaxx Diddi
For Eg we have
T
T
E2
Σ
mh E
S max EGg,
Σ
?h exp( -2),5k)
k=1
k=1
S max Eun
exp(- (P2p-22
2x;
2 2mmaxx Dzidis
where the second inequality uses Lemma
Putting the upper bounds of E1, and &g,s together proves the lemma.
Now we are ready to prove Theorem 3.
Proof for Theorem 3. According to (28), we have
T
[[((Or)] = Ao +
>(EEs, (As - Ak-1).
k=1
Using Lemma and Lemma 6, we have that
Σ (Ax A
1
7k
exp(-2218k)
E.ii +
2
k=1
where the error bound can be bounded
S
57max
S exp(-22,,,)
15
Mmax
2
Σ Endi.
Putting these together with the expression of Ag leads to the results in Theorem 3.
L
J.2 PROOF FOR THEOREM
Now take expectation over H, Σ and prove Theorem Throughout the proof, use ? to denote the lower incomplete gamma function, ((, = fe y and use I' to denote the gamma function, I'ss) Jo t" t and use I() I(s) (s, to denote the upper incomplete gamma function.
We first present two lemmas gamma functions.
Lemma 7, For all > -1 and C O, we have
X expl (
(29)
Proof. We substitute u CA and have
f X" exp(-CA)dd
/oa
CA
u exp( du
duu
1
Jo
cr
Cat1 a exp(-u)du
7 1
which completes the proof.
42

---

Published as a conference paper at ICLR 2025
Lemma 8. For all > -1 -1 2 0,
T(a, - 2(a 2 1)(at))
Proof. For all 2 0, it holds that e- I (a 1) To see this, it suffices to show that
g(t) == tut - < (a + 1)(at1), Note that the function g(t) is increasing [0, 2(a + 1)] and
decreasing on 2a + 1), +00). When = 2(a + 1), we have
g(2(a + 1)) = (2(a - 1))a= ((+t)) I (2(a + 2))==12((+t)) (a+1) = (a+ + 1)at1.
Therefore, for all / 2 0, we have 9(1) S (a + 1)(at2), Then, for all x 2 0, we have
[(a, x) + - 't dt / which completes the proof.
+00
(a 1)( d/ 2(a + 1)(a+1)e-
Now are ready to prove Theorem 1.
Proof for Theorem /. First, we recap some definitions from Assumption 1. The distribution of X
given by p(A) = z - We define μ == E|-]. By the definition of EΣ A] in Assumption 1,
E[E A| = FF exp(-r>) for some constant > 0, and E[A² A} = D2 - for some constant introduce the notations
D > O. Furthermore,
K, and B 1
p.
By the identity [[-] E [E[ Al], have μ = Jo Fnt P exx((-r\) 1/2 ddA, from which obtain F 1/2 5o 7ーPー exp( AAjAA
1 (1,r, So F 2r ((,rA))
It suffices to prove [[ - Z(t) = (t) + Munxx) only for t = T. Once we prove it
for T, we can easily apply the theorem for Fege, which original schedule truncated the first i steps, to get the result for all 1 s t T.
Based on the estimate of ((Or) given in Theorem 3, we can take the expectation over Xeo En and Oo to obtain the following bound:
15
[[ (0o, E)] <
5mmax
S
+E
Mmax
Ends
2
i=
miQi
=Q
where
E[M E]] = E
0 di
+E
1
170
E,
-
2
2
=;h
=11
T
exp(-22
(x,(-2A,5A)
m - TTK)
1=1
1=1
=:d In the following, bound Q2, I1, 12, la separately with the help of Lemma and Lemma
For Q1, we have
5d D
/ 3-- ex====xp((22
Q = 5d1mmax S,E, _X exp(-22;5i)A2]
Si
2
5d D
S
15.7 (. 2,25)
Z
O 1T
For Q2, we have
15d
Q2
[EA] = 01
2
43

---

Published as a conference paper at ICLR 2025
For I1, we have
d 11 Bpl
d
/// X X exp(-2ASi)d
1)2
2
22
d
D² 7(a, (2)))
=
22
12n((= 2STA)ST
2at12 2
By Lemma 8, we have (a,, 2S = [(a) e -((1) Thus, we can rewrite 1 as
d d()
I = 20+12 7
D2 D² + Ofe -R()))..)
For 12, we have
1 110
PoBp1e
4 1/
Byp -211)
d
d
Fu
A
x- exp( -xxppx-))exx((2225) x YP
4 MoH
4 110
Z
0
dF
4 70/4
Mh (3, (2S) + (2)) +r)
4Z
4 7011 + O (17ax: ).
For I3, we have
d
13
7h- m2) [Ep[E] EE xx((-2AS)))
k=
d
Fy 1/ x- x- -xp(---22) =da
MM--1 - 91)) |
Z Ju
k=1
T
d
Fy
Z (3, +2) (1)
MM-- - 7k)
4
k=1
Replacing F with
2, ((3rA).
then we have
13
d
Mh- - 7k) 1 ((3, + r))) ()
7 + r))-a)
((,
k=1
d
T
(M--1 - m1) (1 7(3, (2Sx + +r)A) 2 (
4
11
(3,
km
Setting the constants Lo. A, a, B, B, as (19), we can summarize our results for Q1 Q2 12, 12
Q1 = 01maxx S1 a-a-1)
Q2 = =0(maxx)
1 AST + 0(e
-((S))), ((-(s))))
1 Lo + ST
T
I3 B (nk-1 Th G(Sk)
where
(13 + +r)A)
G(z)
= :=
(Cx +1-
7(BrAA)
Putting everything together, we have [E(((O.) _(()
- -B|M(), F.)) + -(((8))) + mmmax ST
-8(1))
(Vhmaxx: + 2 + e D(S:) + Mmaxx: 8) ST min(a+1,9)) + max))
which completes the proof.
44

---

Published as a conference paper at ICLR 2025
0.35
0.30
0.25
Value
0.20
0.15
0.10
G(x)
G(x)
0.00
2
8
Figure 21: A comparison of G(z)) defined in (2) and G(x) defined in (18). G(x) follows exact power form, while G(s) follows a power form approximately. The gap between G(x) and G(x)
converges to as x - +oo. Here G(z) is defined as in (18) with parameter
2 1, and G(x))
r(B) }. In both cases, we the exponent ( 0.2. defined as in (2) with parameters C' 7(6,rA)
45


Reviews:
[{"review_id": "OVzmLcmR15", "reviewer": "Reviewer_g7Pt", "summary": "In this paper authors introduce an empirical power law that predicts the pretrianing loss for LLMs as a given training step. Method is tested across the set of learning rate schedules: constant, cosine and step decay. Using the proposed power law authors were able to find the schedule that outperforms widely used cosine learning rate schedule.", "questions": "It would be good to see the quality change on the validation or test set in addition to the training set.\nIf you have quality on the downstream tasks, that would be very important to look at as well for the proposed method. \nFrom these, I would also want to see how marginal the improvements are by looking at the stddev values for the validation loss due to the randomness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Improving the predictability of LLMs at different scales has been a very important task for many years. This paper attempts to make the projections very granular, in particular, authors propose an empirically tuned power law that allows to predict the training loss for a given model at every individual step across the set of learning rates. Within the range of tested LR schedules, authors also demonstrated that their method improves training loss compared to the popular cosine learning rate schedule.", "weaknesses": "I think this work has a lot of potential, but it is important to conduct additional testing and analysis of the proposed method:\n\n1. Given that the training loss demonstrates an improvement for the provided method, it is important to understand how the gain in training loss quality effects validation loss. In case the proposed method has a lower generalization capacity, it might not be beneficial to train with it even provided the improvement in the training loss. \n2. Several prior works tried to control training dynamics and make LLMs more predictable with various levels of granularity. For example, MuP paper: https://arxiv.org/abs/2203.03466 that proposes a list of adjustments in the architecture to improve predictability of the LLMs at scale. It would be very interesting to understand the benefits of the proposed method in the scenario when combined with techniques like MuP. \n3. During optimization of LLMs it is usually very important to look at the batch size together with learning rate. Here is the paper that proposes how the critical batch size should be estimate assuming the perfectly tuned learning rate: https://arxiv.org/abs/1812.06162. I think it is important to understand the temperature of the training in the proposed method and how the benefits would be affected once the batch size is tuned. \n4. One other thing that I struggled with was understanding what is the actionable outcome from this work. If the authors propose a new learning rate schedule derived from the power law equations, it should be rigorously tested across a set of downstream tasks, or at least authors should demonstrate improvements in terms of the validation loss and not just the training loss.  \n5. Finally, I would recommend authors testing their methodology against the existing scaling laws that allow us to understand benefits of this method when scaling the pretraining FLOP budget.", "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": "OL4MlNbBCp", "reviewer": "Reviewer_Uu5m", "summary": "This paper introduces a multi-power law for predicting loss curves in LLMs. The law is empirically derived, with all experiments conducted using Llama-2 models of various sizes (25M, 100M, 400M). The proposed law predicts training loss curves under decaying learning rates across different iterations. Building on Chinchilla's scaling law, it incorporates additional terms to account for the impact of learning rate decay, which Chinchilla’s law does not explicitly address. The authors fit the multi-power law to three loss curves under distinct learning rate schedules, demonstrating that it can predict loss curves for unseen schedules, validated using WSD. Additionally, they applied the law to optimize learning rate schedules, identifying one that resembles WSD but achieves faster convergence.", "questions": "1. How does the proposed law generalize to other model types, such as high-parameter models like Llama-7B? Currently, it has only been tested on models with 25M, 100M, or 400M parameters.\n\n2. The paper claims that the proposed approach is effective for long-horizon predictions, but the validation is limited to training curves with cosine and constant learning rate schedulers. Could you add plots supporting this claim for additional learning rate schedulers?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed multi-power law shows potential for predicting loss curves for unseen learning rate schedules, though it has not been sufficiently validated.\n\n2. The proposed law could be leveraged to optimize learning rate schedules, potentially accelerating LLM training.\n\n3. The paper provides valuable insights into pretraining dynamics, offering guidance for designing learning rate schedules that can simplify the complex task of learning rate tuning.", "weaknesses": "1. The theoretical justification for the proposed multi-power law is limited. The formulation of the law in Equation 2 lacks rigorous derivation or justification, raising doubts about its generalizability to unseen learning schedules.\n\n2. The paper does not clearly state the fundamental assumptions underlying the proposed law. It is unclear whether the additional terms in Chinchilla's law (the second and third terms in Equation 2) should be model-independent, batch-size-independent, or optimizer-independent. The paper appears to assume that the impact of learning rate schedules remains independent of these factors, a point with which the reviewer may disagree.\n\n3. The experimental section does not explore other types of learning rate schedules, such as cyclic or increasing schedules, leaving it unconvincing that the proposed law applies universally to various schedulers.\n\n4. It remains unclear if the learning rate scheduler derived from the proposed law consistently outperforms baseline schedulers. While the paper demonstrates reduced training loss at specific iterations with the proposed scheduler, it does not clarify if this advantage holds at convergence.\n\n5. The scope of application for the proposed law appears restrictive. The parameters of the law need to be fitted under the same task, model, batch size, and optimizer before yielding an improved learning rate schedule. It seems that the model, batch size, and optimizer must remain unchanged to apply the law effectively.\n\n6. There is a typo on line 491, and Figure 8 lacks clarity. The curves for the baseline methods are difficult to distinguish due to missing labels.", "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": "FWXby9Z7gs", "reviewer": "Reviewer_H6ae", "summary": "This paper introduces a multi-power law for loss curve prediction of large language models. By using power law to fit loss curve of different learning rate in each step, it can accurately predict the loss curves for unseen schedules of different shapes and horizons in different sizes of Llama-2 models. Moreover, by minimizing the predicted final loss across learning rate schedules, it is able to find a schedule that outperforms the widely-used cosine learning rate schedule.", "questions": "Please address the questions in weaknesses.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tIt achieves good loss curve prediction performance for 25M, 100M and 400M Llama-2 models and unseen learning rate schedules. \n2.\tIt can also help to find a better learning rate schedule than cosine learning schedule.", "weaknesses": "1.\tIt is empirical law which is not well-proofed as well as fully verified. It would be better to include more kinds of LLMs and larger model sizes since it can benefit more for large model training and new model structure exploration.\n2.\tIt only considers the relationship between learning rate and the loss curve. However, many other hyper-parameters will also affect the loss curve, such as model size, batch size, initialization seeds, etc. It would be better to take a comprehensive view so that the prediction could be more practical.\n3.\tIt is interesting that the multi-power law can also be used to tune optimal learning rate schedule for each case. It would be better to give more results on different kinds of models and larger model sizes.\n4.\tIt seems the parameters A, B, C, \\alpha, \\beta, \\gama,and L0 need to be tuned per model (or even per combination of other hyper-parameters). The tuning process needs to collect data from different learning rates. Each optimization takes over 5 × 10^4 steps. The tuning cost is a little high.", "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": "4RpFKrFl5A", "reviewer": "Reviewer_ytrq", "summary": "The authors formulate and fit an intricate multi-power law (MPL)\nfunction to estimate a subset of learning rate schedules using only\nvery simple staged LR schedules, i.e., those with constant LR over\ndifferent sections of time, for its derivation. The main contribution\nis the inclusion of learning rate decay in the formulation of the\nscaling law using the empirically motivated MPL term. It is important\nto note that the empirical law was fit with extremely few examples,\nrequiring only 9 training runs for strong fits on unseen data at test\ntime.\n\nThe method is used to propose an optimised WSD-like schedule with a\nsmoother transition between the constant (\"stable\") and decay phase.", "questions": "### Please address\n\nPage 1, line 052f:  \nI do not understand: the statement is made here (and in section 5.1)\nthat only monotonously decreasing schedules (\"we restrict our\nattention to the class of LR schedules that decay the LR over time\")\nare used. However, in sections 2 and 3 and appendix C.1, you state\nthat warmup was included in every experiment and taken into\nconsideration in practice.\n\nPage 10, line 516f:  \nYou state that the final learning rate for the \"tuned\" WSD schedule\nwas fixed to 1/10 of its peak. I don't think this makes for a fair\ncomparison, since your proposed optimized version decays to a far\nlower final learning rate. This is consistent with results on learning\nrate schedules in the LLM setting, which found that a decay to very\nlow learning rates (more specifically, zero) is best.\n\n### Minor comments\n\nPage 1, line 053:  \nShould this be \"$\\eta_{t} >= \\eta_{t + 1}$\" instead of \"$\\eta_{t} >=\n\\eta_{t - 1}$\"?\n\nPage 3, line 123:  \nI believe this should this be \"as illustrated in Figure 2(b)\", not \"as\nillustrated in Figure 1(b)\".\n\nPage 10, line 512f:  \n\"So the optimization of the WSD schedule is still disturbing.\"  \nPlease rephrase this sentence, it is unclear to me what you mean here.\nDo you want to imply that optimizing the WSD schedule is difficult?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper's findings and the design of the empirical law are extremely\nexiting, even with limitations such as only being able to work with a\nfixed maximum learning rate. Similarly, being able to extrapolate to\nlonger horizons is both important in practice but a testament to the\ndesign of the MPL.", "weaknesses": "Despite highlighting the flexibitily of the method, it is only tested\non few unseen LR schedules.\n\nIt would be nice to quantify how well/badly extrapolation to other\nmodel sizes works.", "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": "GJRq8sFVdV", "reviewer": "Reviewer_ic14", "summary": "The paper presents an empirical scaling law for loss prediction while considering learning rate schedules. The formula is in a multi-power form, which includes a power law that is applied on the sum of learning rates and another power law related to learning rate decay. The rationale behind the approximation has been explained. Loss curves with both predicted and actual values on 25M, 100M, and 400M parameter models have been obtained from experiments and shown in the paper. An approach for finding an optimal learning rate scheduled based on the empirical law is also presented.", "questions": "- Is it possible to analyze the sensitivity of different terms in the expression (2), and try to simplify the formula while still giving a similar prediction?\n- Could you add benchmark evaluation results, to see whether the loss improvement with optimized learning schedule indeed improves the accuracy?\n- For figures with two y axes in the same plot, such as Figure 2, could you put each of them into two separate plots with one y axis each, to improve the readability?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The inclusion of learning rate schedules in scaling law. \n- Experiments were conducted to study various aspects in this work.", "weaknesses": "- The scaling law expression in (2) is quite complicated, making it difficult to understand the intuitive meaning behind the formula. \n- To apply the result, the initial learning rate still needs to be specified. \n- One way of using the result, as presented in the paper, is to find the optimal learning rate schedule. However, the usefulness of such optimization in practice is not very clear. From Figure 1, the loss difference between the different learning rate schedules seems to be quite small. Only the loss is shown throughout the paper, and no benchmark evaluation result is presented, so it is not clear whether such small changes in loss would actually cause a significant difference in the model accuracy or not. \n- The consideration of models with 25M, 100M, and 400M parameters, adhering to the Llama-2 architecture, may be due to limitations in the GPU resource, which is understandable. However, most practical LLMs have at least 1B parameters, so the current results are insufficient for showing that this scaling law actually works for models with practical sizes. \n- Some figures, such as Figure 2, are quite hard to read and understand, due to the many curves in the figure without a clear labeling of them.", "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)