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

Memory savings utilization. The memory savings achieved by our method can be utilized through methods such as dynamic resource allocation and job scheduling technologies, like those employed by companies such as Run.AI [1]. Several research papers have explored and suggested methods for efficient memory allocation and job scheduling that could leverage the memory savings [2,3].

[1] https://www.run.ai/blog/dynamic-gpu-memory-solving-the-problem-of-inefficient-resource-allocation-in-inference-server [2] Wang, Linnan, et al. "Superneurons: dynamic GPU memory management for training deep neural networks." Proceedings of the 23rd ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. ACM, 2018, pp. 41–53. DOI: 10.1145/3178487.3178491. [3] Tang, Yu, et al. "DELTA: Dynamically Optimizing GPU Memory beyond Tensor Recomputation." arXiv preprint arXiv:2203.15980 (2022). URL: https://arxiv.org/abs/2203.15980.

We apologize for the inconvenience, and thank the reviewer for pointing out these Typos.

(1) We remove the sentence mentioned.

(2) The matrices' dimensions were corrected (in the revised version).

(3) We re-created the plot with informative values in the x-axis (in the revised version).

Lastly, we have uploaded a revised version of the paper, which includes additional runtime and ablation experiments, more detailed methodological explanations, and thorough proofreading to improve clarity.

(η,ϵ) definition Our definition is intentionally general, accommodating cases where a matrix is of low rank and meets the definition with a non-zero epsilon rather than requiring epsilon to be zero. In our specific case, epsilon is zero (for the SVD decomposition, as discussed in section 2), but we chose to keep the definition more widely applicable.

• Time comparison: We empirically measure the average time required to calculate the update of the subspace and compare it to the corresponding one in Galore (truncated SVD calculation). We do that w.r.t the “mrpc” experiment task conducted in Table 2 below. Recall that we run the fine-tuning in the three methods with the number of steps (steps for updating the parameters in the entire experiment being the same) in order to maintain fairness in relation to the accuracy obtained (the one that appears in Table 2 below). We used A100 Nvidia for this time measurement. The measurements were done on MRCP task:

Table 2:The table shows the time measurements for subspace computations on the GULE MRPC dataset, comparing AdaRankGard, Galore, and LoRA methods.

From the table, we observe the following: The average time for searching the subspace (adaptive subspace updating) of AdaRankGrad (which is the average time of conducting Algorithm 2, IASS) is approximately 0.1622 milliseconds. In comparison, GaLore without the search process—relying solely on SVD decomposition with a predefined rank—has a significantly higher time cost of approximately 0.6611 milliseconds. However, the number of updates performed by AdaRankGrad is 106.7, whereas GaLore has a fixed number of updates at 35. The effective overhead ratio can be calculated as:

$\text {Overhead Ratio}=\frac{0.1622\times106.7}{0.6611\times35}\approx0.748.$ (*updated)

This indicates that AdaRankGrad does not introduce substantial overhead for adaptive dimension search while offering notable advantages in terms of accuracy and memory efficiency.

Table 3: The table provides a time comparison of the low-rank optimization step (rank=8 for all methods) across AdaRankGard, Galore, and LoRA.

Please remember that in comparison to LORA, AdaRankGrad achieves better accuracy in fine-tuning tasks while using significantly less memory (see Table 4). Furthermore, unlike LORA, AdaRankGrad also supports pre-training. Although the update step for Lora takes longer, it does not have parallel calculation blocks like those used in Glore and AdaGradRank. As a result, the training time of Lora in this experiment was 0.91 times that of AdaGradRank.

Lastly, we note that the average time of Block 2 in AdaRankGrad Moments subspaces transformation is measured with 0.023425 ms (the overhead of that block to the whole training is about 0.023425*106= 2.48305 sec which is neglected)

We would like to thank the reviewer for the thoughtful and detailed feedback. We are encouraged by the reviewer’s positive assessment that the paper has the potential to make a significant algorithmic contribution, pending some clarifications on the theory and experiments. We believe that our responses below address all of the reviewer’s questions and concerns and are organized according to the section numbering provided.

