Schedulers for Schedule-free: Theoretically inspired hyperparameters
This work addresses the gap between theory and practice in schedule-free optimization for machine learning practitioners, though it is incremental as it builds on existing schedule-free methods.
The authors extended the convergence theory of schedule-free optimization to support any learning rate scheduler and derived optimal convergence rates for specific schedules, achieving an O(1/√T) rate for warmup-stable-decay and designing a new adaptive Polyak schedule that performs well in experiments on model distillation.
The recently proposed schedule-free method has been shown to achieve strong performance when hyperparameter tuning is limited. The current theory for schedule-free only supports a constant learning rate, where-as the implementation used in practice uses a warm-up schedule. We show how to extend the last-iterate convergence theory of schedule-free to allow for any scheduler, and how the averaging parameter has to be updated as a function of the learning rate. We then perform experiments showing how our convergence theory has some predictive power with regards to practical executions on deep neural networks, despite that this theory relies on assuming convexity. When applied to the warmup-stable-decay (wsd) schedule, our theory shows the optimal convergence rate of $\mathcal{O}(1/\sqrt{T})$. We then use convexity to design a new adaptive Polyak learning rate schedule for schedule-free. We prove an optimal anytime last-iterate convergence for our new Polyak schedule, and show that it performs well compared to a number of baselines on a black-box model distillation task.