Extreme Tensoring for Low-Memory Preconditioning 

State-of-the-art models are now trained with billions of parameters , reaching hardware limits in terms of memory consumption . This has created a recent demand for memory-efficient optimizers . To this end , we investigate the limits and performance tradeoffs of memory-efficient adaptively preconditioned gradient methods . We propose extreme tensoring for high-dimensional stochastic optimization , showing that an optimizer needs very little memory to benefit from adaptive preconditioning . Our technique applies to arbitrary models ( not necessarily with tensor-shaped parameters ) , and is accompanied by regret and convergence guarantees , which shed light on the tradeoffs between preconditioner quality and expressivity . On a large-scale NLP model , we reduce the optimizer memory overhead by three orders of magnitude , without degrading performance . 1 INTRODUCTION . Among the most influential and important optimization techniques in machine learning are adaptive learning-rate methods , otherwise known as diagonal-matrix adaptive preconditioning . Essentially all of the most-commonly used incarnations of adaptive preconditioning ( AdaGrad , Adam , RMSprop , Adadelta , etc . ) accumulate second-moment estimators of each coordinate of the gradient , then scale the parameter updates by the square roots of these accumulators . These methods come with an overhead memory cost of storing these accumulators , thereby doubling the memory consumption . In the regime where model size encroaches upon the same order of magnitude as the total amount of RAM , a need has arisen to view memory as a limited resource in large-scale optimization . We address the question of whether the benefits of adaptive preconditioning can be attained without significant memory overhead . To this end , we introduce extreme tensoring , a family of generic modifications to any second-moment-based adaptive optimizer . Our method uses a compressed preconditioner which takes the form of a tensor product of arbitrary order , with simple updates , without necessarily requiring the parameters to be tensor-shaped . In our regret analysis , we quantify how extreme tensoring competes provably with full-memory AdaGrad in the online convex optimization framework , with a multiplicative data-dependent constant that can be measured empirically . In a large-scale language modeling setting , we demonstrate that an optimizer requires very little additional memory to benefit from adaptive preconditioning . Furthermore , the inherent flexibility of our method enables us to conduct , to the first of our knowledge , the first empirical study of the tradeoff between training convergence and memory in the optimizer . 1.1 RELATED WORK . The most widely-adopted form of adaptive preconditioning are second-moment-based : for example , AdaGrad ( Duchi et al. , 2011 ) , Adam ( Kingma and Ba , 2014 ) , RMSprop ( Tieleman and Hinton , 2012 ) , and Adadelta ( Zeiler , 2012 ) . Some recent preconditioning methods are not based on second moments ( Bello et al. , 2017 ; Chen and Gu , 2018 ; Bernstein et al. , 2018 ) , and fall beyond our scope . Tensor-factorized preconditioners . Many works ( Martens and Grosse , 2015 ; Gupta et al. , 2018 ; Martens et al. , 2018 ) investigate tensor-factorized preconditioners . These are presented in the view of restricted full-matrix preconditioning ( vs. diagonal for us ) of tensor parameters ( vs. generalpurpose for us ) . In parameter count regimes relevant to this paper , these full-matrix adaptive methods suffer from prohibitive time and memory overhead issues . Extreme tensoring can be viewed as the diagonal version of Shampoo . The matrix version of diagonal Shampoo was discussed in the appendix of ( Gupta et al. , 2018 ) ; in this paper we extend the algorithm to arbitrary tensor factorizations , and focus on the tradeoff between memory consumption and optimizer performance . The theoretical part of our work follows proof techniques seen in ( Gupta et al. , 2017 ) and ( Gupta et al. , 2018 ) , although our diagonal restriction results in distinct updates and incomparable regret bounds . Optimizer memory reduction in deep learning . Perhaps the most closely related work to ours is Adafactor ( Shazeer and Stern , 2018 ) , an empirical work which achieves sublinear-memory adaptive regularization by restricting preconditioners on matrix-shaped gradients to “ row ” and “ column ” learning rates , with a similar update rule . Similar algorithms have appeared in prior work ( Gupta et al. , 2014 ; Shazeer et al. , 2017 ) . As a special case of our regret analysis , we provide some theory for this line of work , while proposing a more general and versatile memory-reduction method . Concurrent to this work , Anil et al . ( 2019 ) are also motivated by memory efficiency , and the algorithm uses a different approximation to the preconditioner than a tensor product . Their algorithms ’ index partitions are compatible with our tensor indices , and we provide a comparison in Appendix C. Optimizer memory reduction outside of deep learning . Older than our work is the classic LBFGS algorithm ( Liu and Nocedal , 1989 ) , a quasi-Newton method with a tunable memory window hyperparameter . This does not immediately apply to our setting , as the tunable memory overhead is ( window size ) × , rather than ( 1 + ε ) × . Furthermore , second-order methods turn out to be brittle under stochastic gradients ; it remains an open empirical problem to robustly exploit Hessian curvature of the loss surface in modern deep learning settings . Intrinsic dimensionality in deep learning . More broadly , there are various empirical attempts to locate and characterize low-dimensional phenomena hiding within high-dimensional parameter spaces in deep learning ( Li et al. , 2018a ; b ) . Our comparison provides a lens on measuring the intrinsic dimensionality of the preconditioner . Another is ( Agarwal et al. , 2020 ) , a different way to impose parameter bottlenecks on adaptive preconditioners ; however , the proposed method does not reduce memory consumption . Our contributions : • Using the user-selectable degree of memory reduction , we conduct an empirical study of training convergence vs. preconditioner memory usage . We find , in a realistic large-scale language modeling setting , that reallocating memory savings from the optimizer to the model itself can pay off in terms of end-to-end performance . • We propose extreme tensoring as a modification to AdaGrad , Adam , etc . for reducing the overhead memory cost in adaptive preconditioning , accompanying algorithms such as ( Shazeer and Stern , 2018 ; Anil et al. , 2019 ) . • We derive a regret bound for extreme-tensored AdaGrad , and provide the first numerical comparisons of regret bounds in deep learning . Though they obviously do not provide convergence guarantees in the non-convex setting , they act as a proxy for understanding the AdaGrad accumulator overestimation induced by memory constraints . Via this diagnostic lens , we see that the regret competes with AdaGrad within a multiplicative constant , which we measure to be small in practice . 2 PRELIMINARIES . 2.1 STOCHASTIC OPTIMIZATION AND ADAPTIVE METHODS . We will state our algorithmic contribution in the usual framework of stochastic optimization of a differentiable function f ( · ) , equipped with an unbiased stochastic gradient oracle ∇̃f ( · ) . At each of T iterations , the algorithm queries a stochastic gradient gt : = ∇̃f ( xt ) . In practice , our target is the large-scale ERM setting , where f is the population loss on training data , and ∇̃f is a mini-batch stochastic gradient with respect to the model parameters . However , since the introduction of AdaGrad ( Duchi et al. , 2011 ) , the standard analysis of adaptive regularization has been in the online convex optimization ( OCO ) framework ( see , e.g . Gupta et al . ( 2017 ; 2018 ) ) . In this language , we view the xt queried by the learning algorithm as a sequence of online decisions chosen from a convex set K , and the gt as a sequence of gradients of adversarially chosen convex loss functions ft. We are interested in minimizing the regret , defined as T∑ t=1 ft ( xt ) −min x∈K T∑ t=1 ft ( x ) . This is a generalization of stochastic convex optimization ; for a broad survey , see Hazan ( 2016 ) . A framework for reducing the stochastic non-convex setting to the online convex setting is introduced in Agarwal et al . ( 2018 ) , translating a sublinear regret bound to convergence to approximate stationary points . 2.2 TENSORS AND INDEXING CONVENTIONS . We introduce some notation and terminology for clarity and convenience . By [ n ] we denote the index set { 1 , . . . , n } . Id and 0d refer to the d × d identity matrix and ddimensional zero vector , respectively . We use • to denote the trace inner product ( Hadamard product ) between two matrices , and · to denote the element-wise product of two vectors , or a vector and a scalar . As is standard in optimization contexts , a tensor of order p is an element of Rd1× ... ×dp : a pdimensional table of real numbers , indexed by a p-tuple I = ( I1 , · · · Ip ) of integers , where Ii ∈ [ di ] . Our methods will require the user to specify a relabeling of gradient vectors , which we will call a tensor index : Definition 2.1 . Let d1 , . . . , dp be positive integers whose product is d. A tensor index is a bijection I : [ d ] →×pi=1 [ di ] between indices for Rd and Rd1× ... ×dp . We will refer to the “ reshape ” conversion between vectors x ∈ Rd and tensors x ∈ Rd1× ... ×dp specified by the index bijection I . For this , we will use the shorthand notation x : = I ( x ) and x : = I−1 ( x ) . Throughout this paper , we will use square brackets to refer to vector and tensor indices : for example , the previous definition can be restated as enforcing x [ i ] = x [ I ( i ) ] . See Figure 1 for an illustration of such a conversion . Although we do not need them to state the simple algorithm , it will be useful in the analysis to introduce a few other pieces of notation . We denote the tensor ( or Kronecker ) product of matrices A , B by A⊗B . Given a positive definite matrix A ∈ Rd×d , for any x ∈ Rd , define the matrix norm ‖x‖A : = √ x > Ax and its dual ‖x‖∗A : = √ x > A−1x . For a square matrix M , diag ( M ) refers to the diagonal matrix of the same dimensions and diagonal entries as M . 3 PRECONDITIONING BY EXTREME TENSORING . We begin by presenting our proposal for black-box memory reduction in adaptive preconditioning . We call this extreme tensoring , as it takes the view of restricting a high-dimensional preconditioner to an arbitrary-order tensor product . The algorithm implicitly maintains a preconditioner which is a rank-one tensor product of the same dimension as a given tensor index . A second-moment accumulator is maintained for sums of squared gradient magnitudes across each ( p − 1 ) -dimensional tensor slice ; the adaptive step size for each coordinate is the inverse square root of the geometric mean of its p corresponding slice sums . The formal specification of the base algorithm ( AdaGrad with extreme tensoring ) is given in Algorithm 1 . Algorithm 1 AdaGrad with extreme tensoring 1 : Input : Initializer x1 , learning rate schedule { ηt } , tensor index I with dimensions ( d1 , . . . , dp ) , ε > 0 2 : Initialize ( S ( 1 ) , . . . , S ( p ) ) : = ( 0d1 , . . . ,0dp ) 3 : for t = 1 , . . . , T do 4 : Receive stochastic gradient gt 5 : Reshape : gt : = I ( gt ) 6 : Accumulate slice sums : ∀i ∈ [ p ] , S ( i ) [ j ] ← S ( i ) [ j ] + ∑ I : Ii=j gt [ I ] 2 7 : Get step sizes : δt [ I ] : = ( ε+ ∏p i=1 S ( i ) [ Ii ] ) − 12p 8 : Update : xt+1 ← xt − ηt · I−1 ( δt ) · gt 9 : end for We make a few important remarks : • AdaGrad is a special case of Algorithm 1 , with p = 1 , d1 = d. The analogues for Adam , RMSprop , Adadelta , etc . are obtained straightforwardly by decaying the accumulator ( S ( i ) [ j ] ← β2 · S ( i ) [ j ] + . . . ) . Extreme tensoring is compatible with first-moment estimation ( i.e . momentum ) , although the memory savings disappear . In the setting for our main empirical result , removing momentum did not degrade performance , corroborating the findings of Shazeer and Stern ( 2018 ) . • If the tensor dimensions ( d1 , . . . , dp ) in the decomposition are close to equal , the memory overhead scales as O ( p · d1/p ) . • As is standard in matrix and tensor optimization methods ( Gupta et al. , 2018 ; Shazeer and Stern , 2018 ; Martens et al. , 2018 ) , independent copies of Algorithm 1 can be run on each tensor-shaped parameter group ; optimizer APIs in standard deep learning packages promote this convention . This can be viewed as maintaining a tensor sum of preconditioners , each of which is a p-wise tensor product . • The accumulator update step can be implemented concisely and efficiently using reindexing procedures ( typically called reshape or view ) in standard deep learning and numerical linear algebra packages . Lemma 4.3 in the analysis shows that the per-coordinate learning rates are underestimates of those prescribed by AdaGrad . This interpretation is used in our proof , and serves as the basis of the empirical measurements in Section 5.3 .

Summary: This paper considers the problem of the need for memory-efficient optimizers given the increase in model complexity. They study memory-efficient adaptively preconditioned gradient methods, and see the trade-offs among expressivity and preconditioner quality. They show results on a large-scale NLP model, and show that the memory overhead can be reduced by 3 orders of magnitude, without sacrificing performance.