AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning

## Abstract

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. Various techniques have been developed to tackle these challenges, 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 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.

###### Lemma #1 #1 \mathbf{#1} #1 .

## 1 Introduction

Large language models (LLMs) have gained significant attention due to their impressive ability to handle various tasks, such as dialogue-based 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 [, , ].

To tackle the challenges associated with LLM fine-tuning, researchers have developed various optimization techniques to reduce memory usage during model training. One key approach that has emerged is Parameter-efficient fine-tuning (PEFT) [], which enables the adaptation 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) algorithm, introduced by []. LoRA reparameterizes a weight matrix ${\bf W}\in\mathbb{R}^{m\times n}$ into ${\bf W}={\bf W}_{0}+{\bf B}{\bf A}$, where ${\bf W}_{0}$ is a frozen full-rank matrix, and ${\bf B}\in\mathbb{R}^{m\times r}$ and ${\bf A}\in\mathbb{R}^{r\times n}$ are low-rank adaptors. Since $r\ll\min(m,n)$, the low-rank adaptors ${\bf A}$ and ${\bf B}$ require fewer trainable parameters, reducing memory usage. LoRA has been widely adopted for fine-tuning, with several variants emerging, including LoRA+, which uses different learning rates for the two matrices [], Adaptive LoRA, which adapts the rank of the matrices during training [], and Sparse LoRA, which introduces sparsity to the matrices to further reduce computational cost []. These methods have been demonstrated to enhance the efficiency and performance of LLM fine-tuning for various tasks.

Despite its advantages, recent research has identified some limitations of low-rank reparameterization. For example, LoRA may not achieve the same performance levels as full-rank fine-tuning [] and might require initial full-rank model training as a warm-up before effectively utilizing the low-rank subspace []. These issues may stem from the fact that optimal weight matrices are not inherently low-rank or from changes in gradient training dynamics introduced by the reparameterization. In addition, LoRA keeps the tuning layer shapes in the base model static without dynamic adjustments. Another approach by [] dynamically adjusts tuning parameters during training, and [] gradually reduces tuning parameters. Still, neither of these methods improves the inference efficiency of the fine-tuned model.

Until recently, the pre-training of large language models (LLMs) has been primarily limited to corporations 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, we take the following as an example. Training a Mistral $7$B model from scratch poses substantial memory challenges. Given its $7$ billion parameters, a 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 [] introduced a training strategy called GaLore that enables full-parameter learning while being more memory-efficient than traditional low-rank adaptation methods such as LoRA. The core idea behind GaLore is to exploit the slowly changing low-rank structure of the gradient ${\bf G}\in\mathbb{R}^{n\times m}$ of the weight matrix ${\bf W}$, rather than approximating the weight matrix itself as low-rank. GaLore significantly improved memory efficiency, reducing optimizer state memory usage by up to 65.5%. The following noticeable variant is Q-GaLore [], which combines low-rank gradient projection with INT4 quantization to further reduce memory usage, and an additional parallel variant would be ReLoRA [] employed in pre-training-by periodically updating ${\bf W}_{0}$ using previously learned low-rank adaptors. We found that Galore to be suboptimal since it arbitrarily requires pre-defining a fixed low-rank size for the gradient projection/low-rank-approximation, while gradients rank gradually diminish during training down to rank one. Additionally, GaLore uses a fixed window size for the number of iterations 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 at any update 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 a 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, [] showed that the rank of the gradient satisfies $\text{rank}\left({\bf G}^{n\times m}\right)<\min\{n,m\}/2$. Second, as highlighted in previous research [, , ], 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 on accuracy varies significantly. This indicates that the optimization steps have less influence on the model’s performance for certain layers, depending on their depth and architecture type. Noise in these layers has a 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).

Building upon these two insights and to address the limitations of the LoRA variants, we propose a method that enables full-parameter learning while dramatically reducing memory requirements and computational complexity through adaptive low-rank gradient projections during the Adam update step.
For each gradient tensor ${\bf G}_{t}^{j}\in\mathbb{R}^{r\times n}$ at layer $j\in[L]$ and iteration $t$, AdaRankGrad efficiently identifies a unique set of significant projection directions (subspace) ${\bf P}_{t}^{j}\in\mathbb{R}^{r_{j}\times n}$ along which the gradient ${\bf G}_{t}^{j}$ exhibits

Figure: Figure 1: The illustration shows how AdaRankGard 3 is trained. First, the gradients ${\bf G}_{t}$ are projected into a 3D space (in this example), represented as $\hat{{\bf G}}_{t}^{3\times m}={\bf P}^{3\times n}_{t}{\bf G}_{t}^{n\times m}$. As convergence occurs, the gradient’s dimension decreases to a 2D space and then to a 1D space. This dimensionality reduction indicates convergence while efficiently using memory.
Refer to caption: extracted/5948913/Screenshot_2024_09_25_224219.png

the largest changes, where $r_{t}^{j}$ is the lowest possible rank to still maintain a predefined information fraction (given threshold) relative to the original, non-projected gradient. Practically, ${\bf P}_{t}^{j}{\bf G}_{t}^{j}$ is the low-projected-gradient and ${{\bf P}_{t}^{j}}^{\top}{\bf P}_{t}^{j}{\bf G}_{t}^{j}$ is the best low-rank approximation the of the gradients ${\bf G}_{t}^{j}$ that embodies the required fraction of its information. The projections ${\bf P}_{t}^{j}$ are being adaptively updated throughout training (based on the convergence criteria of the gradients on the projected subspace), where their rank $r^{j}_{t}$ is dictated by preserving the given information threshold. The method ensures: (1) The method determines the optimal projection dimension for each layer’s gradient tensor independently, adjusting dynamically throughout training. This rank adjustment leverages property we prove that the effective dimensionality of the full gradients gradually decreases over time, allowing updates to be performed in a lower-dimensional projection space, thereby reducing memory usage. (2) The projection matrix for each layer’s gradients is updated based on a convergence criterion within their respective subspace. This ensures updates occur precisely when needed, avoiding premature or delayed transitions between subspaces, resulting in faster overall convergence.

**Table 1: Comparison between AdaRankGrad, GaLore, and LoRA. Assume ${\bf W}\in\mathbb{R}^{n\times m}(n\geq m)$, constant rank $r$, and adaptive-rank $r_{\text{adap}}$ (with initial-rank $r_{\text{init}}=r$).**
|  | AdaRankGrad | GaLore | LoRA |
| --- | --- | --- | --- |
| Weights | $nm$ | $nm$ | $nm+nr+mr$ |
| Optim States ($r_{\text{adap}}<r$) | $nr_{\text{adap}}+2mr_{\text{adap}}$ | $nr+2mr$ | $2nr+2mr$ |
| Multi-Subspace | $\checkmark$ | $\checkmark$ | $\mathbf{x}$ |
| Adaptive-Subspace-Dimension | $\checkmark$ | $\mathbf{x}$ | $\mathbf{x}$ |
| Adaptive-Subspace-Updates | $\checkmark$ | $\mathbf{x}$ | $\mathbf{x}$ |
| Pre-Training | $\checkmark$ | $\checkmark$ | $\mathbf{x}$ |
| Fine-Tuning | $\checkmark$ | $\checkmark$ | $\checkmark$ |

## 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 optimization algorithms [, ]. One common approach is quantization, which helps decrease the memory footprint of optimizer states []. Additionally, recent advancements have suggested reducing the memory used by weight gradients by integrating the backward operation with the optimizer update [, ]. This characteristic has been leveraged to reduce memory usage during training processes [, , ].

##### Low-rank gradient optimization.

The phenomenon of low-rank gradients naturally arises during the training of neural networks, a subject that has been extensively examined both theoretically and practically, e.g., [, , ]. This characteristic low-rank structure gradients has been leveraged to reduce memory usage during training processes [, , ], and results in a reduced computational complexity as compared to standard gradient descent methods.

##### Adam optimization.

Arguably, among the most popular optimization methods used for training large language models (LLMs) are the *Adam optimizer* [] and its variant, *AdamW* [], which incorporates weight decay for regularization. However, it is well-established that Adam optimization has higher memory complexity compared to other optimization alternatives. To illustrate this, let us briefly review how the Adam optimization algorithm operates. First, we need to establish some notation. Consider a neural network denoted as $\Phi(\cdot;\bm{\theta})$, which consists of $L$ layers and is parameterized by $\bm{\theta}\triangleq\left[{\bf W}_{1}^{d_{1}\times d_{0}},\ldots,{\bf W}_{L-1
}^{d_{L-1}\times d_{L-2}},{\bf W}_{L}^{d_{L}\times d_{0}^{L-1}}\right]$. Here, ${\bf W}_{i}$ represents the weights tensor parameters associated with the $i$-th layer, for $i\in[L]$. In the following, let $t\in\mathbb{N}$ represent the $t$-th step of the Adam optimization algorithm. Then, we recall that the single update step in Adam is given by,

