LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics
Authors: Thomas Robert
Mher Safaryan
Ionut-Vlad Modoranu
Dan Alistarh
Venue: ICLR 2025 Poster (2025)
Keywords: adaptive optimization
memory efficiency
low-rank learning
low-rank compression
convergence rates
Primary area: optimization
TL;DR: We propose a new optimizer, LDAdam, that reduces the optimizer's memory footprint by performing low-rank updates while exploring the full parameter space.

Abstract: We introduce LDAdam, a memory-efficient optimizer for training large models, that performs adaptive optimization steps within lower dimensional subspaces, while consistently exploring the full parameter space during training. This strategy keeps the optimizer's memory footprint to a fraction of the model size. LDAdam relies on a new projection-aware update rule for the optimizer states that allows for transitioning between subspaces, i.e., estimation of the statistics of the projected gradients. To mitigate the errors due to low-rank projection, LDAdam integrates a new generalized error feedback mechanism, which explicitly accounts for both gradient and optimizer state compression. We prove the convergence of LDAdam under standard assumptions, and provide empirical evidence that LDAdam allows for efficient fine-tuning and pre-training of language models.

Published as a conference paper at ICLR 2025
LDADAM: ADAPTIVE OPTIMIZATION FROM Low- DIMENSIONAL GRADIENT STATISTICS
Thomas Robert', Mher Safaryan', Ionut-Vlad Modoranu', Dan Alistarh? Institut Polytechnique de Paris (IPP)
Institute of Science and Technology Austria (ISTA)
ABSTRACT
We introduce LDAdam, a memory-efficient optimizer for training large models, that performs adaptive optimization steps within lower dimensional subspaces. while consistently exploring the full parameter space during training. This strat- egy keeps the optimizer's memory footprint a fraction of the model size. LDAdam relies on a new projection-aware update rule for the optimizer states that allows for transitioning between subspaces, i.e., estimation the statisties of the projected gradients. To mitigate the errors due to low-rank projection, LDAdam integrates a new generalized error feedback mechanism, which explicitly accounts for both gradient and optimizer state compression. We prove the convergence of LDAdam under standard assumptions, and show that LDAdam allows for accurate and efficient fine-tuning and pre-training of language models. Code is available https: /tttps://githune com IST DAS-AAA
INTRODUCTION
Scaling up deep neural networks leads to ever-improving models (Kaplan et al. 2020), but also to over-increasing memory and compute requirements. In this work, we focus on the memory cost of optimizer states, i.c., the additional first- and second-order gradient statistics required during the adaptive optimization process, which in the case of the baseline algorithm Adam (Kingma & Ba, 2017 Loshchilov & Hutter 2019) reaches up twice the size of the model. Our goal is to replicate Adam's 's performance and theoretical guarantees while significantly reducing its memory footprint.
We start from the observation that the optimization landscape of deep networks exhibits low intrinsic dimensionality (Li et al. 2018). Low-rank learning has been successfully leveraged for model fine- luning, with the LoRA technique (Hu al. (02T) and its many extensions (Dettmers et al.) 2023 Liu et al. 2024 Nikdan et al.) (024) becoming standard. However, such adapter-based methods do not extend to end-to-end accurate training (Lialin et al. 2023; Zhao al. (2024), nor even to fine- tuning on more demanding tasks (Zhang et al. 2023] Nikdan et al. [224 Liu et al., 2024). Despite promising results such as GaLore (Zhao et al., 2024). designing memory- -fficient full-parameter optimizers which leverage low-dimensional gradient statistics with provable convergence and good practical performance remains an open challenge.
Contributions. We propose Low-Dimensional Adam (I.DAdam), a memory-efficient adaptive opti- mizer which, for the first time, incorporates low-rank compression of both gradients and optimizer states, with convergence guarantees under standard assumptions. The goal of the algorithm is to perform optimization steps within subspaces of lower dimension, while introducing mechanisms to maintain relevant information from the last seen subspaces and 2) recover information lost due to low-rank compression. To reach this goal, our work is based on three main techniques. First, we formulate projection-aware update rules that allow to perform adaptive optimization on dynamically changing coordinate systems and subspaces. Second, we leverage block power iter- ation (Bentbib & Kanber 2015) at each step to efficiently and accurately project the gradient and optimizer states. Third, we use generalized eиror feedback mechanism for both gradients and op- timizer states. Moreover, we prove the convergence of I.Addm under standard assumptions and provide empirical evidence supporting its training efficiency.
Work performed intern ISTA, Correspondence dan. sstarrhii ac.at.

---

