AFLoRA: Adaptive Freezing of Low Rank Adaptation in Parameter Efficient Fine-Tuning of Large Models

Pre-trained language models such as BERT (Devlin et al., 2018), GPT-3 (Brown et al., 2020), and LLaMA2 (Touvron et al., 2023) have demonstrated commendable performance on various natural language processing (NLP) tasks Kang et al. (2024). However, their zero-shot performance on many downstream tasks often falls short of expectations. One possible solution is full fine-tuning (FFT) of the model on the downstream dataset. However, the large model parameter size makes this process prohibitively costly.

To address this challenge, various parameter-efficient fine-tuning (PEFT) methods including low rank adaptation (LoRA) (Hu et al., 2021), adapter tuning (He et al., 2021), and prompt tuning (Lester et al., 2021) are proposed. These methods add parameters to the trained model for fine-tuning, bypassing the need to adjust the weights of the pre-trained model. In particular, LoRA (Hu et al., 2021) and its variants (Zhang et al., 2023) add a trainable low-rank path consisting of down-projection and up-projection matrices to the model, inspired by Aghajanyan et al. (2020) which showed that such low-rank paths can effectively approximate the trained weight tensors. ELoRA Kopiczko et al. (2024) extends LoRA by adding trainable feature transformation vectors to the output of each project matrix. They showed that SoTA accuracy can be achieved with the projection matrices frozen after random initialization while keeping the two feature transformation vectors trainable. This approach significantly reduces the number of trainable parameters. However, compared to LoRA, ELoRA incurs higher computation costs due to the higher rank needed for the frozen projection matrices. Fig. 1 illustrates LoRA and ELoRA, contrasting them to our proposed method AFLoRA.

Our contributions. To reduce the trainable parameter count and computation costs of fine-tuning, we present Adaptive Freezing of Low Rank Adaptation (AFLoRA). More specifically, we first investigate the rank needed for the frozen LoRA path in ELoRA and observe that reducing the rank of the frozen projection matrices (PM) causes a drop in fine-tuning performance.

Based on this insight, we present AFLoRA, which starts with a low-rank trainable path that includes projection matrices and feature transformation vectors and trains the path for some epochs. We then gradually freeze the projection matrices based on a novel freezing score that acts as a proxy for the trainability requirement of a LoRA tensor. In this way, we not only help alleviate the over-fitting issue but also, improve the computation efficiency. To evaluate the benefit of AFLoRA, we perform extensive evaluations on multiple NLP benchmark datasets and compare accuracy, FLOPs, and training time with several existing alternatives. Specifically, compared to ELoRA we yield $1.86\times$ and $2.96\times$ improvement in runtime and FLOPs, respectively, while remaining comparable as LoRA on these two metrics. Compared to LoRA we require $9.5\times$ fewer average trainable parameters to yield similar or improved performance.

PEFT (Hu et al., 2021; Kundu et al., 2024; Sridhar et al., 2023; Yin et al., 2024) refers to a collection of methodologies that focus on allowing a small number of parameters to fine-tune to yield good performance on a downstream task. For example, prefix-tuning (Li and Liang, 2021) adds trainable prefix tokens to a model’s input or hidden layers while adapter-tuning (Houlsby et al., 2019) inserts small neural network layers, known as adapters, within each layer of a pre-trained model. LoRA (Hu et al., 2021), on the other hand, adds low-rank tensors in parallel to the frozen pre-trained weights. AdaLoRA (Zhang et al., 2023) allows the rank of the LoRA path to be chosen in an adaptive way. Other variants like SoRA Ding et al. (2023) and LoSparse Li et al. (2023) have investigated the impact of sparsity in and alongside the low-rank path, respectively. Recently, efficient low-rank adaptation (ELoRA) (Kopiczko et al., 2024) has proposed to keep the LoRA path frozen, while introducing two trainable feature transformation vectors. Thus, this work only studies an extreme scenario of keeping the LoRA path frozen, and, to the best of our knowledge, no work has investigated the trainability requirement of the projection matrices.

Module Structure. Inspired by the framework proposed by Kopiczko et al. (2024), we design the LoRA module to encompass four components, namely, the down-projection linear layer ($lora_{A}$), the up-projection linear layer ($lora_{B}$), and two feature transform vectors ($s_{d}$, and $s_{b}$) placed before and after $lora_{B}$. However, unlike Kopiczko et al. (2024), we keep both the projection matrices ($lora_{A}$ and $lora_{B}$) and vectors trainable at the beginning and keep the rank very low. The module processes a given input $X$ through these components to produce an output $Y$. The complete operation for a layer $l$ can be described as follows:

$${Y}={W^{l}_{0}}{X}+\Lambda^{l}_{b}B^{l}\Lambda^{l}_{d}A^{l}{X}$$

(1)

Here, $A^{l}$ and $B^{l}$ are the trainable LoRA tensors of $lora^{l}_{A}$ and $lora^{l}_{B}$, respectively. $\Lambda_{d}$ and $\Lambda_{b}$ are the vectors of $s_{d}$, and $s_{b}$, respectively. $W^{l}_{0}$ represents the frozen pre-trained weights. We use Kaiming Uniform initialization for $A^{l}$ and $B^{l}$, and follow Kopiczko et al. (2024) to initialize the vectors.

Adaptive Freezing. In pruning literature (Han et al., 2015; Molchanov et al., 2019; Zhang et al., 2022; Yin et al., 2024; Kundu et al., 2021, 2022), sensitivity is gauged to reflect weight variability, necessitating consideration of both the weights’ magnitudes and their gradients. Small weight values suggest minimal impact, while minor gradient values indicate stability. Taking inspiration from this idea, here we introduce the concept of a "freezing score". However, unlike pruning where both magnitude and gradient play a critical role in identifying insignificant weight, we leverage only gradient as a proxy to compute the freezing score. This is because, we assume large magnitude weights with negligible change has the same priority to be frozen as that for small magnitude weights. This score quantifies the degree to which weights vary throughout the training process. Consequently, when the expected changes to the weights become negligible, we may consider them to be frozen, thereby saving computational resources and energy.

The following equation describes the freezing score evaluation steps for a low-rank tensor $A^{l}$.

$$I_{A^{l}}=\left|\nabla\mathcal{L}(\boldsymbol{\theta)}\right|,\overline{I}_{A^{l}}^{(t)}=\beta_{1}\overline{I}_{A^{l}}^{(t-1)}+(1-\beta_{1})I_{A^{l}}^{(t)}$$

(2)

$$U_{A^{l}}^{(t)}=\left|I_{A^{l}}^{(t)}-\overline{I}_{A^{l}}^{(t)}\right|,\overline{U}_{A^{l}}^{(t)}=\beta_{2}\overline{U}_{A^{l}}^{(t-1)}+(1-\beta_{2})U_{A^{l}}^{(t)}$$

(3)

$$s_{A^{l}}^{(t)}=mean(\overline{I}_{A^{l}}^{(t)}\circ\overline{U}_{A^{l}}^{(t)})$$

(4)

Here, for each projection tensor at iteration $t$, we compute a smoothed gradient ($\overline{I}_{A^{l}}^{(t)}$) and uncertainly tensor ($\overline{U}_{A^{l}}^{(t)}$), as shown in Eq. 2 and 3, respectively. We then evaluate the freezing score $s_{A^{l}}^{(t)}$, as the mean of the tensor generated via Hadamard product ($\circ$) between $\overline{I}_{A^{l}}^{(t)}$ and $\overline{U}_{A^{l}}^{(t)}$.

To apply thresholding on the LoRA freezing scores, we use the cubic schedule as (Zhang et al., 2022). In specific, we keep the projection matrices trainable for the initial $t_{i}$ training steps, and then progressively freeze them by calculating the freezing fraction $r(t)$ as shown in the Eq. 5. Finally, all the projection matrices freeze beyond $T-t_{f}$ steps. Note, at step $t$, for a computed freezing fraction $k$, we freeze the lowest $k\%$ projection matrices.

$$r(t)=\left\{\begin{array}[]{ll}0&\quad 0\leq t<t_{i}\\ 1-\left(1-\frac{t-t_{i}}{T-t_{i}-t_{f}}\right)^{3}&\quad t_{i}\leq t<T-t_{f}\\ 1&\quad\text{otherwise}\end{array}\right.$$

(5)

where $t$ refers to current #step, $T$ is the total number of fine-tuning steps. We set $t_{i}$ to the steps corresponding to one epoch and set $t_{f}$ to 70% of the total training steps.

Models & Datasets. We use the PEFT framework of Mangrulkar et al. (2022) and evaluate the fine-tuning performance of DeBERTaV3-base (He et al., 2020) to fine-tune on our framework on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018). The details of the hyperparameter settings for each dataset are listed in Appendix A.2.

