PEMT : Multi-Task Correlation Guided Mixture-of-Experts Enables Parameter-Efficient Transfer Learning (2024)

Abstract: Parameter-efficient fine-tuning (PEFT) has emerged as an effective method for adapting pre-trained language models to various tasks efficiently. Recently, there has been a growing interest in transferring knowledge from one or multiple tasks to the downstream target task to achieve performance improvements. However, current approaches typically either train adapters on individual tasks or distill shared knowledge from source tasks, failing to fully exploit task-specific knowledge and the correlation between source and target tasks. To overcome these limitations, we propose PEMT, a novel parameter-efficient fine-tuning framework based on multi-task transfer learning. PEMT extends the mixture-of-experts (MoE) framework to capture the transferable knowledge as a weighted combination of adapters trained on source tasks. These weights are determined by a gated unit, measuring the correlation between the target and each source task using task description prompt vectors. To fully exploit the task-specific knowledge, we also propose the Task Sparsity Loss to improve the sparsity of the gated unit. We conduct experiments on a broad range of tasks over 17 datasets. The experimental results demonstrate our PEMT yields stable improvements over full fine-tuning, and state-of-the-art PEFT and knowledge transferring methods on various tasks. The results highlight the effectiveness of our method which is capable of sufficiently exploiting the knowledge and correlation features across multiple tasks.

