Accelerating Stochastic Gradient Descent For Least Squares Regression

Introduction

Stochastic gradient descent (SGD) is the workhorse algorithm for optimization in machine learning and stochastic approximation problems; improving its runtime dependencies is a central issue in large scale stochastic optimization that often arise in machine learning problems at scale (Bottou and Bousquet, 2007), where one can only resort to streaming algorithms.

This work examines these broader runtime issues for the special case of stochastic approximation in the following least squares regression problem:

In the limit of large $n$, this minimax rate is achieved by the empirical risk minimizer (ERM), which is defined as follows. Given $n$ i.i.d. samples $\mathcal{S}_{n}=\{(\mathbf{a}_{i},b_{i})\}_{i=1}^{n}$ drawn from $\mathcal{D}$, define

where $\widehat{\mathbf{x}}_{n}^{\textrm{ERM}}$ denotes the ERM over the samples $\mathcal{S}_{n}$. For the case of additive noise models (i.e. where $b=\mathbf{a}^{\top}\mathbf{x}^{*}+\epsilon$, with $\epsilon$ being independent of $\mathbf{a}$), the minimax estimation rate is $d\sigma^{2}/n$ (Kushner and Clark, 1978; Polyak and Juditsky, 1992; Lehmann and Casella, 1998; van der Vaart, 2000), i.e.:

Let us review results in convex optimization in the exact first-order oracle model. Running $t-$steps of gradient descent (Cauchy, 1847) with an exact first-order oracle yields the following guarantee:

which implies that $\mathcal{O}(\sqrt{{\kappa}_{o}})$ oracle calls are sufficient to achieve a given target accuracy. This matches the oracle lower bounds (Nesterov, 2004) that state that $\Theta(\sqrt{{\kappa}_{o}})$ calls to the exact first order oracle are necessary to achieve a given target accuracy. The conjugate gradient method (Hestenes and Stiefel, 1952) and heavy ball method (Polyak, 1964) are also known to obtain this convergence rate for solving a system of linear equations and for quadratic functions. These methods are termed fast gradient methods owing to the improvements offered by these methods over Gradient Descent.

This paper seeks to address the question: “Can we accelerate stochastic approximation in a manner similar to what has been achieved with the exact first order oracle model?”

2 A thought experiment: Is Accelerating Stochastic Approximation possible?

Let us recollect known results in stochastic approximation for the least squares regression problem (in equation 1). Running $n$-steps of tail-averaged SGD (Jain et al., 2016) (or, streaming SVRG (Frostig et al., 2015b)Streaming SVRG does not function in the stochastic first order oracle model (Frostig et al., 2015b)) provides an output $\widehat{\mathbf{x}}_{n}$ that satisfies the following excess risk bound:

where ${\kappa}$ is the condition number of the distribution, which can be upper bounded as $L/\lambda_{\textrm{min}}(\mathbf{H})$, assuming that $\|\mathbf{a}\|\leq L$ with probability one (refer to section 2 for a precise definition of ${\kappa}$). Under appropriate assumptions, these are the best known rates under the stochastic first order oracle model (see section 1.4 for further discussion). A natural implication of the bound implied by averaged SGD is that with $\widetilde{\mathcal{O}}({\kappa})$ oracle calls (Jain et al., 2016) (where, $\widetilde{\mathcal{O}}(\cdot)$ hides $\log$ factors in $d,{\kappa}$), the excess risk attains (up to constants) the (asymptotic) minimax statistical rate. Note that the excess risk bounds in stochastic approximation consist of two terms: (a) bias: which represents the dependence of the generalization error on the initial excess risk $P(\mathbf{x}_{0})-P(\mathbf{x}^{*})$, and (b) the variance: which represents the dependence of the generalization error on the noise level $\sigma^{2}$ in the problem.

A precise question regarding accelerating stochastic approximation is: “is it possible to improve the rate of decay of the bias term, while retaining (up to constants) the statistical minimax rate?” The key technical challenge in answering this question is in sharply characterizing the error accumulation of fast gradient methods in the stochastic approximation setting. Common folklore and prior work suggest otherwise: several efforts have attempted to quantify instabilities in the face of statistical or non-statistical errors (Paige, 1971; Proakis, 1974; Polyak, 1987; Greenbaum, 1989; Roy and Shynk, 1990; Sharma et al., 1998; d’Aspremont, 2008; Devolder et al., 2013, 2014; Yuan et al., 2016). Refer to section 1.4 for a discussion on robustness of acceleration to error accumulation.

Optimistically, as suggested by the gains enjoyed by accelerated methods in the exact first order oracle model, we may hope to replace the $\widetilde{\mathcal{O}}({\kappa})$ oracle calls achieved by averaged SGD to $\widetilde{\mathcal{O}}(\sqrt{{\kappa}})$. We now provide a counter example, showing that such an improvement is not possible. Consider a (discrete) distribution $\mathcal{D}$ where the input $\mathbf{a}$ is the $i^{\textrm{th}}$ standard basis vector with probability $p_{i}$, $\forall\ i=1,2,...,d$. The covariance of $\mathbf{a}$ in this case is a diagonal matrix with diagonal entries $p_{i}$. The condition number of this distribution is ${\kappa}=\frac{1}{\min_{i}p_{i}}$. In this case, it is impossible to make non-trivial reduction in error by observing fewer than ${\kappa}$ samples, since with constant probability, we would not have seen the vector corresponding to the smallest probability.

On the other hand, consider a case where the distribution $\mathcal{D}$ is a Gaussian with a large condition number ${\kappa}$. Matrix concentration informs us that (with high probability and irrespective of how large ${\kappa}$ is) after observing $n=\mathcal{O}(d)$ samples, the empirical covariance matrix will be a spectral approximation to the true covariance matrix, i.e. for some constant $c>1$, $\mathbf{H}/c\preceq\frac{1}{n}\sum_{i=1}^{n}\mathbf{a}_{i}\mathbf{a}_{i}^{\top}\preceq c\mathbf{H}$. Here, we may hope to achieve a faster convergence rate, as information theoretically $\mathcal{O}(d)$ samples suffice to obtain a non-trivial statistical estimate (see Hsu et al. (2014) for further discussion).

Figure 1 shows the behavior of SGD in these cases; both are synthetic examples in $50-$dimensions, with a condition number ${\kappa}\approx 10^{5}$ and noise level $\sigma^{2}=100$. See the figure caption for more details.

These examples suggest that if acceleration is indeed possible, then the degree of improvement (say, over averaged SGD) must depend on distributional quantities that go beyond the condition number $\kappa$. A natural conjecture is that this improvement must depend on the number of samples required to spectrally approximate the covariance matrix of the distribution; below this sample size it is not possible to obtain any non-trivial statistical estimate due to information theoretic reasons. This sample size is quantified by a notion which we refer to as the statistical condition number $\widetilde{\kappa}$ (see section 2 for a precise definition and for further discussion about $\widetilde{\kappa}$). As we will see in section 2, we have $\widetilde{\kappa}\leq{\kappa}$, $\widetilde{\kappa}$ is affine invariant, unlike ${\kappa}$ (i.e. $\widetilde{\kappa}$ is invariant to linear transformations over $\mathbf{a}$).

3 Contributions

This paper introduces an accelerated stochastic gradient descent scheme, which can be viewed as a stochastic variant of Nesterov’s accelerated gradient method (Nesterov, 2012). As pointed out in Section 1.2, the excess risk of this algorithm can be decomposed into two parts namely, bias and variance. For the stochastic approximation problem of least squares regression, this paper establishes bias contraction at a geometric rate of $\mathcal{O}(1/\sqrt{{\kappa}\widetilde{\kappa}})$, improving over prior results (Frostig et al., 2015b; Jain et al., 2016),which prove a geometric rate of $\mathcal{O}(1/{\kappa})$, while retaining statistical minimax rates (up to constants) for the variance. Here ${\kappa}$ is the condition number and $\widetilde{\kappa}$ is the statistical condition number of the distribution, and a rate of $\mathcal{O}(1/\sqrt{{\kappa}\widetilde{\kappa}})$ is an improvement over $\mathcal{O}(1/{\kappa})$ since $\widetilde{\kappa}\leq{\kappa}$ (see Section 2 for definitions and a short proof of $\widetilde{\kappa}\leq{\kappa}$).

See Table 1 for a theoretical comparison. Figure 2 provides an empirical comparison of the proposed (tail-averaged) accelerated algorithm to (tail-averaged) SGD (Jain et al., 2016) on our two running examples. Our result gives improvement over SGD even in the noiseless (i.e. realizable) case where $\sigma=0$; this case is equivalent to the setting where we have a distribution over a (possibly infinite) set of consistent linear equations. See Figure 3 for a comparison on the case where $\sigma=0$.

On a more technical note, this paper introduces two new techniques in order to analyze the proposed accelerated stochastic gradient method: (a) the paper introduces a new potential function in order to show faster rates of decaying the bias, and (b) the paper provides a sharp understanding of the behavior of the proposed accelerated stochastic gradient descent updates as a stochastic process and utilizes this in providing a near-exact estimate of the covariance of its iterates. This viewpoint is critical in order to prove that the algorithm achieves the statistical minimax rate.

