Combining Parameter-efficient Modules for Task-level Generalisation (2023)

Abstract: A modular design encourages neural models to disentangle and recombine different facets of knowledge to generalise more systematically to new tasks. In this work, we assume that each task is associated with a subset of latent skills from an (arbitrary size) inventory. In turn, each skill corresponds to a parameter-efficient (sparse / low-rank) model adapter. By jointly learning adapters and a routing function that allocates skills to each task, the full network is instantiated as the average of the parameters of active skills. We propose several inductive biases that encourage re-usage and composition of the skills, including variable-size skill allocation and a dual-speed learning rate. We evaluate our latent-skill model in two main settings: 1) multitask reinforcement learning for instruction following on 8 levels of the BabyAI platform; and 2) few-shot fine-tuning of language models on 160 NLP tasks of the CrossFit benchmark. We find that the modular design of our network enhances sample efficiency in reinforcement learning and few-shot generalisation in supervised learning, compared to a series of baselines. These include models where parameters are fully shared, task-specific, conditionally generated (HyperFormer), or sparse mixture-of-experts (TaskMoE).

## introduction task-level generalisation involves training a model on multiple tasks (in parallel or sequentially) and then performing zero-shot inference or few-shot adaptation on new tasks @xcite . however, the training signals from different tasks may interfere with each other @xcite or lead to catastrophic forgetting of previous knowledge @xcite . moreover, when the nature of the tasks varies between training and evaluation, models often struggle to generalise systematically @xcite . a potential solution to these challenges consists in a modular design of neural architectures. in fine-grained skill selection figure 1 : a diagram of our latent-skill model: 1) a row of the task-skill binary matrix is selected according to the active task; 2) the sparse (sft) or low-rank (lora) adapters corresponding to active skills from a layer-specific inventory are combined; 3) the resulting combination is applied to the weights of a neural network, e.g. bart in our multi-task experiments. previous work on task-level generalisation focused on settings where the skills relevant for new tasks are already known a priori @xcite . however, this requires expert knowledge, possibly with sub-optimal granularity and limited to a few domains. in alternative, mixture-of-expert (moe) methods such as task-moe @xcite ) learn a routing function that allocates modules to tasks end-to-end. however, moes mostly focus on scaling large lan-guage models trained from scratch, and their effectiveness for out-of-domain generalisation is a moot point @xcite . in this work, we introduce a modular architecture that can efficiently adapt to new tasks, by virtue of a series of inductive biases that facilitate generalisation. the full model is shown in figure 1 . for each task, we learn a binary vector that specifies which skills (from a fixed inventory) are used to solve it. in the example, skills 2 and 4 are active. we implement each skill as a parameter-efficient adapter, either sparse @xcite or low-rank @xcite . the parameters corresponding to every active skill (red and blue in the example) are then combined, by simple averaging. the resulting combination adapts a pretrained model, such as bart @xcite , towards a specific task during multitask learning. the variable-size allocation of skills encourages their re-usage and re-combination in new tasks. to promote a coarse-to-fine dynamic where skill allocation is determined before skill-specific parameters, we explore a dual-speed learning rate for these two components (see section 2.3). we evaluate our model on reinforcement learning on babyai @xcite , a platform for instruction following in a simulated environment, and supervised learning on crossfit @xcite , a benchmark recasting 160 nlp tasks as text-to-text generation problems. compared to a series of competitive baselines, we obtain higher sample efficiency and higher performance in fewshot adaptation to held-out tasks. in particular, our latent-skill model surpasses non-modular methods for multi-task adaptation, such as hyperformers @xcite and ca-mtl @xcite , and mixture-of-experts methods such as task-moe @xcite . finally, we probe the learnt task-skill allocation matrix and illustrate how our method also enhances interpretability, as it discovers which pairs of tasks require common skills. the code is made available at https://github.com/mcgill-nlp/ polytropon . ## a latent-skill multitask model the goal of multitask learning in modelling a set of tasks @xmath0, . . . , t |t | ) is two-fold: 1) increasing sample efficiency on each seen task by borrowing statistical strength from the others; and 2) attaining systematic generalisation, the ability to adapt robustly to new tasks, possibly based on a few target-domain examples. in particular, in supervised learning, each task t i is associated with a dataset d i ≜ {(x 1 , y 1 ), . . . , (x n , y n )} and a loss function l(ŷ, y), where x is an input and y is an output. in reinforcement learning, each task is characterised by an initial state distribution q(x 1 ), a transition distribution q(x t+1 | x t , a t ), and a loss function l(x 1 , a 1 , . . . , x h , a h ) → r, foot_0 where x is a state, a is an action, and h is the temporal horizon of each episode. thus, each task defines a markov decision process (mdp). fully sharing the parameters of a model across tasks (stickland and murray, 2019) may exhaust its capacity or create interference among the gradients from task-specific losses @xcite . these limitations can be countered by instead composing task-specific adapters @xcite or softly sharing parameters @xcite @xcite . the first method leads to an explosion in parameter count, which grows with the number of tasks. moreover, the second method suffers during fewshot adaptation to new tasks as the entangled (i.e., non-modular) knowledge may overfit the training distribution. instead, we posit that there exists a (possibly small) fixed inventory of skills @xmath1, . . . , s |s| ), where |s| ≪ |t |. each skill is an independent facet of knowledge that is reused across a subset of tasks. these assumptions guarantee both scalability and modularity. in particular, we seek to create a model that jointly learns which skills are active for which task, aggregates the corresponding skill parameters according to some deterministic function, and maximises the multitask log-likelihood: ## conclusions we argued that a modular design is crucial to ensure that neural networks can learn from a few examples and generalise robustly across tasks by recombining autonomous facets of knowledge. to this end, we proposed a model where a subset of latent, discrete skills from a fixed inventory is allocated to each task in an end-to-end fashion. the task-specific instantiation of a neural network is then obtained by combining efficient parameterisations of the active skills, such as sparse or low-rank adapters. we evaluate the sample efficiency of our model on multitask instruction following through reinforcement learning and its few-shot adaptability on multitask text-to-text generation through supervised learning. in both experiments, we surpass competitive baselines such as conditional parameter generation (hyperformer) and mixture of experts (task-moe). finally, we show that modularity helps interpret multi-task models by inferring explicit relationships between tasks according to the skills they share.