AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning
Authors: Yehonathan Refael
Jonathan Svirsky
Boris Shustin
Wasim Huleihel
Ofir Lindenbaum
Venue: ICLR 2025 Poster (2025)
Keywords: low rank adaptation; low rank gradient training; memory efficient fine tuning; memory optimization; adaptive rank; foundation large language models;
Primary area: foundation or frontier models, including LLMs

Abstract: Training and fine-tuning large language models (LLMs) come with challenges related to memory and computational requirements due to the increasing size of the model weights and the optimizer states. To tackle these challenges, various techniques have been developed, such as low-rank adaptation (LoRA), which involves introducing a parallel trainable low-rank matrix to the fixed pre-trained weights at each layer. However, these methods often fall short compared to the full-rank weight training approach, as they restrict the parameter search to a low-rank subspace. This limitation can disrupt training dynamics and may require a full-rank warm start to mitigate the impact. 
In this paper, we introduce a new method inspired by a phenomenon we formally prove: as training progresses, the rank of the estimated layer gradients gradually decreases and asymptotically approaches rank one. Leveraging this, our approach involves adaptively reducing the rank of the gradients during Adam optimization steps, using an efficient online-updating low-rank projections rule. We further present a randomized-svd scheme for efficiently finding the projection matrix. 
Our technique enables full-parameter fine-tuning with adaptive low-rank gradient updates, significantly reducing overall memory requirements during training compared to state-of-the-art methods while improving model performance in both pretraining and fine-tuning. Finally, we provide a convergence analysis of our method and demonstrate its merits for training and fine-tuning language and biological foundation models.

Published as a conference paper at ICLR 2025
ADARANKGRAD: ADAPTIVE GRADIENT RANK AND MOMENTS FOR -FFICIEET LLMs TRAINING AND FINE-TUNING
Yehonathan Refael
Jonathan Svirsky²
Boris Shustin
Wasim Huleihel'
Ofir Lindenbaum
Tel Aviv University BBar-Ilan University
University of Oxford
ABSTRACT
Training and fine-tuning large language models (LLMs) come with challenges re- lated to memory and computational requirements due to the increasing size of the model weights and the optimizer states. Various techniques have been developed to tackle these challenges, such as low-rank adaptation (LoRA), which involves introducing parallel trainable low-rank matrix to the fixed pre-trained weights at each layer. However, these methods often fall short compared to the full-rank weight training approach, they restrict the parameter search space. This limitation can disrupt training dynamics and require a full-rank warm start to mitigate the impact. In this paper, introduce new method inspired by a phenomenon we formally prove: as training progresses, the rank of the es- timated layer gradients gradually decreases, and asymptotically approaches rank one. Leveraging this, our approach involves adaptively reducing the rank of the gradients during Adam optimization steps, using an efficient -upliiin low- rank projections rule. We further present a randomized SVD scheme for efficiently finding the projection matrix. Our technique enables full-parameter fine-tuning with adaptive low-rank gradient updates, significantly reducing overall memory requirements during training compared to state-of- thet---rt methods while improv- ing model performance in both pretraining and fine-tuning. Finally, we provide a convergence analysis of our method and demonstrate its merits for training and fine-tuning language and biological foundation models. The code is available GitHuu
low-rank sub-
INTRODUCTION
Large language models (LLMs) have gained significant attention due to their impressive ability to handle various tasks, such dialogue-hased systems and text completion. Both supervised fine- tuning and additional pre-training can further enhance their performance across tasks and domains. However, training these models presents significant computational and memory challenges. This is because performing the gradient updates requires storing billions of LLM's trainable parameters along with the optimizer state (e.g., gradients and moments). In Adam, for example, the gradients and the estimated first and second moments triple the size of the model itself (Xu et al. 2024 Brown al. 2022 Kim et al. 2023).
To tackle the challenges associated with LLM fine-tuning, researchers have developed various op- timization techniques reduce memory usage during model training. One key approach that has emerged is Parameter-efficien fine-tuning (PEFT) (Han et al.) 2024). which enables the adapta- tion of pre-trained language models (PLMs) to different tasks without the need to fine-tune all model parameters. A prominent method within PEFT is the Low-Rank Adaptation (LoRA) al- gorithm, introduced by Hu et al. (2021). LoRA reparameterizes a weight matrix W E R into W Wo + BA, where Wo is a frozen full-rank matrix, and B x and A I="x" low-rank adaptors. Since r < min (m, n), the low-rank adaptors and B require fewer trainable parameters, reducing memory usage. LoRA has been widely adopted for fine-tuning, with several

---

Published as a conference paper at ICLR 2025
variants emerging, including Adaptive LoRA, which adapts the rank of the matrices during training Wang et al. 2023). LoRA+, which uses different learning rates for the two matrices (Chen et al. 2023), and Sparse LoRA, which introduces sparsity to the matrices to further reduce computational cost Xu et al.
2023). These methods have been demonstrated to enhance the efficiency and perfor-
mance of LLM fine-tuning for various tasks.
Despite its advantages, recent research has identified some limitations of low-rank reparameteri- zation. For example, LoRA may not achieve the same performance levels as full-rank fine-tuning (Meng et al., 2024) and might require initial full-rank model training as a warm-up before effec- tively utilizing the low-rank subspace (Lialin al., 2023b). These issues may stem from the fact that optimal weight matrices are not inherently low-rank or from changes in gradient training dy- namies introduced by the reparameterivation. base model static without dynamic adjustments. Another approach by He et al. (2022) dynamically adjusts tuning parameters during training. and (Zhang et al. 2023 Svirsky et al.] 2024) gradually reduces tuning parameters.
addition, LoRA keeps the tuning layer shapes the
Until recently, the pre-training of large language models (LLMs) has been primarily limited to cor- porations and governments with substantial computational and memory resources. The significant challenges posed by the enormous memory and computational requirements made it impractical for the average home user. To illustrate this challenge, take the following an example. Training Mistral 7B model from scratch poses substantial memory challenges. Given its 7 billion parameters, single update step requires approximately 70 GB of memory: 14 GB for the model parameters, 42 GB for Adam optimizer states and gradients, and 14 GB for activations. Consequently, consumer- level GPUs like the NVIDIA RTX 3090, which has 24 GB of VRAM, are inadequate for handling such a large-scale training task.
To overcome this challenge, the study in (Zhao et al. 2024a) introduced a training strategy called GaLore that enables full-parameter learning while being more memory-efficiee than traditional low-rank adaptation methods such as LoRA. The core idea behind GaLore to exploit the slowly changing low-rank structure of the gradient G E R=*m of the weight matrix W, rather than approx- imating the weight matrix itself as low-rank. GaLore significantly improved memory efficiency, re- ducing optimizer state memory usage by up to 65.5%. The following noticeable variant is Q-GaLore (Dttmerss et al. 2023), which combines low-rank gradient projection with INT4 quantization to fur- ther reduce memory usage, and additional parallel variant would be ReLoRA (Lialin et al. 2023b) employed in pre-training-by by periodically updating Wa using previously learned low-rank adaptors. We found that Galore to be suboptimal since it arbitrarily requires pre-defining fixed low-rank size for the gradient projection/llo--rnk---oo waw-rann-approximatonn while gradients rank gradually diminish during training down rank one. Additionally, GaLore uses fixed window size for the number of itera- tions between updates to the subspace onto which the gradients are projected, keeping this window size constant. Finally, Galore does not transform (adjust) the first and second moments any up- date of projection subspace, which we empirically found degrading the potential performance. We suggest an inner transformation scheme of the moments at any projection updates.
Our approach and theoretical results. In this paper, we introduce new training method aimed at optimizing memory efficiency in the training or fine-tuning of large language models (LLMs) while also improving convergence rates and overall performance. Our method leverages two key properties of LLMs. First, we present a novel theoretical finding that shows how the approximate rank of the LLM gradient matrices decreases progressively throughout the training process (even under basic SGD settings), asymptotically approaching rank one. Note that previous studies have only demonstrated an implicit upper bound on the rank of the gradient, which is far from being tight. For example, Zhao et al. (2024a) showed that the rank of the gradient satisfies rank (G"xm) <
min (n, m2/2. Second, as highlighted in previous research (Gromov et al. 2024 Refael et al., 2024 Jaiswal et al., (2024a), the depth of a layer (how far from input/output) and its architectural design contribute differently to the model's performance. Specifically, when perturbations from the same distribution are applied across various layers, the impact accuracy varies significantly. This indicates that the optimization steps have less influence on the model's s performance for certain layers, depending on their depth and architecture type. Noise in these layers has smaller impact on the overall task, meaning that the gradients in these layers carry less important information. This results in naturally lower-rank update steps (gradients).
2

---

Published as a conference paper at ICLR 2025
Building upon these two insights and to address the limitations of the LoRA variants, propose a method that enables full-parameter learning while dramatically reducing mem- ory requirements and computational complexity through adaptive low-rank gradient projec- tions during the Adam update step. For each gradient tensor G/ j E [4] and iteration i, AdaRankGrad efficiently identifies a unique set of signifi- cant projection directions (subspace) Pj E R along which the gradient Gj exhibits
Rex* at layer
the largest changes, where rf is the lowest possible rank to still maintain a predefined information frac- tion (given threshold) relative to the original, non- projected gradient. Practically, P.G. is the low-
projecte-ggraiie and P P/G/ is the best low- rank approximation the of the gradients G. that em- bodies the required fraction of its information. The projections P/ are being adaptively updated through- out training (based on the convergence criteria of the gradients on the projected subspace), where their rank rf is dictated by preserving the given information threshold. The method ensures: (1) The method de- termines the optimal projection dimension for each layer's gradient tensor independently, adjusting dy- namically throughout training. This rank adjustment leverages property we prove that the effective dimen- sionality of the full gradients gradually decreases over time, allowing updates to be performed in a lower- dimensional projection space, thereby reducing mem- ory usage. (2) The projection matrix for each layer's s gradients is updated based on a convergence criterion within their respective subspace. This ensures up- dates occur precisely when needed, avoiding premature or delayed transitions between subspaces and resulting in faster overall convergence.
We -Wi
Hy
Wy iy, WWy 4R out 10 S5
Figure 1: The illustration shows how AdaRankGard 3 is trained. First, the gra- dients G, are projected into a 3D space (in this example), represented as G p " G. As convergence occurs, the gradient's dimension decreases to a 2D space and then to ID space. This dimen- sionality reduction indicates convergence while efficiently using memory.
Table I: Comparison between AdaRankGrad, GaLore, and LoRA. Assume W E Raxm 2 m). constant rank r, and adaptive-rank aadaa (with initial-rank Tink = r).
ArrRankGrad
GaLore
LoRA
Weights
nm
nMm + 12 + mr
Optim States (Tadap r)
aaa + 2
12 + 2mr
2nr + 2mr
Subsipee
x
X
x
Adaptive-Subspace
X
X
Pre- Tre-Training
X
Fine- Tine-Tuning
2 RELATED WORK AND BACKGROUND
Memory efficient optimizers. Memory-efficient optimization has been a recent focus of research. Multiple studies have aimed to reduce the memory requirements of gradient statistics in adaptive op- timization algorithms (Shazeer & Stern 2018 Anil et al., 2019). One common approach is quanti- zation, which helps decrease the memory footprint of optimizer states (Li et al. (024).. Additionally, recent advancements have suggested reducing the memory used by weight gradients by integrating the backward operation with the optimizer update (Lv al. 2023a b). This characteristic has been leveraged to reduce memory usage during training processes (Gooneratne et al.) 2020 Huang et al.) 2023 Modoranu et al. 2023)..
Low-rank gradient optimization. The phenomenon of low-rank gradients naturally arises during the training of neural networks, subject that has been extensively examined both theoretically and
practically, e.g., Zhao et al. (2022); Cosson et al. (2023); Yang et al. (2023). This characteristic low-rank structure gradients has been leveraged to reduce memory usage during training processes
Gooneratne et al.] (2020); [uang et al.] (2023); Modoranu et al. (2023), and results in a reduced computational complexity as compared to standard gradient descent methods.
3