## introduction fine-tuning pre-trained models (plms) has become an effective way to migrate model capabilities to downstream tasks @xcite @xcite . however, training and storing a full copy of the model parameters for each task becomes expensive as the scale of plm increases. to mitigate this problem, parameterefficient fine-tuning methods @xcite @xcite @xcite have been proposed to reduce the number of trainable parameters. despite their efficiency gains, these methods often sacrifice performance compared to full fine-tuning @xcite @xcite . recent work has proposed to distill knowledge from one or multiple source tasks and adapt it to various downstream target tasks to achieve further improvements @xcite @xcite . despite significant success, there remains a substantial performance gap between these methods and full fine-tuning. the limitations of existing methods can be categorized as follows: (1) most existing methods primarily focus on utilizing shared knowledge across all source tasks, neglecting task-specific knowledge during adaptation to downstream tasks. (2) task-specific representations of source and target tasks are typically trained independently, leading to insufficient exploitation of the correlation between them. as a result, the performance of multi-task transfer often lags behind that of distilling knowledge from a single source task. (3) the formulation of source and target tasks may be inconsistent, hindering cross-task adaptation. (4) the knowledge from source tasks is typically used as an initialization, but during fine-tuning, this knowledge may become intertwined with downstream tasks and gradually forgotten. to mitigate these challenges, we propose pemt, a parameter efficient fine-tuning framework based on multi-task transfer learning. pemt comprises two training stages for source task learning and target task adaptation, respectively. (1) in stage 1, we follow adapter-based tuning (e.g., adapter @xcite and @xcite ) to train task-specific adapters on multiple source tasks. we incorporate a sequence of task-specific prompt vectors to distinguish different source tasks and utilize task descriptions to initialize each task prompt effectively. (2) in stage 2, we train the adapter for the downstream task while incorporating knowledge from source tasks into the model. to enable multi-task transfer learning and prevent knowledge forgetting, we freeze the source adapters and integrate them using a mixture-ofexperts architecture (moe). instead of relying on a single source task, the knowledge of source tasks is incorporated as a soft combination of all adapters trained during stage 1. we employ an moe gated unit to measure the correlation between the target task and each source task, leveraging the taskspecific prompt vectors. to ensure the effective utilization of the specific knowledge from source tasks, we introduce the task sparsity loss, encouraging the moe gate to prioritize the most relevant source expert. we conduct experiments on 17 nlp datasets involving multiple tasks and domains to evaluate the effectiveness of our approach. on all benchmarks, pemt achieves an overall improvement of more than 2 points over full fine-tuning and all the compared peft methods as shown in figure 1 . under the few-shot setting, pemt also proves a significant improvement of 10 points over the compared transfer learning models. further analysis on the weights of different task experts demonstrates that the model tends to incorporate knowledge from the most relevant source task expert, which explains the efficiency and adaptability of our method. overall, this work makes the following contributions: parameter-efficient fine-tuning. parameterefficient fine-tuning freezes the original plm and introduces a small number of additional parameters for fine-tuning. existing works can be categorized into two classes, adapter-based tuning and promptbased tuning. adapter-based methods @xcite incorporate a trainable bottleneck module to each transformer layer. prompt-based tuning @xcite @xcite prepends continuous or discrete prompt vectors to the input. recently, some methods are proposed @xcite @xcite @xcite @xcite @xcite to transfer knowledge of trained adapters to downstream tasks. multi-task transfer learning. transferring knowledge from tasks has been proven to be an effective approach @xcite @xcite @xcite @xcite . many studies @xcite @xcite show the zero-shot or few-shot transferring capabilities of language models through massive multitask training over a broad range of tasks. however, the corresponding overhead could be enormous. to overcome the issue, some more recent works @xcite @xcite propose to transfer the knowledge shared by various tasks using parameter-efficient fine-tuning. among the related works, adamix @xcite and mpt @xcite are the most relevant methods. compared to our pemt, adamix trains the representation of source and target tasks independently and fails to sufficiently leverage the interdependency across tasks. mpt learns a single prompt by distilling the shared knowledge while ignoring the rich task-specific information. overview. to sufficiently exploit the taskspecific knowledge of each source task, we divide the training process of pemt into two stages, which are illustrated in figure 2 and figure 3 respectively. in the first stage, we follow the vanilla adapter-based (e.g., @xcite or adapter @xcite ) to train the source task adapters. for each source task, we freeze the original plm parameters and inject a task-specific adapter to the feed-forward layer (ffn) for each transformer layer. besides, to learn a better representation of each task, we incorporate a taskspecific description prompt which is used to measure the correlation between tasks. in stage 2, we distill the knowledge of source tasks as a weighted combination of the source adapters. the mixtureof-experts architecture (moe) is exploited to integrate the frozen source adapters and the moe gate measures the weight of each source expert using the task prompts learned in stage 1. the task prompt for the target task is a correlation-based combination of the trainable prompt vectors and the frozen prompts of the source tasks. to adapt to a downstream task, another task adapter is injected after the moe module as shown in figure 3 . ## experiment we conduct experiments on a comprehensive range of nlp datasets to demonstrate the effectiveness of pemt. the performance of different methods is compared under both full-dataset and few-shot settings. ## analysis we conduct further analysis to investigate the effectiveness of different components of pemt. weights of source adapters. in order to explore how the weights of source experts change on various target tasks, we collect the outputs of the moe gate and visualize them through histograms as shown in figure 4 . as observed, there are obvious tendencies and priorities in the weight distribution. for glue and superglue benchmarks, the knowledge of the mnli plays a dominant role, with a weighting of more than 50% of all tasks. the contributions of some individual tasks are close to 0 under the constraint of task sparsity loss. contrastively, the distribution of weights on mrqa is totally different, where the two tasks squad and record account for about 80% of the weights. the reason is that all the three datasets mrqa, squad and record belong to the q&a category, which also indicates the correlation guided moe module and the task sparsity loss effectively work as expected. number of source tasks to substantiate the scalability of pemt, we investigate how the performance changes when different numbers of source tasks are used in stage 1. as shown in ## conclusion in this paper, we propose pemt, a new parameterefficient fine-tuning framework that is capable of adapting the knowledge from multiple tasks to the downstream target tasks. pemt is facilitated with the correlation features between tasks and sufficiently leverages the task-specific knowledge of source tasks with prompt tuning and the mixture-ofexperts architecture. we also introduce novel methods to improve prompt initialization and model sparsity. experiments are conducted on a comprehensive range of datasets involving multiple tasks and domains and the results demonstrate pemt significantly outperforms existing sota methods.