PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay

Binghui Li, Zilin Wang, Fengling Chen, Shiyang Zhao, Ruiheng Zheng, Lei Wu

arXiv:2602.06797v3stat.MLcs.LG

TL;DR

The paper asks how to design learning-rate schedules that optimize final-step loss under FSL, fixed training horizons, and stability constraints. It characterizes schedules in power-law kernel regression and studies fractional profiles and one-pass SGD. Easy tasks favor early decay, hard tasks favor WSD-like delayed decay, sufficiently fast terminal decay preserves optimal rates, and FSL-motivated power decay achieves optimal last-iterate rates.

  • Problem

    Learning-rate scheduling lacks a systematic theory explaining when decay should begin and how decay shape affects training under competing signal-learning and noise-forgetting dynamics.

    Source: §1 · pp. 1–2

  • Method

    The paper uses FSL to optimize final-step excess risk over schedules under a fixed training horizon and stability constraint, then analyzes fractional schedules and one-pass SGD.

    Source: §1 · pp. 2–3

  • Results

    Easy tasks favor power decay from the start, hard tasks favor WSD-like delayed decay, sufficiently fast terminal decay attains optimal rates, and power decay achieves optimal last-iterate rates.

    Source: §1 · pp. 2–3, §9 · p. 16

  • Takeaways & Limitations

    Task difficulty determines when decay should begin, while model capacity determines the decay shape; exact shape tuning is unnecessary when terminal decay is sufficiently fast.

    Source: §1 · pp. 2–3, §9 · p. 16

    Slow terminal decay can cause schedule-induced capacity saturation, and extending the framework to practical pretraining requires incorporating additional effects.

    Source: §5.1 · p. 9, §9 · p. 16

Abstract

from arXiv · show

We study optimal learning rate (LR) schedules under the functional scaling law (FSL) framework (Li et al., 2025), which decomposes training dynamics into signal learning and noise forgetting. In power-law kernel regression, these two components are governed by a source exponent $s>0$ and a capacity exponent $q>1$, respectively, with smaller $s$ corresponding to harder tasks. For a fixed training horizon $N$, we characterize the schedules that minimize the final-step loss under a stability constraint and reveal a sharp phase transition. In the easy-task regime $s>1-1/q$, the optimal schedule follows power decay from the beginning of training; in the hard-task regime $s<1-1/q$, it becomes warmup-stable-decay (WSD)-like (Hu et al., 2024), staying at the largest admissible LR for most of training before a final decay. In both regimes, the decay exponent is $2q-1$: task difficulty determines when to decay, while model capacity determines how to decay. Beyond the exact optimum, we study fractional schedules, whose shape is defined over relative training progress. We show that precise tuning of the decay shape is often unnecessary: a broad class of profiles attains the optimal convergence rate, while overly slow terminal decay leads to schedule-induced capacity saturation. Finally, for one-pass SGD in kernel regression, FSL-motivated power-decay schedules achieve optimal last-iterate rates. Experiments support the theoretical predictions and the task-dependent transition between early and delayed decay.

1 Introduction

The paper uses FSL to design learning-rate schedules that balance signal learning against noise injection and forgetting under a fixed horizon and stability constraint. It predicts task-dependent timing of decay, rate-robust fractional schedules, and optimal last-iterate rates for one-pass kernel-regression SGD.

1 Introduction

  • FSL frames learning-rate design as balancing faster progress from large rates against their greater gradient-noise injection.The schedule is optimized through an explicit map from learning-rate profiles to loss dynamics.Source: §1 · pp. 1–2
  • The paper defines the optimal schedule as the one minimizing final-step excess risk under a fixed horizon and stability constraint.Source: §1 · p. 2
  • Easy tasks decay from the beginning, whereas hard tasks use WSD-like schedules that remain at the largest admissible rate before final power-law decay.The transition occurs at s_c = 1 − 1/q; task difficulty determines when decay begins, while capacity determines how it decays.Source: §1 · p. 2, Figure 1 · p. 3
  • Fractional schedules with sufficiently fast terminal decay achieve optimal rates, while slow decay causes schedule-induced capacity saturation.For (s, q) = (0.8, 5), power decay with p = 4.2 follows the optimal rate, whereas cosine decay with p = 2 follows the predicted saturated rate.Source: §1 · p. 3, Figure 1 · p. 3
  • Power-decay schedules achieve optimal last-iterate rates for one-pass kernel-regression SGD, removing logarithmic factors in the easy-task regime.Experiments also support the predicted transition from early to delayed decay, including beyond the kernel setting.Source: §1 · p. 3

