Implicit Regularization or Implicit Conditioning? Exact Risk Trajectories of SGD in High Dimensions
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.
While the empirical success of SGD is often attributed to its computational efﬁciency and favorable generalization behavior, neither effect is well understood and disentangling them remains an open problem.
Even in the simple setting of convex quadratic problems, worst-case analyses give an asymptotic convergence rate for
SGD that is no better than full-batch gradient descent (GD), and the purported implicit regularization effects of SGD lack a precise explanation.
In this work, we study the dynamics of multi-pass SGD on high-dimensional convex quadratics and establish an asymptotic equivalence to a stochastic differential equation, which we call homogenized stochastic gradient descent (HSGD), whose solutions we characterize explicitly in terms of a Volterra integral equation.
These results yield precise formulas for the learning and risk trajectories, which reveal a mechanism of implicit conditioning that explains the efﬁciency of SGD relative to GD.
We also 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 machine learning and has been hailed as a major reason for deep learning’s success [11, 21].
Explanations for the effectiveness of SGD typically refer to its computational efﬁciency and to its favorable generalization properties, 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 work of Robbins and Monro [68] and extending to modern large-scale machine learning applications (see e.g. [10, 12]).
However, despite its widespread adoption and algorithmic simplicity, surprisingly little is known about how SGD performs in the types of high-dimensional optimization problems that occur in practice.
Part of the challenge in deriving robust high-level conclusions about the efﬁciency of SGD is simply that those conclusions can depend on precisely which quantities are measured and what assumptions are leveraged.
For example, in the extreme setting where the samples are one-hot vectors, running SGD on a quadratic function is actually identical to running full-batch gradient descent; 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, where each 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-relevant overparameratized or high-dimensional settings where d (cid:38) n .
Moreover, the practical success of SGD has been so remarkable in recent years that a growing body of literature has suggested that its beneﬁt to generalization extends beyond what any improved efﬁciency might reasonably afford [76, 31, 14, 72, 74].
Some of the myriad explanations for SGD’s favorable generalization properties include the local geometry of minimizers [32, 26, 82, 24], connections to approximate Bayesian inference [49], and the regularization properties of noise [75], among many others.
While some of the these perspectives are intuitive and compelling, they are often difﬁcult to rigorously establish from either an empirical or a theoretical perspective.
Empirically, simulations at large scale command signiﬁcant computational resources, and it can be challenging to push to sufﬁciently late times or sufﬁciently large batches to establish the appropriate baselines [72, 75].
Theoretically, the strongest existing results are again in the single-pass setting, for which a number of works have established excess risk bounds for quadratic problems [6, 18, 20, 81].
Much less is known in the multi-pass setting, though stability results were established by [25], and some recent works have begun examining generalization [39].
In this work, we study the dynamics of multi-pass SGD on high-dimensional convex quadratic functions and derive exact asymptotic predictions for the learning and risk trajectories.
Our analysis establishes an asymptotic equivalence to a stochastic differential equation, which we call homogenized stochastic gradient descent (HSGD), whose solutions we characterize explicitly in terms of a Volterra integral 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 SGD for 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 momentum gradient descent, underscoring the data-dependence of the conclusions.
Moreover, our results also show that SGD does not improve generalization performance, whether measured in-distribution or out-of-distribution, and therefore that SGD does not offer any form of implicit regularization in this setting.
We emphasize that our results do not rule out possible beneﬁts for non-convex problems, but they 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); 2.
Exactly characterize the asymptotic training and risk trajectories as the solutions of a deterministic Volterra integral equation (Theorem 2); 3.
Prove that the noise from SGD negatively impacts generalization performance, both in- and out-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 SGD accelerates 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 of the dynamics of multi-pass SGD (Appendix C).
Preliminaries and background
Problem setting.
We consider high-dimensional (cid:96) 2 -regularized least squares problems deﬁned by, min x ∈ R d
[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 ∈ R d is an initial vector assumed to be independent of all other randomness and having norm at most 1 .
When A is large, computing these updates can be expensive, so an unbiased estimator for the true gradient is often used, where a subset of the data points are selected uniformly at random.
We focus on the setting with batch size equal to one and without momentum, which we refer to as stochastic gradient descent (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, previous work has shown that batch sizes that are vanishingly small as a fraction of the number of samples are equivalent to the single-batch analysis, after appropriately adjusting the time by a factor of the batch size [60, Theorem 1]; similarly, we do not consider high-dimensional SGD with momentum as it degenerates to SGD [58].
See also [30].
Diffusion approximations and homogenized SGD.
A common paradigm for understanding SGD is 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 Brownian motion, and the matrix 0 (cid:22) Σ t ∈ R d × d models the noise covariance.
In many analyses, no concrete connection 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 implies that the dynamics are reversible with respect to a density proportional to e − C γ f ( x ) with C γ > 0 some constant.
Consequently, the process can escape local minima, exhibiting a trade-off between the entropy and depth of minima and thereby 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 minima of f can be analyzed, see [14, 35].
While this type of implicit entropic regularization might ultimately underlie the generalization beneﬁts of SGD for nonconvex problems, currently we lack a precise connection between a concrete SLD and a practical nonconvex learning problem.
As such, the implicit regularization effects of SGD on nonconvex losses remains a largely unsolved problem.
For convex quadratics, however, the implications of Eq.
(5) are quite clear: there is no notion of implicit regularization as the noise in SLD negatively impacts generalization performance.
Note that because 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 (see also [88]):
Lemma 1.
Suppose the objective function is f ( x ) = 12 (cid:0) (cid:107)
{+Ax − b (cid:107) 
 + δ (cid:107) x (cid:107) 2 (cid:1) .+}
Suppose ( X t : t ∈ [0 , ∞ )) is an SLD (i.e. X t solves (5) ) with (cid:107) Σ t (cid:107) op almost surely bounded by some C < ∞ .
Suppose the population risk R : R d → R is a convex function and denote x ∗ def =
lim t →∞ X gf t , then
E [ R ( X t )]
[Equation]
for all t ≥ 0 and hence,
[Equation]
If in addition R is strictly convex, and Σ t → Σ ∞ with Σ ∞ (cid:31) 0 , then the inequality is strict.
Proof.
The mean E [ X t ] , by the linearity of the gradient ∇ f , is GF.
Under the conditions given, the law of X t converges to a Gaussian variable centered at x ∗ .
Hence by Fatou’s lemma and Jensen’s inequality, the inequality follows.
We emphasize that this conclusion applies even under general distribution shifts, so long as the risk remains 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 of interest, 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 each process is not assumed to be coupled.
Instead, the most we can hope for is that statistics of the processes agree.
Speciﬁcally, we might hope that matching the noise structure of SGD with a careful choice 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 is given by homogenized SGD (HSGD), introduced simultaneously in [52, 58].
Both the empirical and population 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, and the initial condition is X {+
+} = x 0 .
Roughly, HSGD is a diffusion approximation to SGD that gains explanatory power when the dimensionality is large.
In particular, it does not require the step size γ to be small, in contrast to the usual paradigm of SLD approximations.
Note that as with other universality results, the details of the noise distribution are not relevant and only 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 of the SGD iterates concentrate around the same statistic evaluated on HSGD.
For understanding generalization 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 ∈ R n × d and target vector b ∈ R n must satisfy some quasi-random assumptions—a set of deterministic conditions on the pair ( A , b ) that are satisﬁed with high probability by natural classes of random matrix-vector pairs (see Appendix B for speciﬁcs).
We use the convention that the target and initialization vectors are bounded independent of n , (cid:107) b (cid:107) 22 ≤ C and (cid:107) x 0 (cid:107) 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 (cid:101) x
+ η for some unknown signal (cid:101) x ∈ R d and independent noise η ∈ 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 [88, 30].
• Gaussian random features with a linear ground truth [51, 3, 66, 4].
Suppose A is given by σ ( 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 polynomial or, equivalently, can be represented by
[Equation]
for some d × d symmetric matrix T , vector u ∈ R d , and scalar c ∈ R .
We further assume that (cid:107)∇ 2 R(cid:107) op ≤ C , (cid:107)∇R (0) (cid:107) 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, bounded function γ ( · ) such that γ ( t ) → γ ∈ [0 , ∞ ) and (cid:98) γ def = sup t ≥ 0 γ ( t ) < ∞ .
Main results
Our main results are analyzable (non-asymptotic) expressions for the empirical risk L and the population 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.
[Figure]
Theorem 1 (Equivalence of SGD and HSGD) .
Suppose the pair ( A , b ) ∈ R n × d × R d satisfy the quasi-random assumptions with d ε ≤ n ≤ d 1 /ε for some ε ∈ (0 , 1] .
Let the iterates x t = x sgd (cid:98) t (cid:99) be generated 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]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] The rigorous proof is given in [61]. For the rest of this paper, we will use homogenized SGD to analyze the behavior of multi-pass SGD.
See also Appendix C where we heuristically extend this to the case of streaming SGD.
While the comparison of SGD to HSGD requires relatively strong assumptions on A and b , the analysis of HSGD can be performed under weaker assumptions (no quasirandomness assumptions are needed).
It sufﬁces to suppose the problem is high dimensional in the following sense:
Assumption 2.
The empirical risk L satisﬁes tr ∇ 2 L = n and 0 (cid:22) ∇ 2 L (cid:22) nd − (cid:15) for some (cid:15) > 0 .
This corresponds to the normalization where ∇ 2 L =
A T A and each row of A is length 1 and hence tr ∇ 2 L = n .
Under Assumptions 1 and 2, the dynamics of the empirical and population risk under HSGD concentrate 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, for any 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 D.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 the population 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 is no 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 ) → ∞ (i.e. the usual
Robbins-Monro setting), then the excess population risk of SGD over GF tends to 0 .
If on the other hand γ ( t ) →
[Equation]
{+∈ (0 , 2( 1 n tr (cid:8) ( A T A ) 
+}
A T A[-) 
-]
+ δ I d (cid:9) ) − 1 , then with Ψ ∞ 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 where Ψ ∞ = 0 , there is no excess risk due to SGD and there is no need to send γ to 0 .
Moreover, if the empirical risk Ψ ∞ is small, the excess risk due to SGD is proportional to γ Ψ ∞ , and hence it is frequently orders of magnitude smaller than other potential sources of error.
Furthermore, the excess risk is affected by how similar the population and empirical risks are, in the large directions.
Ridge regularization can substantially reduce excess risk due to SGD in cases where population risk has many small eigenvalues.
In summary, either by sending γ → 0 , working in the interpolation regime, or otherwise in a regime Ψ ∞ is small, the excess risk incurred by running SGD is minimal . [Figure] i m u m e i g e n v a l u e [Figure]
3.2 Implicit conditioning of SGD
In contrast [Figure] the algorithmic advantages of SGD are substantial.
To simplify the discussion, we consider only the case of constant learning rate γ .
In this case, the kernel in (8) and (9) simpliﬁes to a convolution kernel, which has a much simpler theory.
To characterize the rates, we deﬁne λ min as the smallest non-zero eigenvalue of ∇ 2 L .
Then for generic initial conditions, (in particular almost surely if 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 matrix H .
The rate of convergence of Ψ t to Ψ ∞ can be no faster than the underlying GF, given by 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 ∇ 2 L is ﬁnite dimensional, we have that λ ∗ ≤ 2 γ ( λ min ( ∇ 2 L ) + δ ) , owing to the divergence of the integral as x approaches this value from below.
Note that in principal the Malthusian exponent can be negative, in which case SGD is divergent .
The Malthusian exponent gives the effective rate of convergence of constant learning rate SGD.
Deﬁne
[Equation]
Theorem 4 (SGD convergence rates, average-case) .
Then the rates of convergence of both the empirical and population risk are controlled by this parameter
[Equation]
Furthermore, when γ = n (tr( A T A )) − 1 , we have the rate guarantee Ξ( γ ) ≥ λ min ( ∇ 2 L )+
[Equation]
The major difference between SGD and full batch methods such as momentum gradient descent (MGD; see Appendix F.2 for deﬁnitions) is that they have different sensitivities to the Hessian spectrum of the empirical risk L ( x ) = 12 (cid:107)
Ax − b (cid:107) 2 .
Deﬁne the condition numbers
[Equation]
The ﬁrst of these is the classical condition number of the ridge problem, while the second is the averaged condition number that regulates the behavior of SGD in the high-dimensional limit.
MGD has been long established to have a rate of convergence, with proper tuning, controlled by the square root of the condition number [65], which is known to be optimal amongst ﬁrst order algorithms. [Figure]
Theorem 5 (Convergence rates for MGD) .
For isotropic random initialization x 0 or noisy b , δ > 0 , and strictly convex population risk R (cid:0) L ( x m-gd k ) − L ( x ∗ ) (cid:1) 1 /k a.s. −−−−→ k →∞ (cid:18) √ κ − 1 √ κ +(cid:19) and (cid:0) R ( x m-gd k ) − R ( x ∗ ) (cid:1) 1 /k
[Equation]
(cid:19) .
See Appendix F.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 MGD in that SGD with constant learning rate n/ tr( ∇ 2 L ) trains in an ICR-multiple of the number of epochs that MGD requires (lower is better for SGD).
Problems favor SGD when there are large outlier eigenvalues, a common feature of Hessian spectra in practice [69, 70, 1].
Indeed, if the largest eigenvalues are on the same order as the unnormalized trace , individual SGD iterates are as effective as full-batch gradient.
In contrast, when the Hessian spectrum is tightly packed, which is less common in practice but can occur after some preprocessing techniques or e.g. for uncorrelated Gaussian samples, then MGD is favored.
See Fig.
3.3 Numerical results on ICR
In Fig.
3, we illustrate how ICR affects the relative performance of full batch MGD versus single batch SGD in a synthetic least squares setting where we have tight control over the all the particulars of the problem and in a neural network setting.
We control the ICR by controlling the covariance singular value spectrum of the rows of A , which we take as Pareto distributed with exponent s for varying s > 2 .
This choice allows us to affect the ICR of the problem without changing the minimum curvature of the Hessian.
The results, comparing SGD, MGD, and GD, are given in Figure 3.
We note a few key qualitative observations.
First, even in MGD favored conﬁgurations, SGD will outperform MGD on short time scales.
[-Suppose X 
 = Y 0 . After rescaling time,-]When optimizing the hyperparameters in MGD for longtime performance, minimal curvature (which in this case is just the minimal eigenvalue of AA T ) plays a major role in the choices; being tuned for long-time performance, MGD typically performs suboptimally at initialization.
In contrast, the learning rate in SGD only depends on average curvature, and so it generally performs better at initialization on problems with a larger interval of Hessian spectra.
Second, we note that the problem setup was chosen to hold the minimum curvature roughly constant while varying s .
When s tends to 2 from above, the largest eigenvalues of AA T grows with feature dimension d , but the average and minimum eigenvalue stays bounded with feature dimension.
Hence we can send the ICR to 0 by choosing an s above 2 and increasing d (or n ). On the other hand, by sending s → ∞ , we send the covariance matrix to the identity, which tends to be momentum favored
In Fig.
4, we run SGD on a fully-connected 2-layer neural network on a subset of CIFAR-10 in order to examine the dynamics of the ICR for a non-trivial problem and to see how our insights might play out in practice.
Owing to the non-convexity of this problem, we deﬁne the ICR in terms of the Gauss-Newton approximation to the Hessian, or equivalently in terms of the Neural Tangent
Kernel [29].
By changing the activation function of the network, we can vary the initial ICR from an
SGD-favored to a momentum-favored value.
While the ICR does change over the course of training, we ﬁnd that, at least in this setting, the initial ICR can nevertheless predict the relative performance of SGD versus MGD.
Indeed, for activation functions for which the ICR remains above 1 . 0 , the training remains MGD-favored over sufﬁciently long times, and we observe that MGD with optimal parameters does converge faster than SGD.
In contrast, when the ICR remains below 1 . 0 , we ﬁnd that SGD outperforms MGD.
Conclusion.
Using a speciﬁc type of SLD (called HSGD) that matches the second-order correlations in the noise of SGD, we demonstrated that their empirical and population risks match in the high-dimensional limit.
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 no statistical beneﬁt to generalization from the noise of SGD, in overparameterized, interpolating settings little is lost compared to GD.
Moreover, when computational restrictions are imposed, SGD can be radically 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 deep learning clearly favors SGD.
This should highlight the difﬁculty in studying implicit regularization for SGD empirically: any experiment necessarily has a ﬁnite computational budget and may ﬁnd lower population risks with SGD simply via its improved conditioning.
Finally, we demonstrated limitations in using streaming SGD alone as a tool for studying generalization.
As future work, a major outstanding problem (both theoretically and empirically) is extending the analysis above to non-quadratic losses, both train and test, and especially to other high-dimensional problems not in the kernel regime.
Finally, data augmentation can naturally be considered by randomly augmenting each sample from (cid:98) D n in Eq.
Acknowledgments and Disclosure of Funding
C. Paquette’s research was supported by CIFAR AI Chair, MILA, a Discovery Grant from the Natural
Science and Engineering Council (NSERC), and the FRQNT New University Researcher’s Start-up
Program.
Research by E. Paquette was supported by a Discovery Grant from the Natural Science and
Engineering Council (NSERC).
Additional revenues related to this work: C. Paquette has part-time employment at Google Research, Brain Team, Montreal, QC.