Iterate averaging as regularization for stochastic gradient descent

Introduction

Stochastic gradient methods are ubiquitous in machine learning, where they are typically referred to as SGD (stochastic gradient descentAlbeit in general they might not be descent methods.). Since these incremental methods use little computation per data point, they are naturally adapted to processing very large data sets or streams of data. Stochastic gradient methods have a long history, starting from the pioneering paper by Robbins and Monro (Robbins and Monro, 1951). For a more detailed discussion, we refer to the excellent review given by Nemirovski et al. (2009). In the present paper, we propose a variant of SGD based on a weighted average of the iterates. The idea of averaging iterates goes back to Polyak (1990) and Ruppert (1988), and indeed it is often referred to as Polyak–Ruppert averaging (see also Polyak and Juditsky, 1992). In this paper, we study SGD in the context of the linear least-squares regression problem, considering both finite- and infinite-dimensional settings. This latter case allows to derive results for nonparametric learning with kernel methods—we refer to the appendix in Rosasco and Villa (2015) for a detailed discussion on this subject. The study of SGD for least squares is classical in stochastic approximation (Kushner and Yin, 2003), where it is commonly known as the least-mean-squares (LMS) algorithm. In the context of machine learning theory, prediction-error guarantees for online algorithms can be derived through a regret analysis in a sequential-prediction setting and a so called online-to-batch conversion (Shalev-Shwartz, 2012; Hazan, 2016). For example, results in Vovk (2001); Azoury and Warmuth (2001) and Hazan et al. (2007) directly apply to least squares. Alternatively, one can directly analyze SGD in the stochastic setting, as done by Smale and Yao (2006); Ying and Pontil (2008); Tarres and Yao (2014), where the last iterate and decaying step-size are considered, and more recently by Rosasco and Villa (2015); Lin and Rosasco (2017) where multiple passes and mini-batching are considered. A recently popular approach is combining constant step-sizes with Polyak–Ruppert averaging, which was first shown to lead to strong finite-time prediction guarantees after a single pass on the data by Bach and Moulines (2013). This approach was first studied by Györfi and Walk (1996) and subsequent progress was made by Défossez and Bach (2015); Dieuleveut et al. (2017); Jain et al. (2016, 2017); Lakshminarayanan and Szepesvári (2018).

In this paper we propose and analyze a novel form of weighted average, given by a sequence of weights decaying geometrically, so that the first iterates have more weight. Our main technical contribution is a characterization of the properties of this particular weighting scheme that we call geometric Polyak–Ruppert averaging. Our first result shows that SGD with geometric Polyak–Ruppert averaging is in expectation equivalent to considering SGD with a regularized loss function, and both sequences converge to the Tikhonov-regularized solution of the expected least-squares problem. The regularization parameter is a tuning parameter defining the sequence of geometric weights. This result strongly suggests that geometric Polyak–Ruppert averaging can be used to control the bias-variance properties of the corresponding SGD estimator. Indeed, our main result quantifies this intuition deriving a finite-sample bound, matching previous results for regularized SGD, and leading to optimal rates (Tsybakov, 2008). While averaging is widely considered to have a stabilizing effect, to the best of our knowledge this is the first result characterizing the stability of an averaging scheme in terms of its regularization properties and corresponding prediction guarantees. Our findings can be contrasted to recent results on tail averaging (Jain et al., 2016) and provide some guidance on when and how different averaging strategies can be useful. On a high level, our results suggest that geometric averaging should be used when the data is poorly-conditioned and/or relatively small, and tail averaging should be used in the opposite case. Further, from a practical point of view, geometric Polyak–Ruppert averaging provides an efficient approach to perform model selection, since a regularization path (Friedman et al., 2001) is computed efficiently. Indeed, it is possible to compute a full pass of SGD once and store all the iterates, to then rapidly compute off-line the solutions corresponding to different geometric weights (or tail averages), hence different regularization levels. As the averaging operation is entirely non-serial, this method lends itself to trivially easy parallelization.