We use the operator viewpoint for analyzing stochastic gradient methods, introduced in Défossez and Bach (2015). This viewpoint was also used in Dieuleveut and Bach (2015); Jain et al. (2016).

4 Related Work

Stochastic gradient descent (SGD) and its variants are by far the most widely studied algorithms for the stochastic approximation problem. While initial works (Robbins and Monro, 1951) considered the final iterate of SGD, later works (Ruppert, 1988; Polyak and Juditsky, 1992) demonstrated that averaged SGD obtains statistically optimal estimation rates. Several works provide non-asymptotic analyses for averaged SGD and variants (Bach and Moulines, 2011; Bach, 2014; Frostig et al., 2015b) for various stochastic approximation problems. For stochastic approximation with least squares regression Bach and Moulines (2013); Défossez and Bach (2015); Needell et al. (2016); Frostig et al. (2015b); Jain et al. (2016) provide non-asymptotic analysis of the behavior of SGD and its variants. Défossez and Bach (2015); Dieuleveut and Bach (2015) provide non-asymptotic results which achieve the minimax rate on the variance (where the bias is lower order, not geometric). Needell et al. (2016) achieves a geometric rate on the bias (and where the variance is not minimax). Frostig et al. (2015b); Jain et al. (2016) obtain both the minimax rate on the variance and a geometric rate on the bias, as seen in equation 4.

While there have been several attempts at understanding if it is possible to accelerate SGD , the results have been largely negative. With regards to acceleration with adversarial (non-statistical) errors in the exact first order oracle model, d’Aspremont (2008) provide negative results and Devolder et al. (2013, 2014) provide lower bounds showing that fast gradient methods do not improve upon standard gradient methods. There is also a series of works considering statistical errors. Polyak (1987) suggests that the relative merits of heavy ball (HB) method (Polyak, 1964) in the noiseless case vanish with noise unless strong assumptions on the noise model are considered; an instance of this is when the noise variance decays as the iterates approach the minimizer. The Conjugate Gradient (CG) method (Hestenes and Stiefel, 1952) is suggested to face similar robustness issues in the face of statistical errors (Polyak, 1987); this is in addition to the issues that CG is known to suffer from owing to roundoff errors (due to finite precision arithmetic) (Paige, 1971; Greenbaum, 1989). In the signal processing literature, where SGD goes by Least Mean Squares (LMS) (Widrow and Stearns, 1985), there have been efforts that date to several decades (Proakis, 1974; Roy and Shynk, 1990; Sharma et al., 1998) which study accelerated LMS methods (stochastic variants of CG/HB) in the same oracle model as the one considered by this paper (equation 2). These efforts consider the final iterate (i.e. no iterate averaging) of accelerated LMS methods with a fixed step-size and conclude that while it allows for a faster decay of the initial error (bias) (which is unquantified), their steady state behavior (i.e. variance) is worse compared to that of LMS. Yuan et al. (2016) considered a constant step size accelerated scheme with no iterate averaging in the same oracle model as this paper, and conclude that these do not offer any improvement over standard SGD. More concretely, Yuan et al. (2016) show that the variance of their accelerated SGD method with a sufficiently small constant step size is the same as that of SGD with a significantly larger step size. Note that none of the these efforts (Proakis, 1974; Roy and Shynk, 1990; Sharma et al., 1998; Yuan et al., 2016) achieve minimax error rates or quantify (any improvement whatsoever on the) rate of bias decay.

With regards to notions of optimality, there are (at least) two lines of thought: one is a statistical objective where the goal is (on every problem instance) to match the rate of the statistically optimal estimator (Anbar, 1971; Fabian, 1973; Kushner and Clark, 1978; Polyak and Juditsky, 1992); another is on obtaining algorithms whose worst case upper bounds (under various assumptions such as bounded noise) match the lower bounds provided in Nemirovsky and Yudin (1983). The work of Polyak and Juditsky (1992) are in the former model, where they show that the distribution of the averaged SGD estimator matches, on every problem, that of the statistically optimal estimator, in the limit (under appropriate regularization conditions standard in the statistics literature, where the optimal estimator is often referred to as the maximum likelihood estimator/the empirical risk minimizer/an $M$-estimator (Lehmann and Casella, 1998; van der Vaart, 2000)). Along these lines, non-asymptotic rates towards statistically optimal estimators are given by Bach and Moulines (2013); Bach (2014); Défossez and Bach (2015); Dieuleveut and Bach (2015); Needell et al. (2016); Frostig et al. (2015b); Jain et al. (2016). This work can be seen as improving this non-asymptotic rate (to the statistically optimal estimation rate) using an accelerated method. As to the latter (i.e. matching the worst-case lower bounds in Nemirovsky and Yudin (1983)), there are a number of positive results on using accelerated stochastic optimization procedures; the works of Lan (2008); Hu et al. (2009); Ghadimi and Lan (2012, 2013); Dieuleveut et al. (2016) match the lower bounds provided in Nemirovsky and Yudin (1983). We compare these assumptions and works in more detail.

In stochastic first order oracle models (see Kushner and Clark (1978); Kushner and Yin (2003)), one typically has access to sampled gradients of the form:

As a final remark, there have been results (Shalev-Shwartz and Zhang, 2014; Frostig et al., 2015a; Lin et al., 2015; Lan and Zhou, 2015; Allen-Zhu, 2016) that provide accelerated rates for offline stochastic optimization which deal with minimizing sums of convex functions; these results are almost tight due to matching lower bounds (Lan and Zhou, 2015; Woodworth and Srebro, 2016). These results do not immediately translate into rates on the generalization error. Furthermore, these algorithms are not streaming, as they require making multiple passes over a dataset stored in memory. Refer to Frostig et al. (2015b) for more details.

Main Results

Let $\mathbf{H}$ denote the second moment matrix of the input, which is also the hessian $\nabla^{2}P(\mathbf{x})$ of (1):

Furthermore, let the fourth moment tensor $\mathcal{M}$ of the inputs $\mathbf{a}\sim\mathcal{D}$ is defined as:

Finite second and fourth moment: The second moment matrix $\mathbf{H}$ and the fourth moment tensor $\mathcal{M}$ exist and are finite.

Positive Definiteness: The second moment matrix $\mathbf{H}$ is strictly positive definite, i.e. $\mathbf{H}\succ 0$.

We assume $\mathbf{(\mathcal{A}1)}$ and $\mathbf{(\mathcal{A}2)}$. $\mathbf{(\mathcal{A}2)}$ implies that $P(\mathbf{x})$ is strongly convex and admits a unique minimizer $\mathbf{x}^{*}$. Denote the noise $\epsilon$ in a sample $(\mathbf{a},b)\sim\mathcal{D}$ as: $\epsilon\stackrel{{\scriptstyle\textrm{def}}}{{=}}b-\left\langle\mathbf{a},\mathbf{x}^{*}\right\rangle$. First order optimality conditions of $\mathbf{x}^{*}$ imply

Let $\mathbf{\Sigma}$ denote the covariance of gradient at optimum $\mathbf{x}^{*}$ (or noise covariance matrix),

We define the noise level $\sigma^{2}$, condition number ${\kappa}$, statistical condition number $\widetilde{\kappa}$ below. Noise level: The noise level is defined to be the smallest positive number $\sigma^{2}$ such that

$\mu>0$ by $\mathbf{(\mathcal{A}2)}$. Now, let $R^{2}$ be the smallest positive number such that

. The condition number ${\kappa}$ of the distribution $\mathcal{D}$ (Défossez and Bach, 2015; Jain et al., 2016) is

Statistical condition number: The statistical condition number $\widetilde{\kappa}$ is defined as the smallest positive number such that

$\widetilde{\kappa}$ governs how many samples $\mathbf{a}_{i}$ require to be drawn from $\mathcal{D}$ so that the empirical covariance is spectrally close to $\mathbf{H}$, i.e. for some constant $c>1$, $\mathbf{H}/c\preceq\frac{1}{n}\sum_{i=1}^{n}\mathbf{a}_{i}\mathbf{a}_{i}^{\top}\preceq c\mathbf{H}$. In comparison to the matrix Bernstein inequality where stronger (yet related) moment conditions are assumed in order to obtain high probability results, our results hold only in expectation (refer to Hsu et al. (2014) for this definition, wherein $\widetilde{\kappa}$ is referred to as bounded statistical leverage in theorem $1$ and remark $1$).

2 Algorithm and Main Theorem

Algorithm 1 presents the pseudo code of the proposed algorithm. ASGD can be viewed as a variant of Nesterov’s accelerated gradient method (Nesterov, 2012), working with a stochastic gradient oracle (equation 2) and with tail-averaging the final $n-t$ iterates. The main result now follows:

Suppose $\mathbf{(\mathcal{A}1)}$ and $\mathbf{(\mathcal{A}2)}$ hold. Set $\alpha=\frac{3\sqrt{5}\cdot\sqrt{{\kappa}\widetilde{\kappa}}}{1+3\sqrt{5}\cdot\sqrt{{\kappa}\widetilde{\kappa}}},\beta=\frac{1}{9\sqrt{{\kappa}\widetilde{\kappa}}},\gamma=\frac{1}{3\sqrt{5}\cdot\mu\sqrt{{\kappa}\widetilde{\kappa}}},\delta=\frac{1}{5R^{2}}$. After $n$ calls to the stochastic first order oracle (equation 2), ASGD outputs $\bar{\mathbf{x}}_{t,n}$ satisfying:

where $C$ is a universal constant, $\sigma^{2}$, ${\kappa}$ and $\widetilde{\kappa}$ are the noise level, condition number and statistical condition number respectively.

The following corollary holds if the iterates are tail-averaged over the last $n/2$ samples and $n>\mathcal{O}(\sqrt{{\kappa}\widetilde{\kappa}}\log(d{\kappa}\widetilde{\kappa}))$. The second condition lets us absorb lower order terms into leading order terms.

Assume the parameter settings of theorem 1 and let $t=\lfloor n/2\rfloor$ and $n>C^{\prime}\sqrt{{\kappa}\widetilde{\kappa}}\log(d{\kappa}\widetilde{\kappa})$ (for an appropriate universal constants $C,C^{\prime}$). We have that with $n$ calls to the stochastic first order oracle, ASGD outputs a vector $\bar{\mathbf{x}}_{t,n}$ satisfying:

A few remarks about the result of theorem 1 are due: (i) ASGD decays the initial error at a geometric rate of $\mathcal{O}(1/\sqrt{{\kappa}\widetilde{\kappa}})$ during the unaveraged phase of $t$ iterations, which presents the first improvement over the $\mathcal{O}\left(1/{\kappa}\right)$ rate offered by SGD (Robbins and Monro, 1951)/averaged SGD (Polyak and Juditsky, 1992; Jain et al., 2016) for the least squares stochastic approximation problem, (ii) the second term in the error bound indicates that ASGD obtains (up to constants) the minimax rate once $n>\mathcal{O}(\sqrt{{\kappa}\widetilde{\kappa}}\log(d{\kappa}\widetilde{\kappa}))$. Note that this implies that Theorem 1 provides a sharp non-asymptotic analysis (up to $\log$ factors) of the behavior of Algorithm 1.

3 Discussion and Open Problems

A challenging problem in this context is in formalizing a finite sample size lower bound in the oracle model considered in this work. Lower bounds in stochastic oracle models have been considered in the literature (see Nemirovsky and Yudin (1983); Raginsky and Rakhlin (2011); Agarwal et al. (2012)), though it is not evident these oracle models and lower bounds are sharp enough to imply statements in our setting (see section 1.4 for a discussion of these oracles).

where $\widehat{\mathbf{x}}_{n}^{\textrm{ERM}}$ is the ERM over samples $\mathcal{S}_{n}=\{\mathbf{a}_{i},b_{i}\}_{i=1}^{n}$. Averaged SGD (Jain et al., 2016) and streaming SVRG (Frostig et al., 2015b) are known to achieve these rates for the heteroscedastic case. Refer to Frostig et al. (2015b) for more details.Neglecting constants, Theorem 1 is guaranteed to achieve the rate of the ERM for the homoscedastic case (where $\mathbf{\Sigma}=\sigma^{2}\mathbf{H}$) and is tight when the bound $\mathbf{\Sigma}\preceq\sigma^{2}\mathbf{H}$ is nearly tight (upto constants). We conjecture ASGD achieves the rate of the ERM in the heteroscedastic case by appealing to a more refined analysis as is the case for averaged SGD (see Jain et al. (2016)). It is also an open question to understand acceleration for smooth stochastic approximation (beyond least squares), in situations where the rate represented by equation 7 holds (Polyak and Juditsky, 1992).

Proof Outline

We now present a brief outline of the proof of Theorem 1. Recall the variables in Algorithm 1. Before presenting the proof outline we require some definitions. We begin by defining the centered estimate $\bm{\theta}_{j}$ as:

Recall that the stepsizes in Algorithm 1 are $\alpha=\frac{3\sqrt{5}\cdot\sqrt{{\kappa}\widetilde{\kappa}}}{1+3\sqrt{5}\cdot\sqrt{{\kappa}\widetilde{\kappa}}},\beta=\frac{1}{9\sqrt{{\kappa}\widetilde{\kappa}}},\gamma=\frac{1}{3\sqrt{5}\cdot\mu\sqrt{{\kappa}\widetilde{\kappa}}},\delta=\frac{1}{5R^{2}}$. The accelerated SGD updates of Algorithm 1 can be written in terms of $\bm{\theta}_{j}$ as:

Bias-variance decomposition: The proof of theorem 1 employs the bias-variance decomposition, which is well known in the context of stochastic approximation (see Bach and Moulines (2011); Frostig et al. (2015b); Jain et al. (2016)) and is re-derived in the appendix. The bias-variance decomposition allows for the generalization error to be upper-bounded by analyzing two sub-problems: (a) bias, analyzing the algorithm’s behavior on the noiseless problem (i.e. $\bm{\zeta}_{j}=0\ \forall\ j$ a.s.) while starting at $\bm{\theta}_{0}^{\textrm{bias}}=\bm{\theta}_{0}$ and (b) variance, analyzing the algorithm’s behavior by starting at the solution (i.e. $\bm{\theta}_{0}^{\textrm{variance}}=0$) and allowing the noise $\bm{\zeta}_{\operatorname*{\scalerel*{\cdot}{\bigodot}}}$ to drive the process. In a similar manner as ${\bar{\bm{\theta}}}_{t,n}$, the bias and variance sub-problems are associated with ${\bar{\bm{\theta}}}_{t,n}^{\textrm{bias}}$ and ${\bar{\bm{\theta}}}_{t,n}^{\textrm{variance}}$, and these are related as:

Since we deal with the square loss, the generalization error of the output $\bar{\mathbf{x}}_{t,n}$ of algorithm 1 is:

indicating that the generalization error can be bounded by analyzing the bias and variance sub-problem. We now present the lemmas that bound the bias error.

The quantity that needs to be bounded in the term above is $\mathcal{B}^{t+1}\bm{\theta}_{0}\otimes\bm{\theta}_{0}$. Lemma 4 presents a result that can be applied recursively to bound $\mathcal{B}^{t+1}\bm{\theta}_{0}\otimes\bm{\theta}_{0}$ ($=\mathcal{B}^{t+1}\bm{\theta}_{0}^{\textrm{bias}}\otimes\bm{\theta}_{0}^{\textrm{bias}}$ since $\bm{\theta}_{0}^{\textrm{bias}}=\bm{\theta}_{0}$).

(i) the matrices $\widetilde{\mathbf{G}}$ and $\widetilde{\mathbf{G}}^{\top}$ appearing in $\mathbf{G}$ are due to the fact that we prove contraction using the variables $\mathbf{x}-\mathbf{x}^{*}$ and $\mathbf{v}-\mathbf{x}^{*}$ instead of $\mathbf{x}-\mathbf{x}^{*}$ and $\mathbf{y}-\mathbf{x}^{*}$, as used in defining $\bm{\theta}$. (ii) The key novelty in lemma 4 is that while standard analyses of accelerated gradient descent (in the exact first order oracle) use the potential function $\left\|\mathbf{x}-\mathbf{x}^{*}\right\|_{\mathbf{H}}^{2}+\mu\left\|\mathbf{v}-\mathbf{x}^{*}\right\|_{2}^{2}$ (e.g. Wilson et al. (2016)), we consider it crucial for employing the potential function $\left\|\mathbf{x}-\mathbf{x}^{*}\right\|_{2}^{2}+\mu\left\|\mathbf{v}-\mathbf{x}^{*}\right\|_{{\mathbf{H}}^{-1}}^{2}$ (this potential function is captured using the matrix $\mathbf{Z}$) to prove accelerated rates (of $\mathcal{O}\left(1/\sqrt{{\kappa}\widetilde{\kappa}}\right)$) for bias decay.

We now present the lemmas associated with bounding the variance error:

The covariance of the stationary distribution of the iterates i.e., $\lim_{j\to\infty}\bm{\theta}_{j}^{\textrm{variance}}$ requires a precise bound to obtain statistically optimal error rates. Lemma 6 presents a bound on this quantity.

The covariance of limiting distribution of $\bm{\theta}^{\textrm{variance}}$ satisfies:

A crucial implication of this lemma is that the limiting final iterate $\bm{\theta}_{\infty}^{\textrm{variance}}$ has an excess risk $\mathcal{O}(\sigma^{2})$. This result naturally lends itself to the (tail-)averaged iterate achieving the minimax optimal rate of $\mathcal{O}(d\sigma^{2}/n)$. Refer to the appendix E and lemma 17 for more details in this regard.

