Multi-Head Adapter Routing for Cross-Task Generalization
Authors: Lucas Caccia
Edoardo Ponti
Zhan Su
Matheus Pereira
Nicolas Le Roux
Alessandro Sordoni
Venue: NeurIPS 2023 poster (2023)
Keywords: Parameter Efficient Finetuning
Multitask Learning
Transfer Learning
Natural Language Processing
TL;DR: We investigate the role of learned routing for modular PEFT methods, propose a new multihead routing function and show it leads to better performance by maximizing transfer

Abstract: Parameter-efficient fine-tuning (PEFT) for cross-task generalization consists in pre-training adapters on a multi-task training set before few-shot adaptation to test tasks. Polytropon [Ponti et al., 2023] ($\texttt{Poly}$) jointly learns an inventory of adapters and a *routing* function that selects a (variable-size) subset of adapters for each task during both pre-training and few-shot adaptation. In this paper, we investigate the role that adapter routing plays in its success and design new variants based on our findings.
First, we build on the intuition that finer-grained routing provides more expressivity. Hence,
we propose  $\texttt{MHR}$ (Multi-Head Routing) which combines *subsets* of adapter parameters and outperforms $\texttt{Poly}$ under a comparable parameter budget; by only fine-tuning the routing function and not the adapters ($\texttt{MHR}$-$z$) we achieve competitive performance with extreme parameter efficiency. Second, we find that  $\texttt{Poly}$/$\texttt{MHR}$ performance is a result of better multi-task optimization, rather than modular inductive biases that facilitate adapter recombination and local adaptation, as previously hypothesized. In fact, we find that $\texttt{MHR}$ exhibits high gradient alignment between training tasks. We find that routing is most beneficial during multi-task pre-training rather than during few-shot adaptation and propose $\texttt{MHR}$-$\mu$, which discards routing and fine-tunes the average of the pre-trained adapters on each downstream tasks. This establishes $\texttt{MHR}$-$\mu$ as an effective method for single-adapter fine-tuning. We also show that $\texttt{MHR}$-$\mu$ can be used as an effective zero-shot transfer method by training the average of the pre-trained adapters for a few additional steps on the multi-task training set: this yields gains up to 3\% on absolute accuracy w.r.t. the baselines. Code is available at <https://github.com/microsoft/mttl>.

