Efficient Fine-Tuning of Large Language Models via a Low-Rank Gradient Estimator

Abstract

In this paper, we present a Low-Rank Gradient Estimator (LoGE) to accelerate the finetune-time computation of transformers, especially large language models (LLMs). Unlike Parameter-Efficient Fine-Tuning (PEFT) methods, which primarily aim to minimize the number of fine-tuning parameters, LoGE also significantly reduces the computational load of activation gradient calculations by decomposing pre-trained weights and utilizing low-rank matrices during the backward pass. Our approach includes an effective solution for identifying sensitive and important latent subspaces in large models before training with downstream datasets. As LoGE does not alter the network structure, it can be conveniently integrated into existing models. We validated LoGE’s efficacy through comprehensive experiments across various models on various tasks. For the widely used LLaMA model equipped with LoRA, LoGE achieves up to a 1.3× speedup while maintaining graceful accuracy.

1. Introduction

Recently, large language models (LLMs) [1,2,3,4,5,6,7] have demonstrated exceptional capabilities in various language understanding tasks [8], becoming the backbone of leading chat systems [7]. To investigate the capabilities of LLMs further, fine-tuning them for specific downstream tasks has become a popular approach. Although the dataset size for fine-tuning a large language model is significantly smaller than that required for pre-training, the computational overhead of fine-tuning becomes comparable to pre-training due to the growing number of downstream tasks and their increasing complexity. Using LLaMA-65B [5] as an example, fine-tuning the model on a 10B-token dataset for 300 downstream tasks would generate approximately 24 tons of CO₂, posing significant challenges in terms of computational demands and environmental impact.

To enhance memory efficiency and reduce computational overhead during fine-tuning, Parameter-Efficient Fine-Tuning (PEFT) [9,10,11,12,13,14,15,16,17,18,19] significantly reduces the number of trainable parameters, thereby decreasing the memory required for gradients and optimization states. Specifically, the Low-Rank Adapter (LoRA) [17] employs a low-rank approximation to represent weight updates and integrates these updates during inference, introducing no additional computational overhead. PEFT methods reduce the number of trainable parameters, thereby decreasing the computation of weight gradients. However, they still require full-size matrix computations for forward passes and activation gradients, requiring a large amount of calculation.

To address the above challenge, we propose the Low-Rank Gradient Estimator (LoGE), which utilizes low-rank computation during the backward pass while maintaining full-rank computation during the forward pass. Our key idea is to leverage the low-rank structure of the weight $\mathbf{W}$ rather than approximating the weight updates as low rank. In LoGE, we apply decomposition to the full-rank weights and utilize the resultant low-rank matrices to calculate the input gradient, as illustrated in Figure 1. By decomposing full-rank weights into low-rank matrices (denoted as $\mathbf{B}$ and $\mathbf{A}$), we can interpret weight updates as the update of these low-rank matrices. We fix one low-rank matrix (typically denoted as $\mathbf{B}$) to avoid full-rank updates and maintain stability during training. To reduce the computational load during the forward pass while preserving the adaptor’s functionality, we propose a trainable parameter $\mathbf{C}$ as the update to $\mathbf{A}$ and restrict updates solely to $\mathbf{C}$. At the start of training, $\mathbf{C}$ is initialized to $\mathbf{0}$ to align with the pre-trained weights. To identify task-specific subspaces, we utilize data from the downstream task. This allows us to select the optimal subspace for optimization. By combining both gradient similarity and gradient increment values, we introduce a novel importance metric. This metric aids in the selection of the most relevant subspace for the fine-tuning process.

Overall, our contributions are as follows:

• We propose a new efficient fine-tuning method for LLMs by decomposing the pre-trained weight into a low-rank structure in the backward pass while maintaining the full-rank weight in the forward pass.
• We propose a metric to choose the subspace via task-specific gradient scores and weight salience.
• We evaluate our method on different models and different tasks to demonstrate the effectiveness of our proposed method and provide detailed analyses based on the empirical results. Results show that LoGE achieves better performance with less training time cost.

2. Related Work

Large Language Models (LLMs) [1,2,3,4,5,6,7] have become a leading paradigm in natural language processing, achieving top results in numerous tasks [20,21] and forming the basis of chat systems [22]. To maximize the downstream performance, fine-tuning LLMs is essential. However, the size of LLMs has dramatically increased to enhance performance, exemplified by GPT-3, which boasts up to 175 billion parameters and requires approximately 350 GB of memory in FP16 format. This growth poses significant challenges for fine-tuning these models in a memory- and calculation-efficient manner.

Parameter-Efficient Fine-Tuning (PEFT) [9,10,11,12,13] methods are meant to address this problem to reduce the number of trainable parameters by various methods while maintaining performance. PEFT typically involves adding adapter layers between the existing layers of a neural network. The pre-trained weights are frozen, and only the adapter layers are trained [9]. Prefix tuning [14] adds prefix parameters to the hidden states in every layer of the model. Prompt tuning [15] uses the template to reconstruct prompts and only updates the parameters related to prompt understanding. Low-Rank Adaptation (LoRA) [13] is a popular method that decomposes adapter weights into two low-rank matrices. These matrices are then fused with the pre-trained weights during inference. LoRA has demonstrated comparable performance to full-parameter fine-tuning while using much fewer parameters. AdaLoRA [16] utilizes SVD decomposition and prunes less significant singular values for more efficient updates. Another method uses orthogonal factorization to fine-tune diffusion models. LoRA-FA [17] keeps matrix A fixed with random initialization and only trains matrix B. Child Tuning [18] and SPT [19] only update a subset of parameters by strategically masking the gradient of the parameters not in the subset. Our method uses low-rank matrices to represent the subset of parameters in latent space, which allows for accelerating the calculation of backward computation without masking.

