summaries: 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. 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. The objective of the paper is to provide a theoretical justification for the value of using adaptive learning steps. The paper presents two results. The first is essentially of theoretical interest, and assumes that the noise level indicators defined in Eq. (1) [but which are difficult to understand at this level of the paper] are known. The second is more practical: it shows that a variant of the RMSprop algorithm achieves the same results as the "theoretical" algorithm.
abstractText: We investigate stochastic optimization under weaker assumptions on the distribution of noise than those used in usual analysis. Our assumptions are motivated by empirical observations in training neural networks. In particular, standard results on optimal convergence rates for stochastic optimization assume either there exists a uniform bound on the moments of the gradient noise, or that the noise decays as the algorithm progresses. These assumptions do not match the empirical behavior of optimization algorithms used in neural network training where the noise level in stochastic gradients could even increase with time. We address this nonstationary behavior of noise by analyzing convergence rates of stochastic gradient methods subject to changing second moment (or variance) of the stochastic oracle. When the noise variation is known, we show that it is always beneficial to adapt the step-size and exploit the noise variability. When the noise statistics are unknown, we obtain similar improvements by developing an online estimator of the noise level, thereby recovering close variants of RMSProp (Tieleman and Hinton, 2012). Consequently, our results reveal why adaptive step size methods can outperform SGD, while still enjoying theoretical guarantees.
references.authors: A. Agarwal M.J. Wainwright P.L. Bartlett P.K. Ravikumar
references.title: Information-theoretic lower bounds on the oracle complexity of convex optimization
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: N. Agarwal B. Bullins X. Chen E. Hazan K. Singh C. Zhang Y. Zhang
references.title: The case for full-matrix adaptive regularization
references.venue: arXiv preprint arXiv:1806.02958,
references.authors: Y. Arjevani Y. Carmon J.C. Duchi D.J. Foster N. Srebro B. Woodworth
references.title: Lower bounds for non-convex stochastic optimization
references.authors: M.G. Azar I. Osband R. Munos
references.title: Minimax regret bounds for reinforcement learning
references.venue: In International Conferences on Machine Learning (ICML),
references.authors: O. Besbes Y. Gur A. Zeevi
references.title: Stochastic multi-armed-bandit problem with non-stationary rewards
references.venue: In Advances in neural information processing systems,
references.authors: L. Bottou F.E. Curtis J. Nocedal
references.title: Optimization methods for large-scale machine learning
references.venue: Siam Review,
references.authors: S. Bubeck
references.title: Convex optimization: Algorithms and complexity
references.venue: arXiv preprint arXiv:1405.4980,
references.authors: X. Chen S. Liu R. Sun M. Hong
references.title: On the convergence of a class of ADAM-type algorithms for non-convex optimization
references.venue: In International Conference on Learning Representations (ICLR),
references.authors: Y. Drori O. Shamir
references.title: The complexity of finding stationary points with stochastic gradient descent
references.venue: arXiv preprint arXiv:1910.01845,
references.authors: J. Duchi E. Hazan Y. Singer
references.title: Adaptive subgradient methods for online learning and stochastic optimization
references.venue: Journal of Machine Learning Research,
references.authors: C. Fang C.J. Li Z. Lin T. Zhang
references.title: Spider: Near-optimal non-convex optimization via stochastic path integrated differential estimator
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: S. Gadat F. Panloup
references.title: Optimal non-asymptotic bound of the Ruppert-Polyak averaging without strong convexity
references.venue: arXiv preprint arXiv:1709.03342,
references.authors: S. Ghadimi G. Lan
references.title: Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization i: A generic algorithmic framework
references.venue: SIAM Journal on Optimization,
references.authors: S. Ghadimi G. Lan
references.title: Stochastic first-and zeroth-order methods for nonconvex stochastic programming
references.venue: SIAM Journal on Optimization,
references.authors: K. He X. Zhang S. Ren J. Sun
references.title: Deep residual learning for image recognition
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: H. Huang C. Wang B. Dong
references.title: Nostalgic Adam: weighting more of the past gradients when designing the adaptive learning rate
references.venue: In Proceedings of the 28th International Joint Conference on Artificial Intelligence,
references.authors: A. Jadbabaie A. Rakhlin S. Shahrampour K. Sridharan
references.title: Online optimization: Competing with dynamic comparators
references.venue: In Artificial Intelligence and Statistics,
references.authors: C. Jin Z. Allen-Zhu S. Bubeck M.I. Jordan
references.title: Is q-learning provably efficient
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: C. Jin P. Netrapalli M.I. Jordan
references.title: Accelerated gradient descent escapes saddle points faster than gradient descent
references.venue: In Conference On Learning Theory,
references.authors: A. Jofré P. Thompson
references.title: On variance reduction for stochastic smooth convex optimization with multiplicative noise
references.venue: Mathematical Programming,
references.authors: D.P. Kingma J. Ba
references.title: Adam: A method for stochastic optimization
references.venue: arXiv preprint arXiv:1412.6980,
references.authors: K. Levy
references.title: Online to offline conversions, universality and adaptive minibatch sizes
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: K.Y. Levy A. Yurtsever V. Cevher
references.title: Online adaptive methods, universality and acceleration
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: X. Li F. Orabona
references.title: On the convergence of stochastic gradient descent with adaptive stepsizes
references.venue: In The 22nd International Conference on Artificial Intelligence and Statistics,
references.authors: L. Liu H. Jiang P. He W. Chen X. Liu J. Gao J. Han
references.title: On the variance of the adaptive learning rate and beyond
references.authors: M. Liu Y. Mroueh J. Ross W. Zhang X. Cui P. Das T. Yang
references.title: Towards better understanding of adaptive gradient algorithms in generative adversarial nets
references.venue: In International Conference on Learning Representations (ICLR),
references.authors: J. Ma D. Yarats
references.title: On the adequacy of untuned warmup for adaptive optimization
references.venue: arXiv preprint arXiv:1910.04209,
references.authors: S. Merity N.S. Keskar R. Socher
references.title: Regularizing and optimizing LSTM language models
references.venue: In International Conference on Learning Representations (ICLR),
references.authors: A. Mokhtari S. Shahrampour A. Jadbabaie A. Ribeiro
references.title: Online optimization in dynamic environments: Improved regret rates for strongly convex problems
references.venue: IEEE 55th Conference on Decision and Control (CDC),
references.authors: E. Moulines F.R. Bach
references.title: Non-asymptotic analysis of stochastic approximation algorithms for machine learning
references.venue: In Proceedings of Advances in Neural Information Processing Systems (NIPS),
references.authors: A. Nemirovski A. Juditsky G. Lan A. Shapiro
references.title: Robust stochastic approximation approach to stochastic programming
references.venue: SIAM Journal on Optimization,
references.authors: A.S. Nemirovsky D.B. Yudin
references.title: Problem complexity and method efficiency in optimization
references.authors: A. Rakhlin O. Shamir K. Sridharan
references.title: Making gradient descent optimal for strongly convex stochastic optimization
references.venue: In International Conferences on Machine Learning (ICML),
references.authors: L. Rosasco S. Villa B.C. Vũ
references.title: Convergence of stochastic proximal gradient algorithm
references.venue: Applied Mathematics & Optimization,
references.authors: M. Staib S.J. Reddi S. Kale S. Kumar S. Sra
references.title: Escaping saddle points with adaptive gradient methods
references.venue: In International Conferences on Machine Learning (ICML),
references.authors: A.L. Strehl M.L. Littman
references.title: An analysis of model-based interval estimation for markov decision processes
references.venue: Journal of Computer and System Sciences,
references.authors: T. Tieleman G. Hinton
references.title: Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude
references.venue: COURSERA: Neural networks for machine learning,
references.authors: R. Ward X. Wu L. Bottou
references.title: Adagrad stepsizes: Sharp convergence over nonconvex landscapes
references.venue: In International Conferences on Machine Learning (ICML),
references.authors: J. Zhang S.P. Karimireddy A. Veit S. Kim S.J. Reddi S. Kumar S. Sra
references.title: Why adam beats sgd for attention models
references.authors: J. Zhang T. He S. Sra A. Jadbabaie
references.title: Why gradient clipping accelerates training: A theoretical justification for adaptivity
references.venue: In International Conference on Learning Representations,
references.authors: D. Zhou Y. Tang Z. Yang Y. Cao Q. Gu
references.title: On the convergence of adaptive gradient methods for nonconvex optimization
references.venue: arXiv preprint arXiv:1808.05671,
references.authors: Z. Zhou Q. Zhang G. Lu H. Wang W. Zhang Y. Yu
references.title: Adashift: Decorrelation and convergence of adaptive learning rate methods
references.venue: In International Conference on Learning Representations (ICLR),
references.authors: F. Zou L. Shen
references.title: On the convergence of weighted adagrad with momentum for training deep neural networks
references.venue: arXiv preprint arXiv:1808.03408,
references.authors: F. Zou L. Shen Z. Jie W. Zhang W. Liu
references.title: A sufficient condition for convergences of adam and rmsprop
references.venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
sections.heading: 1 INTRODUCTION
sections.text: 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).
sections.heading: 2 MOTIVATING OBSERVATION: NONSTATIONARY NOISE IN DEEP LEARNING
sections.text: 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.
sections.heading: 3 THE BENEFIT OF ADAPTIVITY UNDER NONSTATIONARY NOISE
sections.text: 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.
sections.heading: 4 ADAPTIVE METHODS: ONLINE ESTIMATION OF MOMENTS
sections.text: From now on, we assume that the moment bounds mk are not given. To address the non-stationarity, we estimate the noise intensity based on an exponential moving average, a technique commonly used in adaptive methods. More precisely, the moment estimator m̂k is constructed recursively as
  m̂2k+1 = βm̂ 2 k + (1− β)‖gk‖2, (ExpMvAvg)
  Algorithm 1 Adaptive SGD (x1, T, c,m)
  1: Initialize m̂21 = ‖g(x1)‖2. 2: for k = 1, 2, ..., T do 3: Evaluate stochastic gradient gk at xk. 4: xk+1 = xk − ηkgk with ηk = cm̂k+m . 5: m̂2k+1 = βm̂ 2 k + (1− β)‖gk‖2. 6: end for 7: return xT = ( ∑T i=1 ηixi)/( ∑T i=1 ηi).
  where gk is the k-th stochastic gradient and β is the decay paramter. Then we choose the stepsizes inversely proportional to m̂k+1, leading to Algorithm 1.
  Algorithm 1 could be viewed as a “norm" version of RMSProp (Tieleman and Hinton, 2012): the exponential moving average is performed coordinate-wise in RMSProp, whereas we use the full norm of gk to update the moment estimator m̂k+1. Such a simplification via a full norm variant has also been analyzed in the uniformly bounded noise setting (Levy, 2017; Levy et al., 2018; Li and Orabona, 2019; Ward et al., 2019)—we leave investigation of the more advanced coordinate-wise version as a topic of future research.
  Another important component of the stepsize is the correction constant m, the term appearing in the denominator. This constant provides a safety threshold when m̂k underestimates mk, which is commonly used in the practical implementation of adaptive methods, and even beyond, in reinforcement learning as so-called exploration bonus (Azar et al., 2017; Jin et al., 2018a; Strehl and Littman, 2008).
  To show the convergence of the algorithm, we need to impose a regularity assumption on the sequence of noise intensities. Otherwise, previous estimate may not provide any information of the next one. Assumption 1. We assume that an upper bound M on mk is given, i.e. maxkmk ≤M such that
  (a) The fourth moment of gk is bounded by M4, namely, E[(‖gk‖2 −m2k)2] ≤M4,∀k. (b) The total variation on mk is bounded, namely ∑ k |m2k −m2k+1| ≤ D2 with D2 = Ω(M2).
  The bounded fourth moment ensures concentration of ‖gk‖2, which is necessary to guarantee the quality of the online estimator. In particular, this assumption is satisfied when ‖gk‖ follows a mk sub-Gaussian distribution. Though stronger than bounded variance, this assumption does not lead to better convergence rates for SGD, as many existing results (including lower bounds) require sub-Gaussian or bounded noise, see Agarwal et al. (2009); Ghadimi and Lan (2013) .
  The bounded total variation assumption should be viewed as a regularity condition on the sequence of noise. It is motivated and common in the dynamic online learning literature (Besbes et al., 2014; Jadbabaie et al., 2015; Mokhtari et al., 2016). A key aspect of it is to avoid infinite oscillation, such as the pathological setting where m2k = 1 and m2k+1 = M , in which case the total variation scales with the number of iterations T . The specific constant in D2 = Ω(M2) depends on the shape of the noise. When mk is increasing in the first half and decreasing in the second half, as in the Transformer experiments and Example 1, the total variation is bounded by D2 ≤ 2M2. More generally, if the noise can be decomposed into K piece-wise monotone fragments, then the bound D2 ≤ KM2 holds. With the above assumptions, we are now ready to present our convergence result. Theorem 3. Under Assumptions 1, with probability at least 1/2, the iterates generated by Algorithm 1 using parameters β = 1− 2T−2/3, m = 4 √ D2 +M2T− 1 9 ln(T ) 1 2 , c = R√
  T satisfy
  f(xT )− f∗ ≤ 2RM√T · 32 M
  1 T ∑T k=1 1 mk+m .
  Remark 4. Our result directly implies a 1− δ high probability style convergence rate, by restarting it 2 log(1/δ) times. An additional log(1/δ) dependency will be introduced in the complexity, as in standard high probability results Fang et al. (2018); Jin et al. (2018b); Nemirovski et al. (2009).
  The key to prove the theorem is to effectively bound the estimation error |m̂2k − m2k| relying on concentration, and on bounded variation in Assumption 1. In particular, the choice of the decay
  parameter β is critical, determining how fast the contribution of past gradients decays. Because of the non-stationarity in noise, the online estimator m̂k is biased. The proposed choice of β carefully balances the bias error and the variance error, leading to sublinear regret, see Appendix C.
  Due to the correction constant m, the obtained convergence rate inversely depends on ∑T k=1 1 mk+m ,
  instead of the idealized dependency ∑T k=1 1 mk
  . This additional term makes the comparison less straightforward and we now discuss different scenarios for obtaining a better understanding.
