# Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 68 references that are in the map.
- First posted on arXiv 2 Apr 2025.

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

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [Knowledge distillation: A good teacher is patient and consistent](https://mlsanity.com/paper/2106.05237/) (2021)
- [Loss landscapes and optimization in over-parameterized non-linear systems and neural networks](https://mlsanity.com/paper/2003.00307/) (2020)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Fast and Faster Convergence of SGD for Over-Parameterized Models and an Accelerated Perceptron](https://mlsanity.com/paper/1810.07288/) (2018)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [Coin Betting and Parameter-Free Online Learning](https://mlsanity.com/paper/1602.04128/) (2016)
- [Prodigy: An Expeditiously Adaptive Parameter-Free Learner](https://mlsanity.com/paper/2306.06101/) (2023)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)
- [Learning-Rate-Free Learning by D-Adaptation](https://mlsanity.com/paper/2301.07733/) (2023)
- [The rising costs of training frontier AI models](https://mlsanity.com/paper/2405.21015/) (2024)
- [DoG is SGD's Best Friend: A Parameter-Free Dynamic Step Size Schedule](https://mlsanity.com/paper/2302.12022/) (2023)
- [Adaptive Gradient Descent without Descent](https://mlsanity.com/paper/1910.09529/) (2019)
- [SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation](https://mlsanity.com/paper/2006.10311/) (2020)
- [Global convergence of the Heavy-ball method for convex optimization](https://mlsanity.com/paper/1412.7457/) (2014)
- [Lipschitz and Comparator-Norm Adaptivity in Online Learning](https://mlsanity.com/paper/2002.12242/) (2020)
- [UniXGrad: A Universal, Adaptive Algorithm with Optimal Guarantees for Constrained Optimization](https://mlsanity.com/paper/1910.13857/) (2019)
- [Revisiting the Polyak step size](https://mlsanity.com/paper/1905.00313/) (2019)
- [Online Adaptive Methods, Universality and Acceleration](https://mlsanity.com/paper/1809.02864/) (2018)
- [Making SGD Parameter-Free](https://mlsanity.com/paper/2205.02160/) (2022)
- [Training Neural Networks for and by Interpolation](https://mlsanity.com/paper/1906.05661/) (2019)
- [Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution](https://mlsanity.com/paper/2205.04583/) (2022)
- [DoWG Unleashed: An Efficient Universal Parameter-Free Gradient Descent Method](https://mlsanity.com/paper/2305.16284/) (2023)
- [Stochastic Quasi-Gradient Methods: Variance Reduction via Jacobian Sketching](https://mlsanity.com/paper/1805.02632/) (2018)
- [MoMo: Momentum Models for Adaptive Learning Rates](https://mlsanity.com/paper/2305.07583/) (2023)

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

- [New Perspectives on the Polyak Stepsize: Surrogate Functions and Negative Results](https://mlsanity.com/paper/2505.20219/) (2025)
- [Taking the Road Less Scheduled with Adaptive Polyak Steps](https://mlsanity.com/paper/2511.07767/) (2025)
- [Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients](https://mlsanity.com/paper/2512.02342/) (2025)
- [Adaptive Sharpness-Aware Minimization with a Polyak-type Step size: A Theory-Grounded Scheduler](https://mlsanity.com/paper/2606.01827/) (2026)
- [Polyak-Type Extragradient Methods for Monotone Root-Finding Problems](https://mlsanity.com/paper/2609.26581/) (2026)

## Abstract (arXiv)

We provide a general convergence theorem of an idealized stochastic Polyak step size called SPS^*. Besides convexity, we only assume a local expected gradient bound, that includes locally smooth and locally Lipschitz losses as special cases. We refer to SPS^* as idealized because it requires access to the loss for every training batch evaluated at a solution. It is also ideal, in that it achieves the optimal lower bound for globally Lipschitz function, and is the first Polyak step size to have an O(1/√t) anytime convergence in the smooth setting. We show how to combine SPS^* with momentum to achieve the same favorable rates for the last iterate. We conclude with several experiments to validate our theory, and a more practical setting showing how we can distill a teacher GPT-2 model into a smaller student model without any hyperparameter tuning.