Quantization-Based PEFT. Although PEFT reduces the trainable parameters in models, the pre-trained weights still require a large amount of memory. To address this challenge, QLoRA [23] compresses the pre-trained weights into 4-bit floats and integrates LoRA. Meanwhile, QA-LoRA [24] and EfficientDM [25] address the quantization challenges in LoRA-fused models. However, QLoRA and QALoRA introduce extra dequantization overhead in the calculation. In contrast, our approach eliminates the need for dequantization in the backward pass, thereby achieving significant speed improvements over these memory-efficient LoRA variants.

Low-Rank Structures in Deep Learning. Low rank is widely used in machine learning, as many machine learning problems [26,27,28,29] have certain intrinsic low-rank structures. Particularly in over-parameterized neural networks, it is observed that the networks tend to exhibit low-rank properties post-training [30]. Furthermore, several studies have explored low-rank updates to frozen models. For instance, ReLoRA [31] extends low-rank updates to encompass the entire training process of the original networks, and GaLore [32] employs low-rank decompositions of weight gradients to reduce memory demands during full training. To our knowledge, no existing work applies low-rank decomposition directly to the weights during training to reduce computational load. In our research, we decompose pre-trained weights to expedite gradient calculations.

3. Background

Naive Full-Rank Training. For a given weight ${\mathbf{W}}_{\mathbf{0}}\in {\mathbb{R}}^{m\times n}$, the conventional update can be expressed as follows:

$${\mathbf{W}}_{\mathbf{T}}={\mathbf{W}}_{\mathbf{0}}+\eta \sum _0^{\mathbf{T}-1}{\mathbf{G}}_{\mathbf{t}}$$

(1)

where ${\mathbf{G}}_{\mathbf{t}}$ represents the gradient added to the weight matrix at step t. In training, the backward pass is both memory-intensive and computationally demanding. For example, with Adam, gradients ${\mathbf{G}}_{\mathbf{t}}$ and additional optimizer states $\mathbf{M},\mathbf{V}\in {\mathbb{R}}^{m\times n}$ are used to regularize the gradient, increasing memory requirements by $3mn$. Gradient calculations involve two full-rank operations:

$$\partial \mathbf{O}/\partial \mathit{x}={\delta }_{\mathbf{O}}\mathbf{W}\partial \mathbf{O}/\partial \mathbf{W}={\delta }_{\mathbf{O}}^T\mathit{x}$$

(2)

Here, $\mathit{x}\in {\mathbb{R}}^{b\times n}$ is the input, and ${\delta }_{\mathbf{O}}\in {\mathbb{R}}^{b\times m}$ represents the output error. This naive approach incurs $4bmn$ FLOPs during the backward pass.

Low-Rank Adaptation. For a linear layer $\mathbf{W}\in {\mathbb{R}}^{m\times n}$, LoRA and its variants model the incremental update of the pre-trained weight with the product of two low-rank matrices $\mathbf{AB}$:

$$\mathbf{W}={\mathbf{W}}_{\mathbf{0}}+\mathbf{BA},$$

(3)

where $\mathbf{W},{\mathbf{W}}_{\mathbf{0}}\in {\mathbb{R}}^{m\times n}$, $\mathbf{A}\in {\mathbb{R}}^{r\times n}$, and $\mathbf{B}\in {\mathbb{R}}^{m\times r}$, with $r\ll min(n,m)$. With Adam, LoRA needs $3mr+3nr$ memory to store optimization states and gradients. The gradient calculation for weight involves two low-rank operations:

$$\partial \mathbf{O}/\partial \mathbf{A}={\left({\delta }_{\mathbf{O}}\mathbf{B}\right)}^T\mathit{x}\partial \mathbf{O}/\partial \mathbf{B}={\delta }_{\mathbf{O}}^T\left(\mathit{x}\mathbf{A}\right)$$

(4)

During backpropagation, LoRA simplifies FLOPs to $2bmn+6bmr+6bnr$, thereby reducing the calculation.

Quantization. Quantization is the process of mapping continuous or high bit-width data types into discrete and often lower bit-width values. For instance, converting a 16-bit floating point (FP16) tensor into a 4-bit integer (INT4) tensor with a range of [−7, 8] can be expressed as follows:

$${\mathbf{X}}^{\mathbf{INT}\mathbf{4}}=\mathrm{round}\left({\displaystyle \frac{15}{\mathrm{absmax}\left({\mathbf{X}}^{\mathbf{FP}\mathbf{16}}\right)}}{\mathbf{X}}^{\mathbf{FP}\mathbf{16}}\right)=\mathrm{round}\left({\displaystyle \frac{{\mathbf{X}}^{\mathbf{FP}\mathbf{16}}}{{\mathit{s}}^{\mathbf{FP}\mathbf{16}}}}\right)$$

(5)

where $\mathit{s}$ is the quantization scale. Dequantization is the inverse:

$$\mathrm{dequant}({\mathit{s}}^{\mathbf{FP}\mathbf{16}},{\mathbf{X}}^{\mathbf{INT}\mathbf{4}})={\mathit{s}}^{\mathbf{FP}\mathbf{16}}{\mathbf{X}}^{\mathbf{INT}\mathbf{4}})={\widehat{\mathbf{X}}}^{\mathbf{FP}\mathbf{16}}$$

(6)