Multi-Head Adapter Routing for Cross-Task Generalization
Lucas Caccia
Edoardo Ponti
Zhan Su
Matheus Pereira
Nicolas Le Roux
Alessandro Sordoni
Microsoft Research,
McGill University, MILA,
University of Edinburgh, Université de Montréal, University of Copenhagen
lucas, page- caccia@mail www.:iaammm.mmm.m............... ca, alsordon@microsoff. ott.r/aae....ao.aaaoooooooorooor com
Abstract
Parameter-efficient tine-tuning (PEFT) for cross-task generalization consists pre-training adapters on a multi-task training set before few-shot adaptation to test tasks. Polytropon [Ponti al. 2023] (Poly) jointly learns an inventory of adapters and a routing function that selects (variable-size) subset of adapters for each task during both pre-training and few-shot adaptation. this paper, we investigate the role that adapter routing plays in its success and design new variants based on our findings. First, we build on intuition that finer-grained routing provides more expressivity. Hence, we propose MHR (Multi-Head Routing), which combines subsets of adapter parameters and outperforms Poly under comparable parameter budget; by only fine-tuning the routing function and not the adapters (MHR-2), we achieve competitive performance with extreme parameter efficiency. Second, find that Poly/MHR performance is result of better multi-task optimization, rather than modular inductive biases that facilitate adapter recombination and local adaptation, previously hypothesized. In fact, find that MHR exhibits high gradient alignment between training tasks. We find that routing is most beneficial during multi-task pre-training rather than during few-shot adaptation and propose MHR-/4, (4, which discards routing and fine-tunes the average of the pre-trained adapters on each downstream tasks. This establishes MHR-/1 for single-adapter fine-tuning. We also show that MHR-u can be used as an effective zero-shot transfer method by training the average of the pre-traince adapters for a few additional steps on the multi-task training set: this yields gains to 3% on absolute accuracy w.r.t. the baselines.
effective method
1 Introduction
The ability train effective models with relatively small number of training data is paramount importance due to the paucity of annotated examples for most tasks. One effective few-shot learning approach leverage large models pre-trained on a vast amount unlabelled data and fine-tune them the few examples available for each downstream task. To reduce the memory cost of duplicating the entire array of parameters for each downstream task, recent approaches resort to parameter-efficient fine-tunins min-ttunig (PEFT) methods, such LoRA (Hu et al. [022]] SFT [Anslll et al.) 2022), or (IA)* [Liu et al. 2022 These only fine-tune adapters while leaving the pre-trained model 'frozen'.
Nevertheless, it remains unclear how to best exploit a of training tasks better generalize set of unseen test tasks in sample-efficient fashion, based on just few examples. One straightforward
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

B
Frodd 2
LoRA
Wystraned
Plxud #
Hend
Head 1
AGS Naaa
IO O
GLUE-Ne
Active Task
Multi-Head Routing (MHR)
Z
Tisks
MHR- Z
Modules
Figure 1: Left: A LoRA adapter with weight AB is trained top of a frozen, pre-trained linear layer W. Our method MHR partitions the A, B parameter indexes into subsets (or heads). For each subset, a separate routing function selects the active modules for the current task among m copies with different parameter values, and combines them via averaging to form a task-specific head. The heads then concatenated to form the LoRA adapter. Using multiple heads allows for more fine- -ra-gnrain mixing of task parameters with a negligible increase in overall parameter count. Right: During few-shot adaptation, one can fine-tune only the multi-head routing parameters (MHR-2), keeping the modules frozen, resulting in highly parameter-efficient adaptation.
solution is to perform multi-task pre-training, i.e. first train the large model on the union of the examples from the training tasks, then fine-tune the obtained model to the test task (Liu et al. 2022 Ye et al., 202T]. However, this solution does not take into account that test tasks may require solving different combinations of sub-problems compared to training tasks Vuu et al. 2020), thus
failing to achieve compositional generalization [Rosenbaum et al. 2019 Ponti 2021]. ]. Moreover, specializing the model towards different tasks during training may result in negative transfer, due to their corresponding gradients being misaligned [Wang et al. 2021].
Several PEFT approaches have been proposed to enable better cross-task generalization by training adapters (or soft prompts) on each task independently [Preiffer et al.) 2021 Vu et al., 2021 Asai
et al. 2022, Chronopoulou et al.) 2023]. J. Given a new test task, parameters from similar training tasks are aggregated, which enables transfer. While solely having task specifce parameters is an effective strategy to mitigate interference across training tasks, it also inhibits any positive transfer within the same task pool. Polytropon (Poly) was recently proposed by Ponti et al. [2023] to address these issues: the model assumes that task-specific adapters are learned combinations of a reusable inventory of basis adapters or modules. In practice, each module is implemented as a LoRA (Hu et al., 2022) adapter. which modifies a large pre-trained model, such T5 [Raffel al. 2020). During both multi-task pre-traiinnte pre-training and few-shot adaptation, Poly learns both the inventory of adapters and (continuously relaxed) binary task-module routing matrix, which determines which module is active for each task. While Poly shows promising results, several questions remain unanswered: 1) Does the expressivity of the routing function matter? 2) Why do routing- based PEFT methods yield superior performance? 3) Is routing useful during both multi-task pre-training and few-shot adaptation?
To answer the first question, we propose a new routing function, MHR, that mixes adapters at a more granular level. Differently from Poly, where routing decisions are made for each adapter as a whole, in MHR we linearly combine subsets of the adapter dimensions (i.e. heads), each with different combination coefficients. We evaluate MHR and series of competitive baselines for few-shot task adaptation on the TO task suite [Sanh et al. 2022) and Super-Nare -Naturl Instructions [SuperNI; WWag al. 2022a Based our results, report that MHR outperforms Poly and single adapter baselines. Additionally, we show that, thanks to the increased expressivity of the routing function, it becomes possible to fine-tune only the parameters of the routing function (and not the adapters) during few-shot adaptation: the resulting method, MHR-2, z, yields competitive performance while requiring orders of magnitude fewer parameters.
Regarding the second and third questions, we uncover that optimization during multitask pretraining plays& key role in explaining the downstream performance of routing- based PEFT approaches.

---

Specifically, we find that MHR exhibits a higher cosine similarity between gradients from different tasks than Poly and single-adapter multi-task training. Hence, routing enables more knowledge transfer and less interference across tasks during multi-task pre-training. This finding led us to investigate whether routing is useful also during few-shot adaptation. It has been hypothesized [Ponti et al., [023] that one of the reasons behind Poly's performance resides in the inductive bias of the modular architecture, which allows test tasks to recombine and locally adapt the most relevant modules. To test this hypothesis, we propose MHR- where the routing function is discarded and all available adapter parameters are averaged before few-shot adaptation. We find that MHR-H can recover the performance of MHR, hinting that Poly/MHR gains are only result of better multi-task optimization. Finally, we show that MHR-U can also be used as an effective zero-shot transfer method by training the average of the pre-trained adapters for a few additional steps on the multi-tas task training set. This yields gains up to 3% on absolute accuracy w.r.t. to strong baselines such as T0-11B.
2 Background
In cross-task generalization, we are given a set of tasks T = {Ti,-., -,, TiTi), with each task Tú dataset containing a set of samples D; = {z1, .... (In, Bn)]- The set of all tasks is partitioned into training and test tasks, T Ttrain U Teval, and the objective is to leverage data in Tarain and transfer knowledge to facilitate learning of the test tasks Teval. For all the methods discussed, learning takes place in two phases, excluding the original unsupervised pre-training of the language model backbone on a separate corpus. The first phase consists of multi-task pre-training, in which either an adapter, such LoRA (IA)², the full backhone is trained the set of training tasks Ttrain- The second phase consists in few-shot adaptation, where the learned adapters are fine-tuned independently on each test task in Teval- We follow the procedure from (Raffel et al.. 2020) and formulate each task as a text-to-text problem, enabling standard eeximum--ikeelhooe training with teacher forcing (Bengio
al. 2015 and a cross-entropy loss.
2.1
Adapters: LoRA & (IA)3
LoRA (Hu et al. 2022 and (IA)a Liu et al. 2022 are two recently proposed adapter architectures
that achieve competitive trade-offs between performance and parameter efficiency [Karimi Mahabadi ett al. T Liu et al. 2022]. J. For each linear transformation corresponding the query (q), key (k),
value (T) and output (o) of the self-attention layers, LoRA modifies the base model parameters as follows:
hakon - + A.k. ( 2,
(LoRA)
where Wo are the (frozen) weights of the pre-trained model (e.g. T5 [Rafeel al. 2020 ). A, B Rdxr are low-rank learnable parameters and & 2 1 is a tunable scalar hyperparameter. (IA)², on the other hand, modifies key and value representations in self-attenioon element-wise, and also modifies the feed-forward MLP (f):
h bo - O (w/""") x); h² -(L' O/on (W(2))
((A)²)
where fh..f E Rd are learnable parameters, W the frozen parameters of the feed-forwan layer in the backbone, and a non-linearity. For clarity, we will drop the superscripts g, k, U, in the rest of the paper.
2.2 Polytropon: Adapter Routing
Typical adapter methods either fully share adapters across tasks or train individual adapters for each task. Poly addresses the multi-task problem by softly sharing adapter parameters across tasks. Each Poly layer contains 1) an inventory of adapter modules M {d1,.... &m.) } with |M] << [7|: routing function r(-) that chooses which subset of the modules to combine for cach task.
Each module corresponds to a LoRA adapter, where Φ; are its associated parameters (1) B(() E Rdxr_ r(-) is implemented as a task-module routing matrix Z [| 2, Zr. E MI is a routing vector of task T-, with cell Zrj being the probability logits of using module Φj for task T, in the current layer. Differently from mixture-of-experts [Fedus et al. [002).. which perform token-level top-k routing. Z converges to a binary matrix, defining soft partition over modules.
3

---

Method
Pre- rre-Trann
Fine- MuninTunigg
Inference
Full FT
xxd d
x
LoRA
d x 2r M +
d X 2r
x 2r
Poly
d x
M
d 2r
|
IM
2r [M|
Poly-z
d x 2r X
M +
X
M
M
MHR-H4
d x 2r x [M]
+ 'TI X
M
d 2r
d x 2r M x h
MHR-=
x 2r x [ + [T X
M
x
[M x h d 2r x (M] + (M x h
MHR
d x 2r X [M] + IT] X M x
x 2r
Table I: Number of parameters (per layer) used for each method. The calculation uses LoRA as the base adapter, modifying linear transform in Rddxd Note that the total number of parameters changed by Full FT is larger, given that the method also changes parameters for layers not modified by LORA.
This is achieved by using Gumbel-sigmoid distribution (Jang al. 2017] during training, with Zr, Gumbel(Zr.j) At each forward pass, Poly can be defined
A = Σ a.A ( B = Σ B
(Poly)
2r, where a; "E, E,,
and A ( A B Bd We normalize the mixing coefficients Zr,
for each task to ensure that the number of active modules does not affect the norm A", B". Overall, this approach enables different subsets of modules to be activated for the current layer and combined in task-specific way. Following LoRA] the output of the Poly layer is added to the output of the original layer of the frozen backbone: h Woz + sAT (BT) 3
During multi-task pre-training, for each query, key, value, and output projection in self-attention layers, the parameters learned by Poly are the adapter parameters, (A,, B.) !MM and the routing matrices Z. During fine-tuning, for each test task т, Poly randomly initialize the routing vector 2, IR L and fine-tunes both 2, and all the modules parameters M.
3 Multi-Head Adapter Routing (MHR)
In Poly, module combination remains coarse: only linear combinations of modules are possible, and thus the resulting aggregated adapter remains linear function of the modules. We propose augment the expressivity of the module combination while keeping the parameter count similar. MHR (Fig. D takes inspiration from multi-head attention [Vaswani al.) 2017]: partitions the input dimensions into different disjoint subsets, performs a separate Poly-style le combination for each of them. and finally concatenates them. This corresponds learning a different routing matrix for each subset of input features, therefore enabling the model to select different adapters for different subsets of the input dimensions. This aggregation approach is piecewise linear (i.e., linear within disjoint intervals), which allows for more expressive combinations of modules.
In each MHR layer, the routing function is a third-order tensor Z | where Z.,.A | x is 2D slice of the tensor Z. A slice represents the routing matrix for each of the h heads. [k] E R the k-th partition along the rows of the matrix W E Rdxrr, The adapter parameters A E ddxr for task t, and for each adapter layer, are computed as (similarly for B"):
Let us denote with W[k]
Zuj.k
A Σ A
with A X
(MHR)
Σ3 Zrjik
A concat(Ar AT))
where concat concatenates along the first dimension. Multi-task pre-training and fine-tuning are similar to Poly. Note that MHR results only a negligible increase in the total amount of parameters, since most of the parameters are contained in the LoRA weights A, B (Tab. I

---

Routing-Only Fine- Fine-Tuning (MHR-2) Prior work [Shao et al. 2023 inter alia] has shown that compositional generalization can be achieved by learning (re- Jcombine in novel ways pre-existing modules. We investigate whether fine-tuning the module parameters is really needed for few-shot adaptation in the context of both Poly and MHR. Therefore, we name Poly-z 2 and MHR-z the variants that, during few-shot adaptation, keep the parameters of the modules learned during multi-task pre-training fixed and just update the routing parameters Z. Crucially, this enables highly parameter- efficient adaptation: for LoRA adapters, A and B matrices constitute the overwhelming majority of parameters. Therefore, by freezing the A, B matrices and only updating Z, we can significantly reduce the parameter cost when transferring knowledge to new task.
Adapter Average Fine-Tuning (MHR-H4 To assess the importance of the routing parameters during few-shot adaptation, we propose an additional variant of MHR, MH which shares the same multi- task pre-training procedure as MHR, but for each test task т. fixes 2, = (1 (1/lM),-.
M/IM) during
few-shot adaptation. This is equivalent to discarding the routing parameters and averaging the module parameters into a single one before fine-tuning. Specifically, the adapter used during fine-tuning initialized with (similarly for BT):
1 M Ac:
A
A/: A E 1ddr
(MHR-p)
where A are the parameters of the adapters after MHR multi-task pre-training. Note that, differently from MHR, MHR-ju fine-tunes the same amount of parameters as the single adapter baseline. Thus, any difference in performance between the single adapter baseline and MHR-yu comes from differences in the adapter initialization and must be due to the optimization process taking place in the multi-task pre-training, before few-shot adaptation.
Routing Granularity In the original Poly. Ponti et al. [2023] showed that learning a routing matrix Z for each model layer gave better performance than sharing single Z matrix across all layers. We therefore investigate whether this holds true also for its multi-head counterpart, MHR. In addition, we explore intermediate approaches between one Z per layer and single one shared for the entire model. In particular, we consider sharing 1) for the adapter layers belonging to the same Transformer block; or 2) for every block of / layers, which enables us to easily trade off expressivity for parameter efficiency. As we will demonstrate section [ this an efficient mechanism to navigate this Pareto front in regimes of very small budgets of parameters per task.
4 Experiments
Our experimental evaluation aims to answer three research questions: 1) Does the expressivity of the routing function matter? 2) Why do routing- based PEFT methods yield superior performance? 3) Is routing useful during both multi-task pre-training and few-shot adaptation? We first present the basclines and datasets used in our evaluation and then discuss each question in turn
4.1 Baselines
In addition Poly, we compare MHR the following baselines for task-level generalization.
(ORA/(IA) trains single adapter common all pre-training tasks, which is then fine-tuned on each test task separately. This is arguably the most widespread approach for parameter-efficient cross-task generalization [Liu et al. 2022, Pfeiffer al. 2023].
AdapterSoup Chronopoulou et al. averages the adapter weights of the training tasks most similar given test task, before proceeding with few-shot adaptation. To compute task relatedness, we measure the cosine similarity of sentence embeddings for each task averaged over their training dataset. Notably, unlike the methods proposed in this paper, there is no knowledge sharing (nor interference) during multi-task pre-trale me-training as task adapters are trained independently.
2023 trains different adapter for each task. The method only
Wee note that all experiments
single NVIDIA A100 GPU.
5

---

TO Dataset
Avg. Test
MHR X
Backhone T5-XI.
(IA)
62.4 62.40.4
MHR-Z
Poly
AdapterSoup
62.1
LoRA
66.0
LoRA-big
65.40.9
Paly
Poly-z
Poly
66.40.3
68.01.0
Accuracy
66-
SOOA
LoRA
MHR-2
68.30.8
MHR
69.11.0
65
Backbone TO-3B
64
T-Few Liu et al. 2022
66.20.5
AdapterSoup
66. 1.66
LoRA
67 67.4p.8
IA3 AdapterSoup
Poly-2
65.31.u
Poly
69.0p.8
62-
MHR-z
68. 68.41.2
5K
500K
SM
50M
MHR
69.31.2
Adaptation Parameters
Figure 2: Left: Results of few-shot adaptation on TO dataset Sanh et al. [2022]. We report the mean of the best validation accuracy for each test task. Subscripts correspond to standard deviation. Right: Accuracy of PEFT methods on the TO dataset when applied on top of T5-XL. The x-axis shows the parameter count during the fine-tuning process.
4.2
Datasets
We test our methods on the TO et al. 2022] evaluation suite, following the same setup as Li
al. 2022]. and SuperNI Wang al. 2022a large-scale dataset with more than 1,600 training
tasks.
TO Tasks We follow the pre-training and fine-tuning procedure discussed in Liu et al. 2022 papee-prrrmeterss and losses suggested in the public codebase for T-Few
using
methods were tested with T5-XI, Raffel al. (2020) and TO-3B Sanh et al. 2022 as the backbone model. Crucially, T5 is simply pre-trained on (masked) language modelling, whereas TO is further instruction tuned: in particular, the full model fine-tuned on examples from multiple training tasks that have been augmented with task instructions. To ensure fairness for all methods, we report the median and standard deviation the best validation accuracy for each test task across seeds, when evaluated every 50 training epochs. We treat each data subset-template pair unique task, yielding a total 313 tasks.
Su To limit computational costs, we report the result on out of 119 test tasks. Tasks were chosen at random, with the requirement that at least 300 examples were available, and were equally split into 100 training, 100 validation and 100 test examples. For every method, we perform early stopping on the validation set. We report results with Rouge-L, averaged across seeds. All methods use T5-XL Raffel et al. 2020) as the backbone and not TO, as TO training tasks and SuperNI test tasks may overlap.
Results and Discussion
5.1 Does the expressivity of the routing function matter?
MHR outperforms PEFT approaches We start our analysis by evaluating the effectiveness of our proposed technique when applied over backbone that has undergone prior training on

---

instruction-followee data (TS-XL). As indicated in the TO benchmark results in the top table of Fig. [ clear that multi-head routing techniques have a distinct advantage, outperforming both single-head routing Poly by 1.1% %, and surpassing standard LoRA approaches by an impressive 3.1%. %. We also study the impact of performing instruction tuning of the full backbone before adapter training. To this end, we also experiment with TO-3B as a backhone. In the bottom table of Fig. 2 we can observe that while the relative gap between MHR and baselines smaller, multi-head routing still manages to yield favourable results. Hence, the gains of MHR compound with other multi-task methods such as instruction tuning. Finally, we turn our attention towards the SuperNI dataset (Tab. 2). Here, MHR continues to surpass analogous baselines.
MHR-2 facilitates extreme parameter efficiency Fig. 2 (right) reveals intriguing findings regarding MHR-2. When restrict training only the routing parameters Z in the original Poly, the results are unfortunately not up to par with its version where both routing and adapters are updated. However, when we apply the same constraint MHR, the performance is significantly closer the optimum achieved in this setting. In fact, MHR-2 surpasses prior baselines while simultaneously necessitating fewer parameters for effective adaptation to new tasks. Moreover, by controlling the number of layers which share the same Z allocation (see sec. 3 MHR-z is able to trade-off performance for parameter efficiency, even surpassing Poly-z in settings with only 3K trainable parameters per test task (see also & A.2.1 for more in-depth analysis). This trend is similarly observed in the SuperNI benchmark (Tab. where updates restricted the routing parameters yield performance on par with standard fine-tuning, We therefore conclude that the MHR-z represents robust approach for achieving extreme parameter efficiency in adaptation.
Additional routing heads is more beneficial than extra modules In the original Poly ap- proach, tradeoff between capacity and param- eter efficiency can be achieved by adding ex- tra modules for each adapter layer. However, this results in a linear increase in the number of multi-task parameters, which can become im- practical. To explore a more effective tradeoff, investigate the option of adding additional routing heads instead of extra modules. Fig 3 (right) presents the comparison between the two approaches. It demonstrates that increasing the number of routing heads leads to better perfor- mance compared to adding more modules. Importantly, the benefit of multi-head routing is twofold: it provides increased expressivity for the model, while also maintaining parameter efficiency. This finding highlights the advantage of multi-head routing as a more effective approach for balancing expressivity and parameter count in few-shot adaptation scenarios.
SuperN1 Dataset
Rouge-L
LORA
67.60
LoRA-big
67 07.2p.7
Poly-2
64.6p.3
Poly
67.8u.s
MHR-2
68.0 68.0p.2
MHR
68.5p..
Table 2:
Results on SuperNI dataset. Sub- scripts are standard deviation.
Routing-basee based methods also excel at the 11B scale We proceed to evaluate if Poly and MHR surpass established PEFT approaches when trained over a larger model back- bone. To accomplish this, we employ the 1IB version of TO. As depicted in Tab. B mouting-based methods once again outshine standard adapter training, surpassing our reproduction of the previous state-of-the-an in Liu al. [022) by over 2%. We observe that Poly and MHR show similar performance in standard fine-tuning, but MHR z-tuning remains more performant in routing-only fine- tuning. Indeed, MHR-2 K params) outperforms Poly-z (3.5K params) by 2.9%. while still remaining more param- eter efficient than Liu et al. 2022 (I.IM params).
TO Dataset
Avg. Test
Backbone TO-1IB
T-Few Liu et al. 2002
72.5p.9
LORA
72.31.0
Poly-z
70.0u.a
Poly
74.90.6
MHR-2
72.90.8
MHR
74.7p.6
Table 3: Few-shot results over 1 11B pa- rameter backbones.
5.2 Why do routing-basee PEFT methods yield superior performance?
While our proposed methods have demonstrated promising results across broad spectrum of datasets and varying adaptation parameter budgets, the question of why routing-based PEFT exhibits superior
7

---

Cumulative Gradient Cosine Similarity
Test Performance w.r.t Trainable Parameter Budget
Cumulative Cosine Similarity
14-
MHR
MHR Poly
12
Poly
TO Test Accuracy
69.0
heads
64 heads
10-
LoRA
68.8
8-
32 heads
128 heads
6-
68.4
4-
2-
68.0
modules
modules
64 modules
67.8
10a
10a
10c
10a
100
120
Training Steps
Multi-Task pre trainirinnngg trainable params (M)
Figure 3: Left: Gradient alignment between tasks during multi-task pretraining. Right: Increasing the number heads offer better scaling properties than increasing the number modules.
performance remains unanswered. In this section, we aim uncover the key components that drive MHR's superior performance.
Learning the Routing Function is essential Given that Poly and MHR have access to more parameters than standard adapters during multi-task pretraining, we investigate whether this, and not the routing mechanism, is responsible for their superior performance. To do so, we compare them to baseline approach. Instead of learning the routing function, we randomly assign binary module allocation to each data point in a minibatch, disregarding task information. This random routing approach, akin to Waag al. [2022b), allows to directly assess the influence of additional parameters during multi-task training. At test time, the learned modules are averaged into a single one before tneettning; we therefore refer to this baseline Random-jı. On the TO benchmark with the T5-XL backbone, Random-ju performs similarly to standard LoRA adapter (66.0%), while Poly and MHR outperform by 2% and 3.1 respectively. Therefore, we conclude that learning routing function is crucial, and merely increasing capacity during training does not directly lead to improvements.
MHR fosters transfer and mitigates interference across pretraining tasks Recognizing the pivotal role of the multi-task pretraining step in bolstering Poly's performance, we explore the extent of transfer and interference across training tasks. By monitoring the average gradient alignment for each task pair (in terms of cosine similarity) throughout the training process, we are able to gauge the level positive transfer. As Fig. 3 (left) shows. MHR displays a greater degree of gradient cosine similarity across tasks compared to other PEFT alternatives, including Poly. This finding suggests that the enhanced flexibility offered by multi-head routing may serve to mitigate interference across tasks to larger extent than standard routing while simultaneously promoting positive transfer.
5.3 Is routing important for task generalization?
TO Dataset
Test Acc.
We assessed the importance of routing during pre- training. We now proceed to verify whether it is important learn routing during few-shot adapta- tion, too. Poly-u and MHR-4+ consistently outper- form LoRA, and match the performance of Poly MHR (Tab. 4. This demonstrates that, for few-shot adap- tation, the average of the pre-trained modules pro- vides a better initialization than learning an adapter shared across all the tasks during pre-training. The consistently superior performance of Poly-μ with respect to Random-ju and AdapterSoup stresses the importance of routing during multi-task pre-training (but not during adaptation), which provides an effec- tive adapter initialization for few-shot learning. This
LORA
66.0
AdapterSoup
62.1
Poly
68.0p.s
Poly-μ
67.8u.6
MHR
69. 69.11.1
MHR-j
69.1p.9
SuperNI
Rouge-L
LORA
67.6p.8
Poly
67.8p.8
Poly-
68.3p.5
MHR
68.5u.6
MHR-14
68.5 68.5u.s
Table 4: Evaluating the impact of modular adaptation at test time.

---

Method
Zero-Shot Test on 11 Held-out Tasks
k =
0
k:
1000
k = 5000 A == 10000
Backbone T5-XL
LORA
56.5
56.0
56.1
55.7
Poly-ı
46.0
53.0
56.8
56.3
MHR-
48.0
58.0
57.1
56.3
Backbone TO-1IB Sanh et al. 2022
61.0
LORA
61.2
61.6
61.5
61.5
Poly-ı
62.1
63.6
63.9
64.4
MHR-14
63.5
64.5
64.5
64.4
Table 5: Zero-shot performance for MHR and the baselines, reported as the average the 11 evaluation datasets from Sanh et al. [2022]. To obtain these zero-shot results, we average the learnt Poly/MHR adapters before performing k additional fine-tuning steps the multi-task pretraining data. This effectively enables zero-shot transfer to downstream tasks using the same amount of parameters/flops the baseline LoRA. MHR outperform baseline LoRA by up to 3% absolute accuracy points on TO-11B.
finding could potentially inspire future work for improving meta-learning and weight-averaging approaches [zmailov et al. 2018).
MHR-/ excels at zero-shot learning For many downstream tasks of interest, additional labelled data may not be available. In such settings, it unclear how to leverage MHR-/u and Poly-u methods. To address this, we fine-tune the average of the multi-task trained adapters on the multi-task pre- training data (instead of using the downstream few-shot data), for an additional l steps. The results presented in Table []] We find that without any additional fine-tuning (k 0), averaging the adapters does not yield good results. This due to a potential mismatch between adapters Icarned via task-specific routing, and the uniform routing strategy. We can observe that when fine-tuning the average of the adapters on the multi-task pre-training data for an additional A steps, MHR-y4 show strong performance when evaluated in zero-shot manner. For fair comparison, additionally fine-tune LoRA for the same number of additional steps. Our best model achieves a zero-shot performance of 64.5 on top of TO-1 IB, achieving absolute gain of 3.5% accuracy points.
also
6 Related Work
Multi-task learning is effective for low-resource tasks (Wei et al. 2022 Aribandi et al. 2022 Sanh
al. 2022]. as knowledge can be borrowed from similar tasks by sharing the model parameters. Multi- lask lerning has also been applied across languages and modalities [Ponti et al. 2019 Bugliarello al. 2022 In the context NLP, several families of methods enable learning new tasks from limited set of labelled examples. Few-shot in in-context learning [ICL; Brown et al. 2020], where examples new task are concatenated into an input prompt, enables models generalize unseen tasks without any gradient-based training. Such approaches are however sensitive to the prompt format and example ordering [Zhao et al. 2021]. More importantly, ICL methods incur significant compute overhead, as for every prediction, the full set of examples must be processed by the model [iu et al. 2022]. To remedy this, many parameter-efie efficcien fine-tuning (PEFT) methods have been proposed as alternative to ICL, where small number of parameters added over the frozen pretrained network. To name a few, LoRA (Hu al.) 2022) injects learnable low-rank matrices into each Transformer layer. Alternatively, the learnable matrix can be sparse, selecting nonzero shifts via the Suttery-Ticker hypothesis [Ansell et al. [0021] or via their approximate Fisher information [Sung et al. 2021]. Finally, wreii--tuning methods [Li and Liang 2021] prepend learnable embeddings to the input or intermediate representations to specialize the model towards a downstream task.
Modular networks partition their parameters into several expert modules, each of them specialized to handle specific sub-tasks [Jacobs et al.) 1991 Kirsch et al. 2018). Modular networks are an appealing solution to the problem of adapting to unseen tasks
CCrona al. 2021] as the model

---

can leverage its existing modules and recombine them in novel way, thus achieving systematic generalization [Bahdanau et al. 2019]. They have also been tested in learning scenarios with data presented sequentially [Ostapenko et al.) 202T]. and with changing environments Goyal et al. [002]] In NLP,, mixture-of-expper (MoH) models [Shazeer et al. 2017 Fedus al, 2022], where learned gating mechanism routes token representations to appropriate experts (Feed-Forwane layers), have shown success in scaling the number of parameters while retaining time efficiency. This results in higher performance when compared to their dense counterparts using similar compute budget.
7 Conclusions
In this paper, we tackle the challenge of generalizing to new tasks based on few examples after multi- task pre-training. Specifically, we focus on Polytropon [Ponti et al. 2023), a model where each task is associated with subset of adapters by routing function. We investigate how varying the level of control afforded by the routing function impacts performance on two comprehensive benchmarks for multi-task learning, TO and Super-Natural Instructions. First, newly proposed variant of the routing function, where multiple heads are responsible for different subsets of input dimensions, improves consistently over all other baselines, including LoRA and (IA)² adapters. Second, we identify the cause of the success of routing in its ability to prevent interference between tasks, as it yields a better alignment between their gradients. Third, we find that simple averaging of all multi-task pre-trained adapters before few-shot adaptation to new tasks provides comparable performance, thus offering state-of-the-ar performance for single-adapter few-shot learning. Multi-head routing demonstrates the importance of fine- geeenegridd adapter selection for to improve other modular methods, such as Mixtures of Experts [MoEs: Fedus et al. 2022] in future research.
generalization and holds promise
References
A. Ansell, E. M. Ponti, J. Pfeiffer, S. Ruder, G. Glavaš, I. Vulić, and A. Korhonen. MAD-G: Multilingual adapter generation for efficient cross-lingual transfer. In Findings of the Association for Computational Linguistics: EMNIP 2021, pages 4762-4781, Nov. 2021. doi: 10. 18653v1vv 2021 122..indings-emmpp.4
hww.hhhollooyy.gg22211111./ 1ort2i..iindii findinggseennnnn 4100
A. Ansell, E. M. Ponti, A. Korhonen, and Vulié. Composable sparse fine-tuning for cross- lingual transfer. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, 2022. URL httth::// btopps://ariiie www.aawwwi......e........ee wwtp/ogg//////o//ee /2110 07560.0000000 .dd
V. Aribandi, Y. Tay, T. Schuster, J. Rao, H. S. Zheng, S. V. Mehta, H. Zhuang, V. Q. Tran, D. Bahri, Ni, J. Gupta, K. Hui, S. Ruder, and D. Metzler. Ext5: Towards extreme multi-task scaling for transfer learning. In International Conference on Learning Representations, 2022. URL http://orumm..ddah1BBBBBB
httpps://ppenrevie
A. Asai, M. Salehi, M. F. Peters, and H. Hajishirzi. Attempt: Parameter-efficient multi-task tuning via attentional mixtures of soft prompts. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 6655 2022.
D. Bahdanau, S. Murty, M. Noukhovitch, T. H. Nguyen, H. de Vries, and A. Courville. Systematic generalization: What is required and can it be learned? In International Conference on Learning Representations, 2019, URL http:ttts:://// :ttpttpps://ppppeeiiieeee wttp://www...eeoooddddddeeeeeeee iid
S. Bengio, O. Vinyals, N. Jaitly, and N. Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901 2020.
E. Bugliarello, F. Liu, J. Pfeiffer, S. Reddy, Elliott, E. M. Ponti, and I. Vulié. IGLUE: A benchmark for transfer learning across modalities, tasks, and languages. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2370-2392
10

---

PMLR, 17-23 Jul 2022. URL
mlr
html
A. Chronopoulou, M. E. Peters, A. Fraser, and J. Dodge. Adaptersoup: Weight averaging to improve generalization of pretrained language models. arXiv preprint arXiv:2302.07027 2023.
R. Corona, D. Fried, C. Devin, D. Klein, and T. Darrell. Modular networks for compositional instruction following. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1033-1040, Online, June 2021, Association for Computational Linguistics. doi: 10.18653/v1/2021 naacc-main.
81. URL
2021. www..aaaccc..a....m maiin. 81
W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, Journal of Machine Learning Research, 23(120): 2022. URL http://jmlr ptgg/paperr/v23/11099959
html
A. Goyal, A. Lamb, J. Hoffmann, S. Sodhani, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021. URL
attps://openreviev.nesss iet.:///ee....eeeoooommmmmmmooom
J. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang. L. Wang, and W. Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL
Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. Averaging weights leads to wider optima and better generalization. CoRR, abs/1 803.05407, 2018. URL. http
//arxiv. (ergg/bbs18800 http:/8833.05550 05407
R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991. htttpp:://ateemrr..slilereraraaa com/ neco 1991.3.1.70 pdf
E. Jang, Gu, and B. Poole. Categorical reparameterization with Gumbel-softmax, International Conference Learning Representations, 2017. URL https:///ppenreveer.ne.tt ntt.:/aneet/paffff/eem TkE3y85ee
R. Karimi Mahabadi, S. Ruder, M. Dehghani, and J. Henderson. Parameter-efficient multi-task fine-tuning for Transformers via shared hypemetworks. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, pages 565-576, Aug. 2021. MtRL:tttt:///ccattotoo
org/
2021. acl acl-long.-4. 47
L. Kirsch, J. Kunze, and D. Barber. Modular networks: Learning to decompose neural computation. Advances in Neural Information Processing Systems, 31, 2018. URL http:tttt///poceediiiii neurips.
X. L. Li and P. Liang. Prefix-tuning: Optimizing continuous prompts for generation. Pro- ceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Pa- pers), pages 4582-4597 Online, Aug. 2021. Association for Computational Linguistics. doi: 10.18653/v1/202/ http:////2222aalooogg..333333 a53- URL https:tpp:/ http:::::/aaaaaallllooooooooo
org/ 2org/2022 202100222222222 acl-1ong. .353
H. Liu, D. Tam, M. Muqeeth, J. Mohta, T. Huang, M. Bansal, and C. Raffel. Few-shot parameter- efficient fine-tuning is better and cheaper than in-context learning, 2022. URL https://axxiv
org/abat a20/ 05638
O. Ostapenko, P. Rodriguez, local module composition.
M. Caccia, and L. Charlin. Continual learning via Advances in Neural Information Processing Syy- tems, 34, 2021. URL https://proceedinn newrine wtt.aniiss..ppppree///...... cc/pppppr//2///tlle
fE527759900066224400233333333-83 pdf

---

J. Pfeiffer, A. Kamath, A. Rücklé, K. Cho, and I. Gurevych. AdapterFusion: Non-destructive task composition for transfer learning. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics, pages 487-503, Apr. 2021. URL
ftais:://alaanthloogy org/2021 2tw.2..c.1.cmmam.... eacl-main.
J. Pfeiffer. S. Ruder, I. Vulić, and E. M. Ponti. Modular deep learning. arXiv:2302.11529, 2023. URL httth::/ttr/rt: www..a.wrrg/pf///3//2220 11529. pdf
arXiv preprint
E. Ponti. Inductive Bias and Modular Design for Sample-Eificien Neural Language Learning PhD thesis, University of Cambridge, 2021. URL pesttps:/assaae http:///aaaaaca..ppppppooeee cam ac uk/
hetpr////ssseeessseeeeetrrrroo electronic..pe
E. M. Ponti, H. O'Horan, Y. Berzak, I. Vulié, R. Reichart, T. Poibeau, E. Shutova, and A. Ko- rhonen. Modeling language variation and universals: A survey on typological linguistics for natural language processing. Computational Linguistics, 45(3):559-601, 2019. URL http:://aateraaa.ssiieeeeeeee
wwwaaairrocmmmeeeee cww. com/coli. a a000055.ppp
E. M. Ponti, A. Sordoni, Y. Bengio, and S. Reddy. Combining parameter-efficient modules for task- level generalisation. In Proceedings of the 17th Conference of the European Chapter of the Associ- ation for Computational Linguistics, pages 687-702, Dubrovnik. Croatia, May 2023. Association for Computational Linguistics. URL httpttttsssaccaatatttt
www.llogg..rgg//// org/2023.30..2....oo eacl maai..an.....mmmm....mmmmaammi.m 49
C. Raffel, N. Shazeer, A. Roberts, K. Lee. S. Narang, M. Matena, Y. Zhou, W. Li, P. J. Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(140):1-61 1-67, 2020.
C. Rosenbaum, I. Cases, M. Riemer, and T. Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint arXiv: 2019. URL https //arxiv. cwrxiv.org/pdff.. otg/pdf/1904..m/0m 12774.pdf.
V. Sanh, A. Webson, C. Raffel, S. H. Bach, L. Sutawika, Z. Alyafcai, A. Chaffin, A. Stiegler, T. Scao, Raja, M. Dey, M. Bari, C. Xu, Thakker, S. Sharma, E. Szezechla, T. Kim, G. Chhablani, N. V. Nayak, D. Datta, J. Chang, M. T. Jiang, H. Wang, M. Manica, S. Shen, Z. X. Yong. H. Pandey, R. Bawden, Wang. T. Neeraj. J. Rozen, A. Sharma,
Santilli, T. Févry, J. A. Fries, R. Teehan, S. Biderman, L. Gao, T. Bers, T. Wolf, and A. M. Rush. Multitask prompted training enables zero-shot task generalization. The Tenth International Conference on Learning Representutions, 2022. URL https://aaax www.ooggg.....oooo... 2210 082007 .df
N. Shao, Cai, H. xu, C. Liao, Y. Zheng, and Yang. Compositional task representations for large language models. In The Eleventh International Conference on Learning Representations, 2023. URL
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. http://:tt///////ooooo/oo/////aa
Y.-L. Sung, V. Nair, and C. Raffel. Training neural networks with fixed sparse masks. In A. Beygelz- imer, Y. Dauphin, P. Liang, and W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL J https: /ttp://///peeeeeieeeeeee wttpereevviieeeeefffuueeee http://forrmmmddwwhhhb...
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. CoRR, abs/1706.037622 2017. http://///aaxii org/ /tse//bs///// 1700 03762
T. Vu, T. Wang, T. Munkhdalai, A. Sordoni, A. Trischler, A. Mattarella-Micke. S. Maji, and M. Iyyer. Exploring and predicting transferability across NLP tasks. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 7882-7926, Online, Nov. 2020. Association for Computational Linguistics. doi: 10. 10.118663/11/2020cmmmmppp maw.mmm..mmiin.5...... URL
2020. eenle-main 635
Vu, Lester, Constant, R. Al-Rfou, and D. Cer. Spot: Better frozen model adaptation through soft prompt transfer. arXiv preprint arXiv:2110.0790. 2021.
12

---

Y. Wang, S. Mishra, P. Alipoormolabashi, Y. Kordi, A. Mirzaei, A. Arunkumar, A. Ashok, A. S. Dhanasekaran, A. Naik, D. Stap. E. Pathak, G. Karamanolakis, H. Lai, I. Purohit, I. Mondal, J. Anderson, K. Kuznia, K. Doshi, M. Patel, K. K. Pal, M. Moradshahi, M. Parmar, M. Purohit, N. Varshney. P. R. Kaza, Verma, R. S. Puri, R. Karia, K. Sampat, Doshi, S. Mishra, S. Reddy, S. Patro, T. Dixit, X. Shen, C. Baral, Y. Choi, N. A. Smith, H. Hajishirzi, and D. Khashabi. Benchmarking generalization via in-context instructions on 1,600+ language tasks, 2022a. URL
07705
Y. Wang. S. Mukherjee, X. Liu, J. Gao, A. H. Awadallah, and J. Gao. Adamix: Mixture-of-adaaee -fdaptteer for parameter-efficient tuning of large language models. arXiv preprint arXiv:220S. /24/0, 2022h.
Z. Wang, Y. Tsvetkov, O. Firat, and Y. Cao. Gradient vaccine: Investigating and improving multi- task optimization in massively multilingual models. In International Conference Learning Representations, 2021. URL
nwt. faae//oorum' http://ormmmidd-d- iidFivEjjFK -1
J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le. Finetuned language models are zero-shot learners. International Conference on Learning Representations, 2022. htttt::://pperreeee
net/forum? tetet//orummddgg2eeeeeeee id=gE2rCCoziqR
Q Ye, B. Y. Lin, and X. Ren. CrossFit: A few-shot learning challenge for cross-task generalization in NLP. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 7163-7189 Nov. 2021. doi: 10. 10.11/111/12211 // /aclanthology org/ 2021.
m7in.572. URL https
Z. Zhao, E. Wallace, S. Feng, D. Klein, and S. Singh. Calibrate before use: Improving few-shot performance of language models. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12697-12706. PMLR, 18-24 Jul 2021. URL /ttts::/proceedings.nt mlr press/ 2ress/v139/zhao1.co zht. html
13

---

A Appendix
A.1
Additional Results
More detailed numbers on the TO Sanh et al. [2022] and SuperNI Wang et al. [202a datasets using different backbones, and different adapter layouts over the base model are found in Table
Multi-Task params is the number of additional parameters that must be conserved after
Model
Multi-Task Params
Adaptation Params
Avg. Test
Dataset
Backhone T5-XL-IM
Multi-Task Full Finetuning LoRA
2.8B
2.2M
68.9
(IA)
540K
540K
62.4n.4
AdapterSoup
84M
2.2M
62.1
LoRA
2.2M
2.2M
66.01.6
LoRA-big
35M
35M
65.4u.9
Poly-2
17M
3.5K
66.4u.3
Poly
17M
2.2M
68.01.0
MHR-z (64 h)
17M
220K
68.30.6
MHR (64 h)
17M
2.2M
69.1
Backbone T0-3B
T-Few Liu et al. 2022
540K
540K
66.2o.5
AdapterSoup
84M
2.2M
66.1o.6
LoRA
2.2M
2.2M
67.4o.s
LoRA-big
35M
35M
68.0o.s
Poly-z
17M
3.5K
65.31.0
Poly
17M
2.2M
69.0p.s
MHR (64 h)
17M
220K
68.41.2
MHR (8 h)
17M
2.2M
69.3
Backbone TO-3B light version: (k, ff layers only)
I-LoRA (rank
934K
934K
66.20.9
I-LoRA (rank 16)
15M
67.61.1
AdapterSoup (I-ooRA)
35M
934K
64.91.0
I-Poly-z
7.5M
2.1K
62.91.2
I-Poly
7.5M
934K
68.0o.s
J-MHR: (32 h)
7.5M
74K
66.81.1
1-MHR (8 h)
7.5M
934K
68.5o.7
SuperNI Dataset
Rouge-L
Backbone T5-XL-LM light version (k, v, ff layers only)
I-LoRA
934K
934K
67.6u.8
I-LoRA-iig
18M
18M
67. 67.2p.7
I-Poly-z
7.5M
2.1K
64.60.3
I-Poly
7.5M
934K
67.8o.s
1-MHRZ (64 h)
7.5M
147K
68.00.2
1-MHR (8 hj
7.5M
934K
68.50.3
Table 6: (top) Results on TO dataset Sanh et al. [2022] we report the mean of the best validation accuracy for each test task, when evaluated every train epochs. T-Few our reproduction results in Liu et al. 2022]. LoRA-big means LoRA adapter with a larger rank. (bottom) Results on SuperNatural Instructions dataset.
multi-task pretraining to enable transfer a downstream task. Adaptation Params refer to the number of parameters required to learn new downstream task. For e.g. Poly and MHR, the multi-task parameters includes the learned modules, but not the routing over the training tasks, as these are not required for transfer on new task. Moreover, variants which average the learned modules prior to fine-tuning (MHR-ju and Poly-u) will have both multi-task and adaptation parameters equal to that of single shared adapter, since after multi-task pretraining
average the modules.
14

---

A.2 Navigating the parameter efficiency performance frade-off of tuning only the routing
Here we provide additional results on how different routing based methods can be more expressive when only learning new routing function (over frozen modules) to adapt to a new task.
Polytropon
Multi-Head Routing (MHR)
Head #1
Actire Task
Tasks
Modules
Modules
Figure 4: Different ways control the expressivity of routing based methods. Left: Polytropon, one can only add additional modules, resulting in a linear parameter increase. Right In MHR, additional heads only introduce routing matrices Z, resulting in a negligible parameter increase.
In Fig. 4 (left), only do so by increasing the number of skills at each layer. However, this has a significant impact on the number of multi-task parameters which much be kept in order to perform few-shot transfer. MHR on the other hand, can increase routing capacity in a much more parameter efficient way.
that in order build more expressive routing functions Z, in Poly
A.2.1 On the granularity of routing tensor in MHR
Here provide additional results when modifying the granularity of Z for MHR. We see that one easily trade-off more parameters for better performance.
Performance when tuning only Z w.r.t Test Parameter Budget 68.3
TO Test Accuracy
67.5
MHR Poly
67.0
66.5
66.0
65.5
10a
10f
10a
Test-task param count
Figure 5: Routing-Oni Fine-Tuning (MHR-2)
15

---

B Broader Impact
In our work, we focus on advancing parameter-efffcienn fine-tuning methods for cross-task general- ization. While our research primarily addresses technical challenges and performance improvements, when applying such methods, it is crucial consider the potential negative societal impacts. Specif- ically, we believe that prior to applying our proposed adaptation method, critically examining the potential biases and ethical implications of the underlying large language model, and the data itself must be properly addressed. This includes issues related to fairness, privacy, and the spread of misinformation.
16


Reviews:
[{"review_id": "KrtIhdB2Tr", "reviewer": "Reviewer_HX6c", "summary": "In this work, the authors proposed MHR for cross-task generalization. To achieve extreme parameter efficiency, MHR- z and MHR-u are proposed to balance the performance and efficiency. Besides, this work emphasized the importance of the routing function, which is very insightful for the community.", "questions": "1. In Figure 3, why MHR with 32 heads performs worse than 8 heads version?\n\n2. In Figure 1, I would suggest the authors to mark where is the task-specific head. In addition, More information is needed for MHR-Z in the Fig. 1 as well.\n\n3. Does each task associated with a task-specific Z matrix? If not, how to select active modules based on the input task?", "limitations": "The authors did not present any limitations in the manuscript. \nFor the performance issue I pointed out in the weakness section, one potential reason could be the subheads are combined via averaging in your case. In some Mixture-of-experts papers, a set of dynamic weights can be generated to combine the experts dynamically based on the input information(task in your case).  This could be helpful to further improve the performance, and related discussion can be conducted in the manuscript as well.\n\n[1] Condconv: Conditionally parameterized convolutions for efficient inference, NeurIPS 2019\n[2] A mixture of h−1 heads is better than h heads, ACL 2020\n[3] Attention over Self-attention: Intention-aware Re-ranking with Dynamic Transformer Encoders for Recommendation, TKDE 2022\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Excellent presentations for motivation, and experimental results.\n\n2. I am impressed by the performance of MHR-z, the accuracy is close to poly while there are only very few parameters that need to be adjusted in the fine-tuning stage. ", "weaknesses": "1. The improvement of the MHR is quite limited, which only outperforms the Poly by 1.1%. Besides, there is no significant test, which is very important for such a marginal improvement.\n\n2. Fig. 1 is not well presented.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. Below we address your concerns point by point. \n\n#### **1. On statistical efficiency and pertinence of results**\nWe refer the reviewer to the global comments.\n\n#### **2. Clarification of the task-module routing matrix $Z$.**\nThe task-module routing matrix has shape `(n_tasks, n_modules)`. Therefore, the information specific to task $t$ resides in the $t$-th row of this matrix. This row has shape `(n_modules,)` and selects the active modules for the $t$-th task. In the case of multiple routing heads, each head has its own routing matrix. Therefore, a given MHR adapter layer has `n_heads` $Z$ matrices, each of size `(n_tasks, n_modules)`. For a given task, we retrieve the appropriate row of each $Z$ matrix, and use the prescribed mixing weights to mix the shared modules in a task-specific way. \n\n#### **3. Improving Figure 1.**\nThank you for the suggestions on our main figure. Let us try to better describe it. In Fig. 1, we are showing MHR with `n_heads=2`. Each head has its own task-module routing matrix $Z$. We assume in Fig. 1 that the input belongs to the first task, therefore highlight the first row of $Z$ as the active task. This row shows that the first and third module are selected, thus we select and average those modules for the first head. A similar process is down for the second head. Lastly, the outputs from the two heads are concatenated to form a LoRA adapter. \nFor MHR-Z, we show an example of transfer to a new downstream task. For each head, the task-module allocation matrix (vector) has shape (`n_tasks == 1 x n_modules`). In this example, only this $Z$ vector is trained (red) while the underlying modules are frozen (blue). \n\nTo improve the figure, we modified the figure to better highlight this process for the second head, and increased visibility to highlight the task specific row of the Z matrices. This version can be found in the updated pdf. \n\nPlease let us know if you have further concerns or suggestions on how to improve the figure. \n\n#### **4. Clarification on Figure 3, performance of 32 heads vs 8 heads.**\nActually, both methods perform very similarly (68.73 vs 68.79).\n\n#### **5. Addressing limitations**\nThank you for pointing out relevant papers. Indeed, the proposed routing method can potentially benefit from leveraging additional input information (rather than task-only information). We agree that this is a promising direction for future work."}, {"review_id": "IpCzEvtbY1", "reviewer": "Reviewer_5omF", "summary": "This paper introduces a method called Parameter-efficient Fine-tuning (PEFT) to improve how pre-trained language models adapt to new tasks. They use small adapters and a routing function to select specific adapters for each task. The authors found that finer-grained routing provides better results and propose a method called Multi-Head Routing (MHR) that outperforms previous approaches. They also discovered that the success of their method is mainly due to improved multi-task optimization rather than specific adapter properties. They introduce a simplified variant called MHR-μ that achieves competitive performance with fewer parameters by discarding routing during fine-tuning.", "questions": "The authors are encouraged to address the mentioned concerns in order to provide a better understanding of their work. It would be beneficial for them to provide further clarification regarding the technical novelty of their proposed parameter decomposition approach in LoRA, highlighting any unique aspects that distinguish it from existing methods such as multi-head self-attention and MoE. Offering additional insights and motivations behind their design choices would also help in assessing the novelty and value of their work.\n\nFurthermore, it is important for the authors to address the lack of specific information about the model parameters used in their experiments. Providing details about the size and configuration of the proposed model in comparison to the baselines would enable a more accurate evaluation of its effectiveness. These additional facts would assist in making a fair assessment and potentially lead to a reconsideration of the paper's overall score.", "limitations": "-- No limitation is discussed in the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "++ The paper demonstrates clear and concise writing, making it easily comprehensible.\n\n++ The authors conducted a thorough ablation study to thoroughly evaluate the effectiveness of their proposed approach.", "weaknesses": "-- The paper lacks significant technical novelty. The approach of decomposing the parameters in LoRA into different sets is not particularly interesting. Additionally, the designed modules bear strong resemblance to multi-head self-attention and MoE (Mixture of Experts). Moreover, the intuition behind these designs seems ad-hoc and lacks novel insight, which is crucial for assessing the value of the paper.\n\n-- The paper fails to provide specific details about the model parameters used in the experiments. Considering that the proposed model is substantially larger than the baselines, it is necessary to have more information in order to properly evaluate the proposal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. We address the raised concerns point by point. \n\n#### **1. On the novelty of the proposed method**\nWe refer the reviewer to the global comments\n\n#### **2. Additional information about model parameters and configuration** \nFor all experiments in the paper, all the reported methods besides ia3 use LoRA adapters, and are inserted **on exactly the same layers**. In other words, the amount of adapted layers are the same. In the main results of the paper, we add adapters to key, value, query and output linear mappings of attention mechanisms, as well as to the two linear layers in the feed-forward transformer block. Except for “LoRA-big”, all methods using LoRA use a rank of 1. \n\nWe refer to the additional parameters that must be conserved after multi-task pretraining as `Multi-Task Params`. We denote the additional parameters that must be kept for each new downstream task as the `Adaption Params`. These results are using T5-XL, the lm-adapted version with 3B params. Importantly, **MHR outperforms other baselines with more parameters** (see LoRA rank 16 and AdapterSoup). \n\n| Model             | Multi-Task Params | Adaptation Params | Avg. Test Performance |\n|--------------------|-------------------|-------------------|-----------------------|\n| IA3                 | 540K              | 540K              | 62.4                 |\n| AdapterSoup  | 84M               | 2.2M              | 62.1                 |\n| LoRA              | 2.2M              | 2.2M              | 66.0                 |\n| LoRA rank 16 | 35M               | 35M               | 65.4                 |\n| Poly-Z            | 17M                | 3.5K              | 66.4                  |\n| Poly               | 17M                | 2.2M              | 68.0                 |\n| Poly-mu         | 2.2M               | 2.2M              | 67.8                 |\n| MHR-z           | 17M                | 220K              | 68.3                 |\n| MHR              | 17M                | 2.2M              | 69.1                 |\n| MHR-mu        | 2.2M                | 2.2M             | 69.1                 |\n\nWe see that methods with MHR routing better optimize the adaptation parameter / test set performance tradeoff. Additional results with different configurations can also be found in the Appendix A.1. \n\n\nPlease let us know if you have any other questions. \n\n"}, {"review_id": "7L8OQzazRe", "reviewer": "Reviewer_1NFM", "summary": "The authors propose a new parameter-efficient routing method, Multi-head routing (MHR), which combines parameter subsets as opposed to averaging all weights together. They find this yields better performance, and even only finetuning the routing matrix after initial training works well. Furthermore, they explore the training dynamics of MHR and find its gradients are more aligned than other approaches, suggesting that routing-based training helps to mitigate negative transfer during multi-task training.", "questions": "- Do you think the adaptersoup baseline is hindered by its use of adapters? What would its performance be like if you used Lora as the PEFT module instead?\n- If averaging the learnt modules is better than routing during few-shot adaptation, why should we use the routing method after the multitask pretraining stage? Isn’t this simpler than the main MHR approach proposed?\n- Which T5 XL variant did you use? It would be useful to specify this, since some T5 variants had multi-task data inserted during pretraining (I’m assuming you used the v1.1 lm-adapted ones, which did not).\n", "limitations": "The authors do not explicitly discuss limitations, although they discuss impacts in the appendix. It would be good to directly add a section on limitations, including e.g., the need to perform training over few-shot data vs using an ICL approach, the fact that routing is not needed during few-shot training (it seems), or any other limitations you can think of.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "A straightforward approach, with well-performed experiments over a wide variety of tasks. Relaxing the routing to allow for different parameter subsets is a nice extension of the Poly work. Results seem consistent across settings. The exploration of why these routing approaches work well is interesting and poses some interesting questions that could be explored by future work.\n", "weaknesses": "- I think the adaptersoup baseline is a bit unfair since it uses a different parameter-efficient adaptation method to MHR/Poly. Ideally, you would apply the AdapterSoup approach to these other methods. Otherwise, it’s hard to say if adaptersoup is truly worse or if it is just that adapters are suboptimal compared to LoRA.\n- Gains are somewhat small over Poly, but this can be the case for parameter-efficient tuning. It would be useful to compute statistical significance.\n- The finding that gradient cosine similarity is enhanced is interesting, but is weak evidence (in my opinion) for task interference/transfer. The suggestion that MHR is helping with negative/positive transfer would be much better served with some experiment directly targeting this (e.g. examining MHR with two tasks known to cause positive/negative transfer, and examining the gains/gradients there)\n- The fact that using routing doesn't help in the few-shot stage (section 5.3) suggests that averaging and throwing away the modules learnt via MHR would be a better strategy than keeping the routing, no? Since keeping one set of PEFT weights and no routing is simpler than continuing the routing strategy.\n\nOverall, I think this paper is solid and well done. While the gains of the method are small, I think the idea is interesting and the results around routing-only finetuning and task gradient alignment are interesting and pose interesting questions for future works.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing a thorough and constructive review. We will address the questions one by one. \n\n#### **1. AdapterSoup Baseline and Backbone Clarifications**\nWe apologize for the lack of clarity regarding this. Indeed, our AdapterSoup baseline uses the same LoRA configuration (same base layers, same LoRA rank). As you pointed out, in order to properly assess the contribution of the routing mechanism of AdapterSoup, we kept other design choices fixed. Moreover, you are correct in our choice of backbone, we opted for the t5-xl-lm-adapt v1.1. Lastly, we agree with the limitations you raised regarding ICL vs adapter tuning. We will add this to the next version. \n\n#### **2. On whether to route vs average during-few shot adaptation**\nWe agree that in settings where transfer is limited to a limited set of test tasks, averaging the learned modules is a more straightforward solution. However, in settings where high levels of parameter efficiency are required per task, options such as MHR-z require that the skills be kept. Moreover, we believe that an interesting future research direction would be to apply MHR to continual learning settings, where modular methods have been shown to work well [3]. In such setting, the multi-task optimization step could be composed of multiple phases, in which case keeping the full set of skills may be beneficial. \n\n#### **3. On statistical efficiency and pertinence of results**\nWe refer the reviewer to the global comments.\n\n#### **4. On gradient alignment as a proxy for transfer / interference during training**\nWe agree with the reviewer that indeed, a more thorough investigation on how MHR effects aids transfers / mitigates interference would be beneficial. While we could not find clear task pairs which are known to cause interference, we can however look at alignment across tasks known to be similar. For this, we looked at the cosine similarity of the resulting adapters for all 36 **summarization tasks**. We found that MHR had an average alignment of **0.76** for MHR and **0.73** for Poly.  We repeated this experiment for all **question answering tasks**, and again found superior alignment for MHR than Poly (**0.77** vs **0.74**). \n\n[3] Ostapenko, Oleksiy, et al. \"Continual learning via local module composition.\" Advances in Neural Information Processing Systems 34 (2021): 30298-30312.\n"}, {"review_id": "ALsmfcG3Ho", "reviewer": "Reviewer_m1AN", "summary": "This paper proposed a new parameter-efficient few-shot fine-tuning method for pretrained language models. The method is a follow-up work of Poly. The authors proposed to fine-tuning the routing function and freezing the multi-head adapters. This way, the number of updated parameters reduced significantly while achieving a similar accuracy level in the down-streaming tasks. ", "questions": "See the weakness", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper studied an important problem in efficient fine-tuning of a pretrained language model, i.e., how to achieve a better trade-off between updated parameters and the final accuracy. The paper is well-motivated. \n- Overall the paper is well-written, especially the experimental section. The authors conducted a comprehensive ablation study on some key questions raised in the paper. \n- Although the method adopted in the paper (updating routing parameters only) was simple and straight forward, it seemed to work well in multiple different down-streaming tasks. ", "weaknesses": "- The main shortage of the paper is the limited novelty. I think in comparison to Poly, the main difference of this work is to fine-tune routing only without the adapters. The authors conducted some heuristic analysis in the experimental section to understand the intuition behind. However, the paper may still not distinguish quite clear to the prior work. I would like to see more quantitative results on why fine-tuning routing is essential and necessary, e.g., the learned multi-head adapters are orthogonal in the task space?\n\n- Some method parts are not fully clear to me. For example, the authors discussed (IA)^3 in Sec.2.1. However, the definition of h^{k,v} and h^f are not properly introduced. \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We address below your concerns point by point.\n\n#### **1. On the novelty of the proposed method**\nWe refer the reviewer to the global comments. \n\n#### **2. On why finetuning (only) the router is necessary** \nWe do not mean to claim that fine-tuning only the router is necessary. However, we show that, contrary to Poly, MHR offers this possibility: one of the motivations for fine-tuning only the router is to enable higher parameter efficiency for new tasks. If we misunderstood your question, please let us know and we will be happy to engage. \n\n#### **3. Further clarifications**\nWe apologize for not properly introducing $h^{k,v}$ and $h^f$. These vectors represented the output of the the keys and values in attention mechanisms and the inner activations in position-wise feed-forward networks.\n \n"}]

Author rebuttal:
Dear Reviewers, we thank you for your valuable feedback! We really appreciate the time you spent providing us constructive advice to improve our submission. We list below our general response to all of you, and address individual concerns in separate threads.


#### **1. Novelty w.r.t to prior work**

*i)* MHR is an application of multi-head splitting for routing-based approaches like Polytropon. Although multi-head splitting has been used in other contexts (notably in self-attention), its application for combining adapters is, to our knowledge, novel. Moreover, we show that it unlocks the ability to trade-off parameter efficiency and performance. This is not something that Polytropon enables by itself, as Poly-Z severely underperforms given that it relies on linear combinations of adapters. 

*ii)* We also provide valuable insights on the role of routing in cross-task setups. We showcase:

*a)* **when** routing is critical. Indeed, routing is critical to multi-task optimization, but not few-shot adaptation. In fact, we can average modules, which makes MHR-$\mu$ a better initialization for LoRA than standard multi-task training. We show that MHR-$\mu$ also outperforms AdapterSoup, which relies on fixed routing. This is a new finding, and no prior work has shown that this is possible with, e.g. MoEs.

*b)* **why** routing-based methods such as Poly and MHR work, by showing that routing yields stronger gradient alignment during multi-task optimization. This finding opens the door to the design of new routing methods leveraging prior work on gradient alignment in multi-task learning [1-2]. 


#### **2. Statistical significance of reported improvements**

We employed a statistical analysis to assess the presence of a significant difference in performance between MHR/MHR-z and Poly. Specifically, we executed a matched-pairs Wilcoxon signed rank (WS) test. Our assessment protocol, as outlined in our primary findings (Figure 2, T5-XL backbone), was extended to encompass five separate test seeds. The outcomes distinctly reveal that MHR/MHR-z exhibit a statistical improvement (p < .05) in relation to Poly, as evidenced by the WS-test results.


[1] Yu, Tianhe, et al. "Gradient surgery for multi-task learning." Advances in Neural Information Processing Systems 33 (2020): 5824-5836.

[2] Wang, Zirui, et al. "Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models." International Conference on Learning Representations (2021)


Meta-review: The paper proposes Multi Head Routing (MHR) that combines subsets of parameters for a previous method called Poly, for parameter-efficient fine-tuning. The method achieves better performance with comparable number of parameters. Reviewer ratings for this paper are one accept, one weak accept and two borderline accepts. Reviewers mentioned the method is well-motivated, the empirical evidence is clear, and the presentation is high quality. The AC is pleased to recommend acceptance of the paper.
Decision: Accept (poster)