Faster Convergence of Stochastic Accelerated Gradient Descent under Interpolation

A continuing trend in machine learning is the adoption of powerful prediction models which can exactly fit, or interpolate, their training data [1]. Methods such as over-parameterized neural networks [2, 3], kernel machines [4], and boosting [5] have all been shown to achieve zero training loss in practice. This phenomena is particularly prevalent in modern deep learning, where interpolation is conjectured to be key to both optimization [6, 7] and generalization [8].

Recent experimental and theoretical evidence shows stochastic gradient descent (SGD) matches the fast convergence rates of deterministic gradient methods up to problem-dependent constants when training interpolating models [9, 10, 11]. With additional assumptions, interpolation also implies the strong [12] and weak [13, 14] growth conditions, which bound the second moment of the stochastic gradients. Under strong/weak growth, variance-reduced algorithms typically exhibit slower convergence than stochastic gradient methods despite using more computation or memory [15, 10], perhaps because these conditions already imply a form of “automatic variance reduction” [6]. A combination of interpolation and growth conditions has been used to prove fast convergence rates for SGD with line-search [14], with the stochastic Polyak step-size [16, 17], for mirror descent [18], and for model-based methods [19].

While these results show interpolation is sufficient to break the $\Omega(\epsilon^{-4})$ complexity barrier for computing stationary points of smooth, convex functions with stochastic, first-order oracles [20], significantly less work has been done to obtain the accelerated rates possible in the deterministic setting [21]. Vaswani et al. [22] analyze a stochastic version of Nesterov’s accelerated gradient method (AGD) [23] under the strong growth condition, but their bounds have a linear dependence on the strong growth constant and can be slower than SGD [24]. In contrast, Liu and Belkin [24] propose a modified version of stochastic AGD and extend the statistical condition number approach of Jain et al. [25] to the interpolation setting. However, their results apply primarily to quadratics and are not accelerated for general convex functions.

In this work, we apply the estimating sequences analysis developed by Nesterov [26] to the interpolation setting. Our approach hinges on a simple, in-expectation progress guarantee for SGD, which we prove is a sufficient condition for generic acceleration of stochastic algorithms. This proof technique is completely different from that used by Vaswani et al. [22] and yields an improved dependence on the strong growth constant. In the worst-case, the improvement is proportional to the square-root of the condition number and guarantees stochastic AGD is always at least as fast as SGD.  In what follows, all proofs are deferred to the corresponding section of the appendix.

Consider the unconstrained minimization of a smooth, convex function $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$. We assume $f$ has at least one minimizer $w^{*}$ and denote the optimal set by $\mathcal{W}^{*}$. At each iteration $k$, we sample a stochastic gradient $\nabla f(w_{k},z_{k})$ such that,

meaning the stochastic gradient is unbiased. We assume that $z_{k},z_{j}$ are independent for $k\neq j$, but they do not need to have the same distribution. The stochastic gradients satisfy the strong growth condition when there exists $\rho\geq 1$ for which

holds given any sequence $\left\{w_{k},z_{k}\right\}$. We say that interpolation is satisfied if

for all $z_{k}$. That is, stationarity of $f$ implies stationarity of the stochastic gradients. Although the strong growth condition implies interpolation, we will see that the converse requires further assumptions on $f$ and on the stochastic gradients.

We assume that $f$ is $L$-smooth, meaning $\nabla f$ is $L$-Lipschitz continuous. $L$-smoothness of $f$ implies the following quadratic upper-bound for all $w,u\in\mathbb{R}^{d}$ [26]:

Similarly, we will sometimes require the stochastic gradients to be $L_{\text{max}}$-individually smooth, meaning they satisfy

almost surely for all $k$. We also assume that $f$ is $\mu$-strongly convex, by which we mean

