CURLoRA: Stable LLM Continual Fine-Tuning and Catastrophic Forgetting Mitigation

Catastrophic forgetting is a big challenge in machine learning, particularly in the context of continual learning [5]. Various approaches have been proposed to address this issue:

• •

  Elastic Weight Consolidation (EWC) [7] uses Fisher information to measure the importance of parameters and selectively slow down learning on important parameters.

• •

  Progressive Neural Networks [8] propose to freeze the network trained on previous tasks and add lateral connections to new columns for new tasks.

• •

  Memory-based approaches like Experience Replay [9] store and replay examples from previous tasks during training on new tasks.

As LLMs have grown in size, efficient fine-tuning methods have become crucial:

• •

  Adapter layers [10] introduce small trainable modules between layers of a pre-trained model.

• •

  Low-Rank Adaptation (LoRA) [4] decomposes weight updates into low-rank matrices, significantly reducing the number of trainable parameters.

• •

  Prefix-tuning [11] prepends trainable continuous prompts to the input, allowing for task-specific adaptations.

CUR decomposition has been applied in various domains for its interpretability and efficiency:

• •

  In data analysis, CUR has been used for feature selection and dimensionality reduction [6].

• •

  In scientific computing, CUR has been applied to accelerate large-scale matrix computations [12].

• •

  In machine learning, CUR has been explored for model compression and interpretation [13].

However, to the best of our knowledge, CUR decomposition has not been previously applied to the problem of fine-tuning large language models or addressing catastrophic forgetting in this context.

The core idea is to decompose the pre-trained weight matrices using a modified CUR approach and then fine-tune only the U matrix. This approach constrains the parameter space of possible adaptations keeping the fine-tuned parameters as small as possible to keep $\|W_{\text{adapted}}-W\|_{F}$ close to the original weight matrix frobenius norm ($\|W\|_{F}$) i.e. $W+W_{\text{adapted}}$ is so close to $W$ to avoid the deviation of the adapted output.

Given a weight matrix $W\in\mathbb{R}^{m\times n}$, we first compute the probability of each column:

$$p_{j}=\frac{\|W_{:j}\|_{2}^{2}}{\|W\|_{F}^{2}}$$

(3)

where $W_{:j}$ is the $j$-th column of $W$, while $\|\cdot\|_{2}^{2}$ denotes the square of the L2 norm of the column and $\|\cdot\|_{F}^{2}$ denotes the square of the Frobenius norm of $W$. This will give us the probability of each column. For instance, if $W$ has three columns with norms 2, 3, and 5, the probabilities are 4/38, 9/38, and 25/38 respectively.

We then invert these probabilities:

$$\tilde{p}_{j}=\frac{1/p_{j}}{\sum_{i=1}^{n}1/p_{i}}$$

(4)

where $\tilde{p}_{j}$ is the inverted probability of the $j$-th column of $W$. The same steps are followed for rows. Inverted probabilities are used to sample columns and rows with lower leverage scores, which implicitly regularize the model and limit the magnitude of fine-tuning adjustments.

Then, we sample $r$ columns and rows, where $r<n$, according to these inverted probabilities to construct $C$ and $R$, which will always be fixed, with columns and rows with lower original probabilities. This trick plays a major role in the approach as it serves two purposes:

• •

  It acts as a form of regularization, preventing the model from overfitting or moving too much towards the task and limiting the adaptation of the $U$ matrix stopping it from growing so big in magnitude.

• •

  It preserves the model’s original behavior by focusing adaptations on less influential parts of the weight matrix. In addition, since $C$ and $R$ contain actual columns and rows from the original matrix, they contribute to the stability of the fine-tuning process.

CURLoRA’s approach differs significantly from other initialization methods. Unlike LoRA’s random initialization using Kaiming-uniform or Gaussian for weight A and zeros for weight B [4], or the SVD-based initialization [14], CURLoRA offers more controlled adaptation. While these other methods ensure starting from the base model, they don’t inherently limit the growth of the adaptation matrix (AB in LoRA), potentially leading to significant deviations during training. In contrast, CURLoRA initializes the U matrix as zeros, and importantly, constructs C and R matrices using columns and rows with lower original probabilities (i.e., lower values). This unique combination ensures that the fine-tuning process not only starts from the base configuration but also remains constrained throughout training. The low-value C and R matrices act as natural limiters on the growth of U, thereby preventing large deviations and contributing to enhanced model stability during the fine-tuning process.

$$C=\text{SampleColumns}(W,r,\tilde{p})$$

(5)

$$R=\text{SampleRows}(W,r,\tilde{p})$$

(6)

$$U_{\text{init}}=0$$

(7)