---

Published as a conference paper at ICLR 2025
Adam optimization. Arguably, among the most popular optimization methods used for train- ing large language models (LLMs) are the Adam optimizer (Kingma & Bal 2017) and its vari- ant, AdamW (Loshchilov & Hutter, 2019), which incorporates weight decay for regulariza- tion. However, it is well-established that Adam optimization has higher memory complexity compared to other optimization alternatives. To illustrate this, let briefly review how the Adam optimization algorithm operates. First, we need to establish some notation. Consider neural network denoted as ();;), which consists of layers and is parameterized by &
WT xdo "rr"""....,.... w.... "" x Here, W. represents the weights tensor parameters as- sociated with the i-th layer, for i E [L]. In the following, let t N represent the t-th step of the Adam optimization algorithm. Then, we recall that the single update step in Adam is given by, Specifically, time step f, G, denotes the backpropagated gradient matrix, i.e. 4 (0:-1). The exponentially weighted moving averages of the first and second moments are denoted by M, and Vt, respectively, with their bias-corrected coun- terparts given by M, and V.. The AdamW optimizer up- dates the model parameters at step l. according to the rule,
Adam (single update step)
G. = V (0t-1),
M. B: (1 B1) G.,
V. -2V + (1 B2) G?;
M, MM./ (1 - -B1):
M, 0: 0 (Vetee + 10.-1) where À 2 0 is the weight
V. V./ (1 35)
decay rate (for Adam À 0), and all operations are per- formed element-wise. In this equation, 31 and B2 control the decay rates for the moving averages of the moments, is the learning rate, and e is a small constant to avoid division by zero. Notably, since Adam/W requires storing both M, and V, each time step, it incurs additional memory footprint of 2mn.
0. θ - aMie/ (V
While existing approaches (Zhao et al., 2024b Vyas et al. 2024 Okewu et al., 2020) focus on low-rank approximations of the first and second moments with the goal of reducing memory quirements, we propose to approximate the gradients by a low-rank factorization. Consequently, in our scheme the moments are integrally constrained onto this reduced dimension, and thus we gain both benefits.
3 METHOD AND MAIN RESULTS
3.1 THEORETICAL MOTIVATION: GRADUALLY GRADIENT RANK VANISHING
As mentioned in the introduction, few recent em- pirical results (e.g.. (Jaiswal et al. 2024b Zhao al. 2024a| Lialin et al. (023a)). demonstrate that the gradients, when training or fine-tuning LLM's, are "approximately low-rank"". As an example, this phenomenon can be observed in Figure where it is evident that the squared norm of the gradient's 's sin- gular values decay to zero exponentially fast.
50% Sponet Thorset
Tharsed (Gger
6:44
6.00
As hinted above, this phenomenon only true the approximate sense; roughly speaking, only very few eigenvalues hold almost all the information captured by the gradient. Accordingly, a low-rank matrix ap- proximates the underlying gradient up to a negligi- ble approximation error. The practical implication is that while the weight matrices are not necessarily low-rank, training certain high-rank layers with low- rank based-gra -eadddggai updates is possible. To make the above discussion precise, consider the following def- inition for approximate low-rank matrices.
The figure illustrates the exponen- tial decay of eigenvalues in the MLP layer's gradient, at the first iteration of fine-tuning RoBERTa-Base (Liu) 2019) model, on the MRPC task, from GLUE (Wang et | 2019). Notably, the red line indicates that 50% of the gradient information (in terms of squared norm ratio) is captured by the first eigen- value, while the green line shows that 90% is contained within the first two eigenvalues.
Figure 2:
Definition 1 (Approximate low-rank matrix). A matrix A E RR is called (n, E) approximmaell rank-r, if there exist n E [0, 1), O, and a matrix Aapp.r E Rx with rank Appp.r) and r < min(n, n, m), such that,
A Aapp.r || S || + +8.
(1)
4

---

Published as a conference paper at ICLR 2025
As it turns out, it be shown (see, e.g., (oolb & Van Loan 2013 ) that the optimal Aapp.r minimizing the approximation error in the left left-hand-sidee of equation I can be obtained by applying an SVD on A and then retaining only the top r singular values and their corresponding singular vectors. Mathematically, we have Aapp.r = Ex=1 ,.u.V," where {ar]; are the singular values of A, and (u,); and {v,) are the corresponding left and right singular vectors, respectively. The approximation error is turn given by || App..- 1| = Ei-tt of. This construction gives an (7,, 0) approxximatly rank-r matrix, with the minimal V/a possible.
Recently in (Zhao al. 2024a). the structure of the gradient for wide family of nonlinear networks known as "reversible networks' TTan et al. 2021 was studied defined as follows.
Definition (Reversibility (Tianet al.) 2021)) A layerl is reversible if there is a GG( 0) Reexne- so that the pre-activation at layer l satisfies fa( 0) = Ge(æ; f) (x; 0) and back-
propagated gradient after nonlinearity -t-1 Gg (a;; O)P? (a:; g). for some matrix P Pa(x:; 0)
Rnexn,, A network is reversible if all layers are.
For simplicity of notation, we use G; to denote [G:(2; where t E N is the iteration index in the optimization process. Furthermore, when it is clear from the context, omit the layer index & from our notations. Assuming reversibility and SGD weight update (i.e., W, = W + aG,-1), 1), it is shown in Zhao al. (2024a) that for both le and cross entropy losses, the gradient is of the structure form G matrices, and (B., C. are certain positive semi-definite (PSD) matrices. Furthermore, was proven that if the gradient G1, has the above structure for all 2 to, for some N, then, the stable rank (G,) A G.Ill satisfies,
Di= (A; B,WC,), where N the batch size, {A.). are input-dependent
G.ll2
sr (G.) S sr (G!)) + (1 a
1 17/2
2(1-to)
Go G!!
/
a < X2 denote its two smallest distinct eigenvalues, and G!,, is the projection of G. onto the minimal eigenspace Vi of S that corresponds to X₁- Accordingly, as t - 0o, we get that the final stable rank is upper bounded by sr sT(G!!). Under the same gra- dient structure assumption and for the vanilla settings of the SGD weights update |Batassh et al. (2024), we were able to prove the following stronger result. We prove that the approximated sta- ble rank of the gradients approach we need to make a few notations. Let G2 = V. be the SVD decomposition of G., and let P((,r) = U : r), r)] be the corresponding projection matrix. When clear for the context, we omit the index I and use P,(r) = P.(1 = 1, We have the following result.
where -
N E
the training process progresses. To state this result,
Lemma (Asymptotically rank-one). If a neural nerwork is trained using vanilla SGD, then the following holds for the gradient of a reversible layer at iteration I
Effective Rank
<() A |G, PP(()) < <O(CT))
|G&
for some constant C>1 1.
0
2000 4000
6000
8000
Training Steps
Figure 3: The figure presents the effective rank (see Section 44 measured after every 100 update steps on the RTE dataset, from GLUE (Wang et al. 2019).
The above result implies that G, approaches its rank-one approximation P, as the it- eration number increases, namely, Gg becomes rank-one. The proof of Lemma relegated Section B
we demonstrate that for large language model (RoBERTa-base Liu 2019)). ) which contains also non reversible layers, the rank decay evolves as function of the number of update steps, in a fine-tuning task.
Finally, in Fig. and Fig.
It can be shown that this family includes many different kinds of layers, such as, linear layers (MI. and Conv.), and (leaky) ReLU nin-liiy
5

---

Published as a conference paper at ICLR 2025
3.2 SRAPTIVELLOW-RAE SUBSPACE SELECTION
As aforesaid, our training strategy is self-adapting by updating the low-rank projection subspace for the gradient of each layer, leveraging the phenomena raised by Lemma [ The most natural but computationally expensive way search for the subspace on which a gradient has the highest variance is by applying the truncated SVD (Eckart & Young 1936), with a truncated rank that preserves the predefined fraction of information. In this section, we first utilize an efficient method to approximate the gradient, which is more practical than the expansive truncated SVD approximation. We then use this method propose an algorithm that identifies the most suitable projection subspace with the smallest possible span/rank that maintains the given information threshold.
3.2.1 IDENTIFYING PROJECTING SUBSPACE BY POWER ITERATION
Consider a matrix Raxm Finding its "best" low-rank approximation can be framed as the following optimization problem minq.u u || QU I where Q E m and U E R=xrr As discussed above, the matrix Aapp.r QU represents the (n,E) rank-r approximation of A. Com- puting full SVD for large matrices is computationally intensive and memory these challenges, we detail Algorithm T "good" proxy for the optimal low-rank approximation of A.
To address (Halko al., 2010), an efficient technique for producing
Algorithm is designed solve the optimization and ap- problem arg mill A --QQA
Algorithm Subspace selection via random- ized range finder (SSRF)
A
F
proximates the matrix as Aapp.r se QQ A. The computational complexity of the SSRF algorithm is dominated by two main operations: the matrix mul- tiplication ASt whose computational complexity (mmrr), and the QR decomposition of the resulting matrix Y. requiring (mr²). Therefore, the total complexity is Omnr m)
Inputs: Matrix A 12xx target rank S min m). Initialization: Ω RR N(O, Y +
Construct Q R.xxr using the QR de- composition of Y
Return: Q
In scenarios where < m, n, this complexity simplifies to O(mrr), which is significantly more efficient than using SVD, whose complexity is Ommnn(mn²,, m²n)). This low complexity makes the SSRF preferable for extracting leading singular vectors in large-scale data settings.
3.2.2 FAST ADAPTIVE LOW-RANK APPROXIMATION
is clear from Definition that a larger value of the rank r increases the computational op- erations and induces a higher memory foot- print, while smaller value of will result a non-negligible approximation error. To mitigate this, we suggest Algorithm [ an adaptive procedure for rank selection that ad- justs the value of to preserve at least tain fraction of the information in A. Specif- ically, at start we set ro = Tmins and then find the best Aapo.r Rs QQ A using SSRF. At each time step t, given the current rank Tt. we compute the approximation error as,
Algorithm 2 Information based adaptive subspace selection (IASS)
Inputs: Matrix A E RR*xm, minimum rank P'in 2 1, maximum rank Trax threshold VM E (0,1). 1).
min(n, m). and information
Initializations: 1. 4- 0, <- Tmin <- We + while 2 1 and Tiax
2 1 do
Q SSRF(A, r()
{Call Algorithm
A Aspp.r < QQ A.
Ag-Ap Il's
--in (n.m) 2i (G)))
T/s
|AI
Ei o, (G
T/
Th then
Tmu + [
A
A Agani 11
else if nr
< then
A
Ft
Tue
( 0. (A)2
end if
end while Return:
Σ (A)²
Q Pt-
We then apply binary search until find the maximum for which the corresponding condition still holds: 17 < 7a- It is important to note that, from the outset, we know a priori that the rank the approximated gradients is inherently low. Consequently, the initial value of rmax E min min(m, (m, is

---

Published as a conference paper at ICLR 2025
small, and the number of iterations required for the search is correspondingly minimal, specifically (log(TTmxx T'mi)); obviously (Tmax
T'mn) < min n).
3.3 ADAPTIVE LOw-RANK AND MOMENTS GRADIENT OPTIMIZATION
We can now present our main algorithm for adaptive low-rank and moments gradient optimization in Algorithm B The mathematical formulation of the weights update rule proposed in this paper is detailed in Appendix A comprises four main blocks, all contained within an outer loop that terminates once reach convergence. The role of each block is as follows.
Block 1: We select the (approximated) subspace along the directions of the r largest eigenvectors, using Algorithm 2 The number of orthogonal directions r is determined by the information threshold required to preserve the gradient information, according to equation 2
We transform the first and second gradient moments evaluated during the Adam update steps between the previous and the updated subspace. The main reason for this transformation is because, will be seen below, in the fourth block, the first and second moments of the gradients are aligned with the previous projected subspace, and thus, a transformation is needed to convert them from the previous subspace to the current one.
Block 2:
We tune the parameters during the low-rank update step. The stopping condition trig- gering the update of the orthogonal projection matrix is based on the convergence of the projected gradient onto the subspace.
Block 3:
Block 4: The actual pre-trained model parameters are updated, using the low-dimensional and memory-efficient projected gradients on the selected subspace.
Theorem 2 (Convergence of Algorithm 3 For a loss function L, and given architecture 4. suppose that the compositions of f = L (.) is B-smooth non-convex function that are bounded by some M E R... Let G denote the gradient matrix w.r.t the jth reversible layer at time t E N. Assume that ||G/|l= F S D. for all j E [L] and I. E N. where D E R,. Then, for any >> 0, there exists
CE R. such that for all T > S.. it holds that T Ej= EN-1 Ei=T,
T
G F
S In
particular, Algorithm 3 with vanilla SGD weight update 52 22 v Thh |[,
.lpe. where
i is the Block 3 entry counter, and learning rate a < 2/ 2mxxxx achieves an e-critical point i.e.
G < for some E N, and any E [L].
Note that, for clarity, can assume, without loss of generality, that n. In the opposite case, the projection matrix would multiply the gradient from the right side.
The proof of Theorem can be found in Section B, A few important comments are in order. First, to reduce memory usage, we apply in Algorithm a per-layer weight update during backpropagation, as proposed by recent works, see, e.g.. Lv al. (2024). This is in contrast common optimizers which usually update all weights after backpropagation by storing the full gradients in memory. which could be highly inefficient. Second, note that the Adam update step block in Algorithm 3 can be replaced by any quantized Adam variant, e.g., Li et al. 2); Chen et al. (2021); Seok & Kim (202T) and as so allows obtain tasked fine-tuned quantized model or quantized adaptor; this is discussed in more detail below. Finally, we would like to mention here that one can easily apply
bit projected gradient updates, as introduced in Q-GaLore (Zhang al. 2024).
If, after fine-tuning, one wishes to create an adapter (i.c.. a parallel low-dimensional LoRA type model) alongside the original model, this can be done efficiently as follows. First, we calculate the training weights gap A W Winee-Tuned
W Wreeuuined. where W Wrine-Tuned the model weight at the end of the process, and W WPretrained is the original model weight. Then, rank(A), using some matrix ranking algorithm, and finally. we solve a| ABI using any optimization algorithm (e.g., gradient de- scent). Note that any solution this optimization problem a global optimum Kawaguchi 2016).
we find the l'Adapior
min "x BER X
We focus on SGD for the simplicity (as is standard practice in related literature, c.g., 7hao ct al. 2024a))
Asss known e-stationary, see, c.g., (Cosson al. 2023

---

Published as a conference paper at ICLR 2025
Algorithm 3 Adaptive low-rank and moments gradient (AdaRankGrad)
Input: A Layer W R=** from 0 dataset D, loss function L, information thershold 0 mmn initial rank Vini, maximal rank Ymxx, learning rale a, and small numbers <1, <2 >
Initialization: = 0. Sample batch B - {I,, V: M= D
Compute batch gradient G. + E ow L (b(2,, 84)
[B]
for unsupervised L ( 0))
while |G > do
Block Adaptive subspace selection
Qu T2 + IASS(G; G., Pinits Ymax; 17th)
(Approximated ry-dimension subspace: Q projected matrix) }
Block 2: Moments subspaces transformation
R x Q, Q 1 if 2 1, else 0**xr,
M M R,M, 1, 2 else xm
(1 fd forder moment)
VIixm f- V,-1, if 2 1, clse 07
ooddr moment)
Block Low-rank optimization
G Q, G, while ||G|| do
(Projected gradient the approximated ry re-dimension subspace)
Block 4: Adam update step
M. 31M: (1 B1) G.
V: - 32 + (1 - 82) G
M. < M./ (1 (1-P3 Mi)
V. < V./ (1- -B)
W. W. aQ.Mi/ (V
t ++1
Sample batch B - foi, (1 )1 a
Compute batch gradient G2 <- Σ aw L. (* 0),
G. + Q. G.
end while
end while
(Exit by convergence criteria could alternatively be defined by the number of epochs]
Return W.
4 EXPERIMENTS
In this section, we test the performance of our algorithm on real-world datasets. Before discussing the setup we rely on our experiments, we define four measures for memory usage reduction in the rank-adaptive projection matrices. Specifically, for the jth layer, define the effective layer- gradient-rank, by, R.
ST R where R is the rank of the jth layer projection matrix at for some T N, and accordingly the total weighted-average effective rank is defined as Rudap Eim ET (dd Following that, define the average per-layer T.E,=o(d,dd4t.)
time 0 < t T
reduction in memory footprint, when compared to the non-adaptive low-rank fine-tuning, by M A (R/ - -Raaag) (d, + d)+1), where Rí is the time-independent non-adaptive rank (such in Galore). Finally, we define the total memory reduction by Mred A Σ} (R' -Raanp) (d; + dj+1).
Fine-tuning on GLUE benchmark. We evaluate our model on the GLUE benchmark (Wang et al. 2019) by fine-tuning the pre-trained Roberta-base model Liu (2019) on 8 tasks. We compare the results against full fine-tuning, LoRA, and GaLore methods and present them Table 2 We report the overall (matched and mismatched) accuracy for MNLI, Matthew's correlation for CoLA, Pearson correlation for STS-E, B, Fi-score for MRPC, and accuracy for other tasks. As can be seen, method improves the accuracy of fine-tuning while consuming less training memory average. Fig. shows the average memory reduction measured at the end of each epoch (bluc) and the corre- sponding effective rank (green). Empirical analysis of the hyperparameters appears in Appendix [
8

---

Published as a conference paper at ICLR 2025
Table Evaluating Arad comparing state-of-the-ar memory-cfficient fine-tuning meth-
ods on GLUE benchmark using pre-trained RoBERTa-Bann Base. For AdaRankGrad, we present accu- racy results with average effective rank.
Model
Memory
CoLA
STS-B
MIPC
RTE
SST2
MNLI
ONLI
QQI
Full Fine-Tuning
747M
62:24
90.92
91.30
79.42
94.37
87.18
92.33
92.23
(rank=4)
257M
90.37
86.82
92.18
ore (rank-4)
90.73
91.06
AdaRankGrad (Initial mank-4)
202M
61.4(3.71)
90.97(3.79)
92.613.72)
81.23(3655)
94.833911
86.063.99)
92.5(3.9)
90.4(3.67)
LoRA (rank:8)
MAN
61.33
V0.00
91.90
79705
93.46
86.94
12.25
91.22
trank=8)
Addl (Initia tanke5)
237M
62.016.41)
90.897.77)
93.215.33)
81.23(6.64)
94.800.600 69)
86.5(6.00)
92,6(6.54)
89.766.00)
60%
Effestie Rank
M
50%
40%
Menoyy Ralselim
Ellextive Rank
49%
Mimory Rolocioo
91
30%
101
791
1481
2171
2861
1551
200
$581
9999
10451
723
1340
1973
297
5331
3845
1400
5111
6341
0965
7389
5113
8351
9461
0%
Toining Stop
Training
Figure 4: We present the effective rank measured for non-attention layers and corresponding mem-
ory reduction for AdaRankGrad trained on MRPC (left panel) and RTE (right panel) datasets from the GLUE benchmark.
Fine-tuning Geneformer on Biological Omies Tabular Data High-throughput omics technolo- gies, such as next-generation sequencing Reii-Filho (2009) allow for the simultaneous measure- ment of thousands millions of biological molecules, capturing biological processes at the tissue or single-cell level. Recent studies have aimed to construct foundation models for omics data CCui et al. 2024); Theodoris et al. (2023). Unfortunately, the complexity of omics data makes the use of low-rank optimization methods such as LoRa sub-optimal for foundation omies models, and there- fore, training is often based on relatively high ranks (Chen et al. 2024). In this experiment, we demonstrate that the proposed approach projection of gradients. We conducted fine-tuning experiment for cell classification, specifically focusing on the Classification of Cardiomyopathy Disease States. We used the Geneformer gf- /2-- 30M-i2048{I as our pre-trained base model, which is based on the BERT model and pre-trained on 30 million single-cell transcriptomes (Theodoris et al. 2023 Devlin 2018). Following the setup proposed in (Chen et al. (2024), we fine-tuned our method with a maximal rank of 16 on approxi- mately 93,600 samples for three epochs, with a batch size of 16. We evaluated the model on test set of size 17,228. Figure 5 presents the Macro-Fl FI score and Accuracy measured for each epoch. As shown, our method converges faster with a stable improvement over LoRA.
alleviate this problem due to the adaptive low-rank
AdaRankGirad
1oRA(r16)
Galoce
- AdaRisieernd ToRA(t-16) --de-
0.7
0.88
0.68
0.86
0.84
0.64
Accuracy
0.82
Macro-FI
0.62
0.8
0.6
0.78
0.58
0.76
0.56
0.74
0.72
0.52
0.5
2
Epocias
Epochs
Figure 5: We evaluate AdaRankGrad in the Geneformer fine-tuning task and present Accuracy (left panel) and Macro-1 (right panel) measurements. The adaptive low-rank projections improve the model converges compared LoRA and Galore methods.
AA transformer-baste model trained on 30 million single-cell transcriptomes for gene classification and in silico perturbation analysis. It has 12 layers with an input size of 2048 genes per cell, designed to understand gene network dynamics. Huggingface: httpps:/ www.nungggggnnna..n.oo.oooo co/ /tto/cttee heodor a/Gene

---

Published as a conference paper at ICLR 2025
Pre-training LLAMA C4 Dataset Here, repeated the comparison presented in Zhao al. (2024a) [Table 2] to evaluate AdaRankGrad performance to the state-of-the-ar method, in terms of perplexity and memory usage. We evaluate AdaRankGrad by training large LLaMA-baser models on the C4 dataset, a cleaned and massive version of Common Crawl's web corpus (Raffel et al.) 2020). This dataset is primarily used for pre-training language models and learning word repre- sentations. To closely mimic practical pre-training scenarios, we train on a large dataset without repeating data, scaling model sizes up 350 million. The results are presented in Table 4
Table 3: The LLaMA 7B model was pre-trained on the C4 dataset for 120K steps. Validation perplexity and memory usage estimates reported.
Steps/Tokent
8-bit GaLore
8-bit Adam 8-bit AdaRankGrad
40K 5.2B
17.94
18.09
17.86
80K /10,5B
15.39
15.47
15.27
120K /15.7B
14.95
14.83
14.87
Mem
I8G
26G
16.4G
Table 4: comparison of low-rank state-of-the-art algorithms for pre-training LLaMA models of varying sizes on the C4 dataset. We use initial rank Timi for AdaRankGrad, w.r.t the presented in the last table raw. The information threshold 7hh = 0.48 for the three model training. The valida- tion perplexity is presented, along with estimate of the memory required for the total number of parameters and optimizer states in BF16 format. We used NVIDIA 100 for the three experiments.
60M
130M
350M
1B
Full-Rank
(4.06(0.36G)
25.08(0.76G
18.80(2.06G)
(7.06
GaLore
34.88(0.24G)
25.36(0.52G
18.95(1.22G)
15.64(4.38G)
Low-Rank
78.18(0.26G)
(0.51
37.41(1.08G)
142.53(3.57C)
LoRA
(0.99()
33.92(0.800)
25.58(1.76G)
(9.21
ReLoRA
37. 37.04(0.36G)
(0.37()
29.08 (1.76G)
(8.33(6.17G)
AdaRankGrad
34. 34.24 (0.20G)
25.22 (0.49G)
18.91
(1.91(1.16C)
14.71
Training Tokens
I.I
2.2 B
6.4
13.1B
T r/dnndel
128/256
256/768 768
256/1021
512/2048
5 DISCUSSION
In this paper, we present AdaRankGrad, full-parameters efficient optimization scheme that applies adaptive low-rank updates without relying on a parallel low-rank adapter (LoRA), thus preserv- ing the natural training dynamics. Unlike methods such as LoRA, which rely on parallel adapters, GraRankGrad enables full parameter fine-tuning while maintaining low memory costs through ficient low-rank optimization updates. Moreover, unlike Gore AdaRankGrad leverages the nat- ural phenomenon where the dimensionality the approximated gradient decreases as training pro- gresses. This allows adaptive updates of the projection subspace only when the gradients have converged to a lower-dimensional space, ensuring that optimization resources are fully utilized- no unnecessary updates occur before after convergence. As demonstrated in Section this approach results in superior prediction performance. AdaRankGrad offers unique trade-off between memory efficiency and model performance. While it may slightly increase training time due the need to determine an optimal rank for projecting subspace updates, these updates are infrequent. The search for the best subspace rank is conducted within very narrow range, reducing the overall computa- tional cost. To further mitigate this, we proposed an efficient subspace-update technique that sig- nificantly reduces the SVD-bbsed subspace calculations by an order of magnitude. practice, this optimization compensates for the additional time spent in subspace search, making AdaRankGrad a competitive choice compared to methods like GaLore. For future research, we suggest exploring other update steps besides Adam, such AdaFactor, and studying different efficient algorithms for computing the subspace over which gradients are projected. Another promising avenue would be to investigate algorithms that search for the optimal subspace rank, which preserves a similar fraction of information as the full gradient. Finally, evaluating this method's 's effectiveness in knowledge editing (Rozner al. 2024) is an interesting topic for future work.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
The work of OL is supported by the MOST grant number 0007341.
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimiza- tion, Advances in Neural Information Processing Systems, 2019.
Barak Battash, Lior Wolf, and Ofir Lindenbaum. Revisiting the noise model of stochastic gradient descent. In International Conference Artificial Intelligence and Statistics, pp. 4780-4788.
PMLR, 2024.
Amir Beck. First-order methods in optimization. SIAM, 2017.
Tom Brown, Noam Smith, et al. Efficient memory usage in large language models. In Proceedings of the 36th Conference on Neural Information Processing Systems (NeurIPS 2022), 2022.
Congliang Chen, Li Shen, Haozhi Huang, and Wei Liu. Quantized adam with error feedback. ACM Transactions on Intelligent Systems and Technology (TIST), 12(5):1-25 2021.
Han Chen, Madhavan Venkatesh, Javier Gomez Ortega, Siddharth Mahesh, Tarak Nandi, Ravi K Madduri, Karin Pelka, and Christina V Theodoris. Quantized multi-task learning for context-specific representations gene network dynamics. bioRxiv, 2024.
Y. Chen, Y. Li, and X. Liu. Lorat: Improving low-rank adaptation with parameter-specific learning rates, arXiv preprint vxXiv:2355.1604 2023.
Romain Cosson, Ali Jadbabaie, Anuran Makur, Armin Reisizadeh, and Devavrat Shah. Low-Rank Gradient Descent. IEEE Open Journal of Control Systems, 2023.
Haotian Cui, Chloe Wang. Hassaan Maan, Kuan Pang, Fengning Luo, Nan Duan, and Bo Wang. segpt: toward building a foundation model for single-cell multi-omics using generative ai. Nature Methods, pp. 1-11, 2024.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Ilms, 2023.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1810004805 2018.
Carl Eckart and Gale Young. The approximation of one matrix by another of lower rank. Psychome- trika, 1(3):211-218, 1936.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
M. Gooneratne, K. C. Sim, P. Zadrazil, A. Kabel, F. Beaufays, and G. Motta. Low-rank gradient approximation for memory-efficient on-device training of deep neural network. 2020 IEEE International Conference Acoustics, Speech and Signal Processing, ICASSP 2020, Barcelona, Spain, May 4-8, 2020. IEEE, 2020.
Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A. Roberts. The unreasonable ineffectiveness of the deeper layers, 2024.
Nathan Halko, Per-Gumnar Martinsson, and Joel A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions, 2010.
Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. Parameter-efficien fine-tuning for large models: A comprehensive survey, 2024.
Shwai He. Liang Ding, Daize Dong, Miao Zhang, and Dacheng Tao. Sparseadapter: An easy approach for improving the parameter-efficienng of adapters. arXiv preprint arXiv:2210.04284, 2022.
11

---

Published as a conference paper at ICLR 2025
E.J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Li, S. Wang. L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.0960 2021.
Shuang Huang, Brian D. Hoskins, Michael W. Daniels, Matthew D. Stiles, and George C. Adam. Low-Rank Gradient Descent for Memory-Efficient Training of Deep In-Memory Arrays. ACM Journal on Emerging Technologies in Computing Systems, 2023.
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, 2024a.
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. arXiv preprint arXiv:2407.11239 2024b.
Kenji Kawaguchi. Deep learning without poor local minima. In D. Lee, Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29.
Curran Associates, Inc., 2016.
Beomsu Kim, Huishuai Xu, Jinxian Tang, and Song Han. of compressed large language models via sub-4-bit integer quantization. arXiv preprint arXiv:2305. 14152, 2023.
Memory-efficient fine-tuning
Yunscong Kim, Yujeong Choi, and Minsoo Rhu. Paris and elsa: An elastic scheduling algorithm for reconfigurable multi-gpu inference servers, 2022.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Bin Li, Jie Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 2024.
Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Goldstein. Training quantized nets: A deeper understanding. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates, 2023a.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates. 2023b.
Yinhan Liu. Roberta: robustly optimized bert pretraining approach. arXiv:1907.11692, 2019.
arXiv preprint
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimization with adaptive learning rate, 2024.
Ke Lv, Haoyi Yan, Qibin Guo, Hao Lv, and Xipeng Qiu. Adalomo: Low-memory Optimization with Adaptive Learning Rate. ArXiv preprint arXiv:2310.10195, 2023a. SxKi:23110.10999
Ke Lv, Yongwei Yang. Tianyu Liu, Qingyu Gao, Qibin Guo, and Xipeng Qiu. Full Parameter Fine- tuning for Large Language Models with Limited Resources. ArXiv preprint 27X792230..0788. 2023b. arXiv:2306.09782.
Xiangdi Meng, Damai Dai, Weiyao Luo, Zhe Yang, Shaoxiang Wu, Xiaochen Wang, Peiyi Wang, Qingxiu Dong, Liang Chen, and Zhifang Sui. Periodiclora: Breaking the low-rank bottleneck in lora optimization, 2024.
Ionut- Vlad Modoranu, Alexander Kalinov, Ermin Kurtic, Erwin Frantar, and Dan Alistarh. Error Feedback Can Accurately Compress Preconditioners. ArXiv preprint arXiv:2306.06000 2023. arXiv:2306.0666 060988....
12

---

Published as a conference paper at ICLR 2025
Emmanuel Okewu, Sanjay Misra, and Fernandez-San Lius. Parameter tuning using adap- tive moment estimation in deep learning neural networks. In Computational Science and Its Applications-ICCA. 2020: 20th International Conference, Cagliari. Italy, July 1-4, 2020, Pro- ceedings, Part VI 20, pp. 261-272. Springer, 2020.
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. Journal of machine learning research, 21(140): 2020.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimization towards training A trillion parameter models. CoRR, abs/1910.02054, 2019.
Yehonathan Refael, Adam Hakim, Lev Greenberg, Tal Aviv, Satya Lokam, Ben Fishman, and Shachar Seidman. Slip: Securing Ilms ip using weights decomposition, 2024.
Jorge S Reis-Filho. Next-generation sequencing. Breast cancer research, 11(Suppl 3):S12, 2009.
Amit Rozner, Barak Battash, Lior Wolf, and Ofir Lindenbaum. Knowledge editing in language models via adapted direct preference optimization. arXiv preprint 20ti::20000099220, 2024.
Jinwuk Scok and Jeong-Si Kim. Quantized adam with monotonically increasing resolution of quan- tization. In 2021 International Conference on Information and Communication Technology Con- vergence (ICTC), pp. 1719-1721. IEEE, 2021.
Noam Shazeer and Mitchell Stemn Adafactor: Adaptive learning rates with sublinear memory cost. Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stock- holmsmässan, Stockholm, Sweden, July 10-15, 2018, 2018.
Jonathan Svirsky, Yehonathan Refael, and Ofir Lindenbaum. Finegates: Lims finetuning with com- pression using stochastic gates, 2024.
Yu Tang, Chenyu Wang, Yufan Zhang, Yuliang Liu, Xingcheng Zhang, Linbo Qiao, Zhiquan Lai, and Dongsheng Li. Delta: Dynamically optimizing gpu memory beyond tensor recomputation, 2022.
Christina V Theodoris, Ling Xiao, Anant Chopra, Mark D Chaffin, Zeina R Al Sayed, Matthew Hill, Helene Mantineo, Elizabeth M Brydon, Zexian Zeng, X Shirley Liu, et al. Transfer learning enables predictions in network biology. Nature, 61867965):616-644, 2023.
Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning with dual deep networks, 2021.
Nikhil Vyas, Depen Morwani, and Sham M Kakade. Adamem: Memory efficient momentum for adafactor. In 2nd Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ ICML 2024), 2024.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In Advances in Neural Information Processing Systems, volume 32, 2019.
L Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song. Zenglin Xu, and Tim Kraska. Superneurons: dynamic gpu memory management for training deep neural networks. SIGPLAN Not., 53(1):4--53, February 2018. ISSN 0362-1340. doi: 10.11453200691.317844))
Y. Wang, H. Zhang, and Li. Adaptive lora: A rank-adaptive method for efficient Ilm fine-tuning. arXiv preprint 2336v:2306.06 2023.
Jayi Xu, Zhibo Liu, and Cheng Wang. Understanding Ilms: A comprehensive overview from training to inference. arXiv preprint arXiv:2401.00 .20088, 2024.
Z. Xu, J. Zhang, and X. Wang. Sparse lora: Efficient fine-tuning of large language models with sparse low-rank adaptation. arXiv preprint arXiv:2303. 102003..7600 17626, 2023.
Greg Yang, Jacob B. Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310. 1708833. 2023. arXiv:2310.17811
13

---

Published as a conference paper at ICLR 2025
Minchen Yu, Ao Wang, Dong Chen, Haoxuan Yu, Xiaonan Luo, Zhuohao Li, Wei Wang. Ruichuan Chen, Dapeng Nie, and Haoran Yang. Faaswap: Slo-aware, gpu-efficient serverless inference via model swapping, 2024.
Jihun Yun, Aurélie C Lozano, and Eunho Yang. Adaptive proximal gradient methods for structured neural networks. Advances in Neural Information Processing Systems, 34:24365-24378, 2021.
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- arXiv preprint S.Kiv:2303.105122 2023.
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, 2024.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection, 2024a.
Jingzhao Zhao, Frederik T. Schaefer, and Anima Anandkumar. Zero initialization: Initializing neural networks with only zeros and ones. Transactions on Machine Learning Research, 2022.
Pengxiang Zhao, Ping Li, Yingjie Gu, Yi Zheng, Stephan Ludger Kölker, Zhefeng Wang, and Xi- aoming Yuan. Adapprox: Adaptive approximation in adam optimization via randomized low-rank matrices. arXiv preprint arXiv:2403. 14958, 2024b.
14

---

Published as a conference paper at ICLR 2025
A UPDATE STEP RULE FORMULATION
To describe AdaRankGrad update step rule, we need to establish few important notations first. For t IN, let P: Rm m= be an entry-wise gradient update rule (e.g., Adam, AdamW, AdaFac-
tor, etc.). Fort E N and layer j [4], recall the SVD of the gradient G/ = U. (V)) T. and further
define the projection matrices (r() = SSRF(U,]::: ).
SSRF V/l:: r!) where
for given threshold 0 < T/h < 1, we let rf å sup(r E N : 7m ||G 11 ||G - ( 2 0), computed using Algorithm Next, for j E [4], and some > 0, we let (T/ with T3. O,
denote the monotone sequence of integers, for which (P4, ( G R2 (r (rfp))l S s; under certain conditions that we list below, it is shown in Cosson et al. 2023 TThm. 3.4] (as well as in Zhhao et al. (Th4or)T 3.8), for reversible layers) that T,'s are finite. Then, for E N, and 1- the AdaRankGrad weight update is given by,
E
w W
nrP4. (rt))r Pe (PP), G G (R (+))+ E=
Σ
for j E [L], where 7t is the learning rate at time adapted by Pt. and WT. Wo is pre-trained model, in a case of fine-tuning task, or given weights initialization, in case of pre-training task. In the sequel, we let y A ooo denote the initial learning rate, and for simplicity of notation, when clear from the context, we drop the dependency of the various notations on the layer index j.
B PROOFS
B. PROOF OF LEMMA I
In this section, we prove Lemma 1 Consider the SVD decomposition of the gradient G, U,E, VT, at iteration t. For any natural number r < n, we denote H. (r) = U/:,1 1 r]. For simplicity of notation, let us denote P.(r) H, Hr) (r), where P, (r)r) is orthogonal projection matrix, i.c., P. (r)) = Pr(r), and (,(r)= = P (r). Without the loss of generality, we assume that at é 0. the rank of Go is such that rank (Go) r. For reversible networks, it was shown in Zhao et al. 2024a)| (2024a)[Thea [Theorem 3.2] that the gradients have the form G. = N Ei= (A, - B, B,W,C,), with constant matrices {A,): and PSD matrices {B,, C. for f to, where to E N. Recall that in the vanilla SGD weight update, we have W, W + 1/G Let S NV Σ C, B, and Å A be its two smallest distinct eigenvalues. In order to prove our result we rely on several results and arguments in the proof of Lemma 3.3 in (Zhao al. 2024a). Specifically, let G!!, be the projection of Go onto the minimal eigenspace V of S corresponding to Å1. By assumption, we know that the rank of G!!, is L,, and its SVD is G Et= quayt where (z and {ui ) are the orthonormal unit vectors, and {ct]-= are the corresponding singular values. Thus, it was in (Zhao et 2024a) that,
L Σ
d/ vec (G!!
c
with unit vector uy Y1 21 E V1. It is clear that
ii (yt
@ z (BV @ zz)
(yt yr) (zi zv:) (1 = I')
where I is the indicator function. Using this, it was finally shown in Zhao et al. 2024a that,
|Goll2
max
yTGyy
2 max xJ G max (y: 8 x1) gr = max w7 (1 n/SY' go = (1 - --n max oi go.
15

---

Published as a conference paper at ICLR 2025
Using the above results from (Zhao et al. 2024a), we can now that,
|Gell
max
z Gyy
l| 1
(y 09 2) T
max
gt.
|||||==1.||=____=
(y 2 x) T (1 - -1/5)/90 go
max
l|vll2 |
T (1 n))) go
(Uol:, 1] @ Vol:.
t - n))) g0
= (1 - -)=) o go-
Thus, we have
(1 m/a1' T 9u 2 (1 - ) max vj 90.
1
or equivalently
(1 (1-) TT 2 (1 - nAii'' |8!!|
(3)
Now,
G. P.
x(t) A
|Gell F
|.
||G1l
||+||2 2
P
|G+11
||G/
|Grll -1
|Gell2
|G 2
2
(1 m711)
28
1
(1 1122)
|| 112
|G, 2
| 2
21 (1 m711))
00
(1 - 11/22)
11000 12
1111)
21
TT
(1 1/\1 )²2
0 90
S1
(1 )
21
|
12
(1 1122) 21 11 11
(1 n1^1) 2²
0 go
(1 - 1711
go
where the first inequality follows from the fact that |G |/2 >
|. the second inequality is by
using Zhao et al. (2024a) [Lemma 3.3], i.c.,
|Gell s (1-n)2)* 2t 9 112 + n 2r
90
and the third inequality is due to Eq. equation Finally, by defining the constants ei (1-mA:)
(1-nda)
| 12
and C2
< 1, we get that x(() < 2 C2. which concludes the proof.
go
B.2 PROOF OF THEOREM 2
In this section, we prove Theorem We upper bound Frobenius norm of G for any layer j e [L]: in the following, for simplicity of notation, we ignore the index j an use G, instead. By Lemma 3 the low-rank optimization block 3 in Algorithm B is guaranteed by T E N the time index at which we exit block for the εth time (i.e., |GT, S 52), for E N. Furthermore, we recall that G. & Vws If (0.); when clear from the context, we omit j from W and use instead V w (0.) = VW(W)) Consider the SVD decomposition of the
converge; we denote
16

---

Published as a conference paper at ICLR 2025
gradient Vws (01,) = LT, V For t (T T1+1 11, we define the projected gradient as G. 2 PT, (TT, Gr, where PT. = UT, [. rT.] for a given threshold 0 < 7/h S 1, where
TT. sup sup E N: : 7mh | GT, 1l - GT. PT, (T) 1| 2 0)
(4)
obtained by the search presented in IASS Algorithm using exact truncated-S calculation. Note
that since || G, PT, (r s G 11: always, must be the case that 0 7ph Also, because
Pr, (n)G; G,, PT. in equation
well-defined. Next, let h4 - f(W,) - f[WT T++) )) and og
denote the learning rate. Then,
h = f (W ++1 - f (WT:+1)
= (W: a: (G.) - (WT..1)
$ ((,)) f (WT..1) - a.vec (G.) vec (Vf (Vf(W,))
3
G
o
= (WW) - (WTH))) - Vve (G.) vec (G,) + +
3
G
2
2 2
< ((W)) / (WT, - ar (PT:
G. + 0
|Gr|
2
(2)
3 2
< <((W.) -f (WT...) -a tr (GG) PG. (T. G.) + a
2
|
(3)
8
(W,) / (WT..)) 0z PT, (TT, G. ||
+
|G
2 3
(W.) f (WT..))
G
a
|G
2
fW,, (WT..))
04
G
2
+
B(Da))²
h, - az G 2
3(Do..)
(5)
2
where (1) follows by the assumption that f is 8-smooth function and the decent lemma (see, e.g.
Beck 2017 [Definition 5.1]. (2) follows from the identity vec (AB)* vec (B) tr((AB)TD B), in (3) we use the fact that since P, is an orthogonal projection matrix, we have P, = P, and finally, (4) follows from the assumption that the norms of the gradients are bounded by D. Rearranging equation 5 and choosing a = a, for all t 2 O, we readily obtain that,
Tai--
hT. ITAn + (T1+1 - TT)
G
(6)
2
Thus, for N E N,
N T-41-1
Σ
Σ
G 2
1
N-1
lT. II +
(T141 - T,)
<
(7)
F
TN
2
ET;
hTo
hTN
(TN To)
aTN
2TN
VTN
where R A /4M and in the second equality we choose a
AM Tw3D
Now recall that by
the definition of rT, in equation 4 we have,
||T - PT. (rT. (r. < 7hh |GT,
(10)
$ Note, that this choice a implies the guarantees of Lemma because &
AM
T
4M
Tp a af, and of is the optimized learning rate for the ith block, that was chosen in the proof of
Lemma
17

---

Published as a conference paper at ICLR 2025
Moreover, the following clearly holds for any E N.
|| G. _ = [[ (TT, G. l + |P (rT, (rT,)
(11)
P (TT, G. 1/ + |G, (r.((,,
(12)
and thus by plugging equation 10 into equation 12 at t T, for any i N, we get,
(1
rhaa)| 11 S
PT. (r. 1l
(13)
Accordingly,
Pt (TT. GT. GT. S 7hh
(T..
(14)
7/h
Recall from Subsection
3.1 that for the reversible layer,
|G |[[I -[[[[[a0)) aS) 11|
(15)
$ [ - 11
(16)
max 1 - adil G 1|
(17)
where {Ar] are the eigenvalue of S. Thus, using the fact that is positive semi-definite matrix, the learning rate is such that S Amax' where Amax is the maximal eigenvalue of S, then we get
that max; |1 a²,1 < 1, and accordingly. ||Gells < I This means that the Frobenius norm of the gradient is monotonically decreasing as a function of
Now, recall that 22 is any positive number such that 52 < -I- Tuh ||GT._.ll.. In light of equation B3] this necessarily implies that in each block i, we will execute (at least once) the low- rank optimization block (indeed, the condition ||GT, FF > 52,1 is satisfied). This, conjugated with the monotonicity property that |G that
|| GT. [l. for any t E [Ta, Tx41 - 1] and E [N], imply
N-IT
N- TT41-1
|G
GT.
1
(18)
i=0
t=
N-1 N-1Tmpl
5 1
(r.))))
(19)
7hh
TN
R
(20)
(1
Thh VTN
Accordingly, for any 0, and TN
n²
[1 12.
N T41-1
min DSISTN
|G <
Σ
|G Ss,
(21)
TN
t=T;
and thus, there exists an iteration index t E [0, T N for which,
G. | <
(22)
which, by definition, implies that Algorithm 3 achieves an e-critical point.
Lemma 3 (Convergence of low-rank optimization block), Consider the same setting and assump- tions as in Theorem? 2 Then, the time = T2 N at which Algorithm 3 Bexits block for the fth time
is finite, for any l E
Proof of Lemma 3 We follow the same notations as in the proof of Theorem 2 Recall from equa- tion that the following holds true,
2 + (We+1) S SW..) az |.
3(Da))²
(23)
2
18

---

Published as a conference paper at ICLR 2025
for any 2 0. Now, enter for the first time the low-rank optimization block of Algorithm time To = t = O, and we next show that this block converges. Fix E N. Using equation 23 and we have,
following the same arguments in equation --equation 9
Σ
c.
G.
(Wo) - ((W)) + 3D 2
+
(24)
Ta
R
(25)
VT
where R V4MBD2 and in the second equality choose some any 2 0, and
4M T3De
Accordingly, for
g²
we clearly have,
T-1
min 10
Gr
G
S.
(26)
1=0
note that T₁ is the time index which we exist the low-rank block for the first time, and the above guarantees that Tı finite. The same arguments above apply for any block exit time Te, > which concludes the proof.
namely, there exists T1 [0, T
such that GT.
C ABLATION STUDY
IMPORTANCE OF THE ADAPTIVE SUBSPACE DIMENSION
First, we note that comparing to AdaRankGrad reduces memory usage by over 25%-50% on average compared to GaLore (and LoRA), as shown in Figure 4, with respect to the layers in which the meth- ods are applied and compared on, fine-tuning tasks, and with 20% reduction in memorylstorage per the whole final fine-tuned model, as reported in Table 4.
We emphasize that this significant benefit in reducing the memory needed in training is an exclusive consequence of the adaptivity of the (gradients) subspace dimension during training (exploiting the natural phenomenon of the decrease of the dimension in the gradients- preserving the information ratio of any given predefined threshold).
To evaluate the contribution of adaptive subspace dimension solely to the model performance, we conduct the following experiment, in Table where we fix the subspace update interval 200 steps and study the adaptivity of the subspace dimension.
IMPORTANCE THE ADAPTIVE SUBSPACE UPDATE
To assess the impact of the adaptive subspace update on model performance, we conducted the following experiment, shown in Table 6 where we fix the rank to constant value of rank=4 and examine the adaptivity of the subspace updates.
Table 5: The table presents the results of an ablation experiment in which the subspace update is fixed intervals 200 optimization steps, while adaptivity in the subspace dimension remains enabled.
Model
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
GaLore (rank=4)
60.35
90.73
92.25
79.42
94.0
87.0
92.24
91.06
AdaRankGrad (Initial rank=4)
61.4
90.97
92.6
81.23
94.8
86.6
92.5
90.4
Constant subspace time-update
AdaRankGrad (Initial rank=4)
61.2
90.89
92.58
81.18
94.63
86.91
92.37
90.39
19

---

Published as a conference paper at ICLR 2025
Table 6: The table presents the results of an ablation experiment in which the selected subspaces
dimensions is fixed to r 4, the adaptivity in the subspace updating (times) remains enabled.
Model
CoLA
STS-B
MRPC
RTE
SST2
MALI
QNLI
QQP
Galoore (rank=4)
60.35
90.73
92.25
79.42
94.0
87.0
92.24
91.06
AdaRankCrad (Initial rank=4)
61.4
90.97
92.6
81.23
94.8
86.6
92.5
90.4
Constant rank
AdaRankGrad (Constant rank=4)
61.12
90.81
92.17
80.13
94.0
87.12
92.41
91.29
D EFFECTIVE RANK VS Thh ANALYSIS
0.95
0.94
-1001
0.93
0.92
0.91
06
0.9
Figure 6: The left graph presents the effective rank measured for different values of Yuu while training AdaRankGrad on the MRPC dataset. The right graph presents the effective rank measured for different values of Mhh while training AdaRankGrad on the MRPC dataset.
E GENERALIZATION TO OTHER OPTIMIZERS
It is important to clarify that AdaRankGrad is an adaptable framework applicable to any optimization method. For illustration, AdaRankGrad's application with AdaFactor is presented in Algorithm 4 and AdaRankGrad's with Proximal SGD presented in Algorithm 5 AdaRankGrad identifies
dominant now-dimensional subspace of the gradient, preserving relative information for each update. This allows any gradient-basee optimization step to run within this subspace and, before updating weights, projects the gradient-base step back into the full space. This adaptability across optimizers a core advantage of our approach.
F MEMORY SAVINGS UTILIZATION
Optimize system efficiency. Techniques such PARIS and ELSA Kim at. (2022) enable resource reallocation to concurrent tasks, enhancing GPU utilization in shared environments. Approaches such FaaSwap Yu et al. (2024) support memory sharing, facilitating smaller-scale jobs alongside training, while frameworks like ZeRO Rajbhandar et al. (019) leverage savings for parallel model validation fine-ture ane-tuning tasks. Additionally, the resource management techniques proposed in Su- perNeurons Wang et al. (2018) explore how dynamic GPU memory partitioning and job scheduling can improve utilization across multiple jobs by redistributing memory savings other pending or low-priority tasks. This allows for efficient balancing of training and inference workloads in het- erogeneous environments. Similarly, the DELTA framework Tang et al. (2022) provides a dynamic combination of tensor swapping and recomputation, enabling more efficient use of memory savings during training by supporting overlapping workloads and improving overall system throughput. In- tegration with schedulers like Kubernetes https: prioritization and efficient resource distribution, maximizing throughput. These strategies demon- strate how training efficiency can translate into broader system-level gains.
.io/ allows real-time task
20

---

Published as a conference paper at ICLR 2025
Algorithm 4 AdaRankGrad with AdaFactor step
Imput: A Layer W R= from dataset D, loss function č, information thershold 0 V/s initial rank Timi, maximal rank Tmxx, and small numbers S1, > 0, relative step sizes {pt) second moment
such that 821 = 0, regularization constants and 62, clipping threshold d. decay {P.n) J 1=1
Initialization: O.
Sample batch B <- { |1)
Compute batch gradient G. + EIa ow C (#(2,,
{For unsupervised (1 0))
while |Gell > S1 do
Block 1: Adaptive subspace selection Qe:re 4 [ Tinity Timx; 7/th)
(Approximated re-dimension subspace: Q projected matrix)
Block 2: Low-rank optimization
G, Q, G. while ||Gr|| 52 do
Projectedd gradient on the approximated re-dimension subspace)
Block 3: AdaFactor update step
a, max (ez, RMS (W(-1))
R. Bu.. 1 (1 (1-3an 32)
(G) + ealr Im
C, C, + B2.) 1, G 11 1m
V. R,C,//,,RR
U. G./VV,
U. U./ max (1, RMS (U.)
W. W.-1 - ai Q.U.
++1 Sample batch {.,, Mc 1,,1~ Compute batch gradient G, E1a ow
a
L (0(z,, 0).):
G, Q G,
end while
end while Return W.
(Exit by convergence criteria could alternatively be defined by the number of epochs]
G ADDITIONAL COMPARISON TO LORA
For completeness, we repeated the experiment from Table 2 using rank of 8, employing the hyper-
parameters specified in Table 3 of the original LoRA paper Hu et al. (2021) (note that Table used the hyperparameters reported Galore Table 4 Zhao al. (2024a)). The hyperparameters used in this experiment are detailed in Table 7 and the corresponding results are presented in Table 8
Table 7: Hyperparameter settings for different GLUE tasks using RoBERTa-Base, as used in Table 3 in Hu et al. (202T).
Parameter Task
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch Size
16
16
16
32
32
16
32
16
Epochs
30
60
30
80
25
25
80
40
21

---

Published as a conference paper at ICLR 2025
Algorithm 5 AdaRankGrad with Proximal SGD step (ProxGen Yun et al. (2021))
Input: A Layer W R= from 0 dataset D, loss function L, information thershold 0 WM initial rank Finits maximal rank Tmax, and small numbers $2 > 0, decay {pe).-1 regularization constants 61 and
62-
Initialization: O. Sample batch B <- { ::)
Compute batch gradient G. t [ SW
C (9)2,,
{For unsupervised ( 0))
while |G]]| > do
Block 1: Adaptive subspace selection
Q P2 [A,, Tinity F'ax; 77th)
(Approximated re-dimension subspace: Q projected matrix
Block 2: Moments subspaces transformation
R T - <- Q Q ift 2 1, else 0 x
M. Xm R,M,-1, if else x
(1 ff -order moment)
Block Low-rank optimization
G Q G, while ||Gr|r do
(Projected gradient the approximated ry-dimension subspace)
Block Proximal SGD update step
M. P (1 pr) G.
0++1 prox (0: ar Q.M+)
end while
end while Return W.
(Exit by convergence criteria could alternatively be defined by the number of epochs]
Table 8: Evaluation RoBERTa-Bass (LoRA) on the GLUE benchmark using the corresponding
Hyperparameter settings as presented in Table and LoRA.
we present the accu- racy results for AdaRankGrad
Method MNLI
SST2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
LoRA 87.5
95.1
89.7
63.4
93.3
90.8
86.6
91.5
AdaRankGrad
86.5
95.7
93.2
64.2
94.6
89.6
87.3
92.3
22


Reviews:
[{"review_id": "MBWHwIPeEa", "reviewer": "Reviewer_HugF", "summary": "The authors propose AdaRankGrad that learns to adaptively perform low-rank gradient updates in order to reduce the memory requirements needed to train a large language model. They base their idea on a theorem that shows that the gradients rank reach to 1 over time. Their results on datasets like CoLA, STS-B, and MRPC show that AdaRankGrad is promising at significantly reducing memory requirements while retaining competitive performance compared to GaLore and LoRA.", "questions": "Could you address the weaknesses above?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Reducing memory requirements is a crucial part of training LLMs\n- The paper is well-written and the proofs are well designed.\n- The experiments are comprehensive.", "weaknesses": "- The results of AdaRankGrad compared to GaLore and LoRA do not seem significant in Table. 2 How do you know if it is worth having yet another memory reduction method?\n- How does this method compare to BlockLLM [1]? I see a lot of similarities between the two. \n- How does the iteration speed of AdaRankGrad compare to others like LoRA, and GaLoRE. Wouldn't the method be slow given that it needs to estimate the projection matrix?\n- While Adam is popular, we shouldn't always have to use it. How would AdaRankGrad generalize to other optimizers? it seems that AdaRankGrad is limited to Adam.\n\n\n[1] https://arxiv.org/abs/2406.17296", "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": "yhKApbSlrk", "reviewer": "Reviewer_adxZ", "summary": "The paper introduces a method named AdaRankGrad, aims at enhancing the memory efficiency of fine-tuning LLMs. The authors find and formally prove that the rank of estimated gradients gradually decrease to rank one. Inspired by the phenomenon, authors proposes an online updating low-rank projections rule to help reduce memory requirements while improving model performance. The paper provides a convergence analysis and demonstrates the effectiveness of AdaRankGrad through experiments on language and biological foundation models.", "questions": "Questions\n\n1. Have there been any experiments conducted to assess the impact of this fine-tuning method on the complex capabilities of LLMs, such as MMLU and GSM8k benchmarks?\n\n2. How do you determine the key hyperparameters in AdaRankGrad, such as the information threshold η_{th}? Because the influence of η_{th} seems quite severe to effective rank according to Appendix D. \n\n3. If I understand correctly, the rank of the gradients may change. How do you ensure that AdaRankGrad can adapt to these changes without frequent manual parameter adjustments?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths\n\n1. Memory Efficiency: AdaRankGrad significantly reduces memory usage during training, which is crucial for large-scale models.\n\n2. Theoretical Foundation: The paper provides a theoretical basis for the method, leveraging the phenomenon of decreasing gradient rank.", "weaknesses": "Weaknesses\n\n1. Experiment Design: The paper's experiments, while informative, may not be entirely convincing. The use of weak benchmarks (GLUE) and the omission of key benchmarks such as MMLU (which focuses on common sense) and GSM8K (which assesses reasoning) limit the evaluation of AdaRankGrad's effectiveness on complex capabilities.", "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": "a5pJwAxuo1", "reviewer": "Reviewer_xMT8", "summary": "The paper introduces AdaRankGrad, a novel method aimed at optimizing memory efficiency and enhancing convergence rates and performance during the training or fine-tuning of large language models (LLMs). It addresses the issue of high memory and computational demands associated with the increasing size of model weights and optimizer states. AdaRankGrad leverages the observed phenomenon that the rank of layer gradients decreases over time, approaching rank one, to dynamically reduce the rank of gradients during Adam optimization. This is achieved through an efficient online-updating low-rank projections rule and a randomized-SVD scheme for identifying the projection matrix. The method supports full-parameter fine-tuning while significantly reducing memory usage compared to existing state-of-the-art methods.", "questions": "1. The method proposed in this paper requires the computation of the subspace in each iteration of the model training, how much time overhead does this operation incur compared to full parameter training and compared to Galore?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a theoretical foundation for the observation that the rank of the gradients in LLMs tends to decrease over the course of training, contributing to the understanding of gradient dynamics in deep learning.\n2. In contrast to GaLore, which uses fixed time intervals for subspace updating and ignores momentum subspace alignment, the authors propose adaptive subspace selection and momentum subspace transformation, and experimentally illustrate the effectiveness of the two improvements.\n3. The authors propose to use SSRF to replace the SVD used in Galore, which speeds up the search for subspaces.", "weaknesses": "1. According to Table 1, the main difference between this paper and Galore is the adaptive subspace dimension and adaptive subspace update. We consider that the existing experiments can illustrate the effectiveness of adaptive subspace updating but cannot illustrate the necessity of adaptive subspace dimension. Therefore, it is recommended that the authors include additional ablation studies focusing on the adaptive dimensionality to substantiate its importance.", "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": "iX6VGY1fvv", "reviewer": "Reviewer_x7pw", "summary": "The paper studies an adaptive rank selection method for low-rank gradient update. Compared with previous low-rank gradient update methods like Galore, the proposed method involves an adaptively rank-reducing schedule to further reduce training memory. Experimental results and theoretical guarantees provide solid support for the claim in this paper.", "questions": "Based on the previously mentioned weakness, I have the following questions:\n- I also observed similar GLUE fine-tuning results in the Galore paper, but I didn't find the reason for the gap compared with the original LoRA paper. I'm wondering if the author could provide a detailed description of their LoRA implementation and experimental setup, including hyperparameters, training procedure, and evaluation method, and compare the accuracy here with the same setup as the original LoRA paper, which is a more widely used setup.\n- For training efficiency, I know Galore slows down the training process due to the subspace transform process with a subspace frequency of 200. I'm worried the proposed method may further increase the training time due to the frequent update of subspace. Could you provide the concrete timing comparison between AdaRandGrad, LoRA and Galore for some of the GLUE tasks?\n- For the training efficiency, Could you count the wall-clock time for block 1 and block 2 in Algorithm 3, compared with Block 3. I'm wondering about how much additional time is needed for the added parts.\n- If possible, could you provide an ablation study showing how different frequencies of subspce selection affect the trade-off between the computational efficiency and performance? Especially, a comparison between the Galore and AdaGradRank with both a 200 subspace frequency is necessary to make the comparison fair. \n\nI'm happy to further consider my score if the authors could provide a further comparison for the wall-clock time and comparison with Galore under similar subspace frequency.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper provides solid theoretical and numerical analysis for the design of their adaptive rank schedule. The method is effective for reducing the training memory while maintaining the training performance. The writing of the paper is clear.", "weaknesses": "My main concern is about the experiment:\n- For the result in Table 2, the results of LoRA tuning seem far worse (2-3%) compared with the standard baselines. For example, the LoRA fine-tuning result for SST2 task should be around 95% based on the results reported in the original LORA paper. This result has also been validated in many other related works. I'm wondering if there exists any difference in the experiment setup that leads to this degradation.\n- Training Efficiency: I noticed that the authors mentioned the issue of training time in the paper. However, I found no numerical results showing the difference in training time compared with the LoRA or Galore method. Even though the paper provides an approximation method to calculate the SVD. I assume this adaptive subspace selection process happened every iteration may still slow down the training time significantly. Also, I'm wondering if it's possible to avoid doing this subspace selection each iteration.", "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": "X1lGBW6T0C", "reviewer": "Reviewer_TXWw", "summary": "The paper presents a theoretical proof that the effective rank of the gradient matrix converges to 1. Building on this, it proposes dynamically reducing the gradient matrix rank as pre-training or fine-tuning progresses, thereby adaptively lowering memory consumption. Recognizing that SVD decomposition is computationally intensive, the paper introduces a novel randomized algorithm to efficiently approximate this decomposition. Table 1 highlights the advantages of the proposed method over LoRA and GaLore across several dimensions. Additionally, this approach is applicable to both pre-training and fine-tuning.", "questions": "1. What do the memory savings look like when factoring in the increased running time?\n2. Can the savings be fully utilized by other jobs? If so, how?\n3. Why is $\\epsilon$ needed in defining the $(\\eta, \\epsilon)$-approximation?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a theoretical proof that the effective rank of the gradient matrix converges to 1. Building on this, it proposes dynamically reducing the gradient matrix rank as pre-training or fine-tuning progresses, thereby adaptively lowering memory consumption. Recognizing that SVD decomposition is computationally intensive, the paper introduces a novel randomized algorithm to efficiently approximate this decomposition. Additionally, this approach is applicable to both pre-training and fine-tuning.", "weaknesses": "Overall, this paper could be a significant algorithmic contribution, with the caveat for some clarifications on the theory and experiments. Given these clarifications in an author response, I would be willing to increase the score. \n\nFeedback on specific claims and points for clarification:\n\n1. The memory savings reported in Table 2 may be overstated. Since SSRF runs continuously, the runtime could exceed that of vanilla LoRA or pre-training alone. As a result, while less memory is allocated, it may be held for a longer duration. Memory savings should consider increased runtime since allocating 20% less memory for 20% longer results in only a 4% net saving.\n\n2. It may be challenging to realize the claimed memory savings in practice. Figure 4 shows significant fluctuations in memory savings throughout fine-tuning or pre-training, making it difficult to allocate this saved memory to other concurrent tasks. A smoother, more predictable savings trend would simplify practical implementation.\n\n3. The need for two parameters, $\\eta$ and $\\epsilon$, to define $(\\eta, \\epsilon)$-approximation is unclear. Would a single $\\eta$ suffice?\n\n**Suggestions that do not impact the score:**\n\n- Consider removing the sentence on line 064, “still both of the …” since the paper does not seem to offer a solution for improving inference efficiency.\n  \n- In Algorithm 1 (SSRF), where $A$ is $n \\times m$ and $\\Omega$ is $k \\times n$, the computation $A \\Omega$ seems undefined.\n  \n- For Figure 5, should the x-axis be labeled as 11700, 23400, 35100 (for epochs 1, 2, 3)?", "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)