Chain of LoRA: Efficient Fine-tuning of Language Models via Residual Learning

Wenhan Xia, Chengwei Qin, Elad Hazan

Introduction

Pre-trained language models have become instrumental in natural language processing, demonstrating remarkable performance across various fields. Large language model fine-tuning is a process for adapting pre-trained models to specific tasks, allowing for improved performance on various real-world applications, such as machine translation and code analysis (Lewis et al., 2019; Wang et al., 2021; Qin et al., 2023). Despite the notable benefits of full parameter fine-tuning, the computational expenses and memory requirements it entails present significant challenges, particularly in light of the ever-growing size of large language models.

For this reason, parameter efficient finetuning (PEFT) methods have received significant attention (Pfeiffer et al., 2020; He et al., 2021). Instead of adjusting all the parameters of the model, PEFT involves fewer adjustments to the original model parameters to specialize its knowledge for a particular application (Houlsby et al., 2019; Lester et al., 2021). One of the most widely used paradigms in parameter efficient fine turning is Low-Rank Adaptation (LoRA) (Hu et al., 2021). LORA focuses on modifying only a small, low-rank portion of the model’s weights. This is achieved by adding low-rank matrices to the weights of the model during training. The advantage of LORA is that it significantly reduces the computational burden and time required for fine-tuning, making it more efficient and scalable, especially for very large models. Despite the significant computational advantage of LORA, it is inferior to full parameter fine-tuning in terms of generalization error.

In this paper we investigate whether the generalization error gap between LORA and full parameter fine-tuning can be reduced albeit preserving the computational efficiency. We do this by learning a higher rank augmentation of the LLM weights by method of residual learning. The high rank augmentation is composed of several low rank structures. Namely, we use an iterative procedure to learn a low-rank addition to the existing approximation, thereby increasing its rank. Hence, we call the procedure “chain of LORA”, or COLA for short.

This residual learning method is inspired by the Frank-Wolfe algorithm as applied to matrix completion, which augments an existing completion by a rank one addition. Over many iterations, this residual learning procedure can be shown to produce an accurate higher rank completion.

We present an iterative optimization framework, COLA, for parameter efficient fine tuning. COLA is based on the Frank Wolfe method from mathematical optimization, and we formalize this relationship.

We demonstrate the effectiveness of COLA via extensive experiments across datasets and models. COLA consistently outperforms LoRA in terms of generalization error with no additional cost of compute. For example, fine-tuning OPT-1.3B with COLA brings a relative $6.47\%$ test accuracy gain to LoRA on WSC. LLama2-7B experiments shows up to $4.4\%$ relative test score improvement.

We provide theoretical analysis of the iterative learning framework employed in our proposed method, demonstrating the convergence to stationary points in the setting of smooth nonconvex optimization.

Related Work

Conventional full parameter fine-tuning becomes computationally impractical as both model size and the number of downstream tasks increase. In response to this challenge, recent advancements in parameter-efficient finetuning methods suggest modifying only a small portion of parameters while maintaining the majority of pre-trained model parameters unchanged.

Within this domain, a line of research known as adapter based approach involves inserting compact adapter modules between transformer layers. Throughout the fine-tuning process, only the newly introduced lightweight adapters are trained, while the pre-trained model remains frozen and shared across tasks, thus significantly enhancing the practicality and efficiency of adapting large models to diverse tasks. Houlsby et al. (2019) propose a new bottleneck adapter module and position it twice within each transformer (Vaswani et al., 2017) layer. The adapter employs a bottleneck architecture, incorporating a skip connection to effectively constrain the number of parameters involved in the module design. Variant adapter architecture and placements are proposed in concurrent work (Bapna & Firat, 2019; Stickland & Murray, 2019). Building upon the success of adapter-based approaches for single-task adaptation, subsequent studies extend the adapter-based architecture to the realm of multi-task learning scenarios (Mahabadi et al., 2021). AdapterFusion proposes a two-stage learning framework where task-specific adapters are learned and then later combined in a separate knowledge composition step (Pfeiffer et al., 2020).

Prefix tuning methods.