$$ $\textbf{Adam}\left\{\begin{aligned} &\quad{\bf G}_{t}=\nabla\Phi\left(\bm{ \theta}_{t-1}\right),\\ &\quad{\bf M}_{t}=\beta_{1}{\bf M}_{t-1}+\left(1-\beta_{1}\right){\bf G}_{t}, \\ &\quad{\bf V}_{t}=\beta_{2}{\bf V}_{t-1}+\left(1-\beta_{2}\right){\bf G}_{t}^{ 2},\\ &\quad\hat{{\bf M}}_{t}={\bf M}_{t}/\left(1-\beta_{1}^{t}\right),\\ &\quad\hat{{\bf V}}_{t}={\bf V}_{t}/\left(1-\beta_{2}^{t}\right),\\ &\quad\bm{\theta}_{t}=\bm{\theta}_{t-1}-\alpha\hat{{\bf M}}_{t}/\left(\sqrt{ \hat{{\bf V}}_{t}}+\epsilon\right).\end{aligned}\right.$ $$

Specifically, at time step $t$, ${\bf G}_{t}$ denotes the backpropagated gradient matrix, i.e., $\nabla\Phi\left(\bm{\theta}_{t-1}\right)$. The exponentially weighted moving averages of the first and second moments are denoted by ${\bf M}_{t}$ and ${\bf V}_{t}$, respectively, with their bias-corrected counterparts given by $\hat{{\bf M}}_{t}$ and $\hat{{\bf V}}_{t}$. The AdamW optimizer updates the model parameters at step $t$ according to the rule, $\bm{\theta}_{t}=\bm{\theta}_{t-1}-\alpha\left(\frac{\hat{{\bf M}}_{t}}{\sqrt{
\hat{{\bf V}}_{t}}+\epsilon}+\lambda\bm{\theta}_{t-1}\right)$,
where $\lambda\geq 0$ is the weight decay rate (for Adam $\lambda=0$), and all operations are performed element-wise. In this equation, $\beta_{1}$ and $\beta_{2}$ control the decay rates for the moving averages of the moments, $\alpha$ is the learning rate, and $\epsilon$ is a small constant to avoid division by zero. Notably, since Adam/W requires storing both ${\bf M}_{t}$ and ${\bf V}_{t}$ at each time step, it incurs an additional memory footprint of $2mn$.

While existing approaches [, , ] focus on low-rank approximations of the first and second moments with the goal of reducing memory requirements, 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

Figure: Figure 2: The figure illustrates the exponential decay of eigenvalues in the MLP layer’s gradient, at the first iteration of fine-tuning RoBERTa-Base [] model, on the MRPC task, from GLUE []. Notably, the red line indicates that 50% of the gradient information (in terms of squared norm ratio) is captured by the first eigenvalue, while the green line shows that 90% is contained within the first two eigenvalues.
Refer to caption: extracted/5948913/Colab_Screenshot_2024_09_21_010917_2.png

As mentioned in the introduction, a few recent empirical results (e.g., [, , ]), 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 2, where it is evident that the squared norm of the gradient’s singular values decay to zero exponentially fast.

As hinted above, this phenomenon is only true in the approximate sense; roughly speaking, only very few eigenvalues hold almost all the information captured by the gradient. Accordingly, a low-rank matrix approximates the underlying gradient up to a negligible 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-gradient updates is possible. To make the above discussion precise, consider the following definition for approximate low-rank matrices.

###### Definition 1 (Approximate low-rank matrix) .

A matrix ${\bf A}\in\mathbb{R}^{n\times m}$ is called $(\eta,\varepsilon)$-approximately rank-$r$, if there exist $\eta\in[0,1)$, $\varepsilon>0$, and a matrix ${\bf A}_{\mathsf{app},r}\in\mathbb{R}^{n\times m}$ with $\mathsf{rank}({\bf A}_{\mathsf{app},r})=r$ and $r<\min\{n,m\}$, such that,

$$ $\left\|{\bf A}-{\bf A}_{\mathsf{app},r}\right\|_{F}\leq\eta\cdot\|{\bf A}\|_{F }+\varepsilon.$ (1) $$

As it turns out, it can be shown (see, e.g., []) that the optimal ${\bf A}_{\mathsf{app},r}$ minimizing the approximation error in the left-hand-side of (1), can be obtained by applying an SVD on ${\bf A}$ and then retaining only the top $r$ singular values and their corresponding singular vectors. Mathematically, we have ${\bf A}_{\mathsf{app},r}=\sum_{i=1}^{r}\sigma_{i}{\bf u}_{i}{\bf v}_{i}^{\top}$,
where $\{\sigma_{i}\}_{i}$ are the singular values of ${\bf A}$, and $\{{\bf u}_{i}\}_{i}$ and $\{{\bf v}_{i}\}$ are the corresponding left and right singular vectors, respectively. The approximation error is in turn given by $\left\|{\bf A}-{\bf A}_{\mathsf{app},r}\right\|_{F}^{2}=\sum_{i=r+1}^{\min\{m,
n\}}\sigma_{i}^{2}$. This construction gives an $(\eta_{\bf A},0)$-approximately rank-$r$ matrix, with the minimal $\eta_{\bf A}$ possible.

Recently in [], the structure of the gradient for a wide family of nonlinear networks known as “reversible networks” [] was studied,(^1^11It can be shown that this family includes many different kinds of layers, such as, linear layers (MLP and Conv.), and (leaky) ReLU non-linearity.) defined as follows.

###### Definition 2 .

(Reversibility [])
Layer $\ell\in[L]$ is reversible if there is a ${\bf G}_{\ell}(\bm{x};\bm{\theta})\in\mathbb{R}^{n_{\ell}\times n_{\ell-1}}$ such that
the gradient after nonlinearity satisfies $\tilde{\bm{g}}_{\ell-1}={\bf G}_{\ell}^{\top}(\bm{x};\bm{\theta}){\bf P}_{\ell
}^{\top}(\bm{x};\bm{\theta})\tilde{\bm{g}}_{\ell}$, for some matrix ${\bf P}_{\ell}(\bm{x};\bm{\theta})\in\mathbb{R}^{n_{\ell}\times n_{\ell}}$. A network is reversible if all of its layers are reversible.

For simplicity of notation, we use ${\bf G}_{t}^{\ell}$ to denote $\left[{\bf G}_{\ell}(\bm{x};\bm{\theta})\right]_{t}$, where $t\in\mathbb{N}$ is the iteration index in the optimization process. Furthermore, when it is clear from the context, we omit the layer index $\ell$ from our notations. Assuming reversibility and SGD weight update (i.e., ${\bf W}_{t}={\bf W}_{t-1}+\alpha{\bf G}_{t-1}$), it is shown in [] that for both $\ell_{2}$ and cross entropy losses, the gradient is of the structure form ${\bf G}=\frac{1}{N}\sum_{i=1}^{N}\left({\bf A}_{i}-{\bf B}_{i}W{\bf C}_{i}\right)$, where $N$ is the batch size, $\{{\bf A}_{i}\}_{i=1}^{N}$ are input-dependent matrices, and $\{{\bf B}_{i},{\bf C}_{i}\}_{i=1}^{N}$ are certain positive semi-definite (PSD) matrices.
Furthermore, it was proven that if the gradient ${\bf G}_{t}$, has the above structure for all $t\geq\mathsf{t}_{0}$, for some $\mathsf{t}_{0}\in\mathbb{N}$, then,
the stable rank $\operatorname{sr}\left({\bf G}_{t}\right)\triangleq\frac{\|{\bf G}_{t}\|_{F}}{
\|{\bf G}_{t}\|_{2}}$ satisfies,

$$ $\operatorname{sr}\left({\bf G}_{t}\right)\leq\operatorname{sr}\left({\bf G}_{ \mathsf{t}_{0}}^{\|}\right)+\left(\frac{1-\eta\lambda_{2}}{1-\eta\lambda_{1}} \right)^{2\left(t-\mathsf{t}_{0}\right)}\left.\left\|{\bf G}_{\mathsf{t}_{0}}- {\bf G}_{\mathsf{t}_{0}}^{\|}\right\|_{F}^{2}\middle/\left\|{\bf G}_{\mathsf{t }_{0}}^{\|}\right\|_{2}^{2},\right.$ $$

where ${\bf S}\triangleq\frac{1}{N}\sum_{i=1}^{N}{\bf C}_{i}\otimes{\bf B}_{i}$, $\lambda_{1}<\lambda_{2}$ denote its two smallest distinct eigenvalues, and ${\bf G}_{\mathsf{t}_{0}}^{\|}$ is the projection of ${\bf G}_{\mathsf{t}_{0}}$ onto the minimal eigenspace $\mathcal{V}_{1}$ of ${\bf S}$ that corresponds to $\lambda_{1}$. Accordingly, as $t\to\infty$, we get that the final stable rank is upper bounded by $\operatorname{sr}({\bf G}_{\mathsf{t}_{0}}^{\|})$. Under the same gradient structure assumption and for the vanilla settings of the SGD weights update, we were able to prove the following stronger result. We prove that the approximated stable rank of the gradients approach one as the training process progresses. To state this result, we need to make a few notations. Let ${\bf G}_{t}={\bf U}_{t}\Sigma_{t}{\bf V}_{t}^{\top}$ be the SVD decomposition of ${\bf G}_{t}$, and let ${\bf P}_{t}(\ell,r)={\bf U}[:,\ell:r]_{t}{\bf U}[:,\ell:r]_{t}^{\top}$ be the corresponding projection matrix. When clear for the context, we omit the index $\ell$ and use ${\bf P}_{t}(r)\equiv{\bf P}_{t}(\ell=1,r)$. We have the following result.

Figure: Figure 3: The figure presents the effective rank (see Section 4) measured after every $100$ update steps on the RTE dataset, from GLUE [].
Refer to caption: extracted/5948913/rank_fig.png

###### Lemma 2 (Asymptotically rank-one) .

Given a reversible neural network and using the vanilla setting of SGD for weight update. Then,

$$ $\kappa(t)\triangleq\frac{\left\|{\bf G}_{t}-{\bf P}_{t}(1){\bf G}_{t}\right\|_ {F}^{2}}{\left\|{\bf G}_{t}\right\|_{F}^{2}}\leq O(C^{-t}),$ $$

for some constant $C>1$.

The above result implies that ${\bf G}_{t}$ approaches its rank-one approximation ${\bf P}_{t}(1){\bf G}_{t},$ as the iteration number increases, namely, ${\bf G}_{t}$ becomes rank-one.
The proof of Lemma 2 is relegated to Section A.
Finally, in Fig. 3 and Fig. 4, we demonstrate that for a large language model (RoBERTa-base []), which contains also non reversible layers, the rank decay evolves as a function of the number of update steps, in a fine-tuning task.

###### Definition 1 (Approximate low-rank matrix) .

###### Definition 2 .

###### Lemma 2 (Asymptotically rank-one) .

### 3.2 Adaptive-Low-Rank 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 2. The most natural but computationally expensive way to search for the subspace on which a gradient has the highest variance is by applying the truncated SVD [], 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 to 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 ${\bf A}\in\mathbb{R}^{n\times m}$. Finding its “best” low-rank approximation can be framed as the following optimization problem $\min_{{{\bf Q}},{\bf U}}\left\|{\bf A}-{{\bf Q}}{\bf U}^{\top}\right\|_{F}^{2}$, where ${{\bf Q}}\in\mathbb{R}^{n\times r}$ and ${\bf U}\in\mathbb{R}^{m\times r}$. As discussed above, the matrix ${\bf A}_{\mathsf{app},r}={{\bf Q}}{\bf U}^{\top}$ represents the $(\eta,\varepsilon)$ rank-$r$ approximation of ${\bf A}$. Computing full SVD for large matrices is computationally intensive and memory-demanding. To address these challenges, we detail Algorithm 1 [], an efficient technique for producing a “good” proxy for the optimal low-rank approximation of ${\bf A}$. Algorithm 1 is designed to solve the optimization problem
$\arg\min_{\mathbf{Q}\in\mathbb{R}^{n\times r}}\left\|{\bf A}-\mathbf{QQ}^{\top
}{\bf A}\right\|_{F},$
and approximates the matrix ${\bf A}$ as ${\bf A}_{\mathsf{app},r}\approx\mathbf{QQ}^{\top}{\bf A}.$
The computational complexity of the SSRF algorithm is dominated by two main operations: the matrix multiplication ${\bf A}\Omega$, whose computational complexity is $O(mnk)$, and the QR decomposition of the resulting matrix ${\bf Y}$, requiring $O(mk^{2})$. Therefore, the total complexity is $O(mnk+mk^{2})$. In scenarios where $k\ll m,n$, this complexity simplifies to $O(mnk)$, which is significantly more efficient than using SVD whose complexity is $O(\min(mn^{2},m^{2}n))$. This low-complexity makes the SSRF preferable for extracting leading singular vectors in large-scale data settings.

Figure: Algorithm 1 Subspace selection via randomized range finder (SSRF)

#### 3.2.2 Fast adaptive low-rank approximation

It is clear from Definition 1 that a larger value of the rank $r$ increases the computational operations and induces a higher memory footprint, while a smaller value of $r$ will result in a non-negligible approximation error. To mitigate this, we suggest Algorithm 2, an adaptive procedure for rank selection that adjusts the value of $r$ to preserve at least a certain fraction of the information in ${\bf A}$. Specifically, at start we set $r_{0}=r_{\text{min}}$, and then find the best ${\bf A}_{\mathsf{app},r}\approx\mathbf{QQ}^{\top}{\bf A}$ using SSRF. At each time step $t$, given the current rank $r_{t}$, we compute the approximation error as,

Figure: Algorithm 2 Information-based adaptive subspace selection (IASS)

$$ $\displaystyle\eta_{t}=\frac{\left\|{\bf A}-{{\bf A}}_{\text{app,}{r_{t}}} \right\|_{F}^{2}}{\|{\bf A}\|_{F}^{2}}=\frac{\sum_{i=r_{t}+1}^{n}\sigma_{i}({ \bf A})^{2}}{\sum_{i=1}^{n}\sigma_{i}({\bf A})^{2}}.$ (2) $$

We then apply binary search until we find the maximum $r$ for which the corresponding condition still holds: $\eta<\eta_{\text{th}}$. It is important to note that, from the outset, we know a priori that the rank of the approximated gradients is inherently low. Consequently, the initial value of $r_{\text{max}}\ll\min\{m,n\}$ is small, and the number of iterations required for the search is correspondingly minimal, specifically $O(\log(r_{\text{max}}-r_{\text{min}})),$ obviously $(r_{\text{max}}-r_{\text{min}})\ll\min\{m,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 3.
To describe AdaRankGrad update step rule, we need to establish a few important notations first. For $t\in\mathbb{N}$, let $\rho_{t}:\mathbb{R}^{m\times n}\to\mathbb{R}^{m\times n}$ be an entry-wise gradient update rule (e.g., Adam, AdamW, AdaFactor, etc.). For $t\in\mathbb{N}$ and layer $j\in[L]$, recall the SVD of the gradient ${\bf G}^{j}_{t}={\bf U}^{j}_{t}\Sigma^{j}_{t}({\bf V}^{j}_{t})^{\top}$, and further define the projection matrices ${\bf P}^{j}_{t}(r_{t}^{j})=\text{SSRF}({\bf U}^{j}_{t}[:,:r_{t}^{j}]^{\top}),{
{\bf R}}^{j}_{t}(r_{t}^{j})=\text{SSRF}\left({\bf V}^{j}_{t}[:,:r_{t}^{j}]\right)$, where for a given threshold $0<\eta_{\text{th}}\leq 1$, we let $r_{t}^{j}\triangleq\displaystyle\mathop{\mbox{\rm sup}}\{r\in\mathbb{N}:\eta_{
\text{th}}\cdot\|{\bf G}^{j}_{t}\|_{F}^{2}-||{\bf G}^{j}_{t}-{\bf P}^{j}_{t}(r
){\bf G}^{j}_{t}||_{F}^{2}\geq 0\}$, computed using Algorithm 2. Next, for $j\in[L]$, and some $\varsigma>0$, we let $\{\mathsf{T}^{j}_{\ell}\}_{\ell\geq 0}$, with $\mathsf{T}^{j}_{0}=0$, denote the monotone sequence of integers, for which $\|({\bf P}^{j}_{\mathsf{T}_{i}})^{\top}(r^{j}_{\mathsf{T}^{j}_{i}}){\bf G}^{j}
_{\mathsf{T}^{j}_{i}}{{\bf R}}_{\mathsf{T}^{j}_{i}}^{j}(r^{j}_{\mathsf{T}^{j}_
{i}})\|_{F}\leq\varsigma$; under certain conditions that we list below, it is shown in [][Thm. 3.4] (as well as in [][Theorem 3.8], for reversible layers) that $\mathsf{T}_{i}$’s are finite.
Then, for $i\in\mathbb{N}$, and $t\in[\mathsf{T}_{i}+1,\mathsf{T}_{i+1}]$, the AdaRankGrad weight update is given by,

$$ ${\bf W}_{t}^{j}={\bf W}_{\mathsf{T}_{i}}^{j}+\sum_{\ell=\mathsf{T}_{i}+1}^{t} \eta_{t}{\bf P}^{j}_{\mathsf{T}_{i}}(r^{j}_{\mathsf{T}_{i}})\rho_{\ell}\left(( {\bf P}^{j}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}))^{\top}{\bf G}^{j}_{\ell}{{ \bf R}}^{j}_{\mathsf{T}^{j}_{i}}(r^{j}_{\mathsf{T}^{j}_{i}})\right)({{\bf R}}^ {j}_{\mathsf{T}_{i}}(r^{j}_{\mathsf{T}^{j}_{i}}))^{\top},$ $$

