Learning Rate Decay Can Exponentially Accelerate SGD for Global Optimization of Nonconvex Functions
Abstract
Stochastic Gradient Descent (SGD) is a workhorse algorithm for continuous optimization. It is a folk belief among deep learning practitioners that learning rate decay in SGD substantially improves performance. However, it has remained unclear whether learning rate decay offers an asymptotic advantage for global nonconvex optimization. We answer this question in the affirmative by demonstrating a natural class of nonconvex functions for which, we prove that SGD with learning rate decay requires exponentially fewer gradient queries for global optimization than SGD with any fixed learning rate in the small-step size regime. To the best of our knowledge, this is the first exponential advantage that has been demonstrated for learning rate decay. The class of functions we consider includes many popular benchmark nonconvex functions. Our results provide a robust explanatory theory for the empirically observed benefits of learning rate decay for global optimization. We show that annealing performs the non-logconcave sampling task associated to the optimization problem in polynomial time, via a novel application of weak Poincaré inequalities.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper establishes the first provable exponential acceleration of SGD via learning rate decay for global optimization of nonconvex functions. It introduces a novel application of weak Poincaré inequalities to analyze annealed Langevin diffusion and transfers this analysis to discrete SGD, providing a rigorous theoretical explanation for the empirical success of learning rate decay in navigating complex optimization landscapes.
Methodology
The paper proposes a rigorous theoretical framework to analyze Stochastic Gradient Descent (SGD) with decaying learning rates for global optimization of nonconvex functions. The core methodological contribution is the application of weak Poincaré inequalities to the annealed Langevin diffusion, which serves as a continuous-time proxy for discrete SGD. The authors derive a novel perturbation theorem (Theorem 2) that allows bounding the weak Poincaré constant of a target measure by comparing it to a reference measure with a known Poincaré inequality, even when the density ratio is unbounded. This is a significant technical advance over standard Holley-Stroock type arguments. They then bridge the gap between the continuous diffusion and discrete SGD via a discretization analysis that decouples the exploratory and convergent phases. The approach is mathematically sophisticated and addresses a long-standing gap in understanding why learning rate decay works in practice.
Results & evaluation
The experimental section is limited but sufficient for a theory paper. The authors test three specific nonconvex functions: Separable Double-Well, Rastrigin, and Coupled Double-Well. They demonstrate the predicted polynomial-versus-exponential separation in runtime between fixed-step-size SGD and annealed SGD. The results confirm the theoretical predictions on these benchmark functions. However, the experiments do not extend to deep learning architectures or complex real-world datasets, which limits the immediate practical validation of the theory's broad applicability.
Reproducibility
The paper provides detailed mathematical proofs in the appendix. The experimental setup is described with specific function definitions and parameter choices, though the code is not explicitly linked in the provided text. The theoretical results are self-contained and reproducible in a mathematical sense. The use of AI tools for proof correction is disclosed, which is a positive transparency measure, though it raises minor questions about the originality of the proof structure.
Limitations
The primary limitation is the restriction to block-separable functions with constant-size blocks. While this class includes many benchmarks, it excludes highly coupled nonconvex landscapes common in deep learning. The assumption of Gaussian noise is a simplification that may not hold for all practical SGD implementations. The "spectral transfer" assumption (Assumption 1) is a strong condition that needs to be verified for specific function classes. The paper acknowledges that the results do not yet extend to general nonconvex functions with weak inter-variable couplings.
Broader impact
This paper provides a foundational theoretical justification for one of the most ubiquitous heuristics in deep learning: learning rate decay. By proving an exponential acceleration in the small-step-size regime, it elevates learning rate decay from a "folk belief" to a provably optimal strategy for a broad class of nonconvex problems. This could influence the design of optimization algorithms for scientific computing and other fields where global optimization of nonconvex objectives is critical. The techniques developed (weak Poincaré inequalities for annealing) may be applicable to other sampling and optimization problems. The paper establishes the first provable exponential acceleration of SGD via learning rate decay for global optimization of nonconvex functions. It introduces a novel application of weak Poincaré inequalities to analyze annealed Langevin diffusion and transfers this analysis to discrete SGD, providing a rigorous theoretical explanation for the empirical success of learning rate decay in navigating complex optimization landscapes.
Author of this paper? Send a correction.