Here, $\widehat{\mathbf{X}}$ is an estimated value, as quantization will introduce quantization error for original values. Quantization is introduced to reduce the weight memory usage for LLMs, as PEFT reduces the learnable parameters, thereby reducing memory usage for weight gradients.

4. Method

In many cases, the optimal weight may not be low-rank. Therefore, prior works [31,32] choose to maintain the full-rank weight, and update the full-rank weight with a dynamic low-rank approximate to enhance training efficiency. If we can estimate the full-rank weight with a low-rank approximation during the backward pass while maintaining the full-rank weight in the forward pass, we can further enhance training efficiency. This leads to our proposed LoGE method.

4.1. Low-Rank Gradient Estimator

The core idea of LoGE is to replace the full-rank weight in the backward pass with a low-rank representation and update one low-rank matrix. In the backward pass, we substitute the full-rank matrix with its low-rank approximation to compute the gradient of the activation. However, updating both components of the low-rank approximation prevents the weight update from being directly represented by the updates of the components. Given the low-rank approximation $\mathbf{B}\mathbf{A}$, we can formulate the update of weight ${\delta }_{\mathbf{W}}$ as follows:

$${\delta }_{\mathbf{W}}=(\mathbf{B}+{\delta }_{\mathbf{B}})(\mathbf{A}+{\delta }_{\mathbf{A}})-\mathbf{B}\mathbf{A}={\delta }_{\mathbf{B}}\mathbf{A}+\mathbf{B}{\delta }_{\mathbf{A}}+{\delta }_{\mathbf{B}}{\delta }_{\mathbf{A}}$$

(7)

The ${\delta }_{\mathbf{W}}$ requires full-rank storage, making this approximation memory-inefficient. Inspired by [17], we try to fix one low-rank matrix and only update the other. This approach maintains the update in a low-rank format, thereby reducing the memory usage required for optimization states and gradients. For simplicity in computation, we frame the low-rank update in the adapter format, following the design principles of LoRA [13]. Given a full-rank weight $\mathbf{W}\in {\mathbb{R}}^{m\times n}$ and $n\ge m$, the calculation for LoGE is defined as follows:

$$\mathbf{O}=({\mathbf{W}}_0+\mathbf{B}\mathbf{A})\mathit{x}\partial \mathbf{O}/\partial \mathit{x}={\delta }_{\mathbf{O}}\mathbf{B}(\mathbf{C}+\mathbf{A})$$

(8)

Here, $\mathbf{B}\in {\mathbb{R}}^{m\times r}$ and $\mathbf{A}\in {\mathbb{R}}^{r\times n}$ represent the low-rank approximations of $\mathbf{W}$ used during the backward pass, with $\mathbf{A}$ fixed during training. The matrix $\mathbf{C}\in {\mathbb{R}}^{r\times n}$ is the trainable parameters and represents the updates to $\mathbf{A}$. For cases where $n<m$, we fix the $\mathbf{A}$ and reformulate the calculations to the following:

$$\mathbf{O}=(\mathbf{W}+\mathbf{C}\mathbf{A})\mathit{x}\partial \mathbf{O}/\partial \mathit{x}={\delta }_{\mathbf{O}}(\mathbf{B}+\mathbf{C})\mathbf{A}$$

(9)

The subsequent sections will concentrate on the case where $n\ge m$ to explore LoGE further. Orthogonal decomposition matrices can simplify the calculation of importance for structured sparsity. Therefore, we utilize Singular Value Decomposition (SVD) to decompose the matrix. Given the pre-trained weight ${\mathbf{W}}_0$, this process can be formulated as follows:

$${\mathbf{W}}_0=\mathbf{U}\mathsf{\Sigma }{\mathbf{V}}^T\mathbf{B}={\mathbf{U}}_{\left[idx\right]}\mathbf{C}={\mathsf{\Sigma }}_{\left[idx\right]}{\mathbf{V}}_{\left[idx\right]}^T$$

(10)

Here, $idx$ refers to the selected index vector. The select strategy is discussed in the next subsection.

4.2. Salient Space Choose Strategy

For fine-tuning, the objective is to integrate new knowledge into the pre-trained weights. Therefore, identifying weights that are sensitive to the fine-tuning dataset is crucial. Following the core principle from Child Tuning [18], our approach considers the singular values and the importance of gradients. This dual consideration helps us design a metric to select the most relevant subspaces for updates.

Obtaining the gradient importance of each subspace is challenging, as it requires training models to converge, which is time-consuming. Following prior work [18], we approximate $\mathbf{C}$ using the gradient of $\mathbf{A}$ after a single step of gradient descent.

Since the latent space corresponds to a pair from the weight decomposition, using the full-rank decomposition for gradient descent calculations could lead to large calculation and memory demands. To solve this problem, we sparse the $\mathbf{U},\mathsf{\Sigma }$, and ${\mathbf{V}}^T$ by selecting the top $2r$ singular values and then compute the gradient-based importance metric.

While Child Tuning [18] selects subspaces based on element-wise importance, our method focuses on identifying important channels from the SVD decomposition. We use an importance score for the entire channel rather than for each individual element. Inspired by structural pruning [33], we utilize the second-order Taylor expansion to determine the importance of each channel. Given the channel gradient $g_{\mathbf{A}}=\left\{g_{a_i}\right\}$, the gradient importance score can be formulated as follows:

$$s_{a_i}=g_{a_i}^T{g}_{a_i}$$

(11)

As low-rank matrices estimate the gradients of activations, we also consider the magnitude of singular values. Therefore, we use the energy of the singular values (the square of the singular values) to represent their importance. We then normalize these two metrics and combine them using a temperature parameter:

$$S_i={\displaystyle \frac{\left|\right|{\lambda }_i^2\left|\right|}{max\left(\right|\left|{\lambda }_i^2\right|\left|\right)}}+{\displaystyle \frac{1}{\alpha }}{\displaystyle \frac{\left|\right|s_{a_i}\left|\right|}{max\left(\right|\left|s_{a_i}\right|\left|\right)}}$$

(12)

Here, $\alpha$ serves as a temperature parameter that helps balance differing normalization values. We present a detailed ablation study to compare the performance of different importance metrics and different $\alpha$.

The entire fine-tuning process with LoGE is outlined in Algorithm 1. In practice, the batch size b is set to 32. The time cost for calculating the importance score is approximately 20 s for the OPT-1.3B model on the Alpaca dataset.

Algorithm 1 LoGE for fine-tuning. $\theta$ represents the model parameters, and ${\mathcal{D}}_k$ is the dataset for task k.

Require:  $\theta$, ${\mathcal{D}}_k$ 1: for layer in model layers do 2:    if layer is linear then 3:        $\mathbf{U},\mathsf{\Sigma },{\mathbf{V}}^T\leftarrow {\mathbf{W}}_{\mathbf{0}}$ 4:        layer ← LoGE(${\mathbf{W}}_{\mathbf{0}},{\mathbf{A}}^{\prime },\mathbf{U},\mathsf{\Sigma }{\mathbf{V}}^T$) 5:    end if 6: end for 7: Sample b samples $d_k$ from ${\mathcal{D}}_k$. 8: Calculate gradient of ${\mathbf{A}}^{\prime }$ with $d_k$. 9: Calculate importance index $S_i$ and select $idx$. 10: for layer in model layers do 11:     if layer is linear then 12:        layer ← LoGE(${\mathbf{W}}_{\mathbf{0}},\mathbf{A},{\mathbf{U}}_{\left[idx\right]},{\mathsf{\Sigma }}_{\left[idx\right]}{\mathbf{V}}_{\left[idx\right]}^T$) 13:     end if 14: end for 15: for s in training steps do 16:     Update $\mathbf{A}$ {Training step with LoGE} 17: end for 18: return  $\theta$

4.3. Memory Usage and Calculation

Since LoGE limits the trainable parameters with a low-rank matrix rather than a full-rank one, the memory usage of the optimization states and gradient is reduced to $3nr$. Combined with the parameters, the total memory usage, excluding activation for linear layers, is $mn+mr+5nr$. Furthermore, the smaller trainable parameters also reduce the communication costs for gradient broadcast in distributed training. For the calculation in the backward pass, the calculation for the gradient of the input is reduced from $bmn$ to $bmr+bnr+nr$. For better comparison, we list the exact memory usage and FLOPs in

Table 1. Comparative analysis of training techniques: Naive, LoRA, GaLore, and LoGE. Assume $\mathbf{W}\in {\mathbb{R}}^{m\times n}$, $\mathit{x}\in {\mathbb{R}}^{b\times n}$, and rank r.

	Naive	GaLore	LoRA	LoGE
Weights	$mn$	$mn$	$mn+mr+nr$	$mn+mr+2nr$
Gradients	$mn$	$mn$	$mr+nr$	$nr$
Optim States	$2mn$	$mr+2nr$	$2mr+2nr$	$2nr$
Memory Usage	$4mn$	$2mn+mr+2nr$	$mn+4mr+4nr$	$mn+mr+5nr$
$FLOP{s}_{forward}$	$2bmn$	$2bmn$	$2bmn+2bmr+2bnr$	$2bmn+2bmr+2bnr$
$FLOP{s}_{backward}$	$4bmn$	$4bmn$	$2bmn+4bmr+4bnr$	$2bmr+4bnr+nr$
$FLOP{s}_{total}$	$6bmn$	$6bmn$	$4bmn+6bmr+6bnr$	$2bmn+4bmr+6bnr+nr$

. In the case of fine-tuning an LLaMA-7B model with a rank of 128, the FLOPs for a linear layer under traditional full training are approximately 100.67 M. With the application of LoRA, these FLOPs decrease to 73.4 M. Notably, when employing LoGE and fine-tuning the layers with a larger rank of 384, the FLOPs were further reduced to 49.28 M, which is significantly lower than those achieved with LoRA.

4.4. QLoGE: LoGE with Quantization

To enhance memory efficiency, we implement quantization for the frozen weights. Although quantization training may increase overall training time [23], the trade-off between memory usage and training speed makes it particularly suitable for fine-tuning. Consequently, we apply QLoGE within the fine-tuning methods presented in this paper.

For quantization-based memory-efficient fine-tuning, a significant drawback is the additional computational overhead introduced by dequantizing weights back to the calculation data type. Since fine-tuning often involves long sequence lengths for a single linear layer, this process becomes compute-bound [23], consequently slowing down training speed. The extra dequantization overhead occurs in two areas: during forward calculations and when calculating activation gradients, which can be expressed as follows:

$$\begin{array}{c}\hfill {\mathbf{O}}^{\mathbf{FP}\mathbf{16}}=(\mathrm{dequant}({\mathit{s}}^{\mathbf{FP}\mathbf{16}},{\mathbf{W}}^{\mathbf{NF}\mathbf{4}})+{\mathbf{B}}^{\mathbf{FP}\mathbf{16}}{\mathbf{A}}^{\mathbf{FP}\mathbf{16}}){\mathit{x}}^{\mathbf{FP}\mathbf{16}}\\ \hfill \partial {\mathbf{O}}^{\mathbf{FP}\mathbf{16}}/\partial {\mathit{x}}^{\mathbf{FP}\mathbf{16}}={\delta }_{{\mathbf{O}}^{\mathbf{FP}\mathbf{16}}}(\mathrm{dequant}({\mathit{s}}^{\mathbf{FP}\mathbf{16}},{\mathbf{W}}^{\mathbf{NF}\mathbf{4}})+{\mathbf{B}}^{\mathbf{FP}\mathbf{16}}{\mathbf{A}}^{\mathbf{FP}\mathbf{16}})\end{array}$$

(13)

By replacing the quantized weight in the activation gradient calculation, our method can eliminate one source of this overhead. Since the task loss is calculated with quantized weight, we decompose the simulated weight $\widehat{\mathbf{W}}$ into a low-rank form. The strategy for QLoGE can be formulated as follows:

$$\begin{array}{c}\hfill \widehat{\mathbf{W}}=\mathbf{U}\mathbf{S}{\mathbf{V}}^T{\mathbf{B}}_0={\mathbf{U}}_{\left[idx\right]}{\mathbf{C}}_0={\mathbf{S}}_{\left[idx\right]}{\mathbf{V}}_{\left[idx\right]}^T\\ \hfill \partial {\mathbf{O}}^{\mathbf{FP}\mathbf{16}}/\partial {\mathit{x}}^{\mathbf{FP}\mathbf{16}}={\delta }_{{\mathbf{O}}^{\mathbf{FP}\mathbf{16}}}{\mathbf{B}}^{\mathbf{FP}\mathbf{16}}({\mathbf{C}}^{\mathbf{FP}\mathbf{16}}+{\mathbf{A}}^{\mathbf{FP}\mathbf{16}})\end{array}$$

(14)

5. Experiment

5.1. Implementation Details

We conducted experiments to evaluate the effectiveness of our proposed method and compared it with representative methods.

Datasets and Models: We initially conducted experiments on small models including OPT-1.3B and OPT-2.7B [34], fine-tuning them on complete conversations. We used two distinct datasets: Self-Instruct [35] and Alpaca [36]. Following the benchmark in QLoRA [23], we also used the five-shot performance of the fine-tuned models on the Massively Multitask Language Understanding (MMLU) task [37]. To evaluate the efficiency and effectiveness of QLoGE, we scaled up our experiments to include larger models such as LLaMA2-7B and LLaMA2-13B [5]. We used two distinct datasets: OASST1 [38] and Alpaca [36]. The Self-Instruct dataset [35], generated by GPT-3, contains approximately 52K instructions. The OASST1 dataset [38] comprises human-created and annotated dialogues from ChatGPT-like assistants, featuring 161,443 messages in 35 languages and 461,292 quality assessments. The Alpaca dataset [36] includes 52,000 instructions and demonstrations generated using OpenAI’s text-davinci-003 engine.

To evaluate LoGE with more complex tasks, we fine-tuned LLaMA-2-7B [5] on the MetaMathQA dataset [39] to assess their mathematical problem-solving capabilities on the GSM8K [40] validation sets. Additionally, the models were fine-tuned on the CodeFeedback dataset [41] and evaluated for coding proficiency using the Human-eval datasets [42]. Furthermore, training was extended to the WizardLM-Evol-Instruct dataset [43] with subsequent testing for conversational abilities on the MT-Bench dataset [44]. The GSM8K dataset [40] consists of 8.5K high-quality elementary school math problems, all authored by human writers. The MT-bench dataset [44] contains 3.3K expert-level pairwise human preferences for model responses generated by six models in response to 80 MT-bench questions. The Human-eval dataset [42] comprises 164 original programming problems, evaluating language comprehension, algorithms, and basic mathematics, some of which resemble simple software interview questions. These experiments were conducted using subsets containing 100K data points and were limited to one epoch of training to minimize overhead.

To evaluate LoGE alongside other structured transformer-based large language models, we fine-tuned RoBERTa-large [2] on a subset of the GLUE benchmark [45], including MNLI, QNLI, SST-2, and MRPC.

Training Details: For the OPT models trained on the Alpaca and Self-instruct datasets, we adhered mostly to the hyperparameters used in QLoRA. We employed the Adam optimizer with a learning rate of $8\times {10}^{-5}$. The moving average betas were set at [0.9, 0.999], and the weight decay was maintained at 0. And we used a constant learning rate scheduler. For LLaMA models trained on the Alpaca and Self-instruct datasets, we also adhered mostly to the hyperparameters used in QLoRA. We employed the Adam optimizer with a learning rate of $2\times {10}^{-4}$. The moving average betas were set at [0.9, 0.999], and the weight decay was maintained at 0. Additionally, we implemented a constant learning rate scheduler, incorporating a warmup ratio of 0.3. For fine-tuning LLMs with complex tasks, we used the Adam optimizer with a learning rate of $2\times {10}^{-5}$. The moving average betas were set at [0.9, 0.999], with no weight decay. Additionally, we utilized a cosine learning rate scheduler with a warmup ratio of 0.3. The hyperparameters for RoBerta-large were almost the same as the hyperparameters from LoRA [13], and is detailed in

Table 2. Hyperparameters of our methods on GLUE. BS means batch size and GA means gradient accumulation.

Dataset	Epoch	BS	GA	LR
MNLI	10	8	4	3 × 10−4
SST-2	10	8	8	4 × 10−4
QNLI	10	8	4	2 × 10 −4
MRPC	20	8	4	3 × 10 −4