sections.heading: 4.1 DISCUSSION OF THE CONVERGENCE RATE
sections.text: To illustrate the difference between convergence rates, we first consider the synthetic noise model introduced in Example 1. The detailed comparison is presented in Table 1, where we observe two regimes regarding the exponent α:
  • When 0 ≤ α ≤ 19 , the rate of the adaptive algorithm matches (idealized baseline) up to logarithmic dependency, and is Tα better than the (constant baseline).
  • When 19 ≤ α, the adaptive convergence rate no longer matches the (idealized baseline). Nevertheless, it is always T 1 9 faster than the (constant baseline).
  In both cases, the adaptive method achieves a non-trivial improvement, polynomial in T , compared to the (constant baseline). Even though the improvement T 1 9 might seem in-significant, it is the first result showing a plausible non-trivial advantage of adaptive methods over SGD under nonstationarynoise. Further, note that the adaptive convergence rate does not always match the (idealized baseline) when α is large. Such a discrepancy comes from the correction term m, which makes the stepsize more conservative than it should be, especially when mk is small.
  The above comparison relies on the specific choice on the noise model given in Example 1. Now we formalize some simple conditions allowing comparison in more general settings.
  Corollary 5. If the ratio M/(minkmk) ≤ T 1 9 , then adaptive method converges in the same order as the (idealized baseline), up to logarithmic dependency.
  This result is remarkable since the adaptive method does not require any knowledge of mk values, and yet it achieves the idealized rate. In other words, the exponential moving average estimator successfully adapts to the variation in noise, allowing faster convergence than constant stepsizes.
  Corollary 6. Let m2avg = ∑ m2k/T be the average second moment. If M/mavg ≤ T 1 9 , then adaptive method is no slower than the (constant baseline), up to logarithmic dependency.
  The condition in Corollary 6 is strictly weaker than the condition in Corollary 5, which means even though an adaptive method may not match the idealized baseline, it could still be non-trivially better than the constant baseline. This case happens e.g., when α > 19 in Table 1, where the adaptive method is O(T 19 ) faster than the constant baseline. Indeed, O(T 19 ) is the maximum improvement one can expect according to our current analysis.
  Corollary 7. Recall that M is an upper bound on mk, i.e., maxmk ≤M . Therefore 1. The convergence rate of the constant baseline is no slower than O(2RM/ √ T ).
  2. The convergence rate of the adaptive method is no faster than Õ(2RM/T 12 + 19 ).
  The order of maximum improvement O(T 19 ) is determined by the specific choice of m in Theorem 3, which is chosen to be Õ(MT− 19 ). Indeed, the correction term is helpful when the estimator m̂k underestimates the true value mk, avoiding the singularity at zero. Hence, the choice of m is related
  to the average deviation between m̂k and mk. Under a stronger concentration assumption, we can strengthen the maximum improvement to O(T 16 ), as shown in Appendix E. The noise model in Example 1 provides a favorable scenario where the maximum improvement is attained. However, in some scenarios, the convergence rate of an adaptive method can be slower than the constant baseline. Adversarial scenario. If mk = 1/Tα for all i ∈ [1, T ] except at T/2 it takes the value mT/2 = 1 with α > 1/9, then the convergence rate of both constant and idealized baselines are O(T− 1 2 ), while the adaptive method only converges in Õ(T− 1−2α2 ). The subtle change at iteration T/2 amplifies the exponential moving average estimator and requires a non-negligible period to get back to the constant level. It is clear that the estimator becomes less meaningful under such a subtle change.
  Overall, it is hard to provide a complete characterization of the variation in noise. In Corollary 5 and 6, we show that when the ratio between the maximum and the minimum/average second moment is not growing too fast, adaptive methods do improve upon SGD.