for $j\in[L]$, where $\eta_{t}$ is the learning rate at time $t$ adapted by $\rho_{t}$, and ${\bf W}_{\mathsf{T}_{0}}={\bf W}_{0}$ is a pre-trained model, in a case of fine-tuning task, or a given weights initialization, in case of pre-training task. In the sequel, we let $\eta\triangleq\eta_{0}$ 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$.

Our algorithm comprises four main blocks, all contained within an outer loop that terminates once we 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 (2).
- •
Block 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, as 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 3: We tune the parameters during the low-rank update step. The stopping condition triggering the update of the orthogonal projection matrix is based on the convergence of the projected gradient onto the subspace.
- •
Block 4: The actual pre-trained model parameters are updated, using the low-dimensional and memory-efficient projected gradients on the selected subspace.

Figure: Algorithm 3 Adaptive low-rank and moments gradient (AdaRankGrad)

###### Theorem 3 (Convergence of Algorithm 3 ) .

For a loss function $\mathcal{L},$ and given architecture $\Phi$, suppose that the compositions of $f\equiv\mathcal{L}\left(\Phi(\cdot)\right)$ is $\beta$-smooth non-convex function that are bounded by some $M\in\mathbb{R}_{+}$. Let ${\bf G}_{t}^{j}$ denote the gradient matrix w.r.t. the $j$th reversible layer at time $t\in\mathbb{N}$. Assume that $\|{\bf G}_{t}^{j}\|_{F}\leq D$, for all $j\in[L]$ and $t\in\mathbb{N}$, where $D\in\mathbb{R}_{+}$. Then, for any $\varepsilon>0$, there exists $\mathsf{C}\in\mathbb{R}_{+}$ such that for all $\mathsf{T}_{N}>\frac{\mathsf{C}}{\varepsilon^{2}}$, it holds that $\frac{1}{\mathsf{T}_{N}L}\sum_{j=1}^{L}\sum_{i=0}^{N-1}\sum_{t=\mathsf{T}_{i}}
^{\mathsf{T}_{i+1}-1}\left\|{\bf G}_{t}^{j}\right\|_{F}^{2}\leq\varepsilon$. In particular, Algorithm 3, with vanilla SGD weight update(^2^22We focus on SGD for the simplicity (as is standard practice in related literature, e.g., []).), $\varsigma_{2}\triangleq\varsigma_{2,i}=\sqrt{1-\eta_{\text{th}}}\cdot\left\|{
\bf G}_{\mathsf{T}_{i-1}}\right\|_{F}^{2}$, where $i$ is the Block $3$ entry counter, and learning rate $\alpha<2/\lambda_{\max}$, achieves an $\varepsilon$-critical point,(^3^33Also known as $\varepsilon$-stationary, see, e.g., [].) i.e., $\left\|{\bf G}_{t}^{j}\right\|_{F}^{2}\leq\varepsilon$, for some $t\in\mathbb{N}$, and any $j\in[L]$.

