F LATTER , FASTER : SCALING MOMENTUM FOR OPTIMAL SPEEDUP OF SGD
Anonymous authors
Paper under double-blind review
A BSTRACT
Commonly used optimization algorithms often show a trade-off between good generalization and fast training times.
For instance, stochastic gradient descent (SGD) tends to have good generalization; however, adaptive gradient methods have superior training times.
Momentum can help accelerate training with SGD, but so far there has been no principled way to select the momentum hyperparameter.
Here we study training dynamics arising from the interplay between SGD with label noise and momentum in the training of overparametrized neural networks.
We ﬁnd that scaling the momentum hyperparameter 1 − β with the learning rate to the power of 2 / 3 maximally accelerates training, without sacriﬁcing generalization.
To analytically derive this result we develop an architecture-independent framework, where the main assumption is the existence of a degenerate manifold of global minimizers, as is natural in overparametrized models.
Training dynamics display the emergence of two characteristic timescales that are well-separated for generic values of the hyperparameters.
The maximum acceleration of training is reached when these two timescales meet, which in turn determines the scaling limit we propose.
We conﬁrm our scaling rule for synthetic regression problems (matrix sensing and teacher-student paradigm) and classiﬁcation for realistic datasets (ResNet-18 on CIFAR10, 6-layer MLP on FashionMNIST), suggesting the robustness of our scaling rule to variations in architectures and datasets.
I NTRODUCTION
The modern paradigm for optimization of deep neural networks has engineers working with vastly overparametrized models and training to near perfect accuracy (Zhang et al., 2017).
In this setting, a model will typically have not just isolated minima in parameter space, but a continuous set of minimizers, not all of which generalize well.
Liu et al.
(2020) demonstrate that depending on parameter initialization and hyperparameters, stochastic gradient descent (SGD) is capable of ﬁnding minima with wildly different test accuracies.
Thus, the power of a particular optimization method lies in its ability to select a minimum that generalizes amongst this vast set.
In other words, good generalization relies on the implicit bias or regularization of an optimization algorithm.
There is a signiﬁcant body of evidence that training deep nets with SGD leads to good generalization.
Intuitively, SGD appears to prefer ﬂatter minima (Keskar et al., 2017; Wu et al., 2018; Xie et al., 2020), and ﬂatter minima generalize better (Hochreiter & Schmidhuber, 1997).
More recently, a variant of SGD which introduces “algorithmic” label noise has been especially amenable to rigorous treatment.
In the overparametrized setting Blanc et al.
(2020) were able to rigorously determine that SGD with label noise converges not just to any minimum, but to those minima that lead to the smallest trace norm of the Hessian.
However, Li et al.
(2022) show that the dynamics of this regularization happen on a timescale proportional to the inverse square of the learning rate η – much slower than the time to ﬁrst converge to an interpolating solution.
Therefore we consider the setting where we remain near the local minima, which is responsible for signiﬁcant regularization after the initial convergence of the train loss (Blanc et al., 2020).
With the recent explosion in size of both models and datasets, training time has become an important consideration in addition to asymptotic generalization error.
In this context, adaptive gradient methods such as Adam (Kingma & Ba, 2015) are unilaterally preferred over variants of SGD, even though they often yield worse generalization errors in practical settings (Keskar & Socher, 2017;
Wilson et al., 2017), though extensive hyperparameter tuning (Choi et al., 2019) or scheduling (Xie et al., 2022) can potentially obviate this problem.
These two constraints motivate a careful analysis of how momentum accelerates SGD.
Classic work on acceleration methods, which we refer to generally as momentum, have found a provable beneﬁt in the deterministic setting, where gradient updates have no error.
However, rigorous guarantees have been harder to ﬁnd in the stochastic setting, and remain limited by strict conditions on the noise (Polyak, 1987; Kidambi et al., 2018) or model class and dataset structure Lee et al.
In this work, we show that there exists a scaling limit for SGD with momentum (SGDM) which provably increases the rate of convergence.
Notation.
In what follows, we denote by C n , for n = 0 , 1 , . . .
the set of functions with continuous n th derivatives.
For any function f , ∂f [ u ] and ∂ 2 f [ u, v ] will denote directional ﬁrst and second derivatives along directions deﬁned by vectors u, v ∈ R D .
We may occasionally also write ∂ 2 f [Σ] = (cid:80) Di,j =1 ∂
f
[ e i , e j ]Σ ij .
Given a submanifold Γ ⊂ R D and w ∈ Γ , we denote by T w Γ the tangent space to Γ in w , and by P L ( w ) the projector onto T w Γ (we will often omit the dependence on w and simply write P L ).
Given a matrix H ∈ R D × R D , we will denote by H (cid:62) the transpose, and H † the pseudoinverse of H .
1. H EURISTIC E XPLANATION FOR O PTIMAL M OMENTUM -B ASED S PEEDUP
Deep neural networks typically posses a manifold of parametrizations with zero training error.
Because the gradients of the loss function vanish along this manifold, the dynamics of the weights is completely frozen under gradient descent.
However, as appreciated by Wei & Schwab (2019) and
Blanc et al.
(2020), noise can generate an average drift of the weights along the manifold.
In particular, SGD noise can drive the weights to a lower-curvature region which, heuristically, explains the good generalization properties of SGD.
Separately, it is well-known that adding momentum typically leads to acceleration in training Sutskever et al.
Below, we will see that there is a nontrivial interplay between the drift induced by noise and momentum, and ﬁnd that acceleration along the valley is maximized by a particular hyperparameter choice in the limit of small learning rate.
In this section we will illustrate the main intuition leading to this prediction using heuristic arguments, and defer a more complete discussion to Sec. 3.
We model momentum SGD with label noise using the following formulation
[Equation]
where η is the learning rate, β is the (heavy-ball) momentum hyperparameter (as introduced by
Polyak (1964)), w ∈ R D denotes the weights and π ∈ R D denotes momentum (also called the auxiliary variable).
L : R D → R is the training loss, and σ : R D
→ R D × r is the noise function, whose dependence on the weights allows to model the gradient noise due to SGD and label noise.
Speciﬁcally this admits modeling phenomena such as automatic variance reduction (Liu & Belkin, 2020), and expected smoothness, satisfying only very general assumptions such as those developed by Khaled & Richt´arik (2020).
Finally, ξ k ∈ R r is sampled i.i.d.
at every timestep k from a distribution with zero mean and unit variance, and (cid:15) > 0 .
We will now present a heuristic description of the drift dynamics that is induced by the noise along a manifold of minimizers Γ = { w : L ( w ) = 0 } ⊆
R D , in the limit (cid:15) → 0 .
In practice, this limit corresponds to choosing small strength of label noise and large minibatch size.
Let us assume that the weights at initial time w 0 are already close to Γ , and π 0 = 0 .
Because of this, the gradients of L at w 0 are very small, and so only ﬂuctuations transverse to the manifold will generate systematic drifts.
Denoting by δw k = w
− w k the displacement of the weights after k timesteps, let us Taylor expand the ﬁrst equation in (1) to get π k
+
=
βπ k − ∇ 2 L ( w 0 )
[ δw k ] + (cid:15)σ ( w 0 ) ξ k .
By construction, the Hessian ∇ 2 L ( w 0 ) vanishes along the directions tangent to Γ , while in the transverse direction we have an Ornstein-Uhlenbeck (OU) process.
The number of time steps it takes to this process to relax to the stationary state, or mixing time, is τ 1 =
Θ (1 / (1 − β )) as β → 1 , which can be anticipated from the ﬁrst equation in (1) since π k +1 − π k ∼ − (1 − β )
π k (see Sec.
3.3 for a more detailed derivation).
After this time, the variance of this linearized
OU process becomes independent of the time step, and can be estimated to be (see Appendix
F): (cid:104) ( δw Tk ) (cid:62) δw Tk (cid:105) =
Θ (cid:0) (cid:15) 2 η/ (1 − β ) (cid:1) , where (cid:104)· · · (cid:105) denotes the noise average.
To keep track of the displacements in the longitudinal directions, δw L k , we need to look at the cubic order in the Taylor expansion of L ( w {+
+} + δw k ) , i.e. ∂ 2 ( ∇ L )[ δw k , δw k ] .
Let P L be the projector onto the tangent space.
The expectation value of momentum, upon applying the longitudinal projector, is P L (cid:104)
π k +1 (cid:105) =
− P L (cid:80) kj =0 β k − j (cid:104)∇ L ( w j ) (cid:105) = − 12 (cid:80) kj =
β k − j ∂
( P L ∇ L )[ (cid:104) ( δw j ) (cid:62) δw j (cid:105) ]
= − 12 1 1 − β ∂
( P L ∇ L )[ (cid:104) ( δw j ) (cid:62) δw j (cid:105) ] .
The variance of the transverse displacements therefore generates longitudinal motion, and the latter will then scale as P L (cid:104) π k +1 (cid:105)
=
Θ (cid:0) (cid:15) 2 η/ (1 − β ) 2 (cid:1) .
Using the second equation in (1), we see that P L ( (cid:104) δw k +1 (cid:105) − (cid:104) δw k (cid:105) )
=
P L (cid:104)
η π k +1 (cid:105) .
Deﬁne τ 2 to be the number of time steps it takes so that the displacement of w k along Γ becomes a ﬁnite number as we take (cid:15) → 0 ﬁrst, and η → 0 afterward.
From the above considerations, we then ﬁnd that
[Equation]
The key observation now is that, when the weights are initialized near the valley, the convergence time is controlled by the largest timescale among τ 1 and τ 2 .
While τ 2 decreases as β → 1 , τincreases.
Therefore, the optimal speedup of training is achieved when these two timescales intersect τ 1 = τ 2 , which happens for 1 − β
=
Cη 2 / 3 .
More generally, we consider a double scaling where β → 1 as η → 0 according to
[Equation]
We consider this scaling limit in this paper, and we thus ﬁnd that the scaling power γ = 2 / 3 achieves the optimal speed-up along the zero-loss manifold.
1. L IMIT D RIFT -D IFFUSION
We now describe the rationale for obtaining the limiting drift-diffusion on the zero-loss manifold, for a process of the form (1) which foreshadows the rigorous results presented in section 3.
As discussed above, the motion along the manifold is slow, as it takes Θ( (cid:15) − 2 ) time steps to have a ﬁnite amount of longitudinal drift.
We want to extract this slow longitudinal sector of the dynamics by projecting out the fast-moving components of the weights.
For stable values of the optimization hyperparameters, the noiseless ( (cid:15) 2 = 0 ) dynamics (1) will map a generic pair ( π, w ) , as k → ∞ , to (0 , w ∞ ) , where w ∞ ∈ Γ .
Deﬁne Φ : R D × D
→ R D to be this mapping, i.e. Φ( π, w ) = w ∞ .
As we now show, when (cid:15) > 0 , Φ can be used precisely to project onto the slow, noise-induced longitudinal dynamics.
Let us collectively denote x k =
( π k , w k ) and write eq.
(1) as x k +1 = x k + F ( x k ) + (cid:15) ˜
σ ( x t ) ξ t .
We can perform a Taylor expansion in (cid:15) to obtain Φ( x t +1 ) − Φ( x t ) =
∂
Φ( x t )[ (cid:15) ˜ σ ( x t ) ξ t ] + ∂
Φ( x t )[ (cid:15) ˜ σ ( x t ) ξ t , (cid:15) ˜ σ ( x t )
ξ t ] + · · · .
Therefore, denoting Y ( t = (cid:15) 2 k ) =
Φ( x k ) , the limit dynamics as (cid:15) → 0 can be well-approximated by the continuous time equation
[Equation]
where we interpret the time increment dt = (cid:15) 2 , and (cid:104) ξ 2 t (cid:105)
= 1 .
Note that until here we have not taken a small learning rate limit.
The learning rate can be ﬁnite, as far as the map Φ( π, x ) exists.
The small noise limit is sufﬁcient to allow a continuous-time description of the limit dynamics because the noise-induced drift-diffusion along the valley requires Θ( (cid:15) − 2 ) timesteps to lead to appreciable longitudinal displacements.
A similar approach to what we just described was used in Li et al.
(2022), although in our case the limit drift-diffusion is obtained in the small noise limit, rather than small learning rate.
The reason for our choice is that, since we scale β according to (2), the deterministic part of eq.
(1) becomes degenerate as we take η → 0 , in which case it would not possible to apply the mathematical framework of Katzenberger (1991) on which our results below rely.
To further simplify our analysis, particularly the statement of Theorem B.4, we will further take η → 0 after taking (cid:15) → 0 , and retain only leading order contributions in η .
The main contributions of this paper are:
We develop a general formalism to study SGD with (heavy-ball) momentum in Sec. 3, extending the framework of Li et al.
(2022) to study convergence rates and generalization with momentum.
We ﬁnd a novel scaling regime of the momentum hyperparameter 1 − β ∼ η γ , and demonstrate a qualitative change in the noise-induced training dynamics as γ is varied.
We identify a special scaling limit, 1 − β ∼ η 2 / 3 , where training achieves a maximal speedup at ﬁxed learning rate η .
In Sec. 4, we demonstrate the relevance of our theory with experiments on toy models (2layer neural networks, and matrix sensing) as well as realistic models and datasets (ResNet18 on CIFAR10).
R ELATED W ORKS
Loss Landscape in Overparametrized Networks
The geometry of the loss landscape is very hard to understand for real-world models.
Choromanska et al.
(2015) conjectured, based on empirical observations and on an idealized model, that most local minima have similar loss function values.
Subsequent literature has shown in wider generality the existence of a manifold connecting degenerate minima of the loss function, particularly in overparametrized models.
This was supported by work on mode connectivity (Freeman & Bruna, 2017; Garipov et al., 2018; Draxler et al., 2018; Kuditipudi et al., 2019), as well as on empirical observations that the loss Hessian possesses a large set of (nearly) vanishing eigenvalues (Sagun et al., 2016; 2017).
In particular, Nguyen (2019) showed that for overparametrized networks with piecewise linear activations, all global minima are connected within a unique valley.
The Implicit Regularization of SGD Wei & Schwab (2019), assuming the existence of a zeroloss valley, observed that SGD noise leads to a decrease in the trace of the Hessian.
Blanc et al.
(2020) demonstrated that SGD with label noise in the overparametrized regime induces a regularized loss that accounts for the decrease in the trace of the Hessian.
Damian et al.
(2021) extend this analysis to ﬁnite learning rate.
HaoChen et al.
(2021) study the effect of non-isotropic label noise in SGD and ﬁnd a theoretical advantage in a quadratic overparametrized model.
Wu et al.
(2022) show that only minima with small enough Hessians (in Frobenius norm) are stable under SGD.
The speciﬁc regularization induced by SGD was found in quadratic models (Pillaud-Vivien et al., 2022), 2-layer Relu networks (Blanc et al., 2020), linear models (Li et al., 2022), diagonal networks (Pesme et al., 2021).
Additionally, Kunin et al.
(2021) and Xie et al.
(2021) studied the diffusive dynamics induced by SGD both empirically and in a simple theoretical model.
Momentum in SGD and Adaptive Algorithms Momentum is a general term applied to techniques introduced to speed up gradient descent.
Popular implementations include Nesterov (Nesterov, 1983) and Heavy Ball (HB) or Polyak (Polyak, 1964).
We focus on the latter in this paper, which we refer to simply as momentum.
Momentum provably improves convergence time in the deterministic setting.
Intuitively, introducing β in 1 gives the motion in parameter space an effective “inertia” or memory, which promotes motion not strictly following the local gradient, but moving rather along the directions which persistently decrease the loss function across iterations Polyak (1964);
Sutskever et al.
Less is known rigorously when stochastic gradient updates are used.
Indeed, Polyak (1987) suggests the beneﬁts of acceleration with momentum disappear with stochastic optimization unless certain conditions are placed on the properties of the noise.
See also (Jain et al., 2018; Kidambi et al., 2018) for more discussion and background on this issue.
Nevertheless, in practice it is widely appreciated that momentum is important for convergence and generalization (Sutskever et al., 2013), and widely used in modern adaptive gradient algorithms Kingma & Ba (2015).
Some limited results have been obtained showing speedup in the mean-ﬁeld approximation (Mannelli & Urbani, 2021) and linear regression Jain et al.
Modiﬁcations to Nesterov momentum to make it more amenable to stochasticity (Liu & Belkin, 2020; Allen-Zhu, 2017), and near saddle points (Xie et al., 2022) have also been considered.
T HEORETICAL R ESULTS 3. G ENERAL SETUP
Following the line from section 1.2 In this and the following section, we will rigorously derive the limiting drift-diffusion equation for the weights on the zero-loss manifold Γ , and extract the timescale τ 2 associated to this noise-induced motion.
In Sec.
3.3 we will then compare τ 2 to the timescale τ 1 associated to the noiseless dynamics and evaluate the optimal value of γ discussed around Eq.
We will use Eq.
(1) to model momentum SGD.
As illustrated in Sec.
1.1, the drift is controlled by the second moment of ﬂuctuations, and we thus expect the drift timescale to be Θ( (cid:15) 2 ) .
We will then rescale time k = t/(cid:15) 2 , so that the motion in the units of t is O (1) as (cid:15) → 0 .
More explicitly, take (cid:15) n to be a positive sequence such that (cid:15) n → 0 as n → ∞ .
For each n we consider the stochastic process that solves Eq.
[Equation]
with
[Equation]
and where X ( t = (cid:15) 2 k ) =
( π k , w k ) , ˜ σ ( X ) =
( σ, ησ ) and F ( X ) =
(( β − 1) π − ∇ L ( w ) , η ( βπ − ∇ L ( w ))) .
(cid:98) x (cid:99) denotes the integer part of a real number x .
See Appendix B.2 for a proof of equivalence between (1) and (4).
Assumption 3.1.
The loss function L : R D
→ R is a C 3 function whose ﬁrst 3 derivatives are locally Lipschitz, σ is continuous, and Γ = { w ∈ R D : L ( w ) = 0 } is a C 2 -submanifold of R D of dimension M , with 0 ≤ M ≤ D .
Additionally, for w ∈ Γ , rank ( ∇ 2 L ( w ))
=
D − M .
Assumption 3.2.
There exists an open neighborhood U of { 0 } × Γ ⊆
R D × R D such that the gradient descent starting in U converges to a point x =
( π, w ) ∈ { 0 } × Γ .
More explicitly, for x ∈ U , let ψ ( x, 0) =
x and ψ ( x, k + 1) =
ψ ( x, k ) + F ( ψ ( x, k )) , i.e. ψ ( x, k ) is the k th iteration of x + F ( x ) .
Then Φ( x ) ≡ lim k →∞ ψ ( x, k ) exists and is in Γ .
As a consequence, Φ ∈ C 2 on U (Falconer, 1983).
3. L IMITING DRIFT - DIFFUSION IN MOMENTUM SGD
In this section we shall obtain the explicit expression for the limiting drift-diffusion.
The general framework is based on Katzenberger (1991) (reviewed in Appendix B).
Before stating the result, we will need to introduce a few objects.
Deﬁnition 3.3.
For a symmetric matrix H ∈ R D × R D , and W H = { Σ ∈ R D × R D : Σ = Σ (cid:62) , HH † Σ
= H † H Σ
= σ } , we deﬁne the operator ˜
L H : W H → W H with ˜ L H S ≡ { H, S } + 12 C − 2 η 1 − 2 γ
[[ S, H ] , H ] , with [ S, H ] = SH − HS .
It can be shown that the operator ˜
L H is invertible (see Lemma C.3).
Consider the process in Eq.
Note that, while at initialization we can have X n (0) / ∈ Γ , the solution X n ( t ) → Γ as n → ∞ , i.e. it becomes discontinuous.
This is an effect of the speedup of time introduced around Eqs.
To overcome this issue, it is convenient to introduce
Y n ( t ) ≡ X n ( t ) − ψ ( X n (0) , A n ( t ))
+ Φ( X n (0)) , so that Y n (0) ∈ Γ is initialized on the manifold.
Theorem 3.4 (Informal) .
Suppose the loss function L , the noise function σ , the manifold of minimizers Γ and the neighborhood U satisfy assumptions (3.1) and (3.2), and that X n (0) =
X (0) ∈ U .
Then, as (cid:15) n → 0 , and subsequently taking η → 0 , Y n ( t ) converges to Y ( t ) , where the latter satisﬁes the limiting drift-diffusion equation
[Equation]
where W ( t ) is a Wiener process.
A rigorous version of this theorem is given in section
B.
The ﬁrst term in Eq.
(6) induces diffusion in the longitudinal direction.
The second term is of geometrical nature, and is necessary to guarantee that Y ( t ) remains on Γ .
The second line describes the drift induced by the transverse ﬂuctuations.
Eq.
(6) resembles in form that found in Li et al.
(2022), although there are two crucial differences.
First, time has been rescaled using the strength of the noise (cid:15) , rather than the learning rate.
The different rescaling was necessary as the forcing term F in Eq.
(4) depends non-homogeneously on η , and thus the theory of Katzenberger (1991) would not be directly applied had we taken the small learning rate limit.
Second, and more crucially, the drift terms in Eq.
(6) are proportional to η 2 − 2 γ , which is a key ingredient leading to the change in hierarchy of the timescales discussed in Sec. 1.1.
One ﬁnal difference, is that the last term involves the operator ˜
L H instead of the Lyapunov operator.
For γ < 12 , ˜ L H reduces to the Lyapunov operator L H at leading order in η , with L H S ≡ { H, S } .
For γ > 12 , however, we cannot neglect the η -dependent term in ˜ L H (see discussion at the end of
Appendix C).
Corollary 3.5.
In the case of label noise, i.e. when, for w ∈ Γ , Σ = c ∇ 2 L , for some constant c > 0 , Eq.
(6) reduces to
[Equation]
where we have rescaled time back to t
= k , i.e. we performed t → t(cid:15) 2 .
3. S EPARATION OF TIMESCALES AND OPTIMAL MOMENTUM SCALING
The above results provide the estimate for the timescale τ 2 of the drift along the zero-loss valley.
As discussed in Sec. 1.1, training along the zero-loss manifold Γ is maximally accelerated if this time scale is equal to the timescales τ 1 for relaxation of off-valley perturbations.
As we take (cid:15) → 0 , this relaxation is governed by the nonzero eigenvalues of the Hessian as well as by the learning rate η and momentum β .
Therefore we expect τ 1 =
Θ( (cid:15) 0 ) , and this will be conﬁrmed by the analysis below.
It will be therefore sufﬁcient to obtain the leading order expression of τ 1 by focusing on the noiseless (cid:15)
= 0 dynamics.
Additionally, since we are interested in local relaxation, it will sufﬁce to look at the linearized dynamics around Γ .
Working in the extended phase space x k = ( π k , w k ) , and linearizing Eq.
(1) around a ﬁxed point x ∗ = (0 , w ∗ ) , with w ∗ ∈ Γ , the linearized update rule is δx k +
=
J ( x ∗ ) δx k , where δx k = x k − x ∗ and J ( x ∗ ) is the Jacobian evaluated at the ﬁxed point (with the explicit form given in Eq.
Denote by q i the eigenvector and λ i the corresponding eigenvalue of the Hessian.
We show in
Appendix E that the Jacobian is diagonalized by the eigenvectors
k i ± = (cid:0) µ i ± q
i , q i (cid:1) with eigenvalues
[Equation]
and µ i ± = βηκ i ± − 1 − ηλ i .
We proceed to study the decay rate of the different modes of the Jacobian to draw conclusions about the characteristic timescales of ﬂuctuations around the valley.
Longitudinal motion: On the valley, the Hessian will have a number of “zero modes” with λ i = 1 .
These lead to two distinct modes in the present setting with momentum, which we distinguish as pure and mixed.
The ﬁrst pure longitudinal mode is an exact zero mode which has κ i + = 1 with k i
+ = (0 , q i ) , corresponding to translations of the parameters along the manifold, and keeping π = 0 at its ﬁxed point value.
The second mode is a mixed longitudinal mode with κ i − =
β with k i − =
( − (1 − β ) / (2 βη ) q
i , q i ) .
This mode has a component of π along the valley, which must subsequently decay because the equilibrium is a single point π = 0 .
Therefore, this mode decays at the characteristic rate β for π , gleaned directly from Eq.
Transverse motion: When the w and π are perturbed along the transverse directions q i with positive λ i , the relaxation behavior exhibits a qualitative change depending on β .
Using the scaling function β = 1 − Cη γ , for small learning rate, the spectrum is purely real for γ < 1 / 2 , and comes in complex conjugate pairs for γ > 1 / 2 .
This leads to two distinct scaling behaviors for the set of timescales.
Deﬁning a positive c 1 ≤ min { λ i | λ i > 0 } , we ﬁnd: 1) For γ < 1 / 2 , transverse modes are purely decaying as (1 − Cη γ ) k ≤ | δx Tk | ≤ (1 − ( c 1 /C )
η 1 − γ ) k , with the lower bound set by the mixed longitudinal mode.
For γ > 1 / 2 , the transverse modes are oscillatory but with an envelope that decays like | δx T,envk | ≈ (1 − Cη γ ) k/ 2 .
We leave the derivation of these results to Appendix (E).
Collecting these results, we can describe the hierarchy of timescales τ 1 in the deterministic regime as a function of γ (excluding the pure longitudinal zero mode):
[Equation]
These are illustrated schematically in Fig.
1(a), where the ﬁnite timescales are shown as a function of γ .
We compare these “equilibration” timescales τ 1 , i.e. characteristic timescales associated with relaxation back to the zero-loss manifold, with the timescale τ 2 ∼ η 2( γ − 1) associated with driftdiffusion of the noise-driven motion along the zero-loss manifold Eq.
For small γ , the timescale associated with the drift-diffusion along the valley is much faster than that associated with the relaxation of the dynamics toward steady state.
Transverse and mixed longitudinal ﬂuctuations relax much faster than the motion along the valley, and produce an effective drift toward the minimizer of the implicit regularizer.
However, the timescales collide at γ = 2 / 3 , suggesting a transition to a qualitatively different transport above this value, where the transverse and the mixed longitudinal dynamics, having a long timescale, will disrupt the longitudinal drift Eq.
This leads us to propose γ = 23 as the optimal choice for SGD training.
We consistently ﬁnd evidence for such a qualitative transition in our experiments below.
In addition, we see that speedup of SGD with label noise is in fact maximal at this value where the timescales meet.
3. A SOLVABLE EXAMPLE
In this section we analyse a model that will allow us to determine, on top of the optimal exponent γ = 23 , also the prefactor C .
We will specify to a 2-layer linear MLP, which is sufﬁcient to describe the transition in the hierarchy of timescales described above, and is simple enough to exactly compute
C .
We will show in Sec.
4.1 that C depends only mildly on the activation function.
We apply a simple matching principle to determine C , by asking that the deterministic timescale τ 1 is equal to the drift-diffusion timescale τ 2 .
In the previous section, we found the critical γ = 2 / 3 by requiring these timescales have the same scaling in η .
In order to determine C , we need more details of the model architecture.
Deﬁnition 3.6 (UV model) .
We deﬁne the UV model as a 2-layer linear network parametrized by f ( x ) = 1 √ n UV x ∈ R m , where x ∈ R d , V ∈ R n × d , and U ∈ R m × n .
For d = m = 1 , we refer to this as the vector UV model (Rennie & Srebro, 2005; Saxe et al., 2014; Lewkowycz et al., 2020).
For a training dataset D = { ( x a , y a ) (cid:12)(cid:12) a = 1 , ..., P } , the dataset covariance matrix is Σ ij = 1 P (cid:80) Pa
=1 x ai x aj , and the dataset variance is µ 2 = trΣ .
For mean-squared error loss, it is possible to explicitly determine the trace of the Hessian (see Appendix D).
SGD with label noise introduces y a → y a + (cid:15)ξ t where (cid:104)
ξ 2 t (cid:105) = 1 , from which we identify σ µ,ja = P − 1 ∇ µ f j ( x a ) , where µ runs over all parameter indices, j ∈ [ m ] , and a ∈ [ P ] .
With this choice, the SGD noise covariance satisﬁes σσ T
=
P − 1 ∇ 2 L .
Equipped with this, we may use Corollary 3.5 with c =
P − 1 to determine the effective drift (presented in Appendix D).
For the vector UV model, the expression simpliﬁes to dY = − τ − 1 2 Y dt , with
[Equation]
The timescale of the fast initial phase τ − 1= ( C/ 2) η γ follows from the previous section.
Then requiring τ 1 = τ 2 implies not only γ = 2 / 3 , but
[Equation]
[Equation]
One particular feature to note here is that C will be small for overparametrized models and/or training with large datasets.
E XPERIMENTAL V ALIDATION 4. 2- LAYER MLP WITH L INEAR AND N ON -L INEAR A CTIVATIONS
The ﬁrst experiment we consider is the vector UV model analyzed in Sec. 3.4.
Our goal with this experiment is to analyze a simple model, and show quantitative agreement with theoretical expectations.
Though simple, this model shows almost all of the important features in our analysis: the 2(1 − γ ) exponent below γ = 23 , the γ exponent above 23 , and the constant C theoretically evaluated in Sec. 3.4.
To this end, we extract the timescales at different values of γ and show them in Fig.
We train on an artiﬁcially generated dataset D = { ( x a , y a ) } 5 a =1 with x ∼ N (0 , 1) and y = 0 .
We use the full dataset to train.
From Eq.(124) we know that the norm of the weights follows an approximately exponential trajectory as it approaches the widest minimum ( U = V = 0 ).
We therefore measure convergence timescale, T c , by ﬁtting an exponential ae − t/T c to the squared distance from the origin, | U | 2 + | V | 2 .
To extract the scaling of T c with γ we perform SGD label noise with learning rates η ∈ [10 − 3 , 10 − 1 ] and corresponding momentum parameters β
= 1 − Cη γ .
We ﬁt the timescale to a power-law in the learning rate T c ( η, γ ) =
T
η − α ( γ ) (see Fig. 1(b)).
Imposing that T 0 be independent of γ , as predicted by theory, we found the numerical value C ≈ 0 .
2 , which is consistent with the theoretical estimate of C = 0 .
17 from Sec.
We ﬁnd consistency with prediction across all the values of γ we simulated.
Note that, for γ > 23 the timescale estimate ﬂuctuates more which is a consequence of having a slower timescale for the transverse modes.
As discussed at the end of
Sec. 3.3, such slowness disrupts the drift motion along the manifold.
γ = 23 is clearly the optimal scaling.
We repeated the same experiments using nonlinear activations, speciﬁcally we considered tanh and
ReLU acting on the ﬁrst layer.
The timescales for tanh are shown in Fig.
1(c), and we refer the reader to Appendix A for the ReLU case.
The optimal scaling value is still γ = 23 for both tanh and
ReLU and the optimal C remains close to 0 .
4. R ES N ET 18 ON CIFAR
We now verify our predictions on a realistic problem, which will demonstrate the robustness of our analysis.
We focus on ResNet18 (He et al., 2016), speciﬁcally implemented by Liu (2021), classiﬁer trained on CIFAR10 Krizhevsky et al.
We aim to extrapolate the theory by showing optimal acceleration with our hyperparameter choice once training reaches an interpolating solution.
To this purpose, we initialize the network on the valley, obtained starting from a random weight values and training the network using full batch gradient descent without label noise and with a ﬁxed value β = 0 .
9 until it reaches perfect training accuracy.
With this initialization, we then train with
SGD and label noise for a ﬁxed number of epochs multiple times for various values the momentum hyperparameter β .
Finally, we project the weights back onto the valley before recording the ﬁnal test of the optimal momentum β a function of η . We perform a ﬁt exponent γ theory γ we also accuracy.
This last step can be viewed as noise annealing and allows us to compare the performance of training the drift phase for the different values of β .
From this procedure we extract the optimal momentum parameter β ∗ ( η ) that maximizes the best test accuracy during training as a function of the learning rate, which we can then compare with the theoretical prediction.
As shown in Fig. 2(b), 1 − β ∗ follows the power law we predicted almost exactly.
The optimal choice for speedup does not have to coincide with the optimal choice for generalization.
Strikingly, this optimal choice of scaling also leads to the best generalization in a realistic setting!
This can be easily interpreted if we assume that the more we decrease the Hessian the better our model generalizes and by applying the fact that our scaling leads to the fastest transport along the manifold.
The second important point is the value of the constant C ≈ 0 .
1 found as the coefﬁcient of the power-law ﬁt.
If we set η = 1 this corresponds to setting β ∗ = 0 .
9 which is the traditionally recommended value.
The result here, can therefore be viewed as a generalization of this common wisdom.
For more experiments, we refer the reader to Appendices A.2 and A.4.
C ONCLUSION
We studied the implicit regularization of SGD with label noise and momentum in the limit of small noise and learning rate.
We found that there is an interplay between the speedup of momentum and the limiting diffusion generated by SGD noise.
This gives rise to two characteristic timescales associated to the training dynamics, and the longest timescale essentially governs the training time.
Maximum acceleration is thus reached when these two timescales coincide, which lead us to identifying an optimal scaling of the hyperparameters.
This optimal scaling corresponds not only to faster training but also to superior generalization.
More generally, we have shown how momentum can signiﬁcantly enrich the dynamics of learning with SGD, modulating between qualitatively different phases of learning characaterized by different timescales and dynamical behavior.
It would be interesting to explore our scaling limit in statistical mechanical theories of learning to uncover further nontrivial effects on feature extraction and generalization in the phases of learning (Jacot et al., 2018; Roberts et al., 2021;
Yang et al., 2022).
For future work, it will be very interesting to generalize this result to adaptive optimization algorithms such as Adam and its variants, to use this principle to design new adaptive algorithms, and to study the interplay between the scaling we found and the hyperparameter schedule.
A E XPERIMENTAL D ESIGN A. UV M ODEL
In our experiments with the vector UV model we aim to extract how the timescale of motion scales with η for each γ , therefore we train the model over a sweep of both of these parameters.
As a reminder, the loss of the UV model is
[Equation]
where P is the dataset size and u , v are n -dimensional vectors.
We set y i = 0 for simplicity, and with label noise this becomes y i ( t ) = (cid:15) · ξ i ( t ) for i.i.d standard Gaussian distributed ξ i .
We initialize x i once (not online) as i.i.d.
random standard Gaussian variables.
For all experiments we choose (cid:15) = 12 .
We initialize with u i , v i i.i.d.
standard Gaussian distributed and keep this initialization constant over all our experiments to reduce noise associated with the speciﬁc initialization.
For each value of η, β we train with label noise SGD and momentum until | u | 2 + | v | 2 < εn with ε = 0 . 1 , thereby obtaining a time series for each of u ( t ) and v ( t ) .
We extract the timescale by ﬁtting log (cid:0)
| u | 2 + | v | 2 (cid:1) to a linear function and taking the slope.
A. M ATRIX S ENSING
We also explore speedup for a well understood problem: matrix sensing.
The goal is to ﬁnd a low-rank matrix X ∗ given the measurements along random matrices
A i : y i =
Tr A i X ∗ .
Here
X ∗ ∈ R
d × d is a matrix of rank r (Soudry et al., 2017; Li et al., 2018).
Blanc et al.
(2020) analyze the problem of matrix sensing using SGD with label noise and show that if X ∗ is symmetric with the hypothesis X = UU (cid:62) for some matrix U , then gradient descent with label noise corresponds not only to satisfying the constraints y i =
Tr A i X , but also to minimizing an implicit regularizer, the Frobenius norm of U , which eventually leads to the ground truth.
In the analogous UV matrix model (with X ∗ is an asymmetric matrix of low rank r ), we demonstrate a considerable learning speedup by adding momentum, and show that this speedup is not monotonic with increasing β ; there is a value β ∗ at which the acceleration appears optimal.
This non-monotonicity with an optimal β ∗ is observed for both the Hessian trace and the expected test error.
Assuming that in this setting we also have γ = 2 / 3 , we can extract C ∗ =
(1 − β ∗ ) /η 2 / 3 ≈ 0 .
24 P − 1 / 3 , which compares favorably to the upper bound we may extract from Appendix D of ≈ 0 .
12 P − 1 / 3 .
In the experiments with matrix sensing we aim to demonstrate the beneﬁt of momentum in a popular setting.
Matrix sensing corresponds to the following problem:
Given a target matrix X ∗ ∈ R
d
× d of low rank r (cid:28) d and measurements { y i =
Tr A i X ∗ } Pi =1 how can we reconstruct X ∗ ?
One way to solve this problem is to write our guess X = UV the product of two other matrices, and do stochastic gradient descent on them, hoping that the implicit regularization induced by this parametrization and the learning algorithm will converge to a good low rank X .
A.2. E XPERIMENTAL D ETAILS
In our experiments we study the d = 100 , r = 5 , P = 5 rd = 2500 case.
We draw ( A i ) ij ∼ N (0 , 1) as standard Gaussians and choose X ∗ by drawing ﬁrst ( X 0 ) ij ∼ N (0 , 1) and then performing
SVD and projecting onto the top r singular values by zeroing out the smaller singular values in the diagonal matrix.
We intitalize U
=
V = I d .
We perform SGD with momentum on the time dependent loss (with label noise depending on time)
[Equation]
where (cid:15) 2 = 0 .
1 , ξ i ( t ) ∼ N (0 , 1) .
We choose η = 0 .
1 for all of our experiments.
The hessian of the loss is deﬁned to be the Hessian averaged over the noise.
Equivalently we may just set ξ i ( t )
= 0 when we calculate the Hessian because averaging over the noise decouples the noise.
Similarly when we deﬁne the expected test loss we deﬁne it as an average over all A i setting ξ i ( t )
=
in order to decouple the noise.
Averaging over ξ i ( t ) and A i would simply lead to an additional term (cid:104) ξ i ( t ) 2 (cid:105) which would simply contribute a constant.
We remove this constant for clarity.
As a result, the expected test that we plot is proportional to the squared Frobenius norm of the difference between the model UV and the target X ∗ ,
[Equation]
It is also interesting to note that we observe epoch-wise double descent Nakkiran et al.
(2020) in this problem.
In particular, we observe that the peak in the test error can be controlled by the momentum hyperparameter, and becomes especially pronounced for β → 1 .
A. R ES N ET 18 ON CIFAR
We train our model in three steps: full batch without label noise until 100% test accuracy, SGD with label noise and momentum, and then a ﬁnal projection onto the interpolating manifold.
The model we use is the ResNet18 and we train on the CIFAR10 training set.
The ﬁrst step is full batch gradient descent on the full CIFAR10 training set of 50,000 samples.
We train with a learning rate η = 0 .
1 and momentum β = 0 .
9 and a learning rate schedule with linearly increases from 0 to η over 600 epochs, after which it undergoes a cosine learning rate schedule for 1400 more epochs stopping the ﬁrst time the network reaches 100% test accuracy which happened on epoch 1119 in our run.
This model is saved and the same one is used for all future runs.
The loss function we use is cross cross entropy loss.
Because we will choose a label noise level of p = 0 .
2 which corresponds to a uniformly wrong label 20% of the time, during this phase of training we train with the expected loss over this randomness.
Notice that this loss is actually linear in the labels so taking the expectation is easy.
The second step involves starting from the initialization in step 1 and training with a different learning rate and momentum parameter.
In this step we choose the same level of label noise p =
.but
take it to be stochastic.
Additionally we use SGD instead of gradient descent with a batch size of 512.
This necessitates decreasing the learning rate because noise is greatly increased as demonstrated in the main text.
In this step we train for a ﬁxed 200 epochs for any learning rate momentum combination.
We only compare runs with the same learning rate value.
We show an example of the test accuracy as we train in phase 2 for η = 0 .
001 in ﬁgure 4.
Notice that for both too-small and too-large momentum values that the convergence to a good test accuracy value is slower.
The initial transient with the decreased test accuracy happens as we start on the valley and adding noise coupled with momentum causes the weights to approach their equilibrium distribution about the valley.
For wider distributions the network is farther from the optimal point on the valley.
As training proceeds we see that the test accuracy actually increases over the baseline as the hessian decreases and the generalization capacity of the network increases.
This happens most quickly for the momentum which matches our scaling law.
The ﬁnal step is a projection onto the zero loss manifold.
This step is necessary because the total width of the distribution around the zero loss manifold scales with 1 1 − β , and this will distort the results systematically at larger momentum, making them look worse than they are.
We perform this projection to correct for this effect and put all momenta on an equal footing.
This projection is done by training on the full batch with Adam and a learning rate of η = 0 .
in order to accelerate training.
We do not expect any signiﬁcant systematic distortion by using Adam for the projection instead of gradient descent.
To determine the optimal value of β we sweep several values of β and observe the test accuracy after the previously described procedure.
To get a more precise value of β instead of simply selecting the one with the highest test accuracy we ﬁt the accuracy A ( β ) to
[Equation]
to the parameters a max , a 1 , a 2 , and β ∗ , thereby extracting β ∗ for each η .
A. MLP ON F ASHION MNIST
We perform a similar experiment as in section A.3 but with different model and dataset: a 6-layer
MLP trained on FashionMNIST (Xiao et al., 2017) as our dataset.
We perform the experiment with a 6 layer MLP with Relu activation after the ﬁrst 4 layers, tanh activation after the ﬁfth layer, and a mapping to logits.
We use cross entropy loss label noise p = training initialized zero manifold momentum
After this we sweep η ∈ [10 − 5 , 10 − 4 ] and β ∈ [ . 95 , 1 − 2 η ] and train for 600 epochs with label noise and 400 without label noise.
This allows us to obtain a test accuracy as a function of η and β and therefore we can obtain the best momentum hyperparameter, β ∗ , as a function of η as in A.3.
We extract the scaling exponent by doing a linear ﬁt between log(1 − β ∗ ) and log η .
The scaling analysis is shown in Fig.
5, and shows that the exponent is consistent with our theory.
B R EVIEW OF RELEVANT RESULTS FROM K ATZENBERGER (1991)
In this Appendix we summarize the relevant conditions and theorems from Katzenberger (1991) that we use to prove our result on the limiting drift-diffusion.
We refer to Katzenberger (1991) for part of the deﬁnitions and conditions cited throughout the below.
In what follows, (Ω n , F n , {F nt } t ≥ 0 , P ) will denote a ﬁltered probability space, Z n an R r -valued cadlag {F nt } -semimartingale with
Z n (0) = 0 , A n a real-valued cadlag {F nt } -adapted nondecreasing process with A n (0) = 0 , and ˜
σ : U → R D × r a continuous function, where U is a neighborhood of { 0 }× Γ as deﬁned in the main text.
Also, X n is an R D valued cadlag {F nt } -semimartingale satisfying
[Equation]
[Equation]
for all t ≤ λ n ( K ) and all compact K ⊂ U , where
[Equation]
[Equation]
be the stopping time of X n ( t ) to leave ˚ K , the interior of K .
For cadlag real-valued seimimartingales
X, Y let [ X, Y ]( t ) be deﬁned as the limit of sums
[Equation] [Equation]
where
=
t 0 < t 1 < · · · < t n = t and the limit is in probability as the mesh size goes to zero.
If
X is an R D -valued semimartingale, we write
[Equation]
Condition B.1.
For every T > (cid:15) > 0 and compact K ⊂ U
[Equation]
as n → ∞ where the inﬁmum of the empty set is taken to be ∞ .
Condition B.2.
For every compact K ⊂ U { Z λ n ( K ) n } satisﬁes the following: For n ≥ 1 let Z n be a {F nt }− semimartingale with sample paths in D R d [0 , ∞ ) .
Assume that for some δ > 0 allowing δ = ∞ and every n ≥ 1 , there exist stopping times { τ kn | k ≥ 1 } and a decomposition of Z n −
J δ ( Z n ) into a local martingale M n plus a ﬁnite variation process F n such that P [ τ kn ≤ k ] ≤ 1 /k , { [ M n ]( t ∧ τ kn ) + T t ∧ τ kn ( F n ) | n ≥ 1 } is uniformly integrable for every t ≥ 0 and k ≥ 1 and
[Equation]
for every (cid:15) > 0 and T > 0 .
Also as n → ∞ and for any T >
[Equation] [Equation]
Condition B.3.
The process
[Equation] [Equation]
exists, is an {F nt }− semimartingale, and for every compact K ⊂ U , the sequence { ¯
Z λ n ( K ) n } is relatively compact and satisﬁes Condition 4.1 in Katzenberger (1991).
Theorem B.4 (Theorem 7.3 in Katzenberger (1991)) .
Assume that Γ is C 2 and for every y ∈ Γ , the matrix ∂F ( y ) has D − M eigenvalues in D (1) .
Assume (B.1),(B.2) and (B.3) hold, Φ is C 2 (or F is LC 2 ) and X n (0)
⇒ X (0)
∈ U .
Let
[Equation] [Equation]
and, for a compact K ⊂ U , let
[Equation] [Equation]
Then for every compact K ⊂ U , the sequence { ( Y µ n ( K ) n , Z µ n ( K ) n , µ n ( K )) } is relatively compact in D R 2 D × r [0 , ∞ ) ×
[0 , ∞ ] (see Katzenberger (1991) for details about the topology).
If ( Y, Z, µ ) is a limit of this sequence then ( Y, Z ) is a continuous semimartingale, Y ( t )
∈ Γ for every t almost surely, µ ≥ inf { t ≥ 0 | Y ( t ) / ∈ ˚ K } almost surely, and
[Equation]
[Equation]
B. A PPLYING T HEOREM B.
Recall the equations of motion of stochastic gradient descent
[Equation]
[Equation]
where σ ( w ) ∈
R D × r is the noise function evaluated at w ∈ R D , and ξ k ∈ R r is a noise vector drawn i.i.d.
at every timestep k with zero mean and unit variance.
We now show that this equation satisﬁes all the properties required by Theorem B.4.
The manifold Γ is the ﬁxed point manifold of (non-stochastic) gradient descent.
{ 0 } × Γ is a C 2 manifold because Γ is C 2 , which follows from assumption 3.1.
The ﬂow F ( w, π ) =
( η ( βπ − ∇ L ( w )) , βπ − ∇ L ( w )) .
As shown in Appendix E, dF has exactly M zero eigenvalues on Γ ∩ K .
F inherits the differentiable and locally Lipschitz properties from ∇ L , and therefore satisﬁes the conditions of B.4.
Next, notice that the noise function ˜ σ : R 2 D
→ R 2 D × r is continuous because σ is.
Now we deﬁne A n and Z n (as in the main text) so that X n reproduces the dynamics in equation (26), except with a new time parameter t = k(cid:15) 2 n .
[Equation]
So that, with these choices, Eq.
(15) precisely corresponds to (26), up to the rescaling t =
k(cid:15) 2 n .
Now we show that A n , Z n satisfy the conditions of B.4.
Clearly A n (0) =
Z n (0) by deﬁnition.
Then
[Equation] [Equation]
when we take (cid:15) n → 0 , thus recoverying condition B.1.
By deﬁnition Z n is a martingale.
Notice also by the deﬁnition of Z n , because ξ k is i.i.d.
with variance 1, that Z n ( t ) has variance
A n ( t ) (cid:15) 2 n ≤ t which is uniformly bounded and hence Z n ( t ) is uniformly integrable for stopping times τ kn = 2 k > k .
Also note that (cid:12)(cid:12) ∆ Z n ( k(cid:15) 2 n ) (cid:12)(cid:12) =
| (cid:15) n ξ k | which goes to zero in probability as (cid:15) becomes small because ξ k has bounded variance, and ∆ Z n ( t ) is zero otherwise.
This shows that Z n satisﬁes condition B.2.
Because Z n and A n are discontinuous at the same time we automatically satisfy condition B.3 as pointed out by Katzenberger.
This shows that we satisfy the conditions of Theorem B.4, therefore we have the following
Lemma B.5.
The SGD equations formulated as in (26) satisfy all the conditions of Theorem B.4.
B.
E QUIVALENCE BETWEEN EQ . (1) AND EQ . (4)
Eq.
(4) is the rewriting of eq.
(1) in the form presented in Katzenberger (1991), on which our theory is based.
We show the equivalence below, and include the deﬁnitions of A and Z to keep this section self-contained.
The statement is that Eq.
(1), i.e.
[Equation]
can be rewritten in the form of eq.
[Equation]
where X n ( t ) =
( π k , w k ) , with the correspondence t =
k(cid:15) 2 n , where k denotes the SGD time step.
Also, ˜ σ ( X ) =
( σ, ησ ) and F ( X ) =
(( β − 1) π − ∇ L ( w ) , η ( βπ − ∇ L ( w ))) .
Before we begin the proof, we will need the following fact (see Sec.
2 of Katzenberger (1991)):
[Equation]
where f and g are c`adl`ag functions, in particular they are right-continuous and have left limits everywhere.
The integral above is done with respect to the measure dg , the differential of a function g .
The sum is taken over all r ∈
( s, t ] where g is discontinuous and the notation ∆ g ( r )
=
g ( r ) − g ( r − ) indicates the discontinuity of g at r , where g ( r − ) =
lim u
→ r − g ( u ) indicates the left limit of g at r .
Finally g c denotes the continuous part of g
[Equation]
We now show by induction that X n ( t = k(cid:15) 2 n ) solves the ﬁrst equation above.
Note that
[Equation]
For brevity we will drop the subscript n and let k(cid:15) 2 = t .
Consider
[Equation]
where in the last step we use eq.
(32) as A, Z, ˜ σ and F are c`adl`ag, and that dA c = dZ c = 0 .
The sums are taken over all s ∈ ( k(cid:15) 2 , ( k +1) (cid:15) 2 ] where Z ( s ) and A ( s ) are discontinuous.
The only point of discontinuity of A and Z in this interval is at s = t + (cid:15) 2 , so
[Equation]
because the jumps of Z and A at s = t + (cid:15) 2 are (cid:15)ξ k +1 and 1 , respectively.
Now we must determine the left limit of X at t + (cid:15) 2 .
Notice that for 0 < δ < (cid:15)
[Equation]
because A and Z are continuous on ( t, t + δ 2 ] .
Hence the left limit of X (( t + (cid:15) 2 ) − ) =
X ( t ) .
Putting these equations together
we ﬁnd that
[Equation]
which, using the deﬁnition of X ( k(cid:15) 2 )
=
( π k , w k ) , is eq.
(29), thus proving equivalence.
C E XPLICIT EXPRESSION OF LIMITING DIFFUSION IN MOMENTUM SGD
In this section we provide the proof of Theorem 3.4.
Recall that Φ satisﬁes
[Equation]
where x = ( π, w ) , and F given in (132).
To obtain the explicit expression of the limiting driftdiffusion, according to Theorem (B.4), and keeping into account Assumption (3.2), we need to determine Φ up to its second derivatives.
To this aim, we shall expand Eq.
(42) up to second order in the series expansion in F .
In components, this reads:
[Equation]
subject to the boundary condition
[Equation]
Here, g i =
∂ i L .
In Eq.
(43), we already substituted β = 1 − Cη γ .
In what follows, we shall ﬁnd Φ to leading order in η as η → 0 .
We will solve the above problem by performing a series expansion in Φ around a point ¯ w ∈ Γ and π =
up to second order:
[Equation]
For example, we have
[Equation] [Equation]
More precisely, we regard this as an expansion in powers of δw and π : Φ 00 is zeroth order, Φ 10 , Φare ﬁrst order, and the remaining terms are second order.
We will occasionally write explicitly the index of Φ , e.g. Φ k,ij 02 = ∂
Φ k ∂w i ∂w j .
It will be useful to introduce the longitudinal projector onto Γ ,
P L ( w ) :
R D → T w Γ , deﬁned such that for any vector v ∈ T w Γ : P L v = v .
The transverse projector is then P T
= Id − P L .
We will also decompose various tensors using these projectors, e.g.
[Equation]
[Equation]
with Φ ij 11 LT = Φ kl 11 P ikL P jlT .
Note that the Hessian H = ∇
g ∈ R D
× D satisﬁes HH †
= P T , where
H † denotes the pseudoinverse.
At zeroth order, we obviously have Φ 00 ( w ) = w .
Lemma C.1.
The ﬁrst order terms in the series expansion (45) are given by
[Equation]
Proof.
Suppose ˆ w ( s ) ∈
R D is a curve lying on Γ .
Then due to the boundary condition (44), ∂ s Φ
[Equation]
π = 0 , ˆ
[Equation] [Equation]
ˆ w i =
∂ s ˆ w i .
This means that Φ 01 L = P L .
Now from (43),
[Equation]
This condition should hold for any π and δw , therefore we arrive at
[Equation] [Equation]
and
[Equation]
Decomposing into longitudinal and transverse components, and noting that the Hessian satisﬁes
P L H = 0 , the above equation becomes
[Equation]
which together with (50) and the above discussion gives
[Equation]
which concludes the ﬁrst order analysis.
We now need a Lemma, which requires Deﬁnition (3.3) in the main text.
We report it here for convenience:
Deﬁnition C.2.
For a symmetric matrix H ∈ R D × R D , and W H = { Σ ∈ R D × R D : Σ = Σ (cid:62) , HH † Σ
= H † H Σ
= σ } , we deﬁne the operator ˜
L H : W H → W H with ˜ L H S ≡ { H, S } + 12 C − 2 η 1 − 2 γ
[[ S, H ] , H ] , with [ S, H ] = SH − HS .
Lemma C.3.
The inverse of the operator ˜
L H is unique.
Proof.
Let us go to a basis where H is diagonal, with eigenvalues
λ i .
In components, the equation ˜
L H S =
M reads
[Equation]
54) which has a unique solution, with
[Equation]
Lemma C.4.
The second order terms in the series expansion (45) are given by
[Equation]
where
[Equation]
Proof.
Consider a path ˆ w ( s ) lying on Γ .
From P T H
=
H we have
[Equation]
and thus, using HH † = P T , we ﬁnd
[Equation] [Equation]
where the second equation is obtained from the ﬁrst by taking the transpose.
Putting the last two relations together, we ﬁnd
[Equation] [Equation]
From (48) we can then write
[Equation]
66) where ∂
Li = P ijL
∂ j .
This leads to
[Equation] [Equation]
Also note that Φ j,ik 02 LT = Φ j,ki 02 TL , so the only component still to be determined in Φ 02 is Φ 02 TT .
The next step is to expand Eq.
(43) to second order:
[Equation] [Equation]
where A ( ij )
= 1 2 ( A ij + A ji ) denotes the symmetric part.
Neglecting various terms that are subleading at small η , gives
[Equation]
The ﬁrst line immediately gives
[Equation]
Taking the second line of (70) and projecting onto the longitudinal part the index k gives
[Equation]
and using (67),(68),
[Equation]
Projecting the second line of (70) on the transverse part of the index
k we have
[Equation]
Using (71) and keeping (68) into account
[Equation]
Eq.
(76) becomes
[Equation]
neglecting subleading terms in η ,
[Equation]
Using (75) in (78),
[Equation]
79) and simplifying,
[Equation]
which determines Φ 11 LT in close form.
Indeed
[Equation]
the above has the form, in matrix notation
[Equation]
and can be immediately inverted to solve for Φ 11 LT .
The only two undetermined components now are Φ 11 TT and Φ 02 TT .
One condition is obtained from (77) which gives, keeping (71) into account,
[Equation] [Equation]
and thus
[Equation]
Further taking symmetric and antisymmetric part in ik of the above gives
[Equation]
[Equation] [Equation]
The other condition comes from the third line of (70) which, using (48) and (71), and neglecting subleading terms in η
[Equation]
gives
[Equation]
Projecting the k and l indices on the longitudinal part
[Equation]
gives
[Equation]
, which is an identity that can be checked from (63).
Projecting k on the longitudinal part and l on the transverse part gives
[Equation]
which is implied by (75), indeed
[Equation]
therefore (87) does not give a new condition.
The only new condition comes from projecting the k and l indices of (86) on the transverse direction, giving
[Equation]
Plugging in (83),
[Equation]
Neglecting subleading terms in η this can be rewritten as
[Equation]
where we recall that the last term is symmetric in k and l .
In matrix notation this reads
[Equation]
where
[Equation]
is the longitudinal derivative of the Hessian, projected on the transverse directions.
Decomposing Φ j 11 TT into symmetric and antisymmetric parts Φ j 11 TT = S ( j ) + A ( j ) , Eq.
(92) reads (suppressing the index j )
[Equation]
The ﬁrst term is subleading in η , therefore we have
[Equation]
This equation together with (85), which in matrix form reads
[Equation]
[Equation]
determine S and A , and thus Φ 11 TT .
Note that M j has only a longitudinal component in the index j , therefore the transverse parts of S and A vanish, i.e.
[Equation]
Eq.
(97) is natural as the slow degrees of freedom are the longitudinal coordinates.
Note that these two equations admit a unique solution.
Indeed, solving (96) for A gives
[Equation]
Plugging in (95), we ﬁnd
[Equation]
where ˜ L H S ≡ { H, S } + 12 C − 2 η 1 − 2 γ
[[ S, H ] , H ] is introduced in deﬁnition C.2.
By Lemma C.3, (99) admits a unique solution.
Then
[Equation] [Equation]
The other components are (see eqs.
[Equation] [Equation]
The only contributing term at leading order in η to the limiting diffusion equation is (106).
Splitting it into longitudinal and transverse components, we ﬁnd:
[Equation]
and, using (100) and (93) we have, in matrix notation,
[Equation]
To write things more compactly, the following Lemma will be useful:
Lemma C.5.
For any transverse symmetric matrix T :
[Equation]
Proof.
From (61), Φ 20 TT is proportional to the symmetric part of Φ 11 TT which was denoted by S below Eq.
(93) and satisﬁes Eq.
Therefore Φ 20 TT also satisﬁes Eq.
(99), up to an overall factor:
[Equation]
Then, for any T :
[Equation]
Moreover,
[Equation]
Since the two above equations hold for any T and ˜ L H is invertible, this implies
[Equation]
114) which is the statement of the lemma.
To leading order in η we then have, for a symmetric matrix V , ∂ 2 Φ[ V ] = (cid:80) Di,j =1 ˜ ∂ i ˜ ∂
j Φ V ij , and thus
[Equation]
where V LL =
P L V P L , V TL
= P T V P L , and V TT =
P T V P T are transverse and longitudinal projections of V .
We also have, using (48), and to leading order in η ,
[Equation] [Equation]
where dW is a Wiener process.
Applying Theorem B.4, and keeping into account that, to leading order in dt , d [ Z i , Z j ] =
δ ij dt ,
we ﬁnd
[Equation]
where Σ = σσ T .
For γ < 12 , ˜ L H reduces to the Lyapunov operator at leading order in η , i.e. ˜ L H S
= { H, S } .
For γ > 12 , from (55), it is easy to see that the role of the divergent term proportional to η 1 − 2 γ , when acting ˜ L − 1 H on S , is to set to zero the off-diagonal entries of S ij at O ( η 1 − γ ) , i.e.
[Equation]
Using Lemma B.5, we ﬁnally conclude the following Corollary, which is the formal version of
Theorem 3.4 in the main text:
Corollary C.6.
Consider the stochastic process deﬁned in Eq.
(26) parametrized by (cid:15) n , with initial conditions ( π 0 , w 0 ) ∈
U , under assumptions 3.1 and 3.2.
Fix a compact K ⊂ U .
Then the conclusions of Theorem B.4 apply, and Y ( t ) satisﬁes the limiting diffusion equation
[Equation]
where W ( t ) is a Wiener process.
Let us now see the special case of label noise.
In this case Σ = cH , so that Σ is only transverse.
Moreover, using (55),
[Equation]
and
[Equation]
This proves Corollary 3.5 in the main text.
D E FFECTIVE D RIFT IN UV MODEL
We start with mean-square loss
[Equation]
with the data covariance matrix Σ as deﬁned in the text.
The trace of the Hessian on the zero loss manifold ( L ( w ∗ ) = 0 ) is given explicitly by
[Equation]
Taking gradients of this and plugging into Corollary 3.5, repeated in Eq.
(121) leads to an explicit expression for the drift-diffusion along the manifold
[Equation]
The simpliﬁcation we cite in the main text in Sec. 3.4 is due to the fact that for input and output dimension d = m = 1 , we have that Σ =
Tr Σ = µ 2 , and ˆ S is proportional to the identity.
For matrix sensing, in order to compute the trace of the Hessian, we use (12) with ξ t = 0 and with slightly different notation ( V (cid:62) instead of V ).
The loss is then
[Equation] [Equation]
Then the trace of the Hessian is
[Equation]
We ﬁnd the noise function is
[Equation]
where f ( A ) =
Tr ( UV (cid:62) A ) .
Since the Hessian on the zero loss manifold is H µν =
Pd (cid:80) i ∇ µ f ( A i ) ∇ ν f ( A i ) , we see that σσ T =
(2 /Pd ) H .
Therefore, we get
[Equation]
To get a crude estimate of the coefﬁcient C in the main text, we approximate the top eigenvalue of ˆ
S with 1 d Tr ˆΣ 1 .
With this, we get for
[Equation]
+
[Equation]
γ
[Equation]
Here we have denoted by a ij an arbitrary element of the data matrices A , with brackets signifying an average over the distribution of these elements.
Assuming the fast initial fast remains the same, and τ − 1≈ ( C/ 2) η γ , we get
[Equation]
For the values used in our experiments, this gives C ≈ 0 .
12 × P − 1 / 3 .
E L INEARIZATION A NALYSIS OF M OMENTUM G RADIENT D ESCENT IN THE SCALING LIMIT
Here we elaborate on the discussion in Sec. 3.3, providing derivations of various results.
We take a straightforward linearization of the deterministic (noise-free) gradient descent with momentum.
Working in the extended phase space x
= ( π, w ) , the dynamical updates are of the form
[Equation]
The ﬁxed point of the evolution x ∗ = (0 , w ∗ ) will have the momentum variable π = 0 , and the coordinate satisfying ∇ L ( w ∗ )
= 0 .
Linearizing the update (132) around this point
[Equation]
Note ∇ 2 L ( w ∗ ) is the Hessian of the loss function at the ﬁxed point.
The spectrum of the Jacobian can be written in terms of the eigenvalues of the Hessian λ i .
This is accomplished by using a straightforward ansatz for the (unnormalized) eigenvectors of the Jacobian k i =
( µ i q i , q i ) , where q i are eigenvectors of the Hessian with eigenvalue λ
i .
Solving the resulting coupled eigenvalue equations for eigenvalue κ
i :
[Equation]
For a ﬁxed λ i , there will be two solutions given by
[Equation]
For the set of zero modes λ i = 0 , we get the following modes: κ + = 1 , corresponding to motion only along w , with eigenvector k i
=
(0 , q i ) .
In addition, there is a mixed longitudinal mode which includes a component of π along the zero manifold k i =
( µ − q i , q i ) , and has an eigenvalue κ − = β .
On the zero loss manifold, we can assume the Hessian is positive semi-deﬁnite, and that the positive eigenvalues satisfy
[Equation]
137) for constants c 1 , c 2 independent of η, β .
We now analyze the spectrum of the Jacobian one eigenvalue at a time, and then use these results to informally control the relaxation rate of off-manifold perturbations.
It is useful ﬁrst to consider the conditions for stability, i.e. | κ i | < 1 , which are stated in (138,139) below:
Case 1 : If ηλ i < (1 − (cid:112) β ) 2 , then κ i ± ∈ R and | κ i ± | < 1 iﬀ 0 <
ηλ i < 2(1 + β ) .
Case 2 : If ηλ i > (1 − (cid:112) β ) 2 , then κ i ± ∈ C and | κ i ± | = (cid:112) β < 1 .
Proof of Case 1: The condition ηλ i < (1 − √ β ) 2 implies A 2 > 4 β , where
A = 1 + β − ηλ i .
The condition for stability then requires − 1 < κ i <
+1 .
We satisfy both sides of this inequality:
If 1 + β − ηλ i =
A > 0 , then | κ i − | < 1 ,and κ i + > 0 , so we simply require κ i + < 1 , i.e.
[Equation]
[Equation] [Equation]
[Equation] [Equation]
β
[Equation]
If 1 + β − ηλ i < 0 , then | κ i + | < 1 , and κ i − < 1 , so we only require κ
i − > −−
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
⇒
[Equation]
Proof of Case 2: the condition ηλ i > (1 − √ β ) 2 implies A 2 < 4 β , where
A = (1 + β − ηλ i ) .
The corresponding eigenvalues of the Jacobian can be written κ ± i =
(1 / 2)( A ± (cid:112)
A 2 − 4 β ) =
(1 / 2)( A ± i (cid:112) 4 β − A 2 ) .
Computing the absolute value then gives | κ ± i | 2 = (1 / 4)( A 2 +4 β − A 2 ) =
β
These results show us when the GD+momentum is stable.
Next, assuming stability, we want to estimate the rate of convergence to the ﬁxed point.
More precisely, we would like to determine the fastest mode as well as the slowest mode.
To this end, we deﬁne two quantities
[Equation] [Equation]
Using the explicit scaling for momentum, and in the limit of small learning rate, we prove the following for ρ 1 :
Lemma E.1.
Let β = 1 − Cη γ , and
η sufﬁciently small:
For γ < 1 / 2 , the condition for Case 1 (138) holds,
[Equation]
For γ > 1 / 2 , the condition for Case 2 (139), and
[Equation]
Proof: For small η , we ﬁnd that
[Equation]
When γ > 1 / 2 , this expression tends to zero as η → 0 .
Therefore, for sufﬁciently small η , the condition for Case 2 (139) will be satisﬁed and | κ i | = √ β for all λ i > 0 .
Since β < √ β , this implies ρ 1 = √ β and ρ 2 = β .
The scaling behavior in the Lemma follows by substitution.
For γ < 1 / 2 , (157) diverges as η → 0 , which means Case 1 (138) will obtain for all λ i .
Next, since for small η , 1 + β − ηλ i > 1 + β − ηc 2 = 2 − Cη γ − ηc 2 > 0 , since C and c 2 are order one constants, and η → 0 .
In this case, the largest contribution from the nonzero eigenvalues λ
i will come from κ i + .
In particular, we ﬁnd
[Equation] [Equation]
[Equation]
For γ < 1 / 2 , we have the hierarchy η γ > η 2 γ > η > η γ +1 > η 2 .
This allows us to simplify the upper bound
[Equation] [Equation]
Next, we ﬁnd a lower bound for ρ 2 .
This will be controlled by κ i − .
We may use then that
[Equation] [Equation]
Finally, note that since η 1 − γ < η γ , we have that 1 − c 1 C η 1 − γ > 1 − Cη γ
= β , so indeed ρ 2 < ρ 1 .
Equipped with the upper and lower bounds on the spectrum leads naturally to bounds on the relaxation rate.
For the purely decaying modes at γ < 1 / 2 , we use ρ t 2 ≤ | δx Tt | ≤ ρ t 1 , δx Tt represents the projection of the ﬂuctuations δx t onto the transverse and mixed longitudinal modes.
After applying
Eqs.
(153) and (154)
, we arrive at the result quoted in the main text in Sec. 3.3.
For γ > 1 / 2 , the modes are oscillatory.
However, the eigenvalues within the unit circle have norm either √ β , β , as reﬂected in Eqs.
(155) and (156).
This implies that we can estimate the decay rate of the envelope of the transverse and mixed longitudinal modes in this regime, thereby arriving at second expression quoted in the main text in Sec. 3.3.
F L INEARIZED SGD AND O RNSTEIN -U HLENBECK PROCESS ON THE VALLEY
In this appendix, we provide a derivation of some of the statements quoted in Sec. 1.1.
To get there, we start with the basic model for momentum SGD (1) but linearize around a point on the valley w 0 ∈ Γ where L ( w 0 ) =
∇ L ( w 0 ) = 0 .
Let w k = w
+ δw k , and deﬁne the Hessian
H ( w 0 )
= ∇ 2 L ( w 0 ) .
Then
[Equation]
Consider the projector along transverse nonzero eigemode λ of H ( w 0 ) , P Tλ , and deﬁne P Tλ x = X , and P Tλ π = Π .
Let ¯ σ = P Tλ σ , and P Tλ H = λP
Tλ .
Let ¯ σ ¯ σ (cid:62) = Λ .
Then
[Equation]
This is a simple OU process
[Equation]
and we can easily compute the second moments.
Deﬁne the second moments
[Equation]
We ﬁnd by taking the equations above, squaring them, then averaging over the noise,
[Equation]
Next, assuming a stationary distribution implies C ( k +1)
=
C ( k )
[Equation]
which then allows us to solve for equilibrium variance, and extract the main quantity of interest, which is the variance of the weights.
We ﬁnd then
[Equation] [Equation]
In the limit of small η we extract the scaling behavior quoted in Sec.(1.1).
We can see how the mixing timescale τ 1 , discussed in Sec. 1.1, arises from this linearized analysis.
By taking the expectation value of the OU process, the noise will vanish and we ﬁnd that the average values follow the linearized GD dynamics analyzed in E. Thus, from this linearized GD analysis we can extract the characteristic timescale for the OU process to approach its mean value.