sections.heading: 5 EXTENSIONS OF THM 3
sections.text: In this section, we discuss several extensions to Thm 3. The results are nontrivial but the analysis is almost the same. Hence we defer the exact statements and proofs to appendices. Addressing the variance oracle. So far, we have focused on the noise oracle based on the second moment mk and made the connection with existing adaptive methods. However, there is some unnaturalness underlying the non-stationary oracle on mk. Indeed, it is hard to argue that mk is iterate independent since m2k = σ 2 k + ‖∇f(xk)‖2. Even though the influence of ‖∇f(xk)‖2 might be minor when the variance σ2k is high (e.g. as in Figure 1), it is still changing mk. In contrast, the variance σk is an intrinsic quantity coming from the noise model, which could be iterate independent. Hence the variance oracle is theoretically more sound. We now present the necessary modifications in order to adapt to the variance oracle.
  First, in order to estimate the variance, we need to query two stochastic gradients gk and g′k at the same iterate, then we construct the estimator following the recursion
  σ̂2k+1 = βσ̂ 2 k + (1− β)‖gk − g′k‖2.
  Second, the smoothness condition on f is required, i.e., L-Lipschitzness on the gradient of f . In this case, it is necessary to ensure that the step-size being not larger than 1/2L. This translates to an additional constraint on the correcting constant m. More precisely, the stepsize is given by
  ηk = c σ̂k+m with m ≥ 2cL.
  Note that the L-smoothness condition is not required in the second moment oracle. This is why the second moment oracle is more suitable to nonsmooth setting (see Section 6.1 of Bubeck (2014)). A complete algorithm for the variance oracle is provided in Algorithm 2. The convergence results are essentially the same by replacing mk with σk, see Appendix H. Extension to nonconvex setting. We also provide an extension of our analysis to non-convex smooth setting. In which case, we characterize the convergence with respect to the gradient norm ‖∇f(xk)‖2, i.e., convergence to a stationary point. The conclusions are very similar to the one in the convex setting and the results (Theorems 12 and 14) are deferred to Appendix F. Variants on stepsizes. To go beyond the second moment of noise, one could apply an estimator of the form m̂pk+1 = βm̂ p k + (1− β)‖gk‖p when the p-th moment of the gradient is bounded. This allows stepsize of the shape ηk ∝ 1/(m̂pk +mp) 1/p as in ADAM, Adamax Kingma and Ba (2014).
