Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules
Authors: Binghui Li
Fengling Chen
Zixun Huang
Lean Wang
Lei Wu
Venue: NeurIPS 2025 spotlight ()
Keywords: Scaling Laws
Learning Rate Schedule
Kernel Regression
LLM Pre-Training
Loss Dynamics
Primary area: theory
TL;DR: We introduce a Functional Scaling Law that predicts full SGD loss dynamics under arbitrary learning rate schedules.

Abstract: Scaling laws have emerged as a unifying lens for understanding and guiding the training of large language models (LLMs).  However, existing studies predominantly focus on the final-step loss, leaving open whether the entire $\textit{loss dynamics}$ obey similar laws and, crucially, how the $\textit{learning rate schedule}$ (LRS) shapes them. We address these gaps in a controlled theoretical setting by analyzing stochastic gradient descent (SGD) on a power-law kernel regression model. The key insight is a novel $\textbf{intrinsic-time}$ viewpoint, which captures the training progress more faithfully than iteration count. We then establish a $\textbf{Functional Scaling Law (FSL)}$ that captures the full loss trajectory under arbitrary LRSs, with the schedule’s influence entering through a simple convolutional functional. We further instantiate the theory for three representative LRSs---constant, exponential decay, and warmup–stable–decay (WSD)---and derive explicit scaling relations in both data- and compute-limited regimes. These comparisons explain key empirical phenomena: (i) higher-capacity models are more data- and compute-efficient; (ii) learning-rate decay improves training efficiency; and (iii) WSD-type schedules outperform pure decay. Finally, experiments on LLMs ranging from 0.1B to 1B parameters demonstrate the practical relevance of FSL as a surrogate model for fitting and predicting loss trajectories in large-scale pre-training.

Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules
Binghui Li
Fengling Chen Zixun Huang?*
Lean Wang
Lei Wu 12.4.
Center for Machine Learning Research, Peking University
School of Mathematical Sciences, Peking University
State Key Laboratory of Multimedia Information Processing,
School of Computer Science, Peking University AI for Science Institute, Beijing
(libinghui, lean) edu. en, fichen_lwyccestt pku. edu. cn alexpku@stu, pku. edu. cn, leiwu@math, pku edu. cn
We strongly recommend reading the arXiv version of this paper, available at https://axiv..or/ab//2500...
Abstract
Scaling laws have emerged as unifying lens for understanding and guiding the training of large language models (LL.Ms). However, existing studies predomi- nantly focus on the final-step loss, leaving open whether the entire loss dynamics obey similar laws and, crucially, how the learning rate schedule (LRS) shapes them. We address these gaps in controlled theoretical setting by analyzing stochastic gradient descent (SGD) on a power-law kernel regression model. The key in- sight is novel intrinsic-time viewpoint, which captures the training progress more faithfully than iteration count. We then establish a Functional Scaling Law (FSL) that captures the full loss trajectory under arbitrary LRSs, with the sched- ule's influence entering through a simple convolutional functional. We further instantiate the theory for three representative L.RSs-constant, exponential decay, and carmuu-staabe-deese data- and compute-limited regimes. These comparisons explain key empirical phenomena: (i) higher-capacity models are more data- and compute-efficient; (ii) learning-rate decay improves training efficiency; and (iii) WSD-type schedules outperform pure decay. Finally, experiments on LLMs ranging from O.IB to IB parameters demonstrate the practical relevance of FSL as surrogate model for fitting and predicting loss trajectories in large-scale pre-training.
(WSD)-anc and derive explicit scaling relations in both
Introduction
well established that the training of large-scale deep learning models mysteriously follows scaling laws, which describe how model performance scales predictably with available resources such as compute or data [19]. In particular, the landmark study by Kaplan al. [25] demonstrated that. in LLM pre-training, the loss L decreases with model size M and dataset size D according to pawer-law relation:
(1) where as and ap are the scaling exponents, La denotes the irreducible loss, and Ca., Cp are some constants. Such empirical relations have proven remarkably robust across scales, architec-
D) Lo + Ca + CDD-an
"quual contribution. Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

