arXiv · 2504.01898
Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation
Abstract
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/\sqrt{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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Robert M. Gower, Guillaume Garrigos, Nicolas Loizou, Dimitris Oikonomou, Konstantin Mishchenko, Fabian Schaipp. 2025-04-02. Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation. https://arxiv.org/abs/2504.01898
Cite the original work for its findings. Save a collection to share your selection of sources.