The rest of the paper is organized as follows. In Section 2, we introduce the necessary background and present the geometric Polyak–Ruppert averaging scheme. In Section 3, we show the asymptotic equivalence between ridge regression and constant-stepsize SGD with geometric iterate averaging. Section 4 presents and dicuss our main results regarding the finite-time prediction error of the method. Section 5 describes the main steps in the proofs. The full proof is included in the Appendix. We conclude this section by introducing some basic notation used throughout the paper.

Preliminaries

We study the problem of linear regression under the square loss, more commonly known as linear least squares regression. The objective in this problem is to minimize the expected risk

and satisfies $R(w^{*})=\inf_{w\in\mathcal{H}}R(w)$. We assume $w^{*}$ to exist, since in general this might not be true in infinite dimensions. Also we abuse the notation in Eq. (2) since in general $\Sigma$ is not invertible and a pseudoinverse should be considered. This choice is made only to ease the notation.

We study algorithms that take as input a set of data points $\left\{(x_{t},y_{t})\right\}_{t=1}^{n}$ drawn identically and independently from $\mathcal{D}$ and output a weight vector $w$ to approximately minimize (1). The quality of an approximate solution is measured by the the excess risk

To compute a solution from data, we consider the stochastic gradient method, a.k.a. SGD, that for least squares takes the form

where $(\eta_{t})_{t}>0$ is a sequence of stepsizes (or learning rates), and $w_{0}\in\mathcal{H}$ is an initial point. Typically, a decaying stepsize sequence is chosen to ensure convergence, see Nemirovski et al. (2009) and references therein. A result relevant to our study is obtained by Bach and Moulines (2013) for a finite dimensional setting ($\cal H$ is of dimension $d$). Unlike previous results, here it is shown that convergence for a constant stepsize $\eta$ can be proved if Polyak–Ruppert (PR) averaging

is considered. This result was later strengthened in various respects by Défossez and Bach (2015) and Dieuleveut, Flammarion, and Bach (2017), notably by weakening the assumptions in Bach and Moulines (2013) and separating error terms related to “bias” and “variance”. We highlight one result from Dieuleveut et al. (2017) which considers the sequence

and under technical assumptions discussed later, proves the excess-risk bound

where $\sigma^{2}>0$ is an upper bound on the variance of the label noise. The iteration (4) is only of theoretical interest since the covariance $\Sigma$ is not known in practice. However, the obtained bound is simpler to present allows easier comparison with our result. A bound slightly more complex than (5) can be obtained when $\Sigma$ is not known (Dieuleveut et al., 2017, Theorem 2).

In this paper, we propose a generalized version of Polyak–Ruppert averaging that we call geometric Polyak–Ruppert averaging. Specifically, the algorithm we study computes the standard SGD iterates as given by Equation (3) and outputs

after round $n$, where $\lambda\in[0,1/\eta)$ is a tuning parameter. That is, the output is a geometrically discounted (and appropriately normalized) average of the plain SGD iterates that puts a higher weight on initial iterates. It is easy to see that setting $\lambda=0$ exactly recovers the standard form of Polyak–Ruppert averaging. Our main result essentially shows that the resulting estimate $\widetilde{w}_{n}$ satisfiesThe bound shown here concerns an iteration similar to the one shown on Equation (4), and is proved in Appendix C. We refer to Theorem 4.1 for the precise statement of our main result.

under the same assumptions as the ones made by Dieuleveut et al. (2017). Notably, this guarantee matches the bound of Equation (5), with the key difference being that the factor $\frac{1}{n}$ in the first term is replaced by $\frac{1}{n}+\frac{\eta\lambda}{2}$. This observation suggests the (perhaps surprising) conclusion that geometric Polyak–Ruppert averaging has a regularization effect qualitatively similar to Tikhonov regularization. Before providing the proof of the main result stated above, we first show that this similarity is more than a coincidence. Specifically, we begin by showing in Section 3 that the limit of the weighted iterates is exactly the ridge regression solution on expectation.

Geometric iterate averaging realizes Tikhonov regularization on expectation

between $w_{t}$ and $w_{t-1}$, which can be iteratively applied to obtain

In contrast we also define the iterates of regularized SGD with stepsize $\gamma>0$ and regularization parameter $\lambda$ as