holds for all $w,u\in\mathbb{R}^{d}$ and some $\mu\geq 0$. When $\mu=0$, strong convexity reduces to convexity of $f$. If $f$ is strongly convex with $\mu>0$, the stochastic gradients are $L_{\text{max}}$-individually smooth, and interpolation holds, then the strong growth constant is bounded as $\rho\leq L_{\text{max}}/\mu$ (see Lemma 13). Recalling that the ratio $\kappa=L/\mu$ is the condition number of $f$, we see that the strong-growth constant is bounded by a quantity like the condition number of the worst-conditioned stochastic function. ¹¹1Note that $f(u,z_{k})$ is typically not strongly convex, so this analogy is not formal.

Our analysis in this section builds on the estimating sequences approach developed by Nesterov [44]. However, we consider variable step-sizes $\eta_{k}$ and allow the iterates $w_{k}$ to be set using a general update scheme. The procedure takes $\gamma_{0}>0$ as input and uses the following updates, which we call generalized stochastic AGD:

where $m$ is an as-yet unspecified update for the “primal sequence” $w_{k}$ and $v_{0}=w_{0}$ (which implies $y_{0}=w_{0}$). Note that the step-size $\eta_{k}$ is required at the start of the iteration to compute $\alpha_{k}$. As a result, local search methods like the Armijo line-search [45] must re-evaluate $\gamma_{k+1},y_{k}$, and $w_{k+1}$ for each candidate step-size.

Choosing $m$ to be one step of SGD with step-size $\eta_{k}$ yields the familiar updates of the standard version of stochastic AGD (see Nesterov [26, Eq. 2.2.20]):

Our approach hinges on the fact that, under strong growth, stochastic gradient updates for $w_{k}$ give a similar per-step progress condition as deterministic gradient descent. Since descent in the primal step is the only link between $w_{k}$ and the “dual sequence” $y_{k}$, generalized stochastic AGD with any primal update obtaining similar per-iteration progress can be analyzed in the same fashion as stochastic AGD. This allows us to derive a fast convergence rate for the general scheme in Eq. 5.

We start by deriving the progress condition for SGD.  It is straightforward to prove the following bound using $L$-smoothness, the definition of $w_{k+1}$, and strong growth:

Substituting any fixed step-size $\eta_{k}\leq 1/\rho L$ into Eq. 7 gives the following equation, which we call the expected progress condition:

which is equivalent to the progress made by gradient descent with step-size $\eta_{k}\leq 1/L$ up to a factor of $\rho$ [46]. In order to make use of the expected progress condition, we now introduce the estimating sequences framework.

Unlike the standard definition, we permit $\lambda_{k}$ and $\phi_{k}$ to depend on $z_{0},\ldots,z_{k-1}$, making them random variables. Typically the first function $\phi_{0}$ is deterministic and chosen to satisfy $\phi_{0}(w)\geq f(w)$ for all $w$ near $w_{0}$. Since Eq. 9 guarantees $\lim_{k}\phi_{k}(w)\leq f(w)$, $\phi_{k}$ can be interpreted as a sequence of relaxing upper-bounds, where the rate of relaxation is controlled by $\lambda_{k}$. The next condition captures when $\phi_{k}$ is a good local model of $f$.

In what follows, we use $\mathbb{E}$ without subscripts to denote the total expectation with respect to $z_{0},\ldots,z_{k-1}$. That is, all randomness in the procedure up to iteration $k$. If $\phi_{k}$ maintains the local upper-bound property in expectation, then Eq. 9 guarantees

which shows that generalized stochastic AGD converges in expectation at a rate controlled by $\lambda_{k}$. As a result, the bulk of our analysis focuses on establishing the local upper-bound property for a suitable choice of estimating sequences.

Following Nesterov [26], choose $\lambda_{0}=1$, $\phi_{0}(w)=f(w_{0})+\frac{\gamma_{0}}{2}\|w-w_{0}\|_{2}^{2}$, and