.

For smaller models like RoBERTa-large [2] and OPT [34], we applied LoGE only to the q-proj and v-proj layers to preserve model performance. For larger models like LLaMA-7B, we extended the use of LoGE to all linear layers within the attention module, as larger models exhibit greater robustness to model compression.

Implement Details: All experiments were conducted on RTX 3090 GPUs (NVIDIA, Santa Clara, CA, USA) using PyTorch (v2.2.0). To ensure better reproducibility, we have provided the prototype code in the Supplementary Materials.

5.2. Main Results

5.2.1. Results on OPT Models

We first applied LoGE to fine-tune the OPT models. Given that the MMLU results for OPT-1.3B and OPT-2.7B were about 25%, we used the perplexity of the trainable datasets for comparison.

Table 3. PPL on evaluation datasets across OPT models: comparing sizes and datasets. Experiments were run with one RTX3090. The symbol “↓” indicates that lower values are better.

Model	Dataset	Method	Rank	Memory	Metric (PPL ↓)	Runtime
OPT1.3B	Self-Instruct	LoRA	128	3697 MiB	2.73	1.98 h
OPT1.3B	Self-Instruct	GaLore	128	3893 MiB	2.99	1.59 h
OPT1.3B	Self-Instruct	LoGE	256	3769 MiB	2.71	1.53 h
OPT1.3B	Alpaca	LoRA	128	3715 MiB	4.82	2.05 h
OPT1.3B	Alpaca	GaLore	128	3976 MiB	5.00	1.71 h
OPT1.3B	Alpaca	LoGE	256	3713 MiB	4.81	1.63 h
OPT2.7B	Self-Instruct	LoRA	128	6834 MiB	2.44	2.71 h
OPT2.7B	Self-Instruct	GaLore	128	7100 MiB	2.67	2.32 h
OPT2.7B	Self-Instruct	LoGE	256	6629 MiB	2.43	2.09 h
OPT2.7B	Alpaca	LoRA	128	6808 MiB	4.36	2.90 h
OPT2.7B	Alpaca	GaLore	128	7209 MiB	4.55	2.54 h
OPT2.7B	Alpaca	LoGE	256	6969 MiB	4.37	2.19 h

summarizes the results with respect to different model sizes and fine-tuning datasets. For models fine-tuned with Alpaca, we observe that the performance was almost the same for LoRA and LoGE, with OPT-1.3B performing better with LoGE. Consequently, we included Self-instruct datasets for further evaluation. The memory usage of LoGE was slightly higher than that of LoRA, as LoGE needed to store three low-rank matrices $\mathbf{A}$, $\mathbf{B}$, and $\mathbf{C}$. Since only $\mathbf{A}$ was trainable, the optimizer state for LoGE was not significantly larger than that for LoRA. Additionally, we applied LoGE to only a few layers to preserve model performance. Therefore, LoGE did not reach the ideal speedup ratios.

5.2.2. Results on Complex Tasks

We then fine-tuned LLaMA2-7B on more complex tasks, with results detailed in

Table 4. Results of fine-tuning Llama 2-7B using LoRA, LoRArs, PiSSA, and LoGE, tested on MT-Bench, GSM8K, and Human-eval. Experiments were run with two RTX3090.

Dataset	Method	Rank	Memory	Metric	Runtime
MT-Bench	LoRA	128	26,894 MiB	5.61	5.70 h
MT-Bench	rsLoRA	128	26,928 MiB	5.25	5.58 h
MT-Bench	PiSSA	128	26,986 MiB	5.30	5.95 h
MT-Bench	LoGE	256	26,680 MiB	5.47	4.18 h
GSM8K	LoRA	128	21,648 MiB	42.08	7.18 h
GSM8K	rsLoRA	128	22,744 MiB	51.55	7.96 h
GSM8K	PiSSA	128	21,904 MiB	52.31	7.73 h
GSM8K	LoGE	256	22,454 MiB	48.74	5.20 h
Human-Eval	LoRA	128	26,634 MiB	14.76	10.95 h
Human-Eval	rsLoRA	128	26,618 MiB	18.29	11.72 h
Human-Eval	PiSSA	128	25,786 MiB	18.76	11.61 h
Human-Eval	LoGE	256	26,476 MiB	16.46	9.17 h

. According to these results, our LoGE method achieved superior performance with reduced training times compared to LoRA. The performance improvements were consistent across various models, demonstrating the versatility of our approach. The training loss plot in Figure 2 shows that LoGE converged faster than LoRA in the initial steps and maintained a smaller training loss throughout. This accelerated convergence can be attributed to the fact that while the parameters of LoRA were randomly initialized and required several steps to establish convergence directions, LoGE leveraged the gradients of the low-rank approximation of weights, which had already been aligned with the correct convergence direction. We also compared our method with two LoRA variants [46,47] that focus on enhancing model performance. While our approach lagged slightly in performance, it demonstrated a significant reduction in training time. This is primarily because our method prioritizes training efficiency by compressing pre-trained weights, which inevitably introduces some gradient errors.

5.2.3. Results on GLU Tasks

We then fine-tuned RoBerta-large on a subset of GLUE, with results detailed in

Table 5. Results of fine-tuning RoBerta-large using Full-FT, LoRA, and LoGE on a subset of GLUE. Experiments were run with one RTX3090. The symbol “↑” indicates that upper values are better.

Dataset	Method	Rank	Memory	Acc. ↑	Runtime
MNLI	FT	-	10,710 MiB	90.2%	12.25 h
	LoRA	8	4636 MiB	90.6%	8.54 h
	LoGE	256	5052 MiB	89.9%	7.96 h