This latter definition can be seen as an empirical version of the iteration in Eq. (4).

By our assumption on $\gamma$, we have $\gamma\Sigma\preccurlyeq I$, which implies that then the limit on the right-hand side exists and satisfies $\sum_{t=0}^{\infty}\left(I-\gamma\Sigma-\gamma\lambda I\right)^{t}=\left(\gamma\Sigma+\gamma\lambda I\right)^{-1}$. Having established the existence of the limit, we rewrite the regularized SGD iterates (7) as

Let $\eta$, $\gamma$ and $\lambda$ be such that $\gamma\lambda<1$ and $\eta=\frac{\gamma}{1-\gamma\lambda}$. Then,

This proposition is proved using the same ideas as Proposition 3.1; we include the proof in Appendix A for completeness.

Main result: Finite-time performance guarantees

While the previous section establishes a strong connection between the geometrically weighted SGD iterates with the iterates of regularized SGD on expectation, this connection is clearly not enough for the known performance guarantees to carry over to our algorithm. Specifically, the two iterative schemes propagate noise differently, thus the covariance of the resulting iterate sequences may be very different from each other. In this section, we prove our main result that shows that the prediction error of our algorithm also behaves similarly to that of SGD with Tikhonov regularization. For our analysis, we will make the same assumptions as Dieuleveut et al. (2017) and we will borrow several ideas from them, as well as most of their notation.

We now state the assumptions that we require for proving our main result. We first state an assumption on the fourth moment of the covariates.

This assumption implies that $\mbox{tr}\left[\Sigma\right]\leq R^{2}$, and is satisfied, for example, when the covariates satisfy $\left\|x\right\|\leq R$ almost surely. We always assume a minimizer $w_{*}$ of the expected risk to exist and also make an assumption on the residual $\varepsilon$ defined as the random variable

This assumption is satisfied when $\left\|x\right\|$ and $y$ are almost surely bounded, or when the model is well-specified and corrupted with bounded noise (i.e., when $\varepsilon$ is independent of $x$ and has variance bounded by $\sigma^{2}$). Under the above assumptions, we prove the following bound on the prediction error—our main result:

Suppose that Assumptions 1 and 2 hold and assume that $\eta\leq\frac{1}{2R^{2}}$ and $\lambda\in[0,1/\eta)$. Then, the iterates computed by the recursion given by Equations (3) and (6) satisfy

The (rather technical) proof of the theorem closely follows the proof of Theorem 2 of Dieuleveut et al. (2017). We describe the main components of the proof of our main result in Section 5. For didactic purposes, we also present a simplified version of our analysis where we assume full knowledge of $\Sigma$ in Appendix C.

We next discuss various aspects and implications of our results.

Apart from constant factorsBy enforcing $\gamma\lambda\leq 1/2$, the $1-\gamma\lambda$ and $2-\gamma\lambda$ terms in the denominator can be lower bounded by a constant., our bound above precisely matches that of Theorem 1 of Dieuleveut et al. (2017), except for an additional term of order $\gamma\lambda\sigma^{2}\mbox{tr}\left[\Sigma^{2}\left(\Sigma\right)^{-2}\right]$. This term, however, is not a mere artifact of our proof: in fact it captures a distinctive noise-propagating effect of geometric PR averaging scheme. Indeed, the regularization effect of our iterate averaging scheme is different from that of Tikhonov-regularized SGD in one significant way: while Tikhonov regularization increases the bias and strictly decreases the variance, our scheme may actually increase the variance for certain choices of $\lambda$. To see this, observe that setting a large $\lambda$ puts a large weight on the initial iterates, so that the initial noise is amplified compared to noise in the later stages, and the concentration of the total noise becomes worse. We note however that this extra term does not qualify as a serious limitation, since the commonly recommended setting \lambda=\mathcal{O}\bigl{(}\frac{1}{\eta n}\bigr{)} still preserves the optimal rates for both the bias and the variance up to constant factors.

Optimal excess risk bounds.