Performance Comparison. We benchmark the performance with AFLoRA and present comparison with LoRA and its variants. For ELoRA, we reproduce the results at our end while the results for other methods are sourced from Ding et al. (2023). As shown in Table 1, AFLoRA can achieve SoTA performance on the majority of datasets and on average while requiring similar and $9.5\times$ fewer average trainable parameters as compared to ELoRA and LoRA, respectively.

Runtime & FLOPs Comparison. Fig. 3 shows the comparison of the normalized average training runtime, normalized FLOPs, and normalized trainable parameters. For AFLoRA, we average the training time, FLOPs, and trainable parameters over six GLUE datasets (except the MNLI and QQP datasets). Note, for LoRA and ELoRA, the trainable parameters and FLOPs remain fixed for all the datasets. We compute their average runtime the same way as ours. Compared to ELoRA we can yield up to $1.86\times$ and $2.96\times$ runtime and FLOPs improvement while remaining comparable with LoRA in these two metrics. Compared to LoRA we yield $9.5\times$ parameter reduction, while remaining comparable with ELoRA. These results clearly demonstrate AFLoRA as PEFT method that can yield similar parameter efficiency as ELoRA while costing no training overhead in FLOPs or time.

Results with Large Language Models (LLMs). We now demonstrate the AFLoRA fine-tuning performance with two popular LLM variants, namely, LLaMA-7B Touvron et al. (2023) and BART-Large Lewis et al. (2019) on GSM8k complex reasoning and CNN/Daily mail summarizing task, respectively. As demonstrated in Table 2, on GSM8k, AFLoRA yields improved accuracy of $1.09\%$ while requiring $3.15\times$ fewer trainable parameters as compared to that with LoRA. On CNN/DailyMail Summarizing task (Table 3), AFLoRA requires $1.69\times$ fewer trainable parameters to reach similar or improved rouge score values.

We conducted our ablation studies on six GLUE benchmark datasets, omitting QQP and MNLI, the two most computationally demanding datasets.

Do we really need adaptive freezing? We conducted experiments with all the LoRA PMs frozen (same as ELoRA), all the LoRA PMs trainable, and with our adaptive training of LoRA PMs. We use, $r=4$ for the LoRA path, for all. As we can see in Table 4, keeping the projection matrices trainable yields better average performance compared to keeping them frozen throughout. However, AFLoRA with adaptive freezing yields even better performance than keeping them trainable throughout, potentially highlighting its ability to regularize the fine-tuning against overfitting.

Do we need to keep the PMs trainable for all layer types? There are two major layer types, FFN and the attention layers. We place the PMs in both along with the feature transformation vectors. We then study the necessity of keeping the PMs trainable in these two layer types. Note, here, we keep the vectors trainable for all throughout. As shown in Table 5, keeping the PMs trainable (and then adaptive freezing) in the FFN yields better performance compared to the alternatives. Note we keep the PMs in the attention layers frozen to random values. Interestingly, allowing all PMs to initially train and then adaptively freeze yields poorer performance than allowing them only in MLP. This may hint at the FFN weights to play a more important role in fine-tuning performance.

Ablation with sensitivity choices. Fig. 4 presents ablation with three sensitivity scores based on three different sensitivity choices, namely, $|grad(p)|$ (adopted in AFLoRA), $|p*grad(p)|$, and $|grad(p)/p|$. On average, the freezing score adopted in AFLoRA, consistently yields better accuracy over the other two.

Discussion on Freezing Trend. We use the RTE dataset as a case study, to understand the freezing trend of the PMs across different layers. Specifically, we illustrate the specific number of iterations required before freezing each component in Fig. 5. Interestingly, as can be seen from the figure, analysis reveals that the down-projection matrix parallel to the intermediate linear layer requires longer training duration prior to being frozen, as compared to the other PMs. This may potentially hint at the low approximation ability of the intermediate layer as compared to the second MLP in the FFN.

In this paper, we presented AFLoRA, an adaptive freezing of LoRA adapters that allow near-optimal trainability of the LoRA projection matrices and freezes them driven by a "freezing score" after certain fine-tuning steps. Compared to LoRA, AFLoRA can reduce the trainable parameters by up to $9.5\times$ while yielding $0.85\%$ average improved performance as evaluated on the GLUE benchmark.

In the ablation study with various freezing score metrics, we discovered that alternative scoring methods outperform ours on certain datasets, suggesting possible room for research in refining the freezing scores. This can further improve performance with AFLoRA. Additionally, the integration of AFLoRA in the adaptive rank evaluation framework can potentially open a new direction of PEFT that we consider as future research.