PaperlayerView exact revision on arXiv

Source-linked AI summary

Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules

Binghui Li, Fengling Chen, Zixun Huang, Lean Wang, Lei Wu

arXiv:2509.19189v4cs.LGstat.ML

TL;DR

Existing theory has limited systematic characterization of how capacity, task difficulty, and learning-rate schedules shape scaling efficiency. This paper analyzes power-law kernel regression with intrinsic time and a Functional Scaling Law, finding that task accumulation progressively extends a power-law regime.

  • Problem

    Existing theory lacks a systematic characterization of how model capacity, task difficulty, and learning-rate schedules shape scaling efficiency.

    Source: §1 · p. 2

  • Method

    The paper analyzes power-law kernel regression using intrinsic time and a Functional Scaling Law whose schedule influence enters through a convolution-type functional.

    Source: §7 · p. 17

  • Results

    As task count increases, the power-law regime progressively extends and, in the idealized M → ∞ limit, spans the entire training process.

    Source: §4.4 · p. 12

  • Takeaways & Limitations

    The framework offers a theoretical account of how learning-rate schedules shape scaling and of WSD's effectiveness in LLM pre-training.

    Source: §7 · p. 17

    The optimal learning-rate schedule under resource constraints remains largely unresolved, and FSL needs refinement through extensive large-scale LLM experiments.

    Source: §7 · p. 17

Abstract

from arXiv · show

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 loss dynamics obey similar laws and, crucially, how the 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 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 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.

1 Introduction

Existing theory leaves open how capacity, task difficulty, and learning-rate schedules shape scaling efficiency, and whether scaling laws describe the full loss trajectory. The paper studies these questions in power-law kernel regression, introducing an intrinsic-time Functional Scaling Law and deriving schedule-specific scaling results.

1 Introduction

  • Prior work lacks a systematic account of how model capacity, task difficulty, and learning-rate schedules shape scaling efficiency, or whether scaling extends across the full loss trajectory.Classical scaling laws address end-of-training loss, while theoretical justification for trajectory-wide scaling remains missing.Source: §1 · p. 2

1.1 Our Contribution

  • The Functional Scaling Law uses intrinsic time to capture full loss dynamics, with learning-rate schedules entering through a convolutional functional for gradient-noise accumulation.The forgetting kernel describes how injected noise dissipates during training.Source: §1.1 · pp. 2–3
  • WSD has the best scaling efficiency, followed by exponential decay and then constant schedules.The analysis derives final-loss scaling relations for these schedules in both data-limited and compute-limited regimes.Source: §1.1 · p. 3
  • Higher-capacity models are more efficient in both compute and data, while compute-optimal training scales data more than model size and requires peak learning rate to scale with budget.The peak learning-rate result applies to training budgets measured by data or compute.Source: §1.1 · p. 3

1.2 Related Work

  • The framework unifies prior power-law-kernel analyses by extending from specific schedules and final-step loss to general schedules and full loss dynamics.The paper also connects its FSL to trajectory fitting in LLM pre-training and notes a near-equivalence to the multi-power-law model after integration by parts.Source: §1.2 · p. 4

2 Power-Law Kernel (PLK) Regression

The paper studies a width-M model that projects features into a lower-dimensional subspace and learns a target function under power-law spectral assumptions. It distinguishes model capacity from task difficulty, including how a fixed target can appear easier under a higher-capacity model.

2 Power-Law Kernel (PLK) Regression

  • The regression setup uses a feature map with covariance H and independent Gaussian label noise, and studies width-M models formed by projecting N-dimensional features.The analysis considers both top-M feature selection and random projections.Source: §2 · pp. 4–5

2.1 Model Capacity and Task Difficulty

  • Model capacity is represented by the feature-spectrum exponent β, with eigenvalues satisfying λ_j ≂ j^-β for β > 1.Model size M controls the number of retained features, while β controls their decay in importance.Source: §2.1 · p. 5
  • The target's intrinsic difficulty is α = sβ, independent of the model spectrum; for a fixed target, higher capacity (smaller β) increases relative difficulty parameter s and makes the task easier.The target's spectral energy decays as a power law under the stated assumptions.Source: §2.1 · pp. 5–6

3 One-Pass SGD and Intrinsic-Time SDE

The section formulates one-pass mini-batch SGD and approximates it with a continuous-time stochastic differential equation. It then reparameterizes time to reflect learning-rate-adjusted training progress, simplifying how schedules enter the dynamics.