The proof of Theorem 3 can be found in Section A. A few important comments are in order. First, to reduce memory usage, we apply in Algorithm 3 a per-layer weight update during backpropagation, as proposed by recent works, see, e.g., []. This is in contrast to 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., [, , ], and as so allows to 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 4-bit projected gradient updates, as introduced in Q-GaLore [].

If, after fine-tuning, one wishes to create an adapter (i.e., 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 $\Delta\triangleq{\bf W}_{\text{Fine-Tuned}}-{\bf W}_{\text{Pretrained}}$, where ${\bf W}_{\text{Fine-Tuned}}$ is the model weight at the end of the process, and ${\bf W}_{\text{Pretrained}}$ is the original model weight. Then, we find the $r_{\text{Adaptor}}\triangleq\mathsf{rank}(\Delta)$, using some matrix ranking algorithm, and finally, we solve
$\min_{{\bf A}\in\mathbb{R}^{n\times r_{\text{Adaptor}}},{\bf B}\in\mathbb{R}^{
r_{\text{Adaptor}}\times m}}\left\|\Delta-{\bf A}{\bf B}\right\|_{F}^{2}$, using any optimization algorithm (e.g., gradient descent). We note that any solution to this matrix factorization optimization problem is well-known as a global optimum [].

###### Theorem 3 (Convergence of Algorithm 3 ) .

## 4 Experiments

In this section, we test the performance of our algorithm on real-world datasets. Before discussing the setup we rely on in our experiments, we define four measures for memory usage reduction in the rank-adaptive projection matrices. Specifically, for the $j$th layer, we define *the effective layer-gradient-rank*, by,
$\mathcal{R}^{j}_{\text{adap}}\triangleq\frac{\sum_{t=0}^{\mathsf{T}-1}\mathsf{
R}^{j}_{t}}{\mathsf{T}}$, where $\mathsf{R}_{t}^{j}$ is the rank of the $j$th layer projection matrix at time $0\leq t\leq\mathsf{T}-1$, for some $\mathsf{T}\in\mathbb{N}$, and accordingly the *total weighted-average effective rank* is defined as $\mathcal{R}_{\text{adap}}\triangleq\frac{\sum_{j=1}^{L}\sum_{t=0}^{\mathsf{T}-
1}\mathsf{R}^{j}_{t}(d_{j}+d_{j+1})}{\mathsf{T}\cdot\sum_{j=0}^{L}(d_{j}\cdot d
_{j+1})}.$ Following that, we define the average per-layer reduction in memory footprint, when compared to the non-adaptive low-rank fine-tuning, by $\mathcal{M}_{\text{red}}^{j}\triangleq(\bar{\mathsf{R}}^{j}-\mathcal{R}^{j}_{
\text{adap}})\cdot(d_{j}+d_{j+1}),$ where $\bar{\mathsf{R}}^{j}$ is the time-independent non-adaptive rank (such as in Galore). Finally, we define the total memory reduction by $\mathcal{M}_{\text{red}}\triangleq\sum_{j=1}^{L}(\bar{\mathsf{R}}^{j}-\mathcal
{R}^{j}_{\text{adap}})\cdot(d_{j}+d_{j+1}).$

### 4.1 Fine-tuning on GLUE benchmark

We evaluate our model on the GLUE benchmark [] by fine-tuning the pre-trained Roberta-base model [] on 8 tasks. We compare the results against full fine-tuning, LoRA, and GaLore methods and present them in Table 2. We report the overall (matched and mismatched) accuracy for MNLI, Matthew’s correlation for CoLA, Pearson correlation for STS-B, F1-score for MRPC, and accuracy for other tasks.
As can be seen, our method improves the accuracy of fine-tuning while consuming less training memory on average. In addition, in Figure 4, we present the average memory reduction measured at the end of each epoch (blue) and the corresponding effective rank (green). Finally, empirical analysis of the hyperparameters is provided in Figures 5.

**Table 2: Evaluating AdaRankGrad comparing to state-of-the-art memory-efficient fine-tuning methods on GLUE benchmark using pre-trained RoBERTa-Base. For AdaRankGrad, we present accuracy results with average effective rank.**
| Model | Memory | CoLA | STS-B | MRPC | RTE | SST2 | MNLI | QNLI | QQP |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Full Fine-Tuning | 747M | 62.24 | 90.92 | 91.30 | 79.42 | 94.57 | 87.18 | 92.33 | 92.28 |
| LoRA (rank=4) | 257M | 61.38 | 90.57 | 91.07 | 78.70 | 92.89 | 86.82 | 92.18 | 91.29 |
| GaLore (rank=4) | 253M | 60.35 | 90.73 | 92.25 | 79.42 | 94.0 | 87.0 | 92.24 | 91.06 |
| AdaRankGrad (Initial rank=4) | 202M | 61.4(3.71) | 90.97(3.79) | 92.6(3.72) | 81.23(3.65) | 94.8(3.91) | 86.6(3.69) | 92.5(3.9) | 90.4(3.67) |
| LoRA (rank=8) | 264M | 61.83 | 90.80 | 91.90 | 79.06 | 93.46 | 86.94 | 92.25 | 91.22 |
| GaLore (rank=8) | 257M | 60.06 | 90.82 | 92.0 | 79.78 | 94.38 | 87.17 | 92.2 | 91.11 |
| AdaRankGrad (Initial rank=8) | 237M | 62.0(6.41) | 90.89(7.77) | 92.5(5.33) | 81.23(6.64) | 94.80(5.69) | 86.5(6.00) | 92.6(6.54) | 89.7(6.00) |

Figure: Figure 4: We present the effective rank measured for non-attention layers and corresponding memory reduction for AdaRankGrad trained on MRPC (left panel) and RTE (right panel) datasets from the GLUE benchmark.
Refer to caption: extracted/5948913/MRPC_r8.png

Figure: Figure 5: The left graph presents the effective rank measured for different values of $\eta_{th}$ while training AdaRankGrad on MRPC dataset. The right graph present the effective rank measured for different values of $\eta_{th}$ while training AdaRankGrad on MRPC dataset.
Refer to caption: extracted/5948913/effective_rank_xi.png

### 4.2 Fine-tuning Geneformer on Biological Omics Tabular Data

High-throughput omics technologies, such as next-generation sequencing [], allow for the simultaneous measurement of thousands to 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 [, ]. Unfortunately, the complexity of omics data makes the use of low-rank optimization methods such as LoRa sub-optimal for foundation omics models, and therefore, training is often based on relatively high ranks []. In this experiment, we demonstrate that the proposed approach can alleviate this problem due to the adaptive low-rank projection of gradients.
We conducted a fine-tuning experiment for cell classification, specifically focusing on the Classification of Cardiomyopathy Disease States. We used the *Geneformer gf-12L-30M-i2048*
(^4^44A transformer-based 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: [https://huggingface.co/ctheodoris/Geneformer](https://huggingface.co/ctheodoris/Geneformer))
as our pre-trained base model, which is based on the BERT model and pre-trained on 30 million single-cell transcriptomes [, ]. Following the setup proposed in [], we fine-tuned our method with a maximal rank of 16 on approximately 187,000 samples for three epochs, with a batch size of 16. We evaluated the model on a test set of size 17,228. Figure 6 presents the Macro-F1 score and Accuracy measured for each epoch. As shown, our method converges faster with a stable improvement over LoRA.

Figure: Figure 6: 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 to LoRA and Galore methods.
Refer to caption: extracted/5948913/cell_class_accuracy.png

### 4.3 Pre-training LLAMA on C4 Dataset

Here, we repeated the comparison presented in [] [Table 2] to evaluate AdaRankGrad performance to the state-of-the-art method, in terms of perplexity and memory usage. We evaluate AdaRankGrad by training large LLaMA-based models on the C4 dataset, a cleaned and massive version of Common Crawl’s web corpus []. This dataset is primarily used for pre-training language models and learning word representations. To closely mimic practical pre-training scenarios, we train on a large dataset without repeating data, scaling model sizes up to $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/Tokens | 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 | 18G | 26G | 16.4G |

**Table 4: A 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 $r_{\text{init}}=r$ for AdaRankGrad, w.r.t the $r$ presented in the last table raw. The information threshold $\eta_{\text{th}}=0.48$ for the three model training. The validation perplexity is presented, along with an estimate of the memory required for the total number of parameters and optimizer states in BF16 format. We used NVIDIA A100 for the three experiments.**
|  | $\mathbf{60M}$ | $\mathbf{130M}$ | $\mathbf{350M}$ |
| --- | --- | --- | --- |
| Full-Rank | $34.06(0.36\mathrm{G})$ | $25.08(0.76\mathrm{G})$ | $18.80(2.06\mathrm{G})$ |
| GaLore | $34.88(0.24\mathrm{G})$ | $25.36(0.52\mathrm{G})$ | $18.95(1.22\mathrm{G})$ |
| Low-Rank | $78.18(0.26\mathrm{G})$ | $45.51(0.54\mathrm{G})$ | $37.41(1.08\mathrm{G})$ |
| LoRA | $34.99(0.36\mathrm{G})$ | $33.92(0.80\mathrm{G})$ | $25.58(1.76\mathrm{G})$ |
| ReLoRA | $37.04(0.36\mathrm{G})$ | $29.37(0.80\mathrm{G})$ | $29.08(1.76\mathrm{G})$ |
| AdaRankGrad | $\mathbf{34.24}(0.206\mathrm{G})$ | $\mathbf{25.22}(0.497\mathrm{G})$ | $\mathbf{18.91}(1.106\mathrm{G})$ |
| Training Tokens | 1.1 B | 2.2 B | 6.4 B |
| $r/d_{\text{model}}$ | $128/256$ | $256/768$ | $256/1024$ |

## 5 Discussion

In this paper, we present AdaRankGrad, a full-parameters efficient optimization scheme that applies adaptive low-rank updates without relying on a parallel low-rank adapter (LoRA), thus preserving the natural training dynamics. Unlike methods such as LoRA, which rely on parallel adapters, AdaRankGrad enables full parameter fine-tuning while maintaining low memory costs through efficient low-rank optimization updates. Moreover, unlike GaLore, AdaRankGrad leverages the natural phenomenon where the dimensionality of the approximated gradient decreases as training progresses. 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 or after convergence. As demonstrated in Section 4, this approach results in superior prediction performance.

AdaRankGrad offers a unique trade-off between memory efficiency and model performance. While it may slightly increase training time due to 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 a very narrow range, reducing the overall computational cost. To further mitigate this, we proposed an efficient subspace-update technique that significantly reduces the SVD-based subspace calculations by an order of magnitude. In 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 as 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. These directions could further enhance the efficiency and performance of optimization schemes like AdaRankGrad.

## Appendix A Proofs

### A.1 Proof of Lemma 2

In this section, we prove Lemma 2.
Consider the SVD decomposition of the gradient ${\bf G}_{t}={\bf U}_{t}\Sigma_{t}{\bf V}_{t}^{\top},$ at iteration $t$. For any natural number $r<n,$ we denote ${\bf H}_{t}^{n\times r}(r)={\bf U}[:,1:r]$. For simplicity of notation, let us denote ${\bf P}_{t}(r)={\bf H}_{t}(r){\bf H}_{t}^{\top}(r),$ where ${\bf P}_{t}(r)$ is an orthogonal projection matrix, i.e., ${\bf P}_{t}^{\top}(r){\bf P}_{t}(r)={\bf P}_{t}(r),$ and ${\bf P}_{t}(r)={\bf P}_{t}^{\top}(r)$. Without the loss of generality, we assume that at $t=0$, the rank of ${\bf G}_{0}$ is such that $\text{rank}\left({\bf G}_{0}\right)>r$. For reversible networks, it was shown in [][Theorem 3.2] that the gradients have the form ${\bf G}_{t}=\frac{1}{N}\sum_{i=1}^{N}\left({\bf A}_{i}-{\bf B}_{i}{\bf W}_{t}{
\bf C}_{i}\right)$, with constant matrices $\{{\bf A}_{i}\}_{i}$, and PSD matrices $\{{\bf B}_{i},{\bf C}_{i}\}_{i}$, for $t\geq\mathsf{t}_{0}$, where $\mathsf{t}_{0}\in\mathbb{N}$. Recall that in the vanilla SGD weight update, we have ${\bf W}_{t}={\bf W}_{t-1}+\eta{\bf G}_{t-1}$. Let ${\bf S}\triangleq\frac{1}{N}\sum_{i=1}^{N}{\bf C}_{i}\otimes{\bf B}_{i}$, and $\lambda_{1}<\lambda_{2}$ 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 []. Specifically, let ${\bf G}_{\mathsf{t}_{0}}^{\|}$ be the projection of ${\bf G}_{\mathsf{t}_{0}}$ onto the minimal eigenspace $\mathcal{V}_{1}$ of $S$ corresponding to $\lambda_{1}.$ By assumption, we know that the rank of ${\bf G}_{\mathsf{t}_{0}}^{\|}$ is $L$, and its SVD is ${\bf G}_{\mathsf{t}_{0}}^{\|}=\sum_{l=1}^{L}c_{l}\bm{z}_{l}\bm{y}_{l}^{\top}$, where $\left\{\bm{z}_{l}\right\}_{l=1}^{L}$ and $\left\{\bm{y}_{l}\right\}_{l=1}^{L}$ are the orthonormal unit vectors, and $\left\{c_{l}\right\}_{l=1}^{L}$ are the corresponding singular values. Thus, it was in [] that,

$$ $\displaystyle g_{0}^{\|}=\operatorname{vec}\left({\bf G}_{\mathsf{t}_{0}}^{\|} \right)=\sum_{l=1}^{L}c_{l}\left(\bm{y}_{l}\otimes\bm{z}_{l}\right)\triangleq \sum_{l=1}^{L}c_{l}\bm{v}_{l},$ $$

with unit vector $\bm{v}_{l}\triangleq\bm{y}_{l}\otimes\bm{z}_{l}\in\mathcal{V}_{1}$. It is clear that

$$ $\displaystyle\bm{v}_{l}^{\top}\bm{v}_{l^{\prime}}$ $\displaystyle=\left(\bm{y}_{l}^{\top}\otimes\bm{z}_{l}^{\top}\right)\left(\bm{ y}_{l^{\prime}}\otimes\bm{z}_{l^{\prime}}\right)$ $\displaystyle=\left(\bm{y}_{l}^{\top}\bm{y}_{l^{\prime}}\right)\left(\bm{z}_{l }^{\top}\bm{z}_{l^{\prime}}\right)$ $\displaystyle=\mathbb{I}\left(l=l^{\prime}\right),$ $$

where $\mathbb{I}$ is the indicator function.
Using this, it was finally shown in [] that,

$$ $\displaystyle\left\|{\bf G}_{t}\right\|_{2}$ $\displaystyle=\max_{\left\|\bm{y}^{\prime}\right\|_{2}=1,\left\|\bm{z}^{\prime }\right\|_{2}=1}\bm{z}^{\prime\top}{\bf G}_{t}\bm{y}^{\prime}$ $\displaystyle\geq\max_{l}\bm{z}_{l}^{\top}{\bf G}_{t}\bm{y}_{l}=\max_{l}\left( \bm{y}_{l}\otimes\bm{z}_{l}\right)^{\top}g_{t}=\max_{l}\bm{v}_{l}^{\top}(1- \eta S)^{t}g_{0}=\left(1-\eta\lambda_{1}\right)^{t}\max_{l}\bm{v}_{l}^{\top}g_ {0}.$ $$

Using the above results from [], we can now see that,

$$ $\displaystyle\left\|{\bf G}_{t}\right\|_{2}$ $\displaystyle=\max_{\left\|\bm{y}\right\|_{2}=1,\left\|\bm{z}\right\|_{2}=1} \bm{z}^{\prime\top}{\bf G}_{t}\bm{y}$ $\displaystyle=\max_{\left\|\bm{y}\right\|_{2}=1,\left\|\bm{z}\right\|_{2}=1} \left(\bm{y}\otimes\bm{z}\right)^{\top}g_{t}$ $\displaystyle=\max_{\left\|\bm{y}\right\|_{2}=1,\left\|\bm{z}\right\|_{2}=1} \left(\bm{y}\otimes\bm{z}\right)^{\top}(1-\eta S)^{t}g_{0}$ $\displaystyle=\left({\bf U}_{0}[:,1]^{\top}\otimes{\bf V}_{0}[:,1]\right)^{ \top}(1-\eta S)^{t}g_{0}$ $\displaystyle=\tilde{\bm{v}}^{\top}(1-\eta S)^{t}g_{0}$ $\displaystyle=\left(1-\eta\lambda_{1}\right)^{t}\tilde{\bm{v}}^{\top}g_{0}.$ $$

Thus, we have

$$ $\left(1-\eta\lambda_{1}\right)^{t}\tilde{\bm{v}}^{\top}g_{0}\geq\left(1-\eta \lambda_{1}\right)^{t}\max_{l}\bm{v}_{l}^{\top}g_{0},$ $$

or equivalently

$$ $\left(1-\eta\lambda_{1}\right)^{t}\tilde{\bm{v}}^{\top}g_{0}\geq\left(1-\eta \lambda_{1}\right)^{t}\left\|g_{0}^{\|}\right\|_{2}.$ (3) $$

Now,

$$ $\displaystyle\kappa(t)$ $\displaystyle\triangleq\frac{\left\|{\bf G}_{t}-{\bf P}_{t}(1){\bf G}_{t} \right\|_{F}^{2}}{\left\|{\bf G}_{t}\right\|_{F}^{2}}$ $\displaystyle=\frac{\left\|{\bf G}_{t}\right\|_{F}^{2}-\left\|{\bf G}_{t} \right\|_{2}^{2}}{\left\|{\bf G}_{t}\right\|_{F}^{2}}$ $\displaystyle\leq\frac{\left\|{\bf G}_{t}\right\|_{F}^{2}-\left\|{\bf G}_{t} \right\|_{2}^{2}}{\left\|{\bf G}_{t}\right\|_{2}^{2}}=\frac{\left\|{\bf G}_{t} \right\|_{F}^{2}}{\left\|{\bf G}_{t}\right\|_{2}^{2}}-1$ $\displaystyle\leq\frac{\left(1-\eta\lambda_{1}\right)^{2t}\left\|g_{0}^{\|} \right\|_{2}^{2}}{\left\|{\bf G}_{t}\right\|_{2}^{2}}+\frac{\left(1-\eta \lambda_{2}\right)^{2t}\left\|g_{0}^{\perp}\right\|_{2}^{2}}{\left\|{\bf G}_{t }\right\|_{2}^{2}}-1$ $\displaystyle\leq\underbrace{\frac{\left(1-\eta\lambda_{1}\right)^{2t}\left\|g _{0}^{\|}\right\|_{2}^{2}}{\left(1-\eta\lambda_{1}\right)^{2t}\tilde{\bm{v}}^{ \top}g_{0}}}_{\leq 1}+\frac{\left(1-\eta\lambda_{2}\right)^{2t}\left\|g_{0}^{ \perp}\right\|_{2}^{2}}{\left(1-\eta\lambda_{1}\right)^{2t}\tilde{\bm{v}}^{ \top}g_{0}}-1$ $\displaystyle\leq\frac{\left(1-\eta\lambda_{2}\right)^{2t}\left\|g_{0}^{\perp} \right\|_{2}^{2}}{\left(1-\eta\lambda_{1}\right)^{2t}\tilde{\bm{v}}^{\top}g_{0 }}=\frac{\left(1-\eta\lambda_{2}\right)^{2t}\left\|g_{0}^{\perp}\right\|_{2}^{ 2}}{\left(1-\eta\lambda_{1}\right)^{2t}\tilde{\bm{v}}^{\top}g_{0}},$ $$

where the first inequality follows from the fact that $\left\|{\bf G}_{t}\right\|_{F}^{2}>\left\|{\bf G}_{t}\right\|_{2}^{2},$ the second inequality is by using [][Lemma 3.3], i.e.,

$$ $\left\|{\bf G}_{t}\right\|_{F}^{2}\leq\left(1-\eta\lambda_{2}\right)^{2t}\left \|g_{0}^{\perp}\right\|_{2}^{2}+\left(1-\eta\lambda_{1}\right)^{2t}\left\|g_{0 }^{\|}\right\|_{2}^{2},$ $$

and the third inequality is due to Eq. (3). Finally, by defining the constants $c_{1}\triangleq\frac{\left(1-\eta\lambda_{2}\right)}{\left(1-\eta\lambda_{1}
\right)}$ and $c_{2}\triangleq\frac{\left\|g_{0}^{\perp}\right\|_{2}^{2}}{\tilde{\bm{v}}^{
\top}g_{0}}<1$, we get that $\kappa(t)<c_{1}^{2t}\cdot c_{2}$, which concludes the proof.

### A.2 Proof of Theorem 3

In this section, we prove Theorem 3. We upper bound Frobenius norm of ${\bf G}^{j}_{t}$, for any layer $j\in[L]$; in the following, for simplicity of notation, we ignore the index $j$ an use ${\bf G}_{t}$ instead. By Lemma 4, the low-rank optimization block 3 in Algorithm 3 is guaranteed to converge; we denote by $\mathsf{T}_{\ell}\in\mathbb{N}$ the time index $t$ at which we exit block 3 for the $\ell$th time (i.e., $\|\hat{{\bf G}}_{\mathsf{T}_{\ell}}\|\leq\varsigma_{2}$), for $\ell\in\mathbb{N}$. Furthermore, we recall that ${\bf G}^{j}_{t}\triangleq\nabla_{\bm{{\bf W}^{j}}}f\left(\bm{\theta}_{t}\right)$; when clear from the context, we omit $j$ from ${{\bf W}^{j}}$, and use instead $\nabla_{\bm{{\bf W}^{j}}}f\left(\bm{\theta}_{t}\right)=\nabla f\left({\bf W}_{
t}\right)$.
Consider the SVD decomposition of the gradient $\nabla_{\bm{{\bf W}^{j}}}f\left(\bm{\theta}_{\mathsf{T}_{i}}\right)={\bf U}_{
\mathsf{T}_{i}}\Sigma_{\mathsf{T}_{i}}{\bf V}_{\mathsf{T}_{i}}^{\top}$. For $t\in[\mathsf{T}_{i},\mathsf{T}_{i+1}-1]$, we define the projected gradient as $\hat{{\bf G}}_{t}\triangleq{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}
_{t},$ where ${\bf P}_{\mathsf{T}_{i}}={\bf U}_{\mathsf{T}_{i}}\left[:,:r_{\mathsf{T}_{i}}
\right]^{\top},$ for a given threshold $0<\eta_{\text{th}}\leq 1$, where

$$ $\displaystyle r_{\mathsf{T}_{i}}\triangleq\displaystyle\mathop{\mbox{\rm sup}} \left\{r\in\mathbb{N}:\eta_{\text{th}}\cdot\|{\bf G}_{\mathsf{T}_{i}}\|_{F}^{2 }-\left\|{\bf G}_{\mathsf{T}_{i}}-{\bf P}_{\mathsf{T}_{i}}(r){\bf G}_{\mathsf{ T}_{i}}\right\|_{F}^{2}\geq 0\right\},$ (4) $$

obtained by the search presented in IASS Algorithm 2, using exact truncated-SVD calculation.
Note that since $\left\|{\bf G}_{t}-{\bf P}_{\mathsf{T}_{i}}(r){\bf G}_{t}\right\|_{F}^{2}\leq
\|{\bf G}_{t}\|_{F}^{2}$ always, it must be the case that $0<\eta_{\text{th}}\leq 1$. Also, because ${\bf P}_{\mathsf{T}_{i}}(n){\bf G}_{t}={\bf G}_{t}$, $r_{\mathsf{T}_{i}}$ in (4) is well-defined. Next, let $h_{t}\triangleq f({\bf W}_{t})-f({\bf W}_{\mathsf{T}_{i+1}})$, and $\alpha_{t}$ denote the learning rate. Then,

$$ $\displaystyle h_{t+1}$ $\displaystyle=f\left(\mathbf{W}_{t+1}\right)-f\left(\mathbf{W}_{\mathsf{T}_{i+ 1}}\right)$ $\displaystyle=f\left(\mathbf{W}_{t}-\alpha_{t}\left(\hat{{\bf G}}_{t}\right) \right)-f\left(\mathbf{W}_{\mathsf{T}_{i+1}}\right)$ $\displaystyle\underset{(1)}{\leq}f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W }_{\mathsf{T}_{i+1}}\right)-\alpha_{t}\text{vec}\left(\hat{{\bf G}}_{t}\right) ^{\top}\text{vec}\left(\nabla f\left(\mathbf{W}_{t}\right)\right)+\alpha_{t}^{ 2}\frac{\beta}{2}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2}$ $\displaystyle=f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W}_{\mathsf{T}_{i+1} }\right)-\alpha_{t}\text{vec}\left(\hat{{\bf G}}_{t}\right)^{\top}\text{vec} \left({\bf G}_{t}\right)+\alpha_{t}^{2}\frac{\beta}{2}\left\|\hat{{\bf G}}_{t} \right\|_{F}^{2}$ $\displaystyle\underset{(2)}{\leq}f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W }_{\mathsf{T}_{i+1}}\right)-\alpha_{t}\operatorname{tr}\left(\left({\bf P}_{ \mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_{t}\right)^{\top}{\bf G}_{t}\right) +\alpha_{t}^{2}\frac{\beta}{2}\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle\underset{(3)}{\leq}f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W }_{\mathsf{T}_{i+1}}\right)-\alpha_{t}\operatorname{tr}\left(\left({\bf G}_{t} ^{\top}{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}})\right){\bf G}_{t}\right)+ \alpha_{t}^{2}\frac{\beta}{2}\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle=f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W}_{\mathsf{T}_{i+1} }\right)-\alpha_{t}\left\|{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_ {t}\right\|_{F}^{2}+\alpha_{t}^{2}\frac{\beta}{2}\left\|{\bf G}_{t}\right\|_{F }^{2}$ $\displaystyle=f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W}_{\mathsf{T}_{i+1} }\right)-\alpha_{t}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2}+\alpha_{t}^{2} \frac{\beta}{2}\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle\underset{(4)}{\leq}f\left(\mathbf{W}_{t}\right)-f\left(\mathbf{W }_{\mathsf{T}_{i+1}}\right)-\alpha_{t}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2} +\frac{\beta(D\alpha_{t})^{2}}{2}$ $\displaystyle=h_{t}-\alpha_{t}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2}+\frac{ \beta(D\alpha_{t})^{2}}{2},$ (5) $$

