SwitchCIT: Switching for Continual Instruction Tuning of Large Language Models

Large language models (LLMs) have demonstrated remarkable capabilities across numerous domains, as highlighted by OpenAI, (2023) and Bubeck et al., (2023). However, whereas LLMs pre-trained on extensive language data excel in general language understanding, they may not be optimized for every specific task of interest prompted by instructions. Therefore, there is need for continual instruction learning to adapt LLMs to evolving tasks and domains. Indeed, continual instruction learning is essential for LLMs such as GPT (Radford et al.,, 2019) to maintain their effectiveness and relevance in handling a wide range of tasks and domains.

Such models are trained on vast amounts of text data and fine-tuned for specific applications, often by learning tasks sequentially (Luo et al.,, 2023), i.e. learning on datasets pertaining to one task all at once, before moving on to the next task. The challenge lies in their ability to continually learn and adapt as they encounter new tasks and information. However, in continual instruction learning scenarios, where models are sequentially trained on different tasks or datasets, catastrophic forgetting occurs when the model’s parameters are updated to accommodate new information, leading to degradation or complete loss of performance on previously learned tasks.

A typical way to balance new learning with the retention of previously acquired capabilities in LLMs is through replaying old data. However, with the rapid iterations of LLMs for diverse and complex use cases, retaining old data becomes exceptionally challenging. Moreover, continually tuning an LLM with a large number of parameters is highly costly in terms of both computation and memory usage. Parameter-efficient fine-tuning (PEFT) such as low-rank adaptation (LoRA) (Hu et al.,, 2022) provides an option of lightweight with portable parameters, which could be paired with an LLM to perform specific tasks. Therefore, in this work, we focus on alleviating catastrophic forgetting during continual instruction tuning of LLMs, particularly with minimal data retention and its interplay with PEFT.

We propose a novel continual instruction tuning method, SwitchCIT, that alleviates forgetting of previously seen tasks by introducing a switch network to identify a task given an instruction, leveraging the clustering phenomenon of task-specific instruction vectors (Wu and Varshney,, 2024). For each new task, we fine-tune the task performance by including extra parameters created by PEFT methods such as LoRA (a self-expansion process), making the method more practical.

Catastrophic forgetting in neural networks is related to the palimpsest phenomenon that new memories rapidly overwrite old ones  (Zenke and Laborieux,, 2024). SwitchCIT may be considered as a way to avoid the need to overwrite the old memories of previously learned tasks by introducing extra parameters for a new task. Moreover, there exists a line of methods inspired by synaptic consolidation in brains that reduces the learning rate on specific weights based on their importance to previously encountered tasks. Our method has the advantage of enabling the full tuning of all weights to adapt to a specific task without restricting any particular set of weights.

We summarize our contributions as follows:

• •

  We propose a novel continual instruction-tuning approach to alleviate catastrophic forgetting by using a switch network for task routing to different specialized models tuned via PEFT.

• •

  We conduct experiments for instruction-tuning on five continual natural language generation tasks, demonstrating the effectiveness of our method compared to several baselines.

Continual Learning and Catastrophic Forgetting. The study of continual learning focuses on developing algorithms that learn from a continuous stream of data, enabling a model to acquire new knowledge while retaining previously learned information without catastrophic forgetting (Wang et al., 2024b, ). Catastrophic forgetting happens when LLMs forget previously learned information as new tasks are learned (Luo et al.,, 2023). Anonymous, (2024) provides an insightful study empirically showing that pre-trained LLMs may forget domain knowledge and tasks that were not included in the fine-tuning process, while supervised fine-tuning offers substantial benefits to the models. To counter this effect, here we use a switch network to classify tasks and route computations from their instructions. By doing so, we can fine-tune task performance by including extra parameters created by PEFT methods such as LoRA for each task.

Understanding Transformers. Prior studies have offered insightful understandings of Transformer models with focuses on the internal representations (Wu and Varshney,, 2024, 2023; Nanda et al.,, 2023) and attention mechanisms (Sun and Marasović,, 2021; Olsson et al.,, 2022). Inspired by Wu and Varshney, (2024), we design a novel method for continual instruction tuning for LLMs via switching instead of concentrating on understanding of Transformers.

Instruction Tuning. Instruction tuning is the process of tuning a model from specific instructions or prompts that will guide the model toward behaving in the desired fashion. A major issue with LLMs has been the mismatch between the training objective of the LLM and users’ objectives. Instruction tuning has been developed, in part, to address this issue. This method of training aims to align language models with human intent (Ouyang et al.,, 2022; Stiennon et al.,, 2020; Zhang et al., 2023b, ). We concentrate our work on the specific case of continual instruction tuning across different tasks, which presents unique challenges such as catastrophic forgetting.