2 Related work

The paper connects learning-rate schedule design to scaling-law and kernel-regression theory while addressing unresolved differences between cosine, WSD, and polynomial schedules. Its FSL analysis characterizes schedules directly rather than merely verifying prescribed profiles.

2 Related work

  • The work extends scaling-law analysis by using Li et al. (2025)'s FSL to characterize optimal learning-rate schedules.It positions schedule design within broader neural scaling-law research and contrasts with FSL-based batch-size scheduling by Wang et al. (2026).Source: §2 · pp. 3–4
  • Unlike prior optimal-control analyses that simplify SGD dynamics or verify prescribed schedules, FSL makes full schedule optimization tractable under power-law structure.Source: §2 · p. 4
  • The fractional-schedule analysis generalizes duration-aware properties attributed to cosine schedules and identifies when such profiles attain optimal rates.Source: §2 · p. 4
  • The paper explains differing empirical preferences for linear and concave decay by tying optimal decay shape to the forgetting-kernel profile.It also provides theoretical support for an approximately (1 − x)^1.5 WSD shape found by Luo et al. (2025).Source: §2 · p. 4
  • For one-pass kernel-regression SGD, FSL-derived power decay removes logarithmic factors previously incurred by last-iterate methods.Source: §2 · p. 4

3 Learning setup and functional description

The learning setup models one-pass SGD in power-law kernel regression and uses FSL to decompose excess risk into signal learning and noise forgetting. This decomposition turns schedule design into choosing both total intrinsic time and its temporal allocation.

3.1 Power-law kernel regression

  • The setup studies supervised kernel regression in feature space, trained by one-pass SGD and evaluated using excess risk.The model uses a feature map and an associated kernel covariance structure.Source: §3.1 · p. 5
  • Power-law structure uses capacity exponent q > 1 for spectral decay and source exponent s > 0 for target structure, with smaller s indicating harder tasks.The assumptions also include mixed fourth-moment control, satisfied for example by centered Gaussian features.Source: §3.1 · p. 5

3.2 Functional scaling laws

  • FSL decomposes expected excess-risk dynamics into a signal-learning function and a forgetting-kernel term for injected stochastic noise.The signal term describes noiseless target learning, while the kernel term describes how noise effects decay with subsequent intrinsic time.Source: §3.2 · pp. 5–6
  • The source exponent governs signal-learning speed, while the capacity exponent governs noise forgetting; higher-capacity models with smaller q forget noise more slowly.Source: §3.2 · p. 6
  • Intrinsic time is the cumulative learning-rate area, and the final-step objective exposes how each schedule controls signal learning, noise injection, and subsequent forgetting.For fixed total intrinsic time, schedule design reduces to distributing intrinsic time over training, although scale and temporal profile remain coupled in the full optimization.Source: §3.2 · p. 6

4 Optimal learning rate schedules

Under a fixed horizon and stability constraint, the optimal schedule minimizes final-step excess risk and undergoes a task-dependent transition between immediate and delayed decay. The decay profile is set by model capacity, while task difficulty determines when decay begins.

4 Optimal learning rate schedules

  • For easy tasks, s > 1 − 1/q, the optimal schedule follows power decay from the beginning of training.Source: §4 · p. 7
  • For hard tasks, s < 1 − 1/q, the peak learning rate saturates at the stability limit, producing a WSD-like schedule with a long stable phase before final power-law decay.Source: §4 · p. 7
  • At the boundary s = 1 − 1/q, scaling analysis does not distinguish immediate from delayed decay because problem-dependent constants determine whether stability is active.Source: §4 · p. 7
  • The stability-constrained optimization exhibits a sharp transition at s_c = 1 − 1/q between immediate and delayed decay.The transition occurs when the unconstrained peak learning rate reaches the stability threshold.Source: §4 · pp. 6–7
  • The decay exponent is 2q − 1 and is determined by the capacity exponent q, whereas the source exponent s determines when decay begins.Higher task difficulty means slower signal learning and favors maintaining a large learning rate longer.Source: §4 · p. 7