Where $W$ is the original weight matrix, $r$ is the rank (number of columns/rows to sample) and $\tilde{p}$ represents the inverted probabilities used for sampling.

During fine-tuning, we update only the $U$ matrix, keeping $C$ and $R$ fixed as they play a crucial role in ensuring the stability of the process by limiting the increase of $U$:

$$W_{\text{adapted}}=CUR$$

(8)

To understand how CURLoRA helps mitigate catastrophic forgetting, we analyze its properties mathematically:

To understand why CURLoRA is expected to perform better than standard LoRA in terms of catastrophic forgetting, we can analyze the shift in the output during fine-tuning.

For a given input $x$, the original output is $y=xW$. After fine-tuning:

For LoRA: $y_{\text{adapted}}=x(W+AB)$

For CURLoRA: $y_{\text{adapted}}=x(W+CUR)$

We can quantify the shift using the Frobenius norm of the difference:

$$\|y-y_{\text{adapted}}\|_{F}=\|xW-x(W+W_{\text{adapted}})\|_{F}=\|xW-xW-xW_{\text{adapted}}\|_{F}=\|xW_{\text{adapted}}\|_{F}$$

(15)

For LoRA: $\|x(AB)\|_{F}$

For CURLoRA: $\|x(CUR)\|_{F}$

This equation measures the shift in the model’s output after fine-tuning. $y$ is the original output, and $y_{\text{adapted}}$ is the output after fine-tuning. After fine-tuning for a different task, the adapted output $y_{\text{adapted}}$ might shift. We use the Frobenius norm to quantify this shift. If the shift is small, it means that the model’s predictions haven’t changed much, indicating that the model has retained its original knowledge. As shown, the shift depends on $W_{\text{adapted}}$ i.e. to make sure the shift isn’t so big, we need to keep $W_{\text{adapted}}$ as small (in magnitude or size) as possible.

CURLoRA’s main aim is to minimize $W_{\text{adapted}}$ while ensuring that the difference $\|W-W_{\text{adapted}}\|_{F}$ remains close to $\|W\|_{F}$. By focusing on minimizing $W_{\text{adapted}}$, CURLoRA effectively controls the shift in the output, thereby preserving the model’s original behavior and mitigating catastrophic forgetting.

Theoretically, CURLoRA should result in a smaller shift because:

1. 1.

   The $C$ and $R$ matrices are directly sampled from $W$, maintaining some structure of the original matrix.

2. 2.

   The $C$ and $R$ matrices are sampled from columns and rows with lower values.

3. 3.

   Only $U$ is trained, which is constrained by $C$ and $R$.

4. 4.

   The initialization of $U$ as a zero matrix.

This constrained adaptation in CURLoRA is expected to lead to better preservation of the model’s original knowledge, thereby reducing catastrophic forgetting.

CURLoRA offers significant memory savings compared to full fine-tuning and even LoRA. For a weight matrix $W\in\mathbb{R}^{m\times n}$, the number of trainable parameters for each method, considering rank $r$ where $r<n$, is:

• •

  Full fine-tuning: $mn$

• •

  LoRA (rank $r$): $mr+nr$

• •

  CURLoRA (rank $r$): $r^{2}$

The memory savings can be substantial, especially for large matrices. In our Mistral experiment, with rank 16, the trainable parameters were:

• •

  Full fine-tuning: 7,248,023,552 parameters

• •

  LoRA: 9,437,184 parameters

• •

  CURLoRA: 24,576 parameters

This reduction in trainable parameters not only saves memory but also potentially leads to faster training and inference times.

In conclusion, CURLoRA provides multiple mathematical mechanisms that can help mitigate catastrophic forgetting:

• •

  It constrains the parameter space of possible adaptations.

• •

  It provides implicit regularization towards the original weights.

• •

  It preserves important directions from the original weight matrix.

• •

  It reduces the degrees of freedom in adaptation, limiting potential deviation.

• •

  It allows for direct control and analysis of weight stability through the $U$ matrix.

These properties suggest that CURLoRA can indeed help in reducing catastrophic forgetting while still allowing for meaningful and good adaptation to new tasks. The effectiveness of these theoretical mechanisms are validated through our experiments on various tasks and datasets, as detailed in the following sections.

Our CURLoRA implementation consists of the following steps:

1. 1.

   Decomposition: For each weight matrix $W$ in the layers we want to apply CURLoRA to, we perform the following:

   • •

     Compute column probabilities: $p_{j}=\frac{\|W_{:j}\|_{2}^{2}}{\|W\|_{F}^{2}}$

   • •

     Invert probabilities: $\tilde{p}_{j}=\frac{1/p_{j}}{\sum_{i=1}^{n}1/p_{i}}$

   • •

     Sample columns and rows according to $\tilde{p}_{j}$ to construct $C$ and $R$

   • •

     Initialize $U$ as a zero matrix