where (1) follows by the assumption that $f$ is $\beta$-smooth function and the decent lemma (see, e.g. [][Definition 5.1], (2) follows from the identity $\text{vec}\left({\bf A}{\bf B}\right)^{\top}\text{vec}\left({\bf B}\right)=
\operatorname{tr}(({\bf A}{\bf B})^{\top}{\bf B})$, in (3) we use the fact that since ${\bf P}_{t}$ is an orthogonal projection matrix, we have ${\bf P}_{t}={\bf P}_{t}^{\top}$, and finally, (4) follows from the assumption that the norms of the gradients are bounded by $D$.
Rearranging (5), and choosing $\alpha_{t}=\alpha$, for all $t\geq 0$, we readily obtain that,

$$ $\displaystyle\sum_{t=\mathsf{T}_{i}}^{\mathsf{T}_{i+1}-1}\left\|\hat{{\bf G}}_ {t}\right\|_{F}^{2}$ $\displaystyle\leq\frac{h_{\mathsf{T}_{i}}-h_{\mathsf{T}_{i+1}}}{\alpha}+\frac{ (\mathsf{T}_{i+1}-\mathsf{T}_{i})\beta D^{2}\alpha}{2}.$ (6) $$

Thus, for $N\in\mathbb{N}$,

$$ $\displaystyle\frac{1}{\mathsf{T}_{N}}\sum_{i=0}^{N-1}\sum_{t=\mathsf{T}_{i}}^{ \mathsf{T}_{i+1}-1}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2}$ $\displaystyle\leq\frac{1}{\mathsf{T}_{N}}\sum_{i=0}^{N-1}\left[\frac{h_{ \mathsf{T}_{i}}-h_{\mathsf{T}_{i+1}}}{\alpha}+\frac{(\mathsf{T}_{i+1}-\mathsf{ T}_{i})\beta D^{2}\alpha}{2}\right]$ (7) $\displaystyle=\frac{h_{\mathsf{T}_{0}}-h_{\mathsf{T}_{N}}}{\alpha\mathsf{T}_{N }}+\frac{(\mathsf{T}_{N}-\mathsf{T}_{0})\beta D^{2}\alpha}{2\mathsf{T}_{N}}$ (8) $\displaystyle=\frac{R}{\sqrt{\mathsf{T}_{N}}},$ (9) $$