Conclusion

This paper introduces an accelerated stochastic gradient method, which presents the first improvement in achieving minimax rates faster than averaged SGD (Robbins and Monro, 1951; Ruppert, 1988; Polyak and Juditsky, 1992; Jain et al., 2016)/Streaming SVRG (Frostig et al., 2015b) for the stochastic approximation problem of least squares regression. To obtain this result, the paper presented the need to rethink what acceleration has to offer when working with a stochastic gradient oracle: these thought experiments indicated a need to consider a quantity that captured more fine grained problem characteristics. The statistical condition number (an affine invariant distributional quantity) is shown to characterize the improvements that acceleration offers in the stochastic first order oracle model.

In essence, this paper presents the first known provable analysis of the claim that fast gradient methods are stable when dealing with statistical errors, in contrast to negative results in statistical and non-statistical settings (Paige, 1971; Proakis, 1974; Polyak, 1987; Greenbaum, 1989; Roy and Shynk, 1990; Sharma et al., 1998; d’Aspremont, 2008; Devolder et al., 2013, 2014; Yuan et al., 2016). We hope that this paper provides insights towards developing simple and effective accelerated stochastic gradient schemes for general convex and non-convex optimization problems.

Sham Kakade acknowledges funding from Washington Research Foundation Fund for Innovation in Data-Intensive Discovery and the NSF through awards CCF-$1637360$, CCF-$1703574$ and CCF-$1740551$.

References

Appendix A Appendix setup

We will first provide a note on the organization of the appendix and follow that up with introducing the notations.

In subsection A.2, we will recall notation from the main paper and introduce some new notation that will be used across the appendix.

In section B, we will write out expressions that characterize the generalization error of the proposed accelerated SGD method. In order to bound the generalization error, we require developing an understanding of two terms namely the bias error and the variance error.

In section C, we prove lemmas that will be used in subsequent sections to prove bounds on the bias and variance error.

In section D, we will bound the bias error of the proposed accelerated stochastic gradient method. In particular, lemma 4 is the key lemma that provides a new potential function with which this paper achieves acceleration. Further, lemma 16 is the lemma that bounds all the terms of the bias error.

In section E, we will bound the variance error of the proposed accelerated stochastic gradient method. In particular, lemma 6 is the key lemma that considers a stochastic process view of the proposed accelerated stochastic gradient method and provides a sharp bound on the covariance of the stationary distribution of the iterates. Furthermore, lemma 20 bounds all terms of the variance error.

Section F presents the proof of Theorem 1. In particular, this section aggregates the result of lemma 16 (which bounds all terms of the bias error) and lemma 20 (which bounds all terms of the variance error) to present the guarantees of Algorithm 1.

A.2 Notations

We begin by introducing $\mathcal{M}$, which is the fourth moment tensor of the input $\mathbf{a}\sim\mathcal{D}$, i.e.:

With this definition in place, we recall $R^{2}$ as the smallest number, such that $\mathcal{M}$ applied to the identity matrix $\mathbf{I}$ satisfies:

Moreover, we recall that the condition number of the distribution ${\kappa}=R^{2}/\mu$, where $\mu$ is the smallest eigenvalue of $\mathbf{H}$. Furthermore, the definition of the statistical condition number $\widetilde{\kappa}$ of the distribution follows by applying the fourth moment tensor $\mathcal{M}$ to $\mathbf{H}^{-1}$, i.e.:

Parameter choices: In all of appendix we choose the parameters in Algorithm 1 as

where $c_{1}$ is an arbitrary constant satisfying $0<c_{1}<\frac{1}{2}$. Furthermore, we note that $c_{3}=\frac{c_{2}\sqrt{2c_{1}-c_{1}^{2}}}{c_{1}}$, $c_{2}^{2}=\frac{c_{4}}{2-c_{1}}$ and $c_{4}<1/6$. Note that we recover Theorem 1 by choosing $c_{1}=1/5,c_{2}=\sqrt{5}/9,c_{3}=\sqrt{5}/3,c_{4}=1/9$. We denote

Furthermore, we denote by $\bm{\Phi}_{k}$ the expected covariance of $\bm{\theta}_{k}$, i.e.:

Next, let $\mathcal{F}_{k}$ denote the filtration generated by samples $\{(\mathbf{a}_{1},b_{1}),\cdots,(\mathbf{a}_{k},b_{k})\}$. Then,

Without loss of generality, we assume that $\mathbf{H}$ is a diagonal matrix. We now note that we can rearrange the coordinates through an eigenvalue decomposition so that $\mathbf{A}$ becomes a block-diagonal matrix with $2\times 2$ blocks. We denote the $j^{\textrm{th}}$ block by $\mathbf{A}_{j}$:

where $\lambda_{j}$ denotes the $j^{\textrm{th}}$ eigenvalue of $\mathbf{H}$. Next,

Thus implying the following relation between the operators $\mathcal{B},\mathcal{D}$ and $\mathcal{R}$:

Appendix B The Tail-Average Iterate: Covariance and bias-variance decomposition

We begin by considering the first-order Markovian recursion as defined by equation A.2:

We refer by $\bm{\Phi}_{j}$ the covariance of the $j^{\text{th}}$ iterate, i.e.:

Consider a decomposition of $\bm{\theta}_{j}$ as $\bm{\theta}_{j}=\bm{\theta}_{j}^{\textrm{bias}}+\bm{\theta}_{j}^{\textrm{variance}}$, where $\bm{\theta}_{j}^{\textrm{bias}}$ and $\bm{\theta}_{j}^{\textrm{variance}}$ are defined as follows:

Before we prove the decomposition holds using an inductive argument, let us understand what the bias and variance sub-problem intuitively mean.

Note that the bias sub-problem (defined by equation 13) refers to running algorithm on the noiseless problem (i.e., where, $\bm{\zeta}_{\operatorname*{\scalerel*{\cdot}{\bigodot}}}=0$ a.s.) by starting it at $\bm{\theta}_{0}^{\textrm{bias}}=\bm{\theta}_{0}$. The bias essentially measures the dependence of the generalization error on the excess risk of the initial point $\bm{\theta}_{0}$ and bears similarities to convergence rates studied in the context of offline optimization.

The variance sub-problem (defined by equation 14) measures the dependence of the generalization error on the noise introduced during the course of optimization, and this is associated with the statistical aspects of the optimization problem. The variance can be understood as starting the algorithm at the solution ($\bm{\theta}_{0}^{\text{variance}}=0$) and running the optimization driven solely by noise. Note that the variance is associated with sharp statistical lower bounds which dictate its rate of decay as a function of the number of oracle calls $n$.

Now, we will prove that the decomposition $\bm{\theta}_{j}=\bm{\theta}_{j}^{\textrm{bias}}+\bm{\theta}_{j}^{\textrm{variance}}$ captures the recursion expressed in equation A.2 through induction. For the base case $j=1$, we see that

Now, for the inductive step, let us assume that the decomposition holds in the $j-1^{st}$ iteration, i.e. we assume $\bm{\theta}_{j-1}=\bm{\theta}_{j-1}^{\textrm{bias}}+\bm{\theta}_{j-1}^{\textrm{variance}}$. We will then prove that this relation holds in the $j^{th}$ iteration. Towards this, we will write the recursion:

This proves the decomposition holds through a straight forward inductive argument.

In a similar manner as $\bm{\theta}_{j}$, the tail-averaged iterate ${\bar{\bm{\theta}}}_{t,n}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\frac{1}{n-t}\sum_{j=t+1}^{n}\bm{\theta}_{j}$ can also be written as ${\bar{\bm{\theta}}}_{t,n}={\bar{\bm{\theta}}}_{t,n}^{\textrm{bias}}+{\bar{\bm{\theta}}}_{t,n}^{\textrm{variance}}$, where ${\bar{\bm{\theta}}}_{t,n}^{\textrm{bias}}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\frac{1}{n-t}\sum_{j=t+1}^{n}\bm{\theta}_{j}^{\textrm{bias}}$ and ${\bar{\bm{\theta}}}_{t,n}^{\textrm{variance}}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\frac{1}{n-t}\sum_{j=t+1}^{n}\bm{\theta}_{j}^{\textrm{variance}}$. Furthermore, the tail-averaged iterate ${\bar{\bm{\theta}}}_{t,n}$ and its bias and variance counterparts ${\bar{\bm{\theta}}}_{t,n}^{\text{bias}},{\bar{\bm{\theta}}}_{t,n}^{\text{variance}}$ are associated with their corresponding covariance matrices $\bm{\bar{\Phi}}_{t,n},\bm{\bar{\Phi}}_{t,n}^{\text{bias}},\bm{\bar{\Phi}}_{t,n}^{\text{variance}}$ respectively. Note that $\bm{\bar{\Phi}}_{t,n}$ can be upper bounded using Cauchy-Shwartz inequality as:

The above inequality is referred to as the bias-variance decomposition and is well known from previous work Bach and Moulines (2013); Frostig et al. (2015b); Jain et al. (2016), and we re-derive this decomposition for the sake of completeness. We will now derive an expression for the covariance of the tail-averaged iterate and apply it to obtain the covariance of the bias ($\bm{\bar{\Phi}}_{t,n}^{\text{bias}}$) and variance ($\bm{\bar{\Phi}}_{t,n}^{\text{variance}}$) error of the tail-averaged iterate.

We begin by writing out an expression for the tail-averaged iterate ${\bar{\bm{\theta}}}_{t,n}$ as:

To get the excess risk of the tail-averaged iterate ${\bar{\bm{\theta}}}_{t,n}$, we track its covariance $\bm{\bar{\Phi}}_{t,n}$:

Note that the above recursion can be applied to obtain the covariance of the tail-averaged iterate for the bias ($\bm{\bar{\Phi}}_{t,n}^{\text{bias}}$) and variance ($\bm{\bar{\Phi}}_{t,n}^{\text{variance}}$) error, since the conditional expectation arguments employed in obtaining equation B.1 are satisfied by both the recursion used in tracking the bias error (i.e. equation 13) and the variance error (i.e. equation 14). This implies that,

B.2 Covariance of Bias error of the tail-averaged iterate

To obtain the covariance of the bias error of the tail-averaged iterate, we first need to obtain $\bm{\Phi}_{j}^{\text{bias}}$, which we will by unrolling the recursion of equation 13:

Next, we recount the equation for the covariance of the bias of the tail-averaged iterate from equation 19:

Now, we substitute $\bm{\Phi}_{j}^{\text{bias}}$ from equation B.2:

There are two points to note here: (a) The second line consists of terms that constitute the lower-order terms of the bias. We will bound the summation by taking a supremum over $j$. (b) Note that the burn-in phase consisting of $t$ unaveraged iterations allows for a geometric decay of the bias, followed by the tail-averaged phase that allows for a sublinear rate of bias decay. ∎

B.3 Covariance of Variance error of the tail-averaged iterate

Next, in order to obtain the covariance of the variance of the tail-averaged iterate, we first need to obtain $\bm{\Phi}_{j}^{\text{variance}}$, and we will obtain this by unrolling the recursion of equation 14:

We will substitute the expression for $\bm{\Phi}_{j}^{\text{variance}}$ from equation B.3.

Equations B, B.2, B.3 wrap up the proof of lemmas 3, 5.

The parameter error of the (tail-)averaged iterate can be obtained using a trace operator $\left\langle{\operatorname*{\scalerel*{\cdot}{\bigodot}}},{\operatorname*{\scalerel*{\cdot}{\bigodot}}}\right\rangle$ to the tail-averaged iterate’s covariance $\bm{\bar{\Phi}}_{t,n}$ with the matrix $\begin{bmatrix}\mathbf{I}&0\\ 0&0\end{bmatrix}$, i.e.

In order to obtain the function error, we note the following taylor expansion of the function $P(\operatorname*{\scalerel*{\cdot}{\bigodot}})$ around the minimizer $\mathbf{x}^{*}$:

This implies the excess risk can be obtained as:

Appendix C Useful lemmas

In this section, we will state and prove some useful lemmas that will be helpful in the later sections.

Since we assumed that $\mathbf{H}$ is a diagonal matrix (with out loss of generality), we note that $\mathbf{A}$ is a block diagonal matrix after a rearrangement of the co-ordinates (via an eigenvalue decomposition).

In particular, by considering the $j^{\text{th}}$ block (denoted by $\mathbf{A}_{j}$ corresponding to the $j^{\textrm{th}}$ eigenvalue $\lambda_{j}$ of $\mathbf{H}$), we have:

Implying that the determinant $\left|\mathbf{I}-\mathbf{A}_{j}^{\top}\right|=(q-c\delta)\lambda_{j}$, using which:

Accumulating the results of each of the blocks and by rearranging the co-ordinates, the result follows. ∎

In a similar manner as in lemma 7, we decompose the computation into each of the eigen-directions and subsequently re-arrange the results. In particular, we note:

Multiplying the above with the result of lemma 7, we have:

From which the statement of the lemma follows through a simple re-arrangement.

In a similar argument as in previous two lemmas, we analyze the expression in each eigendirection of $\mathbf{H}$ through a rearrangement of the co-ordinates. Utilizing the expression of $\mathbf{I}-\mathbf{A}_{j}^{\top}$ from equation 25, we get:

Rearranging the co-ordinates, the statement of the lemma follows. ∎

The matrix $\mathbf{A}$ satisfies the following properties:

Eigenvalues $q$ of $\mathbf{A}$ satisfy $\left|{q}\right|\leq\sqrt{\alpha}$, and

$\left\|\mathbf{A}^{k}\right\|_{2}\leq 3\sqrt{2}\cdot k\cdot\alpha^{\frac{k-1}{2}}\;\forall\;k\geq 1$.

Since the matrix is block-diagonal with $2\times 2$ blocks, after a rearranging the coordinates, we will restrict ourselves to bounding the eigenvalues and eigenvectors of each of these $2\times 2$ blocks. Combining the results for different blocks then proves the lemma. Recall that $\mathbf{A}_{j}=\begin{bmatrix}0&1-\delta\lambda_{j}\\ -c&1+c-q\lambda_{j}\end{bmatrix}$.

Part I: Let us first prove the statement about the eigenvalues of $\mathbf{A}$. There are two scenarios here:

Complex eigenvalues: In this case, both eigenvalues of $\mathbf{A}_{j}$ have the same magnitude which is given by $\sqrt{\det(\mathbf{A}_{j})}=\sqrt{c(1-\delta\lambda_{j})}\leq\sqrt{c}\leq\sqrt{\alpha}$.

Real eigenvalues: Let $q_{1}$ and $q_{2}$ be the two real eigenvalues of $\mathbf{A}_{j}$. We know that $q_{1}+q_{2}=\textrm{Tr}\left(\mathbf{A}_{j}\right)=1+c-q\lambda_{j}>0$ and $q_{1}\cdot q_{2}=\det(\mathbf{A}_{j})>0$. This means that $q_{1}>0$ and $q_{2}>0$. Now, consider the matrix $\mathbf{G}_{j}\stackrel{{\scriptstyle\textrm{def}}}{{=}}(1-\beta)\mathbf{I}-\mathbf{A}_{j}=\begin{bmatrix}(1-\beta)&-1+\delta\lambda_{j}\\ c&-1+{(1-\beta)(1-\alpha)}+q\lambda_{j}\end{bmatrix}$. We see that $((1-\beta)-q_{1})((1-\beta)-q_{2})=\det(\mathbf{G}_{j})=(1-\beta)(1-\alpha)\left((1-\beta)-1\right)+(1-\beta)\left(q-\alpha\delta\right)\lambda_{j}=(1-\beta)\left(1-\alpha\right)\left(\gamma\lambda_{j}-\beta\right)\geq 0$. This means that there are two possibilities: either $q_{1},q_{2}\geq(1-\beta)$ or $q_{1},q_{2}\leq(1-\beta)$. If the second condition is true, then we are done. If not, if $q_{1},q_{2}\geq(1-\beta)$, then $\max_{i}q_{i}=\frac{\det(\mathbf{A}_{j})}{\min_{i}q_{i}}\leq\frac{c(1-\delta\lambda_{j})}{(1-\beta)}\leq\alpha(1-\delta\lambda_{j})$. Since $\sqrt{\alpha}\geq\alpha\geq 1-\beta$, this proves the first part of the lemma.

Part II: Let $\mathbf{A}_{j}=\mathbf{V}\mathbf{Q}\mathbf{V}^{\top}$ be the Schur decomposition of $\mathbf{A}_{j}$ where $\mathbf{Q}=\begin{bmatrix}q_{1}&q\\ 0&q_{2}\end{bmatrix}$ is an upper triangular matrix with eigenvalues $q_{1}$ and $q_{2}$ of $\mathbf{A}_{j}$ on the diagonal and $\mathbf{V}$ is a unitary matrix i.e., $\mathbf{V}\mathbf{V}^{\top}=\mathbf{V}^{\top}\mathbf{V}=\mathbf{I}$. We first observe that $\left|{q}\right|\leq\left\|\mathbf{Q}\right\|_{2}\stackrel{{\scriptstyle(\zeta_{1})}}{{=}}\left\|\mathbf{A}_{j}\right\|_{2}\leq\left\|\mathbf{A}_{j}\right\|_{F}\leq\sqrt{6}$, where $(\zeta_{1})$ follows from the fact that $\mathbf{V}$ is a unitary matrix. $\mathbf{V}$ being unitary also implies that $\mathbf{A}_{j}^{k}=\mathbf{V}\mathbf{Q}^{k}\mathbf{V}^{\top}$. On the other hand, a simple proof via induction tells us that

