[Paper Review] Stochastic Gradient Descent with Polyak's Learning Rate

Adam M. Oberman, Mariana Prazeres|arXiv (Cornell University)|Mar 20, 2019

Stochastic Gradient Optimization Techniques25 references7 citations

TL;DR

This paper proposes a stochastic gradient descent (SGD) method with an adaptive learning rate based on Polyak's step size, which uses an estimate of the optimal function value $f^*$ to dynamically adjust the learning rate. The method achieves a non-asymptotic convergence rate of $\mathcal{O}(1/k)$ with a potentially better rate constant than optimally scheduled SGD, demonstrating improved practical convergence on convex problems and deep neural networks without hyperparameter tuning.

ABSTRACT

Stochastic gradient descent (SGD) for strongly convex functions converges at the rate $\bO(1/k)$. However, achieving good results in practice requires tuning the parameters (for example the learning rate) of the algorithm. In this paper we propose a generalization of the Polyak step size, used for subgradient methods, to Stochastic gradient descent. We prove a non-asymptotic convergence at the rate $\bO(1/k)$ with a rate constant which can be better than the corresponding rate constant for optimally scheduled SGD. We demonstrate that the method is effective in practice, and on convex optimization problems and on training deep neural networks, and compare to the theoretical rate.

Motivation & Objective

• To address the challenge of tuning learning rates in SGD for optimal convergence, especially when strong convexity parameters are unknown.
• To generalize Polyak's learning rate—originally for subgradient methods—to the stochastic gradient setting.
• To achieve a non-asymptotic convergence rate of $\mathcal{O}(1/k)$ with a potentially improved rate constant compared to scheduled SGD.
• To demonstrate practical effectiveness of the method on convex optimization and deep neural network training without manual hyperparameter tuning.

Proposed method

• Adapts Polyak's learning rate $h(x_k) = \frac{f(x_k) - f^*}{\|\partial f(x_k)\|^2}$ to stochastic gradients by estimating $f^*$ via a running best value $f_k^{\text{best}}$.
• Uses the adaptive learning rate $h_k = \frac{f(x_k) - f_k^{\text{best}} + \gamma_k}{\|\nabla_{mb}f(x_k)\|^2}$ with $\gamma_k > 0$, $\gamma_k \to 0$, $\sum \gamma_k = \infty$ for convergence.
• Employs a capped version in practice to prevent numerical instability from inaccurate $f^*$ estimates.
• Derives non-asymptotic convergence bounds of the form $\mathbb{E}[\|x_k - x^*\|] \leq \frac{C}{k + \gamma}$, valid at every iteration $k > 0$.
• Compares the rate constant $\alpha_P$ of Polyak SGD to the optimal scheduled SGD rate constant $\alpha_S$, showing $\alpha_P$ can be smaller under favorable initial conditions.
• Applies the method to quadratic problems and CIFAR-10 with AllCNN, using PyTorch implementation with $f^*$ estimated from a prior run.

Experimental results

Research questions

• RQ1 Can Polyak's learning rate be generalized to stochastic gradient descent to achieve better convergence constants?
• RQ2 Does the adaptive Polyak SGD method achieve a non-asymptotic $\mathcal{O}(1/k)$ convergence rate with a better rate constant than scheduled SGD?
• RQ3 How does the method perform in practice on convex problems and deep learning tasks compared to tuned SGD schedules?
• RQ4 Can the method reduce or eliminate the need for manual learning rate tuning in deep learning, especially when $f^*$ is estimated from prior training?

Key findings

• The Polyak SGD method achieves a non-asymptotic convergence rate of $\mathcal{O}(1/k)$ with a rate constant $\alpha_P = \frac{2\mu^2}{\sigma^2 + 2\mu^2(L - \mu)q_0}$, which can be smaller than the optimal scheduled SGD rate constant $\alpha_S$ when $q_0 \leq \frac{M^2}{2\mu^2(L - \mu)}$.
• In numerical experiments on a quadratic problem, Polyak SGD outperforms both standard and optimally scheduled SGD in average excess loss over 40 runs, with tighter theoretical bounds.
• When initialized close to the optimum, Polyak SGD rapidly reduces the learning rate and avoids overshooting, while scheduled SGD with a fixed schedule fails to exploit the good initialization.
• On CIFAR-10 with AllCNN, Polyak SGD achieves similar training loss and test error to a well-tuned SGD schedule, despite requiring no learning rate tuning.
• The method requires only one hyperparameter—$f^*$—which can be estimated from a prior training run, making it practical for repeated training scenarios.
• The theoretical bounds are non-asymptotic and hold at every iteration $k > 0$, unlike many standard convergence results that only apply in the limit.