where $R\triangleq\sqrt{4M\beta D^{2}}$, and in the second equality we choose $\alpha=\sqrt{\frac{4M}{\mathsf{T}_{N}\beta D^{2}}}$.(^5^55Note, that this choice of $\alpha$ implies the guarantees of Lemma 4, because $\alpha=\sqrt{\frac{4M}{\mathsf{T}_{N}\beta D^{2}}}\leq\sqrt{\frac{4M}{(\mathsf {T}_{i+1}-\mathsf{T}_{i})\beta D^{2}}}\triangleq\alpha_{i}^{\prime}$, and $\alpha_{i}^{\prime}$ is the optimized learning rate for the $i$th block, that was chosen in the proof of Lemma 4.)
Now recall that by the definition of $r_{\mathsf{T}_{i}}$ in (4), we have,

$$ $\left\|{\bf G}_{\mathsf{T}_{i}}-{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){ \bf G}_{\mathsf{T}_{i}}\right\|_{F}^{2}\leq\eta_{\text{th}}\left\|{\bf G}_{ \mathsf{T}_{i}}\right\|_{F}^{2}.$ (10) $$

Moreover, the following clearly holds for any $t\in\mathbb{N}$,

$$ $\displaystyle\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle=\left\|{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_{t} \right\|_{F}^{2}+\left\|{\bf P}^{\perp}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){ \bf G}_{t}\right\|_{F}^{2}$ (11) $\displaystyle=\left\|{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_{t} \right\|_{F}^{2}+\left\|{\bf G}_{t}-{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i} }){\bf G}_{t}\right\|_{F}^{2},$ (12) $$

