Making Small Language Models Better Multi-task Learners with Mixture-of-Task-Adapters

In the first stage of training, our main goal is to learn the correspondence between the adapters and tasks. We abandon the token-level input format in MoE (Fedus et al., 2022) and instead favor the sentence-level input, the entire sentence is taken as a whole and is not broken down into individual tokens to be entered into a module. In our work, we propose the MTA (Mixture-of-Task-Adapters) architecture, which can be seen as parallel adapter blocks. As shown in Figure 1, in the first stage, the MTA module mainly consists of Parallel Adapters and Task Weights Selector.

To establish a correspondence between tasks and adapters, with each adapter being dedicated to a specific type of task, we introduce a bias to the initialization of the weights in the Task Weight Selector. This ensures that the model is capable of enabling task-to-adapter correspondence. Using Softmax-T amplifies the effect of the bias and enhances the model’s adaptation capability. For example, the weights of the type of $i$-th task $W_{i}$ can be initialized as follows:

where $N$ denotes the number of parallel adapters, $\frac{1+\lambda}{N}$ refers to the third gating value, guiding the third adapter to pay more attention to the $i$-th type of task. Note that during model training, these weights are also learnable, making our mechanism self-adaptive. The MTA module during the first phase of training has the following formula:

where $A_{i}(x)$ is the output of the $i$-th adapter network w.r.t. the input $x$. $W$ denotes the manually initialized matrix of gates for all the types of tasks. T is the sharpening coefficient. By introducing the gate values, we guide each adapter to make major contributions to one type of tasks only. This approach enables us to retain the specific characteristics of individual tasks. Take our experimental multi-task datasets as an example. The self-adaptive weights learned during the first stage are shown in Figure 3. In addition, we insert a special token “[START]” to each input sample, aiming to capture the semantics of the whole input sequence.

After completing the first stage of training, we inherit all the weights from the first stage. To further adjust the collaborative relationship between adapters, we introduce additional shared adapters and a gate network, while freezing all training parameters except for the MTA module. The gate network is used to generate adaptive weights, aiming to obtain collaborative relationships among various adapters for different tasks by capturing the “[START]” information in the hidden layer. The formula for the gate network is as follows:

where $A^{*}(x)$ denotes the combined representations of shared and top-$K$ adapter modules. $S(x)$ is the output of the shared adapters. $G$ denotes a gate network with linear layers and activation functions.

Relying only on the first stage parallel adapters leads to direct interactions between tasks and makes it difficult to achieve a better state of collaboration. In the second stage, the top-k selection can eliminate direct interference between different tasks. To compensate for the collaborative effect of synchronous deletion without reintroducing direct interference, a shared adapter structure is added as a transition module for information interaction to moderate the direct interference between tasks.

We employ three common types of NLP tasks (text classification, natural language inference task, and question answering) to construct multi-task datasets. The datasets include:

• •

  Text classification: CR (Hu and Liu, 2004), MR (Pang and Lee, 2005), SST-2 (Socher et al., 2013) and TREC (Voorhees and Tice, 2000);

• •

  Natural language inference: SNLI (Bowman et al., 2015);

• •

  Question answering (text generation): SQuAD (Rajpurkar et al., 2016).

Following the work of T5 (Raffel et al., 2019), we add appropriate prompts for all the tasks to convert them into a uniform input format. For classification and inference tasks, we match the model outputs with the ground-truth and report the prediction accuracy. For generation, we constructed the data in such a way as to generate questions based on the answers, we report the Rouge-L score (Lin, 2004), measuring the similarity between generated and reference texts. In the implementation, we load the pre-trained weights from T5-base, with 24 transformer layers. Baseline models include the base versions of BART (Lewis et al., 2019), GPT-2 (Radford et al., 2019) and Switch Transformer (Fedus et al., 2022). The performance of ChatGPT is also reported for reference. For our method, we also conduct comparative experiments on T5-large.

In the first set of experiments, we mix multiple tasks into a test set, and then obtain four scores via a unified scoring criterion. Several points can be observed from Table 1. Firstly, our model achieves higher overall results than other models, demonstrating the effectiveness of our proposed method. Secondly, when compared to the token-level processing method of Switch Transformer, our sentence-level approach to task processing exhibits a significant performance advantage on this multi-task dataset. Finally, we compare the results of our proposed model with ChatGPT-turbo. The outputs generated by ChatGPT-turbo are more difficult to control due to its intrinsic properties. Therefore, we relax the scoring criteria to include answer words as long as they are correct for both classification tasks and natural language inference tasks. Despite this relaxation, our model still outperforms the ChatGPT-turbo baseline, demonstrating its effectiveness even against larger models after fine-tuning on common multi-task datasets.

To further demonstrate the effectiveness of our proposed two-stage training method and the shared adapter, we conduct a series of ablation experiments. Due to space limitations, we show only the most representative set of ablation experiments. As shown in Table 2, models with multiple adapter structures show significant improvements compared to the original model, regardless of whether we use T5-base or T5-large.

Our experiments have revealed that partial parameter fine-tuning performs better and requires less computational resources than full fine-tuning during the second stage of training. Furthermore, a comparison with the training results of the two-stage model architecture using global parameter fine-tuning shows that partial fine-tuning can optimize the collaboration between adapters and significantly improve model performance.