dz, SGD results
FSL prediction
5 0.5,8=4..
0.5, 10, 4.0
Learning Schedules
10-
0.035
0.030
Risk
Risk
0.025
0.020
const
0.015
wsd
cyclic
wsd
0.010
2000
4000
6000
8000
10000
10a
Step
Data
(a) Loss Dynamics
(b) Scaling Behavior
Figure FSL accurately captures the loss dynamics and scaling behavior SGD PLK regression. In both subplots, solid lines denote the results of SGD, while dashed lines represent the corresponding FSL predictions. (a) FSL accurately tracks the loss dynamics of SGD, averaged over 1000 runs, for three learning rate schedules: cosine, WSD-like, and non-standard cyclic schedule. (b) FSL predictions (dashed) computed using the analytical forms from Section 5, and compared
mean 200 SGD rums (solid).
tures, and training setups [20, 60, 36] and have become foundational principles for guiding L.M development [18, 24, 1, 5, 55, 27]. In practice, they are now routinely used to design optimal resource-allocation strategies [20] and to tune key hyperparameters such as learning rates and batch sizes [36, 29].
Despite their empirical success, the theoretical understanding of scaling laws remains limited. Recent studies have begun to illuminate the underlying mechanisms [54, 22, 39, 63, 23, 42, 43, 2, 14, 3, 7, 35, 50, 8, 70), yet two important gaps persist:
Determinants scaling efficiency. Existing studies lack a systematic characterization of how key factors-such as model capacity, task difficulty, and hyperparameter choices- govern scaling efficiency, as reflected by the exponents 0' and o In particular, learning rate schedules (LRSs) are known to be critical in practice [40, 4, 17]. but their precise role in shaping scaling behavior remains unclear.
Beyond the final-step loss. The scaling law (1) focuses only the end-of-training loss [25, 20), thus leaving open whether the full trajectory follows similar laws. Empirical studies [59,3] 38] suggest this possibility, but the fits there are still crude and lack theoretical grounding.
1.1 Our Contribution
In this paper, take step toward addressing these gaps in controlled yet representative theoretical setting. We study stochastic gradient descent (SGD) training of the power-law kernel (PLK) regression-a widely adopted surrogate for scaling-law analysis [7, 50, 35, 8). The PLK regression characterized by four parameters: the task difficulty 8, the capacity exponent 3, the model size M, and the label-noise level To capture the influence of learning-rate schedules (LRSs), SGD via an intrinsic-time SDE, in which the concept of intrinsic time emerges as a key quantity enabling a unified characterization of how different LRSs shape the loss dynamics. Building this formulation, we establish the Functional Scaling Law (FSL). which provides a unified description of the entire loss dynamics beyond the traditional final-loss prediction. Concretely, for general intrinsic-timm LRS [0, oo) + [0, oo), and under some conditions, the dynamics of the expected (where denotes the intrinsic time) satisfies:
model
loss [x,))
[(R())))
+ <(1)
+
2
M
signal learning
error
approx.
noise accumulation
where e(t) (1 +) and KCtt) = (1 + - Each term in FSL admits clear interpretation:
denotes the irreducible error caused by label noise, - represents the approximation error, e(t) characterizes the signal-learning dynamics under noiseless (full-batch) gradient descent, and
a² 2
the final term captures the injection and dissipation of gradient noise, with the LRS y entering through

---

Table 1: Learning-rate schedule (LRS) strongly influences scaling efficiency in power-law kernel regression.
Efficiency is determined by two key factors: relative task difficulty E (0, 2)) and model capacity 8 > 1. We distinguish between an easy-learning regime (s e 1 - 1/3) and a hard-learning regime (s < 1 1/3).
Data- Optimal Sealing Laws
Compute-Optimal Scaling Laws
Learning Rate Schedule (LRS)
Easy
Hard
Easy
Hard
Constant
D
1-7443
Exponential-docay
D Thक (log DJ
D=o(log Dy
C Tean
(log C*)
C 149
(log C)
Warmup-stable-deaat (WSD)
D (log D)
TRRM
D
(log C) 9E79
C 197
tractable convolutional functional. The function K, a memory kernel, which we also refer to as the forgetting kernel, quantifies how fast the injected noise dissipates during training.
Building on FSL, we derive concrete scaling laws for the final-step loss under three representative LRSs constnnn exponential decay [15], and warmup- stauup--table-dee decay (WSD) [69, 21]-in both data- limited and compute-limited regimes. The results, summarized in Table 1, recover and extend prior analyses [7, 8, 50, 35], and reveal several unifying insights.
Scaling efficiency of different schedules. WSD achieves the best scaling efficiency, followed by exponential decay and then constant schedules. This efficiency hierarchy provides theoretical justification for learning-rate decay and explains empirical success of WSD [69, 21, 58, 36).
Role of model capacity. cighei-ccpacity models are consistently more efficient in both compute and data, highlighting the necessity of scaling model capacity [25].
Data-model trade-off. Compute-optimal training requires scaling data more than model size,
consistent with established heuristics in LLM pre-training [20].
Scaling law for peak learning rate. Optimal scaling requires the peak learning rate (LR) to
scale appropriately with the training budget (data or compute), revealing the importance of careful peak LR tuning [6, 29].
Beyond PLK regression, we further apply the FSL ansatz fit and predict loss trajectories from LLM pre-training experiments with model sizes ranging from O.1B to IB parameters, covering both dense and MoF. architectures. These results highlight the potential of FSL. practical surrogate for understanding and guiding L.M pre-training. To better situate our contribution, we provide a detailed comparison with related work in Appendix B.
Notation. For any n N, let [n] == {1, .,...,,,,, n).
For a positive semi-definite (PSD) matrix S, denote by (S) its j-th largest eigenvalue. and define the S-induced norm |u|ls := Vu Su for any vector u. We write = B (resp. A > B) B A (resp. A B) is PSD. Throughout the paper. we use R to denote equivalence up to constant factor, and S (resp. 2) to denote an inequality up to a constant factor. For two nonnegative functions f.g TR20 + Rgo. we write f(t) ~ g(t if there
exist constants C1, C2 > 0 (independent of t) such that Ctf(t) & g(t) < C2f(1), At > O.
2 Power-Law Kernel (PLK) Regression
Let X' denote the input domain and D the input distribution, and assume labels are generated by y (o(x), 0*) + e, where fx(x) == (Ø(x), , 0*) is the target function, and the label noise A(0, o²) is independent of x. Here Φ : X' + RN with N E N+ U {oo} is feature map, satisfying the following assumption:
Assumption 2.1 (Hypercontractivitty).. Let H := ExNDiΦ(> ((x)((()()))))))))())(((())))))) T] be the feature covariance. There exist constants C1, C2 > such that for any PSD matrix A NNN, C1 tr(HA) A I Ex-D (Ø(x)) (xxx)) (xx) |
HAH C2 C.H)H))
This condition ensures that the feature distribution is sufficiently regular- its fourth-order moments are controlled by the second-order ones [41]. holds, for example, for Gaussian features ((x) ~ (O, H) with C 1, C2 2 (see Lemma F.1). For simplicity, we also assume:
Assumption 2.2. = diag(A1. Ae.... ......... AN) with A1 3 X2 2
AN.
To learn f*, we consider model of width M: v) Σ U/ [ o(x)= =: (v, Wo(x)), where E RM denotes trainable weights and W E M projects the N-dimensional features onto an M dimeennsion subspace. We study two choices of projection W:
=
3

---

Top-M features: w, ej for j E [M]] i.e., selecting the 1op-M features (
Random-M features: w (O, Ix) independently for j E [M]
The top-M setting is particularly simple yet analytically representative case, widely adopted in prior scaling-law studies [43, 13). For random features [3, 7. 50, 35, 8), we will show that, in certain regimes, their scaling behavior closely parallels that of the top-M case. As clarified in Appendix A.3, our setup is equivalent to learning with the kernel Kg xx) Φ(x) Φ(x').
We now formalize the key notions of model capacity and task difficulty. Let Φy == و5/ for {
j E [N]. So } forms an orthonormal basis of (D))
Assumption 2.3 (Model capacity). The spectrum of the feature map satisfies X, ~ j jo 1.
The condition 3 > ensures tr (H) = شド >3 < C for some constant C' i'ndependent of N, making our analysis dimension-free and applicable the infiiit-diimensional setting (N oo).
For the top-M features, the model takes the form f(:v)- = -ー زyرy = E
ΣØ= " -//20 reveals that higher-index (less significant) features are increasingly down- weighte by the factor j-8/2. As & increases, the spectrum decays more rapidly, and the model effectively relies on fewer features. Hence, the model's expressive power is governed by two complementary factors: (i) the model size M, which controls how many features are retained, and (ii) the capacity exponent 3, which controls how quickly these features decay in importance.
Assumption 2.4 (Task difficulty). Suppose 10,12 = ア وー for some >> 0.
Under Assumptions 2.3 and 2.4, the target function admits the expansion f = Eド= ΣN 5-12 //2 ã, ~ > , 13/2 σ.- Since (ő,) } are orthonormal, this assumption implies that the spectral energy of I" decays as power law. The exponent a == x.3 therefore quantifies the task's intrinsic difficulty, which depends only the target function itself and is independent of the model spectrum. In contrast, 8 measures the relative difficulty with respect to a model of capacity 3:: for fixed f (fixed a), adopting higher-capacity model (smaller ()) increases s = a/ B3 making the task relatively easier. In other words, the same task appears easier to higher-capacity model.
0J00 ~
We remark that similar assumptions have been widely used in the analysis of kernel methods [12, 11, 57, 9, 39]. Our work builds upon and extends this line of research.
3 One-Pass SGD and Intrinsic- Time SDE
Given data point (x. y) E XR and model f(; define the loss ((z, v) 1 (f(; v)
Then, the population risk is R(v) = Ez [(( v)] = 2llWTv v 0" f + =: E(v + o² where
E(V) denotes the excess risk. We minimize R(v) via one-pass SGD, given by
Th Σ Vvl(z.va),
Vk+1 Vk
BB
xESt
where Sa = (XL,j )). j=1 is a mini-batch of i.i.d. samples, ?i and B3 are the learning rate and batch size, respectively. The initialization is set to vo = 0.
Throughout, we refer to 17- == (No, 7h .... 7K 1) as the learning rate schedule (LRS). Common choices in practice include the cosine [37, 60]. WSD [21], and multi-step [5] schedules (see Ap- pendix A.2 for details). To analyze the effect of LRS, we rewrite (3) as
Vk+1 Vk= 7Tk (vR(vv) + Ex);
where the gradient noise & = By Exes, Ve(z, vk) (R((v) satisfies [[[x] = 0, EE ]
B. E(vx), with E(-) denoting the noise covariance for batch size 1.
Continuous-time limit. Following prior work [30, 32, 33, 51, 34], we analyze the continuous- time limit of SGD rather than the discrete update (3) or (4). This perspective makes the analysis more tractable and clarifies the emergence of scaling laws. Fix discretization step size h > 0 and

---

let Pp := ma/h for k E N. Then, (4) becomes V&+1 v; -p (R((v)) Pphhk. For sufficiently small this iteration is well approximated by the Itô-type SDE [31, 45]:
h
dv, = -(()) +())
b(T) E(v, dB,.
(5)
where B, E TM is M-dimensional Brownian motion, and sp(.) is the continuous-time LRS satisfying (kh) = m//h for all l E N; b(.) is the continuous-time batch-size schedule satisfying b(kh) Bg for all N. In (5), the learning rate affects both the drift and diffusion terms, thereby coupling the deterministic and stochastic effects.
Intrinsic-time reparametrization. In SDE (5), the physical time serves as the continuous analogue of the discrete step index However, when the learning rate varies over time, the actual training progress is determined not by the number of updates k but by the accumulated step size Σ 7// which more faithfully reflects the total optimization effort. Motivated by this observation, we introduce an intrinsic time variable that rescales the physical time according the LRS:
L (())
1= T(T)
(6)
which measures the LRS -dj-adjested training duration. Let v. = vr-1(1) Applying Øksendal's time change formula [44] to the SDE (5) yields the intrinsic-time SDE:
dv, = - VR(vi) dt + va(() ((v))d) dB, with r(t) = hG( (t)
(7)
b(T-~((t)
Here (t) quantifies the joint effect of learning-rai rate and batch-size scheduling. Compared with (5), the LRS dependence is absorbed from the drift and retained only in the diffusion term, thereby decoupling the deterministic and stochastic effeect. This structural simplification greatly facilitates the subsequent scaling analysis.
For a clearer explanation of the connection between the discrete SGD (4) and the SDE formulations (5) and (7), we refer the reader Appendix A.4.
Intrinsic-Time Functional Scaling Laws
In this section, we present our main results on the Functional Scaling Law (FSL). All proof's are deferred to Appendix D. We begin with assumptions the learning-rate schedule and model size. Assumption 4.1. Suppose Assumptions 2.1, 2.3 and 2.4 hold. Assume both M and N are sufficiently large, and the LRS satisfies sup1zo (t) & C3 for sufficiently small constant C33 0.
Theorem 4.2 (Intrinsic-Time FSL., top-M features, hard-regime). Under Assumption 4.1, let vs denote the solution to the intrinsic-time SDE (7) with top-M features. Then, for f with difficulty E (0. - 1/8) and any z 0, it holds for all t. 2 0 that
B[(():)] - 0 M 3 + +((t) + 1' x[(t=z)(e(z) + (z)
(8)
where e(t) = (1+
= (1+) (1-
This theorem establishes that, for hard tasks with & 1 1/8, the loss dynamics fully charac- terized by the FSL (8). We explain the emergence of power laws in FSL from a multi-task learning perspective in Appendix A.5. Moreover, each term in the FSL (8) admits clear interpretation:
Irreducible error: 3 This term is due to label noise.
Approximation error: M This term corresponds to the error due to finite model size, with the scaling efficiency is determined by the task's intrinsic difficulty s3.
Signal learning: e(t). This term corresponds to learning under full-batch gradient descent, capturing the rate at which SGD extracts the signal f Moreover, the rate depends on the task's relative difficulty& s. For fixed target f" (fixed a = s9), increasing model capacity (smaller ( accelerates its convergence since a/3 becomes larger.
5

---

Noise accumulation: l/ - z))(() + + de. This term characterizes how the learning- rate and batch-size schedules shape the accumulation and dissipation of stochastic noise. The integrand (e(=) + (2)) represents the instantaneous noise magnitude, where e(z) captures mini-batch noise and x² captures label noise. The forgetting kernel K((.) quantifies how noise injected time 2 still affects the loss at time t. Due to K(t)) ~ t- -22-118), higher-capacity model (smaller 3) tends to forget noise more slowly.
Notably, the last two terms together constitute the optimization error and two key factors govern the trade-off between the them: (i) Model capacity: Increasing model capacity (8 4) accelerates signal learning but simultaneously slows noise forgetting. (ii) Learning-rate and batch-size schedules: Smaller learning rates or larger batch sizes suppress noise injection but also shorten the intrinsic training time. However, sufficient intrinsic time is important: the signal- term requires it to effectively reduce the risk, while the noise-forge term relies on to forget noise memorized in early training. Hence, effective schedules must balance these competing objectives- suppressiin injected noise while maintaining enough intrinsic time for both learning and forgetting.
4.1 General Results
The FSL (8) is established for the hard-learning regime where 8 < 1 - 1 We now show that an analogous FSL also holds in the general case. To state the result, we define eM(() (t) ΣØ 2y10312 2Azt, Kar((t)
X36- 22,tt One can verify that both functions exhibit power- Σラ M
law decay for SSIS M
K(r(() = (2-1111) en (1) ~ 1,
ISISM' M
Consequently, (1) ~ et)) and (o) KC()) for > O.
The following theorem provides characterization of the loss dynamics for general case:
Theorem 4.3 (Intrinsic-Time Time FSL. top-M features, general label noise). Suppose Assumption 4.1 holds. Let denote the solution to the intrinsic-time SDE (7) with the top-M features. Define Fm(l;7) el)()) + fh KM =))em(z)) + o²)(()) There exists 0 such that for 0 g & cMP, it holds that
E[R(P=)] - 00²² M- + (F771)
(10)
For all cM8 <t< co, it holds that
M- + (M))))) & [ i M- Fl:71:
(11)
Notably, the constants implicit in S are independent of the noise level
proof sketch is provided in Appendix C.3. The above characterization is uniform with respect to the label-noise level and holds for all s > 0 and 8 1. It asserts that the exact FSL relation (10) (i.e., the FSL (8)) remains valid up to the intrinsic time t < cM2 =: fu- For later times t > t.. although the FSL may longer hold exactly, the loss dynamics remain controlled from both sides in (11).
At the critical time (M, we have eM (M(( < M- indicating that signal learning has reached the approximation-error limit. Beyond this point, further training no longer improves the learned signal; instead, the dynamies become dominated by noise effects. Depending on the interaction between the stochastic gradient noise and the decaying learning rate, additional training may either inject more noise or dissipate it. Thus, is a priori unclear whether the total error will significantly increase or decrease after t M. The upper bound in (I1) ensures that the overall loss remains well-controlled, analogous to the behavior of the infinite- width limit (M = oo).
Nevertheless, an FSL may still hold for all > O, under suitably stronger conditions. In Theorem 4.2, we considered the setting with tasks satisfying s & 1- 1/3. The following result shows that similar characterization extends to general task difficulty with constant label noise.
Theorem 4.4 (Intrinsic-Time FSI, top-M features, constant label noise). Under Assumption 4.1, suppose a Z 1. Let Dt denote the solution to the intrinsic-time SDE (7) with the top-M features. Then, for any > 0 and all t T 0, [[]((,))] 10222 - + (xtttt)).

---

Theorem 4.2 implies that the finite-M M functions eM and KM can be replaced by their anfiite-widdte counterparts Ex and Xx the hard-learning regime. The next result demonstrates that the same FSL. charactteeriittoonn naturally extends to the noiseless case = 0.
Theorem 4.5 (Intrinsic- Time FSL. top-M features, zero label noise), Suppose Assumption 4.1 holds and 0. Let Vs denote the solution to the intrinsic-time SDE (7) with the rop-M features. If E [0.2 1/3,, then for all 3 [[[(())] - + eM + So (t)(t)))emt (z)) dz.
Random-M M features. For the random- featuree case, the modified feature covariance matrix is H = WHW whose cigenvalues we denote by A e do 3 e Am- We similarly define وド= X3e- 22,, The next theorem establishes that the same FSL. also holds when the top-M features are replaced by randomly selected features.
(lr(1) = Σソ= 25109120-22, RM((())
Theorem 4.6 (Intrinsic- Time FSL. random-M features). Suppose Assumption 4.1 holds and E (0, 1]. Let Vi denote the solution to the intrinsic-time SDE (7) with the random-A M features. Then, with probability at least 1 - results of Theorems 4.3, 4.4, and 4.5 continue to hold, after replacing es (-) and Kst (-) with their random-featue feature counterparts (.) and RM(.), respectively.
over the randomness of the projection matrix W, the
Lemma 4.7. With probability at least 1 - exp(- over the randomness of the projection matrix W. it holds that X3- j for any j E [M]]
Theorem 4.6 and Lemma 4.7 together imply that when the task difficulty satisfies & 1, training with random-M M features similar to using the top-M features, up to exponentially small probability, We emphasize, however, that for easier tasks with > 1, the behaviors of random and top feature may diverge and leave this for future investigation.
Learning Rate Schedules Impact Scaling Efficiency
Having established the general FSL, we now instantiate it under three representative PoSss-onstant exponential decay, and warmup-stable-decat (WSD)-to examine how schedule design influences scaling efficiency. All proofs can be found in Appendix E. For clarity, we make:
Assumption 5.1. Assume constant label noise o² > 1 and batch size b(r) = B for all 7 > 0.
Under this assumption, given physical-time LRS function sp(.), Theorem 4.4 implies that the FSL for t R simplifies [[R(v:)] 1o2 = -8 + e
T So (t- ((T)))))(((-((()))
Let EK- [B[R(KK))] f denote the expected excess risk after K training steps. For each LRS, we derive concrete scaling laws describing how ε scales with the model size M, the total step count K, as well as the LRS's hyperparameters. We then reinterpret these results from a resource-allocation perspective by optimizing under two canonical constraints: (i) the data-limited regime, where the total data size == BK is fixed; and (ii) the compute- limited regime [20], where the total compute C== MDr is fixed. For each regime, we further examine how the optimally tuned hyperparameters (e.g., the peak learning rate) should scale with increasing available resources. Finally, for clarity, we distinguish between two task regimes: an easy-learning regime, where > 1 1/3, and hard-learning regime, where s < 1/B.
5.1 Constant LRS
Theorem 5.2 (Scaling law for constant LRS). Under Assumption 5.1, we have Ex R M- (n)))*++ 102
Let 7 == BnBB be the effective learning rate. Then, the scaling law can be rewritten as Ex ~ M- 43
h(y,, M, D), where the excess risk depends the learning rate via 7 = / B. This suggests that we should scale the learning rate linearly with respect to batch size (a.k.a. linear scaling rule) [26, 16, 40].
((
Data-optimal scaling. Clearly, this involves minimizing h(-) while keeping D fixed. A straightfor- ward calculation yields: Topt R D - Mopp Z DaFro E = D- ने.. Notably, both the best achievable excess risk Eopt and optimal learning rate Yopt depend exclusively the task's 's relative
7

---

difficulty s. For a fixed target (fixed a), higher-capacity model (smaller 3) gives a larger =a/B and is therefore more data-efficient.
Compute-optimal scaling. This involves minimizing k(-) while keeping C == DM fixed.
The solution is summarized follows, with the derivation deferred to Appendix E.1: Yopt
CT CTT401 Mopt CT-6413 Dops CTF Eupt ~ CT 1+-0+0
(=+118
This shows that the performance of the compute-optimal model improves with the total compute budget C' in power law. For fixed task (a sB fixed), we have the following observations:
Increasing model capacity ( [ 4) enhances compute efficiency the extra / in the scaling T- quantifies this gain. This explains a well-known empirical observation in LLM pre-training: Large models are more compute-efficien than small models [25, 20].
exponent .8
The optimal learning rate "opt. decreases as C grows, and the compute-optimal allocation favors investing more in data than in model size again consistent with current LLM pre-training practice [5, 55, 20].
Note that [8] also investigated compute-optimal scaling for constant LRS but assumed a fixed learning rate and no label noise. In contrast, consider more realistic scenario where the learning rate is optimally tuned and the irreducible risk is present, leading to a compute-optimal scaling law that matches empirical observations.
5.2 Exponential Decay LRS
For given number of training steps [15, 65], an exponential decay (exp-decay) LRS given by ( = a exp((=Ar) ((Kh) = b, where À is chosen such that ((Kh) = 6. For brevity, we assume 1. Note that the hyperparameters and specify the peak and final learning rates, respectively.
Theorem 5.3 (Scaling law for exp-decay LRS). Under Assumption 5.1, we have EK ~ M- T- + 2
(a min(M,T1/33) where T () K/ /////////////// log(a/b) b) is the total intrinsic time. BT
Let b = a/K. Then the intrinsic time becomes T a(K 1)/ log K, whereas constant LRS with step size 7) = a yields T' = aK. Thus, exp-decay LRS drives the learning rate down to as small as a/ K, yet sacrifices only logarithmic factor of intrinsic time compared to the constant schedule.
Data-optimal scaling, Let ς = a/ B be the effective peak learning rate. By minimizing the right hand side of the exponential decay scaling law with respect to a, b, K, B, M under the constraint KB D (see Appendix E.2), We obtain Mopt oc and
For s 1 Topt (D/ log D) 1440 and Eopt (D/ log D) *841
For s < 1 3 Yopt ~ 1 and Eopt (D/ lo/log)) D)-
Compared with the constant LRS, exp-decay LRS achieves strictly faster decay of the excess risk, justifying the importance of learning-rate decay in stochastic optimization.
Compute-optimal scaling, A straightforward calculation (see Appendix E.2) yields:
For 8 3/ 3. Tapt 2 ( (og
21.0
Mopt
( Dopt R C2 (log C)) शोंग, and
Eopt
C (1oge)
aa 2437
172, Dopt 4 (log C) Cftt,, Eopt (1
For Topt R Mopt
143
In the easy-lcarning regime, the excess-risk rate determined solcly by the intrinsic difficulty a = s3: hence, increasing model capacity alone does not lead to asymptotic gains. The compute-optimal allocation consistently favors data over model and moreover, the optimal compute split depends solely on the task's intrinsic difficulty, with ratio Dopt/Mon R (a/(2+a) decreasing as the task becomes harder. This implies that, for harder tasks, one should allocate more compute to increasing model size. The optimal Tops decreases with the compute budget C, and for fixed a, higher-capacity models (3 !) require smaller Topt.
8

---

In the hard-learning regime, data still dominates compute allocation, but now the optimal split depends only on model capacity. independent of the task difficulty. Moreover, the optimal maximal learning rate remains constant (7opt ~ 1). These results imply that single, universal choice of compute split and learning rate suffices to attain optimal scaling across all tasks satisfying 8 < 11--// -1/9, greatly simplifying hyperparameter tuning. Finally, in this regime, higher-capacity models (B3 4) become strictly more compute-efficient as evidenced by the excess-risk scaling exponent -s8/(1 + B).
5.3 WSD-like LRS
We lastly turn to consider a WSD-like LRS [69, 21], which comprises Ki-step stable phase followed by Kz-step decay phase, for total K = K1 + K2 steps, given by
Sa (p(T) { 1aexp(->(T-Ka Kib)
if K1h;
(12)
if > Kih.
where À is chosen such that (KKh) = b. For brevity, we assume h = 1 and let = K2/K. This schedule thus characterized by three hyperparameters: the peak learning rate a, the final learning rate b, and the decay proportion r, which controls the duration of decay-phase.
Theorem 5.4 (Scaling law for WSD-like LRS). Under Assumption 5.1, we have for the LRS (12):
EK ~ M #3 + (Ti + T2) + o² 1/ + (a - b) (M.
where T1 = aK1 and T/ = (a - b) log(a/b) denote the intrinsic training times of the stable and decay phases, respectively.
172
We see that WSD-like LRS can leverage the initial stable phase to boost the intrinsic training time. For a decay proportion r < 1, we have T = Ti- + +T2 > Ka, which far exceeds the the intrinsic time 7 R aK/ log K achieved by the pure exp-decay LRS. Consequently, WSD removes logarithmic factors in the full-batch GD term, without altering the noise term's 's order as long as r >> O. Building on this insights, we show that WSD can indeed improve the scaling efficiency, detailed below.
Data-optimal scaling. Assuming = a/ K, we have Mopt = 30 and
Itag-a 1440
0-1 (log D) 1-53, Topt E (0, 1), Eapt ~ D * (log D)
vil 1443
For 8 > 1 3. Topt ~ D
For 8 < 3. Topt 5 1, ropt 2 ش1 log D, Eapt ~ D=*
Compared with the exp-decay LRS, both regimes enjoy a logarithmic improvement in excess-risk decay. In particular, for the bard-learning regime, the logarithmic factor disappears. This improvement requires the decay-phase duration only needs scale sublinearly with D, as indicated by loop + as D + oc. This matches the WSD practice in LLM pre-training, where the decay phase typically occupies only 10%-20% of the total training duration. Moreover, our theory suggests that for harder tasks, the decay fraction can be reduced further to enhance compute efficiency.
Compute-optimal scaling. Analogous improvements hold in the compute- limited regime. Assuming b = a/K K and imposing the compute constraint MD = C, the compute-opting satisfies:
For s
>
$. Topt ( (ogge)
1t-d-p 2-23
Topt
E (0, (), Mopt
( ) Dopt
Imaf Cr (log C² C) and Eopt = O 2
48=4 2130.
(log C)
For s < x1-1 3. Yopt ~ Topt 2 D 8-1-88
B log D, Mopt = Citt, CDot ~ Cits, ~ - 1-3
6 Experiments
6.1
Power-Law Kernel Regression
While the FSL is derived in the continuous-time limit, now verify that it also accurately captures the loss dynamics and scaling behavior of the discrete-time SGD (3). Specifically, we consider the PLK regression with difficulty and the results are shown in Figure
0.5 and capacity B = 4, corresponding to a hard-learning regime
FSL accurately captures the loss dynamics of SGD. Figure ((efft) compares the loss dynamics of SGD with the predictions of the FSL under three representative LRSs: cosine, WSD, and an unconventional cyclical schedule [56]. Across all cases, the FSL provides a remarkably accurate description of the SGD's loss evolution. Comparing the WSD and cosine schedules, we observe that the loss under WSD exhibits slower decay during the stable phase but undergoes much

---

LLsMA,
LLaMA, 20B lokens
LR5s 1B OwanMoE
1B OwenMoE. 208 tokans
611
100010
611
CO5
Learning rate
BU000S
2.50
wsd
Risk
2.6
Fllik
(1 010006
2.43
opt
811
811
0L0004
9.3
005
BU0002
wsd
2.30
23-
0.0000
20000
50003
10000
30000
20000
20000
Step
Step
Step
Step
(a) Fitting and prediction using ISL.
(b) The FSL-optimal LRS and its performance
Figure 2: Experiment on LLMs. (a) Fitting and predictive accuracy of the FSL on dense LLaMA
models. (b) Left: comparison of various LRSs. Right: loss trajectories of the FSL-optimal schedule versus baseline LRSs on IB QwenMoE model.
sharper drop once the decay phase begins, ultimately yielding a lower final loss. This seemingly counterintuitive two-phase dynamical behavior of WSD aligns well with empirical observations in practical LLM pre-training [21. 64, 58].
FSL predicts the scaling behavior of SGD. Figure ((right) further validates the scaling laws derived in Section 5 for the three canonical LRSs- constant, exponential, and WSD-like (12). The results show that the final-step loss of SGD closely follows the theoretical predictions of FSL. Among these schedules, WSD yields the best scaling performance, followed by exponential decay, while the constant schedule performs the worst. More experiment details and additional results experiments with varying (s, 3) and other LRSs are provided in Appendix C.I, and exhibit consistent behaviors.
6.2 LLM Pre-training
We now evaluate the practical utility of FSL as a surrogate model for capturing the loss dynamies of LLM pre-training. Specifically, three popular LRSs: cosine, WSD, and the 8-1-1 [5] are considered; see Figure 2b(left) for a visualization. In the 8-1-1 LRS, the learning rate is reduced by a factor vIO at 80% and 90% of the total token budget, yielding final value that is 0.1 times the peak learning rate. For more experiment details, we refer to Appendix C.2.
FSL accurately fit and predict loss curves. We first quantify the descriptive and predictive power of FSL. Following the protocol of [59] and [38], we restrict attention to the post-warmup portion of the loss trajectory. Two Llama [60] models (400 M and 1 B) are trained 20 B tokens under the three LRSs. For each model we (i) fit the FSL parameters on the loss curve obtained using the 8-1-1 LRS and (ii) deploy the fitted FSL to predict the loss curves of the cosine and WSD schedules. Figure 2a demonstrates that FSL not only fits the 8-1-1 trajectory accurately but also generalizes reliably to the unseen WSD and cosine schedules for both model sizes.
The ESL-optimal LRS is WSD-like. We next leverage the fitted FSL to design improved LRSs. Specifically, numerically minimize the final-step loss over the space of LRSs using the fitted FSL. This experiment employs a IB-parameter QwenMoE model [68], trained on 20B tokens using the same three LRSs. We fit the FSL using the trajectory from the 8-1-1 LRS and numerically solve for the FSL-optimal LRS. The model is then trained under this FSL-optimal LRS, using the same compute budget, and compared against the baseline LRSs. Figure 2b(left) shows that surprisingly. the FSL-optimal LRS is WSD-like and the decay phase drives the learning rate far below the conventional 0.1 threshold. This echos recent empirical recommendations by [4. 17]. Furthermore, Figure 2b(right) demonstrates that the FSL-optimal schedule yields a strictly lower final loss than all baselines, substantiating its practical relevance. Taken together, these results suggest that FSL is faithful surrogate for studying LLM training dynamics and a principled tool for interpreting and designing LRSs in large-scale pre-training.
7 Conclusion
In this paper, we present systematic study of how LRS shapes the loss dynamics in kernel regression. Specifically, we establish a novel functional-leve scaling law, which precisely characterizes the loss dynamics of SGD for general learning LRSs. The utility of our FSL is demonstrated through detailed analyses of three widely used LRSss providing theoretical justification for several prevailing practices in L.M pre-training- monii most notably, offering an explanation for the effectiveness of the empirically popular but previously less-understood WSD schedules.
10

---

Acknowledgement
Lei Wu supported by the National Natural Science Foundation of China (NSFC125221200 NSFC92470122, and NSFC12288101). Binghui Li is supported by the Elite Ph.D. Program in Applied Mathematics at Peking University. We are grateful Kaifeng Lyu, Kairong Luo, and Haodong Wen for generously sharing their work [38], which greatly inspired this study. We also thank Tingkai Yan, Yuhao Liu, Yunze Wu, and Zean Xu for many helpful discussions, and the anonymous reviewers for their valuable feedback.
References
[1] Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for generative mixed-modal language models. In International Conference on Machine Learning, pages 265-279. PMLR, 2023. (cited on page 2)
[2] Alexander Atanasov, Jacob A Zavatone- Veth, and Cengiz Pehlevan. Scaling and renormalization high-dimensional regression. arXiv preprint arXiv:2405, 000902, 2024. (cited on pages and 19)
[3] Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):e23118881)) 18781121, 2024. (cited on pages 2, 4, and 19)
[4] Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to zero: Why linearly decaying the learning rate to zero works best for 1lms. arXiv preprint arXiv:2502. 15938, 2025. (cited on pages 2 and 10)
[5] Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong Qiushi Du, Zhe Fu, and Others. DeepSeek LLM: Scaling open-source language models with longtermism. arXiv preprint ar arXiv:2401. 02954, 2024. (cited on pages 2. 4. 8. 10, and 17)
[6] Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal Ir across token horizons. In The Thirteenth International Conference on Learning Representations. (cited on page 3)
[7] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. dynamical model of neural scaling laws. arXiv preprini urXiv:2402.01092, 2024. (cited pages 2, 3, 4, and 19)
[8] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning can improve neural scaling laws. arXiv preprint 14XXv::2499.178555 2024. (cited on pages 2, 3, 4, 8, and 19)
[9] Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves kernel regression and wide neural networks. In International Conference on Machine Learning, pages 1024-1034. PMLR, 2020. (cited page 4)
[10] Sébastien Bubeck and Others. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-357, 2015. (cited on page 23)
[11] Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7:331-368, 2007. (cited on page 4)
[12] Andrea Caponnetto and Ernesto De Vito. Fast rates for regularized least-squares algorithm. 2005. (cited on page 4)
[13] Shihong Ding, Haihan Zhang, Hanzhen Zhao, and Cong Fang. Scaling law for stochastic gradi- descent in quadratically parameterized linear regression. arXiv preprint arXiv:2502.09106 2025. (cited on page 4)
[14] Elvis Dohmatob, Yunzhen Feng, Pu Yang, Francois Charton, and Julia Kempe. A tale of tails: Model collapse as change of scaling laws. arXiv preprint arXiv:2402.0700 2024. (cited on pages and 19)

---

[15] Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: near optimal, geometrically decaying learning rate procedure for least squares. Advances in neural information processing systems, 32, 2019. (cited on pages 3 and 8)
[16] P Goyal. Accurate, large minibatch SGD: training imagenet in hour. arXiv preprint arXiv:1706.02677 2017. (cited on page
[17] Alex Hägele, Elie Bakouch, Atli Kosson, Leandro Von Werra, Martin Jaggi, and Others. Scaling laws and compute-optimal training beyond fixed training durations. Advances in Neural Information Processing Systems, 37:76232-76264, 2024. (cited on pages 10, 17, and 20)
[18] Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, and Others. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701, 2020. (cited page
[19] Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically, arXiv preprini 17Xi
2017. (cited on page 1)
[20] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, and Others. Training compute-optimal large language models. arXiv preprint 2022. (cited on pages 2, 3, 7, 8, and 16)
[21] Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, and Others. MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.0633 2024. (cited on pages 3, 4, 9, 10, 17, and 20)
[22] Marcus Hutter. Learning curve theory. arXiv preprint arXiv:2102.04074 2021. (cited on pages 2 and 19)
[23] Ayush Jain, Andrea Montanari, and Eren Sasoglu. Scaling laws for learning with real and surrogate data. arXiv preprint arXiv:2402.04377 044.74477 2024. (cited on pages 2 and 19)
[24] Arlind Kadra, Maciej Janowski, Martin Wistuba, and Josif Grabocka. Power laws for hyperpa- rameter optimization. arXiv preprint arXiv:2302.004411 2023. (cited on page 2)
[25] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini ar arXiv:2001.0880 2020. (cited on pages 2, 3, and 8)
[26] Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint
2014. (cited page 7)
[27] Tanishq Kumar, Zachary Ankner, Benjamin F Spector, Blake Bordelon, Niklas Muennighoff, Mansheej Paul, Cengiz Pehlevan, Christopher Ré, and Aditi Raghunathan. Scaling laws for precision. arXiv preprint arXiv:241 1.043330 2024. (cited on page 2)
[28] Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch momentum: batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems, 01136944-369577 2022. (cited on page 27)
[29] Houyi Li, Wenzhen Zheng, Jingcheng Hu, Qiufeng Wang, Hanshan Zhang, Zili Wang. Shijie Xuyang. Yuantao Fan, Shuigeng Zhou, Xiangyu Zhang, and Daxin Jiang. Predictable scale: Part I - optimal hyperparameter scaling law in large language model pretraining. arXiv preprint arXiv:2503.04715, 2025. (cited on pages 2 and 3)
[30] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110. PMLR, 2017. (cited on page 4)
[31] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms Mathematical foundations. Journal of Muchine Learning Research, 20(40):1-47, 2019, (cited pages 4 and 5)
12

---

[32] Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate, Advances Neural Information Processing Systems, 53:14454-14555, 2020. (cited on page 4)
[33] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). Advances in Neural Information Processing Systems, 34:12712 12-22222-2225, 2021. (cited on page 4)
[34] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022.
(cited on page 4)
[35] Licong Lin, Jingfeng Wu, Sham M Kakade, Peter Bartlett, and Jason D Lee. Scaling laws in linear regression: Compute, parameters, and data, arXiv preprint vrXiv:2-06.004666 2024. (cited on pages 2, 3, 4, 19, 36, and 37)
[36] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, and Others. DeepSeek- V3 technical report. arXiv preprint arXiv:2412. 19437, 2024. (cited on pages 2, 3, and 20)
[37] Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608s 03983, 2016. (cited on pages 4 and 17)
[38] Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Zhiyuan Liu, Maosong Sun, Kaifeng Lyu, and Wenguang Chen. A multi-power law for loss curve prediction across learning rate schedules. In NeurlPS 2024 Workshop Mathematics of Modern Machine Learning, 2024. (cited on pages 2, 10, 11, 17, and 20)
[39] Alexander Maloney, Daniel A Roberts, and James Sully. A solvable model of neural scaling laws. arXiv preprint
2022. (cited on pages 2, 4, and 19)
[40] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint arXiv: 1812.0666, 2018. (cited on pages 2 and 7)
[41] Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Generalization error of random feature and kernel methods: hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis, 59:3-84, 2022. (cited on page 3)
[42] Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems, 36, 2024. (cited pages and 19)
[43] Yoonsoo Nam, Nayara Fonseca, Seok Hyeong Lee, and Ard Louis. An exactly solvable model for emergence and scaling laws. arXiv preprint arXiv:2404. 1754404..756,, 2024. (cited on pages
and 19)
[44] Bernt Øksendal. Stochastic differential equations. Springer, 2003. (cited on page 5)
[45] Antonio Orvieto and Aurelien Lucchi. Continuous-time models for stochastic optimization algorithms. Advances in Neural Information Processing Systems, 32, 2019. (cited on page 5)
[46] Zhixuan Pan, Shaowen Wang, and Jian Li. Understanding LLM behaviors via compression: Data generation, knowledge acquisition and scaling laws. arXiv preprint arXiv:2504.0959) 2025. (cited on page 19)
[47] Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. Sgd in the large: Average-case analysis, asymptotics, and stepsize criticality. In Conference on Learning Theory, pages 3548 3626. PMLR, 2021. (cited on page 27)
[48] Courtney Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on large- scale, quadratic models. Advances in Neural Information Processing Systems, 34:9229-92400 2021. (cited page 27)
[49] Courtney Paquette, Elliot Paquette, Ben Adlam, and Jeffrey Pennington. Homogenization of sgd in high-dimensions: Exact dynamics and generalization properties. Mathematical Programming, pages 1-90, 2024. (cited on page 27)
13

---

[50] Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of
compute-optimal neural scaling laws. arXiv preprint Stt1:2400.150774 2024. (cited on pages
3, 4, 19, and 27)
[51] Scott Pesme, Loucas Fillaud-Vivien and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218 11211189999 292300, 2021. (cited on page 4)
[52] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, and Others. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019. (cited on page 233)
[53] Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, and Francis Bach, The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. arXiv preprint 2.0.:::2501.19666. 2025. (cited on page 20)
[54] Utkarsh Sharma and Jared Kaplan. A neural scaling law from the dimension of the data manifold. arXiv preprini arXiv:2004. 10802, 2020. (cited on pages 2 and 19)
[55] Xian Shuai, Yiding Wang, Yimeng Wu, Xin Jiang, and Xiaozhe Ren. Scaling law for language models training considering batch size. arXiv preprint arXiv:2412.01505, 2024. (cited on pages
2 and 8)
[56] Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pages 464-472. IEEE, 2017. (cited on page 9)
[57] Stefano Spigler, Mario Geiger, and Matthieu Wyart. Asymptotic learning curves of kernel methods: empirical data versus seacher-studen paradigm. Journal of Statistical Mechanics: Theory and Experiment, (0220(2))124000) 2020. (cited on page 4)
[58] Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, and Others. Kimi K2: Open agentic intelligence. arXiv preprint arXiv:2507.20534 2025. (cited on pages 3, 10, and 20)
[59] Howe Tissue, Venus Wang. and Lu Wang. Scaling law with learning rate annealing. arXiv preprint arXiv:2408. 11029, 2024. (cited on pages 2, 10, 16, 20, and 23)
[60] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, and Others. LLaMA 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09280 2023. (cited on pages 2, 4, 10, 17, and 23)
[61] Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge university press, 2019. (cited on page 41)
[62] Mingze Wang and Lei Wu. theoretical analysis of noise geometry in stochastic gradient descent. arXiv preprint 2rXiv:2310.00062 2023. (cited on page 26)
[63] Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. In International Conference on Machine Learning, pages 33549-23588. PMLR, 2022. (cited on pages 2 and 19)
[64] Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, and Tengyu Ma. Understanding cormup-stable-deca learning rates: A river valley loss landscape perspective. arXiv preprint arXiv:2410.05 192, 2024. (cited on pages 10 and 20)
[65] Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last iter- ate risk bounds of SGD with decaying stepsize for overparameterized linear regression. In International Conference on Machine Learning, pages 24280-24314 PMLR, 2022. (cited on
page 8)
[66] Lei Wu and Weijie J Su. The implicit regularization of dynamical stability in stochastic gradient descent. In International Conference on Machine Learning, pages 37656-37684. PMLR, 2023. (cited on page 29)
14

---

[67] Lei Wu, Mingze Wang, and Weijie Su. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35:4680-4693, 2022. (cited on pages 26 and 29)
[68] An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, and Haoran Wei. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115, 2024. (cited pages 10 and 23)
[69] 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, pages 12104-12113, 2022. (cited on pages 17, and 20)
[70] Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale in pre-training? arXiv preprint arXiv:2410.21676 2024. (cited on page 2)
15

---

Appendix
Table of Contents
A
Miscellanea
A.1
Empirical Fitting of LLM Pre-training Loss Trajectory
A.2
Popular Learning Rate Schedules
A.3
Connections Kernel Regression
A.4
The SDE Modeling
19
A.5
The Emergence of Power Laws
19
B Related Work
19
Experiment Details and Additional Results
20
C.I Power-Law Kernel Regression
20
C.2 LLM pre-training
................................
21
C.3 Key Proof Steps and Core Insights
................................
26
Proofs Section
D.I
Volterra Integral Equation Governing the Loss Dys.i...........................
28
D.2 Power-law Decay of ex and KM Functions
30
D.3 The Case of Top-M M Features
D.4 The Case of Random-M Features
36
E
Proofs Section
42
E.1
Proofs Constant LRS
42
E.2 Proof for The Exponential LRS
................................
43
E.3 Proof for the WSD-Like LRS
48
Auxiliary Lemmas
Miscellanea
A.1 Empirical Fitting of LLM Pre-training Loss Trajectory
The Chinchilla Law [20] describes the final-step loss L as follows:
L(M, D) Lo + A,M x1 + A2D
where Lo. A1, Ag. Ki, and K2 are constants, and D and M represent the amount training data (tokens) and model size (number of parameters), respectively.
Later, [59] proposed the Momentum Law, heuristic rule designed capture the full loss trajectory. Given learning rate schedule == {n, hj. the loss at the k-th step is modeled as
(n) Lo + AST CS2,
where
k = 1ks S2 ΣΣ (ni-1
S
2-2
Here Top, A, C, and x are constants, and E (0, 1) is hyperparameter representing the decay factor for learning rate annealing, which typically ranges from 0.99 to 0.999.
16

---

Subsequently, [38] proposed the Multi-Power Law (MPL), which replaces the S2 in the Momentum Law with additional power laws better capture the progressive loss reduction induced by learning- rate decay. Specifically, the MPI. takes the following form:
(x = Lu + AST LD(k),
(13)
where
&
Σ LD(K): = C 1--1-Ma m)G( St),
S. Σ
G(r)== (C's + 1)-
j=1
Here Lo, A, C, C², KK, K' R" constants.
A.2
Popular Learning Rate Schedules
Here, we introduce some widely used LRSs in the context of LLM pre-training.
Cosine Schedule [37]. The schedule is given by = 1te
2 Mmax +
1- 2 Mmxx cos( x=1) )
where Mmax is the maximum learning rate and the hyper-parameter p is usually chosen as 0.1 such that the minimum learning rate is 7/max/ /0mx 10 [60].
Sarmup-Stable-Deca (WSD) Schedule [69, 21, 171. The schedule consists of three phases: warm-up phase of K warm-up steps, followed by a stable phase maintaining learning rate 7/2 1/max, and finally a decay phase governed by `) h(k - K Mmax for Katable < k < K, where Katable represents the total duration the first two phases. Here, the decay function h(-) (0, can be linear exponential.
Multi-Step Schedule [5]. The entire schedule is divided into S stages, i.c. [Ko K1] U [0, K], where The schedule satisfies that = Л), for K;-1 k & K; (1 & & 5). In our LLM experiments, we consider 8- 1-1 LRS, corresponding to the case where S with
[K1, K2] U... 5K5- Ks]
Kg K1
Ks K.
/max/ V10, and nK, Mmax// /0axx/10, 10, and K₃ 0.8K, K2 = 0.9K.
μK₁ Mmax 7K2
A.3 Connections Kernel Regression
In this section, we explain how our setup in Section are equivalent to learning with kernels.
Definition A.1 (Positive semidefinite (PSD) kernel). function K X x - R is called continuous positive semidefinite (PSD) kernel if it satisfies:
Symmetry: K(x, x') = K(x' x) for all x, x' E X:
Positive semidefiniteness: for any XI,..., Xn X and ai,..., ...,...,, an E R
0:0y K Xi) > O.
Definition A.2 (Reproducing kernel Hilbert space (RKHS)). Given kernel K, the reproducing kernel Hilbert space HK associated with K is Hilbert space of functions f X' - IR such that
(f,K(x, )) f(x), Vf HK, x X.
Kernel methods learn functions from a hypothesis space defined by the associated RKHS. For instance, kernel ridge regression gives estimator:
fx
arg min
(f(x)) - 11)2 + 1
JEHK
==1
Hence, model capacity is determined by the size of the RKHS Hx-
Let D be the input distribution. Given kernel K, define the associated integral operator Tk (D) - (22 by fx)()) Ex~ ((((x)))
17

---

By assuming (x,,x)((xxx)) < co, the operator TK is compact (Mercer's theorem) and consequently the kernel admits the following eigenvalue decomposition
K(x,x) - Σ (x)((x)))(X))
where (Ay)ī=1 and {e,1ʃ=1 denotes the eigenvalues and eigenfunctions, respectively. Moreover, (ei, ()) j. i.e., the eigenfunctions form an orthonormal basis of (D)
Using the spectral decomposition, the RKHS admits the following representation:
HK
ojej
j=l
To better quantify the smoothness of functions, we often consider the interpolation space Hi with $ > defined as
Hi
ajej
Clearly, Hx- = Hk. and
H C HR
V8 82 0.
Hence, the index characterizes the smoothness of function relative the chosen kernel.
In the analysis kernel methods, the following conditions are commonly used to describe the smoothness of the target function and the capacity of the kernel, respectively.
Assumption A.3 (Source condition). There exists some > such that f Hxx-
Assumption A.4 (Capacity condition). There exists some B >>> 1 such that A j-
These conditions yield the following interpretation:
A smaller s indicates that the target function f" belongs to larger space, corresponding to a more difficult learning problem.
A smaller 3 implies slower eigenvalue decay, meaning richer hypothesis space H. and thus higher model capacity.
Our formulation in Section is equivalent to the above setting, but expressed in terms of the feature map Φ. Under Assumption 2.2, have
(') Σ (x)x)))) Σds (x)
In this case, Assumption 2.3 corresponds exactly the above capacity condition, while the task- difficulty assumption in Section 2 can be viewed as power-law version of the source condition. Specifically, under Assumption 2.4,
07
Σ
j
1/2
0j aj Zj.
j=1
Hence, for any arbitrarily small (0, 1), we have S Hxx since
N Σ
- 10 < 0C.
18

---

A.4 The SDE Modeling
The physical-time SDE. In our setup, the SGD update can be written as
K++ V& ip VR(va)) srghEx-
E(v+). By assuming the gradient noise to The term &a is the gradient noise, whose covariance is Bz be Gaussian, the SGD becomes
Vi+1 Va = - N (0, A E(v))
B,
It is exactly the Euler-Manuyama discretization of the Itô-type SDE:
h dv, (R))V p(T))) b(T)
((V))dB.
where B, RM denotes the M -iddimensional Brownian motion, and (p(),, b(-) are the continuous version of LRS function and batch-size schedule function, respectively.
The intrinsic-time SDE. Intuitively, the discrete update (4) can be viewed the Euler-Maruyama discretization of SDE (7) on the non-uniform grid {1x Eド=o Mi JeeN where the effective step size
Ale
T)k
V2 VVR(v))(te+ tx) vt+11 tx N (0, Br I(v))
V&+1
A.5
The Emergence of Power Laws
We illustrate how the power law emerges our setting from multi-task learning viewpoint. For brevity, consider the case of the top-M features and an infinitesimal learning rate, where the
SDE (7) reduces to the gradient flow ODE: dv, = - WHv, dt. Noting that W WH (Xi, A2, ... AM, O,... 0} is diagonal, consequently the ODE is solvable and gives the follow- ing expression of the excess risk:
M
R(v)) - 00202 ~
A 10 22 22,8
Σ
A 071.
j=M+1
learned sub-tasks
unleurned sub Lasks
Intuitively, we can view the learning of each eigenfunction as a sub-task. Due to the limited model size, student model can most learn the top-M eigenfunctions.
(i) Intrinsic-time power law. For each sub-task, the sub-task risk converges exponentially w.r.t. the intrinsic time However, owing to the power-law structure of X , the total multi-
M task risk exhibits a power-law decay for sufficiently large M due to Σ
j==1 Xy10312e 22,t
1 f - 2ut du = ea if M > 1.
(ii) Model-size power law. Approximation error accounts for total risk of the N unlearned sub-tasks, which follows power-law decay due to E/= Ay10912 ~ M x8 if N M
Summary. The emergence of power law arises from the accumulation effect, requiring both the number of learned tasks and unlearned tasks to be large (ideally infinite).
B Related Work
Theoretical explanation of scaling laws. Among the growing body of work seeking to theoretically
explain scaling laws [54, 22, 39, 63, 23, 42, 43, 2, 14, 3, 7, 35, 50, 8, 46), the most closely related are [7, 50, 8, 35), which also analyze PI.K regression (often written in the equivalent linear-regression form). Specifically, [7] studies gradient flow, [50, 8] analyze SGD with constant LRS, and [35]
19

---

considers an exponential-decat LRS. In contrast, we establish a unified scaling law applicable to general LRSs, which not only recovers these prior results as special cases but also substantially extends them by capturing the loss dynamies rather than only the final-step loss. This unification is enabled by introducing the key notion of intrinsic time, which more faithfully captures the effective training progress than the raw number of training steps.
Predicting loss trajectories in LLM pre-training Recent empirical studies [59, 38] have shown that the entire loss trajectory of LLM pre-training-not merely the final loss-can be accurately captured by suitable scaling relations. A detailed description of the corresponding fitting procedures is provided in Appendix A.1. Our theory offers a theoretical explanation for these empirical findings. Interestingly, the multi-pon powli--pwwr--aw (MPL) model proposed by [38] is closely connected to our FSL: through an integration-by -y-rarsss transformation, the FSL expression can be recast into a form that is nearly equivalent the MPL formulation (see Appendix C.2).
Warmup-Stable-Deeay (WSD) LRS. A WSD schedule [69, 21] maintains a constant learning rate for long stable phase, followed by a learning rate decay only near the end of training. Although unconventional, WSD has become popular in LLM pre- tre-training [21, 17] and is already deployed in training industry-scale LLMs such as DeepSeek-V3 [36] and Kimi-K2 [58]. Yet its mechanism remains poorly understood. While recent works [64, 53] offer partial insights, we show-perhapp surprisingly- that even quadratic optimization, corresponding to kernel regression problem, already reproduces the essential advantage of WSD. Furthermore, we quantify this advantage through explicit comparisons of scaling efficiency against constant and exponential-des schedules.
C Experiment Details and Additional Results
In this section, we present the details of our experiments as well as additional results.
C.1 Power-Law Kernel Regression
Physical-time FSL. The FSL (10) is presented in terms of intrinsic time, but in practice, it is often more convenient to use physical time (training steps). By suitable change of times, after 7 steps (equivalently, r/h discrete steps), the FSL maintains the form (10), with adjustments:
1- = =(()
((v,b)) 1' (C(((T)) T(w)) (e(T(w)) + +z2 hu((u)² b(u) du.
Fitting FSL on SGD Average-Risks. To validate that the Functional Scaling law (FSL) can accurately capture the risk curve of SGD, we conducted a series of SGD experiments under different configurations of and 3. Subsequently, we fitted the FSL to these risk curves. Our results demonstrate that FSL indeed provides a close fit to the SGD trajectories.
In each experiment, we adopt a PLKR configuration with M N = 128, a = 3 and employ the top-M M projection matrix, thereby eliminating the approximation error term - We explore range of values for s E [0.5, 1] and B E [1.5, 5]. encompassing both easy- (s > 1- - 1/5) and hard-learning (s < 1 - 1/3) regimes. For each parameter configuration, we execute 200 independent SGD runs with a batch size of over 10,000 steps. The resulting average trajectory across these runs serves as the fitting target. The FSL. fitting is performed using the physical-time FSL formulation.
k
c, + 02 Σ K((())) T()))e -(T
CzO² Σ ((( -(())n?
where Ci. C22 C3 are constants to fit, (m is the learning rate schedule, and T(i) = Eド 'lj
When fitting the SGD trajectory, we minimize the mean squared error (MSE) between the empirical risk trajectory of SGD (without the irreducible risk d² denoted by (k) and the theoretical prediction from FSL, & Formally, we solve the following optimization problem:
min C1,02,63
(EsOD (k) Ex)
k=1
20

---

where K represents the total number of training steps, This minimization is performed using ordinary least squares (OLS), with the integrals in the FSL expression & evaluated numerically via quadrature methods.
We display the learning rate schedules (LRSs) used in the SGD experiments in the top-left panel of Figure 3. Complementing Figure (middle and right), additional experimental results for various values of and are presented in Figure
LBSs troining
0.5.P 1.5, 125
0.5. A -3. 125
0.31
3,002
wod
Wod
Lantning
0.000
Tbol
0.05
2.034
0.05
0.1-
0.002
4yyCle
6.6-
2000
$0.00
5000
30000
3000
(000)
(2401
S000
10000
2000
6200
8300
10000
Step
0.5. 5, M
128
0.8,
128
088 5. 5.M
cyclle
Ri4
11.44
1ttik
Btak
0.022
1116
6.0-
vod
0.01
W0d eyclic
0.01
S000
0260
Step
Step
Stop
Figure 3: Fitting results of FSL SGD trajectories. The shaded curves are the average over 200 independent SGD runs, while the solid curves show the predictions FSL.
Scaling law experiments. These experiments are designed to evaluate the correctness of the scaling laws predicted by our analytical analysis. To this end, we conduct two complementary sets of experiments:
FSL experiments. This experiment is intended validate the theoretical predictions derived from the FSL. We compute the predicted risk by numerically discretizing the FSL (10), with all untracked constants set to 1. For each LRS, following the theoretical analysis, set 7max = 0-05 where = 8/ s/(1 + 8) for the constant learning rate schedule, and r r= for exponential decay and WSD schedules. We fix the batch size B 1; thus, for each data budget D, we compute the intrinsic time and evaluate the final-step loss using the discretized FSL.
This experiment serves to assess whether the scaling behavior predicted by our continuous-time FSL faithfully captures that of discrete-time SGD. We simulate stochastic gradient descent (SGD) with 200 independent trajectories and a fixed batch size For each data budget D, the maximum learning rate set as 7max 0.05D -r using the same theoretical values of as in the FSL experiments. We run SGD for D steps under each corresponding LRS and record the final-step excess risk.
SGD experiments.
C.2 LLM pre-training
Practical FSL Ansatz for LLM pre-training In this section, in order to fit real LLM pre-training loss curves, we will derive an approximation form of the FSL in Theorem 4.2.
First, by the physical- time for of the FSL (10) with h and B(u) = B, we have
1
l K(T(K) T(=)) + e(T(u))) p(u)²² du.
E T(K)* M +
B
Here focus the integral term. Since p(u) .h" ('()) dr + (00), and that T(k) KC((k)) - t) + +e(t) dt. (14)
S K(T(K) - T())) ( (()))(()) du = v(0) /
21

---

Note that this is exactly the SGD noise term at the constant LRS 17(0) for total intrinsic-tine timinssii-tmeeeeee T By results of constant LRS (as seen in the proof of Theorem E.1), we have
(14) ~ +0000 e(T(k)))
As ((0) S 1, we have
& T(k) + M LRD(A).
(15)
where
1
I (IT
- +()ll I. d()'(r)
LRD(k) ==
B
B J. !'" ( [. [ du
1
1
[." (r) (m K((((k) 1) + e(1) dedr.
((()
B
We discretize the outer integral integer nodes y =0, k.
1
((k)
LRD(K)
Σ Mh-1
1n)
K(T(()) - t)(o² + e((t) dt.
B
T(i)
By the integral mean value theorem, we can take (o² + e(t)) outside the integral, which gives
k E
T(k)
LRD(K) =
- (m---1))(-
e(Ei))
K((((:)) - dr,
B
T(
1=1
where & [T(1),((K)) Now since
T(k) KC((()) t) dt
T(K)
1
dt =
T() (1 + 21)--188
(1 (1+) -(1))
T())
then further simplify as
LRD(K) =
5 -
B
1=1
Here, we approximate & as T()) and introduce new parameter to replace for simplicity.
Therefore, combining with (15), when the batch size fixed, after renaming some constants, the final discrete ansatz can be written as
R2 x +
C1 T(k)*
+ 888
(16)
1 C3 (T-11 mn) (ctt T(T))) (1- (L=()
where co. 61, C22 C3 Ca. 65. 8, 3, 7 are constants to fit.
Fitting the Practical FSL The objective of this experiment is to analyze and fit the loss function using our funetional scaling law, by (16), since we do not explore the effect of varying the model size M in our experiments, we drop the term M - and get
(o(k) Lo T(k)*
C1
LRD(K)
where T(k) = Ei= T/ and
1 (1-(+C(((A-TTO))
LRD(K) = 02
(mi-1 1h))
T()+)
22

---

and θ = (Lo, c1. 02, C3 8,
Following [59], utilize the Huber loss the objective function.
K
min
Hubers (log Ce log (g((k)),
where 8 10-3, Leg denotes the ground truth of the validation losses. We adopt the Adam optimizer, with learning rate of 5 x 10 for the index parameters in our law and 5 10 for the coefficient or constant parameters. Each optimization takes over 10,000 steps.
We fit the law on the 400M model and 1B model trained with 20B tokens and an 8-1-1 LRS We then predict the loss curve for the 400M model and IB model with cosine LRS and WSD LRS. The experiment result is present in Figure 2a.
FSL-optimal LRS via numerical variation. We propose to obtain numerical optimal LRS by directly minimizing the final-step loss over the space of LRS using the fitted FSI., termed FSL-optimal LRS.
Step 1: Fitting FSL. Fit FSL on the loss curve of IB QwenMoE model trained on 20B tokens with batch size 288, maximum learning rate yo 0.001, and the 8-1-1 scheduler over total step of K = 33907, following the same procedure described earlier.
Optimize LRS. To improve optimization stability, we reparameterize the learning rate schedule by defining
Step 2:
6, = 7 for i 0, 1..,..... ..... K - 1.
Then, the i-th step learning rate can be recovered by 7/1 = 7/0 Σ =g 5g, which defines a one-to-one correspondence between the learning rate schedule {nH and (
The optimization problem is
R-1
min Le ((m))a)) (6.)
subject to Σ
8n < no, drp > 0, - 0,1, K 1.
(17)
A=
To solve the above constraint optimization, we use the projected gradient descent (PGD) [10]. The learning rate of PGD is searched ranging from 10-8 to x 10-10, and the optimization step number ranges from 50,000 to 100,000.
The resulting FSL-optimal LRS is presented in Figure 2b (left), where cosine, WSD, and 8-1-1 LRSs are also given for a comparison.
Step 3: Evaluate our LRS. We then evaluate the performance of the resulting FSL-optimal LRS, and the three LRSs in Figure 2b (left) are used as baseline. All comparisons are conducted on the same IB QwenMoE model under identical training conditions: 33,907 total steps. batch size 288, and 20B training tokens. Full loss curves are shown in Figure 2b.
Additional Experiments We have further conducted ablation experiments with different model sizes and architectures, different total steps and different WSD schedules.
We validate our functional scaling law in models with various sizes, ranging from 100M to 1B, and diverse architectures including GPT-2 [52], LLaMA [60] and QwenMoE [68]. For each model, first fit the FSL using the 8-1-1 LRS and subsequently employ it to predict the loss curve under a WSD LRS. Next numerically solve the FSL.-ppimal LRS and empirically validate its efficacy by comparing the final pre-training loss against those obtained using other commonly adopted learning rate schedules. We present the results in Figure 4 for the IB LLaMA dense model, Figure 5 for the 100M GPT-2 dense model. The consistent alignment between predicted and observed performance across architectures and sizes underscores the robustness and generalizability FSL.
We further validate the applicability of our functional scaling law (FSL) across varying training durations. Using 100M LLaMA dense model, we conduct experiments with total training steps set to 17k, 34k, 68k, and 134k. As demonstrated in Figures and 7, our FSL accurately models the loss trajectories across all evaluated step counts, confirming its robustness to different total training steps.
Finally, we conduct comprehensive empirical comparison between our FSI--optimal learning rate schedule and various WSD baselines, examining different decay ratios and minimum learning rate
23

---

L.000
1216 Onr
34
14
1276 dor
KSSD
24
Cpt
500
380
15Kg 2X0
2538
3008
3506
5000
LOWE
150M 2000
2508
3000
3008
0999
2000 2001 20X
(a) Prediction on WSD LRS
(b) Optimal Existing LRSs
(c) Loss curve LRSs
Figure 4: Experiment on the 1B LLaMA (dense) model. Figure (a): We fit our functional scaling law on the loss curve IB LLaMA (dense) model with 20B tokens training data and 8-1-1 LRS. Figures (b)(c): The comparison on the IB model between the optimal LRS, cosine LRS, WSD LRS with exponential decay and 8-1-1 LRS.
28
2.00
2771
2.00
27X
270
WSD
245
24
280
280
13:02 15003
2000
1000 1000
2000 19003 2100
Prediction on WSD LRS
(b) Optimal and Existing LRSs
(c) Loss of LRSs
Figure 5: Experiment on the 100M GPT2 (dense) model. Figure (a): We fit our functional scaling law on the loss curve of 100M GPT2 (dense) model with 20B tokens training data and 1-1 LRS. Figures (b)(c): The comparison on the 100M model between the optimal LRS, cosine LRS, WSD LRS with exponential decay and 8-1-1 LRS.
configurations. As evidenced by Figures and 9, our FSL-optimal LRS consistently outperforms all WSD variants, achieving superior final pre-training loss across all experimental conditions. This systematic evaluation demonstrates both the effectiveness of our theoretically-drrine schedule and its practical advantages over conventional heuristic approaches.
17k. R. OM0899
0.09993
17k
34K, RMSE=0. 0.1484
34k
68k, AMSE=0.11104
3.0
134k, RMSE=0.CO 0.08559
MIIS=00.055S
134k
Lo65
2.9
2.6
2.5
Stops
Steps
(a) Prediction result
(b) Optimal LRS
Figure 6: Experiments with different total steps. Figure (a): Fitted functional scaling laws on 100M LLaMA model with different total training steps 17k, 34k, 68k and 134k (corresponding to 10B, 20B, 40B and 80B tokens respectively). Figure (b): Optimal LRSs compared with cosine and WSD LRSs. The solid lines are optimal LRSs, and the dashed lines are cosine/WSD LRSs.

---

2.00
240
COPN
200
250
4000
5004
(a) 17k steps
(b) 68k steps
(c) 134k steps
Figure 7: Experiments with different total steps. We compare loss curves of existing LRSs and
optimal LRS on the 100M LLaMA model with different total training steps 17k, 68k and 134k.
280
CC0V. Coden
2.55
4001 (-13
400V, (4-)
LoeeENA)
2.00
2.43
2.48
2405
28X
2209
34000
28:00
350.0
5000
X20X
$4.0
Slep
Figure 8: WSD with different decay ratios: We train 400M LLaMA (dense) model with 20B tokens of training data and WSD LRSs with the ratios between stable time and decay time of 3:1, 4:1, and 5:1. All WSD LRSs exhibit final loss similar to that of the Cosine LRS, and the optimal LRS derived from our functional scaling law outperforms all other LRSs by loss gap of approximately 0.01.
2330
1899.
Optinal
(950
LoseERA))
2.67
2.580
2.65
2.00
3030
Shee
Figure 9; Comparison between optimal LRS and WSD with near-zero final learning rate: We train a 100M LLaMA (dense) model with 40B tokens training data and various LRSs with the same Tmmax 10-3 including WSD LRS with Mmmi To /max, WSD LRS with 7/min 10-7, cosine LRS with 7/nin 10 7maxx 8-1-1 LRS with Mmin To mmxx. and optimal LRS. The experimental results show that decaying to (near) zero does not result in significant loss reduction.
25

---

C.3 Key Proof Steps and Core Insights
In this section, we outline the main ideas behind the proof of the FSL (10), highlighting the key techniques. Complete proofs of the above theorems are deferred to Appendix D.
We will need the following characterization of the gradient noise structure. Lemma C.1 (Noise structure). For any V E RR,, , it holds that
(2p_E(v) + +2) v²R(v) I Σ(v) I (2p+Elv) + oF) V²R(v),
where V²R(v) WHWT, and the constants p- and p+ are the same as in Assumption 2.1.
The proof is provided in Appendix D.I. Let E(v) denote the gradient noise at Since R(v) ((v + o it follows that for any direction n E SM-1
E ( T n = n E(v)n ~ R(v)n (R)())n,
where n (²R(v)) n represents the local curvature of the risk landscape along n. Hence, the noise energy in each direction is proportional to the product of the population risk and the curvature along that direction. This anisotropic structure of the gradient noise scaling with the risk and shaped by curvature-has also been reported prior work [67, 62].
For clarity, in this section, we focus on the case of top-M features, for which the population risk takes the form
M
28(v)
Σ
07)
Σ
2,10712
(18)
j=1
j=M+1
For the intrinsic-time SDE (7), each coordinate of vi evolves as
M
du,(1) = 07) V 7()
E(vi) (B(
1=1
where ((() is the k-th coordinate of the M-dimensional Browian motion B..
Observe that the summation of white noises is equivalent to a single stochastic term (t) (ty
with variance 9, (t) : Ex= /v ((e)) / E(n)e,, 11 e) E(v))e ej the j-th canonical basis vector for j [M]. Therefore each coordinate D; satisfies
dvy(t) = -,,,, - - dt + v (t)
where (1()) = e, E(v, )e, is the variance of gradient noise along ey. By applying Itô's formula to (v, 0g)² and noting v00) 0, we obtain
(1(r)(1)) - 07) (0 - 0 22,t f -2 (1)
Let & E(v,) and plugging the above equation into (18) gives
M Σ >3 f-" j=1
2B[8.]
Σ e 22,e
Σ 2,1071 (19) J=M+1
By Lemma C.I and noting V²R(v) = diag .A....., Anc), we have
(t) 0 E(vi)e, R( (y(E(()) + o² 02/2).
Let δM = -j M+ A 1, (t) Σソا 1,10712e-22;; and KM = EMI 33e-22,t 2d,t, Plug-
ging them back into (19) gives the following Volterra equation:
B B = 8M (tu(t) L' KMu(t 2){(3) +z²)dz.
(20)
26

---

The above equation characterizes the expected loss dynamics of SGD under general spectrum and has been derived in prior works such [47. 48, 28, 49, 50]. Our key observation that, under the power-law assumptions on {0}), and (Aj}, (Assumptions 2.3 and 2.4), the solution to (20) admits sharp asymptotic characterization, providing explicit upper and lower bounds that precisely capture its scaling behavior.
Let (() == [[E,), g(1) 5m +eM(1) + o² S0 ( 2))(2) dz, and define the linear operator
Tf(t) f[ (t dz.
Then, the Volterra equation (20) can be expressed in the compact form f g Tf. Formally, its solution can be expanded as an infinite series:
f = 1=
(21)
The key observation is that, under Assumptions 2.3 and 2.4, the higher-order terms T for & can be well controlled by the first-order term Tg. This is due to the scale invariance of the forgetting kernel:
Lemma C.2 (Scale invariance). It holds for any S M that
(t/(1/2) 1. Kx(())
Proof. The result follows from the fact that KM exhibits a power-law decay for t M (see (25)). Consequently, (t((t/) ~ (t/2) -2--1/8) t (2-1/8) KM)(t,, which establishes the claim. O
Corollary C.3 (Subconvolution property). For any t S M it holds KM * Kas S Kas (t).
Proof. Noting KM is non-increasing and integrable and applying Lemma C.2, we obtain
(KM * (t)(()= ( 12 Kt - (z) + ("d Kui( - (x) da 1 < 22Kr(112) / (x)(=)dz $ Kmrll22 z Km(()),
which proves the claim.
Let l/v supezo (t)). By Corollary C.3, it holds for any S M that
(2t(t) < 1171 Lf -
x 17 S ((M((=))((e)(((dd=== |l>l)xx
When I sufficiently small, there exists a constant 0 < < such that T gt) (Tg)) holds for any 0 < & M Hence,
f(t) < g(t) + Tg(t) Σ
+ Tg()).
Combining f(1) > (() + Tg(()) with the above upper bound, we conclude f(1) = g(() + Tg((), i.e., we have
E R(()) - 102 = 8as + (t) + J'' Karlt - 2) xz) (2)) dz.
(22)
27

---

The emergence of power-law scaling, We now illustrate how power-law scaling arises in our setting from a multi-task learning viewpoint. Intuitively, we can view the learning of each eigenfunction as sub-task. Due to the limited model size, student model can at most learn the top-M eigenfunctions. Within this view, our FSL. framework reveals three distinct manifestations of power-law behavior:
ii) Approximation error. Approximation enror accounts for total risk of the N - unlearned sub-tasks, which follows power-law decay due to
N δm Σ J=M+1
~ M 83 if - M 2 M.
(23)
(ii) Signal learning, For each sub-task, the sub-task risk converges exponentially w.r.. the intrinsic time t. However, owing the power-law structure of {Ay}. [e)] the total multi- task risk exhibits power-law decay for sufficiently large M due to
M
11 2Ast Sr u--te 2ut du if
1 15
(t)
Σ
(24)
A=1
iiii Noise forgetting, Analogously, for the forgetting kernel, we also have
[ z --- -2st dz (2-1110) if if M. (25)
KM((1)=
X2e 22gt
M-8
Remark C.4 (Task Accumulation Effect). The above derivation reveals that, although single task may not exhibit an exact power-law scaling with respect to intrinsic time, the collective accumulation of many tasks gives rise to such scaling behavior. As the number of tasks increases, the power-law regime progressively extends, and the idealized limit M - oo, it spans the entire training process. In our setting, each sub-task naturally follows exponential learning curve cxp(-At), yet the underlying mechanism appears more universal. We anticipate that similar power-law behavior may arise even when individual task dynamics deviate from this form, which we leave future investigation.
D Proofs for Section
D.1 Volterra Integral Equation Governing the Loss Dynamics
In this section, we derive a Volterra-type integral equation that exactly characterizes the evolution of expected loss under the intrinsic-time SDE. This equation serves as the starting point for all subsequent theoretical analysis. Recall the intrinsic-time SDE:
du = (R(u,) + 7: E(v)) dB.
(26)
where we write = >t) for simplicity.
By the definition of R(v), we have VR(v) WH WT v - Let un W W. 0" Then, we have
& = = u H
To obtain the estimate of E,, we consider the intrinsic-time SDE for u, given by: Lemma D.1. We have
du, -w WWwHH, dt + V% W E.W
(27)
where Σ, = D(v)).
Proof. By Eq, (26),
du, = d(W v, -²) W de, = - WWHHu, dt+ W VYE,dB
28

---

Here B, is an N dimensional standard Brownian motion, we are going to replace it with an M dimensional standard Brownian motion B,,
is easy to see that the diffusion term WT T has the same distribution as W WTW, dB,, hence the SDE can be written in B, as
du, W wwu, dt Vx W E.W dB,,
A key insight for tractability is that the gradient noise exhibits the following anisotropic structure: Our analytic analysis also relies on the noise structure characterized by the following lemma.
Lemma D.2 (Noise Structure). For any RM,, it holds that
E(v) + o²) WHW T 3 Evv) 3 (2p+E(v) + + WHW
Noting V²R(v) = WHWT and R(v) = {vv) o0" this lemma means E(v) ~ R(v)V>R(v)) That the gradient noise scales proportionally with the population risk and aligns with the local curvature. Notably, the noise has two distinct sources: (i) the fit-dependent term & E(v), which arises purely from minibatching and persists even in the absence of label noise; (ii) the 0² term, which captures the contribution from label noise. This anisotropic structure of SGD noise - scaling with risk and shaped by curvature has also been observed in prior work [67, 66].
Proof. Noting f(z;v) = ( Twb(x) - y) we have
Vi(z; = (x) (WTv v 0¹) -Wo(x) VR(v) B[Ve(z v)] WH vWTv-0" -0")
Hence, the covariance matrix of the noise = V((2.v) VR(v) is given by
Σ(v) = E'gg" |v] W (E [x(x)(((x) T un ((x)(())( Hm H) W +
Noting
E (xx())((x) T uu (x))((x) Hun H = E [xx()] un (xx))(())((x Hun HH then applying Assumption 2.1, we have
((v) = p+ Wir(Huu JHW + WHW E2p+E(u) + + WHWT,
where the last step follows from tr(Huu || 2 2E(v). The lower bound follows the same proof.
The excess-risk dynamics is then given by the following Volterra integral equation:
Proposition D.3. For the intrinsic timeiicciimeeeee SDE, we have
2B = Uo A, HA, + L ir(SA HA, (c,,
(28)
where A=:= c- -W WH: S:= W WWW W, and e, 12p=,2p+] for any 7 2 0.
Proof. By Itô's formula,
W W WH: u,) eew wHt (W WHu, di + du,)
wew WH 7 W E, dB.
Integrating both sides, we get
p
WHI u, Ug =
W WWT
WH- V 7r W Σ, dBB.
29

---

Now write = W WHr we have
u = U + W Σ, WdB
Note that the integral with respect to B; always has zero expectation, therefore we have
2 (u Hu,)
=ug A, HA +E S 7
dr.
By Lemma F.2 and Lemma D.2, have
tr (VW E, WA, HA . W W & (2p+&, + (S ,),,
tr (VW E, WAJ HA W E, W > (2p_E, (O) S). (VWEE,WY
Hence there exists some constant O E 12p_, 2p+] such that
Etr (VWTE,WA) W -A W E, W) = (") HS,,,
from which the lemma follows.
D.2 Power-law Decay of the eM and KM Functions
In this section, we rigorously establish the power-law decay behavior of the functions eM and KM- Lemma D.4. For eM (1), KM (1) we have the following estimation:
(t (t) (t--1)))
Moreover,
etr(t) =t (2 1 M
Proof. Recall that
eM(1) Σ
10
22,1 22,1 22
j=1
Since Xo = j- there exists constant such that
A マンー = 22xt < 1 when j > (2ct) $.
We have
(2ct)
M
(t) ~
-221t 1- 11188 +
Σ
22, -22
1=1
j=[(2et à 1+
$ (2ot) maxe -2,
Σ
1-
==((2ct) 3 |+1
Now we bound the two terms separately:
max -
2A, (1)))+t tt <
supe St- e>0
where the last inequality is because sup
2x + is a constant only depending on and 8. uu
We have
0
Σ
j-1-
x-1 11-8 dr [(2ct) =t=
)=((2ot) 7 [+1
30

---

Therefore we have for all t > 0,
ess(t) S (2ct) (2ct)* x
On the other hand, when M 2(((c))) i 1]], i.., M
[(a)) 51
eM(1)
Σ
2 2
Σ
)=((2ct ò 11+1
M Σ
2 --
1-
)=1(2ct) |+1
21(20))
x 1-8 dx
1120))++11
(((2ct) + +1)== (2[(2ct) $ + 1)
2 t where the last line requires 2 1. Therefore we have shown that eM (t) = t when 1 155 M
The proof for K. (t) follows the same way.
Lemma D.5. We can bound the gap between eas (t) and eac (t) as follows:
Coo (t) eM (t) S -
As a consequence, we have
cM (t) -8 ~ co0 (t) + M
Proof. We have
Coc (1) - (t) Σ
2d,t X3 1871
M+1
8
1-88
M+1
T -1
dr ~ M-88
M
D.3
The Case of Top-M Features
First, we prove the general label-noise case of FSL (Theorem 4.3). Applying this result, we then derive the constant label-noise case (Theorem 4.4), the noiseless case (Theorem 4.5) and the hard regime case (Theorem 4.2)
D.3.1 Proof of Theorem 4.3
In the top-M feature case, the matrix W satisfies W, = ej for each j E [M]. therefore we can simplify the equation for Eig,] as follows.
Theorem D.6 (Volterra equation the M case). In the top-M case, we have
-3 23 + (t) + ('" Kas(t - T)r (E[E, ddd dr,
E[E]] = M
(29)
where the function ex and Kas are defined
etr(t) = Σ 2e 2dst, kM((): = X3e 22yt
22gt
(30)
31

---

Proof. By (28) in Proposition D.3, note that W WH = HD:M E RNN is the top- M part of the
matrix H. ie. Ho:M diag(>r ......, Ans. ......0) 0}, we get
B ~ un He M Un + / (r(Ho 2H Jr (1[E)) + o²)
which can be further written in terms of the eigenvalues {A₂} as
M
M
B[E
Σ
22st +
Xy(u
Σ
23e
2
[B[& + d
J=M+1
Note that ug the j-th component of uo, is equal to 03 because of the zero initialization of Po- Therefore by the definition of eM and KM we arrive at the Volterra-type integral equation of E[
Lemma D.7. For the forgetting kernel KM and & MD, there exists constant C independent of such that
Kas KM(1) & CKM((), VI & M
where denotes convolution:
KM * KM)())= = [[M((T)(N((t-)) (t)))
Proof. Observe that KM (1) a monotonically decreasing function. By the symmetry of the convolu- tion, we can write
I (M))))((((T)) dr
KM * KM((()
=2 1/" (M dr.
2//
Since KaM decreasing, for 0 < t/2 we have Kas(t 7) (ar(t/)). This observation yields the upper bound
KM * KM & 2kM / KNr(T) dr & &CM (f),
where the constant finite and given by
c=2 1 C = C
000
(T((T) dr <
Xze 2dg7 2dy7 dr = Σ X = tr(H) < co.
j=1
1=1
remains to show that when & < M
1 2
Kas
(()
for constant C' > 0. Recall that by Lemma D.4 we have
12-
Kar(t) =tr(t)=++00 (2--) KM 5
t
2
Therefore when 1 t S 8 the conclusion follows. When t S 1, the function KCr(() is decreasing and (t)(t) < E X3 < oo, we can always find a constant C" satisfying the condition.
Combining the above estimates, we obtain
(Kar (t)(()) < C.M (1) =M) C'Kar(I).
The proof is complete.
32

---

We now prove Theorem 4.3.
Proof. The lower bound is trivial by B|E,] > (t) and the Volterra equation (29):
B| M ttv(() + /
dr
2 M + elr + /
o dr.
For the upper bounds, we first prove slightly stronger bound with Koo, that is,
Proposition D.8. For all > 0, > 0 and B >> 1, we have
E S ---- + (t)(t) + S:
-o)) dr.
(31)
Proof of D.8. By Equation (29), have
E[ x M 2 + (1) + ( - ()7 +
< M (cM + ( Kcc(t-T),
dr.
Define f(1) == J0 ()) (E[Er) + o²) dr and Tmax sup, 7(), substituting the above
inequality into the definition of f(t), we get
f(1) 2 S
+
1 Kx drr
S Ymmax -8 + (' (C)((t))(em( do?
+ Ymax S. () [L)
-mmax M - + ( Koolt ( +
+ Tmnax (Kioo - ()),, + dr dr
Lemma D.7
Ymax 68 S (t o dr Ymax (t)
Therefore when 7max is sufficiently small (Ymax < [n|H] for some absolute constant c), the constant factor of f() on the right-hand side will be less than f. hence may subtract / from both sides and get
S [' $ Tiaax M- + S' (c((t))((MM(1)) o
Therefore substituting this back to the Volterra equation yields
E[E $ M- (t) + S' ((em(t)))
()( + +0? dr.
33

---

The preceding proposition relies exclusively on two properties of K the identity established in Lemma D.7 and the finiteness of its integral, _o_ (t) dt cc. Since the kernel KM also possesses these properties for & M the derivation of the upper bound
B E M - (t + L (u(t-r)(eMe)) + 0² dr, t & M
follows analogously by substituting Koo with KM..
D.3.2 Proof of Theorem 4.4
Proof. It is clear that, by our assumption, sup, S 1, and that
M
M
(t) ~ j
5-11-88
2x²t
1-88
x8
~
1.
j=1
j=1
By Proposition D.8, we have
B| $ M (t) /
((a(t-))(e(((++o0))
(32)
11+ [( ( ( + dr oo?
(33)
where the last inequality is by 0² 2 1 and that
I (T) /
X3e 275T dr Σ Aj 2
Therefore by the Volterra equation (29), note that W[Er] + 0² x 0² (M((T 0², have
B[E M- + et) + ( Kar(t Kar(t- () o dr.
D.3.3 Proof of Theorem 4.5
Proof. When の² = 0, by FSL in general case (Proposition D.8), we have
E $ M- (t + / Koo ()e-))em(il (T) dr.
In order to get the upper bound with Kas, we first prove lemma.
Lemma D.9. We will bound the gap introduced by Koo and KM:
S. (K - -KM T))em()))))))))) dr S Ammaxl 188,xx-68,-20----
(34)
Proof of Lemma. First note that
2
- (t)))) Σ
X3e 22g1
Σ
j J-29 -
M+1
j= M+1
Therefore we can bound the integral as
S (KCooll
- (K<o(l-T)
$ M 2-2-1 S- eMr(T)+rr
34

---

/ (nr(T)
2841 M2281 mmax M
Ymax M M- 28+1
e -2,T dr
==1
M
Yimax M M-88+
143
Tinax 1
-2441 S Tmax Mmaxi Aymaxi-oB. -08,-29+11
By < 2 1/2 and 7max is sufficiently small, combine the upper bound with (34), and directly conclude that
=8 (t I Kar(t (T) dr.
E[E,) S M
Now with the lower bound in FSL Theorem 4.3, the result follows.
D.3.4 Proof of Theorem 4.2
Proof. In the hard regime s < 1 8. by Propostion D.8, we have
E S M 88 + (t)( + S Kox (TM + + dr.
Lemma D.10. holds for < - y/2 that
Le (Kooll - 7) - KM((t dr $ M-
(35)
Proof of Lemma. First from the definition of K. we obtain
(1x - Kar(1)= = Σ
X3e 241²
f=M+1
Therefore, we have
L (Koo - T) - -MM (t: - -T))r dr
p00
M
Σ
Xze 22,7 dr
Tmax
f=1
Ymax
j $ Mmax
9+1 -8+1
1=1
By 8 < and that Tmax is sufficiently small, can combine the upper bound with (35), (34), and directly conclude that
[ -3 +em(t) + I (=((-TT(e(()+ dr
$ 3B tem( + [
+ -okcll--)-Kx(t-)))ee dr
35

---

- M + (ex + S/" - (e)(())(eM(T) + , dr
(by (344)
+ 2
+
M (t) + 1
) dr.
Combining with the lower bound in FSL Theorem 4.3, the result follows.
D.4 The Case of Random-M Features
Notations. Throughout this section, for a power series P(x) Ei a,dd, we write P(A) Li=o A' to denote the result of substituting square matrix A RRd into P(.), where the power A' is obtained through matrix multiplications.
For a vector v E Rd,, we write the norm ||v|| for ||v||2 not otherwise specified. For a matrix A E Rdxd we use ||A l| to denote the operator norm induced by the vector 2-norm, and μi (A) > (A) > ... pal A) to denote the eigenvalues of A. For positive semi-definite matrix C E Rdxd
and vector E Rd, we define ||vlcc as
First, for the random-M feature setting, we can establish the following Volterra equation. Proposition D.11. Suppose 0 < & 1. Then, with probability at least 1 1----(() similar Volterra equation as derived in Theorem D.6 continues to hold for the randum-feature case; that is,
E[E M +ēm(t) + S( (E
T)
dr,
(36)
where
M
TM (t) = Σ 20103ee-200
RM== = Σ X3e 22,
j=1
j=1
and Ky My (WHWT) ) is the j-th eigenvalue of the random matrix.
Theorem 4.6 then follows by applying exactly the same argument as in the top- M case. It therefore remains prove Proposition D. The key idea is show that the spectrum of the random matrix WHWT E IR MMM closely matches that of the top-M truncation, namely,
Y := 1- (WHWT) Aj+
D.4.1
Concentration Inequalities
Recall that we derived the following recursive equation in Eq. (28):
2EE, u, A H /
tr(SA HA...S) n (c, + dr)dr,
where A; = e W WH: and S = (W" WHW W)
We first introduce the following notation: for integers 0 & < bb N (we allow b co, in this case we regard it as the same as = N).
Ha a.6 diag Aot1. As) E I (b-a)
Und = ((u)o+1,..... (u)b) E R
while
W a W Wa+1..... Wo] Mx(b-a)
the (a + 1) b-th columns of W.
To understand this equation with random projection matrix W, we leverage the following concentra- tion results developed in [35].
36

---

Lemma D.12 (Lemma G.4 in [35]. For H = diagf >2....,Ax' A such that >y= = where
I. there exists B B-dependent constants C1 O2 such that it holds with probability at least
1 - e M(M for all j E [M] that
a & Pyy (WHW E C2
Lemma D.13 (Lemma G.5 in [35]). There exists some B-dependent constant such that for all k 3 1, the ratio between the
2 -th and M-th eigenvalue
14 ( 1:00 H W
PM 1:00 Ho W
1:00
with probability at least 1 - e- R(M)
D.4.2 Upper and Lower Bounds
Lemma D.14. With probability least - (for s > 0 we have
M
trA HA, I S)= = e 2(1 23 =: (t - 7).
Σ
Proof. We can compute that
HA tr( WHW
Ir
W WHW
W Σ
alb! (-(l- (HW w) (
a,1=0
tr
(-t + a)
( WT(WHW T) jo+bb WH
a!b!
a.b=0
tr
(-t
(WHW) T jabb+2
a!b!
0,0=0
(-t + abb
2otb42
a!b!
0,
M
2(1 (((-T)}s
82 X3
which completes the proof.
For the first term in Eq. (28), following [35], we have
Lemma D.15. With probability least 1 e (MM) < 8 & we have
up A HA S M + (1).
Proof. Note
A HA, = Σ (-t)a-b HW WW HWW wH)
1
a,b=0 alb!
HW (WHW joto- WW WW (-t)a-
alb!
a,b=0
= HW (WHW - (WHW))--WH
37

---

where M, = (WWHW with Pt being the power series
P=(x)== = Σ
alb! ( (-t)atb
+)atblatb+)
a+b30
Note that when R, have
(-tx))
(-tz)²
(x(x) Σ
xe e 2tx.
a!
b!
a,620
Hence the eigenvalues of M, is exactly Pe(2g). Since
Hg A H. un (HW) (W "M for any positive integer k < 2 note that WHu Wo.k Ho: Ho.k + W kioo Hizo Ukzons we have
WWH)uo
u A, HA 2() T2);
where
T T
( 0:k
W (WWW M, Wo: o:k Ho.k) Ho:k:
0:k
T2 10:0
(H
1:00
W
(WHWT)- "WW 1:00 H 1n00 M0
Then by Lemma D.18, we can derive an upper bound. Since & 1,
Ti + T2 S
Look 11 +
| l
N
k-(1-s))
1+801-9)
Σ
1-
xp ~
+ k: -88
t
==
By setting k min(t2/0 M -3) we have
Ug A, HA & max(( -0) = 1- + M-
Clearly u, A, HA X 1, therefore by Lemma D.5,
u A, HA < min(t 1] + M +3 Eo0 (1) M- +3 (l) -
Lemma D.16. For it holds with probability at least --((1) that
un A HA Z max M-)))
where DM 20 24,,
Proof. Following the proof of Lemma D.15, we have n A HA, T. (WHWTT) M, WWWT) (WHW WHu,
((WHW) y -M.."M.WWWW) WHu, HW
M Σ
PM-443 (WWHWT) HIHHWT--M,,(NA (M((WWW)-)) μ4 (WHu,u, HWT),
where the last inequality is by Von Neumann's 's trace inequality: For two symmetric matrices A, B E IRdxd with eigenvalues an a2 2 ... 2 ag and by by 2... 2 ba, have
tr(A
a:ba-1-i-
i=1
38

---

Note that M, = Pi (WWWT) ) we then get
ug A, HA
Ms ((WHW " μ4 (WHuou,) HWT)
M
21 " (WHugu HWT).
Note that Un 0", by Assumption 2.3 and 2.4,
un A, H
2t i Ms (WHuou, HWT)
M
2td: 2 1-1--1-8-6
2t 2:10712 (t)
Here in the second line we used Lemma D.12 for matrix Hugu, H
My (Hugu u H) = 1310712 = --1ー My H (WHu.u, HW ~ -
and the third line follows from again from Lemma D.12 that Ny jz On the other hand, we prove the lower bound on M -$3 First, we claim that
Ug A HA > |(1
H (WWHTT) -H H uoll =: T3.
which will be proved in the Lemma D.17
Notice that
T3 IN 1 (WHWT)-
1/2 WHH
1 uou TH H1/)).
where the inner product (A, B) tr(A B) denotes the trace inner product between matrices. Therefore note that pa (H u H i-1-88 by source and capacity conditions.
T3 >
14 IN / - wwH (UN41-; u H²)
i=1
2
(M) (N -1
where the second line follows again from Von Neumann's trace inequality. Since M In T - is a projection matrix such that M = M and rank(Iy -M) = with probability 1, M must have M eigenvalues 0 and N - - eigenvalues 1.
Hence we have
N Σ 1=M M
T3 Z
08 2 M0
Lemma D.17.
- (WHWT)-IWHL 3 L A HA > |[I |
Proof. By the definition of positive semi-definite, we only need to prove that
HA Hà - H (WHW)) 'WH' H
39

---

Notice that
HA. HA, = = HW He WHI
-H I+ E
(-t)a H
(W H
a!b!
a+bia
Notice that H is a positive definite matrix, and now we only need to prove
I+
(-t)a+H H W Ww (WHWT)etb-
(I - H (WWW"")
alb!
44031
Let P WH After simplification, we only need to prove that
(PP {
I+ E
44631 alb! (-t)aa
P 'P.
Notice that, by the definition of matrix exponential, have
I+ Σ 1 (-t)atbp
(
a!b!
atbb1
=I I --- (PPPT)-- PPPPPP (PP Σ 2a+5
PT
(-1)att P
(a b)!
a+b30
= -I-P PP-PPPPPT)) -P P (PP - -2PP 'P.
Notice that the matrix P and -PPP are both positive semi-definite, we have P (PPT)-l 1---e--pp -2PP 'p positive semi-definite. As a result,
I+ Σ (-t)a+bp (PP T PPT)
alb!
atd21
which completes the proof.
Lemma D.18. With probability 1 1 -S((M), we have
|ggk 112
[uy (Work Ho. W (.x)) MM (Wo o:k How W
2
Ti &
T: < U1:so H
Wo:k)
where c is some constant.
Proof. Recall that the eigenvalues of M, (.WWWW is P.(Ag), ) so we have
|M = max P:(d,) = max dSe LEIEM
21
212,
max 2tuze 20A,
sup
DC
2t
LE
2t x>0
2ct
where the last step follows from sup > xe & 11
By the definition of Ti, have
Ti
Hoa W (WWW))-H---... (MW T)
Wos
1 1
[Mllll[[WHW - 0: Losk uo
[[WWHWTT Wo. Ho. 10.
We only need to show
Ha W
(WHWT) WWo &c
kood
MM W Wooe H W
1ss0
40

---

We denote A WA:00 H W Aooe and since WHWT Ho.kHoo W T 0:k
+ A we have
(WHWT)-IW Hock (Ap
& W Wokk H 0:h
+
W
0.4
Wo:k
W
0.A
W Wo.k Hosk
A W 0zk
Ho:k A
W o:k H
W
0:k
A
Wook
W T
0: A Wo.. Ho:k
A W ook Osk
A.
W
oskl
H
Ho:k
0:k
= A Wook H
+
AF
Woak
where the first line uses Sherman -or--MMooo--ooooe suitable shapes, we have
' identity: For matrices A, C, U, V with
(A + AUC)VV- A- A AUCCUCCCI + VA VA-
Since
H 0:k + W Wo Wosk E Wo. Wosk
it follows that
W 0:4 A Wo:k|=11e < W 'Woxxl
Therefore, with probability least -Q(M)
||
Wo. -A-'Wosee
+ W W AWook
112 < | 1l2 Wo:kl2
H
W 0:k A,
-oxll
| ||2 Wo
[ Ag
A Wo:k] 1112
|A lll [ ookkl
Mmin WW 0.k A Wo:k)
Assume and with probability at least -- (9(M) for some constant 0, || || o: l2 <
We may write W 0:k A Wor o.k
Ee As- 8;8e, where 8; iid N(0, I/NN and (2)
eigenvalues of AĄ in non-increasing order. Therefore, for k & M/3,
M Σ
M/2
M/2
clp
8,8, 2
Σ
8,8;
8,8, 2
AM-i
AM-
AM/2
XM/2
with probability at least e (M(M) where in the last inequality we again use the concentration properties of Gaussian covariance matrices (see e.g., Theorem 6.1 in [61]).
IA 112
- Waa 0.k H 0ok + WT - Work 1122 ||A x
Work
Amin
(W 0.1
A Work)
MM/22 (Ag)
Mas (Ag)
Now focus T2, by definition of T2 have
T2 Unooo
HoD
WT
(WHW -- exp(-2/WHW " (WWHW)) -1? W//r knoo
Hivo Mi:oo
1:o0
Ukzoo
H
WT
1:o0
(WHW T)- W 1:00 Hi:0 Ui:oo
H 1/2 W T
(WHWT) W Wooc H 1/2 UR:oo
1:00
1:00
H
H
where the last line follows from
H 1/2 W A:Do
Wi H /2
W Aooo (
H W
W Wivse Ha W. W
- k:00 ll2
Wirse H
0:k
k:0e 1.
The last line is because nonzero projection matrix has norm 1.
41

---

D.4.3 Proof of Proposition D.11
Combining Lemma D.14, D.15 and D.16, we get with probability at least 1 - E --(A))
E un A, HA L. tr(SA, HAr + )²) dr,
(37)
5 +tx(t)++++++++ ( (x dr,
(38)
Here we arrive at the Volterra equation for random feature case. Since we have 23 Xj. the proof in the top-M case also holds for EM and R.M
E Proofs for Section 5
When the condition o 2 holds indicating a constant label-noise level the FSL simplifies to
o² ((p), B
with N(p) ['KM(( KNN(t (C)(t-r)g(TT])) (r)) dr, (39)
B[E(vo)] M
where (l) M 6oo (1) M- 1- M- 6:0 60) - (xr S M-s8 and the fit-dependent noise term 6M is absorbed by the label noise term due to a 2 1. Extending to the full range a > 0 is possible but makes the statements and derivations much more involved. We therefore focus on the above case to streamline the exposition.
E.1 Proofs for Constant LRS
In this section, we prove Theorem 5.2 and present the data-optimal scaling strategy, as well as some results related the compute-optimal allocation.
Theorem E.1 (Restatement of Theorem 5.2). Under Assumption 5.1, when the learning rate 77k) for the top-M selection of the projection matrix W or for the random case with probability at least 1 - e (MM), have
E R
+ 02 + M M-
2
(nK)*
B
Proof. By main Theorem 4.2, when the learning 17(k) 1), denote 3. by (39) we have
[[Ex]
~ 70²
+ +M+* M3
Now we may write
0²
1
E
R
~
70²
+ -0
2
ts
Notice that t = nıK and
2. we have
1 _02 -3
R
+
2
(nK)*
B
Theorem F.2. Given a total data size of D > 1, the optimal strategy for minimizing the final population risk, in terms of the effective learning rate 7 and model size M is:
Topt = D Mopt Disesss
Eopt D 4t.
(40)
Proof. Since we have
[[Ex] x x0² + (>D) + 68
42

---

By weighted AM-GM inequality, we have that when &x is minimized, it must hold that
ya²
(TD)
which gives
Topt ~
Substituting this into the error expression yields
Copt D मी. M To balance the two terms and achieve the optimal rate, we require
Mopt 2 Dave
Consequently, the optimal loss rate becomes
Eopt ~ D 241.
Next we consider the compute optimal strategy for constant learning rates. We define the compute C = M K B to be the product of the model size, training steps and batch size.
Theorem E.3. Given total compute budget of C > 1, the optimal strategy for minimizing the final population risk, in terms of the effective learning rale 7 model size M, and data size D = BK, is:
Topt R C 143-18 Mopt 2 C C14240d Dopt R C 1443448
Proof. Since we have
E[Ex] ~ 10² + (n)
+ M M
substituting K
WB we get
M
E[EK] ~ 70² +
+ M -88
(C>)=
By weighted AM-GM incquality, we have that when EK is minimized, it must hold that (Cy)s²
2a² M
M (Cy)s
M-
which gives
Topt. 2 C 1131
Mopt Cridi
Now we can further compute D BK CM-1 CTTAFN
E.2 Proof for The Exponential-Decee LRS
Recall that the LRS given by
(p() = ac dr with ((K) = b,
where À log(a/b)/K =: 1/K. Note that the intrinsic-time transform is given by
a T(r) = (' p(r)dr= - (1-e-A)).
Thus, have
The total intrinsic time is:
T(K) = a -1-e- log(a/b)
K
(a-b) b) K(a b).
For simplicity, shall write T = T(K) in what follows.
43

---

The LRS- -dj-ussjsted function in intrinsic time is given by
(t(() = (p (t)) = a At.
Lemma E.4. The noise term satisfies (V(s)) b1 + (a - b) with
1 [" M
1 1 e- 222
e 2ul 2ule 2202
2/B du,
1 M
du.
4Tu2+1/3
Proof. We use the integral to approximate the forgetting kernel KM as
J,
M
x 20 -22 -x dx ("
(t) ~
20 E-2 -2
11 1- 1/8 e
2ut du.
M-3
==1
Noticing b= XT and AT = aa- - we have
L' KM (1rp(()) 1'"
( u² 1-110 1100-2uuff 2u(T du (a X1) dl
( (a xt) du
u 2-1180-211 2ul
M-8
A TeeuT
cauT 1
u 11188 2uT
(eauT -1)
du
2u
2u
2u
ae 2uT
XT
X(1 c -2UT)
+
du
M-8
1/
2n
2u /
4u 411B1/B
XI
(a XI + XT)
2nT
A(1 - 2uT
du
2u1/8
2/1/B
4u2+1/B
M-P
1 1-e e-e-2ul 22T
22uf
(a
XT)
du XT
du.
JM-
2u
2/1/3
47 4Tw
++18
Thus, complete the proof.
We next bound 1 and 12 separately.
Lemma E.5. IfT and M sufficiently large, then I 23-1 + or,
Proof. Note that
B(1 M (-(311)
M-B 241/3 du
=: A. 2(3 1)
and
2uT
e dr <
I(1 + 11
JM- 1/1/3 du
22-1/3 71-1/3 =: B.
2(2T)1-1/8 T/Ma r1/B
Then, complete the proof by noting A B. Lemma E.6. IfT and M is sufficiently large, then Emin(M (M, T2/3)
4T
Proof. Letr uT. Then, by a change of variable, we obtain
e-2r 2re 2r
1
(r(r) dr.
14113
4 1/8
is easy to verify that for any B > 1, inf,20 (r( > and (a(r) = 1-10/ large. We refer to Figure 10 for an illustration of (-).)
1/0
when r sufficiently
44

---

14
12
10-
02-
0.
10
Figure 10: Illustration of the function 45 (5(+).
When T/M < and T is sufficiently large such that 1 (r(r) dr 3 and thus we have
1e + 0M,,((1)
47
When T/M > 1, it holds for all exists a Cr.M E [0.5, such that
> that 0.5 < 1 -e 222
2re 2r & 1. Thus, there
1
1 Cr.M 471-1/8
1-1 dr
Cr.MB
T
1/3
T 111///
Cr.M -
471-1/8
MB
4T
Combining the two cases, we complete the proof.
Theorem E.7 (Theorem 5.3 in the main paper). We consider the exponentially decaying learning rate schedule
(p(T) ae Ar with ((K)= =
Under this learning rate schedule, for the top-M projection matrix the random projection with probability at least
e e we have
min (M,,M///))
EK = M 88 + +
b (a
B
T
where T' = ((a K/ log(a/b) is the total intrinsic training time.
Proof. By the functional scaling laws (39),
EK ~ - + + Q² N(O).
The noise term N(y) is estimated by Lemma E,4 and the bound on I1, In as
min M, T2/8)
N(p) = bli + (a b)I2 ~ b (a- 6)
T
which gives
0²
min min(M, (in,M,////)) T1/3)
EK M 38 + +T
b b+ (a -abb)
b)
B
T
so we complete the proof.
Theorem E.8. Given a total data size D) >> 1, the optimal strategy for minimizing the final population
risk when = K is given by Mapt 0c and
1ted-d
If 1 then Topt ~ (D/ log D)
1443
and Eopt (D/ log D) क,.
45

---

Ifs < 1 then Yopt ~ 1 and Eopt ~ (D/ log D
Proof. Denote D
log R then by Theorem 5.3,
EK M - +(DD)-*+ min(M, min(M,(yD)
Case 1.
When M < <D
M
εK R M + (+D)
D
We see that in this case y should as large as possible, since S 1, we set y accordingly.
In this case M sp + M
a2 1468, with equality at M ~ D 1es..
When 3 the above equality condition can be acheived M = DITT < <D. Hence
have that
Mopt Di Topt 1, Eopt D-1 1553.
Note that 7 = and a S which forces B 1, hence D ~ log D
When < - 3. the quantity - + M is decreasing with respect to M, hence the optimal M in D
this case is M (rD) which transfers to case 2.
Case 2. When M > DD,
EK ~ M (xD)
7
51
Clearly in this case Mopt = co, and by AM-GM inequality,
(nD) + x D1-1 2 D
1448
29
with equality at y D 1112
When 1 1. the equality can be achieved, hence have
Mops = 0c, Yopt ~ D
14-8-8 1130
Eopt ~ D D-119
When 3. since y $ 1, we must have
Mopt 00,
Topt R 1,
Eopt 15-5
Similarly, S 1, we have B S 1--1 which means K 2 Dises, hence log K x log D. D
D
log
Combining the two cases together, we see that Mopt can always achieves the optimal rate, hence the conclusion follows.
Summary.
Theorem E.9, Given a large total compute budget C > 1, the optimal strategy for minimizing the final population risk expressed in terms of the effective maximum learning rate 7,, model size M, and data size D is given by:
When 1 1/2 the optimal scaling laws are:
Topt. ~ ( (1o@)) 2150 Mopt F (To))) Dopt C'ETvd (log C'
which leads the following optimal final population risk:
(C) (0)=~(100) (10g)) 2043
46

---

(r monotone decreasing with y, we have that when When < 1 3 the optimal scaling laws are
Topt ~ 1, Mopt (oo Dopt ~ CTi (log C) C)
which leads to the following optimal final population risk:
(Cp ~ (Log )
Proof. Denote D = D/log K. For similar reasons as in the derivation of data-optimal scaling, may assume log K ~ log C to simplify the proof. At this point, the loss can be reformulated as follows.
1 (xD)
min M, (8)
EK ~ M MB
28
D
Case 1. M 7D and we have
1
M
EK ~ M
$8
+ 0²
(yD)* D
As y only appears in the second term, and EK is minimized, it must hold that
M = 175)
When > 1 3. we then consider weighted AM-GM inequality, we have
M =02M
D
Combining with C MD and M (7D)113, have
Topt
(To
1+800-1) 21.0
Mopt (Tog शोω,,
Dopt Catvo (log Civi,
14xd
and
.2 (C) ~ Tog C
2400
When < 11-1 4, since S 1, we set Topt. R accordingly, and proceed as follows:
Mopt (Tog
Dopt C++ (log C") 148
and
(C)
C ~ log
-TFU
Case 2. M > D and we have
(DD) + 2 (70)133
EK R M 88
88
D
As M only appears in the second term, and M- is monotonically decreasing in M, we have that when E is minimized, it must hold that
M (B.)
And then the rest is identical to the first case.
47

---

E.3 Proof for the WSD-Like LRS
To prove Theorem 5.4, we first present the following lemma, which gives an upper bound for the SGD noise induced by the stable phase.
Lemma E.10. For T2 > O, we have
min (M, Kas (T2 t)
T2
Proof. Similar to the previous section, use integral to approximate the forgetting kernel KM and get
/ ("r -2u((a+))
KM + +)
du dt
0 M-B
u-- be -2022 du
2u
u he 2u du.
T2 T
Since the integral 0 be-ee 2n du is convergent, have
1
KM KTr +
x
1-
T2
When T2 > M similarly we have
(Tt ~ 11 f:ee u be 2u To du.
I. Let N/
1, we have
M2
(t) ~
M
u --e -2up du
T2 p
M 8
M T2
2/
u be 2u du
T2
Where the last line is because pe 2up is decreasing in p when u. 2 1.
Theorem E.11 (Theorem 5.4 in the main paper). Suppose the FSL (10) hold and M, K are sufficiently large. Then, we have
b + (a min(M,T, 1/8
EK 88 TT + o²
B
BT4
where T aK1 + (a b)K2/ log(a/b) is the total intrinsic training time, and Tz (a K2/ log(a/b) is the dleaay-phase intrinsic training time.
Proof. By the results of the exponential decay LRS, let A log(a/b)/Ke, we have
T(K)
(T(()() - (1) - Ti
1-'
To
- (T(K)
+ KT t) At) dt,
Hence by the estimation the noise term of the exponential decay LRS (see the proof of Theorem 5.3). we have
Pz KM t) - At) dt = (a b) min M, T
Tb
48

---

Thus, we know
T(K)
/" (M(T(K) t)a + + (a b) min { M, T22
T2
(a
min ( M, T2
~ a
(M + t) + +b+
T3
(a- (a min (., T2
xb
(by using Lemma E.10)
T2
Hence the loss is given by
1 T
o²
min mingM, (M,
εK + M -B
b (a -
T2
Theorem E.12. Assume K then we have the following data data-optimal strategy:
If8 8 2 1/B, we have Topt
1ted-1
(log D) 1+12, (D1)opt: (D2)ops 2 D and
Eopt D #041 (log D) 1448.
Ifs 1/3, we have Yopt R 1, (Di)opt ~ D, (D2)opp 2 Dé-1 log and Eopt ~ D=s
Proof. Since the total intrinsic time T S yD, can always take D₁ D to ensure T 7D. Du then by Theorem E.11,
Denote D2 log
M (+D) + min(M, (TD2)3)
D2
Case 1. When M & DD,
M
EK ~ -3 + (xD) -* +
We that in this y should be large as possible, since with equality at M D.
1, set 1 accordingly.
In this case -B +
Dg 2 D2
When s have that
3. the above equality condition can be achieved as M D 1as Dz. Hence we
Mopt ~ Di Topt 1, Eopt ~ D2
Therefore (D2)opt D. Note that = ~ 1 and a S 1, which forces B = 1, hence D2 log D When 3. the quantity M + M. is decreasing with respect M, hence the optimal M in D2
D
this case is M = (z) which transfers to case 2.
Case 2. When 3 > DD.
EK ~ - + (xD) -* +
D2
Clearly in this case Mopt co, and by AM-GM inequality,
1
20
(xD) ) D
2
D
14. D
49

---

2-1 When > 3 the equality can be achieved, hence we have that (D2)opt D, 80 D2 log
28 with equality at 7 D 1is2 D2
D
Mopt = 20, Topt ~ D 1150 (log K) finb, Eopt ~ D-DDIS
(log K) 1950.
When < 11 3 since S 1, we must have either 7 ~ 1 oor D
odl 1-38
D2
TT
2-1
2 1. To reach the
minimum risk, both cases we require (D2)opt 2 DØ (this gives (D2)apt 2 D log D), and
Mopt 00, Topt ~ 1 Eopt D
Similarly, as a S 1, we have B Slog D4-T=ar T=,,, which means K Zlog Teas which gives the desired rate.
hence log K log D,
Summary. Combining the two cases together, we see that Mopt 0o (case 2) always achieves the optimal rate, hence the conclusion follows.
Theorem E.13, Assume = k. under the compute constraint C > 1, the optimal strategy for minimizing the final population risk expresssd in terms of the effective maximum learning rate model size M, and data size D-is given by:
When 1 1/3, the optimal scaling laws are:
Nopt
Loog
1823- 2447
Mopt
(1oge)
Dopt C 2400
Lizd
(log C) 2- (Di)opt (D2lot D,
which leads the following optimal final population risk:
+8-a
Eopt = C Tean (log C)
When 1 1/8, the optimal scaling laws are:
Topt Mopt C Dopt CTEB (D2)opt (D2)opp 2 D log D,
which leads to the following optimal final population risk:
Eopt
Proof. Since the total intrinsic time T S 7D, we can always take D, D to ensure T = 7D.
Denote D2 = Da the loss can be reformulated as follows.
log K
1
+ oo2 min min(M, (M, (7D2)
EK R M 88 +
(DD)
D2
Case 1. M < DD2 and we have
M D
EK ~ M
D)
As only appears in the second term, and Ex is minimized, it must hold that
bnD is monotone decreasing with we have that when
M 2/8,
When 3 we then consider weighted AM-GG inequality, we have
M
M 48
D2
Combining with M 1/B have
143(6-1)
Topt D
1150
Mopt D
50

---

and
5- vd
Eopt ~ D2
1448 D=
Notice that
2-43
C C D2
D
of
E
21zd
Note that this implies (the max learning rate By S 1).
2 C
= log R log C, and by similar reasons log K ~ log D
Hence when & is optimized, we have D2 ~ D/ log C' and
Topt ~ ( loo
C
1t 2408
Mopt ~ ( log )
2003
Dopt
24ad C2ar8 (log O) 2
and
Eopt (C) ~ ( C 24.8 (log 0)*.
(ogg)
When < 1 3r
since a S I, we set Topt. ~ 1 accordingly, and proceed as follows:
Mopt D
and
Eppt D
Notice that
C D3 D D22
_ = EE6
148
Hence when & is optimized, we have D₂ ~ D/ log C and
Topt I, Mopt C Dopt C
and
Eopt ~ C=1 28
Case 2. M > 7D2 and we have
(7D2)² /
EK M- 08
(nD)²
D2
By AM-GM inequality,
(nD)
x
1
2
144 D2
1
D2
with equality at 7 R D rish D
When 1 1. the equality can be achieved, hence (D2)apt D, and the loss can be written as follows.
EK M + D T T D2
1443
Combining with C = MD. we have the optimal scaling laws follows:
14=8-8 Topt C- 24-3"
(log C*- fros Mopt x Catan (log C'
1-1/0
2t00
Dopt x Czten (log C*)
1448
1-1/0 2-40
which leads to the following optimal final population risk:
Eopt C 24e0 (log C* 84*0.
28 When 1- - 3 since S 1, we must have either 7 R 1 or = D 1is3 D S 1. To reach the minimum risk, in both cases require (D2)ont 2 D (this gives (D2)ope 2 D log D), and Topt F 1, EK = M = 28 D Combining with C = MD, we have the optimal scaling laws follows:
8-1
Topt ~ 1, Mope = C Dopt = CT43, which leads the following optimal final population risk:
Eopt C 143
51

---

Summary. Combining the results of each case, we get the desired optimal scaling strategy stated in the theorem.
F
Auxiliary Lemmas
Lemma F.1. For any PSD matrix A and random gaussian vector (N, H). tr(HA)H 3 E xx Axx -HA tr(HA)H + HAH 2ir(HA)H
Proof. Assume = (Aig) M. The (i,j)-th entry xx Axx
Σ XXXXAAXXXX
A
If t.
E
Σ
XxXXX
2E [A,xtx]] 2A1 2HAH(i,).
If ij
M Σ
Σ
ARXIX
ARXXx
Σ
4m = (
By the trace inequality have Multiplying H both sides,
HA 3 tr(HA).
HAH (HAH)
Combining the results, we have
E|xx Axx 1
tr(HH) + 2HAH (tr(H + HAH.
Lemma F22 Ler P Q be two PSD matrices. Then for any PSD matrix U, we have
tr(vPUvP) < tr( (U(VQU/Q)
Proof. It is clear that Centtrr(PPUUU PUVP) = tr(UP) and tr(UQ) -
P)) > 0,
since and Q - P are both PSD matrices.
52

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Specifically, we state that our work introduces a novel Functional Scaling Law (FSL. that captures the impact of learning-rate and batch-size schedules. Thhese claims are substantiated by rigorous theoretical analysis (e.g., Theorem E. 1. concrete examples (Section 5) and experiments in Section 6.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss several limitations and future directions in Section 7.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
53

---

Answer: [Yes]
The paper presents a complete set of assumptions and rigorous theoretical proofs for all main results. Assumptions 2.1, 2.3 and 2.4 define the problem setup, model capacity, and task difficulty. Detailed proofs of all the theoretical results are provided in Appendix and E.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Appendix C provides detailed specifications of the learning rate schedules, model sizes, number of steps, averaging procedures, and other hyper-parameters. The information provided is sufficient to reproduce the results.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
54

---

reproducible, they some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: While we provide detailed descriptions all experimental setups and hyper- parameters in Section and Appendix C, we do not currently release code or datasets.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips. ccc/pbbicc/eeee wwtcc//guiee//eee Ceeees//ooSSubbsseeee ssionloicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper provides all necessary training details for reproducing the teacher-student kernel regression experiments in Appendix C.I. For LLM experiments, we specify the details in Appendix C.2
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report the standard deviation as the error bar. This is clearly stated in Section and Appendix C.I. While do not explicitly verify the normality of the error distribution, the large number samples ensures that the mean and standard deviation are reliable indicators of statistical trends.
Guidelines:
55

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: While describe the experimental setup in full detail, do not currently report the specific compute resources used.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: : /ttp//euuiip..eee
Answer: [Yes]
Justification: Our work does not involve human subjects, sensitive data, or deployment in real-world applications. All claims rigorously supported by mathematical derivations and empirical validation, and we have taken care to ensure transparency, reproducibility. and fairness throughout the study.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
56

---

Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper is a theoretical work and there is no societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our paper does not involve the release of any pretrained models, generative systems, or scraped datasets that pose a risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The code packages and open-source models used in this paper are all properly credited.
Guidelines:
57

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release any datasets, pretrained models, or external code packages.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the detaset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The research does not involve human subjects crowdsourced data collection. No participant interaction or compensation is involved at any stage of the work.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
58

---

Answer: [NA]
The research does not involve human subjects or any data collection from individuals.
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The paper does not involve LLMs any important, original non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: : htt///uur..cc///ooee// http://sc.cccooeeeenee///////// /LLM) for what should or should not be described.
59


Reviews:
[{"review_id": "TDMK6xWAdD", "reviewer": "Reviewer_93yY", "summary": "This paper considers how learning rate and batch size influence the loss trajectory in a teacher-student kernel regression setup. The setup considers learning teacher function $f^{\\*}(x) = \\langle \\phi(x), \\theta^{\\*} \\rangle$ where $\\phi(x)$ is an $N$ dimensional feature map and $x \\in R^d$ with noise $\\sigma^2$, with a student model $f(x; v) =  \\langle W \\phi(x), v \\rangle $, where W is either a random $M$ dimension projection or the top $M$ feature of the feature map $\\phi$. M is the concept of model size here. The paper then introduce the concept of model capacity by defining $\\phi_j = \\lambda_j^{1/2} e_j$ where $e_j$ is an orthonormal function basis and $\\lambda_j$ follow $j^{-\\beta}$ scaling. With lower $\\beta$, the feature decays slower and hence is a 'richer' model. The authors then assume $\\theta_j$ follows $j^{-1/2}\\lambda_j^{(s - 1)/2}$ for different $j$ so that the function become $j^{-(s\\beta + 1)} e_j$ and $s$ here characterize the relative difficulty. \n\nAfter establishing this setup, the authors consider a continuous time approximation, using SDE to simulate the discrete optimization process. Given the SDE, it is then possible to give a sharp characterization of the loss scaling of the model sizes and data sizes, with a form similar to $1/M^{s\\beta} + 1/t^s + Noise$. They then consider how different learning rate corresponds to different continuous time $t$ (which is an integration of the learning rate) and noise level (which is a convolution of the learning rate). Finally, they use the function form to fit the training loss curve of kernel regression and standard pretraining.", "questions": "Please refer to the weakness section for my main questions. Some of the minor questions include:\n\n1. How will the current theory predict the batch size to scale with training duration?\n\n2. A minor suggestion is that 8-1-1 learning rate schedule is not defined explicitly in the empirical section and is confusing to understand at the beginning.", "limitations": "yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strength\n\n1. The theoretical setup is well-presented and very easy to understand. The concept of model capacity and relative difficulty is well motivated. The final result is also of high theoretical interest.\n\n2. The theoretical analysis is supported by strong empirical evidence, showing the relevance of the analysis presented in this paper. \n\n## Weakness\n\nThis paper lacks a more complete discussion and comparison of previous works. While previous works didn't have an end-to-end theoretical formulation like this paper (which is indeed a unique contribution), the final scaling form presented in this paper is very similar to previous ones, which is worth further discussion. I will list some of the previous works and their similarity below:\n\n1. The idea of using SDE to approximate the discrete optimization process is studied in deep learning literature and the concept of intrinsic time is also introduced in previous works. A very incomplete list includes [1, 2, 3]. Notably this work also considers the interplay between batch size and learning rate (as in this paper).\n\n2. Even in the concept of learning rate scheduling and how it influences loss, [4,5] has already mentioned the decomposition of the integration of learning rate and the noise level (which is determined by a convolution of the learning rate). \n\n3. Finally, it would be worth considering comparing the scaling law formation in this paper with previous works. There should be two comparisons (1) explicitly comparing the formulations and (2) comparing the fitting/prediction quality in real-world experiments. \n\n[1] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations\n[2] On the SDEs and Scaling Rules for Adaptive Gradient Algorithms \n[3] On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs) \n[4] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective\n[5] A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules\n\nAlso, it is worth noticing that while the FSL framework allows varying batch size, it is not studied in depth in this paper. It would improve the clarity of the paper if this is clarified in the abstract and introduction.", "was_revised": true, "final_justification": "This paper build an interesting theoretical model that is reflective of neural network landscape and respond promptly to compare with the baseline in other works. Although the empirical advantage over previous method is slight, the theoretical novelty should be enough for publication in NeurIPS..", "rebuttal": "We sincerely thank the reviewer for the positive support and valuable feedback! We greatly appreciate the insightful review, and the recognition of highlighting the significance of our contribution and the strength of our theory and experiments, as well as the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.\n\n---\n>**[W1]** This paper lacks a more complete discussion and comparison of previous works.\n\n**[E1]** We thank the reviewer for pointing out the relevant works. We will ensure that the discussion of these works is highlighted in the revised version of the paper.\n\n$\\quad$\n>**[W1-1]** The idea of using SDE to approximate the discrete optimization process is studied in deep learning literature and the concept of intrinsic time is also introduced in previous works.\n\n**[E1-1]** \nWe thank the reviewer for the helpful comment and would like to clarify that both the SDE formulation and the concept of intrinsic time have indeed been introduced in prior works. Our contribution does not lie in introducing these notions themselves—we cite many relevant references in Section 3, though we have missed some important ones and will work to address that.\n\nOur contribution, rather, lies in showing that both the SDE formulation and, in particular, the notion of intrinsic time are crucial for deriving precise scaling laws that hold under general learning rate schedules. In the revision, we will add a paragraph to clarify these points and better highlight our unique contribution within the SDE-based framework for analyzing stochastic optimization in deep learning.\n\n$\\quad$\n>**[W1-2]** Even in the concept of learning rate scheduling and how it influences loss, [4,5] has already mentioned the decomposition of the integration of learning rate and the noise level (which is determined by a convolution of the learning rate).\n\n**[E1-2]** We agree that both [4] and [5] have previously leveraged the notion of intrinsic time (i.e., the integral of the learning rate) to analyze the impact of learning rate schedule (LRS). \n\n- [4] provides a heuristic explanation for the sharp loss reduction observed under warmup–stable–decay (WSD) schedules. In contrast, we establish that WSD-like schedules indeed outperform pure decay schedules by deriving precise scaling laws. \n\n- The theoretical part of [5] analyzes the effect of LRS in a quadratic optimization setting and presents a risk decomposition into an intrinsic-time scaling term and a noise term. Compared to it, we consider a more structured learning problem characterized by two exponents, $s$ and $\\beta$, which allows  to characterize the LRS’s influence **quantitatively**. Moreover, we derive precise scaling laws for three LRSs, enabling a direct and quantitative comparison across them. \n\nWe will include a more detailed comparison with [4] and [5] in the revision to better clarify the contributions of our work over them.\n\n$\\quad$\n>**[W1-3]** Finally, it would be worth considering comparing the scaling law formation in this paper with previous works.\n\n**[E1-3]** We thank the reviewer for this insightful suggestion. We will include a comparison of our FSL formulation with previous works in the revised version of our paper.\n\n**Comparison with Chinchilla-style scaling law [6][7].** Chinchilla-style scaling laws take the form of the final loss $\\mathcal{L} = \\mathcal{L}_0 + c_1 M^{-a_1} + c_2 D^{-a_2}$, where $\\mathcal{L}_0$ represents the irreducible loss, and $c_1$, $c_2$, $a_1$, and $a_2$ are parameters that depend on the specific problem. In contrast, our FSL framework incorporates the effect of LRS through an additional noise term. Moreover, while Chinchilla-style scaling laws are limited to estimating the final loss of a given LRS (different LRS may have different exponents), our FSL model can predict the entire training process loss of any LRS.\n\n**Comparison with multi-power law [5]**. Reference [5] developed an empirical law that incorporates a carefully designed loss reduction term to predict the training loss under general LRSs. In contrast, our work provides a theoretical explanation with a different decomposition: in their law, the loss reduction term is negative (subtracted from the first term), while our noise term is positive (added to the first term). In terms of fitting ability, our FSL slightly outperforms the multi-power law, as shown by the following experimental results: \nWe fit the laws on the 8-1-1 training curve and predict on cosine and WSD schedules. The results are shown in the table.\n|       | Cosine   |        | WSD     |        |\n|-------|----------|--------|---------|--------|\n|       | MPL      | Ours   | MPL     | Ours   |\n| $R^2$ | 0.9811   | **0.9825** | 0.9819  | 0.9819 |\n| RMSE ($L^2$ distance) | 0.0528   | **0.0512** | 0.0509  | 0.0509 |\n| MAE ($L^1$ distance) | 0.0276   | **0.0242** | 0.0242  | **0.0240** |\n| WorstE ($L^\\infty$ distance) | 0.5886   | **0.5814** | 0.5764  | **0.5691** |\n\n$\\quad$\n>**[W2]** Also, it is worth noticing that while the FSL framework allows varying batch size, it is not studied in depth in this paper.\n\n**[E2]** We thank the reviewer for raising this point and would like to clarify our rationale. \n\nFirst, the study of the influence of learning rate schedules (LRSs) already involves substantial technique development and takes up significant space. \n\nSecond, while the FSL framework in principle allows for varying batch sizes, it is derived for SGD. However, practical LLM training uses adaptive optimizers such as Adam, where the interaction between learning rate and batch size can differ substantially. Understanding how the insights from  the FSL framework  translate to such settings would require substantial additional theoretical and empirical analysis. \n\nFor these reasons, an in-depth analysis of batch size schedules is better suited for a separate work. We will clarify this point in the revised Conclusion section.\n \n$\\quad$\n>**[Q1]** How will the current theory predict the batch size to scale with training duration?\n\n**[A1]** The final-step risk depends on both the learning rate schedule (LRS) and the batch size. Therefore, a comprehensive analysis requires substantial effort and is beyond the scope of this paper. However, under constant learning rate, i.e., it does not scale with training tokens, a careful calculation suggests that the batch size should scale as  \n\n$$\n\\qquad\\qquad\\qquad\\qquad\\qquad\\qquad\\qquad\\qquad\\qquad\\qquad B \\asymp D^{\\frac{s}{s+1}}.\n$$\n\nThis indicates that **larger batch sizes should be used when training for more tokens**. This theoretical prediction is consistent with empirical Step Law of batch size reported in [8], which also suggests that $B\\asymp D^{\\gamma}$ for some $\\gamma>0$.\n\n$\\quad$\n>**[Q2]** A minor suggestion is that 8-1-1 learning rate schedule is not defined explicitly in the empirical section and is confusing to understand at the beginning.\n\n**[A2]** We thank the reviewer for the helpful suggestion. The definition of the 8-1-1 schedule is currently provided in line 605 of the Appendix. We will move it to the main text in the revised version to improve clarity.\n\n---\n**Reference**\n\n[1] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations.\n\n[2] On the SDEs and Scaling Rules for Adaptive Gradient Algorithms.\n\n[3] On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs).\n\n[4] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective.\n\n[5] A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules.\n\n[6] Scaling laws for neural language models. \n\n[7] Training compute-optimal large language models.\n\n[8] Predictable Scale: Part I, Step Law – Optimal Hyperparameter Scaling Law in Large Language Model Pre-training."}, {"review_id": "tYLZXdfXIY", "reviewer": "Reviewer_SA1x", "summary": "Scaling law analyses usually only consider the final loss after training and overlook the effect of learning rate schedule on the training trajectory. This work derives a functional scaling law from a teacher-student kernel regression setting, where the impact of learning rate and batch size schedules is given via a convolution-type noise term. For the constant, exponential decay, and WSD schedules, the authors compute their scaling exponents under data-limited and compute-limited regimes, where the theory aligns with empirical observations in LLM pretraining.", "questions": "* Do the authors foresee any trends with the optimal FSL LR schedule with model scale? The optimal schedule for the 100M and 1B models appear similar.\n* What is the reason for using the 8-1-1 LRS for the fitting of the FSL as opposed to the other schedulers?\n* The FSL incorporates batch-size scheduling, but there doesn’t seem to be mentioning of batch size scheduling later in the paper. Given that batch size scheduling is common for large-scale model pretraining (eg. Llama 3 increases the batch size twice during training), how could the FSL fit an optimal LR schedule in light of varying batch sizes during training?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* The paper is generally well-written with the main result and subsequent applications to specific schedules presented clearly.\n* Studying the interaction between learning rate schedules and scaling laws is an important problem.\n* This work provides a theoretical derivation of a scaling law that is aware of the learning-rate schedule, where previous work similarly tackling this problem were largely heuristic-based.\n* Empirical results in the teacher-student kernel regression setting show strong agreement with theory and the FSL-optimal schedule appears competitive for language model pretraining.\n\nWeaknesses:\n* The FSL-optimal LR requires fitting a ‘practical’ version of the FSL to an existing loss curve and is not predictive. This has to be done for a specific model scale, duration and maximal learning rate; the method would be more compelling if a trend emerged across model scale that allowed for extrapolation instead of requiring a full training run.\n* The practical FSL used for fitting places some assumptions, particularly assuming a fixed batch size; although the authors mention their FSL accounting for batch size schedule as well, this isn’t explored further.\n* There still is a gap in understanding what is the optimal LRS is for a given training setting– the optimal LRS identified by the practical FSL all seem to be of a similar nature but a more precise characterization depending on model architecture, batch size, dataset size, etc. would be valuable.\n\nI did not read the theoretical proofs in detail. Despite what I’ve written above, I do believe the results presented are a novel step in this direction and the general framework could be of interest to the broader community.", "was_revised": true, "final_justification": "The authors have done additional experiments looking at the predictive power of their FSL framework during the rebuttal period. Although their theory would have to be extended for adaptive optimizers, I believe the result is of important interest and will lead to follow-up work.", "rebuttal": "We sincerely thank the reviewer for the encouraging and insightful feedback, and for highlighting the strength of our theoretical contributions and the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.\n\n---\n>**[W1]** The FSL-optimal LR requires fitting a ‘practical’ version of the FSL to an existing loss curve and is not predictive. This has to be done for a specific model scale, duration and maximal learning rate; the method would be more compelling if a trend emerged across model scale that allowed for extrapolation instead of requiring a full training run.\n\n**[E1]** We thank the reviewer for these insightful suggestions and fully agree that improving the predictive power of FSL across model scales,  durations, and maximal learning rates is essential for making it practical. \n\nIn response, we trained a 100M-parameter model and find that the fitted FSL  generalizes well across different durations. Specifically, we fit a FSL using a 20B-token training run and successfully predict loss curves for 80B tokens, achieving a consistently low RMSE of 0.0120 across durations. \n\nTo extend FSL across model scales, we follow the dependency modeling approach of [2]. On the positive side, we observe some degree of generalization. However, due to the lack of a principled design, the fitted FSL tends to fail when extrapolating to model scales far from those used in fitting—at least in our current experiments. \n\nIn summary, addressing these problems remains challenging and will require a quite amount of effort. We will add a discussion in the revised version to highlight these directions for future research toward making FSL practical.\n\n$\\quad$\n>**[W2]** The practical FSL used for fitting places some assumptions, particularly assuming a fixed batch size; although the authors mention their FSL accounting for batch size schedule as well, this isn’t explored further.\n\n**[E2]** We thank the reviewer for raising this point and would like to clarify our rationale for focusing on a fixed batch size in this paper. \n\nFirst, the study of learning rate schedules (LRSs) already involves substantial technical development and occupies significant space.\n\nSecond,  the FSL framework is derived for SGD but practical LLM training typically relies on adaptive optimizers, where the interaction between learning rate and batch size can differ substantially. Understanding how insights from the FSL framework translate to such settings would require considerable additional theoretical and empirical analysis. For these reasons, we believe that an in-depth study of batch size schedules is best addressed in a separate work. In the revision, we will clarify this point in a discussion section.\n\nThat said, we have already conducted some preliminary exploration of batch size scheduling. When the learning rate is fixed, our analysis reveals that the optimal batch size schedule should be: $B_{k} \\asymp  (K-k)^{2-1/\\beta}$ in the data-optimal regime for the learning problem we considered. This indicates that the batch size should gradually increase over the course of training. To validate this insight, we trained a LLaMA-0.5B model on the C4 dataset for 10B tokens using a constant learning rate of 1e−3, and compared different batch size schedules:\n| setting | final loss |\n| --- | --- |\n| all tokens using 512 batch size | 2.9 |\n| all tokens using 1024 batch size | 2.86 |\n| 32% tokens using 512, 68% tokens using 1024 batch size | 2.83 |\n\nThese results support the theoretical insight that increasing batch size over time can improve performance.\n\n$\\quad$\n>**[W3]** There still is a gap in understanding what is the optimal LRS is for a given training setting– the optimal LRS identified by the practical FSL all seem to be of a similar nature but a more precise characterization depending on model architecture, batch size, dataset size, etc. would be valuable.\n\n**[E3]** We thank the reviewer for the insightful comments. We conducted additional experimental analysis and found that the fitted FSL-optimal LRS consistently follow a WSD pattern with power decay of the form:\n$$\nf(k - K_{stable}) = \\left( \\frac{K_{total} - k}{K_{total} - K_{stable}} \\right)^\\gamma,\n$$\nWe further explored how the optimal exponent $\\gamma$ depends on two factors:\n\n- **Training tokens**: For a 100M-LLaMA model, we fit the optimal $\\gamma$ for different training tokens:\n| Dataset Size | 10B  | 20B  | 40B  | 80B   |\n|-------------|------|------|------|-------|\n| $\\gamma$    | 1.91 | 1.59 | 1.86 | 2.15  |\n\n- **Model size and architecture:** We trained a 1B-LLaMA and 1B-MoE model for 20B tokens. The optimal $\\gamma$ are $1.40$ and $1.28$, respectively.\n\nAlthough the optimal $\\gamma$ varies with tokens, model size, and architectures, it is consistently greater than 1, suggesting that convex decay outperforms concave decay.\n\nA comprehensive analysis of these trends is an important direction for future research.\n\n$\\quad$\n>**[Q1]** Do the authors foresee any trends with the optimal FSL LR schedule with model scale? The optimal schedule for the 100M and 1B models appear similar.\n\n**[A1]** See [E3] for experimental results on the shape of optimal LRS. According to our preliminary experiments, we conjecture that optimal LRS takes the form of power-decay, and the exponent $\\gamma$ is smaller for larger model sizes: for 100M-LLaMA models, the optimal $\\gamma$ is around $2$, while for the 1B LLaMA and MoE, the optimal $\\gamma$ is around $1.5$. The detailed analysis is beyond the scope of our paper, and we consider this an important future direction.\n\n$\\quad$\n>**[Q2]** What is the reason for using the 8-1-1 LRS for the fitting of the FSL as opposed to the other schedulers?\n\n**[A2]** We thank the reviewer for the valuable comment. The 8-1-1 LRS is used due to its simplicity, being a multi-stage piecewise constant schedule. However, we acknowledge that other schedules can also be used for fitting with similar results. For example, we conducted fitting experiments on a 100M model trained with 20B tokens using WSD and cosine schedules, then predicted the loss curve of these 3 schedules (8-1-1, WSD, cosine). The average RMSE of these predictions was 0.04615 and 0.04916 for coefficients fitted with WSD and cosine, respectively, while the prediction RMSE for coefficients fitted with the 8-1-1 schedule was 0.04617. \n\n$\\quad$\n>**[Q3]** doesn’t seem to be mentioning of batch size scheduling later in the paper. Given that batch size scheduling is common for large-scale model pretraining (eg. Llama 3 increases the batch size twice during training), how could the FSL fit an optimal LR schedule in light of varying batch sizes during training?\n\n**[A3]** First, we refer to our response **[E3]** for weakness **[W3]**. Second, as for identifying an optimal LR schedule under varying batch sizes, we agree this is an important direction for future work. However, the current FSL framework is derived for SGD, whereas practical LLM training typically uses Adam, where the interaction between learning rate and batch size can differ significantly. Extending FSL to this setting would require substantial additional theoretical and empirical development, which we believe falls beyond the scope of the current paper and is best pursued in future research.\n\n---\n\n**Reference**\n\n[1] A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules.\n\n[2] Scaling Law with Learning Rate Anneal."}, {"review_id": "ZHbagOgO4X", "reviewer": "Reviewer_HjUT", "summary": "This paper analyzes the scaling behaviour of the risk for a teacher-student regression problem when using different learning-rate schedules. \nThe paper shows that the theoretical model describes well the actual dynamics of SGD in the same teacher-student setting, but also for LLM pretraining (where the assumptions of the theory do not hold). As an application, the paper derives a optimal schedule from fitting their theoretical law, and shows that this schedule also works well in practice.", "questions": "* You claim (line 93) that you provide the first theoretical explanation for different scaling behaviours across LR schedules. I am not sure whether this is accurate: the scaling with data can be seen basically as the convergence rate of SGD. Studying convergence rates for LR schedules has been done previously in the optimization community (even though it is not termed a scaling law). For example, this paper (https://arxiv.org/pdf/2501.18965) shows an improved rate for WSD vs. constant schedule, which essentially also implies a different scaling in data budget. Of course the specific setting of your paper allows for much more insights with respect to model capacity, batch sizes etc, thus the contributions of your submission are still very much useful.\n\n* Regarding the sentence: \n\n> As a result, different LRSs can induce distinct scaling behaviors, with the scaling exponent itself potentially depending on the schedule.\n\nIs there actually empirical evidence for this? The only reference I know on this is https://arxiv.org/abs/2406.19146, where it is shown that the scaling law for constant schedule is almost the same as for cosine schedule (under suitable tuning of other hyperparameters).\n\n* Can you analyse the WSD schedule with a linear cooldown (or other forms that exponential in general)? If not, what are the limiting factors in the analysis? This is relevant as the prevalent cooldown shape in practice seems to be linear or 1-sqrt (see the paper by Hägele et al, 2024).\n\n* Caption of Figure 2 says:\n\n> In all panels, dashed curves represent the theoretical predictions from Section 5. (a) Solid curves show the risk predicted via numerical simulation using the FSL.\n\nThis was confusing, as I thought that the FSL is the theoretical prediction, and the solid curves come from actual SGD runs. What am I missing here?\n\n* The FSL theory applies to SGD, but is then used to fit Adam runs in the LLM setup. Could you comment on this limitation?\n\n* You are fitting the values of $s$ and $\\beta$ in the LLM experiment. Do the scaling coefficients obtained from the fitted values agree with the scaling laws found for LLM pretraining in practice (e.g. the Chinchilla law)?\n\n* The emphasis in the abstract on batch size schedule is very strong, but the main text does not provide many insights into optimal (?) batch size schedules, or empirical validation of the theory regarding the batch size.\n\n* line 245 says that LLM practice is to allocate more compute needs to data. But in contrast, the Chinchilla scaling law concludes that the powerlaw coefficient in data and model is the same?\n\nMinor:\n- line 98: \"up to a constant factor\": multiplicative or additive factor?\n- Equation (1): I think there should be a factor 1/2 in front of the squared norm to be consistent\n- line 225: how is compute $C$ defined here?\n- line 248: typo \"brevity\"", "limitations": "No potential negative societal impacts.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The structure and presentation of the paper are very clear, and the main objective can be easily grasped (even though the mathematical details are technical).\n\nFor the experimental validation, some additional efforts are needed:\n\n* From the LLM experiments it seems that the FSL fit is predictive across schedules (fit on 811, then predict cosine and WSD). This is great, but something that is missing here is the role of the peak learning rate. Usually, for different schedules the optimal peak LR is different. In your experiments, did you tune the peak LR for all schedules? Can the FSL also predict the loss curves across peak learning rate? This discussion is missing, but for practical applications it is highly relevant, as the (peak) learning rate will most likely be the most important hyperparameter to tune.\n\n* The optimal LRS for the LLM experiment has an unfair advantage in the sense that the other schedules do not decay to zero, but a 10% of the peak LR. However, it is known (see your references 7,21) that cooldown to zero gives a benefit, and in fact for WSD (see [21]) this is what was proposed from the beginning. Hence, in order to properly assess whether the optimal schedule brings a benefit, the baseline should be cosine and WSD schedules that decay to zero. It is however a nice finding, that the theoretically optimal schedule decays to zero (but again this has been previously found in other settings, see for example https://arxiv.org/pdf/2310.07831).\n\n* Related to the above, the LLM experiments use the same peak LR for all schedules. It has been shown empirically and theoretically that different schedules will have different optimal peak LRs, hence the missing tuning of the peak LR might confound your findings.\n\n**Ablating on the role of the peak and final learning rate, in particular their effect on the experiments provided, is the main revision that I suggest to improve this submission.**\n\n* A minor weakness regarding clarity: some important information, on which constants are exactly fitted (in Sec 6.2), and which runs form the basis of the fit (in Sec. 6.1) are only provided in the appendix. However, this information would be insightful when reading the main text.\n\n* Another minor point: the connection between the discrete schedule and peak LR, and the continuous LRS $\\phi$ does not become fully clear and needs further explaining. For example, in line 173, the maximal argument for $\\varphi$ seems to be $\\tau$, the number of training steps, but then in line 169 it seems that the argument of $\\varphi$ is the step counter times $h$. So on which time scale is $\\varphi$ defined and how does that scale related to the (integer) SGD iteration counter?", "was_revised": true, "final_justification": "Based on my initial review, concerns regarding practicality and whether the current experiments reveal the full picture remain. The authors were also transparent about this in the discussion period. Therefore my borderline score remains.", "rebuttal": "We sincerely thank the reviewer for the positive support and valuable feedback! We greatly appreciate the insightful review, recognizing the significance of our contribution, the solidity of our theory, and the clarity of our writing. We are delighted to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses.\n\n---\n\n>**[W1]** In your experiments, did you tune the peak LR for all schedules? Can the FSL also predict the loss curves across peak learning rate?\n\n**[E1]** Due to the high computational cost and limited resources, we could not perform extensive tuning in our LLM experiments. The primary goal of these experiments was to demonstrate FSL's potential in predicting pre-training loss curves. Developing a practical version of FSL remains a key focus of our ongoing work.\n\nIn SGD experiments on linear regression, we found that FSL predictions generalize well across different peak learning rates. With hyperparameters $M = 256$, $s = 1.5$, $\\beta = 3$, and $10^4$ training steps, we predicted risk curves for the WSD and cosine schedules with a peak learning rate of $10^{-2}$, using coefficients fitted on WSD with a peak learning rate of $10^{-3}$. The resulting RMSE for these predictions was $1.52 \\times 10^{-4}$ and $2.31 \\times 10^{-4}$, respectively. For comparison, the RMSE with coefficients fitted at $10^{-2}$ was $1.45 \\times 10^{-4}$ and $2.34 \\times 10^{-4}$. These results show that FSL generalizes well across peak learning rates.\n\n$ $\n>**[W2]** Hence, in order to properly assess whether the optimal schedule brings a benefit, the baseline should be cosine and WSD schedules that decay to zero.\n\n**[E2]** We clarify that we have already verified the effect of the minimal learning rate in additional experiments, as shown in Appendix, Figure 10. Our results demonstrate that the optimal schedule provides a benefit, even when compared to WSD which exponentially decays to $10^{-7}$. We will add the discussion of minimal learning rates in the revised version.\n\n$ $\n>**[W3]** It has been shown empirically and theoretically that different schedules will have different optimal peak LRs, hence the missing tuning of the peak LR might confound your findings.\n\n**[E3]** In our SGD experiments, we did tune the peak learning rates for different LRSs, as shown in Figure 2(b) (right). Therefore it is a fair comparison with peak learning rates into consideration, and it shows that good schedules result in a lower final risk.\n\nHowever, for the LLM experiments, tuning the peak learning rate for each LRS would be too costly. As our work primarily focuses on theoretical analysis, we did not have sufficient computational resources to conduct these additional experiments.\n\n$ $\n>**[W4]** A minor weakness regarding clarity:...\n\n**[E4]** We clarify that this work primarily focuses on theoretical analysis, and implementation details were placed in the appendix due to space constraints. In the revised version, we will incorporate the relevant information into Section 6 to improve clarity and accessibility.\n\n$ $\n>**[W5]** Another minor point:.... \n\n**[E5]** There is a typo in line 173 causing the confusion. The $\\tau$ stands for the **continuous** training steps, and the true training step of SGD should be $\\tau / h$. $\\varphi$ is defined on the continuous training steps, which is $h$ times the discrete steps as defined in line 169.\n\n$ $\n>**[Q1]** You claim (line 93) that you provide the first theoretical explanation for different scaling behaviours across LR schedules. \n\n**[A1]** We clarify that [1] shows that the WSD schedule yields a logarithmic factor improvement over constant LR, but only in terms of worst-case upper bounds on the convergence rate for non-smooth convex problems. Since the comparison is between two upper bounds, it does not provide a theoretical justification for the benefit of WSD.\n\nWe also highlight the distinction between scaling laws and “convergence rates” in optimization. The latter typically defines a family of problems (e.g., non-smooth convex problems) and derives an upper bound on the convergence rate that holds uniformly across the family. A matching lower bound may be established to demonstrate minimax optimality. However, such worst-case bounds do not necessarily reflect the actual convergence behavior for a specific problem. In contrast, scaling laws are derived for a fixed problem instance, enabling a more accurate characterization of convergence behavior.\n\n$ $\n>**[Q2]** The only reference I know on this is https://arxiv.org/abs/2406.19146, …\n\n**[A2]** Indeed, in their work [11] they did not compare the scaling exponent of constant and cosine schedule directly, and their training dataset is different from Chinchilla, thus the comparison is unreasonable.\n\n$ $\n>**[Q3]** Can you analyse the WSD schedule with a linear cooldown (or other forms than exponential in general)? \n\n**[A3]**  We can analyze the linear cooldown for the WSD schedule with a fixed ratio of linear cooldown (e.g., a 20% cooldown ratio). If the minimal learning rate is set to zero, we have the following:\n\n* If $\\beta \\geq 2$:  $\\mathcal{E}_K \\asymp M^{-s\\beta} + T^{-s} + \\sigma^2 \\frac{1}{T^{\\frac{1}{2}}}.\n$\n\n* If $\\beta \\in (1,2)$:   $\\mathcal{E}_K \\asymp M^{-s\\beta} + T^{-s} + \\sigma^2 \\frac{\\min\\left(1, (M^{-\\beta} T)^{ \\frac{1}{2} - \\frac{1}{\\beta}}\\right)}{T^{\\frac{1}{\\beta}}}.\n$\n\nFrom these expressions, we can proceed to derive both the data-optimal scaling and the compute-optimal scaling.\n\n$ $\n>**[Q4]** Caption of Figure 2 ...\n\n**[A4]** We apologize for the confusion and will clarify the caption. The derivation of the “theoretical predictions” $D^{-\\alpha}$ consists of two steps: (1) SDE modeling of SGD, and (2) deriving $D^{-\\alpha}$ from FSL. In Figure 2(a), the solid line shows the prediction from simulating FSL, validating the second step. In Figure 2(b), the solid line represents the prediction from simulating SGD, confirming the accuracy of the SDE modeling.\n\n$ $\n>**[Q5]** The FSL theory applies to SGD, but is then used to fit Adam runs in the LLM setup. Could you comment on this limitation?\n\n**[A5]** We clarify this discrepancy from three perspectives:\n\nFirst, while our FSL theory is based on SGD—common in prior studies on scaling laws [2–4]—we view this as a foundational step. Our contribution builds on this by incorporating the influence of learning rate schedules, potentially offering insights into how adaptive methods like Adam improve scaling efficiency.\n\nSecond, the fact that our FSL, derived under SGD assumptions, still fits Adam-based LLM runs reasonably well is a surprising and intriguing empirical observation. It suggests that despite Adam’s adaptive nature, its behavior may resemble that of SGD in some respects. This aligns with recent findings: (1) SGD variants can match Adam's performance on training Transformer models [5–7,10], though scale is still limited; (2) most LLM parameters can be trained effectively using SGD [8].\n\nThird, the main practical limitation is modeling the interaction between learning rate and batch size. For SGD, the learning rate is typically scaled linearly with batch size, while for Adam, square-root scaling is recommended, informed by the signSGD interpretation of Adam. Extending our FSL framework to accurately predict batch size scaling for adaptive optimizers like Adam is complex and remains a key area for future work.\n\nWe will incorporate these into the discussion section of the revision of our paper.\n\n$ $\n>**[Q6]** You are fitting the values of $s$ and $\\beta$ …\n\n**[A6]** We clarify that in our LLM experiments, we focused solely on the learning rate schedule and did not fit the scaling with respect to model size. Our FSL framework, derived from kernel regression, suggests a model-size scaling of $M^{-s\\beta}$, indicating that the scaling exponent for model size depends on data size scaling. However, we expect this relationship may not hold in practice for LLM tasks.\n\nNonetheless, our theoretical results in Section 5 (lines 237, 262, 305) consistently show that, across all three learning rate schedules, the compute-optimal strategy prioritizes increasing data size over model size. Specifically, the data-size scaling exponent slightly exceeds the model size exponent, aligning qualitatively with the Chinchilla law (data: 0.54, model: 0.46).\n\n$ $\n>**[Q7]** The emphasis in the abstract on batch size schedule is very strong…\n\n**[A7]** Due to the space limit, we refer to the response **E2** for **W2** of Reviewer SA1x.\n\n$ $\n>**[Q8]** line 245 says …\n\n**[A8]** We clarify that the Chinchilla law [8] asserts that the power coefficient in data is slightly higher than that in the model, which is consistent with our theoretical results. Specifically, in the Chinchilla paper, the fitted scaling law is (equation (10) in [8])\n$$ \nL(M, D) = L_0 + \\frac{A}{M^{0.34}} + \\frac{B}{D^{0.28}}.\n$$\n\n$ $\n>**[Q9]** Minor typos.\n\n**[A9]** We will correct them in the revised version of our paper. The constant factor in line 98 is multiplicative.\n\n---\n**Reference**\n\n[1] The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training.\n\n[2] Scaling Laws in Linear Regression: Compute, Parameters, and Data.\n\n[3] 4+3 Phases of Compute-Optimal Neural Scaling Laws.\n\n[4] HOW FEATURE LEARNING CAN IMPROVE NEURAL SCALING LAWS.\n\n[5] Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be.\n\n[6] No more Adam: Learning rate scaling at initialization is all you need.\n\n[7] Adam-mini: Use Fewer Learning Rates To Gain More.\n\n[8] Deconstructing what makes a good optimizer for autoregressive language models.\n\n[9] Training compute-optimal large language models.\n\n[10] Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful.\n\n[11] Resolving Discrepancies in Compute-Optimal Scaling of Language Models."}, {"review_id": "yK1jiiyhcU", "reviewer": "Reviewer_Ss9R", "summary": "This paper proposes functional scaling laws, and extension of existing LLM scaling laws, to account for learning rate schedules. This novel scaling law additionally allows analysis of loss during the intermediate-steps during training, whereas existing scaling laws only account for the final-step loss. The authors leverage the SDE view of SGD to provide theoretical analysis based on teacher-student kernel regression setting for three different learning rate schedules.", "questions": "- Is it possible to consider generalization bounds as in [1]? Additionally it would be nice to add discussion on this.\n\n\n[1] Finzi, Marc, et al. \"Compute-Optimal LLMs Provably Generalize Better with Scale.\" https://arxiv.org/abs/2504.15208.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The field of scaling laws is outside the scope of my expertise, so I primarily focused on the theory part of the paper.\n\n**Strengths**: \n- The paper is well-written and straightforward to follow.\n- The theoretical results are well-justified: the proofs are convincing at a high level, even though I did not check every detail.\n- Experimental validation matches the theory.\n- Discussion and limitations of assumptions are well discussed.\n\n\n**Weaknesses**: \n- Empirical results are on relatively small models, but this limitation is discussed and the paper's contribution is mainly theoretical.", "was_revised": true, "final_justification": "Scaling laws are outside of my expertise, but the author's application of SDE analyses to this field is thorough and novel. Having read the other reviews and replies, I'm happy to keep my positive score.", "rebuttal": "We sincerely thank the reviewer for the positive support and valuable feedback! We greatly appreciate the insightful review, and the recognition of highlighting the significance of our contribution and the strength of our theory and experiments, as well as the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.\n\n---\n\n>**[Q]** Is it possible to consider generalization bounds as in [1]? Additionally it would be nice to add discussion on this.\n\n**[A]** We thank the reviewer for the valuable comment. Our work focuses on the online (one-pass) SGD setting, where each data point is used only once. This reflects the standard practice in modern LLM pre-training, where massive datasets are typically passed through only once due to their scale. In this setting, traditional notions of in-distribution generalization—such as the gap between training and validation loss—are less relevant. Instead, generalization is often assessed in terms of out-of-distribution performance on downstream tasks. As a result, there is effectively no distinction between training and validation loss during pre-training, and traditional in-distribution generalization bounds do not directly apply. We hope this clarification addresses the reviewer’s concern.\n\n---\n**Reference**\n\n[1] Compute-Optimal LLMs Provably Generalize Better with Scale."}]

Author rebuttal:

Meta-review: (a) This paper develops a theory to predict the scaling behavior for different learning rate schedules, with applications to LLMs. 

(b) Strenghts: currently, the design of learning rate schedulers is more of an art than a science, making this paper a valueable contribution to the community to improve our understanding of scaling behavior. Theory is complemented with strong empirical evidence. All reviewers find the paper easy to read. 

(c) Weaknesses: some comparisons with recent works were missing, but were provided during rebuttal. There were some concerns about the hyperparameter choices for the experiments, but seeing as this is a theoretical paper mostly, these concerns are minor. 

(d) Accept. The paper is important and on a timely topic, likely with high impact. 

(e) Some of the asked ablations were carried out, and several questions were clarified.
Decision: Accept (spotlight)