5 Fractional LR schedules: shape and rate optimality

Fractional schedules separate learning-rate scale from a shape defined over relative training progress, allowing peak-rate tuning while keeping the profile fixed. Their terminal decay controls whether rates remain optimal, while horizon-independent profiles face a general rate barrier.

5 Fractional LR schedules: shape and rate optimality

  • Fractional schedules define their shape over relative progress z/N and optimize total intrinsic time while holding the profile fixed.This separates the overall learning-rate scale from the schedule shape.Source: §5 · pp. 7–8

5.1 Terminal decay and convergence rates

  • Any fractional profile with terminal exponent p satisfying p + 1 > q attains the optimal convergence rate after tuning its intrinsic-time budget.The exact optimal terminal exponent 2q − 1 is therefore unnecessary for rate optimality.Source: §5.1 · p. 9
  • When p + 1 < q, terminal decay becomes the bottleneck, so the optimized rate no longer improves with increasing model capacity.This schedule-induced capacity saturation is rate-optimal only for s ≤ 1 − 1/(p + 1).Source: §5.1 · p. 9
  • Cosine decay is rate-optimal for all tasks when q < 3, while linear and 1-sqrt decay are rate-optimal for all tasks when q < 2.Source: §5.1 · p. 9
  • For PLK regression with (s, q) = (0.8, 5), power decay with p = 4.2 follows the optimal rate whereas cosine decay with p = 2 shows the predicted saturation.Source: §5.1 · p. 10

5.2 Why relative training progress matters

  • Relative-progress normalization avoids the rate barrier because changing a fixed absolute-step decay exponent changes the total intrinsic-time scaling with N.Source: §5.2 · p. 10
  • Horizon-independent profiles η_N(z) = η_0h(z) remain rate-suboptimal even when their peak learning rate is retuned for each training horizon.Polynomial decay exemplifies the barrier because it cannot simultaneously accumulate enough intrinsic time and suppress late-stage noise.Source: §5.2 · p. 10

6 Discrete-time SGD for kernel regression

The paper derives last-iterate convergence guarantees for one-pass SGD under kernel-regression source and capacity conditions, using FSL-motivated power decay. These schedules attain minimax-optimal or best-possible one-pass rates across easy and hard task regimes.

6 Discrete-time SGD for kernel regression

  • The analysis retains source exponent s and capacity exponent q under general spectral regularity conditions rather than requiring the exact power-law structure.This permits direct comparison with existing kernel-regression results.Source: §6 · p. 11
  • The statistical target is the minimax lower bound, and the analysis asks how closely the last iterate of one-pass SGD can approach it.The supremum ranges over distributions satisfying the stated assumptions.Source: §6 · p. 11
  • Power-decay SGD attains the minimax-optimal rate in the easy-task regime and the best rate achievable by one-pass SGD in the hard-task regime.The theorem considers η_k = η_0(1 − k/N)^p under the stated source and capacity assumptions.Source: §6 · p. 11
  • The easy-task result matches the minimax-optimal rate without logarithmic factors, closing the gap between last-iterate and averaged SGD.The authors identify this as the first such last-iterate result to their knowledge.Source: §6 · p. 11

7 Experiments with nonlinear neural networks

Neural-network experiments test whether task difficulty controls the preferred decay duration beyond the kernel setting. Across nonlinear regression and CIFAR-100 classification, harder tasks favor later decay while easier tasks favor earlier decay.

7 Experiments with nonlinear neural networks

  • The observed pattern is that easier tasks favor longer decay phases, whereas harder tasks favor shorter decay phases.The regression and classification comparisons provide two nonlinear-network tests of this task-dependent trend.Source: §7 · p. 11, §7.1 · p. 12, §7.2 · p. 12
  • The experiments vary decay ratios while holding the main training setup fixed, using one-pass SGD for MLP regression and warmup followed by stable-linear decay for ResNet-18 classification.The MLP experiment fixes peak LR 0.1 for 1000 steps; the CIFAR experiment uses peak LR 0.2 for 5 epochs.Source: §7.1 · p. 12, §7.2 · p. 12

