Accelerating SGD with momentum for over-parameterized learning

1 INTRODUCTION . Many modern neural networks and other machine learning models are over-parametrized ( 5 ) . These models are typically trained to have near zero training loss , known as interpolation and often have strong generalization performance , as indicated by a range of empirical evidence including ( 23 ; 3 ) . Due to a key property of interpolation – automatic variance reduction ( discussed in Section 2.1 ) , stochastic gradient descent ( SGD ) with constant step size is shown to converge to the optimum of a convex loss function for a wide range of step sizes ( 12 ) . Moreover , the optimal choice of step size η∗ for SGD in that setting can be derived analytically . The goal of this paper is to take a step toward understanding momentum-based SGD in the interpolating setting . Among them , stochastic version of Nesterov ’ s acceleration method ( SGD+Nesterov ) is arguably the most widely used to train modern machine learning models in practice . The popularity of SGD+Nesterov is tied to the well-known acceleration of the deterministic Nesterov ’ s method over gradient descent ( 15 ) . Yet , has not not theoretically clear whether Nesterov SGD accelerates over SGD . As we show in this work , both theoretically and empirically , Nesterov SGD with any parameter selection does not in general provide acceleration over ordinary SGD . Furthermore , Nesterov SGD may diverge , even in the linear setting , for step sizes that guarantee convergence of ordinary SGD . Intuitively , the lack of acceleration stems from the fact that , to ensure convergence , the step size of SGD+Nesterov has to be much smaller than the optimal step size for SGD . This is in contrast to the deterministic Nesterov method , which accelerates using the same step size as optimal gradient descent . As we prove rigorously in this paper , the slow-down of convergence caused by the small step size negates the benefit brought by the momentum term . We note that a similar lack of acceleration for the stochastic Heavy Ball method was analyzed in ( 9 ) . To address the non-acceleration of SGD+Nesterov , we introduce an additional compensation term to allow convergence for the same range of step sizes as SGD . The resulting algorithm , MaSS ( Momentum-added Stochastic Solver ) 1 updates the weights w and u using the following rules ( with the compensation term underlined ) : wt+1 ← ut − η1∇̃f ( ut ) , ut+1 ← ( 1 + γ ) wt+1 − γwt + η2∇̃f ( ut ) . ( 1 ) Here , ∇̃ represents the stochastic gradient . The step size η1 , the momentum parameter γ ∈ ( 0 , 1 ) and the compensation parameter η2 are independent of t. We proceed to analyze theoretical convergence properties of MaSS in the interpolated regime . Specifically , we show that in the linear setting MaSS converges exponen- tially for the same range of step sizes as plain SGD , and the optimal choice of step size for MaSS is exactly η∗ which is optimal for SGD . Our key theoretical result shows that MaSS has accelerated convergence rate over SGD . Furthermore , in the full batch ( deterministic ) scenario , our analysis selects η2 = 0 , thus reducing MaSS to the classical Nesterov ’ s method ( 15 ) . In this case our convergence rate also matches the well-known convergence rate for the Nesterov ’ s method ( 15 ; 4 ) . This acceleration is illustrated in Figure 1 . Note that SGD+Nesterov ( as well as Stochastic Heavy Ball ) does not converge faster than SGD , in line with our theoretical analysis . We also prove exponential convergence of MaSS in more general convex setting under additional conditions . We further analyze the dependence of the convergence rate e−s ( m ) t and optimal hyper-parameters on the mini-batch size m. We identify three distinct regimes of dependence defined by two critical values m∗1 and m ∗ 2 : linear scaling , diminishing returns and saturation , as illustrated in Figure 2 . The convergence speed per iteration s ( m ) , as well as the optimal hyper-parameters , increase linearly as m in the linear scaling regime , sub-linearly in the diminishing returns regime , and can only increase by a small constant factor in the saturation regime . The critical values m∗1 and m ∗ 2 are derived analytically . We note that the intermediate “ dimin- ishing terurns ” regime is new and is not found in SGD ( 12 ) . To the best of our knowledge , this is the first analysis of mini-batch dependence for accelerated stochastic gradient methods . We also experimentally evaluate MaSS on deep neural networks , which are non-convex . We show that MaSS outperforms SGD , SGD+Nesterov and Adam ( 10 ) both in optimization and generalization , on different architectures of deep neural networks including convolutional networks and ResNet ( 7 ) . The paper is organized as follows : In section 2 , we introduce notations and preliminary results . In section 3 , we discuss the non-acceleration of SGD+Nesterov . In section 4 we introduce MaSS and analyze its convergence and optimal hyper-parameter selection . In section 5 , we analyze the mini-batch MaSS . In Section 6 , we show experimental results . 1.1 RELATED WORK . Over-parameterized models have drawn increasing attention in the literature as many modern machine learning models , especially neural networks , are over-parameterized ( 5 ) and show strong generalization performance ( 16 ; 23 ; 2 ) . Over-parameterized models usually result in nearly perfect fit ( or interpolation ) of the training data ( 23 ; 18 ; 3 ) . Exponential convergence of SGD with constant step size under interpolation and its dependence on the batch size is analyzed in ( 12 ) . 1Code url : https : //github.com/ts66395/MaSS There are a few works that show or indicate the non-acceleration of existing stochastic momentum methods . First of all , the work ( 9 ) theoretically proves non-acceleration of stochastic Heavy Ball method ( SGD+HB ) over SGD on certain synthetic data . Furthermore , these authors provide experimental evidence that SGD+Nesterov also converges at the same rate as SGD on the same data . The work ( 22 ) theoretically shows that , for sufficiently small step-sizes , SGD+Nesterov and SGD+HB is equivalent to SGD with a larger step size . However , the results in ( 22 ) do not exclude the possibility that acceleration is possible when the step size is larger . The work ( 11 ) concludes that “ momentum hurts the convergence within the neighborhood of global optima ” , based on a theoretical analysis of SGD+HB . These results are consistent with our analysis of the standard SGD+Nesterov . However , this conclusion does not apply to all momentum methods . Indeed , we will show that MaSS provably improves convergence over SGD . There is a large body of work , both practical and theoretical , on SGD with momentum , including ( 10 ; 8 ; 1 ) . Adam ( 10 ) , and its variant AMSGrad ( 17 ) , are among the most practically used SGD methods with momentum . Unlike our method , Adam adaptively adjusts the step size according to a weight-decayed accumulation of gradient history . In ( 8 ) the authors proposed an accelerated SGD algorithm , which can be written in the form shown on the right hand side in Eq.8 , but with different hyper-parameter selection . Their ASGD algorithm also has a tail-averaging step at the final stage . In the interpolated setting ( no additive noise ) their analysis yields a convergence rate of O ( Poly ( κ , κ̃ ) exp ( − t 9 √ κ1κ̃ ) ) compared to O ( exp ( − t√ κ1κ̃ ) ) for our algorithm with batch size 1 . We provide some experimental comparisons between their ASGD algorithm and MaSS in Fig . 4 . The work ( 21 ) proposes and analyzes another first-order momentum algorithm and derives convergence rates under a different set of conditions – the strong growth condition for the loss function in addition to convexity . As shown in Appendix F.3 , on the example of a Gaussian distributed data , the rates obtained in ( 21 ) can be slower than those for SGD . In contrast , our algorithm is guaranteed to never have a slower convergence rate than SGD . Furthermore , in the same Gaussian setting MaSS matches the optimal accelerated full-gradient Nesterov rate . Additionally , in our work we consider the practically important dependence of the convergence rate and optimal parameter selection on the mini-batch size , which to the best of our knowledge , has not been analyzed for momentum methods . 2 NOTATIONS AND PRELIMINARIES . Given dataset D = { ( xi , yi ) } ni=1 ⊂ Rd × C , we consider an objective function of the form f ( w ) = 1 n ∑n i=1 fi ( w ) , where fi only depends on a single data point ( xi , yi ) . Let ∇f denote the exact gradient , and ∇̃mf denote the unbiased stochastic gradient evaluated based on a mini-batch of size m. For simplicity , we also denote ∇̃f ( w ) : = ∇̃1f ( w ) .We use the concepts of strong convexity and smoothness of functions , see definitions in Appendix B.1 . For loss function with µ-strong convexity and L-smoothness , the condition number κ is defined as κ = L/µ . In the case of the square loss , fi ( w ) = 12 ( w Txi − yi ) 2 , and the Hessian matrix is H : = 1 n ∑n i=1 xix T i . Let L and µ be the largest and the smallest non-zero eigenvalues of the Hessian respectively . Then the condition number is then κ = L/µ ( note that zero eigenvalues can be ignored in our setting , see Section 4 ) . Stochastic Condition Numbers . For a quadratic loss function , let H̃m : = 1m ∑m i=1 x̃ix̃ T i denotes a mini-batch estimate of H . Define L1 be the smallest positive number such that E [ ‖x̃‖2x̃x̃T ] L1H , and denote Lm : = L1/m+ ( m− 1 ) L/m . ( 2 ) Given a mini-batch size m , we define the m-stochastic condition number as κm : = Lm/µ . Following ( 8 ) , we introduce the quantity κ̃ ( called statistical condition number in ( 8 ) ) , which is the smallest positive real number such that E [ ‖x̃‖2H−1 x̃x̃ T ] κ̃H . Remark 1 . We note that L1 ≥ L , since E [ ‖x̃‖2x̃x̃T ] = E [ H̃21 ] = H2 + E [ ( H̃1 − H ) 2 ] H2 . Consequently , according to the definition of Lm in Eq.2 , we have Lm ≥ L , ∀m ≥ 1 ; Lm → L , as m→∞ . ( 3 ) Hence , the quadratic loss function is also Lm-smooth , for all m ≥ 1 . By the definition of κm , we also have κm ≥ κ , ∀m ≥ 1 ; κm → κ , as m→∞ . ( 4 ) Remark 2 . It is important to note that κ̃ ≤ κ1 , since E [ ‖x̃‖2H−1 x̃x̃ T ] 1µE [ ‖x̃‖2x̃x̃T ] κ1H . 2.1 CONVERGENCE OF SGD FOR OVER-PARAMETRIZED MODELS AND OPTIMAL STEP SIZE . We consider over-parametrized models that have zero training loss solutions on the training data ( e.g. , ( 23 ) ) . A solution fi ( w ) which fits the training data perfectly fi ( w ) = 0 , ∀i = 1 , 2 , · · · , n , is known as interpolating . In the linear setting , interpolation implies that the linear system { xTi w = yi } ni=1 has at least one solution . A key property of interpolation is Automatic Variance Reduction ( AVR ) , where the variance of the stochastic gradient decreases to zero as the weight w approaches the optimal w∗ . Var [ ∇̃mf ( w ) ] ‖w −w∗‖2E [ ( H̃m −H ) 2 ] . ( 5 ) For a detailed discussion of AVR see Appendix B.2 . Thanks to AVR , plain SGD with constant step size can be shown to converge exponentially for strongly convex loss functions ( 13 ; 19 ; 14 ; 12 ) . The set of acceptable step sizes is ( 0 , 2/Lm ) , where Lm is defined in Eq.2 and m is the mini-batch size . Moreover , the optimal step size η∗ ( m ) of SGD that induces fastest convergence guarantee is proven to be 1/Lm ( 12 ) .

Summary: The authors present a new first order optimization method that adds a corrective term to Nesterov SGD. They demonstrate that this adjustment is necessary and sufficient to benefit from the faster convergence of Nesterov gradient descent in the stochastic case. In the full-batch (non deterministic) setting, their algorithm boils down to the classical formulation of Nesterov GD. Their approach is justified by a well conducted theoretical analysis and some empirical work on toy datasets.