Implicit Regularization or Implicit Conditioning? Exact Risk Trajectories of SGD in High Dimensions
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Stochastic gradient descent (SGD) is a pillar of modern machine learning, serving as the go-to optimization algorithm for a diverse array of problems.
Whilethe empirical success of SGD is often attributed to its computational efﬁciencyand favorable generalization behavior, neither effect is well understood and disentangling them remains an open problem.
Even in the simple setting of convexquadratic problems, worst-case analyses give an asymptotic convergence rate for
SGD that is no better than full-batch gradient descent (GD), and the purportedimplicit regularization effects of SGD lack a precise explanation.
In this work, westudy the dynamics of multi-pass SGD on high-dimensional convex quadratics andestablish an asymptotic equivalence to a stochastic differential equation, whichwe call homogenized stochastic gradient descent (HSGD), whose solutions wecharacterize explicitly in terms of a Volterra integral equation.
These results yieldprecise formulas for the learning and risk trajectories, which reveal a mechanismof implicit conditioning that explains the efﬁciency of SGD relative to GD.
Wealso prove that the noise from SGD negatively impacts generalization performance,ruling out the possibility of any type of implicit regularization in this context.
Finally, we show how to adapt the HSGD formalism to include streaming SGD,which allows us to produce an exact prediction for the excess risk of multi-pass
SGD relative to that of streaming SGD (bootstrap risk).
Introduction
Stochastic gradient descent (SGD) is the algorithm of choice for optimization in modern machinelearning and has been hailed as a major reason for deep learning’s success [9, 15].
Explanations for theeffectiveness of SGD typically refer to its computational efﬁciency and to its favorable generalizationproperties but theoretical understanding of these purported beneﬁts is far from complete.
The efﬁciency of SGD has been the subject of extensive research, dating back to the original workof Robbins and Monro [43] and extending to modern large-scale machine learning applications (seee.g.
However, despite its widespread adoption and algorithmic simplicity, surprisingly littleis known about how SGD performs in the types of high-dimensional optimization problems thatoccur in practice.
Part of the challenge in deriving robust high-level conclusions about the efﬁciencyof SGD is simply that those conclusions can depend on precisely which quantities are measured andwhat assumptions are leveraged.
For example, in the extreme setting where the samples are one-hotvectors, running SGD on a quadratic function is actually identical to running full-batch gradientdescent; as such, any statements about the two algorithms’ relative efﬁciency must be data-dependent.
Furthermore, the majority of prior analyses focus on the streaming or single-pass setting, whereeach sample is seen a single time.
While this setting is appropriate when the number of samples n is much larger than the dimensionality d , it does not adequately describe the practically-relevantoverparameratized or high-dimensional settings where d & n .
Moreover, the practical success of SGD has been so remarkable in recent years that a growing body ofliterature has suggested that its beneﬁt to generalization extends beyond what any improved efﬁciencymight reasonably afford [50, 22, 11, 46, 49].
Some of the myriad explanations for SGD’s favorablegeneralization properties include the local geometry of minimizers [23, 19, 53, 17], connections toapproximate Bayesian inference [30], and the regularization properties of noise [48], among manyothers.
While some of the these perspectives are intuitive and compelling, they are often difﬁcult torigorously establish from either an empirical or a theoretical perspective.
Empirically, simulationsat large scale command signiﬁcant computational resources, and it can be challenging to push tosufﬁciently late times or sufﬁciently large batches to establish the appropriate baselines [46, 48].
Theoretically, the strongest existing results are again in the single-pass setting, for which a numberof works have established excess risk bounds for quadratic problems [6, 12, 14, 52].
Much less isknown in the multi-pass setting, though stability results were established by [18], and some recentworks have begun examining generalization [28].
In this work, we study the dynamics of multi-pass SGD on high-dimensional convex quadraticfunctions and derive exact asymptotic predictions for the learning and risk trajectories.
Our analysisestablishes an asymptotic equivalence to a stochastic differential equation, which we call homogenizedstochastic gradient descent (HSGD), whose solutions we characterize explicitly in terms of a Volterraintegral equation.
These results allow us to deﬁne a precise data-dependent implicit-conditioning ratio(ICR) that determines whether SGD is more efﬁcient than its full-batch cousins.
The ICR favors SGDfor many practical datasets, providing some explanation for the observed superior efﬁciency of SGD;interestingly, we also highlight settings for which SGD is less efﬁcient than full-batch momentumgradient descent, underscoring the data-dependence of the conclusions.
Moreover, our results alsoshow that SGD does not improve generalization performance, whether measured in-distribution orout-of-distribution, and therefore that SGD does not offer any form of implicit regularization in thissetting.
We emphasize that our results do not rule out possible beneﬁts for non-convex problems, butthey do provide some of the ﬁrst explicit negative results in the convex quadratic case.
1.1 Contributions
Our primary contributions are to:
Establish the equivalence of quadratic statistics computed on the iterates of SGD and on a particular stochastic Langevin diffusion process called homogonized SGD (Theorem 1);
Exactly characterize the asymptotic training and risk trajectories as the solutions of a
deterministic Volterra integral equation (Theorem 2);
Prove that the noise from SGD negatively impacts generalization performance, both in- andout-of-distribution (Section 3.1), but explain why the impact is often minimal in practice;4.
Introduce the implicit-conditioning ratio that describes when and by how much SGDaccelerates convergence relative to the best full-batch methods (Section 3.2);5.
Analyze the limit of streaming SGD to show its inability to capture many salient features ofthe dynamics of multi-pass SGD (Section 4).
Preliminaries and background
Problem setting.
We consider high-dimensional ` 2 -regularized least squares problems deﬁned by,
[Equation]
where   0 is the ridge-regularization parameter.
We denote the ridgeless empirical risk as
[Equation]
On the problem (1), the steps taken by gradient decent (GD) can be written recursively as
[Equation]
where  > 0 is the momentum parameter,  k is the learning rate schedule, and x 0 2 R d is an initialvector assumed to be independent of all other randomness and having norm at most 1 .
When A islarge, computing these updates can be expensive, so an unbiased estimator for the true gradient isoften used, where a subset of the data points are selected uniformly at random.
We focus on thesetting with batch size equal to one and without momentum, which we refer to as stochastic gradientdescent (SGD), and for which the iterates can be written recursively as
[Equation]
where the i k ⇠ Unif ([ n ]) iid.
While it would also be possible to consider mini-batch SGD, previouswork has shown that batch sizes that are vanishingly small as a fraction of the number of samplesare equivalent to the single-batch analysis, after appropriately adjusting the time by a factor of thebatch size [37, Theorem 1]; similarly, we do not consider high-dimensional SGD with momentum asit degenerates to SGD [36].
See also [21].
Diffusion approximations and homogenized SGD.
A common paradigm for understanding SGDis through stochastic Langevin diffusions (SLD), i.e. solutions of equations of the form
[Equation]
where  is the step size of SGD, f is the loss function, B t is a d -dimensional standard Brownianmotion, and the matrix 0  ⌃ t 2 R d ⇥ d models the noise covariance.
In many analyses, no concreteconnection between SGD and SLD is developed, and the diffusion is merely used to build intuition.
A common example is the isotropic case ( ⌃ t / I d ), for which the Fokker-Planck equation impliesthat the dynamics are reversible with respect to a density proportional to e  f ( x ) .
Consequently, theprocess can escape local minima, exhibiting a trade-off between the entropy and depth of minima andthereby highlighting a possible mechanism of implicit regularization.
In the general anisotropic case,describing the stationary distribution is more difﬁcult; nonetheless, the local geometry near minimaof f can be analyzed, see [11, 25].
While this type of implicit entropic regularization might ultimately underlie the generalization beneﬁtsof SGD for nonconvex problems, currently we lack a precise connection between a concrete SLDand a practical nonconvex learning problem.
As such, the implicit regularization effects of SGD onnonconvex losses remains a largely unsolved problem.
For convex quadratics, however, the implications of Eq.
(5) are quite clear: there is no notion ofimplicit regularization as the noise in SLD negatively impacts generalization performance.
Note thatbecause the noise is mean zero, any SLD is centered around gradient ﬂow (GF)
X gf t , which solves
[Equation]
(6) leading to the following conclusion for generalization:
Lemma 1.
Suppose the objective function is f ( x ) =
[Equation]
Suppose ( X t : t 2[0 , 1 )) is an SLD (i.e. X t solves (5) ) with k ⌃ t k op almost surely bounded by some C < 1 .
Supposethe population risk R : R d !
R is a convex function and denote x ⇤ def = lim t !
1 X gf t , then
E [ R ( X t )] | {z } pop. risk of SLD  R ( X gf  t ) | {z } pop.
risk of gradient ﬂow for all t  0 and hence, lim inf t !
1 E [ R ( X t )] | {z } limiting pop.
risk of SLD  R ( x ⇤ ) | {z } limiting pop.
risk gradient ﬂow .
If in addition R is strictly convex, and ⌃ t !
⌃ 1 with ⌃ 1  0 , then the inequality is strict.
Proof.
The mean E [ X t ] , by the linearity of the gradient r f , is GF.
Under the conditions given, thelaw of X t converges to a Gaussian variable centered at x ⇤ .
Hence by Fatou’s lemma and Jensen’sinequality, the inequality follows.
We emphasize that this conclusion applies even under general distribution shifts, so long as the riskremains a convex function.
Still, the utility of Lemma 1 may not be immediately clear, as it pertains to SLD and we have not yet established any concrete connection between SLD and the process ofinterest, SGD.
Nor is it evident what form such a connection should take—the agreement between
SGD and an SLD cannot occur at the level of individual states since the randomness from eachprocess is not assumed to be coupled.
Instead, the most we can hope for is that statistics of theprocesses agree.
Speciﬁcally, we might hope that matching the noise structure of SGD with a carefulchoice of SLD will cause relevant statistics, like the population risk, to be equal.
It turns out that such a choice of SLD exists for convex quadratic problems in high dimensions, and isgiven by homogenized SGD (HSGD), introduced simultaneously in [33, 36].
Both the empirical andpopulation risks ( L , R resp.) of HSGD agree with the same of SGD in the high-dimensional limit(see Thm. 1).
Mathematically, HSGD is the strong solution of the stochastic differential equation:
[Equation]
where again B t is a d -dimensional standard Brownian motion,  ( t ) is the learning rate schedule, andthe initial condition is X 0 = x 0 .
Roughly, HSGD is a diffusion approximation to SGD that gainsexplanatory power when the dimensionality is large.
It does not require the step size  to be small.
Note that as with other universality results, the details of the noise distribution are not relevant andonly the second-order correlations contribute, which are carefully matched by HSGD to SGD.
The precise sense of the comparison requires us to evaluate low-dimensional statistics of the highdimensional dynamics; “low-dimensional” must be effective, in that the univariate statistics ofthe SGD iterates concentrate around the same statistic evaluated on HSGD.
For understandinggeneralization or implicit regularization properties, a important statistic is the population risk, R .
Assumptions.
For all parts of our analysis to hold, the pair ( A , b of the data matrix A
2 R n ⇥ d
andtarget vector b 2 R n must satisfy some quasi-random assumptions—a set of deterministic conditionson the pair ( A , b ) that are satisﬁed with high probability by natural classes of random matrix-vectorpairs (
see
Appendix A for speciﬁcs).
We use the convention that the target and initialization vectorsare bounded independent of n , k b k 22  C and k x 0 k 22  C , respectively.
We illustrate some examples below that we have shown to satisfy the quasi-random assumptions.
• Gaussian linear regression.
Here the rows of A are iid and drawn from a Gaussian with norm-bounded covariance ⌃ and the target b is drawn from a generative model, b =
A e x + ⌘
for some unknown signal e x 2 R d and independent noise ⌘ 2 R n .
• Subgaussian linear designs.
In the example above, we can relax the Gaussian assumption to
be of the form x ⌃ 1 / 2 for x a vector of iid centered subgaussian random variables [55, 21].
• Gaussian random features with a linear ground truth [32, 3, 42, 4].
Suppose A is givenby  ( XW ) for an iid standard Gaussian weight matrix W and Gaussian data matrix X .
Suitable assumptions on the activation function  and the covariance ⌃ of X added.
Assumption 1.
The population risk R : R d !
R is a quadratic, that is, it is a degree-2 polynomialor, equivalently, can be represented by
[Equation]
c for some d ⇥ d symmetric matrix T , vector u 2 R d , and scalar c 2 R .
We further assume thatkr 2 R k op  C , kr R (0) k 22  C , and | R (0) |  C .
A natural population risk is given by R ( x ) =
12 E
[( a · x  b ) 2 ] where ( a , b ) ⇠ D .
This distribution
D may or may not be the same as the distribution that generated the data [ A | b ] used in training.
As we work in the high-dimensional limit, we suppose that  k =  ( k/n ) for a smooth, boundedfunction  ( · ) such that  ( t ) !
2 [0 , 1 ) and b  def = sup t  0  ( t ) < 1 .
Main results
Our main results are analyzable (non-asymptotic) expressions for the empirical risk L and thepopulation risk R of SGD at any time t for the high-dimensional least squares problem (1).
To begin,we ﬁrst establish the following equivalence between SGD and HGSD.
Theorem 1 (Equivalence of SGD and HSGD) .
Suppose the pair ( A , b ) 2 R n ⇥ d ⇥ R d satisfy thequasi-random assumptions with d "  n  d 1 / " for some " 2 (0 , 1] .
Let the iterates x t = x sgd b t c begenerated from multi-pass SGD Eq.
(4) and X t be the solution of Eq.
Then for any deterministic
T > 0 and any D > 0 , there is a C > 0 such that
[Equation]
As a result, for the rest of this paper, we can use homogenized SGD to analyze the behavior ofmulti-pass SGD.
We will establish a similar conjecture for streaming SGD in Section 4.
While the comparison of SGD to HSGD requires relatively strong assumptions on A and b , theanalysis of HSGD can be performed under weaker assumptions (no quasirandomness assumptionsare needed).
It sufﬁces to suppose the problem is high dimensional in the following sense:
Assumption 2.
The empirical risk L satisﬁes tr r 2 L = n and 0  r 2 L  nd  ✏ for some ✏ > 0 .
This corresponds to the normalization where r 2 L =
A T A and each row of A is length 1 and hencetr r 2 L = n .
Under Assumptions 1 and 2, the dynamics of the empirical and population risk under HSGDconcentrate around a deterministic dynamical system driven by a Volterra integral equation:
Volterra Dynamics, Multi-pass.
The following deterministic dynamical system is the high-dimensional limit for L ( X t ) and R ( X t ) , respectively
[Equation]
where the integrated learning rate  and kernel K , for any d ⇥ d matrix P , respectively are
[Equation]
Theorem 2 (Concentration of HSGD around Volterra dynamics) .
Under Assumptions 1 and 2, forany T > 0 and for any D > 0 there exists sufﬁciently large C > 0 such that for all d >
[Equation]
where t and ⌦ t solve (8) and (9) .
We give a formal proof of the concentration result in Appendix B.1 in Theorem 11.
3.1 No implicit regularization from SGD
From (9), for convex R we observe immediately that the population risk ⌦ t is only larger than thepopulation risk of GF.
Moreover, we have an explicit formula for the excess risk due to SGD noise,
[Equation]
Note that the population risk of SGD tracks that of GF.
If GF overﬁts, SGD overﬁts as well; there isno statistical regularization due to the noise of SGD applied to empirical risk minimization (ERM).
We can further analyze the long-time behavior of SGD with exact limiting values for this excess risk.
Theorem 3 (Time inﬁnity risk values) .
If  ( t ) ! 0 as t !
but  ( t ) !
1 (i.e. the usual 188 Robbins-Monro setting), then the excess population risk of SGD over GF tends to 0 .
If on the otherhand  ( t ) !
2 (0 , 2( 1 n tr
[Equation]
A T A )
A T A +  I d )  1 , then with 1 given by the limiting empirical risk,
[Equation]
the excess risk due to SGD converges to
[Equation]
There are a few conclusions to draw directly from this.
In the interpolation regime, that is where1 = 0 , there is no excess risk due to SGD and there is no need to send  to 0 .
Moreover, if theempirical risk 1 is small, the excess risk due to SGD is proportional to  1 , and hence it isfrequently orders of magnitude smaller than other potential sources of error.
Furthermore, the excessrisk is affected by how similar the population and empirical risks are, in the large directions.
Ridgeregularization can substantially reduce excess risk due to SGD in cases where population risk hasmany small eigenvalues.
In summary, either by sending  ! 0 , working in the interpolation regime,or otherwise in a regime 1 is small, the excess risk incurred by running SGD is minimal .
3.2 Implicit conditioning of SGD
In contrast, the algorithmic advantages of SGD are substantial.
To simplify the discussion, weconsider only the case of constant learning rate  .
In this case, the kernel in (8) and (9) simpliﬁes to aconvolution kernel, which has a much simpler theory.
To characterize the rates, we deﬁne  min as thesmallest non-zero eigenvalue of r 2 L .
Then for generic initial conditions, (in particular almost surelyif X 0 is nonzero isotropic), GF has the following convergence rate
[Equation]
otherwise .
Here we use the notation that  min ( H ) and  max ( H ) are the smallest and largest eigenvalues of the 206 matrix H .
The rate of convergence of t to 1 can be no faster than the underlying GF, given by 207 the rate above.
On the other hand, for larger  the Volterra term in (8) can frustrate the convergence.
The Malthusian exponent of the convolution Volterra equation is given by
[Equation]
As r 2 L is ﬁnite dimensional, we have that  ⇤  2  (  min ( r 2 L ) +  ) , owing to the divergence ofthe integral as x approaches this value from below.
Note that in principal the Malthusian exponentcan be negative, in which case SGD is divergent .
The Malthusian exponent gives the effective rate ofconvergence of constant learning rate SGD.
Deﬁne
[Equation]
Theorem 4 (SGD convergence rates, average-case) .
Then the rates of convergence of both theempirical and population risk are controlled by this parameter
[Equation]
Furthermore, when  =  min ( r 2 L ) n 2tr( A T A ) , we have the rate guarantee ⌅ (  )   min ( r 2 L )
+ .
The major difference between SGD and full batch methods such as momentum gradient descent(M-GD; see Appendix D.2 for deﬁnitions) is that they have different sensitivities to the Hessianspectrum of the empirical risk L ( x )
= 12 k
Ax  b k 2 .
Deﬁne the condition numbers
[Equation]
and [Equation] The ﬁrst of these is the classical condition number of the ridge problem, while the second is theaveraged condition number that regulates the behavior of SGD in the high-dimensional limit.
M-GDhas been long established to have a rate of convergence, with proper tuning, controlled by the squareroot of the condition number [41], which is known to be optimal amongst ﬁrst order algorithms.
Theorem 5 (Convergence rates for M-GD) .
For isotropic random initialization x 0 or noisy b ,  > 0 ,and strictly convex population risk R
[Equation]
See Appendix D.2 for elaboration.
In light of Theorems 4 and 5, we can deﬁne the implicit-conditioning ratio as
[Equation]
which measures the efﬁciency of SGD over M-GD in that SGD with constant learning rate n/ tr( r 2 L )trains in an ICR-multiple of the number of epochs that M-GD requires (lower is better for SGD).
Problems favor SGD when there are large outlier eigenvalues, a common feature of Hessian spectrain practice [44, 45, 1].
Indeed, if the largest eigenvalues are on the same order as the unnormalizedtrace , individual SGD iterates are as effective as full-batch gradient.
In contrast, when the Hessianspectrum is tightly packed, which is less common in practice but can occur after some preprocessingtechniques or e
Streaming SGD
In this section we introduce constant learning rate (  ( t ) ⌘  ) streaming SGD.
Let ( a k , b k ) 1 k
=1 be iidsamples from a R d ⇥ R -dimensional distribution D .
Streaming SGD using data from D , which wedenote D -SGD, is the algorithm
[Equation]
This naturally describes one-pass SGD, in which data points are used only once.
If R is the expectedrisk (i.e. population risk), and if R is given by 12 E ( a · x  b ) 2 with ( a , b ) ⇠ D , then D -SGD isdirectly solving the population risk minimization.
This is an idealized situation as one does not haveaccess to inﬁnite data in practice.
Deterministic behavior of streaming risks and comparison to HSGD.
D -SGD can encompassmulti-pass SGD by letting ( a k , b k )
nk
=1 be the ﬁrst n iid samples from D and considering b
D -SGD for
[Equation]
This means that a sample from b
D n ( conditionally on the dataset) is distributed like ( a I , b I ) , where Iis an uniformly random choice of index.
To enable a comparison of streaming SGD to multi-pass SGD, we suppose for some n the matrix
[ A | b ] is n ⇥ ( d + 1) matrix whose rows are iid samples from D .
We construct the empirical risk Lfrom these n samples as in (2).
We then deﬁne the streaming loss
[Equation]
where ( a , b ) ⇠ D .
The associated homogenized SGD representation, which we call homogenized D -SGD, is
[Equation]
This naturally leads to Volterra dynamics in which r 2 L is replaced by r 2 S in Eqs.
(8)-(10) whosesolution we denote with st and ⌦ st (see Appendix B).
We prove a weak equivalence between HSGD and homogenized D -SGD in the following theorem.
Theorem 6.
Suppose X
=
Y 0 .
After rescaling time,
[Equation]
uniformly on compact sets of time a.s.
Furthermore, GF converges, X gf  t/n !
X s-gf  t , and the Volterra equations converge, t/n !
st and⌦ t/n !
⌦ st , uniformly on compact sets of time.
We expect this to hold in much greater generality as suggested by Figures 1 and 3.
This is animmediate consequence of the law of large numbers due to which 1 n L converges almost surely to S ;the applications in which we are typically interested would take time large as a function of n (thenumerical results are extremely strong, see Figures 3 and 1), and we leave a deeper mathematicalinvestigation of this point as an open question.
We also note that for t very large with n , there islikely another behavior that takes hold.
For example, when t = n , observe that X t/n will have usedapproximately (1  1 /e ) n samples, whereas Y n will have used n , and hence we expect a breakdownin the connection.
Other works also examined the expected risk of streaming including [7, 54] butnot as the limit of the dynamics of multi-pass SGD as we have done.
Features of generalization.
While there are connections between streaming and multi-pass SGD,certain behavior is only accessible in the multi-pass setting.
For example, early stopping can be auseful ingredient in avoiding overﬁtting when learning overparameterized models.
However, late-timeoverﬁtting is only observable with multi-pass SGD and does not occur for streaming SGD (see thenonmonotonicity in population risk in Fig. 3).
Similarly, comparing streaming and multi-pass SGD has been suggested as a method for analyzinggeneralization.
The bootstrap risk was introduced in [34] as an intepretable component in a decomposition of the population risk.
It is deﬁned as the excess risk of SGD for ERM L with n iid samplesfrom D when compared to D -SGD, i.e.
[Equation]
SGD
[Equation]
Thus the Volterra equations can be used to give an exact expression for the bootstrap risk.
In particular,we can predict the iteration at which the bootstrap risk becomes large as the streaming and multi-passrisks bifurcate.
This allows for quantitative prediction of the stopping time in Claim 1 of [34], acentral conjecture of their paper.
On taking time to inﬁnity, we can further evaluate the SGD bootstrap risk.
Theorem 7.
If   min { 2( 1 n tr( A T A ))  1 , 2(tr r 2 S )  1 } , then the limiting bootstrap risk is givenby
[Equation]
Here 1 is the limiting training losses given by Theorem 3, i.e.
[Equation]
Conclusion.
Using a speciﬁc type of SLD (called HSGD) that matches the second-order correlations in the noiseof SGD, we demonstrated that their empirical and population risks match in the high-dimensionallimit.
Moreover, the risks of HSGD behavior deterministically, as described by a Volterra equation.
With this connection, we investigated the beneﬁts of SGD on a convex objective.
While there is nostatistical beneﬁt to generalization from the noise of SGD, in overparameterized, interpolating settingslittle is lost compared to GD.
Moreover, when computational restrictions are imposed, SGD can beradically faster than GD because of its dependence on a different condition number of the Hessian.
We characterized this speed up using the ICR, which when calculated for datasets common in deeplearning clearly favors SGD.
This should highlight the difﬁculty in studying implicit regularizationfor SGD empirically: any experiment necessarily has a ﬁnite computational budget and may ﬁndlower population risks with SGD simply via its improved conditioning.
Finally, we demonstratedlimitations in using streaming SGD alone as a tool for studying generalization.
As future work, a proper theoretical connection between streaming and multi-pass SGD remains tobe made.
A major outstanding problem (both theoretically and empirically) is extending the analysisabove to non-quadratic losses, both train and test, and especially to other high-dimensional problemsnot in the kernel regime.
Finally, data augmentation can naturally be considered by randomlyaugmenting each sample from b
D n in Eq.