Memory Augmented Optimizers for Deep Learning

1 INTRODUCTION . Gradient-based learning involves minimizing a scalar-valued function F ( θ ) with respect to a parameter vector θ ∈ Rd using an iterative procedure . When training over n examples of input-output pairs , ( xi , yi ) , the optimization problem boils down to solving minθ F ( θ ) , where F ( θ ) = 1 n n∑ i=1 Li ( θ ) , with Li ( θ ) = L ( Mθ ( xi ) , yi ) , for some problem-dependent loss function L and a predictive model M parameterized by θ. Stochastic Gradient Descent ( SGD ) ( Robbins and Monro , 1951 ) is one common method used to tackle this problem , and is often preferred to full-batch Gradient Descent when the quantity of data required to train θ is large , since it can be more efficient to measure a single component of the gradient ( or a mini-batch of component gradients ( Bottou , 1999 ) ) , and move in a noisy direction , than to compute a full gradient at each time step . Several techniques have been proposed to further accelerate the convergence of SGD ( Ghadimi and Lan , 2016 ; Allen-Zhu and Yuan , 2016 ; Zhou et al. , 2018 ; Xu et al. , 2017 ) . These include approaches that maintain a knowledge of previous gradients implicitly by summarizing them in a momentum buffer ( Polyak , 1964b ) , and potentially adapting the learning rate based on the gradient statistics ( Duchi et al. , 2011 ; Hinton et al. , 2012 ; Zeiler , 2012 ; Kingma and Ba , 2015 ) . Optimization techniques such as SGD with Momentum ( Polyak , 1964a ) , AdaGrad ( Duchi et al. , 2011 ) , RMSprop ( Hinton et al. , 2012 ) , AdaDelta ( Zeiler , 2012 ) , and Adam ( Kingma and Ba , 2015 ) maintain a set of buffers that track running moments of the gradients . Such light-weight techniques have shown significant application advantage in practice , whereas in theory , algorithms that store all of the gradients , like SAG ( Roux et al. , 2012 ) and SAGA ( Defazio et al. , 2014 ) , can achieve better convergence . A drawback of full-history methods is that the memory requirement linearly increases with the size of the data . Limited-Memory BFGS ( Nocedal , 1980 ) and its online , stochastic form , oLBFGS ( Schraudolph et al. , 2007 ) aim to exploit this trade-off , but may not converge faster than SGD ( Mokhtari et al. , 2015 ) , and are thus disadvantaged compared to accelerated SGD variants . As a stepping stone towards reaping the advantages of both ends of the spectrum , we propose using memory to augment conventional algorithms . This memory is used to store a small set of critical gradients ( e.g. , gradients with a larger ` 2-norm ) that are occasionally updated with newer gradients . Specifically , rather than storing gradients for all the examples in the dataset , our proposed approach aims to find a smaller set of gradients to store that will aid in the optimization process . Through extensive experiments we show that this marginal increase in the memory requirement ( Figure 1 ) provides accelerated convergence and even improved test performances in majority of the experiments . In this work , we : • Present a framework of memory-augmented optimizers† , § , ‡ compatible with popular gradient-based algorithms . • Prove theoretically that such algorithms converge for smooth strongly convex objectives . • Show that the proposed memory augmented optimizers can lead to faster convergence and better final performance through exhaustive empirical study on eight different architectures among four different tasks ( classification , language modeling , natural language inference , and dialogue ) and six different datasets . • Demonstrate the memory augmented optimizers ’ robustness to gradient selection heuristics , replacement strategies , or the aggregation techniques to summarize information in the memory . 2 MEMORY-AUGMENTED OPTIMIZERS . In this work , we propose augmenting standard optimizers with a fixed size memory of past gradients . The gradients are stored in the limited memory only when they are deemed critical as defined by their ` 2-norm . Gradient Selection to the Memory Letting gt : = ∇θLi ( θt ) denote a component ( or mini-batch ) gradient at time-step t , we use ‖gt‖2 as a scalar indicator of the importance of the gradient , which serves as a proxy for the priority of a gradient to remain in the memory buffer gc . In order to ensure that the buffer eventually flushes out stale gradients , the proxy norms for gradients in the buffer are scaled down by a hyperparameter decay factor , denoted decay ∈ [ 0 , 1 ) . The proposed approach maintains a gradient buffer of fixed capacity C and stores the gradients selected by a chosen heuristic . We refer our heuristic as the critical gradients , which stores the top C gradients by using the ` 2-norm of gt to determine priority . Gradients in the buffer are stored as tuples ( ‖gt‖2 , gt ) , akin to a key-value pair in a dictionary structure , where the key ‖gt‖2 is referred to as the proxy norm , and decayed at each time step t by a decay factor decay ∈ [ 0 , 1 ) . At any iteration t during training , the gradient in a full capacity priority buffer with the smallest ` 2 proxy-norm is replaced by the current mini-batch gradient gt if ‖gt‖2 is greater than the smallest ` 2 proxy-norm in the buffer . The multiplicative decay factor ensures that the buffer is frequently refreshed with more recent gradients . Note that only the proxy-norm used for the heuristic rule is decayed by the decay factor decay ∈ [ 0 , 1 ) , while the gradient itself is not affected . This technique of storing the critical gradients is general enough to be employed in any deep learning model , and can be easily be combined with many existing adaptive gradient optimization algorithms , such as those described in ( Polyak , 1964b ; Zeiler , 2012 ; Duchi et al. , 2011 ; Hinton et al. , 2012 ; Kingma and Ba , 2015 ; Lucas et al. , 2019 ) . Critical Gradient Stochastic Gradient Descent Critical Gradient Stochastic Gradient Descent ( SGDC ) is the explicit integration of critical gradients into SGD . Specifically , the iteration comes in the form θt+1 = θt − α · aggr ( gt , gc ) , ( 1 ) where gc is the set of critical gradients and aggr denotes an aggregation function which is used to combine the critical gradients with the current-iteration gradient . We propose two possible functions for aggr including mean , the average of gt and all critical gradients , and sum , the addition of gt to the average of all critical gradients . Mathematically , for a buffer of size C these are defined as mean ( gt , gc ) = 1 C + 1 ( gt + ∑ gc∈gc gc ) ( 2 ) sum ( gt , gc ) = gt + 1 C ∑ gc∈gc gc ( 3 ) In general , we find optimization to be robust to the specific choice of aggregation function , but we do observe that adaptive learning rate methods ( Adam , RMSprop ) perform best with mean , while convential SGD ( with or without Polyak momentum ) performs best with sum . Critical Gradients with Accelerated Methods Agarwal et al . ( 2019 ) show that curvature in adaptive gradients can be estimated without loss of generality with a finite moving window of past gradients . Following which we propose critical gradient extensions of SGD with Momentum ( SGDM ) , RMSprop , and Adam . The retention of the critical gradients can be naturally extended to more firstorder optimization methods with the replacement of gt in those algorithms with aggr ( gt , gc ) . The update steps for the corresponding C versions are as follows : mt = γ ·mt−1 + aggr ( gt , gc ) θt+1 = θt − α ·mt ( 4 ) E [ g2 ] t = 0.9 · E [ g2 ] t−1 + 0.1 · aggr ( gt , gc ) 2 θt+1 = θt − η√ E [ g2 ] t + · aggr ( gt , gc ) ( 5 ) mt = β1mt−1 + ( 1− β1 ) · aggr ( gt , gc ) vt = β2vt−1 + ( 1− β2 ) · ( aggr ( gt , gc ) ) 2 m̂t = mt 1− βt1 v̂t = vt 1− βt2 θt+1 = θt − η√ v̂t + m̂t ( 6 ) The general algorithm to integrate critical gradients into any optimizer is provided in Algorithm 1 in Appendix A . Replacing gt with aggr allows the optimization function to retain key gradients and optimize the parameters . The acceleration provided by the topC gradients complements the adaptive learning rate methods by not only computing the learning rate with respect to gt but also with respect to the critical gradients from the past . We observe that all C-extended optimizers in our empirical analysis and experiments in § 5 lead to faster convergence than their vanilla , non memory-augmented counterparts . 3 PROOF OF CONVERGENCE . Consider the SGDC update in Equation ( 1 ) , where aggr ( gt , gc ) outputs a linear combination of the critical gradients in memory gk ∈ gc , and the stochastic gradient gt , which is computed at the current parameters θt . When the staleness of gradients in the memory is bounded , we can prove convergence of this method through the lens of linear multi-step methods . Formally , suppose there exists an integer K > 0 such that gk ∈ gc at iteration t , implies t− k < K , and suppose the objective F : Rd → R is twice-continuously differentiable , L-smooth , and µ-strongly convex , with 0 < µ ≤ L. Assume that for all k , the stochastic gradient gk is a random vector satisfying E [ gk ] = ∇f ( θk ) and E [ ‖gk −∇f ( θk ) ‖2 ] ≤ σ2 , for some finite constant σ2 . Letting ζk : = gk −∇f ( θk ) denote the gradient noise at iteration k , we take these gradient noise terms to be mutually independent . For examples of functions satisfying these properties , see , e.g. , ( Nesterov , 2004 ; Bubeck , 2015 ) . Examples of typical tasks satisfying these assumptions are ` 2-regularized logistic regression and ` 2-regularized least-squares regression ( i.e. , ridge regression ) . Taken together , these properties imply that the Hessian ∇2F ( θ ) exists , and for all θ ∈ Rd , the eigenvalues of ∇2F ( θ ) lie in the interval [ µ , L ] . In this case , let θ ? denote the unique minimizer of F . To prove convergence , we view SGDC as a linear multi-step method , i.e. , the parameters at the current time step are updated with a linear combination of gradients from the past K time steps . From this view , we can describe the evolution of the algorithm according to a discrete-time linear dynamical system , the convergence of which is characterized by the spectral properties of the system matrix . Our convergence theorem relies on the largest singular value of this system matrix , but can be strengthened to instead rely on the spectral radius of the system matrix by constructing an appropriate Lyapnuov function containing a time-invariant system matrix satisfying a particular Integral Quadratic Constraint ( IQC ) ( Lessard et al. , 2016 ) . However , one downside of the IQC framework is that one must typically solve a semidefinite program ( e.g. , using a numerical solver ) to obtain explicit convergence rates ( Hu and Lessard , 2017 ) , whereas Theorem 1 provides an analytical convergence rate . Theorem 1 ( Linear Convergence of SGDC ) . Let q ( α , K ) : = sup { wk∈ [ 0,1 ] } , { ηk∈ [ µ , L ] } ρ ( Λ ) , where ρ ( Λ ) are the singular values of the numerical matrix Λ : = [ λ0 λ1 · · · λK−1 0 1 · · · 0 · · · · · · · · · · · · · · · · · · 0 0 · · · 1 ] ∈ RK×K , with λ0 = 1− αη0 and λk = −αwkηk for k > 0 . If the step-size α > 0 is chosen sufficiently small , such that q ( α , K ) < 1 , then E ‖θt+1 − θ ? ‖2 ≤ q ( α , K ) 2t ‖θ1 − θ ? ‖ 2 + α2K 1− q ( α , K ) 2 σ2 , where the expectation on the left hand side is with respect to the joint distribution over the random noise variables ζk , . . . , ζ1 . Theorem 1 shows that SGDC can be made to converge exponentially fast to a neighbourhood of the solution , the size of which is proportional to the variance of the stochastic gradients σ2 and the step-size α2 . If K = 1 ( i.e. , the memory gc , at all times t , only contains the mostly recently computed gradient ) , then SGDC reduces to regular stochastic gradient descent , and the rate q ( α , K ) in Theorem 1 reduces to the well known convergence rate of SGD for smooth strongly convex functions ( maxη∈ [ µ , L ] |1− αηL| ) , and the variance bound reduces to the standard variance bound for SGD with a constant step-size ( α2σ2/ ( 1− q2 ) ) , see , e.g. , ( Bottou et al. , 2018 ; Assran and Rabbat , 2020 ) . For K > 1 , it may be possible to choose the step-size α and the aggregation weights { wk } to obtain accelerated convergence rates , faster than SGD . For example , accelerated gradient methods such as Polyak and Nesterov momentum can be viewed as multi-step methods with K = 2 . Note that although K appears in the numerator of the coefficient of the variance term , it is also present in the rate term in the denominator , where faster convergence rates , smaller q ( α , K ) , directly lead to smaller variance bounds . Note that expressing the convergence rate of a multi-step method as the roots of a polynomial , as we do in Theorem 1 , is not new ; see , e.g. , Polyak ( 1964a ) . It is also straightforward to extend Theorem 1 to replace the bounded variance assumption ( common in the stochastic approximation setting ) with other conditions , such as the Strong Growth Condition or the Weak Growth Condition ( Bottou et al. , 2018 ) , which are more tailored to stochastic finite-sum models .

Summary: The paper proposes an optimisation framework that uses K previous gradients stored in a buffer to compute an aggregated proxy-gradient for the  use in other gradient frameworks (studied are SGD, SGDM, RMSProp, Adam) and evaluates it on CIFAR 10/100 as well as NLP tasks. The paper also presents a convergence proof on strongly convex smooth objectives. In almost all empirical evaluations the method shows an improvement in terms of iteration and performance and ablations are performed to show strong robustness to hyperparameters introduced.