AdaMeM: Memory Efficient Momentum for Adafactor

Authors: nikhil  vyas
depen  morwani
sham  kakade
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-03 07:48:17

Abstract: Adafactor is a memory efficient algorithm which does not maintain momentum and has near 0 memory overhead as compared to gradient descent. However it performs worse than Adam in many setups. Prior works have shown that this gap can be removed by adding momentum to Adafactor. This comes at the cost of increased memory requirements. In this work we use the ideas of low rank optimizers such as LoRA and GaLore to maintain momentum on a low rank subspace of the weights on top of Adafactor to give a new optimizer: AdaMeM. However unlike low rank optimizers we still utilize full rank gradients but maintain momentum only on the top SVD subspace of the gradients. We show results on language modelling for models of size 210M and 550M demonstrating improved performance over Adafactor and GaLore. We also give theoretical arguments supporting the design of AdaMeM.

Main paper: 
 introduction: 
 Adam (Kingma & Ba, 2015) , the default optimizer utilized in language modeling tasks within deep learning, is considerably more memory-intensive compared to stochastic gradient descent (SGD) (Robbins & Monro, 1951) , as it maintains two additional values: the first and second-order momentum for each network parameter. The development of memory-efficient optimizers that retain Adam's performance remains a critical area of research, as evidenced by numerous previous studies.Adafactor (Shazeer & Stern, 2018) , in its original proposal, entirely eliminated the memory overhead by removing the first-order momentum and implementing a highly efficient factorization of the second-order momentum. Nonetheless, Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024).subsequent works have demonstrated that Adafactor's performance lags behind that of Adam (Rae et al., 2021; Zhai et al., 2022) . Reintroducing the first-order momentum has been shown to bridge this performance gap (Zhai et al., 2022) . Similarly, a recently proposed optimization algorithm, Lion (Chen et al., 2023) , also removes the memory overhead associated with the second-order momentum, substituting it with a sign operation.Other research efforts have focused on reducing the number of parameters involved in the overall update, thereby diminishing the memory overhead for both first and secondorder momentum. LoRA (Hu et al., 2022) introduced low-rank updates for fine-tuning network parameters, significantly reducing the optimizer's memory requirements. ReLoRA (Lialin et al., 2023) extended this approach to pretraining by periodically merging the low-rank updates into the parameters. GaLore (Zhao et al., 2024) also maintained low-rank updates for parameters, but specifically within the top singular value decomposition (SVD) subspace of the gradients.In this work, we integrate the memory-efficient preconditioning of Adafactor with the low-rank momentum concepts of GaLore, resulting in AdaMem-a capable optimizer with substantially lower overhead compared to Adam. We empirically compare our method with Adam, Adafactor, and GaLore for language pretraining tasks, demonstrating that our method, which incorporates low-rank momentum, achieves performance significantly closer to Adam than either Adafactor or GaLore.We support our empirical findings with theoretical results. Specifically, we establish a novel connection between lowrank optimization and the Shampoo optimizer (Gupta et al., 2018) , illustrating that the low-rank subspace used for maintaining momentum approximates the top eigenspace of the Hessian, based on the assumptions of the Shampoo optimizer. For linear regression, we further demonstrate that maintaining momentum in the top eigenspace of the Hessian substantially enhances the convergence speed of gradient descent. Figure 1 . Comparison of our algorithm (AdaMeM) to Adafactor and Adam as baselines and the prior work of GaLore. In these plots we plot final validation loss as a function of rank used by AdaMeM and GaLore. We observe that while AdaMeM is always better than Adafactor, GaLore is only better than Adafactor at large ranks. In the right and middle figure the low rank approximation is applied to all layers except first and last layers while in the last plot it is applied to all layers. See Section 5 for additional details. 
