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

In plain words

A learning rate sets how big each training step is. Large steps learn fast but add noise that the model has to undo later; small steps are calm but slow. Training recipes differ: some lower the rate steadily from the start, while warmup-stable-decay recipes keep it high for most of the run and drop it only at the end. The open question is what decides when the drop should begin.

This paper gives one, for a simplified model of learning: a linear model on power-law features, trained one example at a time. The main finding is that the right moment depends on the task. Easy tasks benefit from shrinking the rate from the very beginning. Hard tasks benefit from keeping it at the largest safe value for most of training and dropping it only near the end. Once the moment is chosen, the shape of the drop is set by how quickly the model forgets old noise, not by the task. The authors also show that a fairly broad family of drop shapes works about as well as the exact optimum, as long as the rate falls fast enough near the end.

TL;DR

• Under a functional scaling law, the final-step loss of one-pass SGD depends on the schedule through signal learning and noise forgetting §3.2.
• Optimal schedule, easy tasks (s > 1-1/q): power decay (1 + δ_N - z/N)^(2q-1) from step 0 [thm. 4.1].
• Optimal schedule, hard tasks (s < 1-1/q): hold the stability limit, then decay with the same exponent over the last fraction r_N = o(1) of training [thm. 4.1].
• Power decay (1 - k/N)^p with p > q-1 gives last-iterate rates without log factors for easy tasks [thm. 6.3]. Authors describe this as the first such result §6.
• All results are asymptotic and up to constants; the paper states no numeric constants for the peak rate, stability limit or step-size threshold.

Problem

Schedule design is usually done by trial and error, and classical theory treats the schedule as a fixed input. Classical polynomial decay is rate-suboptimal for the final-step loss in this setting, and exponential decay needs logarithmic factors §1 §2. Empirically, WSD schedules can match or beat cosine decay at large scale, but the reason is unclear §1. The paper assumes a constant label noise, power-law spectra and hypercontractive features, and a stability cap on the learning rate §3.1 §4.

Method

Setting. One-pass SGD on a linear model over features φ, with one fresh sample per step and zero initialisation. The step size at step k is η_k §3.1 eq. 2. Spectra follow λ_j ≍ j^(-q) and target coefficients satisfy λ_j |θ*_j|² ≍ j^(-(1+sq)), with q > 1 (capacity) and s > 0 (source, lower means harder) §3.1 [assump. 3.2]. The paper's Assumptions 3.1 (constant label noise) and 3.3 (hypercontractive features; Gaussian features satisfy it) are also required §3.1.

Functional scaling law (FSL). The paper takes the FSL of Li et al. (2025) as given §3.2. With intrinsic time t_k = Σ_{i<k} η_i, the expected excess risk is, up to constants eq. 3:

E[E(θ_k)] ≍ S(t_k) + Σ_{j<k} K(t_k - t_j) η_j²

with S(t) = (c_S + t)^(-s) and K(t) = (c_K + t)^(-(2 - 1/q)), where c_S, c_K ≍ 1 eq. 4. S is signal learning and K is the forgetting kernel §3.2. In continuous time, with z the step and η(z) the rate, the final-step objective for a fixed horizon N is eq. 5 eq. 6 eq. 7:

F[η] = S(T_N) + ∫_0^N K(∫_z^N η(u) du) η(z)² dz, T_N = ∫_0^N η(z) dz

The total intrinsic time T_N controls signal learning; the local rate controls the noise injected at step z; the remaining area controls how much of that noise is forgotten §3.2.

Optimal schedule. Minimise F[η] over 0 ≤ η(z) ≤ η_stab, with η_stab ≍ 1 the stability limit eq. 8 §4. The proof has three steps §8.1 app. A:

1. Fix T = T_N. Cauchy–Schwarz on the noise term under the constraint ∫_0^T 1/φ = N gives an optimal intrinsic-time profile φ*(τ) ∝ K(T-τ)^(-1/2) eq. 19. Converted to step time, this is power decay with exponent 2q-1.
2. Substitute φ* back. The loss becomes T^(-s) + T^(1/q)/N eq. 20. Balancing gives T* ≍ N^(q/(1+sq)).
3. Check the peak. φ*(0) ≍ T*/N ≍ N^(-(s-1+1/q)/(s+1/q)). If s > 1-1/q this goes to 0 and the constraint is inactive. If s < 1-1/q it diverges, the cap becomes active, and the optimum holds η_stab over an initial plateau before decaying §4 §8.1.