and thus by plugging (10) into (12), at $t=\mathsf{T}_{i}$, for any $i\in\mathbb{N}$, we get,

$$ $\displaystyle(1-\eta_{\text{th}})\|{\bf G}_{\mathsf{T}_{i}}\|_{F}^{2}\leq\left \|{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_{\mathsf{T}_{i}}\right\| _{F}^{2}.$ (13) $$

Accordingly,

$$ $\displaystyle\left\|{\bf P}^{\perp}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G }_{\mathsf{T}_{i}}\right\|_{F}^{2}\leq\frac{\eta_{\text{th}}}{1-\eta_{\text{th }}}\left\|{\bf P}_{\mathsf{T}_{i}}(r_{\mathsf{T}_{i}}){\bf G}_{\mathsf{T}_{i}} \right\|_{F}^{2},$ (14) $$

Recall from Subsection 3.1 that for the reversible layer,

$$ $\displaystyle\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle=\left\|(I-\alpha{\bf S}){\bf G}_{t-1}\right\|_{F}^{2}$ (15) $\displaystyle\leq\left\|(I-\alpha{\bf S})\right\|_{2}^{2}\left\|{\bf G}_{t-1} \right\|_{F}^{2}$ (16) $\displaystyle=\max_{i}|1-\alpha\lambda_{i}|^{2}\left\|{\bf G}_{t-1}\right\|_{F }^{2},$ (17) $$

