PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence

Nicolas Loizou, Sharan Vaswani, Issam Laradji, Simon Lacoste-Julien

arXiv:2002.10542v3math.OCcs.LGstat.ML

TL;DR

SGD step-size selection is central, and the Polyak step-size’s required optimal function values are often available in machine-learning problems. The paper proposes and analyzes SPS, proving convergence guarantees across settings and reporting strong experimental performance, especially for interpolating over-parameterized models.

  • Problem

    Step-size selection is the main parameter for guaranteeing SGD convergence, while Polyak-step-size information is often readily available for standard machine-learning losses.

  • Method

    The paper proposes SPS, a stochastic adaptive learning rate extending the classical Polyak step-size, and theoretically analyzes SGD equipped with it across multiple settings.

  • Results

    The analysis establishes convergence guarantees for strongly convex, convex, and non-smooth convex settings, while experiments show strong performance against state-of-the-art optimization methods.

  • Takeaways & Limitations

    For over-parameterized models satisfying interpolation, SPS supports fast convergence to the true solution without problem-dependent constants or additional computational overhead.

  • Takeaways & Limitations

    SPS analysis must handle correlations from adaptive step-sizes and cannot use the descent lemma when adapting to the Lipschitz constant, complicating the convex proof.

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces SPS, an adaptive learning rate for SGD, and analyzes its convergence across convexity settings. It emphasizes interpolation, where SPS can achieve fast convergence without problem-dependent constants or extra computational overhead, supported by experiments.

  • SPS extends the classical Polyak step-size to stochastic gradient descent as an adaptive learning-rate rule.
  • The paper provides convergence guarantees for SGD with SPS in strongly convex, convex, and non-convex smooth settings.
  • SPS analysis also yields novel convergence results for constant step-size SGD, including convergence to a neighborhood characterized without some additional variance assumptions.
  • Under interpolation, SPS enables SGD to converge to the true solution at a fast rate matching the deterministic case.
  • In the interpolation setting, SPS requires neither problem-dependent constants nor additional computational overhead.
  • Experiments on synthetic and over-parameterized-model tasks show superior convergence for SGD with SPS relative to state-of-the-art optimization methods.

2 SGD and the Stochastic Polyak Step-size

SGD uses stochastic examples and a per-iteration step-size; SPS adapts the learning rate from the current stochastic loss and gradient, extending Polyak’s method while retaining low overhead.

  • SGD selects example i uniformly at random and uses γk > 0 as the step-size at iteration k.
  • The deterministic Polyak step-size: The deterministic Polyak step-size minimizes an upper bound on the next iterate’s distance to the optimum using a subgradient and optimal function value.
  • Stochastic Polyak Step-size: Applying the deterministic rule directly to SGD is impractical because it requires the full function value and full gradient at every iteration.
  • Stochastic Polyak Step-size: SPS replaces those full quantities with the current stochastic loss fi(xk) and stochastic gradient ∇fi(xk), computed without additional update cost.
  • The SPS parameter c controls the step-size, while SPSmax caps it at γb to ensure convergence to a small neighborhood around the solution.
  • SPS requires knowledge of f*i, which is often available for standard unregularized or regularized losses and need not be attained by fi.
  • Optimal Objective Difference: The analysis assumes finite optimal objective difference rather than finite gradient noise; interpolation settings yield σ = 0.

3 Convergence Analysis

The analysis establishes convergence guarantees for SPS across strongly convex, convex, PL, and broader non-convex settings, including interpolation and constant-step-size regimes. The guarantees avoid several common assumptions and characterize convergence either to the optimum or to a controlled neighborhood.

  • SPS provides convergence guarantees for strongly convex, convex, and non-convex smooth objectives, with proofs developed in the main convergence-analysis section.
  • Strongly convex objectives: For strongly convex objectives, the best convergence rate and tightest neighborhood are obtained with c = 1/2.The result uses α := min{1/(2cLmax), γb}.
  • Assumptions: SPS guarantees convergence without bounded-gradient or growth assumptions because the step-size provides a natural bound on stochastic-gradient norms.
  • Strongly convex objectives: Interpolation allows SPS to converge without an upper bound on the step-size, while c = 1/2 yields the stated strongly-convex convergence result.The interpolation setting includes losses that attain their minima at a common solution.
  • Constant step-size: Constant-step-size SGD with SPS has a convergence neighborhood depending on the optimal objective difference σ² rather than the variance z² = E[∥∇f_i(x*)∥²].This is stated as a first result of this form to the authors’ knowledge.
  • Convex objectives: For convex objectives, interpolation with unbounded SPS yields convergence to the optimum at an O(1/K) rate, while bounded variants recover constant-step-size SGD rates.
  • PL objectives: For PL objectives, constant-step-size SGD converges under smoothness without bounded gradient, bounded variance, or interpolation assumptions.Under interpolation, the analysis gives linear convergence to the optimum with a constant step-size.
  • General non-convex objectives: Under the strong growth condition, SPS achieves O(1/K) convergence to a neighborhood governed by δ; with interpolation, it reaches a first-order stationary point as efficiently as deterministic gradient descent.