Parameter-Efficient Fine-Tuning. PEFT addresses the challenge of needing enormous computing resources to fine-tune contemporary LLMs. PEFT reduces the number of fine-tuning parameters and memory usage while still achieving similar results as full fine-tuning (Xu et al.,, 2023). One particularly popular PEFT method is LoRA, which freezes the model weights of the pre-trained model and injects trainable rank decomposition matrices into each layer of the Transformer architecture, allowing training on a small number of additional parameters rather than on the original pre-trained model (Hu et al.,, 2019). Here, we use LoRA to create extra parameters for fine-tuning tasks not yet seen by the LLM.

Mixture of Experts. Mixture of experts (MoE) models integrate multiple sub-models, or experts, to address different parts of the input space  (Jacobs et al.,, 1991; Du et al.,, 2022; Zoph et al.,, 2022). Though the MoE philosophy is similar to ours, SwitchCIT uses different models to handle different parts of a task space represented by instructions, rather than an arbitrary input space as in MoE models. Also, we separate the learning processes for model selection and the models themselves, whereas MoE models learn both simultaneously. SwitchCIT can self-expand its parameters to adapt to new tasks, whereas MoE models typically do not.

We illustrate SwitchCIT in Figure 1. At inference time, a tuple $[I,x]$ is given, where $I$ is an instruction and $x$ is an optional input. Based on the instruction $I$, a switch network routes the computation to a model trained explicitly for the predicted task such that the performance of both previously learned and newly learned tasks is mostly retained. More specifically, the switch network identifies tasks via a multi-class classification from their instructions for the routing by using instruction features extracted by a lightweight LLM, $W_{small}$. We use the last token representation of an instruction from the final layer of $W_{small}$ as the features. This design is inspired by the fact that vector representations of instructions belonging to the same task are clustered together within the hidden representation space, with the task-specific clustering phenomenon becoming more pronounced in later layers (Wu and Varshney,, 2024). Note that effective clustering implies good separability of task representations.

A selected model relies on a concatenation of the instruction and the input, $[I;x]$ to anticipate an output $y$ via an internal representation $h$ produced by a base LLM $W$ and its task-specific weight $\Delta{W}$. For brevity, we omit details about the computations of $h$ and of reaching $y$ from $h$, which involves a causal decoding process; see Vaswani et al., (2017); Hu et al., (2022) for more details. Therefore, the switch network allows tasks to be handled by models dedicated to them. Models tailored to different tasks will not interfere with one another, which consequently alleviates catastrophic forgetting of previously learned task. Here, both $x$ and $y$ could be considered as textual sequences in the context of language generation. All models $M_{1}$-$M_{T}$ are instruction-tuned for different tasks ($1$ to $T$) by introducing extra parameters $\Delta{W}$ through a PEFT method such as LoRA. The switch network may be easily implemented as a multi-layer perceptron (MLP) model with an instruction feature encoder.

In this section, we briefly overview the model implementation and datasets. Experimental setups are further detailed in Appendices B and C.

We use BLOOMZ 1.1B and BLOOMZ 7.1B  (Muennighoff et al.,, 2023) as two choices of our base LLMs and LoRA (Hu et al.,, 2022) to learn task-specific and portable parameters. A switch network is implemented using a two-layer MLP network. We use OPT-125M (Zhang et al.,, 2022), a lightweight LLM with only 125 million parameters to extract features for the switch network. Implementation and training details are in Appendix B.

We consider the benchmark introduced by Scialom et al., (2022) and five continual instruction tasks for natural language generation selected by Luo et al., (2023) to differ from the training and evaluation tasks of BLOOMZ. These tasks are:

• •

  Text Simplification (Simp), (Jiang et al.,, 2020) and (Alva-Manchego et al.,, 2020), paraphrasing a given text into simpler language;

• •

  Empathetic Dialogue Generation (Emdg) (Rashkin et al.,, 2019), generating a dialogue response that offers a reason within a specified emotional context.

• •

  Inquisitive Question Generation (InqQG) (Fan et al.,, 2019), create questions that prompt long-form answers.

• •

  Explanation Generation (Exp), according to Camburu et al., (2018), this involves generating natural language explanations for provided premises, hypotheses, or labels.

• •

  Headline Generation with Constraint (HGen) (Scialom et al.,, 2022), which focuses on producing headlines that include specific keywords, positioned as per the constraints specified. We evaluate a model’s performance on each task using metrics outlined by Scialom et al., (2022).

Specifically, we use SARI for Simp; BERTSore (BS) for Emdg, InqQA, and Exp; Rouge-1 (R1) for HGen. We use the BERTScore version based on DeBERTa-MNLI.

Following Luo et al., (2023), we train a model based on the order of instruction tasks: Simp → Emdg → InqQG → Exp → HGen. We use the specific prompts designed by Scialom et al., (2022) and train on 100,000 data samples, following previous works (Luo et al.,, 2023; Scialom et al.,, 2022). We feed only an instruction without using the prompt template to the switch network for determining task identity, since the prompt does not contain useful information for task identification and this is sufficient for high performance as shown in Table 1.