7.1 Nonlinear regression with two-layer MLPs

  • In nonlinear regression, the harder target γ = 0.6 prefers r = 0.5, whereas the easier target γ = 10 prefers decay from the start at r = 1.The comparison averages 32 runs and agrees with the theoretical prediction.Source: §7.1 · p. 12, Figure 2 · p. 12

7.2 Full versus binary classification on CIFAR-100

  • For CIFAR-100, harder CIFAR-Full prefers r = 0.4, while easier CIFAR-Binary prefers r = 1.CIFAR-Full uses 100 fine-grained classes; CIFAR-Binary groups images into animals and non-animals.Source: §7.2 · p. 12

8 Proof sketches and interpretation

The proof sketches derive optimal schedules by optimizing signal learning and noise forgetting in intrinsic time, then enforcing stability. They also show that fractional schedules need sufficiently fast terminal decay to retain optimal rates.

8.1 Proof sketch of Theorem 4.1

  • For a fixed intrinsic-time budget, minimizing the noise term yields an intrinsic-time profile determined by the forgetting kernel, which converts to a power-decay LR shape.The variational argument uses the intrinsic-time change of variables and an equality condition in the optimization bound.Source: §8.1 · p. 13
  • In the easy-task regime s > 1 − 1/q, the stability constraint is asymptotically inactive, so power decay is feasible from the beginning of training.The unconstrained peak LR vanishes asymptotically in the relevant scaling.Source: §8.1 · p. 14
  • In the hard-task regime s < 1 − 1/q, the unconstrained peak LR exceeds the stability threshold, producing a stability-saturated plateau followed by the same power-law decay.The optimal schedule therefore has a stable-decay structure, with the decay fraction vanishing asymptotically but containing a diverging number of iterations.Source: §8.1 · pp. 14–15
  • The stable and decay phases trade off signal-learning time against noise forgetting, so increasing the decay fraction improves forgetting but reduces intrinsic time for signal learning.For hard tasks, the signal-learning term dominates for constant decay fractions, forcing the optimal decay fraction to shrink with N.Source: §8.1 · pp. 14–15

8.2 Proof sketch of Theorem 5.2

  • For a power-decay tail, the fast regime p/(p + 1) > 1 − 1/q makes the combined exponent p/(p + 1) − (2 − 1/q) > −1, so the limiting integral converges absolutely.In the slow regime p/(p + 1) < 1 − 1/q, the limiting integral diverges as y → 1 and T −1 acts as a cutoff; at the boundary, the tail integral is of order log T.Source: §8.2 · pp. 15–16

9 Conclusion and discussion

The discussion interprets the phase transition as a consequence of stability-constrained signal learning and noise forgetting, while the proof appendix derives the optimal profile and intrinsic-time horizon. It also identifies boundaries for extending the macroscopic analysis to practical pretraining.

9 Conclusion and discussion

  • Task difficulty determines when decay begins, while the forgetting kernel determines the optimal decay shape independently of the target function.As harder tasks raise the unconstrained peak LR beyond stability, the optimal schedule transitions to delayed decay.Source: §9 · p. 16
  • A broad class of schedules with sufficiently fast terminal decay achieves the optimal scaling rate, but slower decay can cause schedule-induced capacity saturation.Thus exact shape optimality and rate optimality are distinct, explaining why practical profiles can perform similarly in some regimes but differ in others.Source: §9 · p. 16
  • The FSL analysis is macroscopic: once the signal-learning function and forgetting kernel are specified, it no longer depends on kernel regression’s microscopic structure.The paper suggests possible extension to LLM pretraining because both settings have been shown to follow FSLs.Source: §9 · p. 16
  • Extending the framework to practical pretraining requires identifying the relevant training regime, which may depend on batch size rather than task difficulty alone.Larger batches reduce noise injection but also reduce optimization steps and accumulated intrinsic time under a fixed data budget.Source: §9 · p. 16