So, we have $\left\|\mathbf{A}_{j}^{k}\right\|_{2}=\left\|\mathbf{Q}^{k}\right\|_{2}\leq\left\|\mathbf{Q}^{k}\right\|_{F}\leq\sqrt{3}k\left|{q}\right|\max\left(\left|{q_{1}}\right|^{k-1},\left|{q_{2}}\right|^{k-1}\right)\leq 3\sqrt{2}\cdot k\cdot\alpha^{\frac{k-1}{2}}$, where we used $\left|{q}\right|\leq\sqrt{6}$ and $\max\left(\left|{q_{1}}\right|,\left|{q_{2}}\right|\right)\leq\sqrt{\alpha}$. ∎

Finally, we state and prove the following lemma which is a relation between left and right multiplication operators.

Let $\mathbf{A}$ be any matrix with $\mathcal{A}_{\mathcal{L}}=\mathbf{A}\otimes\mathbf{I}$ and $\mathcal{A}_{\mathcal{R}}=\mathbf{I}\otimes\mathbf{A}$ representing its left and right multiplication operators. Then, the following expression holds:

Let us assume that $\mathbf{A}$ can be written in terms of its eigen decomposition as $\mathbf{A}=\mathbf{V}\Lambda\mathbf{V}^{-1}$. Then the first claim is that $\mathcal{I},\mathcal{A}_{\mathcal{L}},\mathcal{A}_{\mathcal{R}}$ are diagonalized by the same basis consisting of the eigenvectors of $\mathbf{A}$, i.e. in particular, the matrix of eigenvectors of $\mathcal{I},\mathcal{A}_{\mathcal{L}},\mathcal{A}_{\mathcal{R}}$ can be written as $\mathbf{V}\otimes\mathbf{V}$. In particular, this implies, $\forall\ i,j\in\{1,2,...,d\}\times\{1,2,...,d\}$, we have, applying $\mathbf{v}_{i}\otimes\mathbf{v}_{j}$ to the LHS, we have:

Applying $\mathbf{v}_{i}\otimes\mathbf{v}_{j}$ to the RHS, we have:

The next claim is that for any scalars (real/complex) $x,y~{}\neq 1$, the following statement holds implying the statement of the lemma:

Recall the matrix $\mathbf{G}$ defined as $\mathbf{G}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\begin{bmatrix}\mathbf{I}&\frac{-\alpha}{1-\alpha}\mathbf{I}\\ 0&\frac{1}{1-\alpha}\mathbf{I}\end{bmatrix}\begin{bmatrix}\mathbf{I}&0\\ 0&{\mu}{\mathbf{H}}^{-1}\end{bmatrix}\begin{bmatrix}\mathbf{I}&0\\ \frac{-\alpha}{1-\alpha}\mathbf{I}&\frac{1}{1-\alpha}\mathbf{I}\end{bmatrix}$. The condition number of $\mathbf{G}$, $\kappa(\mathbf{G})$ satisfies $\kappa(\mathbf{G})\leq\frac{4{\kappa}}{\sqrt{1-\alpha^{2}}}$.

Since the above matrix is block-diagonal after a rearrangement of coordinates, it suffices to compute the smallest and largest singular values of each block. Let $\lambda_{i}$ be the $i^{\textrm{th}}$ eigenvalue of $\mathbf{H}$. Let $\mathbf{C}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\begin{bmatrix}1&0\\ \frac{-\alpha}{1-\alpha}&\frac{1}{1-\alpha}\end{bmatrix}$ and consider the matrix $\mathbf{G}_{i}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\mathbf{C}\begin{bmatrix}1&0\\ 0&\frac{\mu}{\lambda_{i}}\end{bmatrix}\mathbf{C}^{\top}$. The largest eigenvalue of $\mathbf{G}_{i}$ is at most $\sigma_{\textrm{max}}\left(\mathbf{C}\right)^{2}$, while the smallest eigenvalue, $\sigma_{\textrm{min}}\left(\mathbf{G}_{i}\right)$ is at least $\frac{\mu}{\lambda_{i}}\cdot\sigma_{\textrm{min}}\left(\mathbf{C}\right)^{2}$. We obtain the following bounds on $\sigma_{\textrm{min}}\left(\mathbf{C}\right)$ and $\sigma_{\textrm{max}}\left(\mathbf{C}\right)$.

where we used the computation that $\det\left(\mathbf{C}\mathbf{C}^{\top}\right)=\frac{1}{1-\alpha}$. This means that $\sigma_{\textrm{min}}\left(\mathbf{G}_{i}\right)\geq\frac{\mu}{2\lambda_{i}}$ and $\sigma_{\textrm{max}}\left(\mathbf{G}_{i}\right)\leq\frac{2}{\sqrt{1-\alpha^{2}}}$. Combining all the blocks, we see that the condition number of $\mathbf{G}$ is at most $\frac{4{\kappa}}{\sqrt{1-\alpha^{2}}}$, proving the lemma. ∎

Appendix D Lemmas and proofs for bias contraction

Let $\mathbf{v}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\frac{1}{1-\alpha}\left(\mathbf{y}-\alpha\mathbf{x}\right)$ and consider the following update rules corresponding to the noiseless versions of the updates in Algorithm 1:

where $\widehat{\mathbf{H}}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\mathbf{a}\mathbf{a}^{\top}$ where $\mathbf{a}$ is sampled from the marginal on $(\mathbf{a},b)\sim\mathcal{D}$. We first note that

To establish this result, let us define two quantities: $e\stackrel{{\scriptstyle\textrm{def}}}{{=}}{\left\|\mathbf{x}-\mathbf{x}^{*}\right\|_{2}^{2}}$, $f\stackrel{{\scriptstyle\textrm{def}}}{{=}}{\left\|\mathbf{v}-\mathbf{x}^{*}\right\|^{2}_{{\mathbf{H}}^{-1}}}$ and similarly, $e^{+}\stackrel{{\scriptstyle\textrm{def}}}{{=}}{\left\|\mathbf{x}^{+}-\mathbf{x}^{*}\right\|_{2}^{2}}$ and $f^{+}\stackrel{{\scriptstyle\textrm{def}}}{{=}}{\left\|\mathbf{v}^{+}-\mathbf{x}^{*}\right\|^{2}_{{\mathbf{H}}^{-1}}}$. The potential function we consider is $e+\mu\cdot f$. Recall that the parameters are chosen as:

with $c_{1}<1/2$, $c_{3}=\frac{c_{2}\sqrt{2c_{1}-c_{1}^{2}}}{c_{1}}$, $c_{2}^{2}=\frac{c_{4}}{2-c_{1}}$. Consider $e^{+}$ and employ the simple gradient descent bound:

$\mathbf{v}=\frac{1}{1-\alpha}\cdot\mathbf{y}-\frac{\alpha}{1-\alpha}\cdot\mathbf{x}$

$\mathbf{z}=\beta\mathbf{y}+(1-\beta)\mathbf{v}=\mathbf{y}+(1-\beta)(\mathbf{v}-\mathbf{y})$. Then substituting $\mathbf{v}$ in terms of $\mathbf{x}$ and $\mathbf{y}$ as in the equation above, we get: $\mathbf{z}=\mathbf{y}+\left(\frac{\alpha\cdot(1-\beta)}{1-\alpha}\right)(\mathbf{y}-\mathbf{x})$

Substituting equations D, D into equation D, we get:

Rewriting the guarantee on $e^{+}$ as in equation D:

By considering $e^{+}+\mu\cdot f^{+}$, we see the following:

Set $\frac{\gamma\mu\alpha}{1-\alpha}=1$ implying $\alpha=\frac{1}{1+\gamma\mu}=\frac{\sqrt{{\kappa}\widetilde{\kappa}}}{c_{2}\sqrt{2c_{1}-c_{1}^{2}}+\sqrt{{\kappa}\widetilde{\kappa}}}$

With these in place, we have the final result:

In particular, setting $\beta=c_{3}\gamma\mu=c_{3}\frac{c_{2}\sqrt{2c_{1}-c_{1}^{2}}}{\sqrt{{\kappa}\widetilde{\kappa}}}$, we have a per-step contraction of $1-\beta$ which is precisely $1-c_{3}\frac{c_{2}\sqrt{2c_{1}-c_{1}^{2}}}{\sqrt{{\kappa}\widetilde{\kappa}}}$, from which the claimed result naturally follows by substituting the values of $c_{1},c_{2},c_{3}$. ∎

For any psd matrix $\mathbf{Q}\succeq 0$, we have:

From Lemma 4, we conclude that \left\langle\mathbf{G},\mathcal{B}^{k}\mathbf{Q}\right\rangle\leq\bigg{(}1-\left(\frac{c_{2}c_{3}\sqrt{2c_{1}-c_{1}^{2}}}{\sqrt{{\kappa}\widetilde{\kappa}}}\right)\bigg{)}^{k}\left\langle\mathbf{G},\mathbf{Q}\right\rangle. This implies that \left\|\mathcal{B}^{k}\mathbf{Q}\right\|_{2}\leq\bigg{(}1-\left(\frac{c_{2}c_{3}\sqrt{2c_{1}-c_{1}^{2}}}{\sqrt{{\kappa}\widetilde{\kappa}}}\right)\bigg{)}^{k}\left\|\mathbf{Q}\right\|_{2}\kappa(\mathbf{G}). Plugging the bound on $\kappa(\mathbf{G})$ from Lemma 12 proves the lemma. ∎