The bound in Theorem 4.1 is essentially the same as the one derived in Dieuleveut et al. (2017)[Theorem 2]. Following their same reasoning, the bound can be optimized with respect to $\lambda,\gamma$ to derive the best parameters choice and explicit upper bounds on the corresponding excess risk. In the finite dimensional case, it is easy to derive a bound of order ${\cal O}(d/n)$, which is known to be optimal in a minmax sense (Tsybakov, 2008). In the infinite dimensional case, optimal minmax bound can again be easily derived, and also refined under further assumptions on $w_{*}$ and the covariance $\Sigma$ (De Vito et al., 2005; Caponnetto and De Vito, 2007). We omit this derivation.

When should we set λ>0𝜆0\lambda>0?

We have two answers depending on the dimensionality of the underlying Hilbert space $\mathcal{H}$. For infinite dimensional spaces, it is clearly necessary to set $\lambda>0$. In the finite-dimensional case, the advantage of our regularization scheme is less clear at first sight: while Tikhonov regularization strictly decreases the variance, this is not necessarily true for our scheme (as discussed above). A closer look reveals that, under some (rather interpretable) conditions, we can reduce the variance, as quantified by the following proposition.

If $\mbox{tr}\left[\Sigma^{-1}\right]>\frac{1}{2}\gamma dn$ there exists a regularization parameter $\lambda^{*}>0$ satisfying

Letting $s_{1},s_{2},\dots,s_{d}$ be the eigenvalues of $\Sigma$ sorted in decreasing order, we have

Taking derivative of $f$ with respect to $\lambda$ gives

so $f^{\prime}(0)>0$ holds whenever $\mbox{tr}\left[\Sigma^{-1}\right]>\frac{1}{4}\gamma dn$. The proof is concluded by observing that $f^{\prime}(0)>0$ implies the existence of a $\lambda^{*}$ with the claimed property.

Intuitively, Proposition 4.2 suggests that the geometric PR averaging can definitely reduce the variance over standard PR averaging whenever the covariance matrix is poorly conditioned and/or the sample size is small. Notice however that the above argument only shows one example of a good choice of $\lambda$; many other good choices may exist, but these are harder to characterize.

Geometric averaging vs. tail averaging.

It is interesting to contrast our approach with the tail averaging scheme studied by Jain et al. (2016, 2017): instead of putting large weight on the initial iterates as our method does, Jain et al. suggest to average the last $n-\tau$ iterates of SGD for some $\tau$. The effect of this operation is that the $\left\|\Sigma^{-1/2}\left(w^{*}-w_{0}\right)\right\|^{2}n^{-2}$ term arising from Polyak–Ruppert averaging is replaced by a term of order $\exp(-\eta\mu\tau)\left\|w^{*}-w_{0}\right\|^{2}$, where $\mu>0$ is the smallest eigenvalue of $\Sigma$. Clearly, this yields a significant asymptotic speedup, but gives no advantage when $\gamma n\leq\mu^{-1}$ (noting that $\tau<n$). Contrasting this condition with our Proposition 4.2 leads to an interesting conclusion: for small values of $n$, geometric averaging has an edge over tail averaging and vice versa.

What is the computational advantage?

The main practical advantage of our averaging scheme over Tikhonov regularization is a computational one: validating regularization parameters becomes trivially easy to parallelize. Indeed, one can perform a single pass of unregularized SGD over the data, store the iterates and average them with various schedules to evaluate different choices of $\lambda$. Through parallelization, this approach can achieve huge computational speedups over running regularized SGD from scratch. To see this, observe that the averaging operation is entirely non-serial: one can cut the (stored) SGD iterates into $K$ contiguous batches and let each individual worker perform a geometrically discounted averaging with the same discount factor $(1-\gamma\lambda)$. The resulting averages are then combined by the master with appropriate weights. In contrast, regularized SGD is entirely serial, so validation cannot be parallelized.

Choosing the right averaging.

Finally, we note that the same method as above can be used to choose the correct parameter $\tau$ for tail averaging. This suggests a simple and highly parallelizable scheme for choosing the right averaging, with the identity of the best scheme depending on the interaction between $n$ and $\Sigma$ as discussed above.

Connections to early stopping.

