Stochastic Optimization with Non-stationary Noise: The Power of Moment Estimation

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is one of the most popular optimization methods in machine learning because of its computational efficiency compared to traditional full gradient methods . Great progress has been made in understanding the performance of SGD under different smoothness and convexity conditions ( Agarwal et al. , 2009 ; Arjevani et al. , 2019 ; Drori and Shamir , 2019 ; Ghadimi and Lan , 2012 ; 2013 ; Nemirovsky and Yudin , 1983 ; Rakhlin et al. , 2012 ) . These results show that with a fixed step size , SGD can achieve the minimax optimal convergence rate for both convex and nonconvex optimization problems , provided the gradient noise is uniformly bounded . Yet , despite the theoretical minimax optimality of SGD , adaptive gradient methods ( Duchi et al. , 2011 ; Kingma and Ba , 2014 ; Tieleman and Hinton , 2012 ) have become the methods of choice for training deep neural networks , and have received a surge of attention recently ( Agarwal et al. , 2018 ; Chen et al. , 2019 ; Huang et al. , 2019 ; Levy , 2017 ; Levy et al. , 2018 ; Li and Orabona , 2019 ; Liu et al. , 2019 ; 2020 ; Ma and Yarats , 2019 ; Staib et al. , 2019 ; Ward et al. , 2019 ; Zhang et al. , 2019 ; 2020 ; Zhou et al. , 2018 ; 2019 ; Zou and Shen , 2018 ; Zou et al. , 2019 ) . Instead of using fixed stepsizes , these methods construct their stepsizes adaptively using the current and past gradients . Despite advances in the literature on adaptivity , theoretical understanding of the benefits of adaptation is very limited . We provide a different perspective on understanding the benefits of adaptivity by considering it in the context of non-stationary gradient noise , i.e. , the noise intensity varies with iterations . Surprisingly , this setting is rarely studied , even for SGD . To our knowledge , this paper is the first work to formally study stochastic gradient methods in this varying noise scenario . Our main goal is to show that : Adaptive step-sizes can guarantee faster rates than SGD when the noise is non-stationary . We focus on this goal based on several empirical observations ( Section 2 ) , which lead us to model the noise of stochastic gradient oracles via the following iteration dependent quantities : m2k : = E [ ‖g ( xk ) ‖2 ] , or σ2k : = E [ ‖g ( xk ) −∇f ( xk ) ‖2 ] , ( 1 ) where g ( xk ) is the stochastic gradient and ∇f ( xk ) the true gradient at iteration k. Notation ( 1 ) provides more fine-grained description of noise behavior than uniform bounds on the variance by permitting iteration dependent noise intensities . It is intuitive that one should prefer smaller stepsizes when the noise is large and vice versa . Thus , under non-stationarity , an ideal algorithm should adapt its stepsize according to the parameters mk or σk , suggesting a potential benefit of adaptive stepsizes . Contributions . The primary contribution of our paper is to show that a stochastic optimization method with adaptive stepsize can achieve a faster rate of convergence ( by a factor that is polynomialin-T ) than fixed-step SGD . We first analyze an idealized setting where the noise intensities are known , using it to illustrate how to select noise dependent stepsizes that are provably more effective ( Theorem 1 ) . Next , we study the case with unknown noise , where we show under an appropriate smoothness assumption on the noise variation that a variant of RMSProp ( Tieleman and Hinton , 2012 ) can achieve the idealized convergence rate ( Theorem 3 ) . Remarkably , this variant does not require the noise levels . Finally , we generalize our results to nonconvex settings ( Theorems 12 and 14 ) . 2 MOTIVATING OBSERVATION : NONSTATIONARY NOISE IN DEEP LEARNING . Neural network training involves optimizing an empirical risk minimization problem of the form minx f ( x ) : = 1 n ∑n i=1 fi ( x ) , where each fi represents the loss function with respect to the i-th data or minibatch . Stochastic methods optimize this objective randomly sampling an incremental gradient ∇fi at each iteration and using it as an unbiased estimate of the full gradient . The noise intensity of this stochastic gradient is measured by its second moments or variances , defined as , 1 . Second moment : m2 ( x ) = 1n ∑n i=1 ‖∇fi ( x ) ‖2 ; 2 . Variance : σ2 ( x ) = 1n ∑n i=1 ‖∇fi ( x ) −∇f ( x ) ‖2 , where∇f ( x ) is the full gradient , To illustrate how these quantities evolve over iterations , we empirically evaluate them on three popular tasks of neural network training : ResNet18 training on Cifar10 dataset for image classification1 , LSTM training on PTB dataset for language modelling2 ; transformer training on WMT16 en-de for language translation3 . The results are shown in Figure 1 , where both the second moments and variances are evaluated using the default training procedure of the original code . On one hand , the variation of the second moment/variance has a very different shape in each of the considered tasks . In the CIFAR experiment , the noise intensity is quite steady after the first iteration , indicating a fast convergence of the training model . In LSTM training , the noise level increases and converges to a threshold . While , in training Transformers , the noise level increases very fast at the early epochs , then reaches a maximum , and turns down gradually . On the other hand , the preferred optimization algorithms in these tasks are also different . For CIFAR10 , SGD with momentum is the most popular choice . While for language models , adaptive methods such as Adam or RMSProp are the rule of thumb . This discrepancy is usually taken as granted , based on empirical validation ; and little theoretical understanding of it exists in the literature . 1Code source for CIFAR10 https : //github.com/kuangliu/pytorch-cifar 2Code source for LSTM https : //github.com/salesforce/awd-lstm-lm 3Code source for Transformer https : //github.com/jadore801120/attention-is-all-you-need-pytorch Based on the observations made in Figure 1 , a natural candidate emerges to explain this discrepancy in the choice of algorithms : the performance of stochastic algorithms varies according the characteristics of gradient noise encountered during training . Despite this behavior , noise level modeling has drawn surprisingly limited attention in prior art . Reference ( Moulines and Bach , 2011 ) studies convergence of SGD assuming each component function is convex and smooth ; extensions to the variation of the full covariance matrix are in ( Gadat and Panloup , 2017 ) . A more fine-grained stochastic oracle assumes that the variances grow with the gradient norm as σ2 + c‖∇f ( x ) ‖2 , or grow with the suboptimality σ2 + c‖x−x∗‖2 ( Bottou et al. , 2018 ; Jofré and Thompson , 2019 ; Rosasco et al. , 2019 ) . Unfortunately , these existing oracles fail to express the variation of noise observed in Figure 1 . Indeed , the norm of the full gradient , represented as the difference between the orange and the blue line , is significantly smaller compared to the noise level . This suggests that noise variation is not due to the gradient norm , but due to some implicit properties of the objective function . This observation motivates us to introduce the following non-stationary noise oracle : Definition 1 ( non-stationary noise oracle ) . The stochasticity of the problem is governed by a sequence of second moments { mk } k∈N or variances { σk } k∈N , such that , at the kth iteration , the gradient oracle returns an unbiased gradient g ( xk ) such that E [ g ( xk ) ] = ∇f ( xk ) , and either ( a ) with second moment E [ ‖g ( xk ) ‖2 ] = m2k ; or ( b ) with variance E [ ‖g ( xk ) −∇f ( xk ) ‖2 ] = σ2k . The non-stationary noise oracle is a relaxation of the standard uniform noise oracle in which case mk or σk are constant . By introducing the time dependency , we aim to understand how the variation of noise influences the convergence rate of optimization algorithms . An example that falls into this category is when the noise is additive to the gradient , namely g ( xk ) ∼ ∇f ( xk ) +N ( 0 , σ2k ) . We emphasize that our goal is to demystify the correlation between the noise intensity and the performance of optimization algorithm , instead of explaining why certain shape of noise occurs . In general , the variation in noise is a consequence of the combination on data distribution , training model and optimization method , which is complex and highly non trivial . We simplify it by assuming that the noise intensity is decoupled from its location , meaning that , the parameters mk or σk only depend on the iteration number k , but do not depend on the specific location where the gradient is evaluated . This is empirically justified as the pattern of the noise is mostly determined by the task instead of the optimization algorithms , see Appendix A . The simplification helps to focus on the shape of noise , taking a first step towards the goal : characterize the convergence rate of adaptive algorithms under non-stationary noise . 3 THE BENEFIT OF ADAPTIVITY UNDER NONSTATIONARY NOISE . In this section , we investigate the influence of nonstationary noise in an idealized setting where the noise parameters mk are known . For brevity , we will first focus on the convex setting and present our results based on the second moment parameters mk . We defer the discussion on nonconvex problems and variants on variance parameters σk to Section 5 . One reason that we prioritize the second moment than the variance is to draw a connection with the well-known adaptive method RMSProp ( Tieleman and Hinton , 2012 ) and Adam ( Kingma and Ba , 2014 ) . One common feature shared by both algorithms is that they scale the step sizes inversely to an exponential moving average of estimated second moments . Below , we start from the idealized case assuming the second moments are known and show why inverse scaling could speed up convergence . Let f be convex and differentiable . We consider the problem minx f ( x ) , where the gradient is given by the nonstationary noise oracle defined in Definition 1 . We assume that the optimum is attained at x∗ and we denote f∗ the minimum of the objective . We are interested in studying the convergence rate of a stochastic algorithm with update rule xk+1 = xk − ηkg ( xk ) , ( 2 ) where ηk are stepsizes that are oblivious of the iterates { xk } k∈N . Theorem 1 . Under the second moment oracle given in Definition 1 ( a ) , the weighted average xT = ( ∑T k=1 ηkxk ) / ( ∑T k=1 ηk ) obtained by the update rule ( 2 ) satisfies the suboptimality bound E [ f ( xT ) − f∗ ] ≤ ‖x1 − x∗‖2 + ∑T k=1 η 2 km 2 k∑T k=1 ηk . ( 3 ) Although the theorem follows from standard analysis , it leads to valuable results as explained below . Corollary 2 . Denote M = maxmk and R = ‖x1 − x∗‖ . We have the following convergence rate : 1 . SGD with constant stepsize : if ηk = η = R√∑T k=1m 2 k , then E [ f ( xT ) − f∗ ] ≤ 2R √∑T k=1m 2 k T = 2RM√ T · √ 1 T ∑T k=1m 2 k M2 . ( constant baseline ) 2 . SGD with idealized stepsize : if ηk = R√Tmk , then E [ f ( xT ) − f∗ ] ≤ 2R √ T∑T k=1 1 mk = 2RM√ T · 1 M 1 T ∑T k=1 1 mk . ( idealized baseline ) To facilitate comparison , we have normalized the convergence rates with respect to the conventional rate 2RM/ √ T ( Nemirovski et al. , 2009 ) . In the standard setting , the values of mk are unavailable , but the uniform bound M is known , in such a case taking mk = M recovers the standard result . When the values of mk are given , both constant baseline and idealized baseline benefits from it , improving upon the conventional rate 2RM/ √ T . The improvement factor in constant baseline has depends on the average of the second moments √∑ m2k/T , while as the improvement factor in idealized baseline depends on the average of the harmonic sum 1T ∑ 1/mk . In particular , from Jensen ’ s inequality E [ X ] −2 ≤ E [ X−2 ] , we have ( 1T ∑ k 1 mk ) −2 ≤ 1T ∑ km 2 k , implying that the idealized baseline is always better than the constant baseline . This result is rather expected , as the stepsizes are adapted to the noise intensity . As a consequence , the accumulations of the parameters mk in different forms governed the convergence rate . To further illustrate such difference , we consider an illustrative synthetic noise model , mimicking the shape of noise we observed in the training of Transformer ( see Figure 1 ( c ) ) . Example 1 . Consider the following piece-wise linear noise model with γ = 5 ( 1− T−α ) /T . mk = 1 Tα if k ∈ [ 1 , T 5 ] ; γ ( k − 2T5 ) + 1 if k ∈ ( T 5 , 2T 5 ] ; 1 if k ∈ ( 2T5 , 3T 5 ] ; γ ( 3T5 − k ) + 1 if k ∈ ( 3T 5 , 4T 5 ] ; 1 Tα if k ∈ ( 4T 5 , T ] . 0.0 0.2 0.4 0.6 0.8 1.0 Iteration # 1e6 0.0 0.2 0.4 0.6 0.8 1.0 m k alpha = 1/9 alpha = 1/6 alpha = 1/4 In this example , the maximum noise intensity M is 1 , and the minimum intensity is 1/Tα , inducing a large ratio of order Tα . Following the bounds developed in Corollary 2 , the performance of the constant baseline maintains the standard O ( 1/ √ T ) convergence rate , while as the idealized baseline converges at O ( 1/T 1 2 +α ) . Hence a nontrivial acceleration of order Tα is obtained by using the idealized stepsize , and this acceleration can be arbitrarily large as α increases . This example is encouraging , showing that the speedup due to adaptive stepsizes can be polynomial in the number of iterations , especially when the ratio between the maximum and the minimum noise intensity is large . However , explicit knowledge on mk is required to implement these idealized stepsizes , which is unrealistic . The next sections of the paper demonstrates that estimating the moment bound in an online fashion can achieve a convergence rate comparable to the idealized setting .

Summary: This paper studies gradient-based stochastic optimization algorithms which incorporate (estimates of) the noise statistics in the adaptive stepsize design. Starting from the standard analysis of SGD with adaptive steps (Thm 1) the authors show in Cor 1 how, using a second-moment-dependent learning rate, one can “accelerate” (see comment later) the convergence of SGD. Next, the authors show (Thm 3) that one can recover a similar result by estimating the second moment using an exponential moving average.