A Proof of Theorem 4.1

  • The appendix derives the optimal schedule by first optimizing the profile at fixed intrinsic time and then optimizing the total intrinsic-time horizon.The variational formulation is rewritten in intrinsic time before the two optimization steps are applied.Source: A Proof of Theorem 4.1 · p. 22, §A.1 · p. 22, §A.2 · p. 23

A.3 Step 3: Incorporating the peak learning rate constraint

This section incorporates the peak learning rate constraint into the schedule analysis and examines fractional-profile tails and noise-scaling regimes. It also bounds the early contribution relative to the tail, including the logarithmic factor arising in the critical case.

A.3 Step 3: Incorporating the peak learning rate constraint

  • The constrained optimum is non-increasing in training time, with decay governed by the unconstrained profile after the stable phase.Monotonicity follows from the stated proposition; in the decay region, complementary slackness sets the multiplier to zero.Source: §A.3 · pp. 24–25
  • For hard tasks s < 1 −1/q, the peak constraint is active, and the optimal schedule has a stable phase followed by decay.For easy tasks s > 1−1/q, the unconstrained peak tends to zero and the stability constraint is eventually inactive.Source: §A.3 · pp. 23, 25
  • For sufficiently large N, Q_N attains its minimum over (0, 1] × (0, 1), while the minimum of G tends to zero.Source: §A.3 · pp. 25, 27

B Proofs for fractional LR schedules

  • If a training-step profile decays as (1−x)^p, its intrinsic-time counterpart decays with exponent p/(p+1) near the terminal point.Lemma B.1 establishes the transformed tail behavior through the cumulative time reparameterization.Source: B Proofs for fractional LR schedules · p. 29

B.1 Proof of Theorem 5.2

  • For fractional schedules, the noise scaling uses α = min{q, p + 1}, with the convolution analysis separating early and tail regions and identifying two tail regimes.The stated cases are a−ρ < −1, yielding TItail ≂ T^−a, and a−ρ > −1, yielding TItail ≂ T^(1−ρ).Source: §B.1 · p. 30
  • In the critical case a −ρ = −1, integrating u−1 yields T_tail ≂ T^(1−ρ) log T.The logarithmic factor arises from the integral in the tail region.Source: §B.1 · p. 31
  • Across all three cases, T_early ≲ T^(1−ρ) is bounded by the tail contribution, so combining both regions proves (32).Source: §B.1 · p. 31

B.2 Proof of Theorem 5.3

The proofs optimize the FSL loss over stability-admissible intrinsic times by comparing signal-learning and noise-forgetting terms. This yields distinct hard-task and easier-regime scalings, with the hard-task optimum reaching the stability boundary.

B.2 Proof of Theorem 5.3

  • The admissible optimization range is 0 < T ≤ cζN, where cζ is fixed by the stability threshold and noise magnitude.This constraint supplies the upper bound on T in the hard-task argument.Source: §B.2 · p. 31
  • In the hard regime, the loss is comparable to N^-s throughout admissible times, so the optimizer satisfies T* ≂ N.The same conclusion holds at q = p + 1 because the extra logarithmic factor is dominated by the strict power gap.Source: §B.2 · p. 31
  • In the balancing regime, choosing T̂ = N^(1/(s + 1/α)) equalizes the decreasing signal and increasing noise terms.Here α = min{q, p + 1}, and the resulting lower and upper bounds match at the chosen scale.Source: §B.2 · p. 31

B.3 Proof of Theorem 5.4

  • Theorem 5.4 rules out any rate exponent ρ > s/(s + 1) by forcing H(N) to diverge, contradicting the assumed bound.The contradiction is established after extending the integer inequality to sufficiently large real arguments and integrating it.Source: §B.3 · pp. 33–34

C Proofs for discrete-time SGD

  • The discrete-time SGD analysis assumes source and capacity conditions that produce the signal-learning and noise-forgetting exponents used in the FSL bounds.The setup allows the target coefficient sequence to be nonsquare-summable, so an ℓ2-valued target parameter need not exist.Source: C Proofs for discrete-time SGD · p. 34, §C.1 · p. 35