Alternative research investigates the incorporation of tunable parameters into both the input and hidden layers, as explored by Li & Liang (2021). These lightweight task-specific vectors, commonly referred to as the prefix, offer a notable reduction in the memory load required for storing task-specific models. Additionally, they outperform full fine-tuning, particularly in scenarios with limited data availability. Efficient prompt tuning further simplifies prefix tuning by concatenating a trainable tensor (”soft prompt”) with the model’s input embeddings (Lester et al., 2021). These ”soft prompts” are learned through backpropagation to perform downstream tasks.

LoRA and its variants.

The most closely related work to ours is LoRA (Hu et al., 2021), which introduces trainable low-rank matrices to approximate weight update during fine-tuning. We elaborate on its technical details in the preliminaries section below. Building upon the foundation laid by LoRA, numerous recent studies have explored its variants from different perspectives. QLoRA (Dettmers et al., 2023) further leverages 4-bit quantization to effectively and efficiently fine-tune LLMs. To enhance parameter efficiency, Tied-LoRA, introduced by Renduchintala et al. (2023), incorporates weight tying and selective training. Chen et al. (2023) propose LongLoRA to extend the context sizes of LLMs with limited computation cost. MultiLoRA (Wang et al., 2023) is designed specifically for better multi-task adaptation. Concurrently, Sheng et al. (2023) introduce S-LoRA, offering a framework that enhances the scalable serving of multiple LoRA adapters.

Optimization for fine tuning of LLM has special challenges, notably memory constraints. For this reason, zero-order optimization methods were proposed (Malladi et al., 2023).

Our Method

In this section we describe our method for fine tuning. It is divided into two parts, in the first we present necessary background for our exposition, and the second gives details of COLA.

Low Rank Adaptation (LoRA) (Hu et al., 2021) aims to improve the efficiency of fine-tuning large language models by training much smaller low-rank decomposition matrices of certain weights. It hypothesizes a low ”intrinsic rank” of weight updates at task adaptation and injects trainable low-rank decomposition matrices into each layer of the Transformer architecture. Consider a weight matrix $W_{frozen}$ from the pre-trained model, the weight update $\Delta W$ for task adaptation is represented with a low-rank decomposition $BA$ . The forward pass with LoRA is as follows:

Frank-Wolfe The Frank-Wolfe method, also known as the conditional gradient method, is an optimization algorithm for solving constrained convex, and more recently nonconvex, optimization problems. The key feature of the Frank-Wolfe method is how it handles the constraints. Instead of projecting onto the constraint set via projections, it uses a linear optimization oracle. Iteratively, the method finds a linear approximation of the objective function within the feasible region and moves towards the minimizer of this approximation.

The Frank-Wolfe algorithm is particularly suited for problems in which linear optimization is easier than Euclidean projections. For this reason, “projection free” methods were considered in the machine learning community (Hazan, 2008; Jaggi, 2013; Hazan & Kale, 2012; Garber & Hazan, 2016). More recently nonconvex optimization was considered using the Frank Wolfe method in (Lacoste-Julien, 2016; Reddi et al., 2016).

2 Chain of LoRA

In this section we give the details of our simple yet effective optimization framework for efficient parameter finetuning of large language models. The key idea of our method is to form a chain of LoRAs and iteratively learn the low-rank adaptation LoRA modules. As illustrated in Figure 1, our method is comprised of three stages: Tune LoRA, Tie a knot, Extend the chain. We first introduce notations, followed by an explanation of the three stages in the workflow. We also provide the detailed step-by-step procedure in Algorithm 1.

where $\mathcal{L}$ is the task-specific objective function. COLA follows an iterative residual learning paradigm. Fine-tuning each $(A_{i},B_{i})$ can be viewed as learning the residual of $\Delta W^{\star}-\sum_{j=1}^{i-1}B_{j}A_{j}$ , which is an easier optimization problem compared to learning $\Delta W^{\star}$ from scratch. We hypothesize that $\sum_{i=1}^{M}B_{i}A_{i}$ approximates $\Delta W^{\star}$ better than a single LoRA update $BA$ , and we design a chaining framework to achieve this with less computation compared to the baseline LoRA.