SST2	FT	-	9444 MiB	96.4%	1.25 h
	LoRA	8	4118 MiB	96.2%	1.22 h
	LoGE	256	4486 MiB	95.7%	1.37 h
QNLI	FT	-	16,488 MiB	94.7%	3.25 h
	LoRA	8	7564 MiB	94.9%	2.48 h
	LoGE	256	8566 MiB	94.3%	2.21 h
MRPC	FT	-	10,194 MiB	90.9%	0.22 h
	LoRA	8	5686 MiB	90.1%	0.16 h
	LoGE	256	6032 MiB	91.1%	0.17 h

. Because smaller models are harder to compress [48], and RoBerta-large has about 350 M parameters, which is much smaller than prior models, we kept the rank at 256 to ensure the compression ratio. Compared to LoRA, LoGE achieved a slight speedup with a marginal increase in memory usage. The reason is that a rank of 256 is relatively large for RoBERTa-large, given its hidden state size of 1024. The slight increase in training time for SST-2 was due to gradient accumulation, as SST-2 used eight steps compared to four in other tasks. Additionally, longer inputs in tasks like SST-2 and MRPC reduced the speedup ratio, as attention calculations became more time-intensive. Additionally, LoGE achieved better performance on smaller subsets like MRPC. This is because LoGE was initialized with SVD, providing it with an advantageous starting point and the correct convergence direction.

5.2.4. QLoGE Performance Comparison

We applied our QLoGE method to fine-tune LLaMA models and evaluated their performance using the MMLU benchmark. As summarized in

Table 6. Five-shot MMLU benchmark accuracy across LLaMA models: comparing sizes and datasets. The symbol “↑” indicates that upper values are better.

Model	Dataset	Method	Rank	Memory	Metric (Acc. ↑)	Runtime
LLaMA-7B	OASST1	QLoRA	64	7836 MiB	33.9%	4.33 h
LLaMA-7B	OASST1	QLoGE	256	9390 MiB	34.0%	3.41 h
LLaMA-7B	Alpaca	QLoRA	64	7836 MiB	38.4%	22.72 h
LLaMA-7B	Alpaca	QLoGE	256	9462 MiB	38.6%	17.67 h
LLaMA-13B	OASST1	QLoRA	64	13,258 MiB	45.5%	7.06 h
LLaMA-13B	OASST1	QLoGE	320	14,542 MiB	45.1%	5.87 h
LLaMA-13B	Alpaca	QLoRA	64	13,986 MiB	47.8%	37.65 h
LLaMA-13B	Alpaca	QLoGE	320	14,604 MiB	47.6%	31.34 h

, QLoGE achieves comparable accuracy to QLoRA, with up to a 1.3× speedup. Additionally, there is an increase in memory usage—about 20% for LLaMA-7B and 10% for LLaMA-13B. This increase is due to the higher rank used in QLoGE compared to QLoRA. A higher rank is necessary in LoGE to avoid performance degradation, as a smaller rank can affect gradient accuracy. Combining quantization with low-rank updates further amplifies the gradient gap, leaving this challenge as an open question for future research. Furthermore, our method reduces the number of trainable parameters and eliminates the significant overhead associated with dequantizing weights for computation, offering another considerable advantage.

5.3. Ablation Studies

The performance reported in this subsection is assessed using perplexity (PPL), where lower values indicate better model performance. For efficient evaluation, we fine-tuned the models with one epoch for comparison in this section.

5.3.1. The Impact of Different $\alpha$ and Different Datasets on Model Performance

We employed a temperature parameter, $\alpha$, to combine two normalized importance scores: the magnitude of the singular values and the contribution to the gradient. To explore how different values of $\alpha$ (ranging from 0.1 to 10) affect model performance, we conducted tests across various datasets using the OPT-1.3B model in

Table 7. Perplexity of different $\alpha$ with different datasets for OPT-1.3B. A value of 0 means using only the magnitude of the singular value, and ∞ means using only the contribution to the gradient. The symbol “↓” indicates that lower values are better.

PPL ↓	$\mathit{\alpha }=0$	$\mathit{\alpha }=0.1$	$\mathit{\alpha }=0.3$	$\mathit{\alpha }=1$	$\mathit{\alpha }=3$	$\mathit{\alpha }=10$	$\mathit{\alpha }=\mathit{\infty }$
Self-instruct	4.18	4.15	4.15	4.18	4.16	4.19	4.69
Alpaca	7.53	7.55	7.52	7.56	7.56	7.54	7.63

. Our findings reveal that models perform optimally when initialized with the top-K largest singular values. This supports findings from Hu et al. [13], which noted a significant overlap with the top singular vectors in LoRA. Interestingly, relying solely on gradient contribution for subspace selection resulted in suboptimal performance, diverging from the results observed in Child Tuning [18] and SPT [19]. This could be due to our method’s reliance on low-rank matrices to estimate input gradients, potentially omitting subspaces with the largest singular values and thus introducing significant estimation errors. Remarkably, we found $\alpha =0.3$ to be the most effective setting for both datasets tested. This suggests that the fusion method and this particular hyperparameter setting are generally applicable, not just dataset-specific.

5.3.2. The Time Cost of SVD

The time cost of SVD is detailed in

Table 8. Time cost of SVD and the total fine-tuning process.

Model	SVD (s)	Fine-Tuning (s)	Total (s)	Ratios
LLaMA-7B	612	39,420	40,032	1.53%
LLaMA-13B	1420	44,100	45,520	3.11%