Response to the reviewer’s noted weaknesses:

Memory and time analysis.

We acknowledge that the SSRF-based subspace selection in AdaRankGrad introduces some additional runtime compared to fixed-rank approaches like LoRA. However, the adaptive mechanism in AdaRankGrad is specifically designed to minimize frequent re-computation by updating the subspace only when the gradient convergence criteria are met. This approach reduces unnecessary computations.

Our results demonstrate that, despite these adaptive updates, AdaRankGrad maintains overall memory efficiency without significantly increasing runtime. While continuous SSRF updates might suggest a trade-off between runtime and memory allocation, we have empirically observed that the adaptive frequency of rank updates allows AdaRankGrad to retain allocated memory for only a marginally longer duration than LoRA. We also plan to explore optimizations that could further reduce the frequency of SSRF updates, which would enhance runtime efficiency while supporting the reported memory gains.

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 methods are applied to, in fine-tuning tasks, and with 20% reduction in memory/storage, per the whole final fine-tuned model, as reported in Table 4. Below, we provide a detailed analysis and comparison of the training times for GaLoRE, AdaRankGrad, and LoRA. Our analysis shows that the total training time does not exceed that of GaLoRE and is competitive with LoRA. Specifically, we compare the total training time ratio of LoRA to AdaRankGrad, across all experiments listed in the first row of Table 2 (in the paper), the average ratio is 0.874972 (a more detailed analysis will be presented below). Therefore, accounting for the reduction in memory usage, the increase in running time (compared to LoRA) does not introduce overhead.

We want to emphasize that AdaRankGrad’s memory reduction has practical implications beyond overall efficiency: it enables feasible computation on GPUs with limited memory capacity, which may otherwise be insufficient for training large models. By reducing the memory footprint, AdaRankGrad allows more users to perform fine-tuning or pre-training on consumer-grade GPUs, thereby making advanced model adaptation accessible to a broader range of hardware setups.

Time analysis. Unlike GaLoRE, AdaRankGrad does not compute the exact singular value decomposition (SVD). Instead, it uses an efficient SVD approximation (Algorithm 1) with a complexity of $O(mnr)$, where $m$ and $n$ represent the dimensions of the layer ($m \times n$), and $r$ is the low-rank approximation ($r \ll m, n$). Specifically, in our experiments, $m, n$ are typically in the range of 1000–4000, and $r$ is set to 4, 8, or 16. This approximation significantly reduces the complexity compared to the exact SVD’s $O(\min(mn^2, m^2n))$, with only a negligible impact on error. AdaRankGrad adaptively determines the minimum rank $r \ll m, n$ required to preserve a specified information threshold from the full gradient. This is achieved using a binary search algorithm with a worst-case complexity of $O(\log(r_{max} - r_{min}))$, where $r_{max} \ll m, n$. Consequently, the complexity of this search is an order of magnitude lower than that of a single exact SVD computation, as performed in GaLoRE.In summary, the overall complexity of the estimated projection in AdaRankGrad is $O(\log(r_{max})rmn)$, which is substantially lower than GaLoRE’s $O(\min(mn^2, m^2n))$.

Time analysis clarification

(1) We would be happy to provide clarification on this matter. As indicated in theTable 2 (in our previous comment), the average overall overhead for subspace update time is 17.307 for AdaRankGrad, which is lower compared to 23.138 for Galore, namely AdaRankGrad requires less time for subspace updates, by a factor of 0.748 (and not 1.34). In our previous comment, the numerator and denominator were flipped, but the result was displayed correctly.

$\text {Overhead Ratio} = \frac{0.1622 \times 106.7}{0.6611 \times 35} \approx 0.748.$

(2) To ensure a fair comparison in the experiments, we made sure that all methods perform the same number of optimization steps (single Adam step), differing only in how they utilize this total step "budget." Each step operates on a batch of samples of the same size and spans the same number of epochs. Following is a clarification.

