Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution

TL;DR

The dynamics and the convergence properties of SGD equipped with new variants of the stochastic Polyak stepsize are studied and a novel modification of SPS is proposed, which guarantees convergence to the exact minimizer - without a priori knowledge of the problem parameters.

Abstract

Recently, Loizou et al. (2021), proposed and analyzed stochastic gradient descent (SGD) with stochastic Polyak stepsize (SPS). The proposed SPS comes with strong convergence guarantees and competitive performance; however, it has two main drawbacks when it is used in non-over-parameterized regimes: (i) It requires a priori knowledge of the optimal mini-batch losses, which are not available when the interpolation condition is not satisfied (e.g., regularized objectives), and (ii) it guarantees convergence only to a neighborhood of the solution. In this work, we study the dynamics and the convergence properties of SGD equipped with new variants of the stochastic Polyak stepsize and provide solutions to both drawbacks of the original SPS. We first show that a simple modification of the original SPS that uses lower bounds instead of the optimal function values can directly solve issue (i). On the other hand, solving issue (ii) turns out to be more challenging and leads us to valuable insights into the method's behavior. We show that if interpolation is not satisfied, the correlation between SPS and stochastic gradients introduces a bias, which effectively distorts the expectation of the gradient signal near minimizers, leading to non-convergence - even if the stepsize is scaled down during training. To fix this issue, we propose DecSPS, a novel modification of SPS, which guarantees convergence to the exact minimizer - without a priori knowledge of the problem parameters. For strongly-convex optimization problems, DecSPS is the first stochastic adaptive optimization method that converges to the exact solution without restrictive assumptions like bounded iterates/gradients.

Figures (14)

• Figure 1: We consider a $100$ dim problem with $n=100$ datapoints where each $f_i = \frac{1}{2}(x-x_i^*)^\top H_i(x-x_i^*) + f_i^*$, with $f_i^*=1$ for all $i\in[n]$ and $H_i$ a random SPD matrix generated using the standard Gaussian matrix $A_i\in\mathbb{R}^{d\times 3d}$ as $H_i = A_iA_i^\top/(3d)$. If $x_i^*\ne x_j^*$ for $i\ne j$, then the problem does not satisfy interpolation (left plot). Instead, if all $x_i^*$s are equal, then the problem is interpolated (central plot). The plot shows the behaviour of SPS$^{\ell}_{\max}$ ($\gamma_b=2)$ for different choices of the approximated suboptimality $\ell_i^*$. We plot (mean and std deviation over 10 runs) the function suboptimality level $f(x)-f(x^*)$ for different values of $\ell_i^*$. Note that, if instead $f_i^*=0$ for all $i$ then all the shown algorithms coincide (right plot) and converge to the solution.
• Figure 2: Dynamics of SPS$_{\max}$ with decreasing multiplicative constant (SGD style) compared with DecSPS. We compared both in the interpolated setting (right) and in the non-interpolated setting (left). In the non-interpolated setting, a simple multiplicative factor introduces a bias in the final solution, as discussed in this section. We consider two dimensional $f_i=\frac{1}{2}(x-x_i^*)^\top H_i (x-x_i^*)$, for $i = 1,2$ and plot the contour lines of the corresponding landscapes, as well as the average landscape $(f_1+f_2)/2$ we seek to minimize. Solution is denoted with a gold star.
• Figure 3: DecSPS ($c_k=c_0\sqrt{k+1}$) sensitivity to hyperparameters on the Synthetic Dataset, with $\lambda=0$. Repeated 10 times and plotted is mean and std.
• Figure 4: Left: performance of DecSPS, on the A1A Dataset ($\lambda = 0.01$). Right: performance of DecSPS on the Breast Cancer Dataset ($\lambda = 1e-1$). Further experiments can be found in §\ref{['app:exp']} (appendix).
• Figure 5: Left: Performance of Adam (with fixed stepsize and no momentum) and Right: AMSgrad (with sqrt decreasing stepsize and no momentum) compared to DecSPS on the A1A and Breast Cancer dataset, respectively. Plots comparing the performance of Adam with DecSPS on the Breast Cancer Dataset can be found in Figure \ref{['Fig_SGDVsAdam_Breast']}, and plots comparing AMSgrad with DecSPS on the A1A Dataset can be found in Figure \ref{['Fig_SGDVsAdam_A1A']}. Plotted is also the average stepsize (each parameter evolves with a different stepsize).

Theorems & Definitions (46)

• Theorem 1: Main result of loizou2021stochastic
• Remark 1: What if $\|\nabla f_{\mathcal{S}_k}\|=0$?
• Theorem 2
• Corollary 1
• Remark 2
• Lemma 1
• Remark 3
• Theorem 3
• Corollary 2
• Remark 4: Beyond bounded iterates