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

Abstract: We present a novel parameter-efficient finetuning (PEFT) method, dubbed as adaptive freezing of low rank adaptation (AFLoRA). Specifically, for each pre-trained frozen weight tensor, we add a parallel path of trainable low-rank matrices, namely a down-projection and an up-projection matrix, each of which is followed by a feature transformation vector. Based on a novel freezing score, we then incrementally freeze these projection matrices during fine-tuning to reduce the computation and alleviate over-fitting. Our experimental results demonstrate that we can achieve state-of-theart performance with an average improvement of up to 1.09% as evaluated on the GLUE and GSM8k benchmark while yielding up to 9.5× fewer average trainable parameters. While compared in terms of runtime, AFLoRA can yield up to 1.86× improvement as opposed to similar PEFT alternatives. Besides the practical utility of our approach, we provide insights on the trainability requirements of LoRA paths at different modules and the freezing schedule for the different projection matrices. Code is released at: https://github.com/ zeyuliu1037/AFLoRA/tree/main.

heading: 1 Introduction
text: 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 parameterefficient 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.
161
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× and 2.96× improvement in runtime and FLOPs, respectively, while remaining comparable as LoRA on these two metrics. Compared to LoRA we require 9.5× fewer average trainable parameters to yield similar or improved performance.

heading: 2 Related Works
text: 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.

heading: 3 Motivational Case Study
text: To understand the high-rank requirement for the frozen projection matrices in ELoRA, we conduct two sets of fine-tuning on SST-2 and MRPC, with ELoRA having rank (r) of 1024 and 4, respectively. As we can see in Fig. 2, the model with r = 4, yields poorer performance, highlighting the need for high rank for the frozen tensors. This high rank causes ELoRA to potentially be FLOPs inefficient.

heading: 4 AFLoRA: Methodology
text: 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 (loraA), the up-projection linear layer (loraB), and two feature transform vectors (sd, and sb) placed before and after loraB . However, unlike (Kopiczko et al., 2024), we keep both the projection matrices (loraA and loraB) 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 l0X + Λ l bB lΛldA lX (1)
Here, Al and Bl are the trainable LoRA tensors of loralA and lora l B , respectively. Λd and Λb are the vectors of sd, and sb, respectively. W l0 represents the frozen pre-trained weights. We use Kaiming Uniform initialization for Al and Bl, 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 Al.
IAl = |∇L(θ)| , I (t) Al = β1I (t−1) Al + (1− β1)I (t) Al (2)
U (t)
Al = ∣∣∣I(t) Al − I(t)Al ∣∣∣ , U (t)Al = β2U (t−1) Al +(1−β2)U (t) Al (3)
s (t)
Al = mean(I
(t) Al ◦ U (t) Al ) (4)
Here, for each projection tensor at iteration t, we compute a smoothed gradient (I(t)Al ) and uncertainly tensor (U (t)Al ), as shown in Eq. 2 and 3, respectively. We then evaluate the freezing score s(t)
Al , as the
mean of the tensor generated via Hadamard product (◦) between I(t)Al and U (t)
Al . 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 ti training steps, and then progressively freeze them by calculating the freezing fraction r(t) as shown in Eq. 5. Finally, all the projection matrices freeze beyond T − tf steps. Note, at step t, for a computed freezing fraction k, we freeze the lowest k% projection matrices.
where t refers to current #step, T is the total number of fine-tuning steps. We set ti to the steps corresponding to one epoch and set tf to 70% of the total training steps.

heading: 5 Experiments
text: 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 a 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× 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, that 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× and 2.96× runtime and FLOPs improvement while remaining comparable with LoRA in these two metrics. Compared to LoRA we yield 9.5× parameter reduction while remaining comparable with ELoRA. These results clearly demonstrate AFLoRA as a 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 BARTLarge (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× fewer trainable parameters as compared to that with LoRA. On the CNN/DailyMail Summarizing task (Table 3), AFLoRA requires 1.69× fewer trainable parameters to reach similar or improved rouge score values.

heading: 6 Ablations and Discussions
text: 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,
Table 4: Ablation study on the trainability impact of the projection matrices (PM) of the AFLoRA module. We keep the vectors trainable throughout for all.
PM #Params. CoLA SST-2 MRPC QNLI STS-B RTE Avg. Trainable 0.45M 70.15 95.99 92.4 94.16 89.90 88.45 88.51 Frozen 0.08M 70.36 94.95 89.22 93.61 91.17 85.92 87.54 AFLoRA (Ours) 0.14M 72.01 96.22 91.91 94.42 91.84 88.09 89.23
Figure 4: A comparison of performance outcomes utilizing three distinct freezing score methodologies.
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.

heading: 7 Conclusions
text: In this paper, we presented AFLoRA, 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× while yielding 0.85% average improved performance as evaluated on the GLUE benchmark.

heading: 8 Limitation
text: 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 for PEFT that we consider as future research.

heading: A Appendix
text: A.1 Dataset The details of train/test/dev splits and the evaluation metric of the GLUE (Wang et al., 2018) dataset are reported in Table 6. We use the Huggingface Transformers library (Wolf et al., 2020) to source all the datasets.
A.2 Hyperparameter configuration Table 7 shows the main hyper-parameter set up in this paper. Besides them, we use the same optimizer, warmup Ratio, and LR schedule as Hu et al. (2021). We use NVIDIA RTX A6000 (maximum GPU memory=49140MB) to measure the training runtime. For all experiments, we run 5 times using different random seeds and report the average results.
A.3 Ablation study on if freezing the two projection matrices in the same layer simultaneously
We study the value of freezing both projection matrices in the same layer simultaneously. The results, depicted in Table 8, demonstrate that freezing the projection matrices separately yields consistently superior performance compared to freezing them simultaneously.

heading: Future Work
text: Additionally, the integration of AFLoRA in the adaptive rank evaluation framework can potentially open a new direction for PEFT that we consider as future research