Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensions s
AfﬁliationAddress email
Abstract
We analyze the dynamics of large batch stochastic gradient descent with momentum(SGD+M) on the least squares problem when both the number of samples anddimensions are large.
In this setting, we show that the dynamics of SGD+Mconverge to a deterministic discrete Volterra equation as dimension increases,which we analyze.
We identify a stability measurement, the implicit conditioningratio (ICR), which regulates the ability of SGD+M to accelerate the algorithm.
When the batch size exceeds this ICR, SGD+M converges linearly at a rate of
O (1 / p  ) , matching optimal full-batch momentum (in particular performing aswell as a full-batch but with a fraction of the size).
For batch sizes smaller than the
ICR, in contrast, SGD+M has rates that scale like a multiple of the single batch
SGD rate.
We give explicit choices for the learning rate and momentum parameterin terms of the Hessian spectra that achieve this performance.
Stochastic learning algorithms are the methods of choice for optimization of high-dimensionalproblems.
Often stochastic learning algorithms incorporate momentum into their stochastic gradientsto improve practical performance.
Perhaps the simplest, stochastic gradient descent with momentum(SGD+M) adds a ﬁxed multiple of the backward difference of iterates to its stochastic gradientestimator, see Section 1 for details.
In the inﬂuential work of [Sutskever et al., 2013], the authorsempirically show augmenting stochastic gradient descent (SGD) with momentum signiﬁcantlyimproves training performance of deep neural networks.
Despite the wide usage of these stochasticmomentum methods in machine learning practice, our understanding of its behaviour is not well–understood.
It has been hypothesized that stochastic-based momentum algorithms improve training because theyare employed on a large batch of a data set [
Kidambi et al
; thereby emulating the speed-upone sees in full-batch settings.
For many learning problems, the “large batch” setting is often pairedwith high-dimensional problems, meaning there are many samples (and likely also many features tohave interesting behavior).
We know of no theoretical analysis that can justify this claim for standard
SGD+M, although for variations of SGD+M and SGD with Nesterov momentum [Nesterov, 2004]there has been some success in proving accelerated rates [Jain et al., 2018, Liu and Belkin, 2020,
Allen-Zhu, 2017].
A reason is that typical approaches for analyzing SGD+M do not distinguishlarge and small batch sizes.
We address this problem in this paper and we introduce a stabilitymeasurement that exactly captures the transition of SGD+M to an accelerated method.
We commentthat in the high–dimensional, vanishing batch fraction setting (the mini-batch size is o ( n ) , where n isthe number of samples), there is work proving in various simpliﬁed settings that SGD+M producesthe same iterates as SGD with a larger learning rate, up to a vanishing error [Paquette and Paquette,2021].
In this paper, we study the dynamics of mini-batch SGD+M (with constant learning rate) on a leastsquares problem when the number of samples n and features d are large (see Section 1 for details).
We are motivated by the setting where the mini-batch size  is proportionate to the number of samplesn and so we deﬁne the ratio ⇣ def =
/n , which we refer to as the batch fraction .
We provide anon-asymptotic comparison for the behavior of the training loss under SGD+M to a deterministicfunction , whose accuracy improves when the number of samples and features are large while thebatch fraction is strictly positive (see Figure 1).
This function solves a discrete Volterra equation:
[Equation]
The forcing term F ( t ) and kernel K ( t ) are explicit functions that depend on the hyperparameters andthe full Hessian spectra (see Section 2 and Appendix A).
They transparently reveal that the dynamicsof SGD+M and of SGD are truly non-equivalent in that there is no mapping of the hyperparameterswhich leads them to have the same training dynamics.
We also note that a similar equation appearsin the vanishing batch setting [Paquette et al., 2021], although in that setting it is a Volterra integralequation, which can be recovered from (1) by sending ⇣ ! 0 . changes on a Gaussian random least squares problem with ratio r def = dn varying. Here  =
[Equation]
and
[Equation]  [Equation] q
[Equation]
[Equation]
When ⇣ goes above a speciﬁc value (see
Proposition 5), the convergence rate freezes at 1 / p  .
Otherwise the rate behaves like C/  with a constant C = ⇣ / (1  ⇣ ) and we get no speed up from SGD+M.
An advantage of the exact loss trajectory is that wegive a rigorous deﬁnition of the large batch and smallbatch regimes which reﬂect a transition in the convergence behavior of SGD+M. To do this we introduce the condition number  , the average conditionnumber ¯  , and the implicit conditioning ratio (ICR)deﬁned as
[Equation]
Here  2 j are the eigenvalues of the Hessian of theleast squares problem with  2max and  2min the largestand smallest (non-zero) eigenvalues.
We refer to thelarge batch regime where ⇣  ICR and the smallbatch regime where ⇣  ICR.
In the large batchregime (2), SGD+M matches the performance of theheavy-ball algorithm: the convergence is linear withrate given by O (1 / p  ) .
In the small batch regime,the performance matches that of SGD, i.e. the convergence rate is O ( ⇣ /  ) .
We give matching lowerbounds, and we provide momentum and learning ratechoices that achieve the claimed performance.
In addition we show there is a saturating batch fraction (see Figure 2), after which increasing the batchfraction does not improve the rate.
It explicitly occurs when ⇣ = ICR.
Moreover this saturating batchfraction occurs before full batch, i.e
Related work.
Recent works have established convergence guarantees for SGD+M in both stronglyconvex and non-strongly convex setting [Flammarion and Bach, 2015, Sebbouh et al., 2021], includingalmost sure convergence [Gadat et al., 2018].
In the work of [Orvieto et al., 2020], they usedstochastic differential equations (SDEs) to obtain convergence of SGD+M. Specializing to the settingof minimizing quadratics, [Loizou and Richtarik, 2020] demonstrated that the iterates of SGD+Mconverge linearly (but not in L2) under an exactness assumption.
Determining batch size has been an important issue in determining the convergence rate of SGD and
SGD+M. There are instances where (small batch size) SGD+M does not necessarily achieve betterperformances than small batch size SGD (see [Zhang et al., 2019, Kidambi et al., 2018, Paquetteet al., 2021]).
As for mini-batch SGD without momentum, [Ma et al., 2018] showed that there is asaturating batch size (roughly  /  ) above which increasing the batch size no longer improves therate.
In [De et al., 2017], the authors implemented an adaptive (increasing) batch size schedule andthey used it to show linear convergence for SGD.
For generalization, [Smith et al., 2018] empiricallyshowed that for SGD and SGD+M, instead of decaying the learning rate, one can increase the batchsize during training to obtain a similar learning curve.
SGD+M has been proven to be useful in practical applications as well, including machine learning.[Sutskever et al., 2013] demonstrated that SGD+M shows an empirical advantage in training deepand recurrent neural networks (DNNs and RNNs respectively).
Many authors have proposed thatlearning rate warmup enables us to scale training efﬁciently to larger batch sizes ([Goyal et al., 2017,
McCandlish et al., 2018, Smith et al., 2018]).
Setting
We consider the least squares problem when the number of samples ( n ) and features ( d ) are large:arg
[Equation]
with b def =
A e x + ⌘
[Equation]
(3) where A 2 R n ⇥ d is a data matrix whose i -th row is denoted by a i 2 R d , e x 2 R d is the signal vector,and ⌘ 2 R n is a source of noise.
The target b =
A e x + ⌘ comes from a generative model corruptedby noise.
We let  21  · · ·   2 n  0 be the eigenvalues of the matrix AA T with  2max and  2min
the 95 largest and smallest (nonzero) eigenvalues.
We apply SGD with momentum (SGD+M) with mini-batches to the ﬁnite sum, quadratic problem(3).
SGD+M iterates by selecting uniformly at random a subset B k ✓ { 1 , 2 , · · · , n } of cardinality  98 and makes the update
[Equation]
where P k def =
X i 2 B k e i e Ti , (4) with P k a random orthogonal projection matrix and e i the i -th standard basis vector.
Here  > 0 isthe learning rate parameter,  is the momentum parameter, and the function f i is the i -th element ofthe sum in (3).
When the stochastic gradient in (4) is replaced with the full-gradient r f ( x ) , the resulting algorithmwith learning rate and momentum optimally chosen yields the celebrated algorithm, heavy-ballmomentum (a.k.a.
Polyak momentum)
[Polyak, 1964].
The optimal learning rate and momentumparameters are explicitly given by
[Equation]
It is well-known that heavy-ball is an optimal algorithm on the least squares problem in that itconverges linearly at a rate of O (1 / p  ) (see [Pedregosa, 2021]).
In this paper, we adhere whenever possible to the following notation.
We denote vectors in lowercaseboldface ( x ) and matrices in upper boldface ( A ) .
The entries of a vector (or matrix) are denoted bysubscripts.
Unless otherwise speciﬁed, the norm k · k 2 is taken to be the standard Euclidean norm if it 111 is applied to a vector and the operator 2-norm if it is applied to a matrix.
1.1 Random least squares problem
To perform our analysis we make the following explicit assumptions on the signal e x , the noise ⌘ , andthe data matrix A .
Assumption 1.1 (Initialization, signal, and noise) .
The initial vector x 0 2 R d is chosen so thatx
e x is independent of the matrix A .
The noise vector ⌘ 2 R n is centered and has i.i.d.
entries,independent of A .
The signal and noise are normalized so that
[Equation]
Next we state assumptions on the data matrix
A as well as its eigenvalue and eigenvector distribution.
Each row a i 2 R d ⇥ 1 is centered and is normalized so that E k a i k
= 1 for all i .
We suggestas a central example the Gaussian random least squares setup where each entry of A is sampledindependently from a standard normal distribution with variance 1 d .
Assumption 1.2 (Orthogonal invariance) .
Let A be a random n ⇥ d matrix.
Suppose these randommatrices satisfy a left orthogonal invariance condition: Let O 2 R n ⇥ n be an orthogonal matrix.
Then the matrix A is orthogonally left invariant in the sense that
[Equation]
This assumption implies that the left singular vectors of A are uniformly distributed on the spherewhich is the strongest form of eigenvector delocalization; many distributions of random matricesincluding some sparse ones (such as random regular graph adjacency matrices) are known to havesome form of eigenvector delocalization.
The classic example of a random matrix which has leftorthogonal invariance is the sample covariance matrix, Z p ⌃ , for an i.i.d.
Gaussian matrix Z andany covariance matrix ⌃ .
Numerical simulations suggest that (6) can be weakened in that the theoryherein can be applied to other ensembles without this orthogonal invariance property.
See Figure 3.
Deterministic Dynamical Equivalent of SGD+M
With these assumptions, we can give an explicit representation of the loss values on a least squaresproblem at the iterates generated by SGD+M algorithm.
We show in this section (see Theorem 1):for any T > 0 ,
[Equation]
where solves (1).
We begin by discussing the forcing and the noise terms of ( t ) and theirrelationship to SGD+M.
Forcing term: problem instance information.
The forcing term represents the mean (with respectto expectation over the mini-batches) behavior of SGD+M. In fact, the forcing term is the loss funder full-batch gradient descent with momentum  but with learning rate ⇣ .
For a small learningrate  , the forcing term F ( t ) in (1) governs the dynamics of ( t ) .
Let w t def =
Ax t  b and observe that 1 / 2 k w t k 22 is the loss f ( x t ) .
One way to get the dynamics of fis by deriving a recurrence for w 2 t,j (for each j 2 [ n ] ) which we get from the updates of SGD+M. Therecurrence is as followed: deﬁne ˜ X t,j def = ( w 2 t,j w 2 t  1 ,j w t,j w t  1 ,j ) T and there exists a matrix M j 145 (see Appendix A.2) so that
[Equation]
The forcing term at iteration t is given by applying a linear recurrence M j operator t  1 timeson a vector containing initialization information at each j 2 [ n ] and then summed up for the ﬁrstcoordinate.
It is clear that the maximum of the eigenvalues of the operator M j is essential to analyzethe convergence behavior of F ( t ) .
We denote  1 ,j =  and  i,j , i = 2 , 3 , the eigenvalues of M jand note that   max i =2 , 3 |  i,j | (see (36) for an explicit formula of  i,j that depends only on  , , and eigenvalues of AA T ).
Let  2 ,j be the eigenvalue of M j with the biggest modulus and let
[Equation]
From this, an equation for the forcing term F ( t ) can be made explicit in terms of  i,j , see Thm.
1below and Appendix A. We can conclude that F ( t ) =
O (  t 2 , max ) .
Kernel term: noise from the algorithm.
The convolution term in (1) is due to the inherentstochasticity of SGD+M. More speciﬁcally, it is given by
[Equation]
where ⌦ j def = 1  ⇣ 2 j +  , and H 2 ( t ) def = 1 n n X j =
[Equation]  [Equation] The presence of (training loss) is due to the fact that the noise generated by the k -th stochasticgradient is proportional to ( k ) (training loss), and the function H 2 ( t  k ) represents the progress ofthe algorithm in sending this extra noise to 0 .
Observe (9) scales quadratically in the learning rate  .
Hence for large learning rates, (9) dominates the decay behaviour of .
Further details discussed in Section 3.1.
We now state the main result:
Theorem 1 (Concentration of SGD+M) .
Suppose Assumptions 1.1 and 1.2 hold with the learningrate  < 1+  ⇣ 2 max and the batch size satisﬁes  /n
=
⇣ for some ⇣ > 0 .
Let the constant T 2 N .
Thenthere exists C > 0 such that for any c > 0 , there exists D > 0 satisfying
[Equation]
(10) for sufﬁciently large n 2 N .
The function is the solution to the Volterra equation
[Equation]
where for k = 0 , 1 , i = 2 , 3 , and j 2 [ n ] ,  i,j def =  i,j ⌦ j / (  i,j +  ) , andh k ( t ) =
1 n
[Equation]
For a more detailed description on h 0 , h 1 , H 2 as well as the proof of Theorem 1 and Corollary 1, see
Appendix A. The expression of highlights how the algorithm, learning rate, batch size, momentum,and noise levels interact with each other to produce different dynamics.
Note that the learningrate assumption will be necessary for the solution to the Volterra equation to be convergent, see
Proposition 2.
When  ! 0 , we obtain the Volterra equation for SGD with mini-batching.
Corollary 1 (Concentration of SGD, no momentum) .
Under the same setting as Theorem 1 and when = 0 , the function values f ( x t ) converge to ( t ) as in (10) where now the limit is a solution tothe Volterra equation
[Equation]
where for k = 0 , 1 , 2 ,
[Equation]
Remark.
Note that H 2 ( t ) reduces to h 2 ( t ) in  = 0 case.
Also when the limit ⇣ !
and whenwe scale time by t/ ⇣ , we have that (1  ⇣ 2 j ) 2 t/ ⇣ !
e  2  2 j t .
This coincides with the result from[Paquette et al., 2021, Theorem 1].
Indeed, this shows not only how our dynamics of SGD+Mincludes the no momentum case (i.e. SGD), but also how the dynamics of SGD+M differ from SGD.
Convolution Volterra analysis
In this section, we outline how to utilize the Volterra equation (11) to produce a complexity analysisof SGD+M.
For additional details and proofs in this section, see Appendix C.
We begin by establishing sufﬁcient conditions for the convergence of the solution to the Volterraequation (11).
Our Volterra equation can be seen as the renewal equation ([Asmussen, 2003]).
Let ustranslate (11) into the form of the renewal equation as follows:
[Equation]
where ( f ⇤ g )( t ) =
P 1 k =0 f
( t  k ) g ( k ) .
Let the kernel norm be k K k =
P 1 t =0 K ( t ) .
By [Asmussen,2003, Proposition 7.4], we see that k K k < 1 is necessary for our solution to the Volterra equation tobe convergent.
Indeed, we have the following result.
Proposition 1.
If the norm k K k < 1 , the algorithm is convergent in that
[Equation]
Note that the noise factor e
R and the matrix dimension ratio d/n appear in the limit.
Proposition 1formulates the limit behaviour of the objective function in both the over-determined and the underdetermined case of least squares.
When under-determined, the ratio d/n  1 and the limiting ( 1 )is 0 ; otherwise the limit loss value is strictly positive.
The result (14) only makes sense when thenoise term K satisﬁes || K || < 1 ; the next proposition illustrates the conditions on the learning rate 195 and the trace of the eigenvalues of AA T such that the kernel norm is less than 1.
Proposition 2 (Convergence threshold) .
Under the learning rate condition  <
[Equation]
⇣ 2 max and tracecondition (1  ⇣ )
[Equation]
, the kernel
[Equation]
The learning rate condition quantiﬁes an upper bound of good learning rates by the largest eigenvalueof the covariance matrix  2max , batch size ⇣ , and the momentum parameter  .
The trace conditionillustrates a constraint on the growth of  2max .
Moreover, for a full batch gradient descent model( ⇣ = 1) , the trace condition can be dropped and we get the classical learning rate condition forgradient descent.
3.1 The Malthusian exponent and complexity
The rate of convergence of is essentially the worse of two terms – the forcing term F ( t ) and adiscrete time convolution
P {+tk =
+}
( k ) K ( t  k ) which depends on the kernel K .
Intuitively, theforcing term captures the behavior of the expected value of SGD+M and the discrete time convolutioncaptures the slowdown in training due to noise created by the algorithm.
Note that F ( t ) is alwaysa lower bound for ( t ) , but it can be that ( t ) is exponentially (in t ) larger than F ( t ) owing to the convolution term.
This occurs when something called the Malthusian exponent , denoted ⌅ , of theconvolution Volterra equation exists.
The Malthusian exponent ⌅ is given as the unique solution to
[Equation]
, if the solution exists.
[Equation]
The Malthusian exponent enters into the complexity analysis in the following way:
Theorem 2 (Asymptotic rates) .
The inverse of the Malthusian exponent always satisﬁes ⌅  1 > 2 , max for ﬁnite n .
Moreover, for some C > 0 , the convergence rate for SGD+M is
[Equation]
Thus to understand the rates of convergence, it is necessary to understand the Malthusian exponent asa function of  and  .
3.2 Two regimes for the Malthusian exponent
On the one hand, the Malthusian exponent ⌅ comes from the stochasticity of the algorithm itself.
Onthe other hand,  2 , max (  ,  , ⇣ ) is determined completely by the problem instance information — theeigenspectrum of AA T .
(Note we want to emphasize the dependence of  2 , max on learning rate,momentum, and batch fraction.)
Let  2max and  2min denote the maximum and minimum nonzero eigenvalues of AA T , respectively.
For a ﬁxed batch size, the optimal parameters (   ,   ) of  2 , maxare
[Equation]
[Equation]
In the full batch setting, i.e. ⇣ = 1 , these optimal parameters   and   for  2 , max are exactlythe Polyak momentum parameters (5).
Moreover, in this setting, there is no stochasticity so the
Malthusian exponent disappears and the convergence rate (16) is  2 , max .
We observe from (17) thatfor all ﬁxed batch sizes, the optimal momentum parameter,   , is independent of batch size.
Theonly dependence on batch size appears in the learning rate.
At ﬁrst it appears that for small batchfractions, one can take large learning rates, but in that case, the inverse of the Malthusian exponent⌅  1 dominates the convergence rate of SGD+M (16) and you cannot take  and  to be as in (17)(See Figure 4).
We will deﬁne two subsets of parameter space, the problem constrained regime and the algorithmicallyconstrained regime (or stochastically constrained regime).
The problem constrained regime is forsome tolerance " >
0problem constrained regime
[Equation]
The remainder we call the algorithmically constrained regime.
To explain the tolerance: for ﬁniten , it transpires that we always have ⌅  1 >  2 , max , but it could be vanishingly close to  2 , max as afunction of n .
Hence we introduce the tolerance to give the correct qualitative behavior in ﬁnite n .
Proposition 3.
If the learning rate   min( 1+  ⇣ 2 max , (1 p  ) 2 ⇣ 2min ) , with the trace condition8(1  ⇣ )   · 1 n tr( A T A ) < 1 , then (  ,  ) is in the problem constrained regime with " = 1 / 2 .
Therefore by (16), we have that
[Equation]
for some D >
[Equation]
19) we note that the expression in the parenthesis is 1  12 (1   2 , max ) +
O ((1   2 , max ) 2 ) .
In the problem constrained regime, it is worthwhile to note that the overall convergence rate isthe same as full batch momentum with adjusted learning rate, i.e., the batch size does not play animportant role as long as we are in the problem constrained regime:
Proposition 4 (Concentration of SGD + M, full batch) .
Suppose ⇣ = 1 and Assumptions 1.1 and 1.2hold with the learning rate  < 1+   2 max .
If we let x full t denote the iterates of full-batch gradient descentwith momentum (GD+M), then
[Equation]
where
[Equation]
The functions h 1 and h 0 are deﬁned in Theorem 1 with ⇣ = 1 .
In particular, let  full denote thelearning rate for full batch GD+M, and  , ⇣ < 1 for the learning rate and batch fraction in SGD+Mwith corresponding in Theorem 1.
Then when  full = ⇣ is satisﬁed, and full share the sameconvergence rate in the problem constrained regime.
Performance of SGD+M: implicit conditioning ratio (ICR)
Recall from (2) the deﬁnition of condition number, average condition number, and the implicitconditioning ratio
[Equation]
and
[Equation]
Moreover recall that we refer to the large batch regime where ⇣  ICR and the small batch regimewhere ⇣  ICR.
We begin by giving a rate guarantee that holds in the problem constrained regime, for a speciﬁcchoice of  and  .
Proposition 5 (Good momentum parameters) .
Suppose the learning rate and momentum satisfy =
(1  p  )
⇣ 2min and  = max
[Equation]
1 p 2 
[Equation]
, where C def = ⇣ / (8(1  ⇣ )) .
Then  2 , max =  and for some C > 0 , the convergence rate for SGD+M is
[Equation]
Remark 1.
We note that for all  satisfying (1 p  ) 2 ⇣ 2min  [Equation] 1+ p  ) 2 2 ⇣ 2max with the learning rate  as in(22) , we have that  2 , max =  .
By minimizing the  (i.e., by ﬁnding the fastest convergence rate),we get the formula for the momentum parameter in (22) .
The exact tradeoff in convergence rates (23) occurs when
[Equation]
As ⇣  1 , this condition is only nontrivial when ICR ⌧ 1 , in which case ⇣ = 8 p 2 ICR, up to vanishingerrors.
Large batch ( ⇣  ICR).
In this regime SGD+M’s performance matches the performance of theheavy-ball algorithm with the Polyak momentum parameters (up to absolute constants).
Morespeciﬁcally with the choices of  and  in Proposition 5, the linear rate of convergence of SGD+Mis 1  c p  for an absolute c .
Note that ⇣ does not appear in the rate, and in particular there is no gainin convergence rate by increasing the batch fraction.
Small batch ( ⇣  ICR).
In the small batch regime, the value of C is relatively small and the ﬁrstterm is dominant in (23), and so the linear rate of convergence of SGD+M is 1  c ⇣ for some absoluteconstant c > 0 .
In this regime, there is a beneﬁt in increasing the batch fraction, and the rate increaseslinearly with the fraction.
We note that on expanding the choice of constants in small ⇣ the choicesmade in Proposition 5 are
[Equation]
This rate can also be achieved by taking  = 0 , i.e. mini-batch SGD with no momentum.
Moreover,it is not possible to beat this by using momentum; we show the following lower bound:
Proposition 6.
If ⇣  min { 12 , ICR } then there is an absolute constant C >
so that for convergent(  ,  ) (those satisfying Proposition 2),
[Equation]
This is a lower bound on the rate of convergence by Theorem 2.
Conclusion and future work
We have shown that the SGD+M method on a least squares problem demonstrates deterministicbehavior in the large n and d limit.
We described the dynamics of this algorithm through a discrete Volterra equation and for a ﬁxed batch fraction.
Moreover we characterized a dichotomy ofconvergence regimes depending on the learning rate and momentum parameters.
Furthermore, weproved that SGD+M shows a distinguishable improvement over SGD in the large batch regime andwe provided parameters which achieve acceleration.
Our theory is also supported by numericalexperiments on the isotropic features model and MNIST data set (see Appendix D for details).
While our analysis focuses on SGD+M algorithm applied to the least squares problems with orthogonal invariant data matrix, Figure 3 suggests that the Volterra prediction might hold in even greatergenerality.
Removing these conditions, we leave as future work.
Another direction of future workconsists in ﬁnding the deterministic dynamics for generalization errors.