COLA forms a chain of LoRAs by iteratively tuning, merging, and extending LoRA modules, as depicted in Figure 1. We denote the length of the chain in COLA as the number of residual LoRA modules optimized. For COLA with a chain length of M, the three sub-steps in Figure 1 are repeated M times. Below we describe the three sub-steps in detail.

Tune LoRA: In this step, we perform standard LoRA tuning, i.e., learning only the A and B matrices and leaving all other model parameters untouched. At initialization of COLA, this step learns LoRA modules $(A_{1},B_{1})$ on top of the frozen pre-trained LLM weights $W_{pretrained}$ . After the initial phase of COLA, the LoRA modules $(A_{i},B_{i})$ are fine-tuned on top of fixed model weights incorporated with previously learned LoRAs’ weights. The fixed model weights at the i-th iteration of COLA is $W_{pretrained}+\sum_{j=1}^{i-1}B_{j}A_{j}$ .

Tie a knot: After the current LoRA modules $(A_{i},B_{i})$ are trained, we merge them into the previously frozen LLM weights and we refer to this step as ”tie a knot”. This way, we incorporate the weight update, approximated by $B_{i}A_{i}$ , into the frozen model weights. The resulting frozen model weights becomes $W_{pretrained}+\sum_{j=1}^{i}B_{j}A_{j}$ . This allows learning only the residual information $\Delta W^{\star}-\sum_{j=1}^{i}B_{j}A_{j}$ for the next iteration. Additionally, merging the LoRA modules into the frozen LLM helps reduce memory burden under limited resource scenarios. Instead of storing a list of LoRA modules introduced in the COLA, merging them to the frozen model weights in a running fashion helps keep the GPU memory consumption the same as training LoRA only once.

Extend the chain: We extend the COLA chain by re-initializing a new set of LoRA module $(A_{i+1},B_{i+1})$ to learn the residual weights update needed to adapt the LLM to certain task. In this step, the newly introduced $A_{i+1}$ adopts Gaussian initialization and $B_{i+1}$ is initialized to zero, following Hu et al. (2021). Additionally, we reset all of the optimizer states, including but not limited to the parameters to be optimized and the gradient history.

The training cost of COLA is determined by the rank of the LoRA modules used to form the chain. The training computation for COLA is the same as LoRA when the rank is the same. In COLA, progressively lowering the rank of the LoRA modules may be an effective strategy to approximate optimal residual weight updates for specific tasks and lower the overall training cost. We explore this direction in our experiment section. At inference, all of the learned $B_{j}A_{j}$ can be integrated into the original model weights. Since $W_{\text{pretrained}}$ has the same shape as $B_{j}A_{j}$ , the final integrated model weight has the same number of parameters as the original pre-trained LLM. Therefore, no latency overhead is introduced during inference.

Convergence of COLA and the Nonconvex Frank-Wolfe method

The COLA algorithm described in figure 1 is motivated by and closely related to the Frank Wolfe algorithm (Frank et al., 1956). To see this, notice that COLA is an iterative algorithm whose iterations are succinctly described by the equation

Taking the linear Taylor approximation we can write

and thus, a constrained minimization over a set $\mathcal{K}\subseteq\mathcal{R}^{d}$ can be seen to be approximately

This is reminiscent of the Frank-Wolfe algorithm, which was historically developed in the context of linear programming. Below we analyze a variant of the Frank Wolfe algorithm for stochastic non-convex smooth optimization. The algorithm pseudo-code is given in Algorithm 2, and it is written in COLA notations as an application to fine tuning of LLM. The stochasticity is captured in equation (1), where it is assumed that the direction of the gradient is approximated up to $\varepsilon$ using a stochastic gradient method.

Specifically, we assume that COLA performs gradient updates such that after every epoch we have that

Notice that we have replaced the low rank matrices $A,B$ with a single matrix $W$ . This deviates from the exact specification of COLA, but can be justified according to the following intuition. Linear optimization over the trace norm ball results in a rank one solution, as shown in the context of the Frank Wolfe method in (Hazan, 2008; Allen-Zhu et al., 2017). In COLA, we perform nonconvex optimization over $A,B$ directly, and their rank can be larger than one.

