# Learning Rate Schedules in the Presence of Distribution Shift

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 44 references that are in the map.
- First posted on arXiv 27 Mar 2023.

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Non-stationary Stochastic Optimization](https://mlsanity.com/paper/1307.5449/) (2013)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [Online Optimization : Competing with Dynamic Comparators](https://mlsanity.com/paper/1501.06225/) (2015)
- [Deep learning with Elastic Averaging SGD](https://mlsanity.com/paper/1412.6651/) (2014)
- [Hessian-based Analysis of Large Batch Training and Robustness to Adversaries](https://mlsanity.com/paper/1802.08241/) (2018)
- [The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares](https://mlsanity.com/paper/1904.12838/) (2019)
- [Tracking Slowly Moving Clairvoyant: Optimal Dynamic Regret of Online Learning with True and Noisy Gradient](https://mlsanity.com/paper/1605.04638/) (2016)
- [On the Fenchel Duality between Strong Convexity and Lipschitz Continuous Gradient](https://mlsanity.com/paper/1803.06573/) (2018)
- [Averaging Stochastic Gradient Descent on Riemannian Manifolds](https://mlsanity.com/paper/1802.09128/) (2018)
- [On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models](https://mlsanity.com/paper/2209.05310/) (2022)
- [On Learning Rates and Schr"odinger Operators](https://mlsanity.com/paper/2004.06977/) (2020)
- [Deep Relaxation: partial differential equations for optimizing deep neural networks](https://mlsanity.com/paper/1704.04932/) (2017)
- [Tracking Moving Agents via Inexact Online Gradient Descent Algorithm](https://mlsanity.com/paper/1710.05133/) (2017)
- [Deep Cytometry: Deep learning with Real-time Inference in Cell Sorting and Flow Cytometry](https://mlsanity.com/paper/1904.09233/) (2019)
- [How do SGD hyperparameters in natural training affect adversarial robustness?](https://mlsanity.com/paper/2006.11604/) (2020)

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

- [Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems](https://mlsanity.com/paper/2305.12102/) (2023)
- [A Stability Principle for Learning under Non-Stationarity](https://mlsanity.com/paper/2310.18304/) (2023)
- [Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay](https://mlsanity.com/paper/2602.06797/) (2026)
- [A call for embodied AI](https://mlsanity.com/paper/2402.03824/) (2024)
- [Understanding the Robustness of Multi-modal Contrastive Learning to Distribution Shift](https://mlsanity.com/paper/2310.04971/) (2023)
- [PriorBoost: An Adaptive Algorithm for Learning from Aggregate Responses](https://mlsanity.com/paper/2402.04987/) (2024)
- [Efficient Hyperparameter Search for Non-Stationary Model Training](https://mlsanity.com/paper/2512.01258/) (2025)

## Abstract (arXiv)

We design learning rate schedules that minimize regret for SGD-based online learning in the presence of a changing data distribution. We fully characterize the optimal learning rate schedule for online linear regression via a novel analysis with stochastic differential equations. For general convex loss functions, we propose new learning rate schedules that are robust to distribution shift and we give upper and lower bounds for the regret that only differ by constants. For non-convex loss functions, we define a notion of regret based on the gradient norm of the estimated models and propose a learning schedule that minimizes an upper bound on the total expected regret. Intuitively, one expects changing loss landscapes to require more exploration, and we confirm that optimal learning rate schedules typically increase in the presence of distribution shift. Finally, we provide experiments for high-dimensional regression models and neural networks to illustrate these learning rate schedules and their cumulative regret.