sections.heading: 6 EXPERIMENTS
sections.text: In this section, we describe two sets of experiments that verify the faster convergence of Algorithm 1 against vanilla SGD. The first experiment is on linear regression with synthetic noise described in Example 1 and the second set of experiments is on neural network training.
sections.heading: 6.1 SYNTHETIC EXPERIMENTS
sections.text: In the synthetic experiment, we generate a random linear regression dataset using the sklearn4 library. We design the stochastic oracle as full gradient with injected Gaussian noise, whose coordinate-wise
  4scikit-learn.org/stable/modules/generated/sklearn.datasets.make_regression.html
  standard deviation σ is shown in the left figure of Fig 2. We then run the four algorithms discussed in this work: standard baseline, idealized baseline, Alg 1 and Alg 2. We finetune the step sizes for each algorithm by grid-searching among 10k, where k is an integer. We repeat the experiment for 10 runs and show the average training trajectory as well as the function suboptimality in Fig 2. We observe that the performance is ranked as follows: idealized baseline, Alg 2, Alg 1 and standard baseline.
sections.heading: 6.2 NEURAL NETWORK TRAINING
sections.text: We demonstrate how the proposed algorithm performs in real-world neural network training. We first tested our algorithm on Cifar10 classification task. We then implement our algorithm into the AWD-LSTM codebase described in Merity et al. (2018). We see from Figure 3 that our proposed algorithm can achieve slightly better performance than baselines. Despite our main contribution being providing theoretical analysis for the fast convergence of adaptive methods with moment estimation, these results show that our analysis can also lead to efficient and practical algorithm design.
  Besides convergence, we also measured noise level during neural network training with different optimizers and found that the noise pattern is mostly determined by the learning task instead of by the optimization algorithm. Details can be found in Appendix A.
sections.heading: 7 CONCLUSIONS
sections.text: This paper discusses convergence rates of stochastic gradient methods in an empirically motivated setting where the noise level is changing over iterations. We show that under mild assumptions, one can achieve faster convergence than the fixed step SGD by a factor that is polynomial in number of iterations, by applying online noise estimation and using adaptive step sizes. Our analysis, therefore provides one explanation for the recent success of adaptive methods in neural network training.
  There is much more to be done along the line of non-stationary stochastic optimization. Under our current analysis, there is a gap between the adaptive method and the idealized method when the noise variation is large (see second row in Table 1). A natural question to ask is whether one could reduce this gap, or alternatively, is there any threshold preventing the adaptive method from getting arbitrarily close to the idealized baseline? Moreover, could one attain further acceleration by combining momentum or coordinate-wise update techniques? Answering these questions would provide more insight and lead to a better understanding of widely used adaptive methods.
  Perhaps a more fundamental question is regarding the iterate dependency: the setting where the momentsmk or the variance σk are functions of the current update xk, not just of the iteration index k. Significant effort needs to be spent to address this additional correlation under appropriate regularity conditions. We believe our work lays the foundation to address this challenging research problem.
sections.heading: A ADDITIONAL EXPERIMENT DETAILS AND RESULTS
sections.text: In this section, we provide more details on our experiment setup. The code for reproducing noise estimation and neural network training on AWS-LSTM model is uploaded in the supplementary. The code for Cifar10 classification can be uploaded upon request.
  A.1 DETAILS ON PTB TRAINING
  Our implementation is based on the author’s github repository 5. The original codebase trains the network using clipped gradient descent followed by an average SGD (ASGD) algorithm to prevent overfitting. As generalization error is beyond our discussion, we focus on the first phase (which takes about 200 epochs) by removing the ASGD training part. Aside of the number of epochs, all other parameters are the same as the default training procedure. For our algorithm, we update the clipping value to be 0.25, and set the hyperparameters of Algorithm 1 as ηk = 5, β = 0.95,m = 0.01.
  A.2 DETAILS ON CIFAR10 TRAINING
  Our implementation is based on a pytorch implementation 6 of the original ResNet paper(He et al., 2016). We train the ResNet18 model on Cifar10 dataset for 200 epochs. The baseline uses SGD with learning rate initialized as 0.1. The learning rate is decayed by 10 at epochs 100, 150. We eventually achieves approximately 95% validation accuracy. For our algorithm, we set ηk = 1, β = 0.99,m = 0.01 and use the same learning rate schedule as the baseline.
  A.3 NOISE BEHAVIOR FOR DIFFERENT ALGORITHMS
  In this subsection, we include one interesting observation on the noise pattern of neural network training. From Figure 1, we see that the noise pattern looks very different for different task. One may naturally wonder whether the difference results from the particular training task or from the optimization algorithms. Noticing that each of these three tasks are trained with a different algorithm (Cifar10 with momentum SGD; PTB with clipped SGD; En-De translation with ADAM), we retrained the Cifar10 task and the PTB language modelling with Adam. and rerun the En-De translation experiment with momentum SGD. We found that the En-De translation task diverges when trained with SGD even when learning rate is 0.001. Hence we only report the results from the other two experiments in Figure 4.
  From Figure 4, we see that the pattern looks very similar to the original plots in Figure 1. Hence these experiments suggest that the noise pattern is mostly determined by the learning task instead of by the optimization algorithms.
  5Code source for LSTM https://github.com/salesforce/awd-lstm-lm 6Code source for CIFAR10 https://github.com/kuangliu/pytorch-cifar
