Learning Rate Annealing Improves Tuning Robustness in Stochastic Optimization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Attia, Amit, Koren, Tomer
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914332188606464
author Attia, Amit
Koren, Tomer
author_facet Attia, Amit
Koren, Tomer
contents The learning rate in stochastic gradient methods is a critical hyperparameter that is notoriously costly to tune via standard grid search, especially for training modern large-scale models with billions of parameters. We identify a theoretical advantage of learning rate annealing schemes that decay the learning rate to zero at a polynomial rate, such as the widely-used cosine schedule, by demonstrating their increased robustness to initial parameter misspecification due to a coarse grid search. We present an analysis in a stochastic convex optimization setup demonstrating that the convergence rate of stochastic gradient descent with annealed schedules depends sublinearly on the multiplicative misspecification factor $ρ$ (i.e., the grid resolution), achieving a rate of $O(ρ^{1/(2p+1)}/\sqrt{T})$ where $p$ is the degree of polynomial decay and $T$ is the number of steps. This is in contrast to the $O(ρ/\sqrt{T})$ rate obtained under the inverse-square-root and fixed stepsize schedules, which depend linearly on $ρ$. Experiments confirm the increased robustness compared to tuning with a fixed stepsize, that has significant implications for the computational overhead of hyperparameter search in practical training scenarios.
format Preprint
id arxiv_https___arxiv_org_abs_2503_09411
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Learning Rate Annealing Improves Tuning Robustness in Stochastic Optimization
Attia, Amit
Koren, Tomer
Machine Learning
Optimization and Control
The learning rate in stochastic gradient methods is a critical hyperparameter that is notoriously costly to tune via standard grid search, especially for training modern large-scale models with billions of parameters. We identify a theoretical advantage of learning rate annealing schemes that decay the learning rate to zero at a polynomial rate, such as the widely-used cosine schedule, by demonstrating their increased robustness to initial parameter misspecification due to a coarse grid search. We present an analysis in a stochastic convex optimization setup demonstrating that the convergence rate of stochastic gradient descent with annealed schedules depends sublinearly on the multiplicative misspecification factor $ρ$ (i.e., the grid resolution), achieving a rate of $O(ρ^{1/(2p+1)}/\sqrt{T})$ where $p$ is the degree of polynomial decay and $T$ is the number of steps. This is in contrast to the $O(ρ/\sqrt{T})$ rate obtained under the inverse-square-root and fixed stepsize schedules, which depend linearly on $ρ$. Experiments confirm the increased robustness compared to tuning with a fixed stepsize, that has significant implications for the computational overhead of hyperparameter search in practical training scenarios.
title Learning Rate Annealing Improves Tuning Robustness in Stochastic Optimization
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2503.09411