4 Experimental Evaluation

Experiments validate the theoretical results on synthetic problems and evaluate SPS against established optimizers across over-parameterized models. SPS is robust across settings and performs strongly without SLS’s backtracking procedure, while unbounded SPS can be sensitive to interpolation misspecification.

  • 4.2 Experiments for over-parametrized models: SPS experiments cover deep matrix factorization, kernel-based binary classification, and multi-class CIFAR classification with deep networks.The optimizer comparisons include Adam, SLS, ALI-G, RADAM, and Look-ahead, with results averaged across five independent runs.
  • 4.1 Synthetic experiments: Synthetic experiments benchmark SPSmax against constant step-size SGD on regularized and unregularized logistic-loss binary classification.The dataset has n = 1k examples and dimension d = 100; experiments use γb = {1, 5, 100} and compare constant step-sizes γ = {0.1, 0.01}.
  • 4.1 Synthetic experiments: Constant step-size SGD is not robust to learning-rate choice, showing good convergence at 0.1 but slow convergence at 0.01.
  • Practical considerations: Unbounded SPS is not robust to misspecified interpolation, while SPSmax may require problem-dependent tuning of γb; smoothing limits step-size fluctuations.The smoothing procedure uses an adaptive iteration-dependent upper bound γb and τ = 2 in the experiments.
  • Deep matrix factorization: SPS is robust to the degree of over-parametrization and matches SLS in deep matrix factorization, without SLS’s expensive backtracking procedure.For rank k = 4 interpolation is not satisfied, whereas k = 10 exactly satisfies interpolation.
  • Kernel and deep-network classification: SPS shows strong performance relative to other optimizers for kernel classification and matches SLS on the reported deep-network experiments.The deep-network experiments use CIFAR-10 and CIFAR-100 with ResNet-34; SPS and SLS also exhibit cyclic step-size behavior there.

5 Conclusion

The paper proposes and analyzes SPS, establishes convergence guarantees across several settings, and derives new results for constant-step-size SGD. Experiments show strong performance across diverse tasks, while extensions to accelerated, mini-batch, sampling, distributed, and decentralized settings remain open.

  • The paper proposes and theoretically analyzes a stochastic variant of the classical Polyak step-size.
  • The analysis covers strongly convex, convex, and non-convex settings and also yields new convergence guarantees for constant-step-size SGD.
  • Experiments across diverse tasks show strong performance for SGD with SPS compared with state-of-the-art optimization methods.
  • Potential extensions include accelerated methods, mini-batching, non-uniform sampling, and distributed or decentralized optimization.

Supplementary Material Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence

The supplementary material supplies definitions, background on the deterministic Polyak step-size, and proofs for SPS convergence results. It also compares theoretical rates and neighborhoods with prior methods under interpolation and related assumptions.

  • Basic definitions: The supplement defines strong convexity, convexity, the Polyak-Lojasiewicz condition, and L-smoothness.
  • Deterministic Polyak step-size: The deterministic Polyak step-size minimizes an upper bound on the next iterate’s distance to the optimum, using the optimal function value.The classical step-size requires knowledge of f*, although some applications have f* = 0.
  • SPS convergence analysis: Under interpolation, SPSmax with c = 1/2 converges at the same rate as the corresponding result with Lmax = L.The comparison corollary assumes every fi is μ-strongly convex and L-smooth, a stronger condition than the main theorem’s assumptions.
  • Comparison with prior methods: Under interpolation, SPS converges linearly to x* while Berrada et al.’s method converges to a neighborhood proportional to δ.

B.2 Proof of Theorem 3.4

The proof analyzes constant-step-size SGD within the SPS framework and compares the resulting smooth-convex guarantees with prior stochastic line-search and adaptive methods. Under interpolation, the paper reports convergence to the optimum at O(1/K), while c = 1 is chosen for a simpler bound rather than optimality.

  • Proof strategy: When α = γb, the analyzed method reduces to constant-step-size SGD with γk = γb.
  • Proof strategy: The proof derives the convergence bound by conditioning on iterates, taking expectations, summing from k = 0 to K − 1, and dividing by K.
  • Parameter choice: c = 1 simplifies the upper-bound expression but is not optimal; the optimal c depends on σ and the desired accuracy ε.
  • Comparison with prior methods: Under interpolation, SPS converges to x* at an O(1/K) rate, whereas Berrada et al.’s method converges to a neighborhood proportional to δ.
  • Comparison with prior methods: In the interpolation setting, SPS has a rate similar to stochastic line search.

B.3 SPS on Methods for Solving Consistent Linear Systems

