Title: Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules
Venue: NeurIPS 2025
Accepted: yes
Keywords: Scaling Laws

Learning Rate Schedule

Kernel Regression

LLM Pre-Training

Loss Dynamics
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.


Peer reviews:
comments: 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.

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.

**Strengths**: 
- The paper is well-written and straightforward to follow.
- The theoretical results are well-justified: the proofs are convincing at a high level, even though I did not check every detail.
- Experimental validation matches the theory.
- Discussion and limitations of assumptions are well discussed.


**Weaknesses**: 
- Empirical results are on relatively small models, but this limitation is discussed and the paper's contribution is mainly theoretical.

QUESTIONS:
- Is it possible to consider generalization bounds as in [1]? Additionally it would be nice to add discussion on this.


[1] Finzi, Marc, et al. "Compute-Optimal LLMs Provably Generalize Better with Scale." https://arxiv.org/abs/2504.15208.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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?\n\nQUESTIONS:\n* 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\"\n\nLIMITATIONS:\nNo potential negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nScaling 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}