Expectigrad: Fast Stochastic Optimization with Robust Convergence Properties

1 INTRODUCTION . Efficiently training deep neural networks has proven crucial for achieving state-of-the-art results in machine learning ( e.g . Krizhevsky et al. , 2012 ; Graves et al. , 2013 ; Karpathy et al. , 2014 ; Mnih et al. , 2015 ; Silver et al. , 2016 ; Vaswani et al. , 2017 ; Radford et al. , 2019 ; Schrittwieser et al. , 2019 ; Vinyals et al. , 2019 ) . At the core of these successes lies the backpropagation algorithm ( Rumelhart et al. , 1986 ) , which provides a general procedure for computing the gradient of a loss measure with respect to the parameters of an arbitrary network . Because exact gradient computation over an entire dataset is expensive , training is often conducted using randomly sampled minibatches of data instead.1 Consequently , training can be modeled as a stochastic optimization problem where the loss is minimized in expectation . A natural algorithmic choice for this type of optimization problem is Stochastic Gradient Descent ( SGD ) ( Robbins & Monro , 1951 ) due to its relatively cheap computational cost and its reliable convergence when the learning rate is appropriately annealed . A major drawback of SGD is that its convergence rate is highly dependent on the condition number of the loss function ( Boyd & Vandenberghe , 2004 ) . Ill-conditioned loss functions are nearly inevitable in deep learning due to the high-dimensional nature of the models ; pathological features such as plateaus , sharp nonlinearities , and saddle points become increasingly probable as the number of model parameters grows—all of which can interfere with learning ( Pascanu et al. , 2013 ; Dauphin et al. , 2014 ; Goodfellow et al. , 2016 ; Goh , 2017 ) . Enhancements to SGD such as momentum ( Polyak , 1964 ) and Nesterov momentum ( Nesterov , 1983 ; Sutskever et al. , 2013 ) can help , but they still largely suffer from the same major shortcoming : namely , any particular choice of hyperparameters typically does not generalize well to a variety of different network topologies , and therefore costly hyperparameter searches must be conducted . This problem has motivated significant research into adaptive methods for deep learning , which dynamically modify learning rates on a per-component basis with the goal of accelerating learning without tuning hyperparameters . AdaGrad ( Duchi et al. , 2011 ) was an early success in this area that ( in its simplest form ) divides each step by a running sum of gradient magnitudes , but this can cause its empirical performance to degrade noticeably in the presence of dense gradients . Later methods such as ADADELTA ( Zeiler , 2012 ) , RMSProp ( Tieleman & Hinton , 2012 ) , and Adam ( Kingma & Ba , 2014 ) remedied this by instead normalizing stepsizes by an exponential moving average ( EMA ) . 1Training on small minibatches can also improve generalization ( Wilson & Martinez , 2003 ) . Such methods are able to increase their learning rates after encountering regions of small gradients and have enjoyed widespread adoption due to their consistent empirical performance . Unfortunately , the EMA has recently been shown to cause divergence on a certain convex optimization problem ( Reddi et al. , 2019 ) that we refer to as the Reddi Problem . This finding has severe implications because it points to an underlying flaw shared by the most widely used adaptive methods . Recent attempts to resolve this EMA-divergence issue have been unsatisfying . Proposed methods invariably begin with Adam , and then apply a minor adjustment aimed at suppressing divergence . Specifically , they either suddenly or gradually transition from Adam to SGD during training ( Keskar & Socher , 2017 ; Luo et al. , 2019 ) , or clip certain quantities in the Adam update rule to limit the EMA ’ s sensitivity ( Zaheer et al. , 2018 ; Reddi et al. , 2019 ) . While these methods technically do prevent divergence on the Reddi Problem , they fail to address the fundamental issue that the EMA can be unreliable , and they do not advance our theoretical understanding of alternatives that are inherently robust to divergence . Furthermore , by intentionally reducing the responsiveness of Adam , these modifications do not always translate into better empirical performance for problems of practical interest—yet they come at the expense of increased complexity . The principal objective of this work is to therefore develop a novel adaptive method that provides stronger convergence guarantees than EMA-based methods while retaining fast learning speed . Towards this , we propose the Expectigrad algorithm , which introduces two major innovations : ( 1 ) normalization by the arithmetic mean instead of the EMA , and ( 2 ) “ outer momentum ” in which bias-corrected momentum is applied jointly to the numerator and denominator . Expectigrad provably converges on all instances of the Reddi Problem that causes Adam to diverge , and minimizes the function significantly faster than related methods using the same hyperparameters . We also derive a regret bound for Expectigrad that establishes a convergence rate comparable to the best known rate for Adam . Our bounds also indicate that Expectigrad is less susceptible to noisy gradients . Finally , we test Expectigrad by training various neural network architectures with millions of learnable parameters ; we show that it consistently outperforms Adam and is competitive with other state-of-the-art methods . 2 PRELIMINARIES . We begin with some notational remarks . We always typeset vectors x , y in boldface to avoid confusion with scalars x , y . We represent the j-th component of vector xi as xi , j . The Euclidean norm is denoted by ‖x‖ and the inner product is denoted by 〈x , y〉 . All arithmetic operations can be assumed to be element-wise : e.g . x± y for addition and subtraction , xy for multiplication , x/y for division , xa for exponentiation by a scalar , √ x for square root , and so on . We now consider the optimization setting that is the focus of this paper . Let l : Rd × Rm 7→ R be a ( nonconvex ) function that we seek to ( locally ) minimize in expectation over some distribution P on Ξ ⊂ Rm . Precisely , we must locate a point x∗ ∈ Rd with the property ∇f ( x∗ ) = 0 , where f ( x ) : = E [ l ( x , ξ ) ] . All expectations in our work are implicitly taken over ξ ∼ P ( Ξ ) . Direct computation of∇f ( x ) is assumed to be infeasible , but repeated calculation of∇l ( x , ξ ) is permitted . We also assume that l has the following properties : Assumption 1. l is bounded below . Assumption 2. l is Lipschitz continuous : |l ( x , ξ ) − l ( y , ξ ) | ≤ L‖x− y‖ , ∀x , y ∈ Rd , ∀ξ ∈ Ξ . Assumption 3. l is Lipschitz smooth : ‖∇l ( x , ξ ) −∇l ( y , ξ ) ‖ ≤ L‖x− y‖ , ∀x , y ∈ Rd , ∀ξ ∈ Ξ . Assumption 1 guarantees the existence of a well-defined set of minima ( Nocedal & Wright , 2006 ) . While Assumptions 2 and 3 need not share the same Lipschitz constant L > 0 in general , we assume that L is sufficiently large to satisfy both criteria simultaneously . These conditions are often met in practice , making our assumptions amenable to the deep learning setting that is the focus of this paper.2 In Appendix B , we prove three lemmas from Assumptions 2 and 3 that are necessary for our later theorems . With these results , we are ready to present Expectigrad in the following section . 3 ALGORITHM . Our motivation for Expectigrad comes from the observation that successful adaptive methods are normalized by an EMA of past gradient magnitudes . This normalization process serves two important 2We note that the commonly used ReLU activation is unfortunately not Lipschitz smooth , but smooth approximations such as the softplus function ( Dugas et al. , 2001 ) can be substituted . This limitation is not specific to our work but affects convergence results for all first-order methods , including SGD . Algorithm 1 Expectigrad Select learning rate α > 0 . Default : α = 10−3 Select momentum constant β ∈ [ 0 , 1 ) . Default : β = 0.9 Select denominator constant > 0 . Default : = 10−8 Initialize parameters x0 ∈ Rd arbitrarily Initialize running sum s0 ← 0 · x0 Initialize running counter n0 ← 0 · x0 Initialize momentum m0 ← 0 · x0 for t = 1 , 2 , . . . , T do Compute gradient estimate : gt ← ∇l ( xt−1 , ξt ) with ξt ∼ P ( Ξ ) Update running sum : st ← st−1 + g2t Update running counter : nt ← nt−1 + sign ( g2t ) Update momentum : mt ← βmt−1 + ( 1− β ) gt + √ st/nt . Define 00 = 0 Update parameters : xt ← xt−1 − α1−βtmt end for return xT benefits . First , stepsizes become scale invariant ; the numerator and denominator are equally affected when the objective function is multiplied by a constant , meaning a particular choice of learning rate is more likely to work well for a wide array of tasks . Second , dividing by the recent gradient magnitude acts as a preconditioner , making the loss function more amenable to first-order optimization . One possible explanation for this is that dividing by the EMA—insofar as it estimates the ( diagonal ) Fisher information matrix for a log-likelihood function—acts as an approximation to natural gradient descent ( Amari , 1997 ; Pascanu & Bengio , 2013 ; Kingma & Ba , 2014 ) . A more recent hypothesis is that the EMA equalizes the stochastic gradient noise ( Balles & Hennig , 2017 ) , and that this isotropic noise is beneficial for escaping saddle points ( Staib et al. , 2019 ) . These properties help explain why EMA-based adaptive methods like RMSProp and Adam are successful on a variety of machine learning problems with relatively little hyperparameter tuning . Nevertheless , the EMA itself has been shown to cause divergence on the Reddi Problem—a synthetic online optimization problem proposed by Reddi et al . ( 2019 ) . The rapidly decaying nature of the EMA reduces the influence of large yet infrequent gradients , which can cause the optimization to progress in an arbitrarily poor direction . Various related works have proposed to clip terms in the Adam update rule to induce convergence on the Reddi Problem ; we discuss them in Section 4.1 . We argue that the approach of constraining EMA-based methods like Adam in an effort to force convergence is counterproductive . Doing so not only risks impacting performance , but it also points to a lack of theoretical understanding . While limiting the adaptation of the EMA may prevent divergence on the Reddi Problem , it does not address the root cause of divergence itself . We are instead interested in algorithms that are automatically robust to rare , high-magnitude gradients that are characteristic of the Reddi Problem . We begin our search by considering the following update rule in which stepsizes are normalized by an arithmetic mean of all gradient samples up to the current timestep t : xt ← xt−1 − α · gt + √ 1 t ∑t k=1 g 2 k ( 1 ) Here , α > 0 is the learning rate and > 0 is a positive constant that prevents division by zero . Crucially , this update retains the scale-invariance property of EMA-based methods discussed at the beginning of this section : multiplying the gradient by a scalar does not affect the stepsizes ( assuming is small ) . Furthermore , the denominator becomes less responsive to new gradients as training proceeds , but never becomes completely static . This is precisely the behavior we want for a method that balances empirical speed with theoretical convergence . To improve the performance of the algorithm , we will want to make some modifications to ( 1 ) . First , note that the explicit summation is unnecessary . If we store the sum st at each time t , then we can efficiently compute it from the previous sum : st ← st−1 + g2t . Second , observe that dividing by t to compute the mean could be problematic when gradients are very sparse—commonly the case when using Rectified Linear Unit ( ReLU ) activations ( Nair & Hinton , 2010 ) . Consider what happens when a component of the gradient is zero for a very long time : as t increments unconditionally , the ratio stt approaches zero . When a nonzero component is suddenly encountered , the algorithm will take a large—and potentially catastrophic—step along that dimension . We can mitigate this by introducing an auxiliary variable nt that counts the number of nonzero gradients that have occurred . Note that we can concisely update nt−1 by simply adding the element-wise sign of the squared gradient : nt ← nt−1 + sign ( g2t ) . Putting these two changes together ( and defining 00 = 0 ) , xt ← xt−1 − α · gt + √ st nt ( 2 ) This new algorithm tracks a per-component arithmetic mean where each component may have a distinct number of counts , thereby ignoring sparsity in the optimization problem . Thus , with some additional memory , we have made the update highly efficient and less prone to numerically instability . Finally , we can also apply momentum to help accelerate learning when the gradient signal is weak . Traditionally , momentum is first applied to the gradient estimator , and then is scaled by the adaptive method ( e.g . Adam ) . In our work , we propose to reverse the order by computing the adaptive steps first and then applying momentum to them . This is important to satisfy the superposition principle : the momentum step “ scheduled ” at the current timestep is not re-scaled by future changes in the denominator . Letting β ∈ [ 0 , 1 ) be the momentum constant , we arrive at our final update : xt ← xt−1 − α 1− βt mt where mt = βmt−1 + ( 1− β ) gt + √ st nt ( 3 ) We call this “ outer ” momentum as it applies to the entire update rule and not just the numerator . The coefficient 1/ ( 1−βt ) is a bias-corrective factor introduced by Kingma & Ba ( 2014 ) . We refer to ( 3 ) as Expectigrad and formally present it in Algorithm 1 . While we have only intuitively justified our design choices so far , we will spend the next few sections rigorously analyzing their theoretical and empirical implications .

Summary: This paper proposes the Expectigrad algorithm that normalizes the exponential moving average (EMA) of first moments on the fly. This avoids normalizing historical gradients by future gradients. The normalization factor is an unweighted average, instead of an EMA, of the historical second moments. For the special case where the EMA constant of the first moment is zero, the paper shows that Expectigrad converges to the optimum on the online convex problems proposed by Reddi et al. (2018) for which the vanilla ADAM fails to converge. For general stochastic smooth convex problems with bounded stochastic gradients, the paper shows that the convergence rate of mini-batch Expectigrad is $O(1 / \sqrt{T} + 1 / b)$ where $b$ is the mini-batch size.