The initial curvature $\gamma_{0}$ is an input parameter; differentiating shows that $v_{k+1}$ is actually the minimizer of $\phi_{k+1}$, while $\nabla^{2}\phi_{k+1}=\gamma_{k+1}I$ (Lemma 14). Thus, the auxillary sequences $\gamma_{k},v_{k}$ can be viewed as arising from our choice of local model. The next lemma proves these are valid estimating sequences when the step-size sequence is well-behaved. In what follows, we use the convention $1/0=\infty$ to cover the case of non-strongly convex functions.

The parameter $\eta_{\text{min}}>0$ is required for $\lambda_{k}$ to decrease sufficiently fast, while the upper-bound $\eta_{k}\leq 1/\mu$ is only necessary when $\mu>0$. In this case, it guarantees $\lambda_{k}\geq 0$. This choice of estimating sequences also satisfies the local error-bound property in expectation when $m(\eta_{k},y_{k},\nabla f(y_{k},z_{k}))$ matches the progress of fixed step-size SGD.

Proposition 5 is our main theoretical contribution and immediately leads to two accelerated convergence rates for the generalized stochastic AGD scheme.

This linear rate of convergence requires knowledge of the strong convexity constant in order to set $\gamma_{0}$. However, we can still obtain a $O(1/k^{2})$ rate without knowing $\mu$ so long as the smoothness constant can be estimated. The following theorem is a generalization of Nesterov [26] to stochastic optimization under interpolation.

Now we compare our rates to those existing in the literature. Throughout this section, we assume that $f$ is individually smooth, interpolation holds, and the strong growth condition is satisfied. Recall that the strong growth constant is bounded above as $\rho\leq L_{\text{max}}/\mu$ under these conditions. This worst-case bound on $\rho$ is critical to understanding when stochastic AGD does or does not accelerate.

Before proceeding, we introduce the notion of statistical condition number proposed by Jain et al. [25] and used by Liu and Belkin [24] to analyze their modified version of stochastic AGD (called MaSS) in the least-squares setting. Let $\mathcal{P}$ be a probability distribution over $(x,y)$ and define the least squares objective as

Define the stochastic functions and gradients to be $f_{\text{ls}}(w,z_{k})=\frac{1}{2}(x_{z_{k}}^{\top}w-y_{z_{k}})^{2}$ and $\nabla f_{\text{ls}}(w,z_{k})=x_{z_{k}}(x_{z_{k}}^{\top}w-y_{z_{k}}),$ where $(x_{z_{k}},y_{z_{k}})\sim\mathcal{P}$. These stochastic gradients are $L_{\text{max}}$-individually smooth with $L_{\text{max}}=\sup\left\{\|x_{z_{k}}\|_{2}^{2}:(x_{z_{k}},y_{z_{k}})\in\text{Supp}(\mathcal{P})\right\}$. Assuming we may interchange expectation and differentiation, the Hessian is $H=\mathbb{E}_{x}\left[xx^{\top}\right]$ and the condition number $\kappa$ and statistical condition number $\tilde{\kappa}$ are defined as,

It is straightforward to prove $\kappa,\tilde{\kappa}\leq L_{\text{max}}/\mu$, similar to the strong growth constant.

Table 1 compares our iteration complexities for stochastic AGD to the complexity of SGD under interpolation, the analysis of stochastic AGD by Vaswani et al. [22], and the complexity of MaSS. Unlike Vaswani et al. [22], who use strong growth to show both the optimality gap and distance to a minimizer decrease in expectation at each iteration, our approach only requires the sufficient progress condition. This allows us to shrink the dependence on the strong growth constant from $\rho$ to $\sqrt{\rho}$, which — since $\rho\leq L_{\text{max}}/\mu$ — can be larger than $\sqrt{\kappa}$ in the worst case. Substituting this into the complexity bound shows stochastic AGD requires $O((\sqrt{LL_{\text{max}}}/\mu)\log(1/\epsilon))$ iterations to reach $\epsilon$-sub-optimality. That is, stochastic AGD is always at least as fast SGD and faster when $L_{\text{max}}\gg L$.