. Relative to the total training duration, the time added by SVD is minimal, constituting approximately 1.53% and 3.11% of the overall fine-tuning times. Including the time for SVD, LoGE still achieves approximately a 20% speedup compared to LoRA.

5.3.3. Different Selection for LoGE

To further evaluate the optimal rank r for LoGE and determine the best insertion points, we conducted two different sets of experiments. The first set involved varying the rank r to observe its influence on model performance and acceleration ratio. The second set applied LoGE at different points within the model to assess the limitations of the low-rank estimator in maintaining performance.

Rank r choice for LoGE. We adjusted the rank r from 16 to 512 for the OPT-1.3B models using the Self-Instruct datasets in

Table 9. PPL on Self-Instruct for OPT-1.3B models after applying different r values of LoGE. The symbol “↓” indicates that lower values are better.

r	16	32	64	128	256	512
PPL ↓	6.09	5.07	4.57	4.15	3.82	3.71
Runtime (h)	1.59	1.57	1.60	1.60	1.62	1.72

. Unlike typical Parameter-Efficient Fine-Tuning (PEFT) methods, in our approach, the rank r directly influences the accuracy of gradient estimation. Our results show that even using the same low rank as LoRA, such as $r=16$, can lead to significant gradient estimation errors, although the training runtime remains comparable. We applied the Low-Rank Gradient Estimator (LoGE) specifically to ${\mathbf{W}}_{q_i}$ and ${\mathbf{W}}_{v_i}$. However, as r increased, the computational demands of LoGE surpassed those of LoRA, especially when r reached 512, at which point the computational complexity approached almost $\mathbf{O}\left(d^2\right)$. To strike an optimal balance between trainable parameters, training time, and performance, we selected $r=d/16$ for our studies.

Linear layer choice for LoGE. We varied the weights within the model to which LoGE was applied. Specifically, we avoided applying LoGE to weights in the Feed-Forward Network (FFN) due to dimensional scaling. However, the number of trainable parameters changed depending on which weights LoGE was applied to. Applying LoGE to more layers resulted in more compressed gradients, but we did not adjust the rank r with the increase in layers. We list the results in

Table 10. PPL on Self-Instruct for OPT-1.3B models after applying LoGE to different linear layers. The symbol “↓” indicates that lower values are better.

Weight Type	${\mathit{W}}_{{\mathit{q}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{k}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{v}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{o}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{q}}_{\mathit{i}}},{\mathit{W}}_{{\mathit{k}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{q}}_{\mathit{i}}},{\mathit{W}}_{{\mathit{v}}_{\mathit{i}}}$	${\mathit{W}}_{{\mathit{q}}_{\mathit{i}}},{\mathit{W}}_{{\mathit{k}}_{\mathit{i}}},{\mathit{W}}_{{\mathit{v}}_{\mathit{i}}},{\mathit{W}}_{{\mathit{o}}_{\mathit{i}}}$
PPL ↓	5.44	5.54	4.36	4.44	4.83	4.15	4.21
Runtime (h)	1.42	1.42	1.52	1.44	1.52	1.60	1.95

. Our findings align closely with existing PEFT methods [10,13]. Applying LoGE only to the layers ${\mathbf{W}}_{q_i}$ and ${\mathbf{W}}_{k_i}$ significantly reduced performance. While the performance of ${\mathbf{W}}_{v_i}$ and ${\mathbf{W}}_{o_i}$ was closer to that of $\{{\mathbf{W}}_{q_i},{\mathbf{W}}_{v_i}\}$, the increase in trainable parameters from using LoGE still resulted in a larger performance gap than reported in [13]. This may be due to the difference in the number of trainable parameters between our methods. Additionally, making ${\mathbf{W}}_{v_i}$ trainable increased the training time by approximately 1.1 times, likely due to the involvement of gradient calculations in the attention mechanism. Our results suggest that although increasing the number of trainable parameters can enhance model performance, excessive compression of gradients can be detrimental. Fine-tuning with just two trainable parameters, ${\mathbf{W}}_{q_i}$ and ${\mathbf{W}}_{v_i}$, appears sufficient, indicating that a balance is crucial for optimal fine-tuning efficiency.

6. Limitations and Future Work

One major limitation of our method is its reduced effectiveness on small models (with fewer than 1 billion parameters), as these models are more significantly affected by low-rank gradient compression. Future work could explore several directions:

• Application to training from scratch: Investigating how this method can be integrated into training from scratch, including addressing changes in the subspace during training and fusing the low-rank update branch into the model weights.
• Improving performance on small models: Developing techniques to enhance the accuracy and effectiveness of the method when applied to smaller models.
• Adapting to edge devices: Exploring how to implement this approach on edge devices to enable lifelong learning, including designing systems for efficient fine-tuning tailored to resource-constrained environments.

7. Conclusions

In this paper, we propose the Low-Rank Gradient Estimator (LoGE), an efficient method for fine-tuning LLMs. LoGE incorporates model compression into low-rank weight update approximations, performing low-rank approximations of pre-trained weights during backward computation while preserving the full-rank pre-trained weights during forward computation. This approach reduces the computational cost of training a linear layer by approximately 60%. Experiments across various models and tasks demonstrated that LoGE not only maintained performance but also significantly accelerated the fine-tuning process by about 20%. Additionally, LoGE can be combined with quantization methods, reducing the training time for quantization-based efficient fine-tuning by approximately 30% without any performance degradation. LoGE advances the field of fine-tuning on edge devices and holds significant practical implications for reducing the carbon footprint of deep learning.