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 the problem of stochastic optimization where the gradient noise process is non-stationary. Based on a general convergence results based on a general sequence for the second moments of the stochastic gradient norms and a general stepsize sequence, the authors propose to use an online estimation procedure for the gradient norm second moments, in order to mimic the behavior of the ``idealized'' stepsize sequence. Finite-time convergence rates are established for the algorithms with adaptive stepsize, leading to an acceleration effect in certain regimes for the non-stationarity.