related works: 
 Low-Rank Optimizers: The training of language models using low-rank methods has become a prominent area of research, experiencing a significant surge in recent times. LoRA (Hu et al., 2022) introduced low-rank finetuning of large language models (LLMs), demonstrating competitive performance across various downstream tasks.ReLoRA (Lialin et al., 2023) extended this approach to pretraining by periodically merging low-rank updates into the weights. Similarly, GaLore (Zhao et al., 2024) proposed using low-rank updates, specifically along the top SVD subspace of the gradients.Low-Memory Optimizers: Adafactor (Shazeer & Stern, 2018 ) is a variant of Adam that maintains factored second moments and does not use momentum. As a result, this optimizer incurs almost no 1 overhead. Given that Adafactor's space requirements are negligible compared to the weights, we will use it as our primary baseline for lowspace optimizers. Additionally, the recently proposed optimizer Lion (Chen et al., 2023) reduces the overall memory footprint compared to Adam by replacing the second-order 1 The factored second momentum statistics occupy m + n space for an m × n layer, which is negligible compared to the mn space required to store the weights themselves. Common implementations in libraries such as PyTorch typically use twice the space of the weights, as they also store gradients for all layers. This can be mitigated by optimizers like LOMO (see Appendix A for further discussion). Throughout this paper we will assume such an implmentation. momentum in the Adam update with a sign operation. There is also a large literature (????) on optimizers using lower precision to save memory. These approaches based on using lower precision to save memory are orthogonal to low rank based approaches and can be combined with them (Zhao et al., 2024) , this also applies to our approach.Adamw(factor) as a baseline: In some of our experiments, Adafactor with momentum outperforms Adam. Therefore, we will report the best results from these two methods as Adamw(factor), which will serve as our baseline for algorithms utilizing full-rank momentum.Benefit of using Momentum: Given that Adafactor does not use momentum, it is expected to perform poorly in scenarios where momentum benefits optimization speed. Previous studies (Kidambi et al., 2018; Lee et al., 2022; Wang et al., 2024) have argued and empirically demonstrated that momentum is advantageous with large batch sizes, with its benefits diminishing at smaller batch sizes. We empirically explore this in Section 6, confirming that these observations hold for pretraining language models. This leads to the following question: With large batch sizes, how can we gain the benefits of momentum without incurring its space cost? 