Below we give an analysis of this algorithm which incorporates the stochastic approximation of the iterates $A_{t},B_{t}$ .

Henceforth, let $h_{t}=\mathcal{L}(W_{t})-\mathcal{L}(W^{*})$ , and

The latter quantity is a metric of convergence in nonconvex optimization, which is sometimes called the Frank-Wolfe gap. Notice that $g_{t}$ is zero if and only if the projected gradient of $\mathcal{L}$ at $W_{t}$ is zero.

The following theorem establishes that Algorithm 2 guarantees average duality gap approaching zero for stochastic smooth nonconvex optimization, as long as the distribution shift is bounded sublinearly with time.

Algorithm 2 applied to a sequence of stochastic gradients of $\beta$ -smooth nonconvex functions that are bounded in $\mathcal{K}$ by $M$ , with step sizes $\eta_{t}=\frac{\sqrt{M}}{D\sqrt{\beta T}}$ attains the following convergence guarantee

We denote $\nabla_{t}=\nabla\mathcal{L}(W_{t})$ . For any set of step sizes, we have

Here we denoted by $D$ the diameter of the set $\mathcal{K}$ . We reached the equation $g_{t}+\varepsilon\leq\frac{h_{t}-h_{t+1}}{\eta_{t}}+\eta_{t}\frac{\beta D^{2}}{2}$ . Summing up over all iterations and normalizing we get ,

Experimental Setup

In this section, we initially outline the tasks and models, followed by an introduction to the methods under comparison in our study. Finally, we provide details on the implementation.

models: We experiment with COLA to fine-tune OPT-1.3B (Zhang et al., 2022) and Llama2-7B (Touvron et al., 2023). Both models’ pre-trained checkpoints are from HuggingFace.

datasets: We evaluate the effectiveness of our method and compare it with the LoRA baseline on task adaptation across seven classification tasks: SST-2, WSC, CB, WIC, BoolQ, MultiRC, and RTE.

methods compared: For the current writeup, we mainly compare with LoRA, a representative PEFT method to train only low-rank matrices while keeping the pre-trained model parameters frozen. For future work, we will add in more baselines.

2 implementation details

We implemented our method with the PyTorch and Transformers library (Wolf et al., 2020). All experiments are carried out on NVIDIA A100 (80G) GPU.

We adopt the experimental configuration outlined in Malladi et al. (2023), where we randomly select 1000 examples for training, 500 for validation, and another 1000 for testing across each dataset under consideration. In COLA raining, we use AdamW (Loshchilov & Hutter, 2019) as the base optimizer and train for a total of 5 epochs. For a fair comparison, we keep the epoch number consistent with our baseline. A linear learning rate schedule is applied with the initial learning rate selected from $\{1\times 10^{-3},8\times 10^{-4},5\times 10^{-4},1\times 10^{-4},5\times 10^{-5}\}$ . The batch size is chosen from $\{4,8\}$ . The reported results represent the best score after hyperparameter grid-search for all experiments, conducted over five random seeds.

In implementing LoRA, we adhere to the practice outlined in Hu et al. (2021), introducing trainable linear low-rank modules to both query and value projections within all self-attention layers. While some research (Zhang et al., 2023) has explored the application of LoRA to all projection matrices or all weight matrices, the specific choice of where to apply LoRA is not a pivotal aspect of our work. For OPT experiments, we incorporate bias into the injected LoRA modules, aligning with the approach taken in Mahabadi et al. (2021). Conversely, in Llama-2 experiments, we deliberately disable bias in LoRA to ensure module key matching with the pre-trained checkpoint ”meta-llama/Llama-2-7b-hf.” In all experiments, we set the rank of LoRA (denoted as ”r”) to 8 and $\alpha$ to 16, where the ratio $\alpha/r$ is employed to scale the weight updates.

Results and analysis

We report the test performance of our method and baseline across various tasks in this section. The experiment results on OPT-1.3B are detailed in Table 1, and the results for Llama2-7B are provided in Table 2. Notably, our method consistently outperforms LoRA on all datasets under the same training budget, showcasing its superior performance.

