The Implicit Regularization of Stochastic Gradient Flow for Least Squares
Alnur Ali, Edgar Dobriban, Ryan J. Tibshirani
Introduction
Stochastic gradient descent (SGD) is one of the most widely used optimization algorithms—given the sizes of modern data sets, its scalability and ease-of-implementation means that it is usually preferred to other methods, including gradient descent (Bottou 1998; Bottou 2003; Zhang 2004; Bousquet & Bottou 2008; Bottou 2010; Bottou et al. 2016).
A summary of our contributions in this paper is as follows.
We give a bound on the excess risk of stochastic gradient flow at time , over ridge regression with tuning parameter , for all . The bound decomposes into three terms. The first term is the (scaled) variance of ridge. The second and third terms both stem from the variance due to mini-batching, and may be made smaller by, e.g., increasing the mini-batch size and/or decreasing the step size. The second term may be interpreted as the “price of stochasticity”: it is nonnegative, but vanishes as time grows. The third term is tied to the limiting optimization error of stochastic gradient flow: it is zero in the overparametrized (interpolating) regime (Bassily et al. 2018), but is positive otherwise, reflecting the fact that stochastic gradient flow with a constant step size fluctuates around the least squares solution as time grows. The bound holds with no conditions on the data matrix . Numerically, the bound can be small, indicating a tight relationship between the two estimators.
Using the bound, we show through numerical examples that stochastic gradient flow, when stopped at a time that (optimally) balances its bias and variance, yields a solution attaining risk that is 1.0032 times that of the (optimally-stopped) ridge solution, in less time—indicating that stochastic gradient flow strikes a favorable computational-statistical trade-off.
We give a similar bound on the distance between the coefficients of stochastic gradient flow at time , and those of ridge regression with tuning parameter , which is also seen to be tight.
Outline.
Next, we review related work. Section 2 covers notation, and further motivates the continuous-time approach. In Section 3, we present our bound on the excess risk of stochastic gradient flow over ridge regression. In Section 4, we present a bound relating the coefficients of the two estimators. Section 5 gives numerical examples supporting our theory. In Section 6, we conclude.
Related Work.
Stochastic Gradient Descent. The statistical and computational properties of SGD have been studied intensely over the years, with work tracing back to Robbins & Monro 1951; Fabian 1968; Ruppert 1988; Kushner & Yin 2003; Polyak & Juditsky 1992; Nemirovski et al. 2009. On the statistical side, a lot of the work has focused on delivering optimal error rates for SGD and its many variants, e.g., with averaging, either asymptotically (Robbins & Monro 1951; Fabian 1968; Ruppert 1988; Kushner & Yin 2003; Polyak & Juditsky 1992; Moulines & Bach 2011; Toulis & Airoldi 2017; Nemirovski et al. 2009), or in finite samples (Cesa-Bianchi et al. 1996; Zhang 2004; Ying & Pontil 2008; Cesa-Bianchi & Lugosi 2006; Pillaud-Vivien et al. 2018; Jain et al. 2018).
Notably, Bach & Moulines 2013; Défossez & Bach 2014; Dieuleveut et al. 2017a; Jain et al. 2017; Babichev & Bach 2018 studied SGD with a constant step size for least squares regression with averaging (obtaining optimal rates, which is not our focus). Good references on inference and computation include Fabian 1968; Ruppert 1988; Polyak & Juditsky 1992; Moulines & Bach 2011; Chen et al. 2016; Toulis & Airoldi 2017 and Recht et al. 2011; Duchi et al. 2015, respectively. Mandt et al. 2015; Duvenaud et al. 2016 interpreted SGD with a constant step size as doing Bayesian inference. Many works have empirically investigated the generalization properties of SGD, mainly in the context of non-convex optimization (Jastrzębski et al. 2017; Kleinberg et al. 2018; Zhang et al. 2018; Jin et al. 2019; Nakkiran et al. 2019; Saxe et al. 2019).
Implicit Regularization. Nearly all of the work in implicit regularization thus far has examined the convergence points of gradient descent, and not the whole path, for specific convex (Nacson et al. 2018; Gunasekar et al. 2018a; Soudry et al. 2018; Vaskevicius et al. 2019) and non-convex (Li et al. 2017; Wilson et al. 2017; Gunasekar et al. 2017; Gunasekar et al. 2018b) problems. Notable exceptions include Rosasco & Villa 2015; Lin et al. 2016; Lin & Rosasco 2017; Neu & Rosasco 2018, who studied averaged SGD with a constant step size for least squares regression, arguing that the various algorithmic parameters (i.e., the step size, mini-batch size, number of iterations, etc.) perform a kind of implicit regularization, by inspecting the corresponding error rates. A few works have investigated implicit regularization outside of optimization (Mahoney & Orecchia 2011; Mahoney 2012; Gleich & Mahoney 2014; Martin & Mahoney 2018).
Stochastic Differential Equations. Several papers have studied the same stochastic differential equation that we do (Hu et al. 2017; Feng et al. 2017; Li et al. 2019; Feng et al. 2019), but without the focus on implicit regularization and statistical learning. Along these lines, somewhat related work can be found in the literature on Langevin dynamics (Geman & Hwang 1986; Seung et al. 1992; Neal et al. 2011; Welling & Teh 2011; Sato & Nakagawa 2014; Teh et al. 2016; Raginsky et al. 2017; Cheng et al. 2019).
Preliminaries
Consider the usual least squares regression problem,
for , where is a fixed step size, is the mini-batch size, and denotes the mini-batch on iteration with , for all . For simplicity, we assume the mini-batches are sampled with replacement; our results hold with minor modifications under sampling without replacement. We assume the initialization .
Now, adding and subtracting the negative gradient of the loss in (2) yields
This may be recognized as gradient descent, plus the deviation between the sample average of i.i.d. random variables and their mean, which motivates the continuous-time dynamics (stochastic differential equation)
with . Here, is standard -dimensional Brownian motion. We denote the diffusion coefficient
where the randomness is due to . We call the diffusion process (4) stochastic gradient flow.
At this point, it helps to recall the related work of Ali et al. 2018, who studied gradient flow,
which is gradient descent for (1) with infinitesimal step sizes. In what follows, we frequently use the solution to (6),
where and denote the matrix exponential and the Moore-Penrose pseudo-inverse of , respectively.
Unlike gradient flow, the continuous-time flow (4) does not arise by taking limits of the discrete-time dynamics (2), and should instead be interpreted as an approximation to (2). To see this, consider the Euler discretization of (4),
Figure 1 presents a small numerical example, where we see a striking resemblance between the paths for SGD, the Euler discretization of stochastic gradient flow, and ridge regression with tuning parameter .
2 Basic Properties of Stochastic Gradient Flow
We begin with an important lemma further motivating the differential equation (4); its proof, as with many of the results in this paper, may be found in the supplement. The result shows that both the first and second moments of the Euler discretization of (4) match those of the underlying discrete-time SGD iteration. This means that any deviation between the first two moments of the continuous-time flow (4) and discrete-time SGD must be due to discretization.
Discretization, i.e., showing that (8) and (2) are close in a precise sense, turns out to be non-trivial, and is left to future work.
Next, with the above motivation in mind, we present a lemma establishing that the solution to (4) exists and is unique. The result also gives a more explicit expression for the solution to (4), which plays a key role in many of the results to come.
Fix , , and . Let . Then
is the unique solution to the differential equation (4).
The result actually holds for any Lipschitz continuous diffusion coefficient , e.g., , as well as the time-homogeneous covariance (Mandt et al. 2017; Wang 2017; Dieuleveut et al. 2017b; Fan et al. 2018). In the former case, (4) reduces to (rescaled) Langevin dynamics.
3 Constant vs. Non-Constant Covariances
The differential equation (4) has been considered previously (Hu et al. 2017; Feng et al. 2017; Li et al. 2019; Feng et al. 2019), but several works (Mandt et al. 2017; Wang 2017; Dieuleveut et al. 2017b; Fan et al. 2018) have found it convenient to work with the simplification
where . Here, . However, we present a simple but telling example revealing that these two processes, i.e., the non-constant covariance process in (4), and the constant covariance process in (10), need not be close in general.
Consider the univariate responseless least squares problem,
Let , for . Then SGD for the above problem may be expressed as
Assume the initial point is a nonzero constant, the follow a continuous distribution, and is sufficiently small. Letting be arbitrary, the basic estimate combined with Markov’s inequality shows that
Summing the right-hand side over , we conclude that converges to zero with probability one, by the first Borel-Cantelli lemma.
Now let . We may calculate for the non-constant process that , where , meaning the non-constant process follows the dynamics (the sign of may be chosen arbitrarily)
which may be recognized as a geometric Brownian motion. It can be checked that both the mean and variance of the geometric Brownian motion tend to zero as time grows, provided that , which certainly holds when .
On the other hand, the constant process is an Ornstein-Uhlenbeck process,
Again, it may be checked (e.g., Chapter 5 in Øksendal 2003) that the process mean goes down to zero, whereas the variance tends to the constant . In other words, the limiting dynamics of the constant process exhibit constant-order fluctuations, whereas those of the non-constant process do not. Therefore, for this problem, the latter dynamics more accurately reflect those of discrete-time SGD. See Figure 2 for an example.
We close this section with a simple result bounding the deviation between solutions to the non-constant and constant processes, in expectation. The result indicates that the two processes can be close when the non-constant process dynamics are close to the underlying coefficients. A thorough comparison of the two processes is left to future work.
Statistical Risk Bounds
Here and throughout, we let the predictor matrix be arbitrary and fixed, and assume the response follows a standard regression model,
Here denotes any potential randomness inherent to (e.g., due to mini-batching). We also consider in-sample risk,
We let denote the sample covariance matrix with eigenvalues and eigenvectors , for , and let and denote the smallest nonzero and largest eigenvalues of , respectively.
2 Risk Bounds
Recall the bias-variance decomposition for risk,
A straightforward calculation using the law of total variance shows (see the proof of Theorem 2 for details)
Fix , , and . Let . Then
which may be manipulated to obtain the result given in the lemma (see the supplement for details).
In either case, set small enough so that . Then,
Lemma 5 can be seen as the continuous-time analog of, e.g., Theorem 4 in Karimi et al. 2016, and may be of independent interest.
We recall a result from Ali et al. 2018, paraphrased below.
Putting Lemmas 4 and 5 together with Theorem 1 yields the following result, relating the risk of stochastic gradient flow to that of gradient flow and ridge regression.
Fix . Set according to Lemma 5. Let .
The analogous results for in-sample risk are similar, and deferred to the supplement for space reasons.
Turning to the variance, the law of total variance and the above calculation implies
As for the trace appearing in (15), we have
Here, the second line followed from Lemma 4. The third followed from Fubini’s theorem. The fourth followed by integrating, using the eigendecomposition and Lemma 5, along with one final application of Fubini’s theorem. This shows the claim for gradient flow. The claim for ridge follows by applying Theorem 1. ∎
The following result gives a more interpretable version of Theorem 2, at the expense of some sharpness.
Fix . Set as in Lemma 5. Let . Define
, and .
Interestingly, the result shows that the risk of stochastic gradient flow may be seen as the ridge bias raised to a power strictly less than 1, plus a time-dependent scaling of the ridge variance—which is quite different from the situation with gradient flow (cf. Theorem 1).
Finally, subtracting the ridge risk from both sides of (13) immediately gives our main result, a bound on the excess risk of stochastic gradient flow over ridge.
Fix . Set as in Lemma 5. Let . Then,
We can understand the influence of the effective variance terms on the risks (12), (13), (18) as follows. As stochastic gradient flow moves away from initialization, the stochastic gradients become smaller, and so their variance decreases, which is captured by the first term in (11), as it goes down with time. As stochastic gradient flow approaches the least squares solution, there are two possibilities, depending on whether the solution is interpolating or not. If the solution is interpolating, then stochastic gradient flow can fit the data perfectly, and hence in (11). Otherwise, stochastic gradient flow fluctuates around the solution, which is captured by the second term in (11), as it grows with time.
It is also interesting to note that the bounds (12), (13), (18) depend linearly on , corroborating recent empirical work (Krizhevsky 2014; Goyal et al. 2017; Smith et al. 2017; You et al. 2017; Shallue et al. 2019).
For space reasons, we compare the excess risk bound (18) to the analogous bound for the time-homogeneous process (10) in the supplement.
Coefficient Bounds
Our main result now follows easily, by putting Lemma 7 together with Lemma 4 from Section 3.
Fix . Set as in Lemma 5. Let . Then,
Lemma 7 gives a bound on the first term in the preceding display. Lemma 4 and the same arguments used in the proof of Theorem 2 give a bound on the second term. Putting the pieces together yields the result. ∎
It is possible to give a similar, albeit less sharp, result for any convex loss satisfying a restricted secant inequality (Zhang & Yin 2013), and noise process satisfying a suitable boundedness condition (Vaswani et al. 2018).
Numerical Examples
We give numerical examples supporting our theoretical findings. We generated the data matrix according to , where the entries of were i.i.d. following a normal distribution. We allow for correlations between the features, setting the diagonal entries of the predictor covariance to 1, and the off-diagonals to 0.5. Below, we present results for , , and . The supplement gives additional examples with different problem sizes and data models (Student-t and Bernoulli data); the results are similar. We set 2.2548e-4, following Lemma 5.
Figure 4 plots the risk of ridge regression, discrete-time SGD (2), and Theorem 2. For ridge, we used a range of 200 tuning parameters , equally spaced on a log scale from to . The expression for the risk of ridge is well-known. For Theorem 2, we set . For SGD, we computed its effective time, using and . As for its risk, following the decomposition given in Section 3, we first computed the bias and variance of discrete-time gradient descent, using Lemma 3 in Ali et al. 2018, and then added in the variance given by Lemma 1. As a comparison, Figure 4 also plots the risks of gradient flow (7), coming from Lemma 5 in Ali et al. 2018, and discrete-time gradient descent (as was just discussed).
Though the risks look similar, there are subtle differences (the supplement gives examples with larger step sizes and smaller mini-batch sizes, where the differences are more pronounced). We also see that Theorem 2 tracks the risk of SGD closely. In fact, the maximum ratio, across the entire path, of the risk of stochastic gradient flow to that of ridge is 2.5614, whereas the same ratio for SGD to ridge is 1.7214. Figure 4 also shows the (optimal) time where each method balances its bias and variance. Choosing a tuning parameter by balancing bias and variance is common in nonparametric regression, and doing so here implies that stochastic gradient flow stops earlier than gradient flow, because the effective variance terms (11) are nonnegative. We find the optimal stopping times chosen by balancing bias and variance vs. directly minimizing risk are generally similar. Moreover, the ratio of the (optimal) risks at these times is 1.0032, indicating that stochastic gradient flow strikes a favorable computational-statistical trade-off.
Discussion
Acknowledgements
We thank a number of people for helpful discussions, including Misha Belkin, Quanquan Gu, J. Zico Kolter, Jason Lee, Yi-An Ma, Jascha Sohl-Dickstein, Daniel Soudry, and Matus Telgarsky. ED was supported in part by NSF BIGDATA grant IIS 1837992 and NSF TRIPODS award 1934960. Part of this work was completed while ED was visiting the Simons Institute.
References
Supplementary Material
Appendix S.1 Proof of Lemma 1
For simplicity, below we will omit the source of the randomness for the various estimators. Implicitly, the randomness is from minibatching in SGD, and from the normal random increments in the discretization of SGD (which we wil can dSGF).
By taking expectations in the SGD iteration, we find
This identity only uses that the stochastic gradients are unbiased estimators for the true gradients. Thus, it is true even more generally for any loss function, not just for quadratic loss. However, for quadratic loss, we have a very special property, namely that the gradient is linear in the parameter. Using this, we can move the expectation inside, and we find
By taking the covariance of the SGD iteration, conditionally on the previous iterate , we find
We can also find the explicit form of the recursion. While this is not required in the statement of the lemma, it is used in our numerical examples.
This gives an explicit linear recursion for the covariance matrices. The first term can be viewed as a covariance matrix of the gradients evaluated at the mean value of the process (i.e., at the value of the GD iteration). The second term depends on the covariance of the previous iteration.
Appendix S.2 Proof of Lemma 2
As the diffusion coefficient is Lipschitz continuous and positive semidefinite, standard results from numerical analysis (e.g., Øksendal 2003) show that the solution to the differential equation (4) exists and is unique.
Plugging in the expression for from (4) and simplifying, we see that
Considering only the first integral above, by arguments similar to those given in Lemma 1 of Ali et al. 2018, we obtain
Appendix S.3 Proof of Lemma 3
Calculations similar to those given in Lemma 4 (appearing below) show
Continuing on, and writing , we have
Using the simple fact that , along with the fact that have nonnegative diagonal entries, now yields the result. ∎
where the absolute value is to be interpreted elementwise.
Using the matrix perturbation inequality given in Lemma A.2 of Nguyen et al. 2019, we see that
Noting the expression for the covariance matrix of the stochastic gradients given in (S.3), we obtain for (S.1) that
Appendix S.4 Proof of Lemma 4
Using Ito’s isometry along with the linearity of the trace, we obtain
Letting , the trace appearing in (S.2) may be expressed as
Appendix S.5 Proof of Lemma 5
In this lemma, it will be helpful to start slightly more generally, with the SDE for SGF on a general loss function . The specific proofs of this lemma are in Sections S.5.1 and S.5.2.
To approximate discrete time SGD with learning rate and batch size , it is not hard to see that the same logic we have used throughout the paper leads to the SDE
where is the covariance of the gradients at parameter value , and .
We derive the SDE for the behavior of the loss function itself, for a general loss. For gradient flow on a loss function , i.e., the dynamics , it is well known that the dynamics induced on the loss function is:
This shows that the loss function is always non-increasing, i.e., that gradient flow is a descent method. In contrast, we will find that the loss for stochastic gradient flow is not always non-increasing. We mention that a related calculation has been performed in (Zhu et al. 2018), under different assumptions (starting from a local min, integrating over time), and for a different purpose (to understand dynamics escaping local minima).
For SGF on a loss function , the value of the loss function evolves according to the following SDE:
where is the covariance of the stochastic gradients at parameter value . Also , where SGF approximates discrete time SGD with learning rate and batch size . This can be written in a distributionally equivalent way as
where is a 1-dimensional Brownian motion.
where is a 1-dimensional Brownian motion.
where is the covariance of the gradients at parameter value . Then, Ito’s rule leads to
For the special case of least squares loss, we have the following. We have already calculated most terms, and we have in addition that the second moment matrix of the gradients is
where is the residual. Plugging in the terms for least squares,
Here is a 1-dimensional Brownian motion, which is obtained by transforming the original diffusion term, which is a linear combination of the entries of , into a distributionally equivalent 1-dimensional process. Letting
Comparing this with the noiseless case, i.e., when , we note that both the drift and the diffusion terms have changed. The drift term is reduced by a term that is proportional to . The diffusion term is new altogether. This shows that for sufficiently large , the drift will not be positive, and hence the process will not converge to a point mass limit distribution.
Let us start with studying the diffusion with the second moment matrix first. We will show a geometric contraction of the loss. We can bound the terms in the drift term as follows. We have (using for elementwise product of two conformable vectors or matrices)
The second inequality holds with being the smallest nonzero singular value of . Why? Because it is easy to see that we always have the decomposition
S.5.2 Underparametrized case
Now, if , then in this case, in general the loss cannot converge to zero, because the number of equations is larger than the number of constraints. Instead, the loss converges close to the loss of the OLS estimator:
Moreover, . Also, letting , , and hence
where . Let . Then
By taking expectations in the SDE for the loss, we find
This shows that converges geometrically to the level , which is higher than the minimum OLS loss. In this case, the additional fluctuations occur because of the inherent noise in the algorithm.
Appendix S.6 Calculations for the In-Sample Risks, for Theorem 2
For in-sample risk, we have the bias-variance decomposition
Following the same logic as in the proof of Theorem 2, we see that
Appendix S.7 Proof of Lemma 6
Focusing on just for now, and noting that for , we see
Here, we used the eigendecomposition , and Lemma 5 in Ali et al. 2018. Hence,
Therefore, putting (S.5) and (S.6) together, along with Lemma 5 in Ali et al. 2018, we obtain
As , we have for each such that ,
Multiplying both sides by and rearranging yields
Now, note that is increasing on , and is decreasing on . Also, note that when , and when . Thus, . So,
which shows the claim for gradient flow. Applying Theorem 1 shows the result for ridge. Turning to in-sample risk, the exact same bounds actually follow by similar arguments, just as discussed before. ∎
Appendix S.8 Calculations for Remark 8
Finally, expanding into its power series representation and using the eigendecomposition shows that , which gives
Comparing the preceding calculations with those given in the proof of Theorem 2, we see that a key simplification occurs in (S.7), above. Here, the (relatively) complicated expression appearing in (16),
is replaced with the comparatively simpler expression in (S.7). This simplification allows the risk expression in (S.8) to hold with equality, though it is evidently less refined than the bound appearing in, e.g., (12).
Appendix S.9 Proof of Lemma 7
Letting be a singular value decomposition, we may express
Now define with domain (let ). Lemma 7 in Ali et al. 2018 shows that attains its unique maximum at , where . Moreover, it can be checked that is unimodal. This means that, for and ,
Similar reasoning shows that, for ,
When , we may simply take