sections.heading: B PROOF OF THEOREM 1
sections.text: Proof. The iterate suboptimality have the following relationship: ‖xk+1 − x∗‖2 = ‖xk − ηkgk − x∗‖2 = ‖xk − x∗‖2 − 2ηk〈gk, xk − x∗〉+ η2k‖gk‖2. Rearrange and take expectation with respect to gk we have 2ηk(f(xk)− f∗) ≤ 2ηk〈∇f(xk), xk − x∗〉
  ≤ E‖xk+1 − x∗‖2 − E‖xk − x∗‖2 + η2km2k Sum over k and take expectation we get
  E[ T∑ k=1 2ηk(f(xk)− f∗)] ≤ ‖x1 − x∗‖2 + T∑ k=1 η2km 2 k
  Then from convexity, we have
  E[f(xT )− f∗] ≤ ‖x1 − x∗‖2 +
  ∑T k=1 η 2 km
  2 k∑T
  k=1 ηk ,
  where xT = ( ∑T i=1 ηixi)/( ∑T i=1 ηi). Corollary 2 follows from specifying the particular choices of the stepsizes.
sections.heading: C KEY LEMMA
sections.text: Lemma 8. Under Assumptions 1, taking β = 1− 2T−2/3, the total estimation error of the m̂2k based on (ExpMvAvg) is bounded by:
  E [ T∑ k=1 |m̂2k −m2k| ] ≤ 2(D2 +M2)T 2/3 ln(T 2/3)
  Proof. On a high level, we decouple the error in a bias term and a variance term. We use the total variation assumption to bound the bias term, and use the exponential moving average to reduce variance. Then we pick β to balance the two terms.
  From triangle inequality, we have T∑ k=0 E [ |m̂2k −m2k| ] ≤ T∑ k=1 E [ |m̂2k − E[m̂2k]| ]︸ ︷︷ ︸ Variance term + T∑ k=1 ∣∣E[m̂2k]−m2k∣∣︸ ︷︷ ︸ Bias term
  (4)
  We first bound the bias term. By definition of m̂k, we have E[m̂2k]−m2k = βE[m̂2k−1] + (1− β)m2k−1 −m2k
  = β(E[m̂2k−1]−m2k−1) + (m2k−1 −m2k) Hence by recursion,
  E[m̂2k]−m2k = βk−1 (E[m̂21]−m21)︸ ︷︷ ︸ =0 +βk−2(m21 −m22) + · · ·+ (m2k−1 −m2k)
  Therefore, the bias term could be bounded by T∑ k=1 ∣∣E[m̂2k]−m2k∣∣ ≤ T∑ k=1 k−1∑ j=1 βk−1−j ∣∣m2j −m2j+1∣∣
  = T−1∑ k=1 ∣∣m2k −m2k+1∣∣ T−1−k∑ j=0 βj
  ≤ 1 1− β T−1∑ k=1 ∣∣m2k −m2k+1∣∣ ≤ D 2
  1− β (From Assumption (1))
  The first inequality follows by traingle inequality. The third inequality uses the geometric sum over β. To bound the variance term, we remark that
  m̂2k = (1− β)g2k−1 + (1− β)βg2k−2 + · · ·+ (1− β)βk−2g21 + βk−1g20 .
  Hence from independence of the gradients, we have E [ |m̂2k − E[m̂2k]| ] ≤ √ Var[m̂2k]
  = √
  Var[(1− β)g2k−1] + Var[(1− β)βg2k−2] + · · ·+ Var[(1− β)βk−2g21 ] + Var[βk−1g20 ] ≤ √ (1− β)2 + (1− β)2β2 + · · ·+ (1− β)2β2(k−2) + β2(k−1)M2,
  where M2 is an upperbound on the variance. The first inequality follows by Jensen’s inequality. The second equality uses independence of gi given g1, ..., gi−1. The last inequality follows by assumption 1.
  We distinguish two cases, when k is small, we simply bound the coefficient by 1, i.e.√ (1− β)2 + (1− β)2β2 + · · ·+ (1− β)2β2(k−2) + β2(k−1) ≤ 1
  When k is large such that k ≥ 1 + γ, with γ = 12(1−β) ln( 1 1−β ), we have β 2(k−1) ≤ 1− β, thus√
  (1− β)2 + (1− β)2β2 + · · ·+ (1− β)2β2(k−2) + β2(k−1)
  ≤ √ (1− β)2 1− β2 + β2(k−1)
  ≤ √ (1− β)2 1− β2 + (1− β)
  ≤ √ 2(1− β)
  The second inequality follows by k ≥ 1 + γ, with γ = 12(1−β) ln( 1 1−β ). Therefore, when k ≥ 1 + γ,
  E [ |m̂2k − E[m̂2k]| ] ≤ √ 2(1− β)M
  Therefore, substitute in the above equation into the
  T∑ k=1 E [ |m̂2k − E[m̂2k]| ] = γ∑ k=1 E [ |m̂2k − E[m̂2k]| ] + T∑ k=γ+1 E [ |m̂2k − E[m̂2k]| ] ≤ (γ + (T − γ) √ 2(1− β))M2
  Summing up the variance term and the bias term yields,
  T∑ k=0 E [ |m̂2k −m2k| ] ≤ D 2 1− β + (γ + (T − γ) √ 2(1− β))M2 (5)
  Taking β = 1− T−2/3/2 yields,
  T∑ k=0 E [ |m̂2k −m2k| ] ≤ 2(D2 +M2)T 2/3 ln(T 2/3) (6)