Open questions.

It is natural to ask whether geometric iterate averaging has similar regularization effects in other stochastic approximation settings too. An possible direction for future work is studying the effects of our averaging scheme on accelerated and/or variance-reduced variants of SGD. Another promising direction is studying general linear stochastic approximation schemes (Lakshminarayanan and Szepesvári, 2018), and particularly Temporal-Difference learning algorithms for Reinforcement Learning that have so far resisted all attempts to regularize them (Sutton and Barto, 1998; Szepesvári, 2010; Farahmand, 2011).

The proof of Theorem 4.1

Our proof closely follows that of Dieuleveut et al. (2017, Theorem 1), with the key differences that

we do not have to deal with an explicit “regularization-based” error term that gives rise to a term proportional to $\lambda^{2}$ in their bound, and

the $\frac{1}{n}$ factors for iterate averaging are replaced by $c_{n}(1-\gamma\lambda)^{t}$ for each round, where

As we will see, this change will propagate through the analysis and will eventually replace the $\frac{1}{\gamma^{2}n}$ and $\left(2\lambda+\frac{1}{\eta n}\right)$ factors in the final bound by $c_{n}^{2}\sum_{t=1}^{n}\left(1-\gamma\lambda\right)^{2t}$ and $\frac{c_{n}^{2}}{\gamma^{2}}$, respectively.

In the interest of space, we only provide an outline of the proof here and defer the proofs of the key lemmas to Appendix B. Throughout the proof, we will suppose that the conditions of Theorem 4.1 hold. The lemma below shows that the factors involving $c_{n}^{2}$ are of the order claimed in the Theorem.

The straightforward proof is given in Appendix B.1.

Now we are ready to lay out the proof of Theorem 4.1. Let us start by introducing the notation

and recalling the definition $\varepsilon_{t}=y_{t}-\left\langle x_{t},w^{*}\right\rangle$. A simple recursive argument shows that

We first show a simple upper bound on the excess risk $\Delta(\widetilde{w}_{n})=\left\|\Sigma^{1/2}\left(\widetilde{w}_{t}-w^{*}\right)\right\|^{2}$:

The proof is included in Appendix B.2. In order to further upper bound the right-hand side in the bound stated in Lemma 5.2, we can combine the decomposition of $w_{t}-w^{*}$ in Equation (8) with the Minkowski inequality to get

The first term in the above decomposition can be thought of as the excess risk of a “noiseless” process (where $\sigma=0$) and the second term as that of a “pure noise” process (where $w_{0}=w^{*}$). The rest of the analysis is devoted to bounding these two terms.

We begin with the conceptually simpler case of bounding $\Delta_{2},t$, which can be done uniformly for all $t$. In particular, we have the following lemma:

The rather technical proof is presented in Appendix B.3. We now turn to bounding the excess risk of the “noiseless” process, $\Delta_{1}$:

The following lemma states a bound on $\Delta_{1}$.

The extremely technical proof of this theorem is presented in Appendix B.4.

The proof of Theorem 4.1 is concluded by plugging the bounds of Lemmas 5.3 and 5.4 into Equation (9) and using Lemma 5.2 to obtain

Now we can finish by using the bounds on $c_{n}^{2}$ and $c_{n}^{2}\sum_{t=0}^{n}(1-\gamma\lambda)^{2t}$ given in Lemma 5.1.

GN is supported by the UPFellows Fellowship (Marie Curie COFUND program n∘ 600387). LR is funded by the Air Force project FA9550-17-1-0390 (European Office of Aerospace Research and Development) and by the FIRB project RBFR12M3AC (Italian Ministry of Education, University and Research). The authors thank Francesco Orabona, Csaba Szepesvári and Gábor Lugosi for interesting discussions.

References

Appendix A The proof of Proposition 3.3

The proof is similar to that of Proposition 3.1, although a little more cluttered due to the normalization constants involved in the definition of $\widetilde{w}_{t}$. The analysis in this case relies on defining the truncated geometric random variable $G$ with law