In the GLUE MRPC experiment:

• Galore: performs 30 epochs, with each epoch consisting of approximately 230 optimization steps (calculated based on the total number of samples, 3,668, divided by the batch size, 16). This results in a total of 6,900 optimization steps, during which subspaces are updated every 200 steps, exactly as described in their paper [1].

• AdaRankGrad: performs a total of 6,900 optimization steps with a batch size of 16, similar to Galore. However, unlike Galore, the number of optimization steps applied to each selected subspace adaptively varies (w.r.t convergence criteria). In practice, in this experiment, AdaRankGrad performs a total of 106 subspace updates.

The calculation of the average running time per layer is expressed as:

$\text{Running Time} = (\text{Step Time} + \text{Backpropagation Time}) \times \text{Number of Batches} \times \text{Number of Epochs} + \text{Overall Overhead of Subspace Updates Time}.$

Since the backpropagation time (the execution time for backpropagation on the layer parameters) is identical for both methods, and $\text{Backpropagation Time} \times \text{Number of Batches}=\text{Total optimization steps number}$, the time difference in time between the methods is determined by:

$\Delta \text{Time} = (\text{AdaRankGrad-Step-Time} - \text{Galore-Step-Time})\times \text{Number of Batches} \times\text{Number of Epochs} + (\text{AdaRankGrad Subspace Overhead} - \text{Galore Subspace Overhead}).$

Substituting the specific values from the tables 3 and 4 in our initial comment :

$\Delta \text{Time} = (0.63194-0.63509) \times 230 \times 30 + (17.3075903062-23.13831488)=-27.5657,$ in favor of AdaRankGrad (namely Galore preformed a bit slower).

[1] Zhao, Jiawei, et al. "GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection." arXiv preprint arXiv:2403.03507 (2024).

Memory savings utilization discussion

The correct link is:

https://www.run.ai/blog/dynamic-gpu-memory-solving-the-problem-of-inefficient-resource-allocation-in-inference-servers

Thank you for your suggestion. We have included a discussion on this topic in Appendix Section F of the revised version, as there is no remaining space in the main body of the paper.

We would like to thank the reviewer for suggesting this evaluation. To address your suggestion, we repeated the experiments detailed in Table 2 of the paper, to evaluate the training times of LoRA and AdaRankGrad. For this comparison, we set the rank for both methods to 8 (initial rank for AdaRankGrad) and used the same hyperparameters as those specified in the Table 2 experiments. The table below provides the fine-tuning times for each dataset and method.

Although AdaRankGrad was approximately 7% slower than LoRA, it achieved a memory consumption of more than 60% less compared to LoRA.

(Q4) Ablation studies

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 (in the paper/revised version), with respect to the layers in which the methods are applied and compared on, in fine-tuning tasks, and with 20% reduction in memory/storage, per the whole final fine-tuned model, as reported in Table 4 (in the paper/revised version). 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 - while 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 5 below, where we fix the subspace update interval to 200 steps and study the adaptivity of the subspace dimension.

Table 5: The table presents the results of an ablation experiment in which the subspace update is fixed to intervals of 200 optimization steps while adaptivity in the subspace dimension remains enabled.

Importance of the adaptive subspace update. To assess the impact of the adaptive subspace update on model performance, we conducted the following experiment, as shown in Table 6, where we fixed the rank to a constant value of rank=4 and examined the adaptivity of the subspace updates.

Table 6: The table presents the results of an ablation experiment in which the selected subspace dimension is fixed to $r=4$, and the adaptivity in the subspace updating remains enabled.

The adaptive approach in our method allows optimization within the subspace to converge fully or to exit early when convergence is reached, reducing redundant steps and runtime.

Please note that setting a different number of constant optimization update steps in each selected subspace would not guarantee convergence of the parameters within the subspace, as this approach does not act as a convergence criterion. However, in line with the reviewer's suggestion, we are pleased to present the following table.