Our convergence rate for stochastic AGD also improves over that for SGD under the strong growth condition [35]. The improvement is by a factor $\sqrt{\rho/\mu}$, indicating that acceleration actually shrinks the dependence on the noise level. This quite different from results in the general stochastic setting, where accelerated methods are typically more sensitive to noise [47]. For example, Schmidt et al. [36] show that the noise level must decrease faster for accelerated methods to converge compared to (proximal) SGD when interpolation does not hold. We conclude that interpolation seems to be key when proving fast rates for stochastic AGD.

Comparing our results against MaSS is more difficult due to the dependence on $\tilde{\kappa}$. To understand the difference in convergence rates, we consider two finite-sum example problems. In what follows, let $e_{1},\ldots,e_{n}$ be the standard basis for $\mathbb{R}^{n}$.

As expected, stochastic AGD accelerates due to the gap between the smoothness and individual smoothness constants. In comparison, the complexity bound for MaSS is not accelerated and only matches that for SGD. The next example considers the opposite setting, where $L\approx L_{\text{max}}$ and we do not expect stochastic AGD to be faster than SGD.

As $n\rightarrow\infty$, stochastic AGD attains the same complexity as SGD and is not accelerated. In comparison, the guarantee for MaSS always matches SGD and is slower than stochastic AGD for every finite $n$. We conclude that while both methods are restricted by lower bounds on stochastic acceleration, AGD can accelerate on some simple problems where MaSS fails.

We derive new convergence rates for a generalized version of stochastic Nesterov acceleration. Our approach extends the estimating sequences framework to the stochastic setting and shows that any update scheme making sufficient progress in expectation can be accelerated. As this sufficient progress condition is satisfied by SGD under the strong growth condition, our proof immediately specializes to give fast rates for stochastic AGD. Compared to previous work, our convergence bounds improve the dependence on the strong growth constant from $\rho$ to $\sqrt{\rho}$. This improvement can be larger than the square-root of the condition number, shows stochastic AGD is at least as fast as SGD, and explains the strong empirical performance of stochastic acceleration shown by Vaswani et al. [22]. We also leverage our generalized algorithm to prove convergence guarantees for stochastic AGD with preconditioning. In particular, we show that preconditioning further speeds-up accelerated SGD when the stochastic gradients are small in the matrix norm induced by the preconditioner.

In addition to these results, the utility of our theoretical approach is further demonstrated by recent literature. Our core result for stochastic AGD (Proposition 5) was previously made available in a master’s thesis [48] and the proof technique has since been leveraged to give optimal bounds for stochastic acceleration in the general setting [49]. Yet, several questions remain unanswered. For example, the convergence of stochastic AGD under relaxed conditions, like weak growth or with a stochastic line-search [22], has not been proved. And while our generalized AGD scheme also suggests accelerating methods like stochastic proximal-point, establishing the expected progress condition appears difficult and new insights may be required. We leave these questions to future work.

Acknowledgments

We would like to thank Frederik Kunstner, Victor Sanchez-Portella, and Sharan Vaswani for many insightful discussions.

Funding

Aaron Mishkin was supported by NSF GRF Grant No. DGE-1656518 and by NSERC PGS D Grant No. PGSD3-547242-2020. Mert Pilanci was supported by the NSF under Grant ECCS-2037304 and Grant DMS-2134248, by an NSF CAREER Award under Grant CCF-2236829, by the U.S. Army Research Office Early Career Award under Grant W911NF-21-1-0242, and by the Stanford Precourt Institute. Mark Schmidt was partially supported by the Canada CIFAR AI Chair Program and NSERC Discovery Grant No. RGPIN-2022-03669.

Data Availability

This paper does not rely any data, public or private. Derivations for all theoretical results are included in the appendix.