where $\{\lambda_{i}\}_{i}$ are the eigenvalue of ${\bf S}$. Thus, using the fact that ${\bf S}$ is positive semi-definite matrix, if the learning rate $\alpha$ is such that $\alpha\leq\frac{2}{\lambda_{\max}}$, where $\lambda_{\max}$ is the maximal eigenvalue of ${\bf S}$, then we get that $\max_{i}|1-\alpha\lambda_{i}|^{2}\leq 1$, and accordingly, $\left\|{\bf G}_{t}\right\|_{F}^{2}\leq\left\|{\bf G}_{t-1}\right\|_{F}^{2}$. This means that the Frobenius norm of the gradient is monotonically decreasing as a function of $t$.

Now, recall that $\varsigma_{2,i}$ is any positive number such that $\varsigma_{2,i}<\sqrt{1-\eta_{\text{th}}}\cdot\|\mathbf{G}_{\mathsf{T}_{i-1}}
\|_{F}$. In light of (13), this necessarily implies that in each block $i$, we will execute (at least once) the low-rank optimization block (indeed, the condition $\|\hat{\mathbf{G}}_{\mathsf{T}_{i}}\|_{F}>\varsigma_{2,i}$ is satisfied). This, conjugated with the monotonicity property that $\left\|{\bf G}_{t}\right\|_{F}^{2}\leq\left\|{\bf G}_{{\mathsf{T}_{i}}}\right
\|_{F}^{2}$, for any $t\in[\mathsf{T}_{i},\mathsf{T}_{i+1}-1]$ and $i\in[N]$, imply that

$$ $\displaystyle\frac{1}{\mathsf{T}_{N}}\sum_{i=0}^{N-1}\sum_{t=\mathsf{T}_{i}}^{ \mathsf{T}_{i+1}-1}\left\|{\bf G}_{t}\right\|_{F}^{2}$ $\displaystyle\leq\frac{1}{\mathsf{T}_{N}}\sum_{i=0}^{N-1}\sum_{t=\mathsf{T}_{i }}^{\mathsf{T}_{i+1}-1}\left\|{\bf G}_{{\mathsf{T}_{i}}}\right\|_{F}^{2}$ (18) $\displaystyle\leq\frac{1}{(1-\eta_{\text{th}})\mathsf{T}_{N}}\sum_{i=1}^{N-1} \sum_{t=\mathsf{T}_{i}}^{\mathsf{T}_{i+1}-1}\left\|{\bf P}_{\mathsf{T}_{i}}(r_ {\mathsf{T}_{i}}){\bf G}_{t}\right\|_{F}^{2}$ (19) $\displaystyle\leq\frac{R}{(1-\eta_{\text{th}})\sqrt{\mathsf{T}_{N}}}.$ (20) $$

Accordingly, for any $\varepsilon\geq 0$, and $\mathsf{T}_{N}>\frac{R^{2}}{(1-\eta_{\text{th}})^{2}\varepsilon^{2}}$,

$$ $\displaystyle\min_{0\leq t\leq\mathsf{T}_{N}}{\left\|{\bf G}_{t}\right\|_{F}^{ 2}}\leq\frac{1}{\mathsf{T}_{N}}\sum_{i=0}^{N-1}\sum_{t=\mathsf{T}_{i}}^{ \mathsf{T}_{i+1}-1}\left\|{\bf G}_{t}\right\|_{F}^{2}\leq\varepsilon,$ (21) $$

and thus, there exists an iteration index $t\in[0,\mathsf{T}_{N}]$ for which,

$$ $\displaystyle\left\|{\bf G}_{t}\right\|_{F}^{2}\leq\varepsilon,$ (22) $$

which, by definition, implies that Algorithm 3 achieves an $\varepsilon$-critical point.

###### Lemma 4 (Convergence of low-rank optimization block) .

Consider the same setting and assumptions as in Theorem 3. Then, the time $t=\mathsf{T}_{\ell}\in\mathbb{N}$ at which Algorithm 3 exits block 3 for the $\ell$th time is finite, for any $\ell\in\mathbb{N}$.

###### Proof of Lemma 4 .

We follow the same notations as in the proof of Theorem 3. Recall from (5) that the following holds true,

$$ $\displaystyle f\left(\mathbf{W}_{t+1}\right)$ $\displaystyle\leq f\left(\mathbf{W}_{t}\right)-\alpha_{t}\left\|\hat{{\bf G}}_ {t}\right\|_{F}^{2}+\frac{\beta(D\alpha_{t})^{2}}{2},$ (23) $$

for any $t\geq 0$. Now, we enter for the first time the low-rank optimization block of Algorithm 3 at time $\mathsf{T}_{0}=t=0$, and we next show that this block converges. Fix $\mathsf{T}\in\mathbb{N}$. Using (23), and following the same arguments as in (6)–(9), we have,

$$ $\displaystyle\frac{1}{\mathsf{T}}\sum_{t=0}^{\mathsf{T}-1}\left\|\hat{{\bf G}} _{t}\right\|_{F}^{2}$ $\displaystyle\leq\frac{f\left(\mathbf{W}_{0}\right)-f\left(\mathbf{W}_{\mathsf {T}}\right)}{\mathsf{T}\alpha}+\frac{\beta D^{2}\alpha}{2}$ (24) $\displaystyle\leq\frac{R}{\sqrt{\mathsf{T}}},$ (25) $$

where $R\triangleq\sqrt{4M\beta D^{2}}$, and in the second equality we choose some $\alpha\leq\sqrt{\frac{4M}{\mathsf{T}\beta D^{2}}}$.
Accordingly, for any $\varsigma\geq 0$, and $\mathsf{T}>\frac{R^{2}}{\varsigma^{2}}$, we clearly have,

$$ $\displaystyle\min_{t\in[0,\mathsf{T}-1]}{\left\|\hat{{\bf G}}_{t}\right\|_{F}^ {2}}\leq\sum_{t=0}^{\mathsf{T}-1}\left\|\hat{{\bf G}}_{t}\right\|_{F}^{2}\leq\varsigma,$ (26) $$

namely, there exists $\mathsf{T}_{1}\in[0,\mathsf{T}-1]$ such that $\left\|\hat{{\bf G}}_{\mathsf{T_{1}}}\right\|_{F}^{2}\leq\varsigma$; note that $\mathsf{T}_{1}$ is the time index at which we exist the low-rank block for the first time, and the above guarantees that $\mathsf{T}_{1}$ finite. The same arguments above apply for any block exit time $\mathsf{T}_{\ell},\;\ell\geq 2$, which concludes the proof.
$\Box$

###### Lemma 4 (Convergence of low-rank optimization block) .

###### Proof of Lemma 4 .

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Could you address the weaknesses above?
soundness: 3
strengths: - Reducing memory requirements is a crucial part of training LLMs
- The paper is well-written and the proofs are well designed.
- The experiments are comprehensive.
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.
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?
- How does this method compare to BlockLLM [1]? I see a lot of similarities between the two. 
- 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?
- 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.


[1] https://arxiv.org/abs/2406.17296

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Questions

1. 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?

2. 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. 

3. 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?
soundness: 3
strengths: Strengths

1. Memory Efficiency: AdaRankGrad significantly reduces memory usage during training, which is crucial for large-scale models.

2. Theoretical Foundation: The paper provides a theoretical basis for the method, leveraging the phenomenon of decreasing gradient rank.
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.
weaknesses: Weaknesses

1. 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 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.
2. 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.
3. The authors propose to use SSRF to replace the SVD used in Galore, which speeds up the search for subspaces.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 4
questions: Based on the previously mentioned weakness, I have the following questions:
- 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.
- 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?
- 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.
- 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. 

I'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.
soundness: 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.
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.
weaknesses: My main concern is about the experiment:
- 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.
- 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. What do the memory savings look like when factoring in the increased running time?
2. Can the savings be fully utilized by other jobs? If so, how?
3. Why is $\epsilon$ needed in defining the $(\eta, \epsilon)$-approximation?
soundness: 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.
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.
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. 

Feedback on specific claims and points for clarification:

1. 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.

2. 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.

3. The need for two parameters, $\eta$ and $\epsilon$, to define $(\eta, \epsilon)$-approximation is unclear. Would a single $\eta$ suffice?

**Suggestions that do not impact the score:**

- 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.
  
- In Algorithm 1 (SSRF), where $A$ is $n \times m$ and $\Omega$ is $k \times n$, the computation $A \Omega$ seems undefined.
  
- For Figure 5, should the x-axis be labeled as 11700, 23400, 35100 (for epochs 1, 2, 3)?

Meta-review: In this paper, the authors proposed a memory-efficient low-rank gradient update algorithm for LLMs training and fine-tuning.

The reviewers raised some concerns and questions regarding the empirical evaluation of the proposed algorithm. Most of the concerns are addressed during the rebuttal. 

The authors are encouraged to conduct experiments on the MMLU task as suggested by a reviewer to make the experimental results more convincing. Rather than that this paper has shown promising results, both empirically and theoretically, of the proposed algorithm. Therefore, I recommend an acceptance.

Decision: Accept (Poster)