The proof follows from Lemma 4. Since $\mathcal{B}=\mathcal{D}+\mathcal{R}$, we have $\left(\mathcal{I}-\mathcal{D}\right){\left(\mathcal{I}-\mathcal{B}\right)}^{-1}=\mathcal{I}+\mathcal{R}{\left(\mathcal{I}-\mathcal{B}\right)}^{-1}$. Since $\mathcal{R},\mathcal{B}$ and ${(\mathcal{I}-\mathcal{B})}^{-1}$ are all PSD operators, we have

Applying Lemma 13 with $\mathbf{Q}=\bm{\theta}_{0}\bm{\theta}_{0}^{\top}$ tells us that $\mathbf{S}_{1}\preceq\frac{4{\kappa}}{\sqrt{1-\alpha^{2}}}\exp\left(-tc_{2}c_{3}\sqrt{2c_{1}-c_{1}^{2}}/\sqrt{{\kappa}\widetilde{\kappa}}\right)\left\|\bm{\theta}_{0}\right\|_{2}^{2}\mathbf{I}$. For $\mathbf{S}_{2}$, we have

Combining the bounds on $\mathbf{S}_{1}$ and $\mathbf{S}_{2}$, we obtain

Plugging the bound for $\kappa(\mathbf{G})$ from Lemma 12 finishes the proof. ∎

For any psd matrix $\mathbf{Q}\succeq 0$, we have:

This corollary follows directly from Lemmas 10 and 13 and using the fact that $1-x\leq e^{-x}$ ∎

The following lemma bounds the total error of ${\bar{\bm{\theta}}}_{t,n}^{\textrm{bias}}$.

We now use lemmas in this section to bound inner product of the two terms in the above expression with $\begin{bmatrix}\mathbf{H}&0\\ 0&0\end{bmatrix}$, i.e. we seek to bound,

For the first term of equation 34, we have

Combining the above and noting the fact that $2\sqrt{{\kappa}\widetilde{\kappa}}R^{2}(q-c\delta)^{2}d<\frac{7d}{\mu}$, we have

This implies, equation 35 can be bounded as:

Consider now a term in the summation in the second term of (34).

Where, $C$ is a universal constant. Plugging (36) and (37) into (34), we obtain

Appendix E Lemmas and proofs for Bounding variance error

Before we prove lemma 6, we recall old notation and introduce new notations that will be employed in these proofs.

We begin with by recalling that we track $\bm{\theta}_{k}=\begin{bmatrix}\mathbf{x}_{k}-\mathbf{x}^{*}\\ \mathbf{y}_{k}-\mathbf{x}^{*}\end{bmatrix}$. Given $\bm{\theta}_{k}$, we recall the recursion governing the evolution of $\bm{\theta}_{k}$:

where, recall, $c=\alpha(1-\beta),\ q=\alpha\delta+(1-\alpha)\gamma$, and $\widehat{\mathbf{H}}_{k+1}=\mathbf{a}_{k+1}\mathbf{a}_{k+1}^{\top}$. Furthermore, we recall the following definitions, which will be heavily used in the following proofs:

And the operators $\mathcal{B},\mathcal{D},\mathcal{R}$ being related by:

Furthermore, in order to compute the steady state distribution with the fourth moment quantities in the mix, we need to rely on the following re-parameterization of the update matrix $\widehat{\mathbf{A}}$:

Considering the expectation of $\widehat{\mathbf{A}}\otimes\widehat{\mathbf{A}}$ with respect to a single draw from the distribution $\mathcal{D}$, we have:

Finally, we let nr and dr to denote the numerator and denominator respectively.

E.2 An exact expression for the stationary distribution

Note that a key term appearing in the expression for covariance of the variance equation (B.3) is ${\left(\mathcal{I}-\mathcal{B}\right)}^{-1}\widehat{\bm{\Sigma}}$. This is in fact nothing but the covariance of the error when we run accelerated SGD forever starting at $\mathbf{x}^{*}$ (i.e., at steady state). This can be seen by considering the base variance recursion using equation (E.1):

This recursion on the covariance operator $\bm{\Phi}_{k}$ can be unrolled until the start i.e. $k=0$ to yield:

E.3 Computing the steady state distribution

We now proceed to compute the stationary distribution. Recall that

Where the expectation is over a single sample drawn from the distribution $\mathcal{D}$. This implies in particular,

Since $\widehat{\bm{\Sigma}}\preceq\sigma^{2}\cdot\begin{bmatrix}\delta^{2}&\delta\cdot q\\ \delta\cdot q&q^{2}\end{bmatrix}\otimes\mathbf{H}$, and ${\left(\mathcal{I}-\mathcal{B}\right)}^{-1}$ is a PSD operator, the steady state distribution $\bm{\Phi}_{\infty}$ is bounded by:

Note that the Taylor expansion above is guaranteed to be correct if the right hand side is finite. We will understand bounds on the steady state distribution by splitting the analysis into the following parts:

Obtain \mathbf{U}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\big{(}\mathcal{I}-\mathbf{V}_{1}\otimes\mathbf{V}_{1}-\mathbf{V}_{1}\otimes\mathbf{V}_{2}-\mathbf{V}_{2}\otimes\mathbf{V}_{1}\big{)}^{-1}\left(\begin{bmatrix}\delta^{2}&\delta\cdot q\\ \delta\cdot q&q^{2}\end{bmatrix}\otimes\mathbf{H}\right) (in section E.3.1).

Combine the above to obtain bounds on $\bm{\Phi}_{\infty}$ (lemma 6).

Before deriving these bounds, we will present some reasoning behind the validity of the upper bounds that we derive on the stationary distribution $\bm{\Phi}_{\infty}$:

This part of the proof deals with deriving the solution to:

This is equivalent to solving the (linear) equation:

Given all these computations, comparing the $(1,1)$ term on both sides of equation E.3.1, we get:

Next, comparing $(1,2)$ term on both sides of equation E.3.1, we get:

Finally, comparing the $(2,2)$ term on both sides of equation E.3.1, we get:

Now, we note that equations E.3.1, E.3.1 are linear systems in two variables $u_{12}$ and $u_{22}$. Denoting the system in the following manner,

For analyzing the variance error, we require $u_{22},u_{12}$:

Substituting the values from equations E.3.1 and E.3.1, we get:

Denominator of $u_{22}$: Let us consider the denominator of $u_{22}$ (from equation 47) to write it in a concise manner.

Plugging in expressions for $q=\alpha\delta+(1-\alpha)\gamma$ and $c=\alpha(1-\beta)$, in $\text{dr}(u_{22})$ we get:

Next, considering $k_{1}-k_{2}$, we have:

Next, considering $\gamma(1-\alpha)k_{1}+\alpha\beta\delta\ k_{2}$, we have:

Consider $c\gamma(1-\alpha)+\alpha\beta q$:

Re-substituting this in the expression for $\gamma(1-\alpha)k_{1}+\alpha\beta\delta k_{2}$, we have:

Substituting equations E.3.1, E.3.1 into equation 49, we have:

We note that the denominator of $u_{12}$ (in equation 48) is just the negative of the denominator of $u_{22}$ as represented in equation E.3.1.

Numerator of $u_{22}$: We begin by writing out the numerator of $u_{22}$ (from equation 47):

We now consider $2cq\delta\gamma(1-\alpha)+(q^{2}+(c\delta)^{2})\alpha\beta\delta$:

Substituting equation E.3.1 into equation E.3.1 and grouping common terms, we obtain:

With this, we can write out the exact expression for $u_{22}$:

Numerator of $u_{12}$: We begin by rewriting the numerator of $u_{12}$ (from equation 48):

We split the simplification into two parts: one depending on $(1+c)$ and the other part representing terms that don’t contain $(1+c)$. In particular, we consider the terms that do not carry a coefficient of $(1+c)$:

Next, we consider the other term containing the $(1+c)$ part:

Substituting equations E.3.1, E.3.1 into equation 57, we get:

With which, we can now write out the expression for $u_{12}$:

Obtaining $u_{11}$: We revisit equation E.3.1 and substitute $u_{22}$ from equation 56:

From which, we consider the numerator of $u_{11}$ and begin simplifying it:

For obtaining a PSD upper bound on $\mathbf{U}_{22}$, we will write out a sharp bound of $u_{22}$ in each eigen space:

Let us consider bounding the numerator of the $2^{\text{nd}}$ term:

We will first upper bound the third term:

Let us consider bounding $\frac{q-c\delta}{2\delta\lambda_{j}}$ :

Substituting the values for $\alpha,\beta,\gamma,\delta$ applying $\frac{1}{1+\gamma\mu}\leq 1$, $c_{3}=\frac{c_{2}\sqrt{2c_{1}-c_{1}^{2}}}{c_{1}}$ and, $c_{2}^{2}=\frac{c_{4}}{2-c_{1}}$ with $0<c_{4}<1/6$ we get:

Now, consider the following bound on $1/c$:

Substituting values of $c_{1}$, $c_{4}$ we have: $1/c\leq 1.5$. This implies the following bound on $u_{22}$:

Alternatively, this implies that $\mathbf{U}_{22}$ can be upper bounded in a psd sense as:

E.3.2 Understanding fourth moment effects

We need to understand $\mathcal{M}\mathbf{U}_{22}$.

where, $s\stackrel{{\scriptstyle\textrm{def}}}{{=}}(6c_{4}+\frac{\delta R^{2}}{2})=23/30\leq\frac{4}{5}$. This implies (along with the fact that for any PSD matrices $\mathbf{A},\mathbf{B},\mathbf{C}$, if $\mathbf{A}\preceq\mathbf{B}$, then, $\mathbf{A}\otimes\mathbf{C}\preceq\mathbf{B}\otimes\mathbf{C}$)),

This will lead us to obtaining a PSD upper bound on $\bm{\Phi}_{\infty}$, i.e., the proof of lemma 6

We begin by recounting the expression for the steady state covariance operator $\bm{\Phi}_{\infty}$ and applying results derived from previous subsections:

Now, given the upper bound provided by equation E.3.2, we can now obtain a (mildly) looser upper PSD bound on $\mathbf{U}$ that is more interpretable, and this is by providing an upper bound on $\mathbf{U}_{11}$ and $\mathbf{U}_{22}$ by considering their magnitude along each eigen direction of $\mathbf{H}$. In particular, let us consider the max of $u_{11}$ and $u_{22}$ along the $j^{th}$ eigen direction (as implied by equations 63, 56):

This implies, we can now consider upper bounding the term in the equation above and this will yield us the result:

Plugging this into the bound for $\max{u_{11},u_{22}}$, we get:

This implies the bound written out in the lemma, that is,

Where, $d$ is the dimension of the problem.

Given these lemmas, we are now in a position to prove lemma 17.

So, we need to understand $\mathcal{R}\mathbf{U}$:

We begin by noting the following while considering the left side of the above expression:

The inner product above is a sum of two terms, so let us consider the first of the terms:

where $(\bm{\Phi}_{\infty})_{j}$ is the $2\times 2$ block of $\bm{\Phi}_{\infty}$ corresponding to the $j^{\textrm{th}}$ eigensubspace of $\mathbf{H}$, $(\bm{\Phi}_{\infty}^{1/2})_{j}$ denotes the $2\times 2d$ submatrix (i.e., $2$ rows) of $\bm{\Phi}_{\infty}^{1/2}$ corresponding to the $j^{\textrm{th}}$ eigensubspace and $\mathbf{A}_{j}$ denotes the $j^{\textrm{th}}$ diagonal block of $\mathbf{A}$. Note that $(\bm{\Phi}_{\infty}^{1/2})_{j}(\bm{\Phi}_{\infty}^{1/2})_{j}^{\top}=(\bm{\Phi}_{\infty})_{j}$. It is very easy to observe that the second term in the dot product can be written in a similar manner, i.e.:

So, essentially, the expression in the left side of the lemma can be upper bounded by using Cauchy-Shwartz inequality:

The advantage with the above expression is that we can now begin to employ psd upper bounds on the covariance of the steady state distribution $\bm{\Phi}_{\infty}$ and provide upper bounds on the expression on the right hand side. In particular, we employ the following bound provided by the taylor expansion that gives us an upper bound on $\bm{\Phi}_{\infty}$:

This implies in particular that $(\bm{\Phi}_{\infty})_{j}\preceq 5\sigma^{2}\mathbf{U}_{j}$ for every $j\in[d]$ and hence, for any vector $\left\lVert\mathbf{a}\right\rVert_{(\bm{\Phi}_{\infty})_{j}}\leq\sqrt{5\sigma^{2}}\left\lVert\mathbf{a}\right\rVert_{\mathbf{U}_{j}}$. The important property of the matrix $\mathbf{U}$ that serves as a PSD upper bound is that it is diagonalizable using the basis of $\mathbf{H}$, thus allowing us to bound the computations in each of the eigen directions of $\mathbf{H}$.

So, let us consider $\begin{bmatrix}\lambda_{j}^{1/2}&0\end{bmatrix}\mathbf{A}_{j}(\mathbf{I}-\mathbf{A}_{j})^{-2}$ and write out the following series of equations:

Next, let us consider $\left\lVert\begin{bmatrix}-c\\ 1\end{bmatrix}\right\rVert^{2}_{\mathbf{U}_{j}}$:

Note that $u_{11},u_{12},u_{22}$ share the same denominator, so let us evaluate the numerator $\text{nr}(c^{2}u_{11}-2cu_{12}+u_{22})$. For this, we have, from equations 63, 61, 56 respectively: Furthermore,

Again, this can be analyzed in each of the eigen directions $(\lambda_{j},\mathbf{u}_{j})$ of $\mathbf{H}$ to yield:

Now, we require to bound the product of equation E.3.2 and E.3.2:

We will consider the four terms within the square root and bound them separately:

Recall the bound on 1/c from equation E.3.1:

We add $T_{1}$ and $T_{2}$ and revisit equation 76:

Where the equation in the penultimate line is obtained by summing over all eigen directions the bound implied by equation E.3.2, and $C$ is a universal constant. ∎

We begin by recounting the expression for the covariance of the variance error of the tail-averaged iterate ${\bar{\bm{\theta}}}_{t,n}^{\textrm{variance}}$ from equation B.3:

The goal is to bound $\left\langle\begin{bmatrix}\mathbf{H}&0\\ 0&0\end{bmatrix},\mathcal{E}_{i}\right\rangle$, for $i=1,..,5$.

For the case of $\mathcal{E}_{2}$, we employ the result from lemma 19, and this gives us:

We will consider bounding $\|\mathbf{A}^{n-t}(\mathcal{I}-\mathcal{B})^{-1}\mathbf{\widehat{\Sigma}}\|$:

Furthermore, for $\|(\mathbf{I}-\mathbf{A}^{\top})^{-2}\mathbf{A}^{\top}\begin{bmatrix}\mathbf{H}&0\\ 0&0\end{bmatrix}\|$, we consider a bound in each eigendirection $j$ and accumulate the results subsequently:

Plugging this into equation E.3.2, we obtain:

To bound $\|\otimes_{2}\begin{bmatrix}-(c\mathbf{I}-q\mathbf{H})\mathbf{H}^{-1/2}\\ (\mathbf{I}-\delta\mathbf{H})\mathbf{H}^{-1/2}\end{bmatrix}\|$, we will consider a bound along each eigendirection and accumulate the results:

Next, we bound $\|\mathcal{B}^{k}(\mathcal{I}-\mathcal{B})^{-1}\mathbf{\widehat{\Sigma}}\|$ (as a consequence of lemma 13 with $\mathbf{Q}=\mathbf{\widehat{\Sigma}}$):

Plugging equation E.3.2 into equation E.3.2, we get:

In a manner similar to bounding $\|\mathbf{A}^{n-t}(\mathcal{I}-\mathcal{B})^{-1}\mathbf{\widehat{\Sigma}}\|$ as in equation E.3.2, we can bound $\|\mathbf{A}^{n-j}(\mathcal{I}-\mathcal{B})^{-1}\mathcal{B}^{j}\mathbf{\widehat{\Sigma}}\|$ as:

Furthermore, we will consider the bound $\|(\mathbf{I}-\mathbf{A}^{\top})^{-1}\mathbf{A}^{\top}\begin{bmatrix}\mathbf{H}&0\\ 0&0\end{bmatrix}\|$ along one eigen direction (by employing equation 26) and collect the results:

Plugging this into equation E.3.2, and upper bounding the sum by $(n-t)$ times the largest term of the series:

Summing up equations E.3.2, 83, E.3.2, E.3.2, E.3.2, the statement of the lemma follows. ∎

Appendix F Proof of Theorem 1

The proof of the theorem follows through various lemmas that have been proven in the appendix:

Section B provides the bias-variance decomposition and provides an exact tensor expression governing the covariance of the bias error (through lemma 3)and the variance error (lemma 5).

Section D provides a scalar bound of the bias error through lemma 16. The technical contribution of this section (which introduces a new potential function) is in lemma 4.

Section E provides a scalar bound of the variance error through lemma 20. The key technical contribution of this section is in the introduction of a stochastic process viewpoint of the proposed accelerated stochastic gradient method through lemmas 6, 17. These lemmas provide a tight characterization of the stationary distribution of the covariance of the iterates of the accelerated method. Lemma 19 is necessary to show the sharp burn-in (up to log factors), beyond which the leading order term of the error is up to constants the statistically optimal error rate $\mathcal{O}(\sigma^{2}d/n)$.

Combining the results of these lemmas, we obtain the following guarantee of algorithm 1: