Eliciting and Understanding Cross-task Skills with Task-level Mixture-of-Experts (2022)

Abstract: Recent works suggest that transformer models are capable of multi-tasking on diverse NLP tasks and adapt to new tasks efficiently. However, the potential of these multi-task models may be limited as they use the same set of parameters for all tasks. In contrast, humans tackle tasks in a more flexible way, by making proper presumptions on what skills and knowledge are relevant and executing only the necessary computations. Inspired by this, we propose to use task-level mixture-of-expert models, which has a collection of transformer layers (i.e., experts) and a router component to choose among these experts dynamically and flexibly. We find that these models help improve the average performance gain (ARG) metric by 2.6% when adapting to unseen tasks in few-shot settings, and by 5.6% in zero-shot generalization settings. Further, we show that the learned routing decisions and experts partly rediscover human categorization of NLP tasks – certain experts are strongly associated with extractive tasks, some with classification tasks, and some with tasks requiring world knowledge.

## introduction pre-trained transformer models @xcite have demonstrated remarkable capabilities in natural language processing (nlp) in recent years. moreover, generative transformers can be viewed as a universal model that can be optimized for any language task primed into text-to-text format @xcite . recently, researchers found that training these transformer models to multi-task on a diverse collection of nlp tasks is beneficial -not only are they better at handling seen tasks @xcite , but also at generalizing and adapting to unseen tasks @xcite . expert 1.1 task a imdb task b qqp does this help multi-task learning? ( §5) does this improve generalization to new tasks? ( §6) what is learned by each expert? ( §7) expert 1.2 expert 1.3 task c commonsenseqa expert 2.1 expert 2.2 expert 2.3 expert 3.1 expert 3.2 expert 3.3 expert 4.1 expert 4.2 expert 4.3 task-level mixture-of-experts in this work, we train such models to multi-task on diverse nlp tasks, aiming at modeling skill sharing explicitly and understanding the learned patterns. however, little is known about how multitasking capabilities and cross-task generalization is achieved, especially that the exact same set of weights is applied, and the same computation is executed, for very different tasks. humans, on the other hand, do not exhaust their brain capacity for every task at hand. humans develop skill sets and accumulate knowledge during learning, and can reuse and recompose them when facing a task. inspired by this, we hypothesize that a model that explicitly emulate skill and knowledge sharing may help improve multi-task performance and generalization to new tasks. a natural fit for this goal would be task-level mixture-of-expert models @xcite , where the model computation is conditioned on the task at hand. more specifically, the model contains a collection of experts and a router that chooses from the experts and composes the final model (fig. 1 2 ). in this paper, we first empirically investigate several key design choices for effectively training tasklevel mixture-of-experts models ( §5). we further test the model's task-level generalization capabilities by testing it on unseen tasks ( §6). compared to a multi-task bart-base @xcite baseline, our final method leads to an 2.6% improvement in the average performance gain (arg) metric when adapting to 18 unseen tasks @xcite in the few-shot learning setting. further, a gain of 5.6% in arg is obtained in the zero-shot setting with p3 dataset @xcite . lastly, we conduct a detailed analysis quantifying the correlations between the learned routes and the characteristics of tasks ( §7). we find that the routing decisions, though learned purely from multi-tasking without prior knowledge, strongly correlate with human understanding of task characteristics, such as the task being a classification task, the task being extractive, or the task requiring world knowledge. task transferability in nlp. @xcite explored supplementary training on intermediate tasks (stilt), i.e., training on a data-rich intermediate task before fine-tuning on the target task. stilt improves performance on the target task and stabilizes the fine-tuning process. @xcite and @xcite further investigated when and why intermediate task transfer works. these studies mainly focus on transferability between specific source-target pairs, while we consider a more general setting of transferring within and beyond a group of nlp tasks. ## problem setting our goal is to better understand multi-task learning with mixture-of-experts models with an explicit routing mechanism. we also hypothesize that such models help improve the model's capability to generalize/adapt to new tasks. our problem setting closely resembles crossfit @xcite . in the following, we introduce data usage ( §3.1), training procedure ( §3.2), and evaluation protocol ( §3.3). ## task-level moe transformers recall that our goal is to better elicit transferable skills during multi-task learning, and understand how those skills contribute the model performance. for this purpose we develop a mixture-of-experts variant of text-to-text transformer models, conditioning on task representations. the model contains two major components: (1) a router that selects and decides which experts to use for each task in each layer, based on its task representation; in the following, we introduce the router and the experts in more details. note that we provide a general description in this section, and leave specific design choices in §5.3 for empirical comparison. ## conclusions inspired by how humans accumulate skills from past experience and re-use them to solve new tasks, in this paper, we develop and conduct extensive experiments with transformer-based task-level mixture-of-expert (moe) models, in hope to provide new insights on multi-task learning and crosstask generalization in nlp. firstly, we empirically investigate importance design choices and quantify their influence on final model. secondly, in both few-shot and zero-shot settings, we demonstrate that task-level mixture-of-expert models are better at generalizing to new tasks. finally, by conducting a detailed analysis on the routing decisions, we find they have strong correlations with human-defined task characteristics, even when the decisions are learned spontaneously without no prior knowledge such as pre-computed task representations. we hope our work provide useful advice on training and interpreting multi-task models in nlp and we hope it will inspire future work in improving multitask learning and cross-task generalization in nlp.