Specifically, for OPT-1.3B experiments, COLA brings a performance boost to LoRA by 3.66 (relative improvement of 6.47 $\%$ ), 1.38 (relative improvement of 1.95 $\%$ ), 1.66 (relative improvement of 2.29 $\%$ ) on tasks WSC, BoolQ and RTE, respectively.

For Llama2-7B experiments, COLA boosts the test score on RTE from 82.09 to 85.70, which corresponds to a 3.61 gain and 4.40 $\%$ relative improvement. On SST-2, the average test scores for both our method and the baseline are the same, possibly due to the relatively low task complexity and the utilization of a subset of test examples.

In our reported results, as detailed in Table 1 and Table 2, we maintain consistency by setting the rank of all injected modules in the sequence to 8, aligning with the baseline LoRA setup. Additionally, we use an equal training epoch budget for different methods and thus ensuring the same training computation cost, as explained earlier in Section 3 .

2 Ablation Study

Different number of LoRAs in the chain: As described in Section 3.2, COLA consists of repeated iterations of LoRA tuning and merging. We denote the length of COLA as the number of LoRAs learned and merged in the fine-tuning process. To investigate the effect of the chain length of COLA on task adaptation performance, we further conduct experiments by varying the number of LoRAs. Specifically, we studied chain length of 1, 2, 3 and present the findings in Figure 2.

Here, chain length of 1 corresponds to the baseline LoRA fine-tuning. All experiments are conducted with a total of 5 training epochs. For example, in COLA experiments with chain length of 2, the first LoRA training phase lasts from epoch 1 to epoch 3. After the first LoRA module merges with the pretrained LLM weights and optimizer states reinitialize, the second LoRA starts from epoch 4 to epoch 5, which in total uses the same 5 total training epochs. All experiments results are reported over five random seeds.

As shown in Figure 2, there is a growing trend of test accuracy as the chain length increases across tasks. This is consistent with our hypothesis that residual learning of LoRA modules will lead to a better approximation of the optimal weight update to the fixed pre-trained LLM for task adaptation. For a majority of tasks, COLA is more robust in terms of generalization error compared to baseline LoRA, as shown by COLA’s smaller standard deviations.

Rank step-down: Since COLA is a residual learning paradigm, we hypothesize that the residual weight update to be learned for task adaptation should be progressively lower in rank. Therefore, instead of using a chain of LoRAs with a fixed rank of eight, as described in Section 6.1, we conduct further studies on lowering the rank.

Here, we consider a simple setting of COLA with length of two. We fix the rank to 8 for the first three epochs and set the rank for the remaining epochs to either 2, 4, 6, or 8. We show the results in Figure 3 and report the test performance in Table 3.

Figure 3 show that COLA with rank step-down outperforms LoRA with a fixed rank of 8 for all tasks (with the exception of one data point–WIC with rank 2). Thus COLA with rank step-down offers both superior generalization ability over standard LoRA and lower computational cost. In addition, our results indicate that the optimal rank to use for COLA is task-dependent. The CB and WSC tasks both benefit from higher rank LoRA modules in the second learning phase. The WIC task, on the other hand, surprisingly shows maximal test accuracy at a rank of 4 for ( $A_{2}$ , $B_{2}$ ).

Computation comparison Table 3 provides a detailed comparison of the training computation cost between COLA of different rank step-down configurations and the baseline.

The training FLOPs are obtained from the HuggingFace trainer state, and are reported as the aggregate over five random seeds. The baseline LoRA uses a fixed rank of 8 throughout training, while COLA starts with rank 8 and continues with different ranks in the residual learning phase. As expected, stepping down the rank in the chain results in higher FLOPs savings. Overall, COLA offers lower generalization error with less compute.

Conclusions and future work

In this work, we introduce Chain of LoRA (COLA) for efficient fine-tuning of large language models. The idea is to use an iterative low rank residual learning procedure to approximate the optimal weight update needed for task adaptation. Preliminary experimental results show that COLA consistently outperforms previous baselines albeit using the same, or less, computational resources.

We are actively working on applying COLA with different base optimizers and further experiments on larger scale LLMs. We are also experimenting beyond classification tasks, namely generation, summarization, and multiple choice.