Finally, we have uploaded a revised version of the paper with thorough proofreading, additional stylistic improvements, expanded descriptions of optimization algorithms incorporating AdaRankGrad, and new experimental results and comparisons in both the experiments section and appendix.

Unlike GaLore, AdaRankGrad does not compute the exact singular value decomposition (SVD). Instead, it uses an efficient SVD approximation (Algorithm 1) with a complexity of $O(mnr)$, where $m$ and $n$ represent the dimensions of the layer ($m \times n$), and $r$ is the low-rank approximation ($r \ll m, n$). Specifically, in our experiments, $m, n$ are typically in the range of 1000–4000, and $r$ is set to 4, 8. This approximation significantly reduces the complexity compared to the exact SVD’s $O(\min(mn^2, m^2n))$, with only a negligible impact on error.

AdaRankGrad adaptively determines the minimum rank $r \ll m, n$ required to preserve a specified information threshold from the full gradient. This is achieved using a binary search algorithm with a worst-case complexity of $O(\log(r_{max} - r_{min}))$, where $r_{max} \ll m, n$. Consequently, the complexity of this search is an order of magnitude lower than that of a single exact SVD computation, as performed in GaLoRE.In summary, the overall complexity of the estimated projection in AdaRankGrad is $\leq O(\log(r_{max})rmn)$, which is substantially lower than GaLoRE’s $O(\min(mn^2, m^2n))$.

• Time comparison: We empirically measure the average time required to calculate the update of the subspace and compare it to the corresponding one in Galore (truncated SVD calculation). We do that w.r.t the “mrpc” experiment task conducted in Table 2 below. Recall that we run the fine-tuning in the three methods with the number of steps (steps for updating the parameters in the entire experiment being the same) in order to maintain fairness in relation to the accuracy obtained (the one that appears in Table 2 below). We used A100 Nvidia for this time measurement. The measurements were done on MRCP task:

Table 2:The table shows the time measurements for subspace computations on the GULE MRPC dataset, comparing AdaRankGard, Galore, and LoRA methods.

From the table, we observe the following: The average time for searching the subspace (adaptive subspace updating) of AdaRankGrad (which is the average time of conducting Algorithm 2, IASS) is approximately 0.1622 milliseconds. In comparison, GaLore without the search process—relying solely on SVD decomposition with a predefined rank—has a significantly higher time cost of approximately 0.6611 milliseconds. However, the number of updates performed by AdaRankGrad is 106.7, whereas GaLore has a fixed number of updates at 35. The effective overhead ratio can be calculated as:

$\text {Overhead Ratio}=\frac{0.1622\times106.7}{0.6611\times35}\approx0.748.$ (*updated)

This indicates that AdaRankGrad does not introduce substantial overhead for adaptive dimension search while offering notable advantages in terms of accuracy and memory efficiency.

Table 3: The table provides a time comparison of the low-rank optimization step (rank=8 for all methods) across AdaRankGard, Galore, and LoRA.

Please remember that in comparison to LORA, AdaRankGrad achieves better accuracy in fine-tuning tasks while using significantly less memory (see Table 4). Furthermore, unlike LORA, AdaRankGrad also supports pre-training. Although the update step for Lora takes longer, it does not have parallel calculation blocks like those used in Glore and AdaGradRank. As a result, the training time of Lora in this experiment was 0.91 times that of AdaGradRank.

Lastly, we note that the average time of Block 2 in AdaRankGrad Moments subspaces transformation is measured with 0.023425 ms (the overhead of that block to the whole training is about 0.023425*106= 2.48305 sec which is neglected)

We would like to thank the reviewer for the thoughtful and thorough review. We also thank the reviewer for acknowledging the paper's experimental results and theoretical guarantees, which provide solid support for the claim in this paper. We believe the following responses address all the reviewer’s questions and concerns, following the provided section numbering.