The resulting schedule is [thm. 4.1]:

• Easy tasks, s > 1-1/q: η*(z) = η_peak (1 + δ_N - z/N)^(2q-1), with η_peak ≍ N^(-(s-1+1/q)/(s+1/q)) and δ_N = o_N(1). Final-step loss E*_N ≍ N^(-sq/(sq+1)), which is the minimax rate eq. 9.
• Hard tasks, s < 1-1/q: η*(z) = η_stab for z ≤ N_1, and η_stab (1 + δ_N - (z - N_1)/(N - N_1))^(2q-1) for z > N_1. The decay fraction is r*_N = (N - N_1)/N ≍ N^(-((1-1/q)-s)/(2-1/q)) = o(1) eq. 10. Final-step loss E*_N ≍ N^(-s).

Section 4 states the boundary s = 1-1/q as undetermined by scaling alone §4.

Fractional schedules. A fractional schedule is η_N(z) = η̄ ζ(z/N), where ζ is a fixed profile on [0,1] with unit integral and η̄ is tuned per N [def. 5.1]. Intrinsic time is T = η̄ N, and the final-step loss is eq. 12:

F_N(T, ζ) = S(T) + (T/N) ∫_0^T K(T - t) ψ(t/T) dt, ψ = ζ ∘ P^(-1)

where P(x) = ∫0^x ζ. Stability reads T ≤ N η_stab / ||ζ||∞ eq. 13. The terminal behaviour is ζ(x) ≍ (1-x)^p near x = 1 eq. 14. Examples: constant (p = 0), cosine 1 + cos(πx) (p = 2), 1-sqrt 3(1-√x) (p = 1), power (p+1)(1-x)^p (p is a design choice) §5.1.

• Scaling for fixed T: F_N ≍ T^(-s) + (T^(1/α)/N)·[log T]^(1{q=p+1}), with α = min{q, p+1} [thm. 5.2] eq. 15.
• Optimising T: if s ≥ 1-1/α, T* ≍ N^(1/(s+1/α)) and E*_N ≍ N^(-sα/(sα+1)) (with a log correction when q = p+1). If s < 1-1/α, T* ≍ N and E*_N ≍ N^(-s) [thm. 5.3].
• So any profile with p+1 > q attains the optimal rate after tuning T. When p+1 < q the rate stops improving with q; the paper calls this schedule-induced capacity saturation §5.1.
• Profiles fixed in absolute steps cannot reach the optimal rate for any non-increasing h, even with a per-N peak tuned, for ρ > s/(s+1) [thm. 5.4] §5.2. Polynomial decay η_0 (1+z)^(-a) with a ∈ (0,1) is an example §5.2.

Discrete SGD with power decay. For η_k = η_0 (1 - k/N)^p, with η_0 chosen per regime [thm. 6.3] app. C:

• Easy tasks: p > q-1 and η_0 ≍ N^(-(sq-q+1)/(sq+1)) give E[E(θ_N)] ≲ N^(-sq/(sq+1)), with no log factor [thm. 6.3].
• Hard tasks: η_0 ≍ 1 and p > s/(1-s) give E[E(θ_N)] ≲ N^(-s) [thm. 6.3].
• In the hard regime, any schedule with max η ≲ 1 has worst-case risk of at least order N^(-s) [prop. 6.4], so the power-decay rate cannot be improved within that bound.

# Final-step-optimal schedule (asymptotic form)  [thm. 4.1]
# Inputs: horizon N, source s, capacity q, peak_lr and eta_stab
# Output: lr(z) for step z in [0, N]
# peak_lr is the eta_peak of the theorem, up to a constant the paper does
# not state: scale ~ N ** (-(s - 1 + 1 / q) / (s + 1 / q)).
# eta_stab is the stability limit, given only as order 1 [eq. 8].

def optimal_lr(z, N, s, q, peak_lr, eta_stab):
    # Easy tasks: power decay from the first step  [thm. 4.1]
    if s > 1 - 1 / q:
        return peak_lr * (1 - z / N) ** (2 * q - 1)
    # Hard tasks: hold the stability limit, then decay  [thm. 4.1]
    # decay_ratio is r_N, given only up to a constant  [thm. 4.1]
    decay_ratio = N ** (-((1 - 1 / q) - s) / (2 - 1 / q))
    decay_start = N * (1 - decay_ratio)
    if z <= decay_start:
        return eta_stab
    progress = (z - decay_start) / (N - decay_start)
    return eta_stab * (1 - progress) ** (2 * q - 1)

