# Theory of Optimal Learning Rate Schedules and Scaling Laws for a Random Feature Model

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 49 references that are in the map.
- First posted on arXiv 4 Feb 2026.

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

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Explaining Neural Scaling Laws](https://mlsanity.com/paper/2102.06701/) (2021)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Sobolev Norm Learning Rates for Regularized Least-Squares Algorithm](https://mlsanity.com/paper/1702.07254/) (2017)
- [A Dynamical Model of Neural Scaling Laws](https://mlsanity.com/paper/2402.01092/) (2024)
- [Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective](https://mlsanity.com/paper/2410.05192/) (2024)
- [Don't be lazy: CompleteP enables compute-efficient deep transformers](https://mlsanity.com/paper/2505.01618/) (2025)
- [Scaling Exponents Across Parameterizations and Optimizers](https://mlsanity.com/paper/2407.05872/) (2024)
- [Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passes](https://mlsanity.com/paper/1805.10074/) (2018)
- [4+3 Phases of Compute-Optimal Neural Scaling Laws](https://mlsanity.com/paper/2405.15074/) (2024)
- [Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit](https://mlsanity.com/paper/2309.16620/) (2023)
- [Scaling Laws in Linear Regression: Compute, Parameters, and Data](https://mlsanity.com/paper/2406.08466/) (2024)
- [How Feature Learning Can Improve Neural Scaling Laws](https://mlsanity.com/paper/2409.17858/) (2024)
- [Optimal Linear Decay Learning Rate Schedules and Further Refinements](https://mlsanity.com/paper/2310.07831/) (2023)
- [Last iterate convergence of SGD for Least-Squares in the Interpolation regime](https://mlsanity.com/paper/2102.03183/) (2021)
- [Scaling Optimal LR Across Token Horizons](https://mlsanity.com/paper/2409.19913/) (2024)
- [Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks](https://mlsanity.com/paper/2507.02119/) (2025)
- [Optimal Protocols for Continual Learning via Statistical Physics and Control Theory](https://mlsanity.com/paper/2409.18061/) (2024)
- [Selecting and Composing Learning Rate Policies for Deep Neural Networks](https://mlsanity.com/paper/2210.12936/) (2022)
- [Dimension-adapted Momentum Outscales SGD](https://mlsanity.com/paper/2505.16098/) (2025)
- [Neural Thermodynamic Laws for Large Language Model Training](https://mlsanity.com/paper/2505.10559/) (2025)
- [Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling](https://mlsanity.com/paper/2510.14717/) (2025)

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

- [Power-Law Spectrum of the Random Feature Model](https://mlsanity.com/paper/2603.14578/) (2026)
- [Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods](https://mlsanity.com/paper/2606.19179/) (2026)
- [Statistical Properties of Training & Generalization](https://mlsanity.com/paper/2606.20299/) (2026)
- [Asymmetric Scaling Laws from Sparse Features](https://mlsanity.com/paper/2605.23591/) (2026)
- [Fast Learning Rate Transfer in Shallow Linear Networks at Growing Training Horizons](https://mlsanity.com/paper/2609.35029/) (2026)
- [From Spectra to Joint Schedules in LLM Pre-training: 3+3(+2) Scaling-Law Regimes](https://mlsanity.com/paper/2609.40148/) (2026)
- [Asymptotic Risk of Learning-Rate Schedules for Least-Squares Regression with Discretely Distributed Features](https://mlsanity.com/paper/2610.16518/) (2026)

## Abstract (arXiv)

Setting the learning rate (LR) for a deep learning model is a critical part of successful training. Choosing LRs is often done empirically with trial and error. In this work, we explore a solvable model of optimal LR schedules for a powerlaw random feature model trained with stochastic gradient descent (SGD). We consider the optimal schedule ηT^star(t) where t is the current iterate and T is the training horizon. This schedule is computed both as a numerical optimization problem and also analytically using optimal control theory. Our analysis reveals two regimes which we term the easy phase and hard phase. In the easy phase the optimal schedule is a polynomial decay ηT^star(t) simeq T-xi (1-t/T)δ where xi and δ depend on the properties of the features and task. In the hard phase, the optimal schedule resembles warmup-stable-decay with constant initial LR and annealing performed over a vanishing fraction of training steps. We investigate joint optimization of LR and batch size and find batch ramps can improve the wall-clock time in the easy phase. Beyond SGD, we derive optimal schedules for momentum parameter β(t) and show that it improves the loss-scaling exponent in the hard phase. We compare our optimal schedule to various benchmarks including (1) optimal constant learning rates ηT(t) ~ T-xi (2) optimal power laws ηT(t) ~ T-xi t-chi, finding that our schedule achieves better rates than either of these. Our theory suggests that LR transfer across training horizon depends on the structure of the model and task. For ResNet image classification on CIFAR-5M, the learning curves exhibit hard-phase behavior where optimal base LRs are constant under sufficient annealing. GPT-2 style transformers trained in language modeling exhibit easy-phase behavior where optimal LRs shift even under annealing.