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

arXiv: https://arxiv.org/abs/2602.06797 · HTML: https://mlsanity.com/paper/2602.06797/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 50 of its 66 references that are in the map.
- First posted on arXiv 6 Feb 2026.

## Builds on (most cited first, 30 of 50)

- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [Scaling Laws for Autoregressive Generative Modeling](https://mlsanity.com/paper/2010.14701/) (2020)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [Explaining Neural Scaling Laws](https://mlsanity.com/paper/2102.06701/) (2021)
- [Spectrum Dependent Learning Curves in Kernel Regression and Wide Neural Networks](https://mlsanity.com/paper/2002.02561/) (2020)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [Scaling Laws for Generative Mixed-Modal Language Models](https://mlsanity.com/paper/2301.03728/) (2023)
- [A Solvable Model of Neural Scaling Laws](https://mlsanity.com/paper/2210.16859/) (2022)
- [Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations](https://mlsanity.com/paper/1811.01558/) (2018)
- [Non-strongly-convex smooth stochastic approximation with convergence rate O(1/n)](https://mlsanity.com/paper/1306.2119/) (2013)
- [A Dynamical Model of Neural Scaling Laws](https://mlsanity.com/paper/2402.01092/) (2024)
- [Fantastic Pretraining Optimizers and Where to Find Them](https://mlsanity.com/paper/2509.02046/) (2025)
- [Scaling Laws for Precision](https://mlsanity.com/paper/2411.04330/) (2024)
- [Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective](https://mlsanity.com/paper/2410.05192/) (2024)
- [Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passes](https://mlsanity.com/paper/1805.10074/) (2018)
- [The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares](https://mlsanity.com/paper/1904.12838/) (2019)
- [4+3 Phases of Compute-Optimal Neural Scaling Laws](https://mlsanity.com/paper/2405.15074/) (2024)
- [Benchmarking Optimizers for Large Language Model Pretraining](https://mlsanity.com/paper/2509.01440/) (2025)
- [Harder, Better, Faster, Stronger Convergence Rates for Least-Squares Regression](https://mlsanity.com/paper/1602.05419/) (2016)
- [The alignment property of SGD noise and how it helps select flat minima: A stability analysis](https://mlsanity.com/paper/2207.02628/) (2022)
- [Scaling Laws in Linear Regression: Compute, Parameters, and Data](https://mlsanity.com/paper/2406.08466/) (2024)
- [The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent](https://mlsanity.com/paper/2305.17490/) (2023)
- [Optimal Linear Decay Learning Rate Schedules and Further Refinements](https://mlsanity.com/paper/2310.07831/) (2023)

## Built on by (most cited first, 8 of 8)

- [Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws](https://mlsanity.com/paper/2602.14208/) (2026)
- [Hyperball May Not Be a Free Lunch](https://mlsanity.com/paper/2607.22444/) (2026)
- [Effective Learning Rate Governs Loss Dynamics in Language Model Pretraining](https://mlsanity.com/paper/2608.24814/) (2026)
- [Same Loss, Same Noise, Opposite Schedules: Noise Structure and Optimizer Normalization Jointly Determine Whether Learning-Rate Cooldown Helps](https://mlsanity.com/paper/2607.12360/) (2026)
- [Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency](https://mlsanity.com/paper/2609.02728/) (2026)
- [Learning with Synthetic Data via SGD in High-Dimensional Linear Regression](https://mlsanity.com/paper/2609.09572/) (2026)
- [Understanding the Weight Averaging Mechanism in LLM Training for Post-Training Quantization](https://mlsanity.com/paper/2610.05329/) (2026)
- [Asymptotic Risk of Learning-Rate Schedules for Least-Squares Regression with Discretely Distributed Features](https://mlsanity.com/paper/2610.16518/) (2026)

## Abstract (arXiv)

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.