# Power-decay SGD schedule for easy tasks  [thm. 6.3] [app. C]
# Inputs: horizon N, source s, capacity q, exponent p > q - 1, prefactor eta0
# Output: step size for SGD step k in [0, N)

def power_decay_lr(k, N, eta0, p):
    # Polynomial decay that reaches zero at step N  [thm. 6.3]
    return eta0 * (1 - k / N) ** p

Hard tasks use the same function with eta0 ≍ 1 and p > s / (1 - s) [thm. 6.3].

Hyperparameters

Results

Author-reported. Setup: theory for one-pass SGD in power-law kernel regression §3.1; small experiments, one run family each §5.1 §7.

Author-reported experiments are small and use one fixed grid of decay ratios per setting. The paper describes them as consistent with the theory, not as tight fits §7.

Paper vs code

No official code. The paper gives no code link in the abstract, the introduction, the experiments or the LaTeX source. Searched the GitHub API for "functional scaling law", "functional scaling law learning rate", "warmup stable decay power decay", "wsd learning rate" and "optimal learning rate schedule kernel regression". No result is owned by an author or their lab; the matches were unrelated projects or third-party reimplementations.

Pitfalls

• Constants are missing everywhere. Every rate, peak and stability value is up to constants or "sufficiently small" [thm. 4.1] [thm. 6.3]. An implementer must tune η_peak and r_N, and the paper gives no finite-N rule for δ_N §4. The MLP and CIFAR experiments tune the decay ratio by grid search, not from the formula §7.
• Exponent 2q-1 vs p > q-1. The optimal schedule of §4 uses exponent 2q-1. The SGD theorem only needs p > q-1 [thm. 6.3]; the paper says exact tuning is unnecessary §5.1. Do not treat 2q-1 as a requirement for the rate.
• Boundary cases differ. §4 says the boundary s = 1-1/q is undetermined by scaling §4, after thm. 4.1. §6 writes the easy regime as s ≥ 1-1/q [thm. 6.3], while §4 writes s > 1-1/q. Treat s = 1-1/q as ambiguous.
• Stability cap is essential in the hard regime. Without it the unconstrained peak diverges as N grows §8.1. The plateau height is η_stab, which the paper leaves as order 1 eq. 8.
• Relative, not absolute, progress. A fixed profile in absolute steps cannot be rate-optimal even after peak tuning [thm. 5.4]. Define the profile over z/N.
• Fractional normalisation. The profile ζ must have unit integral, and the peak is η̄ = T/N; stability then needs T ≤ N η_stab / ||ζ||∞ eq. 13. Compute ||ζ||∞ before choosing T.
• Asymptotic claims. Predictions are scaling laws, up to constants §3.2. The FSL itself is taken from Li et al. (2025) and is not re-derived here §3.2. Finite-N loss values are not predicted.
• Neural-network evidence is narrow. The MLP and CIFAR experiments use stable-linear decay, not the (1-u)^(2q-1) shape of §4 §7. They have 3 to 32 runs per point and no error bars beyond run means fig. 2 fig. 3. Their best decay ratio changes with peak LR and batch size app. D.2 app. D.3, so the "task-dependent" claim is not separated from these other settings.
• Version mismatch in the source. The fetched LaTeX file is named main_arxiv_v1.tex, while the metadata records arXiv v3 [raw/arxiv]. Section, equation and theorem numbers here come from that file. Check them against the v3 PDF before citing them externally.
• Costs the abstract hides. Power decay with p > q-1 needs the peak rate tuned to N, and the easy-task optimum needs η_peak to scale with N. Each N is a separate tuning run in the experiments fig. 1.

Lineage

• Builds on: the functional scaling law of Li et al. (2025) §3.2; averaged SGD minimax rates of Dieuleveut & Bach (2015) and Mücke, Neu & Rosasco (2019) §2; last-iterate exponential-decay results of Ge et al. (2019) and Wu et al. (2022) §1 §2.
• Relates to: duration-aware cosine analysis of Li, Zhuang & Orabona (2021) §2; WSD river-valley view of Wen et al. (2024) and WSD evidence of Hu et al. (2024) §1 §2; numerical multi-power-law WSD schedule of Luo et al. (2024) §2.
• Compared against: cosine decay, linear decay (Bergsma et al. 2025), 1-sqrt decay (Hägele et al. 2024), and classical polynomial decay §5.1 §5.2.
• Superseded by: not stated.