sections.heading: D PROOF OF THEOREM 3
sections.text: On a high level, the difference between the adaptive stepsize and the idealized stepsize mainly depends on the estimation error |m̂2k −m2k|, which has a sublinear regret according to Lemma C. Then we carefully integrate this regret bound to control the derivation from the idealized algorithm, reaching the conclusion.
  Proof. By the update rule of xk+1, we have,
  ‖xk+1 − x∗‖2 = ‖xk − ηkgk − x∗‖2 = ‖xk − x∗‖2 − 2ηk〈gk, xk − x∗〉+ η2k‖gk‖2.
  Noting that the stepsize ηk is independent of gk, taking expectation with respect to gk conditional on the past iterates lead to
  2ηk(f(xk)− f∗) ≤ 2ηk〈∇f(xk), xk − x∗〉 = E[2ηk〈gk, xk − x∗〉|xk, · · · , x1] = −E[‖xk+1 − x∗‖2|xk, · · · , x1] + ‖xk − x∗‖2 + η2km2k.
  Recall that R = ‖x1 − x∗‖, taking expectation and sum over iterations k, we get
  E[2( ∑T k=1 ηk)(f(xT )− f∗)] ≤ R2 + E[ ∑T k=1 η 2 km 2 k].
  Hence by Markov’s inequality, with probability at least 3/4, 2( ∑T k=1 ηk)(f(xT )− f∗) ≤ 4E[2( ∑T k=1 ηk)(f(xT )− f∗)] ≤ 4(R2 + E[ ∑T k=1 η 2 km 2 k]). (7)
  Now we can upper bound the right hand side, indeed
  T∑ k=1 E[η2km2k] = c2 T∑ k=1 E [ m2k (m̂k +m)2 ]
  ≤ c2 (
  T∑ k=1 E [ m2k − m̂2k (m̂k +m)2 ] + T∑ k=1 E [ m̂2k (m̂k +m)2 ])
  ≤ c2 ( 1
  m2 T∑ k=1 E [ |m2k − m̂2k| ] + T
  )
  ≤ c2 ( (M2 +D2)T 2/3 ln(T 2/3)
  m2 + T
  ) ≤ 3c2T (8)
  The last inequality follows by the choice on m. Hence, from Eq. (7), we have with probability at least 3/4,
  2( ∑T k=1 ηk)(f(xT )− f∗) ≤ 4(R2 + 3c2T ) (9)
  Next, by denoting (x)+ = max(x, 0), we lower bound the left hand side,
  1
  c
  ∑ ηk = ∑ 1 m̂k +m
  = ∑ 1 mk +m + ∑( 1 m̂k +m − 1 mk +m ) ≥ ∑ 1
  mk +m − ∑ (m̂k −mk)+ (mk +m)(m̂k +m)
  ≥ ∑ 1 mk +m − ∑ (m̂k −mk)+√ mk +m ·m3/2
  ≥ ∑ 1 mk +m − ∑ 1 2 ( (m̂k −mk)2+ m3 +
  1
  mk +m ) = 1
  2 ∑ 1 mk +m − 1 2m3 ∑ (m̂k −mk)2+ (10)
  Finally, by Markov’s inequality, with probability 3/4∑ (m̂k −mk)2+ ≤ 4E[ ∑ (m̂k −mk)2+] ≤ 4E[ ∑ (m̂k −mk)2] ≤ 8(D2 +M2)T 2/3 ln(T 2/3).
  Following the choice of m = 4 √ D2 +M2T− 1 9 ln(T ) 1 2 , we have
  1
  2m3 T∑ k=1 (m̂k −mk)2+ ≤ T 4(M +m) ≤ 1 4 T∑ k=1
  1
  mk +m
  Consequently, together with (9) and (10), we know that with probability at least 1− 14 − 1 4 = 1/2,
  f(xT )− f∗ ≤ 4(R2 + 3c2T )∑
  k c
  4(mk+m)
  ≤ 2R√ T · 32T∑
  k 1
  (mk+m)
  , (11)
  where the last inequality follows by setting c = R√ T .
  Remark 9. For more general choices of stepsize ηk = 1(m̂pk+mp)1/p , the upper bound in Eq.(8) holds exactly as in the above proof, and the lower bound in Eq.(10) follows from∑
  ηk = ∑ 1
  (m̂pk +m p)1/p = ∑ 1
  (mpk +m p)1/p
  + ∑( 1
  (m̂pk +m p)1/p − 1 (mpk +m p)1/p ) = ∑ 1
  (mpk +m p)1/p
  + ∑ (mpk +mp)1/p − (m̂pk +mp)1/p
  (m̂pk +m p)1/p(mpk +m p)1/p ≥ ∑ 1
  (mpk +m p)1/p
  − ∑ (m̂k −mk)+
  (m̂pk +m p)1/p(mpk +m
  p)1/p (Minkowski inequality)
  ≥ ∑ 1
  (mpk +m p)1/p
  − ∑ (m̂k −mk)+
  m3/2(mpk +m p)1/2p
  ≥ 1 2 ∑ 1 (mpk +m p)1/p − 1 2m3 ∑ (m̂k −mk)2+.
sections.heading: E PROOF WITH CONCENTRATE NOISE
sections.text: In this section, we add additional constraints on noise concentrations.
  Assumption 2. The expected absolute value is not very different from the square root of the second moment, i.e. 2E[‖g(x)‖] ≥ √ E[‖g(x)‖2].
  The constant “2” in the assumption above is arbitrary and can be increased to any fixed constant. The above assumption is satisfied if g(x) follows Gaussian distribution. It is also satisfied if for some fixed constant γ, p(‖g(x)‖ ≥ r) ≤ γE[‖g(x)‖]3r−4, for all r ≥ γE[‖g(x)‖]. We assume that the total variation on the first moment is bounded.
  Assumption 3. We denote λk = E[‖gk‖], and assume that an upper bound M such that
  (a) The second moment of gk is bounded by M2, namely, E[‖gk‖2] ≤M2,∀k. (b) The total variation on the first moment λk is bounded by∑
  k |λk − λk+1| ≤ D = Ω(M). (12)
  Under this stronger assumption, we can perform our online estimator on the first moment E[‖gk‖] instead of the second moment, replacing line 6 of Algorithm 1 by
  m̂k+1 = βm̂k + (1− β)‖gk‖. (13)
  The theorem below shows the convergence rate of the new algorithm.
  Theorem 10. Under assumptions 2, 3,with m = 16(D + M)T−1/6 ln(T ), c = R√ T , Algorithm 1 with update rule (13) achieves convergence rate
  f(xT )− f∗ ≤ 2R√ T · 12T∑
  k 1
  (mk+m)
  , (14)
  With a better concentration of the online estimator, we could allow a less conservative correction constant m, in the order of MT− 1 6 . It is this parameter that controls the maximum attainable improvement compared to the constant baseline. Indeed, we again consider the noise example 1, given in Table 2. In this case, the adaptive method can obtain an improvement of order T 1 6 compared to the constant baseline, while as previously only T 1 9 is achievable.
  The proof of the result follows a similar routine as the proof of Theorem 3. We start by presenting an equivalent lemma of Lemma 8. Lemma 11. Under assumption 3, we can achieve the following bound on total estimation error.
  E[ T∑ k=1 |m̂k − λk|] ≤ 2(D +M)T 2/3 ln(T 2/3)
  Proof. The proof is the same as the proof of Lemma 8, by replacing the second m2k by the first moment λk = E[‖gk‖].
  Proof of Theorem 10. By Assumption 2, we can use first moment of gk to bound the second moment. Hence, Eq. (7) implies that with probability at least 3/4,
  2( ∑T k=1 ηk)(f(xT )− f∗) ≤ 4(R2 + E[4 ∑T k=1 η 2 kλ 2 k]). (15)
  Now we upper bound the right hand side, indeed T∑ k=1 E[η2kλ2k] = c2 T∑ k=1 E [ λ2k (m̂k +m)2 ]
  ≤ c2 (
  T∑ k=1 E [ λ2k − m̂2k (m̂k +m)2 ] + T∑ k=1 E [ m̂2k (m̂k +m)2 ])
  ≤ c2 (
  T∑ k=1 E [ (λk − m̂k)(λk + m̂k) (m̂k +m)2 ] + T )
  ≤ c2 ( 2M
  m2 T∑ k=1 E [|λk − m̂k|] + T
  )
  ≤ c2 ( 4(M +D)MT 2/3 ln(T 2/3)
  m2 + T
  ) ≤ 2c2T
  Hence by Markov inequality, with probability at least 3/4, 2( ∑T−1 k=0 ηk)(f(xI)− f∗) ≤ 4E[2( ∑T−1 k=0 ηk)(f(xI)− f∗)] ≤ 4(R2 + 2c2T )
  Next, we lower bound the left hand side, 1
  c
  ∑ ηk = ∑ 1 m̂k +m
  = ∑ 1 λk +m + ∑( 1 m̂k +m − 1 λk +m ) ≥ ∑ 1
  λk +m − ∑ (m̂k − λk)+ (m̂k +m)(λk +m)
  ≥ ∑ 1
  λk +m − 1 m2
  ∑ (m̂k − λk)+
  By Markov’s inequality and Lemma 11, with probability 3/4, we have∑ (m̂k − λk)2+ ≤ 4E[ ∑ |m̂k − λk|] ≤ 8(D +M)T 2/3 ln(T 2/3).
  Following the choice of m = 16(D +M)T− 1 6 ln(T ), we have
  1
  m2 E[(m̂k − λk)+] ≤
  T 2(M +m) ≤ 1 2 ∑ 1 λk +m
  Consequently, we know that with probability at least 1− 14 − 1 4 = 1/2,
  f(xT )− f∗ ≤ 4(R2 + 2c2T )∑
  k c
  2(λk+m)
  ≤ 2R√ T · 12T∑
  k 1
  (mk+m)
  , (16)
  by setting c = R√ T and the fact that λk ≤ mk.
  F VARIANCE ORACLE AND EXTENSION TO NONCONVEX SETTING
  In this section, we show that we can adapt our algorithm to the variance oracle in Definition 1 (b), where E[‖gk −∇f(xk)‖2] = σ2k. To avoid redundancy, we present the result in the nonconvex smooth setting. We make the following smoothness assumptions. Assumption 4. The function is L−smooth, i.e. for any x, y, ‖∇f(x)−∇f(y)‖ ≤ L‖x− y‖.
  Remark that the L-smoothness condition is not required in the second moment oracle. This is why the second moment assumption is usually imposed in the non-differentiable setting (see Section 6.1 of Bubeck (2014)). We first provide the convergences of SGD serving as the baselines. Theorem 12 (Nonconvex baselines). Under the variance oracle in Definition 1 (b) and Assumption 4, the convergence of SGD using update xk+1 = xk − ηkgk with ηk ≤ 12L satisfies
  E[‖∇f(xI)‖2] ≤ f(x1)− f∗ + L2
  ∑T k=1 η 2 kσ
  2 k∑T
  k=1 ηk , (17)
  where I is an random variable with P(I = i) ∝ ηi.
  This convergence result is very similar to the one shown in the convex setting in Theorem 1. Instead of bounding the function suboptimality, the upper bound is only available on the norm of the gradient, implying convergence to a stationary point. We remark that an additional requirement on the stepsize is required, namely ηk ≤ 1/2L. This is not surprising since taking a large stepsize will not be helpful due to the L-smoothness condition. Hence, we can not take a stepsize inversely depending on σk when the noise is small. This restriction makes the comparison on convergence rate less straightforward. To facilitate the discussion on convergence rate, we make an additional assumption on the lower bound of the variance σk.
  Algorithm 2 Variance Adaptive SGD (x1, T, c,m)
  1: Initialize σ̂1 = ‖g1−g′1‖ 2 2 , where g1, g ′ 1 are two independent stochastic gradients at x1. 2: for k = 1, 2, ..., T do 3: Query two independent stochastic gradient gk, g′k at xk. 4: Update xk+1 = xk − ηk(gk + g′k)/2 with ηt = cσ̂k+m and m ≥ 2cL. 5: Update σ̂2k+1 = βσ̂ 2 k + (1− β) ‖gk−g′k‖ 2 2 6: end for 7: return xI where I is the random variable such that P(I = i) ∝ ηi.
  Assumption 5. For any k ∈ [1, T ], σk ≥ √ 8L(f(x1)− f(x∗))/ √ T .
  We emphasize that the above condition is not necessary to derive convergence analysis, but only for the clarity in terms of comparison on different convergence rate. This assumption help us focus on the case when noise (instead of the shape of the deterministic function) dominates the convergence rate and determines the step size. In otherword, our step size choice under this setting satisfies ηk ≤ 1/2L, leading to the following convergence rate: Corollary 13. Let ∆ = f(x1)− f∗. We have the following two convergence rate bounds for SGD:
  1. SGD with constant stepsize: if ηk = η = √
  2∆ L ∑ k σ 2 k , then
  E[‖∇f(xI)‖2] ≤ √
  2L∆ ∑T k=1 σ 2 k T = √ 2L∆ T · √∑T k=1 σ 2 k T . (constant baseline)
  2. SGD with idealized stepsize: if ηk = 1σk √ 2∆ LT , then
  E[‖∇f(xI)‖2] ≤ √ 2LT∆∑T k=1 1 σk
  = √
  2L∆ T · T∑T k=1 1 σk . (idealized baseline)
  The resulting convergence rate are similar as the convex setting. We now modify the adaptive algorithm to make use of the variance oracle and the smoothness assumption. The through algorithm is presented in Algorithm 2. Particularly, we keep an exponential moving average of the variance instead of the moments, using two stochastic gradients gk and g′k at the same iterate,
  σ̂2k+1 = βσ̂ 2 k + (1− β)
  ‖gk − g′k‖2
  2 .
  In particular, E[‖gk − g′k‖2/2] is an unbiased estimator of σ2k. To provide the convergence analysis of the algorithm, we make the following assumptions on the variation of the variance. Assumption 6. We assume an upper bound on σ2k = E[‖gk −∇f(xk)‖2], i.e. maxk σk ≤M . We also assume that the total variation in σk is bounded. i.e. ∑ k |σ2k − σ2k+1| ≤ D2 = 4M2.
  With the above assumptions, algorithm 2 achieve the following convergence rate. Theorem 14. Under the assumptions 4, 6 and T large enough such that lnT ≤ T 1/3, algorithm 2 with c = √ 2∆ LT , m = 4 √ D2 +M2T−1/9 ln(T )1/2 + 2cL, achieves with probability 1/2,
  ‖∇f(xI)‖2 ≤ √ 2L∆
  T · 32T∑
  k 1
  (σk+m)
  .
  The above theorem is almost the same as Theorem 3, and hence all the remarks for Theorem 3 also applies in the nonconvex case.
sections.heading: G PROOF OF THEOREM 12
sections.text: Proof. By L-smoothness, we have
  f(xk+1) ≤ f(xk)− ηk〈gk,∇f(xk)〉+ Lη2k
  2 ‖gk‖2
  Rearrange and take expectation with respect to gk, we get( ηk −
  Lη2k 2
  ) ‖∇f(xk)‖2 ≤ f(xk)− f(xk+1) +
  Lη2k 2 E[‖gk −∇f(xk)‖2].
  From the condition ηk ≤ 12L , we have,
  ηk 2 ‖∇f(xk)‖2 ≤ f(xk)− f(xk+1) + L 2 η2kσ 2 k.
  Sum over k and take expectation,
  E[ T∑ k=1 ηk‖∇f(xk)‖2] ≤ f(x0)− f(x∗) + L 2 ∑ k η2kσ 2 k.
  Denote I as the random variable such that P(I = i) ∝ ηi. We know
  ( T∑ k=1 ηk)E[‖∇f(xI)‖2] ≤ f(x0)− f(x∗) + L 2 T∑ k=1 η2kσ 2 k.
  This yields the desired convergence rate in (17).
sections.heading: H PROOF OF THEOREM 14
sections.text: The proof is almost identical to the proof of Theorem 3. We start by presenting an equivalent theorem of Lemma 8 below.
  Lemma 15. Under assumption 6, we can achieve the following bound on total estimation error using the estimator (ExpMvAvg):
  E[ T∑ k=γ |σ̂2k − σ2k|] ≤ 2(D2 +M2)T 2/3 ln(T 2/3)
  Proof. This follows by exactly the same proof as Theorem 8 and the fact that E[‖g − g′‖2] = 2E[‖g −∇f‖2].
  Proof of Theorem 14. The first part of the proof follows the schema as in the proof of Theorem 12, when ηk ≤ 1/2L, we know that
  E[( T∑ k=1 ηk)‖∇f(xI)‖2] ≤ ∆ + L 2 ∑ k E[η2kσ2k].
  The rest of the proof is exactly the same as the proof of Theorem 3. We can upper bound the right hand side, indeed
  T∑ k=1 E[η2kσ2k] = c2 T∑ k=1 E [ σ2k (σ̂k +m)2 ]
  = c2 ( T∑ k=1 E [ σ2k − σ̂2k (σ̂k +m)2 ] + T∑ k=1 E [ σ̂2k (σ̂k +m)2 ])
  ≤ c2 ( 1
  m2 T∑ k=1 E [ |σ2k − σ̂2k| ] + T
  )
  ≤ c2 ( 2(M2 +D2)T 2/3 ln(T 2/3)
  m2 + T
  ) ≤ 3c2T
  The last inequality follows by the choice of parameters that M 2+D2
  m2 ≤ 1 16
  T 1/3
  ln(T ) . Hence by Markov inequality, with probability at least 3/4,
  ( ∑T k=1 ηk)‖∇f(xI)‖2 ≤ 4E[( ∑T k=1 ηk)‖∇f(xI)‖2] ≤ 4(∆ + 3Lc2T 2 )
  Next, we lower bound the left hand side as in (10),
  1
  c
  ∑ ηk = ∑ 1 σ̂k +m ≥ 1 2 ∑ 1 σk +m − 1 2m3 ∑ (σ̂k − σk)2+ (18)
  Finally, by Markov’s inequality, with probability 3/4,∑ (σ̂k − σk)2+ ≤ 4E[ ∑ (σ̂k − σk)2+] ≤ 4E[ ∑ (σ̂k − σk)2] ≤ 8(D2 +M2)T 2/3 ln(T 2/3). Following the choice of m = 4 √ D2 +M2T− 1 9 ln(T ) 1 2 + 2cL, we have
  1
  2m3 T∑ k=1 (σ̂k − σk)2+ ≤ T 4(M +m) ≤ 1 4 T∑ k=1
  1
  σk +m
  Together with (18) implies that with probability 3/4,
  ∑ ηk ≥ c
  4 T∑ k=1
  1
  σk +m
  Consequently, we know that with probability at least 1− 14 − 1 4 = 1/2,
  ‖∇f(xI)‖2 ≤ 4(∆ + 3Lc 2T 2 )∑
  k c
  4(σk+m)
  ≤ √ 2L∆
  T · 32T∑
  k 1
  (σk+m)
  , (19)
  by setting c = √
  2∆ LT .