(W1) Results of LoRa from previous papers and settings We note that the finetuning results for LoRA reported in the manuscript were borrowed from the GaLore paper [2] (Table 4). After analyzing the gap between the original LoRa paper [1] and GaLore [2], we found that the differences reside in hyperparameters. For example, Roberta-base was trained on SST-2 task with a learning rate 1E-05 for 30 epochs in [2] (Table 7 in [2]), while the same model was trained on the same task with a learning rate 5e-04 for 60 epochs in [1] (Table 9 in [1]). For completeness, we are working on running our own evaluation of LoRa on these datasets.

[1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." arXiv preprint arXiv:2106.09685 (2021).

[2] Zhao, Jiawei, et al. "Galore: Memory-efficient llm training by gradient low-rank projection." arXiv preprint arXiv:2403.03507 (2024).

As the discussion period is nearing its end, we would like to check if there are any remaining concerns regarding our paper. In response to the reviewers’ feedback, we have conducted a runtime evaluation (detailed in the last Table posted in response to reviewer TXWw) that demonstrates our method reduces the training time of Galore while remaining competitive with LoRA, all while significantly reducing memory usage. To strengthen our evaluation, we have also assessed our method's performance on a complex reasoning task using the GSM8K dataset. In these experiments, we applied a rank of 64 for both LoRA and Galore methods, an initial rank of 64 for AdaRankGrad, a batch size of 32, and fine-tuned over 10 epochs. The experiments were conducted on an NVIDIA A100 GPU. Our results, presented in the table below, show that AdaRankGrad is highly effective for this task, achieving improved accuracy compared to both LoRA and Galore, along with reduced memory usage. We appreciate the reviewers’ valuable suggestions and hope these results adequately address the concerns raised.

The average reduction in memory consumption of AdaRankGrad compared to Galore is 24.3%

Thank you for your detailed response. The additional results and explanation address most of my concerns. I would like to improve my score to 6.

(W3) Iteration speed Unlike GaLore, AdaRankGrad does not compute the exact singular value decomposition (SVD). Instead, it uses an efficient SVD approximation (Algorithm 1) with a complexity of $O(mnr)$, where $m$ and $n$ represent the dimensions of the layer ($m \times n$), and $r$ is the low-rank approximation ($r \ll m, n$). Specifically, in our experiments, $m, n$ are typically in the range of 1000–4000, and $r$ is set to 4, 8. This approximation significantly reduces the complexity compared to the exact SVD’s $O(\min(mn^2, m^2n))$, with only a negligible impact on error.

AdaRankGrad adaptively determines the minimum rank $r \ll m, n$ required to preserve a specified information threshold from the full gradient. This is achieved using a binary search algorithm with a worst-case complexity of $O(\log(r_{max} - r_{min}))$, where $r_{max} \ll m, n$. Consequently, the complexity of this search is an order of magnitude lower than that of a single exact SVD computation, as performed in GaLoRE.In summary, the overall complexity of the estimated projection in AdaRankGrad is $\leq O(\log(r_{max})rmn)$, which is substantially lower than GaLoRE’s $O(\min(mn^2, m^2n))$.

• Time comparison: We empirically measure the average time required to calculate the update of the subspace and compare it to the corresponding one in Galore (truncated SVD calculation). We do that w.r.t the “mrpc” experiment task conducted in Table 2 below. Recall that we run the fine-tuning in the three methods with the number of steps (steps for updating the parameters in the entire experiment being the same) in order to maintain fairness in relation to the accuracy obtained (the one that appears in Table 2 below). We used A100 Nvidia for this time measurement. The measurements were done on MRCP task:

Table 2:The table shows the time measurements for subspace computations on the GULE MRPC dataset, comparing AdaRankGard, Galore, and LoRA methods.

From the table, we observe the following: The average time for searching the subspace (adaptive subspace updating) of AdaRankGrad (which is the average time of conducting Algorithm 2, IASS) is approximately 0.1622 milliseconds. In comparison, GaLore without the search process—relying solely on SVD decomposition with a predefined rank—has a significantly higher time cost of approximately 0.6611 milliseconds. However, the number of updates performed by AdaRankGrad is 106.7, whereas GaLore has a fixed number of updates at 35. The effective overhead ratio can be calculated as:

$\text {Overhead Ratio}=\frac{0.1622\times106.7}{0.6611\times35}\approx0.748.$ (*updated)

This indicates that AdaRankGrad does not introduce substantial overhead for adaptive dimension search while offering notable advantages in terms of accuracy and memory efficiency.

Table 3: The table provides a time comparison of the low-rank optimization step (rank=8 for all methods) across AdaRankGard, Galore, and LoRA.

Please remember that in comparison to LORA, AdaRankGrad achieves better accuracy in fine-tuning tasks while using significantly less memory (see Table 4). Furthermore, unlike LORA, AdaRankGrad also supports pre-training. Although the update step for Lora takes longer, it does not have parallel calculation blocks like those used in Glore and AdaGradRank. As a result, the training time of Lora in this experiment was 0.91 times that of AdaGradRank.

Lastly, we note that the average time of Block 2 in AdaRankGrad Moments subspaces transformation is measured with 0.02342 ms (the overhead of that block to the whole training is about 0.02342*106= 2.48305 sec which is neglected)

We would like to thank the reviewer for this thoughtful and detailed review. We appreciate the recognition of the article's solid theoretical foundation (showing gradients decay to a rank-1 structure over the course of training) and its significant contributions to memory efficiency through cost-effective updates of an adaptive subspace for gradients and moments projection and novel moment subspace transformation adjustment, which the reviewer found also results in competitive performance against the state-of-the-art. We believe the responses below address the reviewer’s questions and the noted weakness.

(W1) additional ablation studies

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 (in the paper/revised version), with respect to the layers in which the methods are applied and compared on, in fine-tuning tasks, and with 20% reduction in memory/storage, per the whole final fine-tuned model, as reported in Table 4 (in the paper/revised version). 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 - while 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 5 below, where we fix the subspace update interval to 200 steps and study the adaptivity of the subspace dimension.

Table 5: The table presents the results of an ablation experiment in which the subspace update is fixed to intervals of 200 optimization steps while adaptivity in the subspace dimension remains enabled.

Importance of the adaptive subspace update. To assess the impact of the adaptive subspace update on model performance, we conducted the following experiment, as shown in Table 6, where we fixed the rank to a constant value of rank=4 and examined the adaptivity of the subspace updates.

Table 6: The table presents the results of an ablation experiment in which the selected subspace dimension is fixed to $r=4$, and the adaptivity in the subspace updating remains enabled.

In summary, the adaptive approach in our method allows optimization within the subspace to converge fully or to exit early when convergence is reached, reducing redundant steps and runtime. Finally, we have uploaded a revised version of the paper with thorough proofreading, additional stylistic improvements, expanded descriptions of optimization algorithms incorporating AdaRankGrad, and new experimental results and comparisons in both the experiments section and appendix.

We have uploaded a revised version of the paper with thorough proofreading, additional stylistic improvements, expanded descriptions of optimization algorithms incorporating AdaRankGrad, and new experimental results and comparisons in both the experiments section and appendix.

(Q2) Impact of threshold $\eta_{th}$. The information threshold $\eta_{th}$ is indeed a hyperparameter that influences the effective rank by controlling the amount of gradient information preserved during updates. Therefore, it is a desired property of this parameter that the effective rank gradually decreases with an increase in $\eta_{th}$ (as shown in Figure 6 in the Appendix). Notice, however, that as seen in Figure 7 in the appendix, changes to the information threshold $\eta_{th}$ have a minimal impact on performance (F1-score), with fluctuations of no more than 1.7 percent. A proper way to use this parameter is to select a memory budget and set $\eta_{th}$ so that a suitable rank is achieved, corresponding to the desired memory.

(Q3) Gradient dimensionality reduction. AdaRankGrad incorporates an adaptive low-rank projection mechanism that dynamically adjusts the approximated-gradient-rank based on preserving-information criteria, allowing the model to respond automatically to changes in gradient rank throughout training. We have demonstrated that in gradient-based optimization, the dimensionality of layer gradients in LLM training naturally decreases, asymptotically approaching 1. Given a selected information threshold, each subspace update (outer loop in Algorithm 3) determines the effective rank of the subspace on which the gradient will be projected, preserving the relative information content of the full gradient. Thus, as the rank of the full gradient naturally reduces during training, the effective rank for the projected gradient automatically decreases as well, with no need for manual adjustment (see Algorithm 2, line $\eta_t \leftarrow \| \mathbf{A} - \mathbf{A}_{\text{app}, r_t} \|_F / \| \mathbf{A} \|_F$ ). The convergence threshold $\varsigma_2$ is set as constant, however, in Algorithm 3 can optionally be automatically adjusted based on the relative adaptive rank. For instance, if the initial rank is 8 and the rank of the selected subspace decreases to 6, the threshold is updated as $\varsigma_2 = \frac{6\varsigma_2}{8}$. This means convergence for each gradient direction is determined by its value falling below a predefined fixed threshold. As a result, the convergence threshold for the gradient norm is scaled according to the dimensionality of the gradient space, eliminating the need for manual parameter adjustment. In conclusion, no manual adjustment of parameters is required, as the adaptive parameters are automatically updated through the described mechanisms.

Finally, we have uploaded a revised version of the paper with thorough proofreading, additional stylistic improvements, expanded descriptions of optimization algorithms incorporating AdaRankGrad, and new experimental results and comparisons in both the experiments section and appendix.

(W3) Iteration speed Unlike GaLore, AdaRankGrad does not compute the exact singular value decomposition (SVD). Instead, it uses an efficient SVD approximation (Algorithm 1) with a complexity of $O(mnr)$, where $m$ and $n$ represent the dimensions of the layer ($m \times n$), and $r$ is the low-rank approximation ($r \ll m, n$). Specifically, in our experiments, $m, n$ are typically in the range of 1000–4000, and $r$ is set to 4, 8. This approximation significantly reduces the complexity compared to the exact SVD’s $O(\min(mn^2, m^2n))$, with only a negligible impact on error.

AdaRankGrad adaptively determines the minimum rank $r \ll m, n$ required to preserve a specified information threshold from the full gradient. This is achieved using a binary search algorithm with a worst-case complexity of $O(\log(r_{max} - r_{min}))$, where $r_{max} \ll m, n$. Consequently, the complexity of this search is an order of magnitude lower than that of a single exact SVD computation, as performed in GaLoRE.In summary, the overall complexity of the estimated projection in AdaRankGrad is $\leq O(\log(r_{max})rmn)$, which is substantially lower than GaLoRE’s $O(\min(mn^2, m^2n))$.

• Time comparison: We empirically measure the average time required to calculate the update of the subspace and compare it to the corresponding one in Galore (truncated SVD calculation). We do that w.r.t the “mrpc” experiment task conducted in Table 2 below. Recall that we run the fine-tuning in the three methods with the number of steps (steps for updating the parameters in the entire experiment being the same) in order to maintain fairness in relation to the accuracy obtained (the one that appears in Table 2 below). We used A100 Nvidia for this time measurement. The measurements were done on MRCP task:

Table 2:The table shows the time measurements for subspace computations on the GULE MRPC dataset, comparing AdaRankGard, Galore, and LoRA methods.

From the table, we observe the following: The average time for searching the subspace (adaptive subspace updating) of AdaRankGrad (which is the average time of conducting Algorithm 2, IASS) is approximately 0.1622 milliseconds. In comparison, GaLore without the search process—relying solely on SVD decomposition with a predefined rank—has a significantly higher time cost of approximately 0.6611 milliseconds. However, the number of updates performed by AdaRankGrad is 106.7, whereas GaLore has a fixed number of updates at 35. The effective overhead ratio can be calculated as:

$\text {Overhead Ratio}=\frac{0.1622\times106.7}{0.6611\times35}\approx0.748.$ (*updated)

This indicates that AdaRankGrad does not introduce substantial overhead for adaptive dimension search while offering notable advantages in terms of accuracy and memory efficiency.

Table 3: The table provides a time comparison of the low-rank optimization step (rank=8 for all methods) across AdaRankGard, Galore, and LoRA.

Please remember that in comparison to LORA, AdaRankGrad achieves better accuracy in fine-tuning tasks while using significantly less memory (see Table 4). Furthermore, unlike LORA, AdaRankGrad also supports pre-training. Although the update step for Lora takes longer, it does not have parallel calculation blocks like those used in Glore and AdaGradRank. As a result, the training time of Lora in this experiment was 0.91 times that of AdaGradRank.

Lastly, we note that the average time of Block 2 in AdaRankGrad Moments subspaces transformation is measured with 0.02342 ms (the overhead of that block to the whole training is about 0.02342*106= 2.48305 sec which is neglected)

We thank the reviewer for referring us to this work. While both BlockLLM [1] and our method aim to reduce the memory requirements in training LLMs, there are significant differences between the methods, as listed below:

• Algorithmically:

  • The dimensionality reduction - in BlockLLM, a subset of the parameters is selected, and the model parameters are updated in this space (by e.g., ADAM) until a stopping criteria prompts trainable subspace update. In our method (AdaRankGrad), a subspace is selected via a projection until a stopping criterion prompts a trainable subspace update. Note that in AdaRankGrad, it is possible to have the same subspace as chosen in BlockLLM, however, the converse is not necessarily true.

  • The sparsity/dimension of the trainable subspace - In BlockLLM, the sparsity is a hyperparameter. In contrast, in AdaRankGrad, the dimension of each subspace is determined by an information threshold, which is also a hyperparameter that directly relates to information loss.

• Theoretically

  • Our method (AdaGradRank) is provided with a theoretical convergence analysis. In particular, the method is motivated by an analysis of gradient rank (which shows that layers rank asymptotically goes to rank one). In contrast, in [1], theoretical analysis of BlockLLM is not provided, though, as mentioned in [1], the method is inspired by block coordinate descent, hence convergence of BlockLLM might be analyzed similarly to block coordinate descent with deterministic importance sampling.

Following the reviewer's suggestion, we present an empirical comparison between our method and BlockLLM. Table 3 in our paper and Table 4 in [1] show the results of pretraining the model by our method on the C4 dataset. As indicated by the results, our method outperforms BlockLLM in terms of perplexity, as shown in the next table:

Table 1: Perplexity for pre-trained models on the C4 dataset:

We thank the reviewer for this thoughtful and comprehensive review! We also thank the reviewer for acknowledging that our paper is well-motivated and well-written, grounded in solid theory, and demonstrates competitive results against state-of-the-art methods. Regarding the weaknesses pointed out, we believe that our responses below thoroughly address all concerns and questions.

(W1) Performance gain by the method

• As for memory reduction: By eliminating the need to predefine both the gradient rank and the fixed cycle length for subspace updates, AdaRankGrad simplifies the training process and reduces memory usage by over 25%-50% on average compared to GaLore (and LoRA), as shown in Figure 4 (in the paper/revised version), with respect to the layers which the methods are applied and compared on, in fine-tuning tasks, and with 20% reduction in memory/storage, per the whole pre-training model, as reported in Table 4 (in the paper/revised version). This adaptability allows AdaRankGrad to allocate memory efficiently where it’s needed during training, achieving significant gains in memory efficiency without compromising model performance and making it a valuable, scalable solution for diverse fine-tuning and training tasks. We believe these improvements will enhance the usability of low-rank gradient-based methods and pave the way for future advancements and memory efficiency of LLM training.
• As for accuracy gain: AdaRankGrad demonstrates a consistently higher accuracy across a range of tasks, including improvements in tasks such as, RTE and MRPC, where adaptive low-rank gradient adjustment is especially effective (Table 2 in the paper/revised version). Significant gains in accuracy are also demonstrated compared to LoRA on the tabular dataset (see, Figure 5).