We compare our method to a direct supervised fine-tuning approach and a continual learning method with a rehearsal mechanism following Scialom et al., (2022). We consider the rehearsal method because it is the only method in the literature that is developed in an experimental setting closest to our continual instruction tuning setting and a representative of continual learning methods based on data replaying. Both our method and the other methods employ the same PEFT using LoRA.

To ensure a fair comparison, we use the same amount of data for rehearsal as we do for training our switch network. To investigate the regime of low data retaining, we reserve only 0.01% of the training data for each task to train the switch network, which is significantly less than the amount of replay data such as the 1% used by traditional continual learning methods for rehearsal. We also evaluate methods vi using 1% of the training data, based on which the rehearsal method was originally developed. Details of the instructions, prompts, and training are in Appendix C. We evaluate these methods on test splits of various task datasets as detailed in Appendix C.

Table 1 presents the progressive performance of task classification by our switch networks trained under different conditions: a low data rate-0.01% and 1% comparable to the replay data used by the rehearsal method of Scialom et al., (2022). Note that after learning each task, we retrain a very lightweight switch network to accommodate the newly learned task. It is evident that switch networks of different settings achieve very high classification accuracy at every learning stage, even when using a lightweight LLM like the OPT-125M for feature extraction. The performance is scaled up by using more training data. Performance does not obviously degrade when including more tasks, demonstrating good robustness. Notice that competitive classification performance is reached even with 100X less data, which may be explained by good separability of different tasks due to task clustering shown by Wu and Varshney, (2024).

Table 2 demonstrates that SwitchCIT not only improves upon the performance of the base LLM (”Initial”) but also outperforms other methods on most tasks after the continual learning process under the same settings. The gap becomes more pronounced when retaining only a very limited amount of data from previous tasks (0.01%). Our method also surpasses the rehearsal method using the same replay data rate (1%) studied by Scialom et al., (2022). Additionally, we compare SwitchCIT (BLOOMZ-1.1B, 0.01%) relying on PEFT to a much larger fully supervised fine-tuned approach based on the BLOOMZ-7.1b model (Luo et al.,, 2023; Muennighoff et al.,, 2023), which is approximately seven times larger than our base LLM. Surprisingly, our method still achieves better performance on most tasks, underscoring its effectiveness. In contrast to the rehearsal method, our approach experiences notably less performance degradation when reducing replay data from 1% to 0.01%. We hypothesize that due to our high-performing switch network, our method is able to specifically select a tailored sub-model for a task, yielding impressive performance.

We calculate a relative gain, a normalized score using the performance achieved by a model when fine-tuned only on one specific task by following Scialom et al., (2022). A high relative gain indicates effective retention of performance on a specific task. From Figure 2, notice that SwitchCIT experiences minimal catastrophic forgetting compared to other approaches, as shown by almost perfect retention of performance for various tasks. The less perfect retention on the InqQG task is due to imperfect task identification by the switch network. In contrast to our method, the rehearsal method retains only some performance on previous tasks. Along with SFT, they both exhibit noticeable levels of catastrophic forgetting as they continue to learn additional tasks.

Many existing works overcome catastrophic forgetting by imposing constraints on the existing parameters of a model, so their model sizes will not change. SwitchCIT introduces new parameters for each additional task. For example, when using BLOOMZ 1.1B as the base LLM, these additional parameters account for only 0.878% of the total parameters. However, only extra parameters specific to a task are loaded during inference. Considering five continual tasks, the additional parameters amount to just 4.39% in exchange for minimal catastrophic forgetting, demonstrating their lightweight and practical feasibility. Note that only the additional parameters specific to a single task are loaded during inference. We anticipate further improvements in these numbers as parameter-efficient methods continue to advance.

Separating the development of the switch network from the instruction-tuned models greatly enhances SwitchCIT’s portability. For instance, to improve task identification by our switch network using more data (from 0.01% to 1.0%, as shown in Table  1), we only need to retrain the switch network and plug it in the existing instruction-tuned models. Conversely, we can also use existing switch networks for better instruction-tuned models as shown in Table  2, where we leverage the same switch network for models with larger base LLMs such as BLOOMZ 7.1B.

We proposed a novel continual instruction-tuning approach to alleviate catastrophic forgetting by using a switch network to identify tasks and then route computations to parameter-efficient tuned models. Experiments conducted on five instruction-based continual natural language generation tasks demonstrate the effectiveness of our method compared to several baselines.

Because of computational constraints, we could only tested our method on relatively small-scale LLMs. However, according to our design, our high-performing switch network is independent of the base LLM and can be paired with larger-scale LLMs that offer superior performance. The remarkable performance of our switch network showcases the effectiveness of our method. It is worth noting that our task classifier in the switch network is incredibly lightweight (154K parameters) and requires minimal data (0.01% of the training data), making it highly practical and easy to integrate. The parameters introduced by LoRA account for less than 2% of the total parameters of the base LLM, contributing to the overall lightweight nature of our method.

Our current approach does not facilitate learning transferable knowledge across continual tasks. Exploring methods to enable our model to leverage transferable knowledge across tasks will be an important future direction for improvement.