C.2 Analysis of one-pass SGD

The one-pass SGD analysis bounds excess risk through intrinsic time and separates signal learning from noise forgetting. Under source and capacity conditions, these quantities yield power-decay risk bounds and matching optimal-rate arguments.

C.2 Analysis of one-pass SGD

  • The expected excess-risk bound is expressed using intrinsic time, signal learning, and a forgetting kernel.Intrinsic time is the cumulative learning-rate sum, t_k = Σ_{i<k} η_i.Source: §C.2 · pp. 36, 38
  • Noise variance in each spectral direction scales with curvature λ_j, while its amplitude depends on current excess risk and label-noise variance.This structure is the basis for the noise-forgetting term in the risk inequality.Source: §C.2 · p. 37

C.3 Properties of signal-learning function and forgetting kernel

  • The signal-learning function and forgetting kernel are nonnegative and non-increasing, with power-law decay controlled by source and capacity assumptions.The kernel exponent is ρ = 2 − 1/q > 1.Source: §C.3 · p. 38

C.4 Excess risk upper bound

  • Schedules satisfying the step-size condition keep expected excess risk uniformly bounded and yield a risk bound in terms of signal learning and noise forgetting.The uniform bound holds over all steps k ≤ N.Source: §C.4 · p. 40

C.4.1 Power-decay schedules

  • For power-decay schedules, the noise contribution is governed by μ = min{1 − 1/q, p/(p + 1)}, with a logarithmic factor at p = q − 1.This identifies the decay exponent threshold at which the noise bound changes form.Source: §C.4.1 · p. 41, §C.5 · p. 41

C.5 Proof of Theorem 6.3

  • In the easy-task regime s ≥ 1 − 1/q, choosing p > q − 1 and balancing terms gives the rate E[E(θ_N)] ≲ N^(-sq/(sq + 1)).The schedule can use a sufficiently small prefactor to satisfy stability uniformly in N.Source: §C.5 · p. 42
  • In the hard-task regime s < 1 − 1/q, a constant-order initial learning rate and sufficiently large p make the noise exponent exceed s.The resulting bound is therefore dominated by the signal-learning rate.Source: §C.5 · p. 42

C.6 Proof of Proposition 6.4

  • The lower-bound construction uses a power-law spectrum and a target supported on one eigen-direction to show no stable schedule uniformly beats the signal-limited rate.The construction selects λ_j = j^-q and chooses the relevant eigen-index as a function of intrinsic time.Source: §C.6 · pp. 42–43

D.1 PLK regression

Additional experiments test fractional schedules across power-law regimes and extend the task-dependent decay pattern to MLP and CIFAR-100 training. Across these settings, harder tasks favor later decay while schedule families reproduce the predicted scaling behavior.

D.1 PLK regression

  • Across fractional schedules and all tested (s, q) settings, observed log-log slopes agree with theoretical predictions in the blue, red, and green regimes.The experiments also show schedule-induced capacity saturation beyond the cosine-versus-power-decay comparison.Source: §D.1 · p. 44
  • Figure 4 compares schedule families across four regime columns, including optimal easy- and hard-task rates and capacity-saturated rates governed by α = min{q, p + 1}.Peak learning rates are tuned separately for each N, and losses average 500 independent SGD runs.Source: Figure 4 · p. 44

D.2 MLP regression

  • For MLP regression, the harder target γ = 0.6 prefers intermediate decay, while the easier target γ = 10 is nearly indifferent between r = 0.9 and r = 1.With peak LR 0.05, the respective grid minimizers are r = 0.4 and r = 0.9.Source: §D.2 · p. 44, Figure 5 · p. 45

D.3 CIFAR-100 classification

  • For CIFAR-100, CIFAR-Full prefers delayed decay ratios r = 0.4 or 0.2, whereas CIFAR-Binary prefers immediate post-warmup decay at r = 1.These results are reported for peak LRs 0.1 and 0.2 under different batch-size settings.Source: Figure 6 · p. 45, Figure 7 · p. 45

Related papers

Loading 2602.06797v3…