A Continuous-Time View of Early Stopping for Least Squares
Alnur Ali, J. Zico Kolter, Ryan J. Tibshirani
INTRODUCTION
Given the sizes of modern data sets, there is a growing preference towards simple estimators that have a small computational footprint and are easy to implement. Additionally, beyond efficiency and tractability considerations, there is mounting evidence that many simple and popular estimation methods perform a kind of implicit regularization, meaning that they appear to produce estimates exhibiting a kind of regularity, even though they do not employ an explicit regularizer.
Research interest in implicit regularization is growing, but the foundations of the idea date back at least 30 years in machine learning, where early-stopped gradient descent was found to be effective in training neural networks (Morgan and Bourlard, 1989), and at least 40 years in applied mathematics, where the same idea (here known as early-stopped Landweber iterations) was found ill-posed linear inverse problems (Strand, 1974). After a wave of research on boosting with early stopping (Buhlmann and Yu, 2003; Rosset et al., 2004; Zhang and Yu, 2005; Yao et al., 2007), more recent work focuses on the regularity properties of particular algorithms for underdetermined problems in matrix factorization, regression, and classification (Gunasekar et al., 2017; Wilson et al., 2017; Gunasekar et al., 2018). More broadly, algorithmic regularization plays a key role in training deep neural networks, via batch normalization, dropout, and other techniques.
In this paper, we focus on early stopping in gradient descent, when applied specifically to least squares regression. This is a basic problem and we are of course not the only authors to consider it; there is now a large literature on this topic (see references above, and more to come when we discuss related work shortly). However, our perspective differs from existing work in a few important ways: first, we study gradient descent in continuous-time (i.e., with infinitesimal step sizes), leading to a path of iterates known as gradient flow; second, we examine the regularity properties along the entire path, not just its convergence point (as is the focus in most of the work on implicit regularization); and third, we focus on analyzing and comparing the risk of gradient flow directly, which is arguably what we care about the most, in many applications.
Our contributions in this paper are as follows.
We prove that, in finite samples, under very weak assumptions on the data model (and with no assumptions on the feature matrix ), the estimation risk of gradient flow at time is no more than 1.69 that of ridge regression at tuning parameter , for all .
We show that the same result holds for in-sample prediction risk.
We show that the same result is also true for out-of-sample prediction risk, but now in an average (Bayes) sense, with respect to a spherical prior on the underlying signal .
For Bayes risk, under optimal tuning, our results on estimation, in-sample prediction, and out-of-sample prediction risks can all be tightened. We prove that the relative risk (measured in any of these three ways) of optimally-tuned gradient flow to optimally-tuned ridge is in between 1 and 1.22.
We derive exact limiting formulae for the risk of gradient flow, in a Marchenko-Pastur asymptotic model where (the ratio of the feature dimension to sample size) converges to a positive constant. We compare these to known limiting formulae for ridge regression.
We support our theoretical results with numerical simulations that show the coupling between gradient flow and ridge can be extremely tight in practice (even tighter than suggested by theory).
Related Work.
After completing this work, we became aware of the interesting recent paper by Suggala et al. (2018), who gave deterministic bounds between gradient flow and ridge regularized estimates, for problems in which the loss function is strongly convex. Their results are very different from ours: they apply to a much wider variety of problem settings (not just least squares problems), and are driven entirely by properties associated with strong convexity; our analysis, specific to least squares regression, is much more precise, and covers the important high-dimensional case (in which the strong convexity assumption is violated).
There is also a lot of related work on theory for ridge regression. Recently, Dobriban and Wager (2018) studied ridge regression (and regularized discriminant analysis) in a Marchenko-Pastur asymptotics model, deriving limiting risk expressions, and the precise form of the limiting optimal tuning parameter. Dicker (2016) gave a similar asymptotic analysis for ridge, but under a somewhat different problem setup. Hsu et al. (2012) established finite-sample concentration bounds for ridge risk. Low-dimensional theory for ridge dates back much further, see Goldenshluger and Tsybakov (2001) and others. Lastly, we point out an interesting risk inflation result in that is vaguely related to ours: Dhillon et al. (2013) showed that risk of principal components regression is at most four times that of ridge, under a natural calibration between these two estimator paths (coupling the eigenvalue threshold for the sample covariance matrix with the ridge tuning parameter).
Outline.
Here is an outline for the rest of the paper. Section 2 covers preliminary material, on the problem and estimators to be considered. Section 3 gives basic results on gradient flow, and its relationship to ridge regression. Section 4 derives expressions for the estimation risk and prediction risk of gradient flow and ridge. Section 5 presents our main results on relative risk bounds (of gradient flow to ridge). Section 6 studies the limiting risk of gradient flow under standard Marchenko-Pastur asymptotics. Section 7 presents numerical examples that support our theoretical results, and Section 8 concludes with a short discussion.
PRELIMINARIES
Consider gradient descent applied to (1), with a constant step size , and initialized at , which repeats the iterations
for . Letting , we get a continuous-time ordinary differential equation
over time , subject to an initial condition . We call (3) the gradient flow differential equation for the least squares problem (1).
To see the connection between (2) and (3), we simply rearrange (2) to find that
and setting at time , we recognize the left-hand side above as the discrete derivative of at time , which approaches its continuous-time derivative as .
In fact, starting from the differential equation (3), we can view gradient descent (2) as one of the most basic numerical analysis techniques—the forward Euler method—for discretely approximating the solution (3).
where is a tuning parameter. The explicit ridge solution is
Though apparently unrelated, the ridge regression solution path and gradient flow path share striking similarities, and their relationship is our central focus.
2 The Exact Gradient Flow Solution Path
Thanks to our focus on least squares, the gradient flow differential equation in (3) is a rather special one: it is a continuous-time linear dynamical system, and has a well-known exact solution.
Fix a response and predictor matrix . Then the gradient flow problem (3), subject to , admits the exact solution
for all . Here is the Moore-Penrose generalized inverse of a matrix , and is the matrix exponential.
This can be verified by differentiating (6) and using basic properties of the matrix exponential. ∎
3 Discretization Error
In what follows, we will focus on (continuous-time) gradient flow rather than (discrete-time) gradient descent. Standard results from numerical analysis give uniform bounds between discretizations like the forward Euler method (gradient descent) and the differential equation path (gradient flow). In particular, the next result is a direct application of Theorem 212A in Butcher (2016).
For least squares, consider gradient descent (2) initialized at , and gradient flow (6), subject to . For any step size where is the largest eigenvalue of , and any ,
The results to come can therefore be translated to the discrete-time setting, by taking a small enough and invoking Lemma 2, but we omit details for brevity.
BASIC COMPARISONS
To compare the ridge (5) and gradient flow (6) paths, it helps to rewrite them in terms of the singular value decomposition of . Let be a singular value decomposition, so that is an eigendecomposition. Then straightforward algebra brings (5), (6), on the scale of fitted values, to
2 Underlying Regularization Problems
Given our general interest in the connections between gradient descent and ridge regression, it is natural to wonder if gradient descent iterates can also be expressed as solutions to a sequence of regularized least squares problems. The following two simple lemmas certify that this is in fact the case, in both discrete- and continuous-time; their proofs may be found in the supplement.
Fix , and let be an eigendecomposition. Assume that we initialize , and we take the step size in gradient descent to satisfy , with denoting the largest eigenvalue of . Then, for each , the iterate from step in gradient descent (2) uniquely solves the optimization problem
where .
Fix , and let be an eigendecomposition. Under the initial condition , for all , the solution of the gradient flow problem (3) uniquely solves the optimization problem
The optimization problems that underlie gradient descent and gradient flow, in Lemmas 3 and 4, respectively, are both quadratically regularized least squares problems. In agreement with the intuition from the last subsection, we see that in both problems the regularizers penalize the lower-variance directions of more strongly, and this is relaxed as or grow. The proof of the continuous-time is nearly immediate from (8); the proof of the discrete-time result requires a bit more work. To see the link between the two results, set , and note that as :
MEASURES OF RISK
For an estimator (i.e., measurable function of ), we define its estimation risk (or simply, risk) as
Next we give expressions for the risk and Bayes risk of gradient flow; the derivations are straightforward and found in the supplement. We denote by , and , the eigenvalues and eigenvectors, respectively, of .
Under the data model (9), for any , the risk of the gradient flow estimator (6) is
and under the prior (10), the Bayes risk is
where . Here and henceforth, we take by convention when .
Compare (11) to the risk of ridge regression,
and compare (12) to the Bayes risk of ridge,
where . These ridge results follow from standard calculations, found in many other papers; for completeness, we give details in the supplement.
and observe that this is clearly minimized at .
2 Prediction Risk
We now define two predictive notions of risk. Let
For space reasons, in the remainder, we will focus on out-of-sample prediction risk, and defer detailed discussion of in-sample prediction risk to the supplement. The next lemma, proved in the supplement, gives expressions for the prediction risk and Bayes prediction risk of gradient flow. We denote .
Under (9), (15), the prediction risk of the gradient flow estimator (6) is
and under (10), the Bayes prediction risk is
Compare (16) and (17) to the prediction risk and Bayes prediction risk of ridge, respectively,
These ridge results are standard, and details are given in the supplement.
RELATIVE RISK BOUNDS
For all , we have (a) and (b) .
Fact (a) can by shown via Taylor series and (b) by numerically maximizing . ∎
A bound on the relative risk of gradient flow to ridge, under the calibration , follows immediately.
The inequality in part (a) holds for the Bayes risk with respect to any prior on .
The results in parts (a), (b) also hold for in-sample prediction risk.
For part (a), set and compare the th summand in (11), call it , to that in (13), call it . Then
where in the second line, we used Lemma 7. Summing over gives the desired result.
Part (b) follows by taking an expectation on each side of the inequality in part (a). Part (c) follows similarly, with details given in the supplement. ∎
For any , gradient flow is in fact a unique Bayes estimator, corresponding to a normal likelihood in (9) and normal prior , where is as in Lemma 4. It is therefore admissible. This means the result in part (a) in the theorem (and part (b), for the same reason) cannot be true for any universal constant strictly less than 1.
2 Relative Prediction Risk
We extend the two simple inequalities in Lemma 7 to matrix exponentials. We use to denote the Loewner ordering on positive semidefinite matrices, i.e., we use to mean that is positive semidefinite.
For all , we have (a) and (b) .
All matrices in question are simultaneously diagonalizable, so the claims reduce to ones about eigenvalues, i.e., reduce to checking that and , for , and these follow by manipulating the facts in Lemma 7. ∎
With just a bit more work, we can bound the relative Bayes prediction risk of gradient flow to ridge, again under the calibration .
Consider the matrices inside the traces in (17) and (19). Applying Lemma 8, we have
The Bayes perspective here is critical; the proof breaks down for prediction risk, at an arbitrary fixed , and it is not clear to us whether the result is true for prediction risk in general.
3 Relative Risks at Optima
We present one more helpful inequality, and defer its proof to the supplement (it is more technical than the proofs of Lemmas 7 and 8, but still straightforward).
For all , it holds that .
We now have the following result, on the relative Bayes risk (and Bayes prediction risk), of gradient descent to ridge regression, when both are optimally tuned.
Consider the data model (9), prior (10), and (out-of-sample) feature distribution (15).
The same result as in part (a) holds for both in-sample and out-of-sample prediction risk.
where in the second line, we applied Lemma 9 (to the case of scalar ). Summing over gives the desired result.
Parts (b) follows similarly, with details in the supplement. ∎
ASYMPTOTIC RISK ANALYSIS
The sample size and dimension both diverge, i.e., , with .
The spectral measure of the predictor covariance converges weakly as to some limiting spectral measure .
Under the above assumptions, the seminal Marchenko-Pastur theorem describes the weak limit of the spectral measure of the sample covariance .
Assuming Assumption A1–Assumption A3, almost surely, the spectral measure of converges weakly to a law , called the empirical spectral distribution, that depends only on .
In general, a closed form for the empirical spectral distribution is not known, except in very special cases (e.g., when for all ). However, numerical methods for approximating have been proposed (see Dobriban 2015 and references therein).
2 Limiting Gradient Flow Risk
The limiting Bayes risk of gradient flow is now immediate from the representation in (12).
Assume Assumption A1–Assumption A3, as well as a data model (9) and prior (10). Then as with , for each , the Bayes risk (12) of gradient flow converges almost surely to
where , and is the empirical spectral distribution from Theorem 4.
Note that we can rewrite the Bayes risk in (12) as , where we let , . Weak convergence of to , from Theorem 4, implies for all bounded, continuous functions , which proves the result. ∎
A similar result is available for the limiting Bayes in-sample prediction risk, given in the supplement. Studying the the limiting Bayes (out-of-sample) prediction risk is much more challenging, as (17) is not simply a function of eigenvalues of . The proof of the next result, deferred to the supplement, relies on a key fact on the Laplace transform of the map , and the asymptotic limit of a certain trace functional involving , from Ledoit and Peche (2011).
where is the inverse Laplace transform of the function
and is the Stieltjes transform of (defined precisely in the supplement).
An interesting feature of the results (20), (21) is that they are asymptotically exact (no hidden constants). Analogous results for ridge (by direct arguments, and Dobriban and Wager 2018, respectively) are compared in the supplement, for space reasons.
NUMERICAL EXAMPLES
We give numerical evidence for our theoretical results: both our relative risk bounds in Section 5, and our asymptotic risk expressions in Section 6. We generated features via , for a matrix with i.i.d. entries from a distribution (with mean zero and unit variance), for three choices of : standard Gaussian, Student with 3 degrees of freedom, and Bernoulli with probability 0.5 (the last two distributions were standardized). We took to have all diagonal entries equal to 1 and all off-diagonals equal to (i.e., ), or . For the problem dimensions, we considered , and , . For both gradient flow and ridge, we used a range of 200 tuning parameters equally spaced on the log scale from to . Lastly, we set , where is the noise variance in (9) and is the prior radius in (10). For each configuration of , we computed the Bayes risk and Bayes prediction risk gradient flow and ridge, as in (12), (14), (17), (19). For , the empirical spectral distribution from Theorem 4 has a closed form, and so we computed the limiting Bayes risk for gradient flow (20) via numerical integration (and similarly for ridge, details in the supplement).
DISCUSSION
In this work, we studied gradient flow (i.e., gradient descent with infinitesimal step sizes) for least squares, and pointed out a number of connections to ridge regression. We showed that, under minimal assumptions on the data model, and using a calibration —where denotes the time parameter in gradient flow, and the tuning parameter in ridge—the risk of gradient flow is no more than 1.69 times that of ridge, for all . We also showed that the same holds for prediction risk, in an average (Bayes) sense, with respect to any spherical prior. Though we did not pursue this, it is clear that these risk couplings could be used to port risk results from the literature on ridge regression (e.g., Hsu et al. 2012; Raskutti et al. 2014; Dicker 2016; Dobriban and Wager 2018, etc.) to gradient flow.
Acknolwedgements. We thank Veeranjaneyulu Sadhanala, whose insights led us to completely revamp the main results in our paper. AA was supported by DoE CSGF no. DE-FG02-97ER25308. ZK was supported by DARPA YFA no. N66001-17-1-4036.
Supplementary Material
Appendix S.1 Proof of Lemma 3
Let be an eigendecomposition of . Then we can rewrite the gradient descent iteration (2) as
Furthermore applying the assumption that the initial point yields
with the second equality following after a short inductive argument.
Compare this to the solution of the optimization problem in Lemma 3, which is
Equating the last two displays, we see that we must have
Inverting both sides and rearranging, we get
and an application of the matrix inversion lemma shows that , so
Appendix S.2 Proof of Lemma 4
Recall that Lemma 1 gives the gradient flow solution at time , in (6). Compare this to the solution of the optimization problem in Lemma 4, which is
To equate these two, we see that we must have
i.e., writing as an eigendecomposition of ,
Inverting both sides and rearranging, we find that
Appendix S.3 Proof of Lemma 5
For fixed , and any estimator , recall the bias-variance decomposition
For the gradient flow estimator in (6), we have
In the second line, we used the fact that and are simultaneously diagonalizable, and so they commute; in the third line, we used the fact that is the projection onto the row space of , and the image of is already in the row space. Hence the bias is, abbreviating ,
where in the second line we used the fact that and are simultaneously diagonalizable, and hence commute, and also the fact that . Putting together (S.2) and (S.3) proves the result in (11).
When follows the prior in (10), the variance (S.3) remains unchanged. The expectation of the bias (S.2) (over ) is
which leads to (12), after the appropriate definition of .
Appendix S.4 Derivation of (13), (14)
As in the calculations in the last section, consider for the ridge estimator in (5),
where we have again abbreviated . The bias is thus
the second equality following after adding and subtracting to the second term in parentheses, and expanding. For the variance, we compute
the second equality following by noting that and are simultaneously diagonalizable, and therefore commute. Putting together (S.5) and (S.6) proves the result in (13). The Bayes result (14) follows by taking an expectation of the bias (S.5) (over ), just as in the last section for gradient flow.
Appendix S.5 Proof of Lemma 6
First, observe that for fixed , and any estimator ,
where . The bias-variance decomposition for out-of-sample prediction risk is hence
For gradient flow, we can compute the bias, from (S.1),
Putting together (S.7) and (S.8) proves the result in (16). The Bayes result (17) follows by taking an expectation over the bias, as argued previously.
We note that the in-sample prediction risk is given by the same formulae except with replaced by , which leads to
Appendix S.6 Derivation of (18), (19)
For ridge, we can compute the bias, from (S.4),
Putting together (S.11) and (S.12) proves (18), and the Bayes result (19) follows by taking an expectation over the bias, as argued previously.
Again, we note that the in-sample prediction risk expressions is given by replacing replaced by , yielding
Appendix S.7 Proof of Theorem 1, Part (c)
As we can see from comparing (11), (13) to (S.9), (S.13), the only difference in the latter in-sample prediction risk expressions is that each summand has been multiplied by . Therefore the exact same relative bounds apply termwise, i.e., the arguments for part (a) apply here. The Bayes result again follows just by taking expectations.
Appendix S.8 Proof of Lemma 9
As in the proof of Lemma 8, because all matrices here are simultaneously diagonalizable, the claim reduces to one about eigenvalues, and it suffices to check that for all . Completing the square and simplifying,
Now observe that, for any constant ,
the last line holding because the basic inequality implies that , for . We see that for the above line to hold, we may take
which has been computed by numerical maximization, i.e., we find that the desired inequality (S.15) holds with .
Appendix S.9 Proof of Theorem 3, Part (b)
The lower bounds for the in-sample and out-of-sample prediction risks follow by the same arguments as in the estimation risk case (the ridge estimator here is the Bayes estimator in the case of a normal-normal likelihood-prior pair, and the risks here do not depend on the specific form of the likelihood and prior).
For the upper bounds, for in-sample prediction risk, we can see from comparing (12), (14) to (S.10), (S.14), the only difference in the latter expressions is that each summand has been multiplied by , and hence the same relative bounds apply termwise, i.e., the arguments for part (a) carry over directly here.
And for out-of-sample prediction risk, the matrix inside the trace in (17) when is
and the matrix inside the trace in (19) when is
Appendix S.10 Proof of Theorem 6
almost surely, where denotes the Stieltjes transform of the empirical spectral distribution ,
For the Bayes prediction risk of gradient flow (17), the connection is less clear. However, the Laplace transform is the key link between (17) and (S.16). In particular, defining , it is a standard fact that its Laplace transform (meaning elementwise integration) is in fact
Using linearity (and invertibility) of the Laplace transform, this means
Therefore, we have for the bias term in (17),
where in the second line we again used linearity of the (inverse) Laplace transform. In what follows, we will show that we can commute the limit as with the inverse Laplace transform in (S.20), allowing us to apply the Ledoit-Peche result (S.16), to derive an explicit form for the limiting bias. We first give a more explicit representation for the inverse Laplace transform in terms of a line integral in the complex plane
for all , we can take limits in (S.21) and apply the dominated convergence theorem, to yield that almost surely,
As for the variance term in (17), consider differentiating with respect to , to yield
with the second line following because the column space of matches that of . The fundamental theorem of calculus then implies that the variance equals
where the equality is due to inverting the Laplace transform fact (S.18), as done in (S.19) for the bias. The same arugments for the bias now carry over here, to imply
Putting together (S.22) and (S.23) completes the proof.
Appendix S.11 Supporting Lemmas
For any real matrices and , define
Appendix S.12 Asymptotics for Ridge Regression
Under the conditions of Theorem 5, for each , the Bayes risk (14) of ridge regression converges almost surely to
This is simply an application of weak convergence of to (as argued the proof of Theorem 5), and can also be found in, e.g., Chapter 3 of Tulino and Verdu (2004).
The limiting Bayes prediction risk is a more difficult calculation. It is shown in Dobriban and Wager (2018) that, under the conditions of Theorem 6, for each , the Bayes prediction risk (19) of ridge regression converges almost surely to
where is as defined in (S.16). The calculation (19) makes use of the Ledoit-Peche result (S.16), and Vitali’s theorem (to assure the convergence of the derivative of the resolvent functional in (S.16)).
It is interesting to compare the limiting Bayes prediction risks (S.25) and (21). For concreteness, we can rewrite the latter as
We see that (S.25) features and its derivative, while (S.26) features the inverse Laplace transform and its antiderivative.
In fact, a similar structure can be observed by rewriting the limiting risks (S.24) and (20). By simply expanding in the numerator in (S.24), and using the definition of the Stieltjes transform (S.17), the limiting Bayes risk of ridge becomes
By following arguments similar to the treatment of the variance term in the proof of Theorem 6, in Section S.10, the limiting Bayes risk of gradient flow becomes
where denotes the density of the empirical spectral distribution , and its Laplace transform. We see (S.27) features and its derivative, and (S.28) features and its antiderivative. But indeed , since we can (in general) view the Stieltjes transform as an iterated Laplace transform. This creates a symmetric link between (S.27), (S.28) and (S.25), (S.26), where in the former plays the role of in the latter.
Appendix S.13 Additional Numerical Results
We thus calibrate according to the square root of the quantities above (this is what is plotted on the x-axis in the left columns of all the figures). The above expressions have the following limits under the asymptotic model studied in Theorem 5:
Furthermore, we note that when , the empirical spectral distribution from Theorem 4 abbreviated as , sometimes called the Marchenko-Pastur (MP) law and has a closed form. For , its density is