2. 2.

   Fine-tuning:

   • •

     Objective:

     • -

       The primary objective of the experiment is to evaluate catastrophic forgetting during continual learning, rather than to optimize accuracy for each individual task.

   • •

     Model Specific Adjustments:

     • -

       For GPT-2 and Mistral, the model’s "lm_head" is replaced with a task-specific output layer. During training, only the $U$ matrix is continually updated, while $C$ and $R$ remain fixed.

     • -

       Replacing the "lm_head" ensures that each task has its own task-specific output layer that remains untouched when the model is being fine-tuned on a different task, contributing to the mitigation of task knowledge degradation.

   • •

     Continual Learning Strategy:

     • -

       Once a weight matrix is decomposed, $C$ and $R$ are fixed permanently. The $U$ matrix is continually updated for each new task to facilitate continual learning.

   • •

     Application of CURLoRA:

     • -

       CURLoRA is applied to the attention layers (Query, Key, Value). [15]

3. 3.

   Inference: Use the adapted weight matrix $W_{\text{adapted}}=CUR$ for forward passes along with the original $W$ matrix i.e. $x(W+CUR)$.

We used the following datasets for our experiments:

• •

  GLUE-MRPC: Microsoft Research Paraphrase Corpus for paraphrase detection [16]

• •

  GLUE-SST-2: Stanford Sentiment Treebank for binary sentiment classification [17]

  These datasets are part of the General Language Understanding Evaluation (GLUE) benchmark [18], which includes a diverse set of tasks for evaluating natural language understanding systems.

• •

  Sentiment140: A large-scale sentiment analysis dataset [19]

• •

  WikiText-2: A dataset that we use to measure language model perplexity [20]

The datasets were selected for their diverse task requirements and common use in benchmarking.

We used Mistral 7B (v0.3) [21] and GPT-2 Large [22] as our base models. For both LoRA and CURLoRA, we used the following hyperparameters:

• •

  Ranks: [8, 16, 24]

• •

  Alpha: 1

• •

  Optimizer: AdamW

• •

  Learning rate: 2.5e-4

• •

  Scheduler: Cosine with 500 warmup steps

• •

  Training epochs: 3

• •

  Batch size:

  • -

    Mistral: 8

  • -

    GPT-2: 32

• •

  Max length:

  • -

    Mistral: 512

  • -

    GPT-2: 256

We used the following metrics for evaluation:

• •

  Accuracy: For classification tasks (MRPC, SST-2, Sentiment140)

• •

  Perplexity: For language modeling capability (WikiText-2)

Our experimental procedure was as follows:

1. 1.

   Measure initial perplexity of the base model on WikiText-2 concatenating the whole dataset into a single string.

2. 2.

   Fine-tune on MRPC and evaluate.

3. 3.

   Fine-tune on SST-2 and evaluate, then re-evaluate on MRPC.

4. 4.

   Fine-tune on Sentiment140 and evaluate, then re-evaluate on MRPC and SST-2.

5. 5.

   Re-calculate perplexity on WikiText-2.

This procedure was carried out for both LoRA and CURLoRA independently.

The experimental results align with our theoretical analysis:

• •

  Parameter Space Constraint: The stability of CURLoRA’s performance across tasks supports our hypothesis that constraining adaptations to the subspace spanned by $C$ and $R$ helps preserve original knowledge.

• •

  Implicit Regularization: The maintained perplexity on WikiText-2 suggests that CURLoRA’s implicit regularization effectively prevents overfitting to specific tasks.

• •

  Reduced Interference: The consistent performance across tasks indicates that CURLoRA successfully reduces interference between task-specific adaptations.

While CURLoRA shows promising results, there are several areas for future research:

• •

  Scalability: While CURLoRA shows promising results, its scalability to larger models needs further investigation. Further studies are needed to assess CURLoRA’s performance on larger models and more diverse tasks like instruction tuning and datasets.

• •

  Computational Complexity: Conducting detailed analysis of time and space complexity compared to full fine-tuning and LoRA.

• •

  Implicit Regularization Limitation: Implicit regularization via zero initialization of $U$ has to be further studied especially in highly dynamic environments where more flexible adaptations are needed.

• •

  Optimal Rank and Alpha Selection: Investigating methods for automatically selecting the optimal rank and alpha for CURLoRA could further improve performance.

• •

  Combination with Other Techniques: Exploring the integration of CURLoRA with other continual learning techniques could yield even better results.

• •

  Quantization Support: Exploring the implementation of CURLoRA on quantized model which may lead to QCURLoRA