For consistent linear systems, SPS with c = 1/2 reduces to the theoretically optimal constant step-size and achieves the corresponding optimal convergence rate. The section situates this result within stochastic reformulations and sketch-and-project methods.

  • Problem reformulation: The stochastic optimization reformulation has the same solution set as the original consistent linear system.Its minimizers equal the system’s solution set L.
  • Method connections: Randomized Kaczmarz and related sketch-and-project algorithms arise as special cases of the general SGD update under different sampling distributions.The framework also includes Gaussian Kaczmarz and randomized coordinate descent variants.
  • SPS reduction: SPS with c = 1/2 becomes the theoretically optimal constant step-size for solving consistent linear systems.The reduction applies to the stochastic optimization reformulation of a consistent linear system.
  • Convergence guarantee: SPS achieves an optimal convergence rate when solving consistent linear systems.The stated rate matches the expression previously proposed for sketch-and-project methods.
  • Proof strategy: The analysis derives the result by selecting a step-size that minimizes an upper bound and recursively unrolling the resulting expected-error relation.The proof uses smoothness, conditional expectations, and a geometric-series argument.
  • Assumptions: The convergence proof requires parameter restrictions on the upper step-size bound and c to ensure positive contraction quantities.The section analyzes separate parameter cases and imposes conditions such as c > Lρ/(4Lmax) in one result.

B.5.1 Additional convergence result for nonconvex smooth functions: Assuming independence of step-size and stochastic gradient

Under independence between the adaptive step-size and stochastic gradient, the analysis provides an additional convergence theorem for smooth functions. The section also extends convergence results to non-smooth convex functions and non-interpolated settings through mini-batch growth.

  • Independence assumption: The additional smooth-function theorem assumes independence between γk and the stochastic gradient at every iteration.The proof uses this assumption together with bounds on SPS rather than its exact form.
  • Convergence theorem: The independence-based result applies to SGD with SPSmax under restrictions on c and the upper step-size bound γb.The stated condition includes c > ρL/(4Lmax) and an upper bound on γb.
  • Non-smooth interpolation: For non-smooth convex functions in the interpolated setting, stochastic subgradient descent with the SPS counterpart has an O(1/K) convergence rate.The result assumes convexity, non-smoothness, interpolation, and uniformly bounded subgradients.
  • Beyond interpolation: Without interpolation, progressively increasing the mini-batch size can guarantee convergence to the optimal solution.The section presents this strategy as an extension beyond the interpolated case.
  • Non-smooth assumptions: The non-smooth analysis replaces stochastic gradients with stochastic subgradients and relies on a G-Lipschitz assumption.A constant G bounds every subgradient norm.

C.2 Increasing Mini-batch Size

The section studies progressively increasing mini-batch sizes as a way to obtain convergence to the optimal solution without interpolation. Results cover strongly convex component functions and functions satisfying the PL inequality.

  • Motivation: Progressively increasing the mini-batch size can guarantee convergence to the optimal solution without interpolation.The section presents this approach as an alternative to restricting the analysis to interpolated problems.
  • Strong convexity: For strongly convex component functions, SGD with SPS and progressive batch-size growth satisfies a stated convergence guarantee.The batch size bk is required to follow the theorem’s prescribed condition.
  • PL condition: For PL component functions, SGD with SPS and progressive batch-size growth also satisfies a stated convergence guarantee.The result uses the variance quantity z2 = supx E||∇fi(x) − ∇f(x)||2.
  • Variance control: The batch-size conditions depend on controlling stochastic-gradient variance through z2.Both theorem statements define z2 using the deviation between stochastic and full gradients.

D Computing f ∗ i for ℓ2-regularized standard surrogate losses

The paper derives closed-form expressions for optimal individual loss values in ℓ2-regularized binary surrogate-loss problems. The derivations use the r-Lambert function, with the classical Lambert W function as a special case.

  • Computation: The required f∗i values can also be computed efficiently to machine precision by solving one-dimensional strongly convex problems with Newton’s method.The paper notes this as an alternative to numerical routines for Lambert-type functions.
  • Special functions: The r-Lambert function generalizes the classical Lambert W function and solves equations of the form x e^(cx) + r x = a.The classical function is recovered when r = 0.
  • Scope: Closed-form expressions for f∗i are derived for ℓ2-regularized standard binary surrogate losses.The section focuses on binary logistic and exponential losses.
  • Binary log-loss: The ℓ2-regularized logistic-loss minimum is expressed using the r-Lambert function.The derivation reduces the problem to a one-dimensional strongly convex minimization and uses r = 1.
  • Binary exponential loss: The ℓ2-regularized binary exponential-loss minimum is expressed using the classical Lambert W function.This corresponds to the r-Lambert formulation with r = 0.

E Additional Experiments

Additional experiments evaluate SGD with SPS on over-parameterized models across deep matrix factorization, kernel-based binary classification, and deep-network multi-class classification.

  • The appendix further evaluates SGD with SPS when training over-parameterized models.
  • Experiments cover deep matrix factorization, binary classification using kernels, and multi-class classification using deep neural networks.
Loading 2002.10542v3…