algorithm: 
 We begin by describing the intuition behind AdaMeM (Algorithm 1). Low-rank algorithms such as LoRA (Hu et al., 2022) , ReLoRA (Lialin et al., 2023) , and GaLore (Zhao et al., 2024) perform gradient descent with momentum in a low-dimensional space. Our approach builds upon the methodology of GaLore. For a gradient matrix G ∈ R m×n with m ≤ n, GaLore maintains momentum within the top r rank subspace corresponding to the left singular vectors of G. Specifically, let P ∈ R m×r denote the projection matrix to the top r left singular vectors of G. Momentum is then maintained for P T G, with P itself being updated every T = 200 steps.Our improvement is based on a straightforward concept: in addition to running gradient descent with momentum in the low-rank space, we can also perform gradient descent (without momentum) on the residual space using the gradient of the current batch. This requires using full-rank gradients while maintaining momentum solely in the low-rank space. Although this outlines the core idea behind our algorithm, further complexities arise because we aim to implement Adam-like algorithms instead of SGD.We address these complexities as follows: First, for the low-rank component of the gradient, we run Adafactor with momentum in the low-rank space, similar to the Adafactor variant of GaLore. A natural choice for handling the residual gradient would be to apply Adafactor (without momentum) in the complementary space. However, we lack a basis for the complementary space. Consequently, we could opt to run Adafactor (without momentum) for the residual gradient in the original space. This approach has a drawback: while the low-rank gradient and the residual gradient are orthogonal by definition, preconditioning the low-rank gradient in the low-rank space and the residual gradient in the original space with Adafactor can disrupt this orthogonality. To maintain orthogonality, we employ a one-sided version of Adafactor, as described in Algorithm 4. We prove below that this preserves orthogonality. In Section 6, we demonstrate that using the standard Adafactor instead results in suboptimal performance.Claim 3.1. Let G ∈ R m×n be the gradient matrix with m ≤ n and P ∈ R m×r be the projected matrix. Then preconditionng the projected gradient (P T G) for Adafactor (Algorithm 3) and residual gradient ((I -P P T )G) by one sided Adafactor (Algorithm 4) preserves their orthogonality.Proof. We begin by noting that precondition P T G by Adafactor (Algorithm 3) results in a matrix of the form A = P D 1 P T GD 2 where D 1 ∈ R r×r , D 2 ∈ R n×n are diagonal matrices. We note that A ∈ R m×n is in the original space. Preconditioning residual gradient ((I -P P T )G) by one sided Adafactor (Algorithm 4) results in a matrix of the form B = (I -P P T )GD 3 where D 3 ∈ R n×n is a diagonal matrices.We are interested in the inner product of A and B when they are treated as vectors i.e. Tr(A T B).Tr(A T B) = Tr(D T 2 G T P D T 1 P T (I -P P T )GD 3 ) = 0 since P T (I -P P T ) = 0.The pseduocode for AdaMeM is given in Algorithm 1.Memory Requirements. We note here that our memory requirement (beyond the weights themselves) for maintaining rank r momentum is mr(for projection matrix P ) + nr(for momentum in the low rank space) + r + n + n(for Adafactor's second moments) ≈ (m + n)r same as the Adafactor version of GaLore.Algorithm 1 AdaMeM. Bias correction and regularization constant ϵ not described for simplicity. Require: A layer weight matrix W ∈ R m×n with m ≤ n.Step size η, relative step size δ = 1, decay rates β 1 = .9, β 2 = .95, rank r, subspace change frequency T = 200. 1: Initialize first-order moment M 0 ∈ R n×r ← 0 2: Initialize auxiliary variables Z 0 , Z os 0 to be used by Adafactor preconditioner (Algorithm 3) and one-sided Adafactor preconditioner (Algorithm 4). 3: Initialize step t ← 1, auxillary 4: repeat 5:G t ∈ R m×n ← -∇ W ϕ t (W t ) 6:R t ← Project(G t , r, T, t) {Project gradient into compact space using Algorithm 2} 7:S t ← G t -P t R t {Residual gradient outside of the compact space} 8: M t ← β 1 • M t-1 + (1 -β 1 ) • R t 9: N t,1 , Z t = Adafactor-Preconditioner(M t , Z t-1 , β 2 ) 10: N t,2 , Z os t = OS-Adafactor-Preconditioner(S t , Z os t-1 , β 2 ) 11: W t ← W t-1 + η • (N 1,t + 
theory: 
 In this section, we provide theoretical arguments to support the implementation of AdaMem.Algorithm 3 Adafactor Preconditioner (Shazeer & Stern, 2018) Operations such as squaring matrices and division of matrices are done elementwise.Require: Momentum matrix M t ∈ R r×n , previous step moment estimates Z t-1 , decay rate β 2 = .95 1: R t-1 , C t-1 ← Z t-1 2: R t = β 2 R t-1 + (1 -β 2 )(M 2 t )1 n {Update R t , C t as done in Adafactor} 3: C t = β 2 C t-1 + (1 -β 2 )1 ⊤ r (M 2 t ) 4: Vt = R t C T t /(1 ⊤ r R t )= β 2 R t-1 + (1 -β 2 )(S 2 t )1 n 3: Z t ← R t 4: Vt = R t 1 T m /m output St √ Vt , Z tConsider a quadratic loss landscape given by L(w) = w T Hw, where H is a postive-definite matrix. Let the eigenspectrum of H be given by λ i for i ∈ [d] where λ 1 λ 2 . . . ≥ λ d . Let κ denote the condition number of H given by λ 1 /λ d . Then it is well known (Nesterov, 1983 ) that gradient descent with momentum can achieve a convergence rate of ( √ κ-1)/( √ κ+1), which is a quadratic speed up of convergence rate of gradient descent, which is (κ -1)/(κ + 1). For maintaining momentum in the top-k eigenspace with power-law decay in the eigenspectrum, that is representative of deep learning tasks (Murray et al., 2023) , the following result holds: Theorem 4.1. Consider a quadratic loss landscape given by L(w) = w T Hw, where the eigenspectrum of H follows a power-law decay, i.e, λ i ∝ i α for some α > 0. For k = d 2/3 , there exists learning rates η 1 and η 2 , and momentum β, such that using learning rate η 1 and momentum β in the top-k eigenspace and learning rate η 2 with no momentum in the complement, leads to a convergence rate of (κ 1/3 -1)/(κ 1/3 + 1).Proof. For k = d 2/3 , consider the top-k eigenspace denoted by S k . For power-law decay, the condition number of S k is given by 1 d 2α/3 = κ 2/3 . Thus, by the convergence rate of gradient descent with momentum, we can set learning rates η 1 and momentum β in this space to get a convergence rate of (κ 1/3 -1)/(κ 1/3 + 1).Let the complement of S k be denoted by C k . The condition number of C k is given by 1 d α/3 = κ 1/3 . Thus, by the convergence rate of gradient descent, we can set learning rate η 2 in this space to get a convergence rate of (κ 1/3 -1)/(κ 1/3 + 1).Note that as d → ∞, with k = d 2/3 , k d → 0. Thus, in higher dimensions, we need to maintain momentum in a very small subspace to get the benefits of momentum. Although notice that this small subspace should correspond to the top eigenspace of the Hessian.Given Theorem 4.1 we know that momentum is only needed in the top eigenspace of the Hessian. This means we need to identify the top eigenspace. Various optimizers such as K-FAC and Shampoo approximate the Hessian by Kronecker products. We will be using the following claim from Gupta et al. ( 2018):Claim 4.2. Under Shampoo's Hessian approximation, eigenvectors of Hessian are of the form Vector(uv T ) where u and v are the eigenvectors of E[GG T ] and E[G T G] respectively.] and E[G T G] by a single sample of G, we see that the top singular vectors of G indeed are an estimate of the top eigenvectors of the Hessian. This shows that the low-rank subspace used in Algorithm 1 is a natural estimate of the top eigenspace of the Hessian. We hope that this connection between Shampoo and low rank training can be generally useful in further understanding the dynamics of low rank training. 
experiments: 
 Models and Dataset. Starting from the OLMo codebase (Groeneveld et al., 2024) , we train decoder-only transformer models with sequence length of 512 in two sizes: 210m and 550m. The models have widths of 1024, 1280, and depths of 12, 24. The batch sizes are 1024, 2048. All models, except those using GaLore, are trained for 8000 steps with 2500 steps of warmup followed by cosine decay. For GaLore, we train for longer to take into account the fact that one gradient step of GaLore can be implemented with fewer FLOPs (see Appendix B). All of our models approximately match chinchilla (Hoffmann et al., 2022) style 20x scaling of tokens with respect to model size. The MLP hidden dimension is 4x the width. The activation function is GeLU (Hendrycks & Gimpel, 2016) . We use RoPE positional encodings (Su et al., 2024) . Attention heads are always of dimension 64. We use PyTorch's default LayerNorm. Following Wortsman et al. (2024) , we do not learn biases for the linear layers or LayerNorms. We train in mixed precision with bfloat16. We use the T5 tokenizer (Raffel et al., 2020) and train on the C4 dataset.Optimizer Hyperparameter Sweeps. Following Touvron et al. (2023) ; Biderman et al. (2023) , we fix β 1 = .9, β 2 = .95, and ϵ = 1e -8 (though we study ablating β 2 in Section 6.1). We use weight decay of .01. For Adafactor and Adam, we sweep over the learning rates [1.0e -4, 3.16e -3, 1.0e -3, 3.16e -3, 1.0e -2]. For GaLore, we sweep over learning rates [3.16e -4, 1.0e -3, 3.16e -3, 1.0e -2, 3.16e -2] and [.125, .25, .5 ] for the α parameter. For AdaMeM, we sweep over learning rates [3.16e -4, 1.0e -3, 3.16e -3, 1.0e -2] and [.5, 1] for the δ parameter.In Figure 1 , we compare our algorithm (AdaMeM) to Adafactor and Adam as baselines and the prior work of GaLore. In these plots we plot final validation loss as a function of rank used by AdaMeM and GaLore. We observe that while AdaMeM is always better than Adafactor, GaLore is only better than Adafactor at large ranks. In the GaLore paper, the low rank approximation is not applied to the first and last layer. We do the same in the right and middle figures by training the first and last layers with Adafactor for AdaMeM and Adam for GaLore. Here we see that while GaLore and AdaMeM perform nearly equally at the highest rank we considered (half of full rank), at lower ranks, AdaMeM performs significantly better. In the last plot, low rank approximation is applied to all layers, and we see that AdaMeM outperforms GaLore at all ranks considered. 
ablations: 
 Other Design Choices. While designing AdaMeM, we tried some other variants which were less performant than AdaMeM. Variant 1 corresponds to directly adding the low rank momentum and the residual gradient and preconditioning in the original space by Adafactor, and Variant 2 is the same as AdaMeM except we use the standard Adafactor instead of one sided Adafactor for the residual gradient. We show their performance in Figure 2 . We note that at smaller ranks, all of these variants outperform GaLore, which shows the importance of using full rank gradients.Momentum and Batch Size. Prior works (Kidambi et al., 2018; Lee et al., 2022; Wang et al., 2024) have given empirical and theoretical evidence that momentum is only helpful at large batch sizes. In Figure 3 , we confirm this for our setup where we vary batch size while keeping overall token The GaLore paper found that GaLore nearly matches Adam even with rank set to half of the full rank. We did not find this to be the case in Figure 1 . We reproduce their experiments in their codebase and show that this was due to non-optimal value of β 2 . They used β 2 = .999, while the standard value used in language modelling (Biderman et al., 2023; Touvron et al., 2023; Wortsman et al., 2024) is β 2 = .95, which is what we have used for our experiments. Table 1 gives the performance of GaLore and Adam for β 2 = .95 and β 2 = .999. We see that• β 2 = .95 outperforms β 2 = .999 for both algorithms.• For β 2 = .95, Adam outperforms GaLore with rank set to half of the full rank.Since we are reproducing these experiments from the Ga-Lore (Zhao et al., 2024 ) work, we did not run GaLore for additional steps as described in Appendix B. 
limitations: 
 The strongest limitation of our work and of other low rank optimizers such as ReLoRA and GaLore is that they are only practically useful to reduce memory in situations where momentum is helpful. This is because if momentum is not helpful then one can use Adafactor which has negligible memory overhead2 above the weights themselves. In particular, the other memory usage is due to activations which grows with batch size and as prior works (as well our results in Section 6) have found, momentum is more beneficial with large batch sizes. In this work we have focused on optimizer memory usage, we leave the exploration of regimes in which optimizers such as ours can reduce total memory usage to future work. 
conclusion: 
 In this work, we propose a memory efficient variant of Adafactor called AdaMem. It combines the preconditioning of Adafactor with low rank momentum ideas of works like LoRA and GaLore. We empirically show that this leads to significantly better performance for low-rank momentum as compared to GaLore, and also provide a theoretical justification for the method. Our method provides a specific recipe for combining preconditioning and momentum in a low-rank subspace, and further exploration within this space is a promising research direction. 
 
 Appendix: This holds under the AdaLomo implementation of Adafactor as discussed in Appendix A. 


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!