Published as a conference paper at ICLR 2025
When performing adaptive optimization within low-dimensional subspaces. key challenge is to adapt the optimizer states to changes in the projection map. Adapting gradient statistics to new subspaces is not obvious and comes with two challenges. First, standard adaptive optimization algorithms are anisotropic: that is, the parameter-wise learning rate tuning implies that the choice of coordinate system used to express the gradient matrix impacts the optimization process. Second, the second-order statistics are nonlinear, which complicates estimation of the gradient statistics and error compensation. We address the first challenge by carefully selecting the coordinate system used for low-rank compression, i.e. the basis given by the most significant singular vectors. The second challenge is addressed by approximating the gradient statistics in the current subspace given the raw moments estimates in the previous subspace. Thus, with LDAdam, optimizer step is performed from the projection-awan aware intermediate optimizer states and the projected gradient, rather than from the optimizer states at the previous iteration and the full-rank gradient (Kingma & Ba 2017). In turn, this leads to stronger convergence guarantees, and better practical performance.
The choice of projection map each step is key to low-rank compression. Performing singular value decomposition (SVD) at each step would add prohibitive overhead: instead, we use block power iteration with warm initialization (Bentbib & Kanber 2015 Vogels et al. 2020), which can achieve a tight approximation of the most significant singular vectors of successive gradients, at low cost. To preserve the information from the last iterations kept within the optimizer states, we show that one can perform the block power iteration on a barycenter of the newly computed gradient and an average of the previous gradients. Finally, mitigate the errors due to low-rank projection, we introduce a new variant of the error feedback mechanism (Seide et al. 2014; Karimireddy et al. 2019) Alistarh al. 2018) that supports compression of both gradient and optimizer states. Moreover, we propose a memory efficient implementation of it by reusing the space allocated to gradient accumulation.
The resulting LDAdam method is memory- cost of optimizer states is nr + 2rm, where is the projection rank, while the cost of Adam's 's state 2nm. This matches the efficiency of the popular Galeore algorithm (Zhao et al. 2024), the closest prior work in this area; and can bring major gains in practice, as usually < min(m,, n). Relative to GaLore, LDAdam has the advantages of tracking gradient evolution in the adapted subspace at each step, and correcting for the projection error. This leads to better practical performance, but also to theoretical guarantees of convergence under standard assumptions. Specifically, for smooth non-convex objectives, LDAdam can preserve the asymptotic convergence rate of AMSGrad (Reddi version of Adam (see Theorem D.. Moreover, we show even
for a weight layer of shape n x m, the space
eet al. 2019), the faster rates for objectives that obey the fllaak-Lojasiewizz condition (see Theorem 2).
We validate the practical performance of LDAdam via an efficient Py PyTorch implementation. We apply LDAdam for fine-tuning RoBERTa (Liu et al. (019) and Llama-family (Touvron et al.) 2023) models on the GLUE (Wang et al. 2018) and Grade-School Math (GSM) (Cobbe et al. 2021 benchmarks, respectively. In addition, we provide pre-training results for Llama-type models with 130M and 350M parameters. Across all tasks and models, LDAdam shows competitive accuracy relative to Adam, but with much smaller optimizer states relative to the baseline, and comparable runtime. Relative to GaLore, LDAdam shows consistently higher accuracy, for matching memory costs and similar runtime. Ablations show that this is a consequence of both our projection-awes intermediate updates and error correction.
In summary, we show that it is possible to achieve both theoretical guarantees and strong practical performance when performing adaptive optimization from low-dimensional gradient statisties. Our techniques may also be relevant to areas such as distributed optimization, while our implementation is relevant in reducing the practical memory overheads of training large language models (LLMs).
2 RELATED WORK
Memory-Effiiiere First-Order Methods in Optimization. Gradient statistics are used to deter- mine the descent direction (Rumelhart et al., 1986) and to adapt the learning rate per parameter (Duchi et al., 2011 Hinton, 2012). Adam (Kingma & Ba 2017) combines both ideas in a bias- corrected design, and AdamW (Loshchilov & Hutter 2019) adds weight decay. Although it is the de facto optimizer for deep neural networks, Adam/Adam W has been the subject of several exten- sions aimed at reducing its memory footprint. For example, Adafactor (Shazeer & Stern 2018), CAME (Luo et al.) (023) and Adam-min (Zhang et al., (024a) rely on a factorized representation of the second-order statistic reduce memory. Quantization has also been proposed to reduce the
2

---

Published as a conference paper at ICLR 2025
footprint of optimizer states (Dettmers et al. 2022), and fusing computation of the backward pass and the optimizer step (Lv et al. 2024b a has also proven effective.
Error Feedback in Distributed Optimization. Several methods have been proposed to reduce the communication overhead in distributed optimization: quantization (Alistarh et al. (017). sparsifica-
tion (Alistarh et al., 2018 Lin et al. 2020) and low-rank approximation (Vogels 2020).. Error
feedback (Seide et al. 2014) has been shown to be effective in improving the empirical performance
of biased gradient compression methods (Karimireddy et al. 2019; Vogels al.) 2020) methods and
key to ensuring their convergence (Stich et al., 2018 A Alistarh et al. 2018). MicroAdam (Modoranu al. 2024) only stores sparse gradients from each optimization step, and uses error feedback for
correction, but reconstructs the full optimizer state at each step, instead of compressing it.
Parameter-Efficiet Fine-Tuning (PEFT), natural approach for reducing memory overheads is
to train less parameters (Han et al. 2024). This idea led to additive PEFT, where new trainable modules are added to the model architecture (Houlsby et al., [2019] Li & Liang] 2021); selective PEFT, where only certain parameters of the original model are retrained (Guo al. [021]; and
reparametrized PEFT, where new modules are trained and merged with the original parameters (Hu
et al. 2021] Nikdan et al.) 2024) Liu et al. 2024). However, methods that train significantly less parameters tend to underperform for pre-training and fine-tuning on challenging tasks.
Optimization via Low-Rank Gradient Projection. GaLore (Zhao et al. 2024), which partly mo- tivated our work, performs adaptive optimization in a now-dimensional subspace that is updated periodically. The adaptive optimizer is fed with low-rank representation of the gradients obtained via SVD, thus saving memory. and produces a low-rank representation of the descent direction, which is then up-projected to update the model parameters. The idea is similar to PEFT in that low-rank gradient projection yields selective training in the coordinate system induced by SVD. The addition of low-rank components to achieve full-rank training is also proposed in ReLoRA (Lialin iet al. (023) or COLA (Xia et al. 2024). However, unlike PEFT methods, GaLore computes and exploits the full gradient of the original model. The final algorithm allows pre-training of 7B param- eter language models on consumer GPUs. However, it can only do so with very small batch size and without the ability to accumulate gradients over multiple mini-batches.
GaLore Comparison. Relative to GaLore, we propose the following new techniques: (1) we in- troduce a new projection-aware update rule (see Sections 3.1 and 3.2). which avoids accumulating low-rank gradients from different subspaces into the optimizer states, and performs frequent updates; (2) to correct for projection errors in both gradients and optimizer states, we introduce a new gener- alized eиror feedback mechanism (see Section 3.4). The convergence of GaLore is only guaranteed under a very strong "stable-rank" assumption, whereas we show convergence under standard con- ditions. Experimentally,
show that these improvements lead to consistently improved accuracy, and that they can be implemented efficiently in terms of both memory and runtime.
Several orthogonal improvements such quantization (Dettmers al. 2022 Zhang et al. 2024b), per-layer weight update (Lv et al. 2024b), and weight factorization (Jaiswal et al. 2024) have been proposed to further improve projection based methods. Many of these improvements compatible with LDAdam, and we plan to investigate them in future work.
3 THE LOW-DIMENSIONAL ADAM (LDADAM) ALGORITHM
3.1 THE NEED FOR A PROJECTION-AWARY UPDATE RULE
Standard Adaptive Optimization. We denote the loss function to minimize, Or the model param- eters at step i, G. = (f. E RR the gradient at 0,, St its mini-batch stochastic counterpart, ?) the learning rate, and a positive scalar used for numerical stability. The standard adaptive optimiza- tion relies on the gradient statistics estimates m, and Uz, which are obtained from the exponential moving average the gradients rate 31 and the squared gradients at rate B2, respectively. With their unbiased counterpart, Adam's 's update is: 0++1 0. 1/: Vorte
vi, I-3 and ü, 1-33
The Challenges Behind Low-Rank Updates. When adapting the learning subspaces to capture the low-rank gradient structure, one expects the optimizer states to retain information from previous iterations. However, frequent updates of the low-rank projection causes significant challenges: first, projections lossy and gradient information is lost; second, since the projections may be different
3

---

Published as a conference paper at ICLR 2025
between steps, subspace adaptation may alter the representation of the low-dimensional optimizer states. To illustrate this second point, we briefly discuss the GaLore solution (Zhao et al. (024).
Let be the compression rank, and U4, and Uc be the projection and back-projection matrices use at step t. 24, E æ yields a truncation to its first r rows of the gradient matrix written in an adapted coordinate system, leading
low-dimensional representation of the gradient with shape r x m. Multiplication with the matrix Ur is then used to express the compressed gradient in the canonical high-dimensional coordinate system. From now on, let my and it (resp. fi, and 6t) be the first and second moment estimates (resp. unbiased estimates) in the low-dimensional space induced step by 24,", and let M, Ue my denote the back- projeeeee low-rank counterpart of m
In this context, Equation I describes the GaLore dynamics (Zhao et al.
2024 Algorithm 2
O2+1 = 0r 17144
1-4 di = 0g v0 1- Σ Vortc
8f U4 u 9r-
(I)
Notice that, the above equation, the change in the low-dimensional projection basis between steps is ignored: can be seen in the right-hand sum, gradients are projected and brok-projected using different maps, i.c. Ue and 24,", respectively. Moreover, the statistics of gradients projected into different subspaces are accumulated on the same momentum buffers. The issue we raise here (highlighted in bold in Equation is partially mitigated in practice by heuristics: occasional updates of the projection map and exponential decay rates lead to only a fraction of the steps being effectively resulting from accumulation of compressed gradients in different subspaces. However, we will see experimentally that this leads to a drop in accuracy relative to the "correct" projections.
3.2 THE LDADAM UPDATE RULE
Overview. With LDAdam, we address two key challenges in low-rank updates: we introduce a projection-aware update rule that accounts for the change in the low-dimensional representation of the optimizer states, and we provide a generalized error feedback mechanism that accounts for the loss of gradient accuracy due low-rank compression. Compared to standard adaptive optimization step, optimizer states are replaced by low-dimensional intermediate optimizer states, denoted by m 1/-1/2; and the gradient is replaced by the accumulator 0₂ U, A, obtained by low- rank projection of the high-dimensional accumulator A, = + Er, where & is the (generalized) error feedback buffer, storing projection errors. In the following, we build the LDAdam algorithm step-by-step. The end-to-end construction is provided in Algorithm
Our first objective is find mechanism transfer gradient information from one lower dimen- sional vector space to another. For this, we take inspiration from orthogonal projections, which pro- vide the best approximation in terms of f2 error. The projection-ave awaree update rule for the first mo- ment estimate (A (Algorithm) | [8ine(8) follows from the multiplication with the matrix 4," -M--1, which stands for both the projection matrix and the change of basis matrix. The matrix U, -U-1 E Rexr is low-dimensional and thus allows efficient transition between subspaces, no intermediate high- dimensional vector is created. We obtain the following update:
m 31 PROMU (M+-) + (1- 31)a, = 31 S. -4-1 1 +(1 -31aa.
(2)
Optimizer States Compression. Note already that projection yields compression. Yet, the discrep- ancy between successive subspaces leads to loss of information contained in the optimizer states. We address this by enforcing slight shifts between subspaces, and by generalizing error feedback to account for the updating of optimizer states. The information lost on gradient momentum due to compression, and that needs to be reintroduced via the error feedback mechanism, is given by:
M U.,, M = Ut-1 m Ui m4-1/22-
Estimating the Statistics of Projected Gradients. Although reliable, the orthogonal projection of the optimizer state does not apply to the estimation of the second-order statistic. In particular, Adam does not rely solely on linear operations and is anisotropic (see Appendix A). To overcome these challenges, we interpret Adam's 's optimizer states as statistical estimates of the first two moments of each gradient coordinate. Indeed, Adam's optimization step holds for an approximation of (Ge)) / &,,,, (G (T)ss. where & (e1,...,e,, en) is the basis provided by the canonical
4

---

Published as a conference paper at ICLR 2025
parameterization of the model, G (G, es) a column of the gradient matrix, and Er, is the exponential weighte expectation from time with decay-rate 3.
We rely on this observation to rewrite Adam's update in any given coordinate system Ee+1: we build M4+1/2 and 0/41/2 as statistical estimates of (EEr,o, Ge and (E,, ((G Di)gn. It follows from G Ejgn't ej)es that the matrix of change of basis ((éi, ey))i.sn provides transition between coordinate systems. Its truncation i, I induces rank-r orthogonal projection of the gradients, and thus holds for the transition between subspaces:
R G Σド= (cs, ( Ra,s, (G$) Ej=( ey)()) = (U,+1 Ub m.)):
(4)
Linearity of expectation leads to an autoregressive update rule for the first moment estimate (see Equation 4. With e))k, = (24++1 5,) the estimation of the first-order statistic in the coordinate system ε resolves to an orthogonal projection when considering only the first r com- ponents (i.e., after truncation to i,j < r), and the above discussion on compression of optimizer states remains valid. Following our statistical approach, we obtain the projection-awars update rule for the second moment estimates (Alg (Algorithm [Line
9
from the approximation below:
((G
(E,, e, Er ((G))))))))) En ex) Ex) (GGGG)
= E (b)) + En ew) ei) (i))))
(5)
The approximation of the second-order statistics in coordinate system requires the estimation of the covariance between the gradient coordinates (see Equation [5) Although too large to store, the covariance can be approximated by the product of first-order moment estimates, assuming dependence between gradient coordinates. The latter assumption is enforced on average over the gradient matrix columns by the nature of the coordinate system inherited from SVD. In practice, we guarantee positive estimate of the second-orde -oder statistics by clipping value to zero if necessary.
3.3 LEARNING SURSPACE ADAPTATION
LDAdam's 's memory savings stem from the low-dimensional structure of the moments estimate. The compression strategy, defined by a choice of truncated coordinate systems, sets the exploration of the parameter space, and its design meets two criteria: (i) the learning subspace must be adapted at each step to integrate error feedback; (ii) the compression operator must be set to approximate the current gradient and to preserve the information stored in the optimizer states.
Projection Map Computation. Although singular value decomposition (SVD) provides an opti- mal low-rank approximation, its computational cost makes it prohibitive in
We follow PowerSGD (Vogels et al.) 2020) to efficiently approximate the most significant singular vectors: we perform single block power iteration (Benthib & Kanber 2015) initialized with the approximation from the previous optimization step; and apply the Gram-Schmidt process to derive the projection map (Algorithm [ Line 6). In addition, we propose to start the whole process by computing SVD of the first gradient.
Learning A Smooth Subspace Transition. LDAdam relies each step compression of the gradient and the optimizer states via a unique low-rank projection map. Instead of fitting two pro- jection maps and interpolate between them, we propose to fit a single projection map the interpo- lation of the gradient and an average of the previous gradients (Algorithm [ Line 5). We introduce an interpolation factor and perform a block power iteration to approximate SVD of the matrix: B. p 44-1 m + (1 At. The interpolation factor helps to balance the accuracy of the gradient compression with the preservation the optimizer states. We suggest setting p = B1, i.c. computing the optimal subspace for gradient momentum compression.
3.4 THE GENERALIZED ERROR FEEDBACK MECUANISM
Error Feedback was introduced for distributed optimization (Side et al. 2014) to account for gradi- ent compression errors by reintroducing them into the next iteration gradient. We extend this mech- anism to account for loss of information both the gradient and the optimizer states. We keep the same structure of the error feedback mechanism: the optimizer is fed by a unique accumulator that sum of the gradient and the error buffer.
5

---

Published as a conference paper at ICLR 2025
Unbiased Error Buffer Loading. With gradient and optimizer states decayed at different rates within the optimizer step, we must adjust the error buffer loading strategy the optimizer states projection enror by the factor B1 B1) to recover the first moment estimate. Furthermore, the condition under which the error buffer loading strategy is consistent with the second moment esti- mate is given by: R2 = (1 32)( - )² We suggest setting 32 = 0.99 and B1 R 0.908, as the slight decrease in the decay rate of the second moment estimate helps to adapt faster to a new learn- ing subspace. Equationn]desse describes the loading strategy of the generalized error feedback mechanism
(Algorithm Line I6].
(.(82) Eu+1 = [A, U4 ael (P,) U4-1 M4-1 - U4 m2-1/2]-
Implementing Error Feedback. A naive implementation of the error feedback mechanism would add memory overhead equal to the model size. Since the gradient and the error buffer are added together before the optimization step is performed, we store the error buffer in the vari- able used for gradient accumulation (Algorithm [ Lines [6 and (e.g. in PyTorch, modify the opt immiiiie zer. zero grad function to store the eπror buffer in the p graa variable). This results in an error feedback mechanism that does not require any further memory for optimizer states.
Algorithm L.Dddam (o Practical View Only, 9t E 1R oo Analytical View Only, Dr. E IRd)
Hyperparameters: step size yı: decay rates 81, ,32
L.A Hyperparameetrs: projection rank r; interpolation factor p
LDAdam Hyperparameters: contraction factor q, E [0, 1); p B1
1: Initialization: mg = 0; uo = 0; Ao = Ud = (VD(go); 20 = 0; E1 =
2: for 2 T) do
Gradient accumulation and error buffer unloading
3:
A, A, + gt
4:
o A, = & + St
Learning subspace adaptation
5:
B, - (1- A)
6:
o U4 "BRMMSCHHIDT(BEE. U4-1)
7:
Ui is any d x orthogonal matrix such that I1(7 Uil,. Bell S s.
Optimizer states projection-aware update
8:
M4-1/2 U, U4 1m-1-
9:
1/-1/2 (1 - 8
(21 11-1)2 (δc-1 - (m)-1)22 (21, Ue-1 m
Optimizer states Adam-ty update
at 11 A,
11:
= f1m,-1/2 1/2 + (1 B1 )ar V: 1/210-1/2 + (1 B2)
12:
13:
E. = -maax Ut. 10 maxe)
Model update
14:
st+1 = 0; mill
th
Ve.te
0.41 0; molle Vi.te
Error buffer loading
16:
A++1 = (Ae- Us ar) + 1-8, (4--1 m4-1 - U4 m4-1/2) o E141
17: end for
CONVERGENCE GUARANTEES FOR LDADAM
Analytical Overview. We now present our theoretical convergence results for LDAdam. We first provide/discuss "analytical" perspective on approach in Algorithm Then we introduce and discuss the analytical assumptions used in the theory, along with two theoretical convergence rates.
The outlined steps closely resemble those in the practical view of the algorithm. One notable dif- ference is the introduction
new AMSGrad-type normalization, é = max Ve. أ mmxx) in line 13 This represents uniform version of the original AMSGrad technique (Reddi et al. 2019),
& max(e,, 1, which enforces coordinate-wn monotonicity for the diagonal preconditioning. In our setup, transition between different low-dimensional spaces, the preconditioning matrix

---

Published as a conference paper at ICLR 2025
is no longer diagonal, and scaling factors may change after such transition. To ensure that adaptive preconditioning remains monotonic, we enforce a monotonic spectrum of the preconditioning.
To avoid unnecessary complications in the notations and dcrivations, in this section assume that the stochastic gradients 9t and other variables derived from the gradients, such as momen- tum, are vectors. We remind the reader that, in practice and in our implementation, gradients are viewed/represented as 2D matrices. In our analysis, vector gradients should be treated as the columns of the gradient matrix. Nevertheless, all steps in Algorithm I are still applicable when dealing with matrix gradients. Specifically, given a matrix gradient St. we obtain a matrix B. for which we assume that, regardless of the low-rank compression method used (e.g., SVD or power iterations), the contractive inequality (I Belle S 9r B holds with respect to the Frobe- nious norm. We also incorporate the debiasing factor /I / B4) in T/.
Convergence Guarantees for General Smooth Non-convex Functions. We first outline the ana- lytical assumptions under which we establish LDAdam's convergence guarantees.
Assumption 1 (Lower bound and smoothness). The loss function f: Rd - R is L-smooth and lower bounded by some f E R: || (0) ff()))l - L|@ flll, for any 0, E Rd.
-R
Assumption 2 (Unbiased and bounded stochastic gradient). For all iterates 2 1, the stochastic gradient gr. at Og unbiased and uniformly bounded by some constant G 2 0:
Eigel= v.(.)). 19: S G.
Assumption 3 (Bounded variance). For all iterates 2 1, the variance of the stochastic gradient Ii at 0, is uniformly bounded by some constant 0² 0: E Billg ff())l o
All three assumptions, including the bounded gradient condition, are standard and commonly used in adaptive optimization literature (Reddi et al.] 2019; Chen et al., 2019; Défossez al., 2022 Modoranu et al. 2024). Attempting to relax the bounded gradient condition, some works (Shi et al, 2020 Zhang et al.] 2022, Wang et al., 2024) resorted to using strong growth condition (which is not necessarily a relaxation over bounded gradient condition) or bounded stochastic noise conditions (Li et al., 2023 Hong Lin 2024). Taniguchi et al. (2024) achieved an optimal rate by modifying Adam and only slightly relaxing the bounded gradient assumption by requiring the expected gradient norm be bounded. The bounded gradient condition in Assumption 2 directly implies a bounded variance condition with a constant 0² 2G2. Therefore, for the purpose of asymptotic analysis, one can omit Assumption 3 However, we will show that in the non-convex convergence rate, the -2-term decays at a rate T/VT. while all other G-terms decay at faster rate of 1/T. Theorem states general non-convex convergence result.
Theorem 1 (Non-convex convergence rate). Let Assumptions 7 12 and 3
hold. Then, choosing
step-size 1) min( 41COVV1C5 VT), L.DAdam (Algorithm // satisfies
+ET V (0.)
20h VT
(1(01) - Lo²
+ +
G
1+8 (1 8 (1-8 with constants Co- := V I-32 (1-3,) G2 + e and C2 = (1 (1--1)(1-00))
B1+(1-B1)0
Discussion. Compared to the standard stochastic gradient descent (Gaadimi & Lan 2016). the leading term 1 1/VT of the obtained rate recovers the optimal non-convex convergence speed. Addi- tionally, the asymptotic rate 0 +) aligns with the rate of uncompressed AMSGrad in the stochastic non-convex setup (Zhou et al., 2024). Therefore, the introduced low-rank compression framework along with the error feedback mechanism does not hurt the asymptotic convergence with respect to T. The slowdown caused by compression in the leading term is only through the constant qr (0, 1] measures the aggressiveness of the compression. Note that our theory applies to any low-rank compression with r 2 More compression corre- sponds to smaller rank for the projections and a smaller factor 1 4r. For instance, when using low-rank compression via SVD decomposition, one can prove 9r (
Co ( (T..)) ), where the factor
2) for any rank r 2 In this case, the leading term of the rate becomes O } indicating that the convergence slows down in the worst case by factor of d/r due to compression. The proof can be found in Appendix E.2
Convergence Rate for Non-Convex Functions under the PL Condition. Next, we show faster convergence for LDAdam when the loss function satisfies the PL condition.
7

---

Published as a conference paper at ICLR 2025
Assumption 4 (PL-condition). For some constant μ. 0 the loss function satisfies the Polyak-
Lojasiewica (PL) inequality:
[[]f 2 2 -f), for any E R
Theorem 2. (PL convergence rate) Let Assumptions 1 andd hold. Then, choosing step-size
7 min 2Co log LDAdam (Algorithm
!
satisfies
af
E[f(0T+1)] -
T
2LC p
6Co(1+C,) AVC
+
with constants C == A1+(1-91) and 1o = min( 16r.Co (1--))))--))
Co(1-8: (111-8)(1-0) (1---))
2p
CT,/0000
Discussion. Notably, the leading term of the convergence rate under the PL condition improves to in the general non-convex setting. Faster rates of ( for adaptive optimization algorithms under the PL condition have been established when 82 - 1 (He et al.) 2023) through the of the AMSGrad trick (Modoranu et al. 2024). In LDAdam, we apply uniform version of AMSGrad normalization with arbitrary 32 < and derive the rate for compres- sion with any rank r 2 Hence, similar to the general non-convex case, we show the best-known convergence rate in the leading term, up logarithmic factor. While the last term has higher-order constant dependencies, these are negligible due to the T damping effect. Thus, the theory sug- gests that the algorithm's 's convergence rate remains comparable to that of the uncompressed version, with only a constant factor affected by the compression rank. To quantify the slowdown caused by compression, again consider low-rank compression via SVD decomposition, lier. As noted before, C. = O -r- = O()), and similarly, from the expression of C1. we have G = ( = ( Therefore, the leading term of the rate becomes (( with compression slowdown factor of ($)². The detailed proof is deferred Appendix E.3
0 ( compared to O( [
discussed ear-
TRAINING MEMORY FOOTPRINT AND RUNTIME
LDAdam's 's memory savings over Adam (Kingma & Ba 2017) (see Table for comparison) follow from the low dimensionality of the optimizer states: for gradients of shape m with S the gradient statistics a learning subspace of rank r are stored as tensors of shape r X m, and the additional projection matrix is of shape n X r. As result, the memory footprint of LDAdam's s optimizer states only a fraction of the size of the model. For an n x n layer, the compression ratio of LDAdam relative to Adam is of 3 Table and Figure 3 reports peak memory for fine-tuning and pre-training tasks.
nn
With the error buffer stored in the gradient variable, the error feedback mechanism is memory tral. However, this comes at the cost of not supporting gradient clipping and per-layer weight updates (Lv al., 2024b Zhao et al. 2024). The latter allows memory savings proportional to the model size. but at the expense of gradient accumulation. To be precise, the standard GaLore algorithm does not support gradient accumulation (see Appendix D.2] for detailed explanation), while its newer "retaining grad" version (Zhao et al.) 2024) does, but at the expense of a memory overhead equal to the size of the model. Although applicable to most optimizers, per-layer weight updates is not widely adopted, due to the large batch size used for large language model training (Touvron et al.) 2023 and the memory-intensine nature of activations in attention layers.
Compared to Gol,ore I.DAdam adds optimizer states projection-aware update and the generalized emorr feedhack mechanism. Yet, the practical runtime differences are not substantial: on Ilama- 350M pre-training, LDAdam has 10% longer runtime compared to GaLore, and is 15% slower than Adam. Table 7 and Figure report runtimes, and Table 5 reports theoretical complexity in terms of standard matrix operations.
EXPERIMENTS
We empirically evaluate LDAdam for fine-tuning and pre-training large language models against baseline Adam (Kingma & Ba 2017) and the memory-efficient GaLore (Zhao et al. 2024) (Algo- rithm 2). Similar most parameter-efficient fine-tuning (PEFT) methods and to GaLore, low-rank compression only to the two-dimensional matrices of the self-attention layers, e.g. wide embedding and output layers are trained using standard Adam optimizer. For fair comparison, in
apply
8

---

Published as a conference paper at ICLR 2025
Table 1: Optimizer comparison: parameter count during training for weight layer of shape m with n. S m (i.e., left projection), training capabilities, and estimates of optimizer states memory footprint in half precision (models architecture described in Table 99.
Adam
LDAdam
Gal (retaining grad)
Gal
Token count
Weights
Yum
YU7
TLTT
Y07TL
Gradients
7m7 2mm
1077
71773 mmm
Optimizer States
2rm
2rm
Gradient Clipping
X
X
Gradient Accumulation
/
X
Memory estimates
RoBERTa -aase (r=8)
0.46 GB
0.15 GB
0.15 GB
0.15 GB
Llama 350M (r=256)
1.37 GB
0.95 GB
0.95 GB
0.95 GB
Llama-2 2 7B (r=32)
25.1 GB
1.22 GB
1,22 GB
1.22 GB
Llama-2 7B (r=512)
25.1 GB
4.87 GB
4.87 GB
4.87 GB
all experiments we do not apply additional layer-wise learning rate scaling, and apply hyperparam- eters suggested the papers introducing the algorithm. To ensure reproducibility, the full list of hyperparameters used is provided in the Appendix
Fine-tuning on GLUE. We evaluate LDAdam for fine-tuning RoBERTa-base model (Liu et al, 2019) the General Language Understanding Evaluation (GLUE) benchmark (Wang et al. (018). RoBERTa-bee base is a 125M parameters large language model, built upon the encoder only attention- based BERT architecture (Devlin et al.) 2019), and pre-trate mreitrineed for natural language understanding (NLU). GLUE benchmark includes wide variety of NLU tasks with training samples of various size. We mimic memory-constrained setups by training for only 3 epochs with a batch size of 16, and used rank r=8 compression for LDAdam and GaLore. For a fair comparison, we tune the learning rate over the set (le-5, 2c-5, ..., 5c-5). reports the best average over 3 seeds (Table 8 reports standard deviation), and Figure shows that the error bufferr norm is of the same order of magnitude as the gradient norm, and that their ratio stable throughout training.
Table 2: Results of fine-tuning RoBERTa- model on the GLUE benchmark.
MRPC
STS-B
CoLA 8.5k
SST-2
QNLI 105k
QQP 364k
MNLI 393k
Avg
Training Samples
3.7k
7k
67k
Adam
88.97
90.08
58.57
94.34
92.81
91.38
87.85
86.28
LDAdam
88.40
90.11
59.91
95.00
92.87
91.28
87.81
86.48
LDAdam no-EF
88.00
89,88
56.86
95.00
92.32
89.75
86.99
85.54
GaLore
86.19
88.97
55.12
94.15
92.01
89.86
86.80
84.73
Fine-tuning on GSM8K. We evaluate LDAdam for fine-tuning Llama-2 7B model (Touvron al.
2023) on the GSM8K mathematical reasoning dataset (Cobbe et al. 2021) containing grade school math word problems. We run experiments for ranks r=32 and r= 15512 2, on 3 epochs, with batch size of 32, and tune the learning rate over (5e-5, 6e-5, .... 9e-5, le-4, 2e-4, .... 5e-4J. Table 3 reports the best average accuracy. We also provide peak memory usage, but also the fixed memory cost of storing model, gradient and activations.
Table 3: Results of fine-tuning Llama-2 7B on the GSM8K dataset.
Adam
LDAdam
LDAdam no-EF
GaLore
r=32
5-512
r=32
r=512
r=32
r=512
Accuracy
34.72
32.53
35.86
30.70
35.78
27.07
35.18
Peak Memory
55.34 GB
32.08 GB
35.58 GB
32.08 GB
35.58 GB
32.52 GB
35.58 GB

---

Published as a conference paper at ICLR 2025
Across all fine-tuning experiments, LDAdam outperforms GaLore, and is either par better than Adam. When ablating the LDAdam techniques, we observe that LDAdam without error feedback also outperforms GaLore, supporting LDAdam's continuous subspace adaptation strategy over Ga- Lore's ' sequential update approach. The fact that we can match Adam once error feedback enabled suggests that this technique is key.
Pre-training on C4. We evaluate LDAdam for pre-training Llama models (Touvron et al. 2023) on the C4 dataset (Raffel et al., 2023). Due to limited compute resources, we pre-train smaller 130M, 350M and 1.3B parameter variants (Liainn et al., 2023) Zhang & Sennrich 2019) Shazeer 2020) of the Llama-2 models. C4 is clean version of Common Crawl's web crawl corpus. We adjust the number of training steps according to the Chinchilla scaling law (Hoffmann et al. (022), i.c. 20 training tokens per model parameter; and use batch size 512 and sequence length 256, accounting to 131 072 tokens per step, thus using the same settings GoLore. For a fair comparison, we tune the learning rate over the set (5e-4, 1e-3, Se-3), and report the best results for ranks r=8 and r=256 in Table Figure depicts the pre-training dynamics.
Table 4: Results (validation perplexity 4) of pre-training Llama models on the C4 dataset.
Adam
LDAdam
LDAdam no-EF
GaLore
r=8
r-256
r=8
x=256
x=8
r=256
Llama 130M
24.64
23.82
22.65
45.17
25.32
67.78
26.04
Llama 350M
18.08
18.37
17.30
19.85
20.03
Furthermore, we scale the experiments up to the Llama 1.3B parameter model, using the same hyperparameters but fixing the learning rate 5e - due to limited computational resources. Adam achieves a validation perplexity of 14.86, while LDAdam with rank r= 16 6 achieves 14.09.
Figure I: Pre-training dynamics for Llama 350M (left) and Llama 1.3B (right) on the C4 dataset.
Adam
Validation Perplexiity (4)
Adam
Validation Perplexity
25
LDAdam
30
LDAdam (r-16)
LDAdam (r=256)
30
LDAdam no-EF (r-256)
25
Galore (r-256)
25-
20
20-
15
10k
20k
30k
40k
50c
23k
75k
125k
175k
Update Steps
Update Steps
Pre-training experiments also show that LDAdam has comparable performance to Adam, both in terms of convergence and validation accuracy. and confirm the improvements over GaLore. Further, for very high compression rates (i.e., r=8), LDAdam still converges, but its no-EF variant does not, nor does Gore. This highlights the need for an error correction mechanism for pre-training tasks, which are known require models updates of high rank. For lower compression (i.e. r=256), LDAdam outperforms the no-EF version and GaLore, which is correlated to their consistent pro- jection errors. For more details the impact of the rank on training, see Figure LDAdam even appears to slightly outperform Adam. We attribute this small improvement to compression, as sparse updates might enhance implicit regularization (Zhang al. 2017 Neyshabur 2017).
DISCUSSION
We proposed new low-rank learning method with the memory-efficiett LDAdam optimizer, an- alyzed its convergence under standard assumptions, and provided empirical evidence of its ability to match the performance of Adam at fraction of its memory cost. LDAdam relies Adam to estimate gradient statistics, but its development required the design of a specific intermediate opti- mizer state update rule. Therefore, efficient implementations of Adam not directly applicable to LDAdam, and extending our work to other optimizers requires further work. On the experimental side, natural next step would be to execute large-scale billion-parameter pre-training experiments.
10

---

Published as a conference paper at ICLR 2025
8 ETHICS STATEMENT
This paper presents work that aims to advance the field of machine learning. We believe that memory-efficient optimization is step toward democratizing large-scale model training, and thus provides opportunities to foster both the development of new applications and the research in the field. There are many societal concerns about the rapidly growing use of artificial intelligence, feel that none of them is specifically related to our work.
REFERENCES
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Communication--ffiieenn SGD via Gradient Quantization and Encoding, 2017. URL https:
QSGD:
/aaxiv.org/abss /1600. 02132
Dan Alistarh, Torsten Hoefler, Mikael Johansson, Sarit Khirirat, Nikola Konstantinov, and Cédric Renggli. The convergence sparsified gradient methods, 2018. URLhttps:/arxiv.oom
abs/1800 10505
Abdeslem Bentbib and A. Kanber. Block Power Method for SVD Decomposition. Analele Stüintifice ale Universitatii Ovidius Constanta, Seria Matematica, 23:45-58, 06 2015. doi:
20015//aoom2215-0002
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the Convergence of A Class of Adam- Type Algorithms for Non-Convex Optimization. International Conference on Learning Repre- sentations, 2019. URL httpttsss///xii org abg/ httaabbs/1..... 02941
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichrro Nakano. Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv
brg//bs/2110.
Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Reseurch, 2022. ISSN 2835-8856. URL
www.:eee/fforrie https: pettps://openreviee www.:nee/ooommmidd/22geeee TSWA
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization, 2022. URL
http://aaaxv..og/abb/111..0.00 02861
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Ilms, 2023. URL http::tttp/:aattaae http:/wwwww..pg/bbb///...011
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. 2019. URL https::/arxiv..rrr
abs/18101 04805
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):1111-215, 2011. URL
http:// wttppppppppp..../////////o.
hwtp://///ccaaaaa.eeeeaaaaaaaaaa html
Saeed Ghadimi and Guanghui Lan. Accelerated gradient methods for nonconvex nonlinear and
stochastic programming. Mathematical Programming, 2016. ISSN 156(1-2):59-9) 99. URL
httts:://axi:: http://///axxi.ooo////sse wttp////////......3.3 3787
Demi Guo, Alexander M. Rush, and Yoon Kim. Parameter-efficient transfer learning with diff pruning, 2021. URL https: //axivv. httgg//bb//222... 07463
Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. Parameter-efffiient tuningg for large models: A comprehensive survey, 2024. /tL.ttpps////xxi...g//aa///////a
14608
Meixuan He, Yuqing Liang, Jinlan Liu, and Dongpo Xu. Convergence of adam for non-convex ob- jectives: Relaxed hyperparameters and non non- ergodic case, Journal of Machine Learning Research, 2023. URL htt.hhtts:///aat. wwwprrwwwwoovooooee /2300b//2332222 /230 11782

---

Published as a conference paper at ICLR 2025
Geoffrey Hinton. Lecture 6e rmsprop:
Divide the gradient by rumning average of its recent mag- nitude. COURSERA: Neural Networks for Machine Learning, 2012. https: /wwww cs. aaronto.edu/-Et. wwte://mae//ssieeeeeeeeeeeee /eettureeeerrreeeeee slides 1ec6.pdf
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL
2203 15556
Yusu Hong and Junhong Lin. On Convergence of Adam for Stochastic Optimization under Relaxed Assumptions. Advances in Neural Information Processing Systems, 2024. URL https://
arxiv. ottlbborggbbb.. htwpg//gaa/..a.../////// 03982
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp, 2019. BRthtsts://arte
org abs 000.000751
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Alu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
7/
09685
Ajay Jaiswal, Lu Yin, Zhenyu Zhang, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. From galore to welore: How low-rank weights non-uniformly emerge from low-rank gradients, 2024. URL/ttps:: /ttppps://aaxi..../// org abs/ 2407 11239
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL http::tps////c /ttp:s///axxvvv.o/ org/ abs 2001. 08361
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian U. Stich, and Martin Jaggi. Error feedback fixes signsgd and other gradient compression schemes, 2019. URL https://aaxxii org/
abs Abs/1901.c 09847
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL https://arxiv: wttp:bbbb/bb//////// http:ss//rra.v.oorr///////111
6980
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes, 2018. URL hnttps://arxxvv.or///ee//
08838
Haochuan Li, Ali Jadbabaie, and Alexander Rakhlin. Convergence of Adam under relaxed assump- tions. Advances in Neural Information Processing Systems, 2023. URL https://arxiv:
org hac//bb//200 2304 13972
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation, 2021. URL
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky, Relora: High-rank training through low-rank updates, 2023. URL http://///aaaaooor///////0000/5/
Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training, 2020. URL BtRppttss:///ttt
org/ abs///b//aa 712. 01887
Shih- Yang Liu, Chien-Ya Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu- Yu-Chiang Frank Wang. Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight decompose low-rank adaptation, 2024. URL
https://a/ee www.rriwvvrgooo/// haagg/as//20 wwwp///2../22..22
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019. URL htL.https:/ puttts://arxive
12

---

Published as a conference paper at ICLR 2025
Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https:
Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang. Xin Jiang, and Yang You. Came: Confidence-
guided adaptive memory efficient optimization, 2023. URL https://arxiv.orr/aaa/
2307. 02047
Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimization with adaptive learning rate, 2024a. http::://arx....r////aaaaaaa///a
Kai Lv. Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources, 2024b. URL https://arxxx
org/ arg/aba// 2306. 09782
Nonut-Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtic, Thomas Robert, Peter Richtarik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space over- head and provable convergence, 2024. URL
htwp://w..r//aas//20.115553.3 155993
Behnam Neyshabur. Implicit regularization in deep learning, 2017. Utt.httss:/axxxivoo
abs/ 2t55/109..0990
Mahdi Nikdan, Soroush Tabesh, Elvir Crévvvćć and Dan Alistarh. Rosa: Accurate parameter- efficient fine-tuning via robust adaptation, 2024. URL
04679
Pytorch Tutorials. How save memory by fusing the optimizer step into the backward pass, n.d. URI.ttsp..//ptooocc.or// hutgriutoriaa lwaalssinntn s/intermeeddeeeeeeeeeee/ee/e ate/ cttte//pptiiii.. zer_step in backward__tttorill..
Accessed: 2024-11-16.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL
10683.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of Adam and beyond, arXiv preprint 4rXiv:1904.09223, 2019. URL
09237
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back- propagating errors. Nature, 323(6888) 1986.
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs, In Fifteenth Annual Conference of the International Speech Communication Association, 2014.
Noam Shazeer. Glu variants improve transformer, 2020. URL https://arxiv.org////
0002.05200
Noam Shazeer and Mitchell Stemn Adafactor: Adaptive learning rates with sublinear memory cost, 2018. UtRlhhttps///ariii
/abs/ 1804. 04235
Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun. Rmsprop converges with proper hy- perparameter. International Conference on Learning Representations, 2020. URL https:
URL http://wwwwie.eeeeeaffi.d-----d.
Sebastian U. Stich, Jean- Cordonnier, and Martin Jaggi. Sparsified sgd with memory, 2018. 07599
Shohei Taniguchi, Keno Harada, Gouki Minegishi, Yuta Oshima, Seong Cheol Jeong, Go Nagahara, Tomoshi liyama, Masahiro Suzuki, Yusuke Iwasawa, and Yutaka Matsuo. ADOPT: Modified Adam Can Converge with Any 3n with the Optimal Rate. Advances in Neural Information Pro- cessing Systems, 2024. ht
htt.v//orggab/bb/aeeeeoo org abs /2411. 02885
13

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL
http:////aaaaa...coaa/////////// /200///230.. 09288
Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. Powersgd: Practical low-rank gradient compression for distributed optimization, 2020. URL httpp://axx.v..oo///bb//////////
13727
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE:
A multi-task benchmark and analysis platform for natural language understanding. In Tal Linzen, Grzegorz Chrupala, and Afra Alishahi (eds.). Proceedings of the 2018 EMNLP Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.1 10/33//1/11-55444 URL DELps://acLLantoooooyy.ooo org/N18 5446
Bohan Wang, Yushun Zhang, Huishuai Zhang, Qi Meng, Zhi-Ming Ma. Tin-Yan Liu, and Wei Chen. Provable adaptivity in adam. International Conference on Knowledge Discovery and Data Min- ing, 2024. URL https: /arxiv org/ abs 09900
Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of language models via residual learning, 2024. URL
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https://
orx.v.org//bs//1100 07467
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization, 2017. URL
1.611.0553
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao Adalora: Adaptive budget allocation for parameter- efficient fine-tuning, 2023. httttt/ss//aaxii..rr//asass//sss
a23///33b//3333.... 10512
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. Advances in Neural Information Processing Systems, 2022. URL /tttts::/aarii..r/////aa// www.:iiivorrraaaoo wbtp://///////e/ 220000
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adnm-mini: Use fewer learning rates to gain more, 2024a. URL https:
16793
Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients, abs/24077 08296
2024b. URL htll/ttts//axxi///
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection, 2024. URL
wwtp://////////..///ooo https://arxizz httt::://axxi.ooo///////222
Dongruo Zhou, Jinghui Chen, Yuan Cao, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. Transactions on Machine Learn- ing Research, 2024. ISSN 2835-8856. URL aattps://ppenzevze net/ htt//forrmmic.
GhCcxhbz3c Featured Certification.
14

---

Published as a conference paper at ICLR 2025
A CHALLENGES FOR ESTIMATING THE SECOND-ORDER STATISTIC OF PROJECTED GRADIENTS
Adam Does Not Rely Solely Linear Operations. Adam follows the proposition from AdaGrad Duchi al. 2011) to use diagonal approximation to estimate: ((1-32) Σ " This leads to 2tg a:Aaam = (1 32) Σr= $2 9r @ 9r and to the matrix inverse square root being replaced by coordinate wise inverse square root operations. Such approximation saves both memory and computation, but at the cost of leaving the linear framework. Therefore, one cannot expect to draw an autoregressive linear algebra formula to update the second moment estimate.
Aam is Anisotropic. Adaptive optimization learning rate tuning for each parameter is equivalent direction-wer srect rescaling of the training space. It implies that the choice of coordinate system affects the optimization steps. Let / be a fixed orthogonal matrix which acts as rotation in the parameter space and consider the following two minimization problems:
(Pr) arg minge f(0) (Pz)) arg mi f(2
(8)
These two problems are equivalent up to a change of parametrization. However, applying Adam to each of them will lead to optimization process that are not equivalent: adaptive estimate of the second-order statistics of the gradient leads to 0, # U 0.:
(1 11) E 81 v((.))
B [V)(0.)]
0c+1 0,
1/ 1-31
0: 11p Ve, ((t:)))
(9)
v(1 92) E 83 (VS(0,))2+
0111 = 01 ?t
1-85
(1-81) E:
V)((2-0))
= 0g VE,,e, (0 (
B W V
(10)
v 1-8
(1- 3.
V)(--.)))++
Low-rank adaptive optimization follows from (Pz) where E B is the truncation of the rotation matrix W to its first rows, and thus yields reparametrization and projection to a lower-dimensional space. However, this leads to 0 _g and the optimization set is restricted to low-rank subspace. To enable full- cill-parameter training, GaLore's strategy is to keep the model parameter high-rank while performing low-dimensional updates, leading to:
@GALORE (GALORE-LL -U.
V1 V1-8 I-8, V(1-B2)
(1 B V (iffator B2) Σ B (UT V/(8fatmm))24e
(GALORE
-ut (To
the Voite
(11)
With the projection map L1 U4 occasionally updated, GaLore's low-rank gradient statistic estimate my (resp. ε,) fails to approximate E,,g, [24" (((fiaT.O)))] (resp. E.,sI1(4, V (((G....)))) )))) be- cause gradients projected into different subspaces accumulated the same momentum buffers. simple idea would be to fix a coordinate system for the momentum buffers. However, this would have be high-dimensional and thus would prevent memory savings. The purpose of LDAdam's 's projection-aware update rule is to adapt the optimizer states to the new learning sub- space and coordinate system, leading to m;-1//2 R E,-1,9, [ [.(((.-AAAM)) ), and 0;-1/2 as E,-1,0, [(4, V (()--DAAAM)) and enabling gradient accumulation within now-dimensional momen- tum buffers. One can thus rewrite LDAdam's 's heuristic performing the following update:
@LDADAM so OLDADAM Udg O
Er B. V (OLDDADS
(12)
EE,,, (14, V (ILLDAADO
As result, while Adam is anisotropic, LDAdam is isotropic since it doesn't rely on any specific coordinate system but instead always uses the one induced by the SVD of its gradient.
15

---

Published as a conference paper at ICLR 2025
B ADDITIONAL MEASUREMENTS
B.I TRAINING MEMORY FOOTPRINT AND RUNTIME
Table 5 reports theoretical complexity in terms of standard matrix operations. We write ORTHO(r; m) the complexity of orthogonalizing a family of vectors of dimension m. our im- plemention, similar to powerSGD (Vogels et al. 2020), we apply power iteration the Gram-Schmidt process which has complexity O mm).
Table 5: Optimizer comparison: matrix operation count for a weight layer of shape m with n S m (i.e., left projection).
Adam
LDAdam
Learning Subspace Adaptation Projection-awee aware Update
3
ORTHO(r,
3 m + r
Adam-type Update
m
Descent Direction
3 x +m
+
Error Buffer loading
+3*
Table&ojeport peak memory for fine-tuning and pre-training on a single NVIDIA H100 80GB GPU with micro batch size and without activation checkpointing.
Table 6: Peak memory for fine-tuning (FT) and pre-training (PT) tasks for micro batch size of 1.
Model
Task
Rank
Optimizer
Adam
LDAdam
RoBERTa-base
MNLI (FT)
r=8
2.39 GB
1.60 GB
Llama 350M
C4 (PT)
r=8
4.20 GB
2.67 GB
Llama 350M
C4 (PT)
r 2-2
4.20 GB
3,01 GB
Llama 1.3B
C4 (PT)
r=1 16
15.20 GB
8.67 GB
Llama-2 7B
GSM8K (FT) GSM8K (FT)
r=32
57.62 GB
31.98 GB
Llama-2 7B
r=512
57.62 GB
35.90 GB
Figure 2 depicts the pre-training dynamics over time, allowing comparison of training efficiency with respect to time.
Figure 2: Pre-training dynamics over time for Llama 350M (left) and Llama 1.3 (right) the C4 dataset.
falidation Perplexity (1)
Adam (r-8)
Adam
LDAdam (r=B)
LDAdam (r=256)
30
no-EF (r=256)
25
GeLore (r=256)
25
20
20
15
12
24
21
48
72
96
120
Training Time (hour)
Training Time (hour)
Figure shows the throughput (token per second) and peak memory (GB) of Adam and LDAdam with respect to rank for pre-training the Llamma 350M model on the C4 dataset. We used a miero batch size of for both to replicate the memory constrained setting. This helps to isolate the effect of the optimisation algorithm itself. See Table 7 and Table 6 for results closer to the standard use case.
16

---

Published as a conference paper at ICLR 2025
Figure 3:
Throughput (token per second) and peak memory (GB) of Adam and LDAdam with respect to rank for pre-training the Llamma 350M model on the C4 dataset, on single NVIDIA H100 80BG GPU, using micro batch size of
Adam
secand
3000
LDAdam
Peak Memory (GB)
2000
LDAdam
SGD (no optimizer state)
32
04
128
250
512
10
64
128
250
512
Rank
Rank
Table reports runtime for fine-tuning and pre-training on a single (except for Llama 1.3B pre- training were four) NVIDIA H100 80GB GPU. We don't activation checkpointing. We use micro batch size of 1 for GSM8K fine-tuning, a micro batch size of 128 for Llama 350M model, and miero batch size of 64 for Llama 1.3B model.
Table 7: Runtime for fine-tuning (FT) and pre-training (PT) tasks.
Model
Task
Rank
Optimizer
Adam
LDAdam
LDAdam no-EF
GaLore
RoBERTa-base
MNLI (FT)
r=8
37m
Olh 07m
Olh 0lm
56m
Llama 350M
C4 (PT)
r=8
21h 59m
23h 12m
LJama 350M
C4 (PT)
r=256
21h 59m
24h 50m
24h 44m
23h 09m
Llama 1.3B
C4 (PT)
r=16
04d 13h 09m
05d 04h 24m
1lama-2 7B
GSM8K (FT)
r=32
50m
56m
55m 01h 07m
01h 08m
Llama-2 7B
GSM8K (FT)
51512
50m
0lh 08m
0lh 09m
B.2 IMPACT OF THE RANK ON TRAINING
Figure 4: Training dynamics and validation perplexity for various rank when pre-training Llama 350M model. For training dynamics we used a single learning rate of 5e to allow comparison between runs and provide results for the first 10000 optimization steps. We report the best validation perplexity for learning rates tuned over the set {5e 4, 1e 3, 5e - 3).
LDAdam (r-8)
Volidation Perplexity (1)
Adam LDAdam
19
Validation loss
5.0
LDAdam (r=32)
LDAdam (r=64)
LDAdam (r#1
LDAdam (5-256)
4.0
17.
2k
4k
8k
16
32
G4
128
250
Update Steps
17

---

Published as a conference paper at ICLR 2025
B.3
FINE-TUNING STANDARD DEVIATION
Table 8: Standard deviation over 3 seeds for the results of fine-tuning RoBERTa-base model on the
GLUE benchmark.
MRPC
STS-B
CoLA 8.5k
SST-2 67k
QNLI 105k
QQP 364k
MNLI
Training Samples
3.7k
393k
Adam
1.068
0.226
2.604
0.517
0.074
0.050
0.167
LDAdam
0.617
0.211
0.633
0.463
0.275
0.080
0.177
LDAdam no-EF
0.849
0.555
0.903
0.066
0.285
0.687
0.087
GaLore
0.861
0.330
2.348
0.413
0.166
0.034
0.111
B.4 ERROR BUFFER NORM DURING TRAINING
Figure 5: Error buffer norm and gradient norm during the fine-tuning of the RoBERTa-base model on the GLUE benchmark.
1000
750
Norm
500
LDAdam no-EF 1gr1
LDAdam |E.ll
20k
40k
60k
Update Steps
18

---

Published as a conference paper at ICLR 2025
C MATERIALS FOR REPRODUCIBILITY
C.1 OPTIMIZER STATES MEMORY ESTIMATES
Table 2 describes the architecture of the trained models architecture and the low-rank structure of the low-dimensional optimizer states, It allows optimizer states token count, thus memory estimates in GB by applying factor 2 for half precision training.
10243
Table 9: Weights and low-dimensional optimizer states shape for used models.
Weight
Low-rank optimizer states
RoBeRTa-base
Token Embedding
50265+768+768
Positional Embedding
768
Attention Head
12*4*768*768
12+4+3+768+r 12x3* 72+3072+
MLP Block
12+3+3072+768
Normalization
2+ 20766422
Dense Layer
768+7684768
Output Layer
108mmllabeeln.labee
Llama 130M
Embedding Layer
32000+768
Attention Head
12*4x768*768
1
MLP Block
12+3+2048+768
12*3* (2-+2444-447))
Layer Normalization
(12*2+1) +768
Output Layer
768+32000
Llama 350M
Embedding Layer
32000x1024
Attention Head
24x4x1024+1024
*44x4***1024*4
MLP Block
24+3+2736+1024
24+3+ (22377
Layer Normalization
(24+2+1)+102)
Output Layer
1024x32000
Llama-2 7B
Embedding Layer
32000x4096
Attention Head
32+4+4096+4096
32+4
MLP Block
32+3+4096+11008
32+3* (2+1100 (2-410024008
Layer Normalization
(32x2+1)x4096
Output Layer
4096x32000
C.2 FINE-TUNING HYPERPARAMETERS
Tables 10 and II detail all the hyperparameters when fine-tuning respectively RoBERTa-base model on the GLUE benchmark and Llama-2 7B model on the GSM8K dataset.
19

---

Published as a conference paper at ICLR 2025
Table 10: Hyperparameters used for fine-tuning RoBERTa-base model on the GLUE benchmark.
Adam
LDAdam
GaLore
Epochs
3
Warm-up
X
Batch Size
16
Maximum Length
128
Data Type
bfloat
Learning Rate
{le
5, 2e 5, ...,, 5e
5)
Learning Rate Scheduling
linear to 0%
Decay Rate 31
0.9
0.908
0.9
Decay Rate $2
0.999
0.99
0.999
Weight Decay
X
X
X
Dropout
X
X
X
Gradient Clipping
X
X
X
Interpolation Factor p
X
0.908
X
Error Feedback
X
/
X
Subspace Frequency
X
1
500
Learning Rate Scaling
X
X
X
Table 11: Hyperparameters used for fine-tuning Llama-2 7B model the GSM8K dataset.
Adam
LDAdam
GaLore
Epochs
3
Training Steps
702
Warm-up Steps
20
Batch Size
32
Maximum Length
512
Data Type
bfloat16
Learning Rate
{5e 5, 5, ..., De 5)
[le
4, 2e 4, ..... 5e
4)
Warm-up Scheduling
linear from 0%
Learning Rate Scheduling
linear to 0%
Decay Rate 5
0.9
0.908
0.9
Decay Rate 32
0.999
0.99
0.999
Weight Decay
X
X
X
Dropout
X
X
X
Gradient Clipping
1.0
X
X
Interpolation Factor P
X
0.908
X
Error Feedback
X
X
Subspace Frequency
X
1
200
Learning Rate Scaling
X
X
X
20

---

Published as a conference paper at ICLR 2025
C.3 PRE-TRAINING HYPERPARAMETERS
Table 12 details all the hyperparameters we use when pre-training Llama models on the C4 dataset.
Table 12: Hyperparameters used for pre-training Llama models on the C4 dataset.
LDAdam
Llama 1.3B
621 440 000
bf oatt6
5e linear from 0%
cosine to 10%
0.908
0.98
X
X
X
0.908
X
200000
131 072
10000
256
512
Adam
0.999
0.9
X
X
1.0
X
X
X
X
0.
GaLore
0.999
33
0.9
X
X
X
200
X
Llama 350M
208 960 000
(5e 4, le 3, Se.
linear from 0%
cosine 10%
bfloat15 1
LDAdam
55000
5500
256
512
171
0.908
0.98
0.908
X
X
X
X
Adam
0.999
0.9
X
1.0
X
GaLore
0.999
(5e 4, le - 3, 5e - 3}
0.9
X
X
X
X
X
200
X
Llama 130M
221 440 000
linear from 0%
cosine 10%
LDAdam
20000
2000
256
512
131 072
bfloatts
0.908
0.98
0. 0.908
X
X
X
/
X
Adam
0.999
0.9
X
1.0
X
X
Learning Rate Scheduling
Total Training Tokens
Warm-up Scheduling
Learning Rate Scaling
Interpolation Factor
Subspace Frequency
Maximum Length
Token Batch Size
Gradient Clipping
Warm-up Steps
Error Feedback
Training Steps
Learning Rate
Decay Rate 81
Decay Rate 822
Weight Decay
Batch Size
Data Type
Dropout
21

---

Published as a conference paper at ICLR 2025
D GALORE ALGORITHM
Algorithm 2 describes Galooe's (Zhao et al. 2024) proposal for performing an Adam-type update from a low-rank gradient projection. For experiments we use the author's implementation of the Ga- Lore algorithm accessible at: https: / com/ jiaweizzhao/ Ga Other than layer-wise learning rate rescaling, in our experiments we use the suggested hyperparameters.
Algorithm GaLore
Hyperparameters: step size nıs decay rates 81, .2
GaLore Hyperparameters: projection rank r; subspace change frequency T, scale factor a Initialization: my Up 0; Uu 0
2: for tort= = (1,2.... T do
3:
i mod T then
4:
Ur. E, V S (V)
5:
Ue = Uel: 1
6:
else
7:
U1 = U
8:
end if
9:
at u,
M4 Bim + (1- 31 )a
11:
Ur B2v2 + (1 - )o
12:
1c+1 = 0₂ aoTil vorte
mr
13: end for
D.1
LAYER-WIL WISE LEARNING RATE RESCALING
Note that the learning rate rescaling induced by multiplying by a # is equivalent to using different learning rates for different layers. This level of hyperparameter tuning is not usual, so for a fair comparison we use a single learning rate for all layers. However, our preliminary experiments suggest that LDAdam would also benefit from layer-wise learning rate tuning.
D.2 INCOMPATIBILITY OF GRADIENT ACCUMULATION AND PER-LAYER WEIGHT UPDATE
GaLore's improvement in memory efficiency over the standard Adam implementation comes from three distinct additions. namely: low-rank gradient projection (Zhao et al. 2024). 8-bit quantization (Dettmers et al.] 2022]. and per-layer weight update (Lv et al. (024b). b). The latter saves memory by releasing the variable used to store the gradient after the model layer has been updated rather than after the entire model has been updated. A practical implementa- tion of per-layer weight updates to add gradient hook that triggers the model update and releases the gradient variable immediately after the gradient for that particular layer has been computed. For example, in PyTorch, the per-layer weight update is built by adding a gra- dient hook using p. register peattterpost..ccee www..accumulltteeeccmommmmm grad hook (optimized hook), with the optimize minizer_hodl function implementing both optimizer. step () and opt optimizes meriizer.oo zero.grad [Pytorch.Tutorials) Tutorials n.d.).
Therefore, when using per-layer weight update, one never has access to the gradient for the entire model, and furthermore the model update (e.g. in PyTorch, run aptimiee zer step () only after accumulation. For the same reason, gradient clipping is not possible when using per-layer weight update, since the norm of the gradient for the entire model is not computable.
cannot accumulate the gradient over multiple miero batches and perform
To enable gradient accumulation for GaLore, the per-layer weight update has be abandoned. The method is called by the author GaLore (no retaining grad) (Zhao et al.) 2024 and it results in an additional memory overhead equal to the size of the model compared GaLore's claim.
22

---

Published as a conference paper at ICLR 2025
E DEFERRED PROOFS
Algorithm 3 LDAdam: Analytical View
1: Hyperparameters: step size T]t decay rates 33 and [3, projection rank r with contraction qr-
2: Initialization: error E Od, moments mg Up 0o O, and Uo Odxr-
3: for = 21. . do
Compute error corrected gradient and momentum
4:
A. - gt. + Lte where 9r - Va. (0:) is a mini-batch stochastic gradient at Oz + (1-11 A1
5: B, --144,,11,,
Update the projection matrix
6:
U, is any d X orthogonal matrix such that [| -U--4,1; B:ll S 9r Bell with gr < 1
Intermediate updates to adjust to the new low-dimensional space
7:
at U3 A,
8:
m/-1/2 = U, 1
9:
1/21/2 (1 11) 41-1)
(1 (1-5 - 1 )2) + (4 U-
me-1 )²
1-81
Adam updates in the low-dimensional space
10:
m -1m-- 1/2 +
B1)at
11:
U B24 1/2 (1 32
12:
ie max(1;, U2; | |l
AMSGrad-type normalization
Update the main model
13:
0.111 8 nelee Vee+e
Update the error feedback
14:
E++1 = (Ac - Ue at) + I- (We-1 m - Ue . 1421/22
91
15:
end for
E.1 KEY LEMMAS
Lemma 1. With Er = To² +
Σ
E
[ for any 2 the following bounds hold:
G
T
1
|B&|l
E
[B,11
(13)
9r
(1 - gr)
1=1
G
T
1&
5
Σ
E || 5 (1- B1 )2 yr
ET
(14)
(1 Bi)(1 - nr))
1=1
B
31 + (1- (1-Bi)ar G
B, + (1 31)
5 (1 3i)(1 -4r))
(15)
31
T
B1 B.- Bi)Es
1/ < (1 B1 B1) -2
B1+ + (1 31)4²
Σ
E
Σt.
(16)
1=1
Proof. Let us start with the proof of the first bound on Bt. Denote qr = (1 - B1 )r + 31 E [(rr.) 1).
|B++ill
M. (1 Ar)
= M, + (1 81) + (1 81)
|1-(1 Bi)M, + -B1)941 Bell
| B1))u22, B, +(1 +1-11) B1)) B.
= [|(1 - (B) -u24," B,) + (L-31)4:+1 + B,
< (1 - B;)1BB: - -4144 Bell + (1 3, 11 + B,1
< ((1 11-)))g---- + 81) B (1 31)119t111
= B. + (1 - 91111
6+1
++1
= G + (1- - 31) Σ
++1
1+1- = (1 - (1-
Σ
of
- [|
2=2
7=1
23

---

Published as a conference paper at ICLR 2025
Using the bounded gradient assumption, we get
(1
G
|B - (1- 3. of
dr
To derive the bound with expectation, we apply Cauchy-Schward inequality and the bounded vari- ance assumption:
T E
T
27
< (1- 31-
Σ
Σ o
T
111
1=1
Σ t 11g-11 }
(1 - (1-3)² Σ R
Σ Φー r=
(1- 31)²
9/- - E
[[lg~] 11a]
1 4
(1 - 3,)2
E [19.11
(1 ),)
T
Σ
E [19 Vf(0.) + V
(1 - )
T
E [V(()))))
(1 )r)
To²
E [[V)
qr))
(1
9r)
To bound the norm of the error ||E=]]. notice that T
||
1 1- 31
[B, u,, B.
9r
|B.ll <
G
(17)
1- B1
(1 B B1) gr)
To get the bound with expectation, we apply previous inequality on B, and get
T
T-1
Σ E []
E
[
gf E [[[BB]]]]
B1)²
To²
Σ
[IVS(())
31 - g2
The fifth bound 15) follows from the triangle inequality and combining the obtained two bounds. For the last bound with expectation, we have
E
31
B. + (1 81)6
31
E.
[
1 + (1 - 31) [1661(]]
B1
B,
[=1
81
To²
1
E
(1 (1-31) 1-1
(1- )2)
(1 9r)
+ (1 B1)(1 - )2
To² +
[[V
31 + (1 -1)
T
To² + Σ
E [IV
(1 - B1 (1) )3
t=1
24

---

Published as a conference paper at ICLR 2025
Lemma 2. 1 and y 2/1 B1)(1 gr), then
E
B1
B1-1 + (1 (1- BI)E
27
s 2C2
o²
+
Σ
" (18)
31
1 1-7
where C2
11 11-11110022 (1-8))(1-.))
Proof. From the condition on y and the notation q. = (1 -31)gr B1 E [9-, 1) from the previous
proof, we have 1-7-- -- < -- or equivalently y q. Q. Using this inequality on y and the
previous bound 16 for EE |. we have
27
Σ E [1]]
S (1- 31)² 2
T-
E
Σ or
119-11
T
Σ कー Σ -6 112 }
$ (1- 31)² 2
TT
(1- - 31)
T
T tal - [1|9t]] |+]
1 q
1=1==1
(1 - 31)²
ΣΣ
T T- T E [119-11]
4r
(1 (1- 31)² (1-
or
T-
[1|=~11]]
1- 4r
(1- T
(g./7) (T./>)--r0 --41
-T [ll9~ [|
1- qr
1 1 0r/7
(1 - 31)3
T
y T-E
E [ll9 12]
(1 9.) 41/1)
1=1
2(1 )
T-TE E [1lar| 2
(1 (
T-TE [1g VS(A.) + VF
(1- )r)
T
(0 + (
(1-
T
[[] f(0)) )
+
7
T- E
(1 9r)2
1 1-1
1=1
Using the bound
17b for 1. we have
T
2q
0²
T
Σ
&
2
+ Σ t=l
T
E [111()))11]]]
(1 - 11)
)
1-7
t=1
25

---

Published as a conference paper at ICLR 2025
Combining these two bounds get
T
31 B1
E
B + (1- B1)5 (1-5
31
Σ
T-t
E.
[1|
B
113
(1
B1)
[1 19]
(1
31)
1=1
2
0²
T
+
E
(1 )2 (1 - --)
1 1-1
2(1 - 31
0²
T
E
((0)))
-1- (1 9-)
t=1
2(31 + (1 - 31)
o²
T
E
[V
(1 - 3.) - 9.)2
t=1
Lemma 3. For AT, - De have
2
2
|AT
Σ |AT/| & Ve
Proof. From the definitions of I 19 and ū. max(Ue) 1/0e- (llmx) we have
Pr
Uia 1/2 E,, 16 + +e)l
U1Diaa
1/2 (/2(1l + +e)u
min + Unl
/Telllmin +
1
Uhe 11-111_1
V /10 11max +
v Vl +
Ut 1iaiDiag -1/2 mmmax + e)(E
3 ú4 Diag -1/2 (0x-1 + e, | < 0,-1111111 Tt-1;
which implies that Ar, P r is positive semidefinite. Hence, ||],|| = Amax Using the convexity of Amax over symmetric matrices, we get
Amax(Al) (AP,) 2 0.
|rril
Amax (Tz-1 -Tr)
T
Amax (Tz-1) + (-Tx) Σ
Xmaa (T+-1) (Tin(Te)
t=1
T
1 Immn +e
Vo +e
1
1
1
1 Ilmin +e
V
max +e
T Imxx te
v Tūr +
1
+
1 Ilmin te
Vi te
v V
mmax te
Vl | +
t=
1
+
Voollmin +e
/ Illrrmmmn te
V1lia te
VIBir + e
2
1 min +e
Ve
26

---

Published as a conference paper at ICLR 2025
For the second sum of squared norms, notice that for scalars b 2: 0. it holds that
(a b)² v b) 12.
Therefore, the ahove derivation be repeated without the square roots follows:
T Σ |TT 1=1
A (T
Amax (Tz-1) + ( =
Amax (T,-1) Amin (T)))²
(Amax (T (Tz-1))² (Amin (T)))²
1 -1 Im + E
Ue max +
1
+
| Immi +
10 mlmax + c
11 Ilmax +
1/ I + e
1
| Immin
Uh min
18--1 maxe +
@e max
1=
+ |lin te
00 I mia te
1 +e
1 mma te
2
00 min +
which completes the proof.
Lemma 4. For all iterates 2 1 the following bound holds
1 Bg (1 31 9,))2
| mmxx
G22
1 - B2 (1 - 31)a (1 qr))
Proof. First, let us bound the a, term using Lemma D and Assumption
1 (1 9.)81 -= CG.
||all| = 114, A s ||ul S |& + ||
s (1 - 31)(1 - ạr)
Next, we bound momentum m.:
||mll s B B1]]m;-1/211 + (1 [axll
< B + - CG
t-1
< Mmoml + (1- - CG
3 = (1- -1- Σ
r=0
Next, we bound the intermediate term ||0,1/2211- Note that by the triangle inequality we have the following direct bound for
- - (44Tu--1) 2 + (1=m)
>"1 + +1u4/ 1- 81
Mi-1
2
27

---

Published as a conference paper at ICLR 2025
Now let us bound each term separately. For the first term have
(1,)) 22222222111
Σ
Ut 1-
max 1815
j=1
= ||
||
max
|,, Uc- |::1]]
1555"
max 15JS"
|14 |: = 11
since columns Up(;, and (4-1(11)) have unit length by construction. Similarly, for the second term
we have
(u4 (uTT)-1))
( m4-1
2
|
14-1, 2 u1 11,,41111133))
1- 87-1
31-1
m 1.j
1 81-1
max 1S 1=1
j=l
M4-1
14 14, U1
[3 [311
m4-1
C
1 8
1515P
1
81-
8
Finally, for the third term we have
U1 (1 m - 8
m4-1 U, U1-1
mc-1
C
5
81-
8
Combining all three bounds together, we arrive
|//-1/2211 S || + G2.
From this get the bound for e using the initialization 1/p 0:
|eellmax S ||vell
32 12/11 + (1-13)) |0
Pb(1l==-11+ + (1 - B2) C2
S Bg (1
t-1
1+ 32
S 3 polla +
02 81) 62
Σ 3 r=0
C²
1 - 32
Hence, using the update rule of i, and initialization
dgg = 0, we conclude
1 +
82 c²
1 32 (1 (1 qr )112
| |lax S Inax[1 maxe | llmxx)
G2.
1 32
1- 32 (1 - 31) - gr)
E.2 NON-CONVEX ANALYSIS
Theorem (Non-convex convergence rate). Let Assumptions 22 and 3
Bhold. Then, choosing step-
size 1) = min J with TTo "I..VVIT5.' Algorithm satisfis
f(0,)-F* Lo²
V(())111 & 2Co
T
VT
eVT
f(01) -f"
L'C2
(1+C VET
(1 (1+22 2C
+4C
TnoT
2T
3/3/2
with constants Co- = V
14.8 (1 (1-8,(1-0.))
G te, C =
A (1-A)((14.)
C2 0= 11+(1-814
1-3 (1- (1-81)(110)
(1-8))(119.))
28

---

Published as a conference paper at ICLR 2025
Remark 1. Further ignoring absolute constants, the bound becomes
Co (100))
Lo²
Co
( [(01)--- C2CoC202 GG CRIGA
+ T
VE +
VT
%
2
cVT
Let us assume that for the low-rank compression of the algorithm we have 4r = (2)) (e.g. SVD option). Then Co O C - C2 = ( and 1 -0 th G). Plugging this
To
asymptotic expressions into the bound and ignoring other parameters (e.g. o² L, E), we get
d G
dG
d2G
&
dG²
d²G
+
+
V
TT
T2
33
r
22VT
or equivalently
d G
d
4
G2
d
2
G
d
G
+
+
r V
T
7
T3/2¹
Proof. Let I', := 1,Diaag-1// 1/2 (ū, +e)U, be a preconditioning matrix and M, := U2m, the exponen- tial moving averages in the full space. With these notations, the update rule of the model becomes 0(+1 = 0₂ - mTT, Mr. As it will be used later, we need to make sure that our preconditioning is positive definite. Due to the structure of rr, it is positive semi-definite only. To make positive definite, notice that rill, Tilk, where the full-rank preconditioner I' takes the form
(19) where Ug E Rdddd is an orthogonal matrix with the same first r columns as U4 E Rd and the diag- onal matrix in the middle has been extended to meet the sizes of L7 by adding the values || |lii + on the diagonal. ryl4;- = 1224; comes from the observation that 21, 14, E Raxr is composed of two blocks: upper r block of the identity matrix and (d r) r block of the zero matrix. Hence, the added (d r) elements on the diagonal do not really affect, so does the last (d - r) columns of Uh Therefore, we can write the model update rule as
r, = U,,iia 1/2
(ē; 102 e, (| Ilmin + +,22,,"
Oc+1 O2- me Mc,
(20)
with full-rank preconditioning Ir, Next, note that
B, M + (1 - B1)g. + (1- - B1)ee,
(21)
m = U, M,-1-
(22)
Then, for the low dimensional momentum and the erorr we get
my = 1/22-1/2 + (1 - --1)a = 4, (31 + (1 31)gc + (1 - --))))))==-- = U, B,,
(23)
(1 - -:81) = (I - U12 )B. B. Mc.
(24)
Letting Bo = 0, we define virtual iterates xe as follows:
B1
T++1 = 0++1- - n/L'r (1 - 11-5 + 1- 31
B₁
(25)
In particular, x1 01. Then, we derive the recurrence relation for the new sequence follows:
$1
Ø(+1 = Pct1 - nTI'' (1- +
-1 B,
24
02 - nlI'M, - n1l's Br 38M +
3
B.
1 Bi
11 TeBe
0.
1- 83
21
0:
To
(3 + (1 (111-111) B1)g + - -Bi)Ea)
1 B1
24
0, rT,
B1
(B,-1 -11- (1--8186 B1)&.) + +Ee) nnTT::
I
Bi
B
= 0, - nTI',
B-- + (1 - (1-81)): nnT'ig
Bi
25
84 TTT + A
(-31 (T- B1
B1
B, + (1 - -B1)E(
29

---

Published as a conference paper at ICLR 2025
where AT, 11-1- Te.
Next we apply smoothness (Assumption gradient Lipschitzness we have
I of the loss function f over the iterates 24. From the
f(1+11) S f(xx) + (Vf(z)) - zτ) 11 -
Taking expectation, we obtain
[[f(z(+1)] - B[f(z)] <
BB (VV((())
K
11 Bi-1 (1 -115-))))) }
31
+nB
ff(z)), Ar,
2
| AA,
51
1
+
B1-1 (1
2
1- B1
-n
KV
(26)
{ (f(a)) AI'
( 31
Bi-1 + (1- Bi)EE) )
(1-
+ nB
1-
n +
Ti9e Ar, (1
B1
E
B + (1 (1- (1-11)
2
(I- B1
111
+ B (V) - Vf(x)), r.9.)
(27)
In the following, we bound all the four terms mentioned above.
Bounding term I. Let
||al|| be the operator norm (with respect to Ea norm) of the matrix AT. We
have
I EE (1771100..TT11)--27V777100).
< 1
< -mmi Thhn
(28)
Co
where we use Assumption 2 and Lemma to bound
Bg (1- B1 311--111-1 9,))
-1/2
Amin (T--1) 2 (1lxx_1llmax + e) -)/
1/2
G2
+
B2 (1 (1) g,)
Co
Note that the purpose of making I' matrix positive definite is have negative term in Bounding term 11. Then we have
28
Vf(O.). Ar, (1 B1
B + (1 (1- )
II < nE
1- 31
(vr(o)) vf(()), (1- 31 Bi-1 + (1 Si)&s (1- )
B1
+ nB
[1T5(0)11
Ar'; (1-)), Bi- + (1 - 1-1))k )1
31
< 7)
31
+ n LE
Tr
B1
B + (1 - 3,)5. (1-81)s || Ar, ( 3n 31
(1-B, B + (1-
(1--81) --:1)/
||
31
1/ CCC2.G
,C,
E AT ||]
(29)
Ve
30

---

Published as a conference paper at ICLR 2025
where C₁
B1+(1 (1--1111))1 )ar and we used the fact that the largest eigenvalue (Tmx ||T,|| (11-8)1---- 9.)
(Ilūx]lmi, + < -1/2 The second inequality is because of smoothness of (()), and the last inequality is due to Lemma Assumption 2 and the property of norms.
Bounding term II This term can be bounded as follows:
AT, (1 (1- 31
Br- + +1 /
111 S LE [|Tege]]]]
+ LE
L Blllg () V/(0))) 1T | (1 (I- 31
Bi- + - ---))))))) )1²
L [E][ + +02) o/ Ar
1 E[I
Lo +1 02 LG EI
(30)
where we used Assumption that 9r is unbiased with bounded variance a².
Bounding term IV. We have
IV - - (1/700/))1/////TT-T-0))1111111/1 Vf((z),, Arr
||
31
B + (1- |
< [VE[()
1
B1
C T.Q
1
+ 2 BIIV f(0.) vf(x))]1]
2c
2p
ve
(
ne
L2 E
81
C
[E[| 1000)
I'
Br-1 +(1 (1- B1)&,
2c
2p
31
Ve
<
n [[IV/(0))1
L2 E
B1
B + (1-11)E
|
22
1 LC
2c
2pc
31
(31)
where (a) due Young's inequality and (b) based on Assumption) Now integrating
(28
29
30
31 into 277
I S - 1
II < CCGGBTIAATT| + G2 B||
Ve
1PL ( + L2
III &
+
IV < TPP
L2
31
27
LC
E
B- (1 -0-880 -:)6
+
E
2
2,e
1- 31
Ve
31

---

Published as a conference paper at ICLR 2025
and taking the telescoping summation over t 1, 1,..., T, obtain (xT+))+1) - f(x1)]
+ r2L
+
90
B [(
1 T
B
B1
B
(1
-
B1)6;
CD
2c
2pc
$1
1=1
1=1
(1 +
(11+ Ci)G2
BII + O
Σ
B BIIAATII7
Ve
16
11
nP
NTT222
TTpLo?
TTP
+
+
Co
2
2pe
e
2pc
(((++)) ((
CI)C
T B[ + CF Σ
VC
where we used
(16 of Lemma with constant C2 (1-8))
Choosing
2Co and
1l0 == ILCoVI+C, and using Lemma 3 we get
Top T
Blf(oT+1) f(ai)]
7
Σ
B(I
2C, 1=1
6²
2n (((++2C)CCLC* 2n + C1)
Ve
Re-arranging terms, we get that
T
((01)-
'FI'"
[())) & 2Cp
Tn
2
((1+C))G2 TVe
+ n1(1 + 221+C)CC,LG'
+4Cp
Te
where in the last inequality used xı 02 and the lower bound f S f(0) for all 0 e R.. Finally, choosing = min(70. VT and considering the two cases, we arrive at the following rate
T
1 11,-T
1(0)) -" +
Lo²
L²CoC2o2
T
E ((())11
& 2Co max 1,
TOVT
VT
EVT
+
2T
1))
+ (1 2C2)
+4C
VeT
3/3/2
f(01) -f fr
Lo²
< 2C.
VT
eVT
(f(0))
- f
C20
(1 (1+
(1 (1+,0))C,LC
+4C
2npoT''
222T
Ver
TT3/2
which completes the proof of the theorem.
E.3 ANALYSIS UNDER PL CONDITION
As in the non-convex analysis, here we derive the convergence rate with fixed step-size n.
Theorem 4. (Convergence rate under PL) Let Assumptions
7
3
and hold. Then, choosing
step-size = min(70.
2Co log T with 170 min( I6LC.
Co(1
9c)
GLVCoCz ), Algorithm
3/4
uT
2n
satisfies
([ - logT T
12LC302 1.2c
6Co(1 +1(())))))))) G9G9)
6* T
+ +
T
uve
32

---

Published as a conference paper at ICLR 2025
Proof. We start from descent lemma
[[f(=++1)] B[f(a+))
n KV (axe), ((ge)
31 (1
n
( Ar,
Ar,
BL-1 + (1-11151) -15x))
1-B
17² L
( 11
||
Tig -rr,
Bc-1 + (1
(32)
1 31
111
We bound part /1 and part II1 in the same way as done in the non-convex analysis. We now provide bound for part I':
r' T (VV J (3r2), ( r99h
1/ LVV (x+);
Al'g:))
= =-yjE
-,)E ((V]( AFm)))
We further expand and bound this equation as follows:
1' 1' - C [IV
(V](())) VS(0.) + (--0),,r-1(((()) vs(x,)))
w/E -VE[())((>) VS(0,) + A
Co 21 [1V[(=))1
(V](x)) -/)
-V (V/(0.) -V/())))))((
7E (T]()))))
S 17
-E E [V[f(=)))*
+ 7 E
C.
Ve
7E -V fz()))))
+ DE
Ar + EB (VV AT.
To bound the second and third terms above we reuse derivation done in 31 to have
B (V (V v((aa)))
NP [
(00))
(33)
2c
1 L2
31
27
+ 2pe
E
B, + (1- 31)E
1-31 3r
TT2
B1
E [Vf(z)) [(0))1]]
E
B + +1 8.)&
(34)
1-
Next, we use the Cauchy-Schwartz inequality to bound inner products above, L-smoothness in- nLC,G and the inequality -||a||2 < Ve
equality to bound [/,f(=;) f(()))l < L| Oul
33

---

Published as a conference paper at ICLR 2025
- -b1l + ||a bll for the first term:
I' - "E [
"
VV
f(00)1
Co
V
"P [[[[IV [()))) + L² E.
||
81 Bi-1 + (1 - B1)E:
2c
2pc
1- B1
+ ny V[ - vf(e))
[1] + A
2C, B [|]
" E [V
2Co
+ 7P
[ f(O:))))
+
L2 2pe
18
3
B (1 - B:)5:
27
2E
VE
B
LCIG
+ r
R.
||| + G ||AT,||]
Vc
S
"1
[
- 7
El| V [()0))
"
B[IV]((x+) - v(((0))1
2C
4Co
2C
AEE [V s(0:)1
+ 12²1,2
E
31
B (1 - BI)E:
2E
2pe
1- B1
L2
||
G AT
Ve
E
[IV
" E
2Cu
4Co
7P
+ 2VE
37
[V](x;) - ((0))1
L'L2 E
81
B + (1 81)5
-1-
2c
2pc
B1
1 G
BMArF +
Ve
E [[V((2)))
- "
EII +
7P
E1IV ((.)
2C,
4Co
2e
3n L2
31
1
+ L2
B1
+
E
Br + (1 - (1-)))
E
B + (1- 11--B1). B:)E
2/e
Bi
2pe
31
+ 1 GC
BIl + G EIGAII
Ve
< "
E
[[+((1)
77
7P ((0))))+
B ((0))11
2Ch
4Co
2c
122
3
1
B1
LF
E
B + (1-
G [|]T.]]]
2c
Ve
p
1 31
Ve
Plugging the obtained bound for I' with previously obtained bounds for II and III
11 S TC T2RPLC2
Ve
III S I "1 +
Lo2 + LG
34

---

Published as a conference paper at ICLR 2025
into 32 choosing p =- and using the step-size bound S 1610% get 8Co
[[((+++))] [[(((a)) -B[[)))))
"
E
[
11
"P
F V [(.)]
20.
4Co
2c
1
3
1
B1
B + (1 8:)6 1
L2
+
E
E [[
2E
VE
1-31
31
Ve
C2
AT.Il]
B
Ar
||TA||
Ve
LO
E ((
n
V f(0.))1 L
S
17
E [1V
17
BIll
2Cp
8Co
1p²
3
31
IF
B- (1- BI)
2c
Ve
8Co
1- B1
71 +
C2)
+n(1 +G) ||, +
B| + +
Ve
" [[I 1
S
111
(Bif(z))]
-
C.
8C,
B1
E
B,-1 + (1 3:)5
6/2
81
11 + BIG + 7 + 0,)
BM + EILLEEI
Ve
where the last inequality we applied PL condition from Assumption 4 the terms, we obtain the following recursion:
After some reshuffling of
7 Elll V f(0.) L0
np Co
fr
1
(B[f(x+)]
8Co
277² L2
31 B + (1 - (1
E
3/2
+ (1 + G + T2(1 + C.)
+
Ve
Now we unroll the obtained recursion and invoke Lemma 2 with 1- 1/ From 17 S ALCo 4u conclude that 7 ne
Co
Co
1). To satisfy the condition on 7 of Lemma 1 we enforce the 1).
35

---

Published as a conference paper at ICLR 2025
bound 17 2μ
Ca (1 - B1)(1 qr)) on the step size. Therefore,
T
T-t
B[f(xr+1))
-
A/4
(((1))-F)
7
Σ
1
7/1 C.
E
B
Co
8C t=1
L T-1
mp
Co
2n L²
T
T-t
31
E
B + (1 Bi)
3/2
Co.
B1
T
T-1
+ m1 [1++
Σ
1
EIA
Co
+(++G)GLG
T
T-
Σ
1
VC
Co
T-t
P
Σ
1
mp
||
(35)
Co
For the second sum above we upper bound it by its infinite sum as
T-1
m
Σ
7//
Co
C.
Co
np
t=
1=
For the third sum, we apply Lemma 2 and for the other three sums we bound the bounds in Lemma 3
1
T/ Co
< 1 and apply
T EIAAT,I1 S
Σ
Σ BMA S
Ve
T
Σ
EIJ
ET 11 < Σ
B|JAT,|115 2
36

---

Published as a conference paper at ICLR 2025
cancelling terms involving gradients with step size restriction 7 - BL./CoO2 and noticing that #1 01, we finally get
Plugging all this bounds into (35
3/4
T
T
T-t
B[((2++1] -
n/ (s(0)) - 8C0
1
np
B
Co
Co
1p Lo Co
mn
2 6/22 20
(Goz
T
7/11
T-1
+
1
E B
Co
np
1=1
+ +1 + C2)G2
2
2
+ +n C2
Ve
Ve
Ve
(1 1 Co
nH T
(f(0))
L|Coo pe
CrC20
163/2
2n(1 C1) + + 2
Ve
nH T
(f(1)) - f*)
Co
+ 2(1 + -.)C2)) + 1/ LCgo2
(ALPCoCr0² 2(1 + 2C1) 3/3/2
+
Ve
uee
The obtained rate above is with respect to the virtual iterates I, that we defined for the purposes of analysis. To convert this rate with respect to the iterates O₂ of the algorithm, we apply L-smoothness to bound the functional difference:
(s(x)) - ((0)) < (VS(0)), x. 0.)) + 2 1,, - 0,12 G22
Ve + LCCG2
20
which implies
m T
- -
(s(O)) - -5)
Co
LC
3(1 + G1)2
(AL2COC202
5(1 + L
+ 17
+
+ 1??
uee
Ve
pe3/2
Plugging 77 = min(70.
(oolog?" Wf
with n7u = min( 16LC%
Co(1-BiX(1
2/
6LL
3/4
we get
1
[[[ - S
max
T2
1
Co C--
(1(0)) 1*)
log T 2Co
[LCno2 3(1 3(+
T 1
μEE
Ve
log T 4C8
ALLCOC2202
+ + CG
T2 112
He3/2
log T (2T.C802
6Co(1 + G1)
+
+ T
T
112c
uve
which completes the proof.
37


Reviews:
[{"review_id": "PlYjeumjD2", "reviewer": "Reviewer_hwuW", "summary": "The paper proposes a new optimization algorithm to train large (language) models.\nLike the recent GaLore algorithm, the proposed algorithm is a variation of ADAM, that performs updates in a lower-dimensional space, to store lower-dimensional quantities to be more memory-efficient.\nAs I understand (and as written by the authors), the real meat behind this kind of algorithm is the choice of the projection to obtain a meaningful lower-dimensional subspace that is both efficient to compute/store and a good enough approximation of the gradient. The proposed algorithm retains the desirable memory-efficient properties of GaLore, while almost matching the generalization performance of ADAM. Convergence proofs and exhaustive experiments on text benchmarks are provided.", "questions": "see Weaknesses", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is overall easy to follow, the idea is interesting, and seems to empirically works great.", "weaknesses": "I see that the authors provide some running times and memory metrics in Appendix B, but it seems to me that this part should be more exhaustive in order to prove the proposed algorithm useful.\nI know this is hard to monitor for algorithms on GPU, but could authors provide graphs similar to Figure 1, but with running times? and memory usage? with ADAM, GaLore, and LDADAM (proposed).\nIn other words, could authors provide graphs with perplexity as a function of runtime and perplexity as a function of memory usage? Even if these required graphs will be imperfect for a lot reasons, I think they are crucial to gain insights and potentially prove usefulness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qy7ADEaF3w", "reviewer": "Reviewer_HAQR", "summary": "The paper proposes a new memory-efficient training algorithm based on Adam via efficient projection onto subspaces. The proposed method enjoys non-asymptotic theoretical rate of convergence and demonstrate good experimental results in training large language models (LLMs).", "questions": "Overall I think this is a work with interesting observations and methods. I have the following concerns and questions (I'm happy to increase my evaluation if the authors address the issues):\n\n1. To my understanding, GaLore doesn't employ SVD step for every iteration in order to save computation overheads. On the contrary, the proposed method would need a Gram-Schmidt process and several extra steps for updating the momentums and gradients. A major conern is the computation overhead of this proposed method. I udnerstand that the memory consumption is roughly the same as GaLore, yet I'm curious about the time required for training LLMs both in theory and in all the experiment settings using the proposed method.\n2. The convergence results in Theorem 1 and 2 are similar to that of full-rank Adagrad/Adam. To my understanding, the low rank parameter $r$ only plays a role in the constant $q_r$? If so, can the author comment on how to choose such $r$ and if the theory has any implication on the choice of $r$. In particular, if $r=m$ and the method reduce back to the original Adam, can we say any matching result with the analysis on the original Adam?\n3. This question is related to the previous one. The experiments show that LDAdam is even better than original Adam in some cases (Table 3 with $r=512$, Table 4 with $r=256$). Can the authors expand on the discussion of \"regularization effect of compression\" or provide some references? It's still not very clear to me right now.\n4. Any experiment on even larger LLMs such as Llama 1b or 7b? I'm curious what will happen if the model is large and the training is conducted in parallel on multiple GPUs.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and the contribution is clear. The algorithm is novel with theoretical convergence. The experiments indicates that the new approach is performing not only better than GaLore, but also better than the vanilla Adam for some parameter settings.", "weaknesses": "(Please respond to the \"Questions\" directly) The computational time of proposed method is not clear; Some discussion on the theoretical results is needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7HDC03PCwL", "reviewer": "Reviewer_2kSc", "summary": "This paper on LDAdam addresses a crucial challenge in training large models: reducing memory consumption of optimizers without sacrificing model quality. The key innovation is the introduction of a lower-dimensional subspace optimization approach, where adaptive optimization occurs within constrained dimensions. The optimizer, though operating in a subspace, ensures exploration across the entire parameter space, thereby maintaining model expressivity while being memory efficient.", "questions": "The original GaLore paper supports gradient accumulation, contrary to the authors' claim. I suggest they clarify this point, either correcting the claim or explaining any differences in their interpretation or implementation of GaLore to avoid potential misunderstandings.\n\nGood paper, I suggest acceptance.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Projection-Aware Update Rule**: LDAdam uses a projection-aware update rule to transition between subspaces. This rule allows it to estimate the gradient statistics even after dimensional reduction, adapting efficiently to changes in the subspace basis without losing essential gradient information.\n\n2. **Generalized Error Feedback Mechanism**: To address inaccuracies from low-rank projections, LDAdam introduces a unique error feedback mechanism that accounts for both gradient and optimizer state compression. This mechanism reintroduces projection errors in subsequent iterations, ensuring stability in updates and helping LDAdam match the performance of uncompressed optimizers.\n\n3. **Convergence and Memory Efficiency**: The paper provides a proof of LDAdam’s convergence under standard assumptions. With this design, LDAdam achieves a significantly reduced memory footprint. In practical terms, LDAdam’s optimizer states require only a fraction of the memory needed by Adam, making it feasible to train large language models in resource-constrained environments.", "weaknesses": "Liang et al. [1] provide a convergence analysis for GaLore and related algorithms without relying on a 'stable-rank' assumption. Since LDAdam aligns with this framework, I recommend citing [1] for its convergence insights. Additionally, a comparison of LDAdam’s convergence analysis with [1], specifically on the impact of removing the 'stable-rank' assumption, would help clarify the authors’ theoretical contributions.\n\n[1] Memory-Efficient LLM Training with Online Subspace Descent, Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NmdJHcGONn", "reviewer": "Reviewer_tZ3Y", "summary": "Thin paper introduces a novel memory-efficient low-dimensional Adam optimizer for training large models. Specifically, the authors propose a new projection-aware update for the optimizer states that allows for transitioning between subspaces. They also perform block power iteration to at each step to efficiently and accurately project the gradient and optimizer states. To reduce the error caused by the low-rank projection, this work leverages a new generalized error feedback mechanism that accounts for both gradient and optimizer state compression. The authors theoretically analyze the convergence rates for different objectives. To validate the proposed algorithm, the authors present extensive experimental results on different datasets with the comparison to different baselines.", "questions": "1. How to relax the bounded gradient assumption in the study?\n2. How to guarantee a feasible $\\mathcal{U}_t$ if Gram-Schmidt process fails?\n3. How to validate the theoretical conclusions from both theorems?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The investigated topic is quite interesting and critical, particularly when LLMs are popular in every area. A memory-efficient optimizer is required to conduct computationally efficient fine-tuning.\n2. The paper is easy to follow and the presentation in this work is clear.\n3. This work presents in-depth theoretical analysis for the proposed algorithm and show the explicits the convergence rate.\n4. The empirical evidence to validate the proposed algorithm look extensive and thorough.\n5. The comparison between the proposed method and baselines is technically convincing.", "weaknesses": "1. Assumption 2 is a quite strong assumption for the gradient, which may limit the applicability of the proposed algorithm. Until now, most existing algorithms have relaxed this assumption in their works such that this assumption imposes some weaknesses to the proposed algorithm.\n2. In this work, the authors have mentioned that GaLore is the closest work. Though we have seen the authors have shown in detail how the proposed algorithm differs, it would be great to include a technical discussion for the difference between them in terms of convergence rate. That way, it is clearer to directly show the exact theoretical difference.\n3. In the practical view of the algorithm, the authors use Gram-Schmidt process to calculate $\\mathcal{U}_t$. However, the process has numerical instability, computational intensity, and even inaccuracies. If this process fails, how can the authors guarantee that there always exists a feasible and reliable $\\mathcal{U}_t$ for the subsequent steps?\n4. Although the experimental results look promising, how to validate the theoretical conclusions from Theorem 1 and Theorem 2, particularly, the impact of the compression ratio? The authors need to present some ablation studies for the proposed algorithm.\n\n\n**********************Post-rebuttal Comment***************************\nAfter carefully reading the rebuttal from the authors and taking other reviewers' comments into consideration, I think the authors have addressed my major concerns. The new results also look interesting to me. Thanks for the efforts. I will also raise my score accordingly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)