[Paper Review] Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence

Nicolas Loizou, Sharan Vaswani+2 more|arXiv (Cornell University)|Feb 24, 2020

Stochastic Gradient Optimization TechniquesComputer Science57 references37 citations

TL;DR

Introduces SPS, a stochastic Polyak step-size for SGD that adapts learning rates using f_i* and f_i(x) to achieve fast convergence across strongly convex, convex, and non-convex settings, with strong results in interpolation regimes.

ABSTRACT

We propose a stochastic variant of the classical Polyak step-size (Polyak, 1987) commonly used in the subgradient method. Although computing the Polyak step-size requires knowledge of the optimal function values, this information is readily available for typical modern machine learning applications. Consequently, the proposed stochastic Polyak step-size (SPS) is an attractive choice for setting the learning rate for stochastic gradient descent (SGD). We provide theoretical convergence guarantees for SGD equipped with SPS in different settings, including strongly convex, convex and non-convex functions. Furthermore, our analysis results in novel convergence guarantees for SGD with a constant step-size. We show that SPS is particularly effective when training over-parameterized models capable of interpolating the training data. In this setting, we prove that SPS enables SGD to converge to the true solution at a fast rate without requiring the knowledge of any problem-dependent constants or additional computational overhead. We experimentally validate our theoretical results via extensive experiments on synthetic and real datasets. We demonstrate the strong performance of SGD with SPS compared to state-of-the-art optimization methods when training over-parameterized models.

Motivation & Objective

• Motivate and address step-size selection for SGD in finite-sum learning problems.
• Introduce the stochastic Polyak step-size (SPS) as an adaptive learning rate for SGD.
• Provide theoretical convergence guarantees for SPS across strong convexity, convexity, and non-convexity.
• Show that SPS can converge to the true solution efficiently in interpolation settings.
• Demonstrate empirical performance of SPS on synthetic and real datasets across varied models.

Proposed method

• Define SPS as gamma_k = (f_i(x^k) - f_i^*) / (c ||∇f_i(x^k)||^2) and its bounded variant SPS_max.
• Relate SPS to the classical deterministic Polyak step-size and discuss required knowledge of f_i^* and choice of c.
• Provide theoretical convergence results for SPS_max under strongly convex, convex, and non-convex (PL) conditions, and for constant step-size regimes.
• Analyze interpolation and over-parameterization settings where SPS achieves fast convergence to the true solution.
• Present extensions to non-smooth and streaming settings, and connect SPS to linear system solvers.

Experimental results

Research questions

• RQ1 Does the stochastic Polyak step-size (SPS) guarantee convergence of SGD across strongly convex, convex, and non-convex objective functions?
• RQ2 How does SPS compare to constant step-size SGD and other adaptive methods in interpolation (over-parameterized) regimes?
• RQ3 What are the convergence rates and neighborhood sizes for SPS and SPS_max under various smoothness and convexity assumptions?
• RQ4 Can SPS achieve convergence to the exact solution in interpolation settings without knowledge of problem-dependent constants?
• RQ5 How does SPS perform empirically on synthetic data and over-parameterized models compared to state-of-the-art optimizers?

Key findings

• SPS provides convergence guarantees for SGD in strongly convex, convex, and non-convex smooth settings.
• SPS_max yields a linear convergence rate to a neighborhood whose size depends on the upper bound γ_b and the optimal objective difference σ^2.
• In interpolation regimes, SPS enables SGD to converge to the true solution at a fast rate without needing problem-dependent constants or extra overhead.
• For constant step-size regimes, SPS reveals convergence behavior that matches or improves over traditional constant-step SGD under certain bounds.
• Empirical results show SPS outperforms several optimization methods on over-parameterized models across synthetic data, deep matrix factorization, kernel-based binary classification, and deep networks.