where we used $\sum_{k=0}^{t}(1-\gamma\lambda)^{k}=\frac{1-\left(1-\gamma\lambda\right)^{t}}{\gamma\lambda}$ and the definition of $\widetilde{w}_{t}$ in the last step. This concludes the proof.

Appendix B Tools for proving Theorem 4.1

where the first inequality uses $1-x\leq e^{-x}$ that holds for all $x\in$ and the second one uses $\frac{1}{1-e^{-x}}\leq\frac{1}{x}+1$ that holds for all $x>0$. The second statement is proven as

where the first inequality again uses $1-x\leq e^{-x}$ that holds for all $x\in$ and the second one uses $\frac{e^{-x}}{1-e^{-x}}\leq\frac{1}{x}$ that holds for all $x>0$. \jmlrQED

B.2 The proof of Lemma 5.2

To handle the second term, we first notice that for any $t$ and $k>t$, we have

Noticing that the last term matches the first term on the right-hand side of Equation (10), the proof is concluded. \jmlrQED

B.3 The proof of Lemma 5.3

The proof of this lemma crucially relies on the following inequality:

Assume that $\eta\leq\frac{1}{2R^{2}}$. Then, for all $k<t$, we have

The result follows from noticing that $\Sigma\preccurlyeq\Sigma^{2}\left(\Sigma+\lambda I\right)^{-1}$.

B.4 The proof of Lemma 5.4

holds since the sum only has positive elements. Following Dieuleveut et al. (2017) again, we define the operator $\mathcal{T}$ acting on an arbitrary Hermitian operator $A$ as

and also introduce the operator $\mathcal{S}$ defined as

so that $\mathcal{T}A=\Sigma A+A\Sigma-\eta\mathcal{S}A$. We note that $\mathcal{S}$ and $\mathcal{T}$ are Hermitian and positive definite (the latter being true by our assumption about $\eta$). Finally, we define $\mathcal{I}$ as the identity operator acting on Hermitian matrices. With this notation, we can write

Thus, defining $E_{0}=\left(w_{0}-w^{*}\right)\otimes\left(w_{0}-w^{*}\right)$, we have

where the last step holds true if $\left\|\mathcal{I}-\eta\mathcal{T}\right\|<1$. For a proof of this fact, we refer to Lemma 5 in Défossez and Bach (2015). Let us define $\mathcal{T}_{\lambda}=\mathcal{T}+\lambda\mathcal{I}$ and $W=\mathcal{T}_{\lambda}^{-1}\left[\Sigma\left(\Sigma+\lambda I\right)^{-1}\right]$, so that it remains to bound $\gamma^{-1}\mbox{tr}\left[WE_{0}\right]$. We notice that, by definition, $W$ satisfies

Also introducing the operators $\mathcal{U}_{L}$ and $\mathcal{U}_{R}$ as the left- and right-multiplication operators with $\Sigma$, respectively, we get after reordering that

Using the fact that $\mathcal{U}_{L}+\mathcal{U}_{R}+\lambda I$ and its inverse are Hermitian, we can show

Furthermore, by again following the arguments This result is proven for finite dimension by Dieuleveut et al. (2017), but can be easily generalized to infinite dimensions. of Dieuleveut et al. (2017, pp. 28), we can also show

Since $\mathcal{S}W$ is positive, this leads to the bound

so it remains to bound $\mbox{tr}\left[\mathcal{S}W\right]$. On this front, we have

by our assumption on the covariates. Also, by Equation (11), we have

by crucially using the assumption $\eta\leq 1/2R^{2}$. Plugging into Equation (12) and using $\eta R^{2}\leq 2$ again proves the lemma. \jmlrQED

Appendix C Analysis under additive noise

In this setting, stochastic gradient descent takes the form

which is the unregularized counterpart of the iteration already introduced in Section 2 as Equation (4). Again, we will study the geometric average

We prove the following performance guarantee about this algorithm:

Recalling the notation $c_{n}=\left(\sum_{t=0}^{n}(1-\gamma\lambda)^{t}\right)^{-1}$, the geometric average can be written as

Now, exploiting the assumption that the noise is i.i.d. and zero-mean, we get

The proof is concluded by appealing to Lemma 5.1.