3 One-Pass SGD and Intrinsic-Time SDE

  • The method minimizes R(v) via one-pass SGD, using mini-batches of i.i.d. samples with learning-rate schedule η and batch sizes Bk; initialization is v0 = 0.Source: §3 · p. 6
  • In the small-step-size limit, the discrete SGD iteration is approximated by an Itô-type SDE, where learning rate affects both drift and diffusion.This continuous-time analysis is used to make the scaling-law analysis more tractable.Source: §3 · pp. 6–7
  • Intrinsic time tracks accumulated learning-rate progress rather than update count, and the reparameterized SDE absorbs schedule dependence from the drift into its diffusion term.The schedule factor γ(t) captures the joint effect of learning-rate and batch-size scheduling, decoupling deterministic and stochastic effects.Source: §3 · p. 7

4 Intrinsic-Time Functional Scaling Laws

The section develops intrinsic-time functional scaling laws for SGD loss trajectories, first interpreting their terms and then extending the characterization across task difficulty, label noise, and feature selection. Its proof links power-law loss behavior to scale-invariant noise forgetting and the accumulation of many learning tasks.

4.1 An Illustrative Setting and the Underlying Scaling Behavior

  • In the hard-learning regime, the FSL characterizes loss through label-noise irreducible error, finite-size approximation error, signal learning, and schedule-shaped noise accumulation.The noise term convolves the learning-rate schedule with the instantaneous noise magnitude and a forgetting kernel.Source: §4.1 · pp. 7–8
  • Increasing model capacity accelerates signal learning but slows noise forgetting, so effective schedules balance noise suppression against enough intrinsic time for learning and forgetting.For a fixed target, higher capacity also makes the task relatively easier by increasing s = α/β.Source: §4.1 · p. 8

4.2 General Results: The Top-M Feature Case

  • For top-M features, the FSL extends to all s > 0 and β > 1, uniformly in label-noise level, exactly through intrinsic time cM^β and with two-sided control thereafter.At that critical time, signal learning reaches the approximation-error limit; later dynamics may add or dissipate noise while remaining bounded.Source: §4.2 · pp. 8–9
  • In the zero-label-noise case, a top-M FSL holds for s ∈ [0, 2 − 1/β], and the hard-regime finite-M functions can be replaced by infinite-width counterparts.Both eM and KM exhibit power-law decay for 1 ≲ t ≲ M^β.Source: §4.2 · pp. 8–9

4.3 General Results: The Random-M Feature Case

  • With random-M features, the top-M FSL characterization transfers with probability at least 1 − exp(−Ω(M)) for s ∈ (0, 1], while easier tasks may differ.The transfer follows alongside random-feature eigenvalues matching the power-law spectrum up to constant factors; behavior for s > 1 is left open.Source: §4.3 · p. 9

4.4 Key Proof Steps and Core Insights

  • Under power-law assumptions, the expected-loss equation admits sharp asymptotic bounds, and scale invariance of the forgetting kernel controls higher-order Volterra terms by the first-order term.The kernel’s power-law decay yields a subconvolution property over the relevant time range.Source: §4.4 · p. 11
  • Collective task accumulation produces power-law behavior even though each individual task has an exponential learning curve; the power-law regime expands with task count.In the idealized limit M → ∞, it spans the entire training process.Source: §4.4 · p. 12

5 Learning Rate Schedules Impact Scaling Efficiency

The paper derives data- and compute-limited scaling laws for constant, exponential-decay, and WSD schedules. The comparisons show efficiency gains from learning-rate decay and an additional improvement from WSD.

5.1 Constant LRS

  • Higher-capacity models are more data-efficient under constant schedules, and compute-optimal training favors data over model size.The optimal effective learning rate decreases as compute grows.Source: §5.1 · p. 13

5.2 Exponential Decay LRS

  • With exponential decay ending at a/K, intrinsic training time is only logarithmically smaller than under a constant schedule with step size a.The schedule can lower the learning rate to a/K while sacrificing only that logarithmic factor.Source: §5.2 · p. 14
  • Exponential decay yields strictly faster excess-risk decay than a constant schedule in the data-limited regime.In the hard-learning regime, the optimal maximal learning rate remains constant and the compute-optimal split depends on model capacity, not task difficulty.Source: §5.2 · p. 14
  • For exponential decay, the easy-learning excess-risk rate depends only on intrinsic task difficulty, whereas higher capacity improves compute efficiency in the hard-learning regime.In the easy regime, increasing model capacity alone gives no asymptotic gain.Source: §5.2 · p. 14

5.3 WSD-like LRS

  • A WSD stable phase boosts intrinsic training time and removes logarithmic factors from the full-batch gradient-descent term without changing the noise-term order when r > 0.The decay phase need only scale sublinearly with data; the paper notes that harder tasks may favor a smaller decay fraction.Source: §5.3 · p. 15
  • WSD improves excess-risk decay by an additional logarithmic factor over exponential decay in both task regimes.In the hard-learning regime, it eliminates the logarithmic factor entirely.Source: §5.3 · p. 15

6 Experiments

Experiments test FSL against discrete-time SGD in power-law kernel regression and evaluate it as a surrogate for LLM training. FSL captures trajectories and scaling behavior, fits and predicts LLM loss curves, and identifies a schedule with lower final loss than the baselines.

6.1 Power-Law Kernel Regression

  • The kernel-regression test uses difficulty s = 0.5 and capacity β = 4, a hard-learning regime.The comparison evaluates FSL against discrete-time SGD.Source: §6.1 · p. 15
  • FSL accurately describes discrete-time SGD loss trajectories under cosine, WSD, and cyclical schedules, including WSD’s slower stable-phase decline and sharper decay-phase drop.WSD ultimately reaches a lower final loss than cosine in this comparison.Source: §6.1 · p. 15
  • WSD gives the best final-loss scaling in power-law kernel regression, followed by exponential decay and then the constant schedule.The observed final-step losses closely follow FSL predictions.Source: §6.1 · p. 16

6.2 LLM Pre-training

  • On 400M- and 1B-parameter LLaMA models trained on 20B tokens, FSL fitted to 8-1-1 loss curves reliably predicts unseen cosine and WSD curves.The fitting and prediction use post-warmup trajectories.Source: §6.2 · p. 16

7 Concluding Remarks

The paper presents FSL as a framework for analyzing how learning-rate schedules affect scaling in power-law kernel regression, with analyses of three schedules supporting prevailing LLM training practices. It identifies unresolved theory and experimentation needed to extend the framework.

7 Concluding Remarks

  • The framework characterizes learning-rate schedule effects on scaling in power-law kernel regression and uses analyses of three schedules to explain WSD’s effectiveness.The authors describe this as theoretical support for practices in LLM pre-training.Source: §7 · p. 17
  • Optimal learning-rate schedules under resource constraints remain largely unresolved, and making FSL practical requires extensive large-scale LLM experiments.The authors also identify batch-size scheduling as an open topic.Source: §7 · p. 17

A.1 Empirical Fitting of LLM Pre-training Loss Trajectory

The appendix introduces empirical loss-trajectory models and common LLM learning-rate schedules, then relates kernel assumptions to function smoothness and model capacity. It also describes an SDE formulation of SGD and notes that the appendix contains experiment details and additional results.

A.1 Empirical Fitting of LLM Pre-training Loss Trajectory

  • The Chinchilla Law models final-step loss, while the Momentum Law [61] and Multi-Power Law [40] model the full training trajectory.The Multi-Power Law adds power laws to better capture progressive loss reduction under learning-rate decay.Source: §A.1 · p. 24

A.2 Popular Learning Rate Schedules in LLM Pre-training

  • The appendix describes cosine, WSD, and multi-step schedules used in LLM pre-training.WSD has warm-up, stable, and decay phases; the 8-1-1 schedule divides training into three stages.Source: §A.2 · p. 24

A.3 Connections to Kernel Regression

  • Kernel source and capacity conditions connect target-function smoothness and kernel eigenvalue decay to task difficulty and model capacity.Smaller s indicates a more difficult problem, while smaller β implies higher model capacity.Source: §A.3 · p. 26

A.4 The SDE Modeling

  • Assuming Gaussian gradient noise turns the SGD update into an Euler–Maruyama discretization of an Itô SDE.The continuous formulation uses learning-rate and batch-size schedule functions.Source: §A.4 · p. 26

B Experiment Details and Additional Results

B.1 Power-Law Kernel Regression

The experiments test whether the continuous-time Functional Scaling Law (FSL) captures discrete-time SGD risk trajectories and predicted scaling behavior. They fit the law to averaged SGD runs and evaluate its predictions under different learning-rate schedules.

Fitting FSL on SGD Average-Risks

  • FSL closely fits the average risk trajectories from 200 independent SGD runs across the tested configurations.The fits use the physical-time FSL and average trajectories from runs with varying s and β.Source: §B.1 · p. 27, Figure 3 · p. 28

Scaling law experiments

  • The scaling-law experiments numerically discretize the FSL and test predictions for constant, exponential-decay, and WSD schedules.The experiments are designed to evaluate the scaling laws predicted by the analysis.Source: §B.1 · pp. 27–28

SGD experiments

  • The SGD experiments use batch size 1 and run for D steps, with ηmax = 0.05D^-r under each corresponding schedule.The exponent r uses the same theoretical values as in the FSL experiments.Source: §B.1 · p. 28

B.2 LLM pre-training

The section develops a practical FSL for fitting and predicting LLM loss curves, then tests its fits and numerically optimizes learning-rate schedules. It also presents the Volterra loss equation and power-law kernel estimates underlying the theory.

B.2 LLM pre-training

  • The practical FSL ansatz approximates the theoretical law for fitting LLM pre-training loss curves, with constants and exponents fitted to observed losses.For fixed batch size, the discrete ansatz uses cumulative learning-rate time and a learning-rate-decay functional.Source: §B.2 · pp. 28–29
  • The FSL-optimal LRS is obtained by minimizing fitted-FSL final-step loss over schedules, then evaluated against cosine, WSD, and 8-1-1 baselines.The QwenMoE evaluation holds model and training conditions fixed across schedules.Source: §B.2 · p. 30
  • FSL accurately models loss trajectories for a 100M LLaMA model across 17k, 34k, 68k, and 134k training steps.The tested durations correspond to 10B, 20B, 40B, and 80B tokens, respectively.Source: §B.2 · pp. 30–31, Figure 6 · p. 32, Figure 7 · p. 32
  • Across models from 100M to 1B and GPT-2, LLaMA, and QwenMoE architectures, FSL predictions align with observed performance, while FSL-optimal schedules beat tested WSD variants in final loss.The comparisons include 1B LLaMA and 100M GPT-2 models, and WSD schedules with different decay ratios and minimum learning rates.Source: §B.2 · pp. 30–31, Figure 4 · p. 31, Figure 5 · p. 31

C.1 Volterra Integral Equation Governing the Loss Dynamics

  • The intrinsic-time SDE yields a Volterra integral equation that characterizes expected-loss evolution; its gradient-noise covariance scales with risk and aligns with local curvature.The noise includes a minibatch-dependent fit term and a label-noise term.Source: §C.1 · pp. 31, 34

C.2 Power-law Decay of the eM and KM Functions

  • For 1 ≲ t ≲ M^β, the kernel function eM(t) decays as t^-s, and the section also bounds the gap between eM and e∞.The analysis establishes power-law decay behavior for eM and KM.Source: §C.2 · pp. 35–36

C.3 The Case of Top-M Features

  • In the top-M feature case, the analysis derives a Volterra equation and applies the general FSL to constant-noise, noiseless, and hard-regime cases.The derivation uses the top-M projection structure.Source: §C.3 · p. 36

C.3.1 Proof of Theorem 4.3

These appendix sections derive and bound the Volterra formulations for top-M and random-M features, then use concentration results to support the random-feature analysis. The proofs also specify the constant-label-noise condition used to simplify one FSL case.

C.3.1 Proof of Theorem 4.3

  • For top-M features, the loss equation follows from the projected dynamics and is expressed using eM and KM; the proof bounds the convolution terms to establish the result.The argument uses properties of the decreasing kernel and its integral.Source: §C.3.1 · pp. 37–38

C.3.2–C.3.4 Proofs of Theorems 4.4, 4.5, and 4.2

  • The constant-label-noise proof uses σ ≳ 1, while the noiseless proof sets σ² = 0; the hard-regime proof requires β and γmax to satisfy its stated bounds.The constant-noise condition allows the fit-dependent noise term to be absorbed by the full-batch GD term.Source: §C.3.2 · p. 39, §C.3.3 · p. 39, §C.3.4 · p. 41

C.4 The Case of Random-M Features

  • For random-M features, a similar Volterra equation holds with probability at least 1 − e^-Ω(M), using eigenvalues of the random matrix WHW⊤.The proof strategy relies on the random-matrix spectrum closely matching the top-M truncation.Source: §C.4 · p. 41

C.4.1 Concentration Inequalities

  • The random-feature argument invokes concentration results for power-law eigenvalues and random projections, including bounds that hold with probability at least 1 − e^-Ω(M).The cited concentration lemmas assume λj ≂ j^-β with β > 1.Source: §C.4.1 · p. 42

C.4.2 Upper and Lower Bounds

  • The upper- and lower-bound analysis compares random-feature terms with eM(t), including a bound of order min{t^-s, 1} + M^-sβ.The displayed bound is stated with probability at least 1 − e^-Ω(M).Source: §C.4.2 · pp. 43–44

C.4.3 Proof of Proposition C.11

  • Combining the random-feature bounds yields the stated high-probability estimates, after which the top-M proof carries over because the random eigenvalues satisfy bλj ≂ λj.The Volterra equation for random features follows from this spectral comparison.Source: §C.4 · p. 41, §C.4.3 · p. 47

D Proofs for Section 5

  • The constant-label-noise case simplifies the FSL by absorbing the fit-dependent noise term into the full-batch GD term when σ ≳ 1.The authors focus on this condition to streamline the exposition, noting that extending the result to all σ ⩾ 0 is more involved.Source: D Proofs for Section 5 · p. 48

D.1 Proofs for Constant LRS

The section derives data- and compute-optimal strategies for constant learning rates by minimizing final population risk. The resulting rates balance data, model size, and effective learning rate under the specified budgets.

Data-optimal strategy

  • Eopt ≂ D^-s/(s+1) + M^-sβ for the data-limited setting, combining a data-dependent term with a model-size term.The optimal effective learning rate is γopt ≂ D^-s/(s+1).Source: §D.1 · pp. 48–49
  • For constant learning rates, the data-optimal loss rate is Eopt ≂ D^-s/(s+1).The rate follows from balancing the terms in the final-risk expression.Source: §D.1 · p. 49

Compute-optimal strategy

  • Under compute budget C, the compute-optimal strategy scales γopt ≂ C^-sβ/(1+β+sβ) and Mopt ≂ C^1/(1+β+sβ).Compute is defined as C = MKB, and the strategy also optimizes data size D := BK.Source: §D.1 · p. 49

D.2 Proof for The Exponential-Decay LRS

For exponential-decay learning rates, the analysis uses total intrinsic training time to express the schedule and derives data- and compute-optimal population-risk scaling. The compute-optimal risk decreases at a rate proportional to (C/log C)^(-sβ/(2+sβ)).

D.2 Proof for The Exponential-Decay LRS

  • The total intrinsic training time for exponential decay is T = (a − b)K/log(a/b).The intrinsic-time transform is tied to ζ = log(a/b)/K.Source: §D.2 · pp. 49, 52
  • When s ⩽ 1 − 1/β, the data-optimal strategy has γopt ≂ 1 and Eopt ≂ (D/log D)^(-s).For the other data-regime expression, the supplied theorem gives Eopt ≂ (D/log D)^(-sβ/(1+sβ)).Source: §D.2 · p. 52
  • The data-optimal analysis concludes that Mopt = ∞ achieves the optimal rate across the two cases.The proof also gives γopt ≂ 1 for the case s ⩽ 1 − 1/β.Source: §D.2 · pp. 52–53
  • Eopt(C) ≂ (C/log C)^(-sβ/(2+sβ)) under the compute-optimal strategy.The theorem formulates the strategy in terms of effective maximum learning rate, model size, and data size.Source: §D.2 · p. 53

D.3 Proof for the WSD-Like LRS

The WSD-like schedule analysis derives data- and compute-budget scaling strategies and requires sufficiently large intrinsic training times. It also presents auxiliary PSD-matrix lemmas used in the proofs.

D.3 Proof for the WSD-Like LRS

  • The WSD-like schedule uses total intrinsic time T = aK1 + (a−b)K2/log(a/b), with decay-phase time T2 = (a−b)K2/log(a/b); both must be ≳ 1.Source: §D.3 · p. 55
  • For s < 1 − 1/β, the data-optimal strategy has γopt ≂ 1, (D1)opt ≂ D, and (D2)opt ≳ D.The proof summary states Mopt = ∞ achieves the optimal rate, with Eopt ≂ D^-s.Source: §D.3 · pp. 56–57
  • Eopt(C) ≂ (C/log C)^(-sβ/(2+sβ))(log C)^s for the WSD-like compute-budget analysis.The derivation sets the stable-phase data allocation to D1 ≂ D and defines D̃2 = D2 log K.Source: §D.3 · pp. 57–58

E Auxiliary Lemmas

  • The auxiliary lemmas use PSD ordering to support trace comparisons with a PSD weighting matrix.The stated setup is P ⪯ Q and PSD U; the proof notes U and Q − P are PSD.Source: E Auxiliary Lemmas · p. 60

Related papers

Loading 2509.19189v4…