SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhan Su
Matheus Pereira
Alessandro Sordoni
Microsoft Research, McGill University, MILA,
University of Edinburgh, Université de Montréal, University of Copenhagen
lucas.page-caccia@mail.mcgill.ca,alsordon@microsoft.com
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] (Po1y) 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 MHR (Multi-Head Routing), which combines subsets of adapter parameters and outperforms Po1y under a comparable parameter budget; by only fine-tuning the routing function and not the adapters (MHR- $z$ ), we achieve competitive performance with extreme parameter efficiency. Second, we find that Poly/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 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- $\mu$ , which discards routing and fine-tunes the average of the pre-trained adapters on each downstream tasks. This establishes MHR- $\mu$  as an effective method for single-adapter fine-tuning. We also show that 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.
The ability to train effective models with a relatively small number of training data is of paramount importance due to the paucity of annotated examples for most tasks. One effective few-shot learning approach is to leverage large models pre-trained on a vast amount of unlabelled data and fine-tune them on 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-tuning (PEFT) methods, such as LoRA [Hu et al., 2022], SFT [Ansell et al., 2022], or (IA) $^3$  [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 set of training tasks to better generalize to a set of unseen test tasks in a sample-efficient fashion, based on just a few examples. One straightforward
Figure 1: Left: A LoRA adapter with weight  $AB^{\top}$  is trained on top of a frozen, pre-trained linear layer  $W$ . Our method MHR partitions the  $A, B$  parameter indexes into  $h$  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 are then concatenated to form the LoRA adapter. Using multiple heads allows for more fine-grained 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- $z$ ), 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., 2021]. However, this solution does not take into account that test tasks may require solving different combinations of sub-problems compared to training tasks [Vu 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 [Pfeiffer et al., 2021, Vu et al., 2021, Asai et al., 2022, Chronopoulou et al., 2023]. Given a new test task, parameters from similar training tasks are aggregated, which enables transfer. While solely having task-specific parameters is an effective strategy to mitigate interference across training tasks, it also inhibits any positive transfer within the same task pool. Polytropon (Po1y) 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 as T5 [Raffel et al., 2020]. During both multi-task pre-training and few-shot adaptation, Po1y learns both the inventory of adapters and a (continuously relaxed) binary task-module routing matrix, which determines which module is active for each task. While Po1y 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 a series of competitive baselines for few-shot task adaptation on the T0 task suite [Sanh et al., 2022] and Super-Natural Instructions [SuperNI; Wang et al., 2022a]. Based on our results, we 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-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 a 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., 2023] 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- $\mu$ , where the routing function is discarded and all available adapter parameters are averaged before few-shot adaptation. We find that MHR- $\mu$  can recover the performance of MHR, hinting that Poly/MHR gains are only a result of better multi-task optimization. Finally, we show that MHR- $\mu$  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-task training set. This yields gains up to  $3\%$  on absolute accuracy w.r.t. to strong baselines such as T0-11B.
In cross-task generalization, we are given a set of tasks  $\mathcal{T} = \{\mathcal{T}_1,..,\mathcal{T}_{|\mathcal{T}|}\}$ , with each task  $\mathcal{T}_i$  dataset containing a set of samples  $\mathcal{D}_i = \{(x_1,y_1),\dots,(x_n,y_n)\}$ . The set of all tasks is partitioned into training and test tasks,  $\mathcal{T} = \mathcal{T}_{train} \cup \mathcal{T}_{eval}$ , and the objective is to leverage data in  $\mathcal{T}_{train}$  and transfer knowledge to facilitate learning of the test tasks  $\mathcal{T}_{eval}$ . 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 as LoRA or (IA)<sup>3</sup>, or the full backbone is trained on the set of training tasks  $\mathcal{T}_{train}$ . The second phase consists in few-shot adaptation, where the learned adapters are fine-tuned independently on each test task in  $\mathcal{T}_{eval}$ . We follow the procedure from [Raffel et al., 2020] and formulate each task as a text-to-text problem, enabling standard maximum-likelihood training with teacher forcing [Bengio et al., 2015] and a cross-entropy loss.
2.1 Adapters: LoRA & (IA) $^3$
LoRA [Hu et al., 2022] and (IA) $^3$  [Liu et al., 2022] are two recently proposed adapter architectures that achieve competitive trade-offs between performance and parameter efficiency [Karimi Mahabadi et al., 2021, Liu et al., 2022]. For each linear transformation corresponding to the query  $(q)$ , key  $(k)$ , value  $(v)$  and output  $(o)$  of the self-attention layers, LoRA modifies the base model parameters as follows:
$$
\boldsymbol {h} ^ {q, k, v, o} = \boldsymbol {W} _ {0} ^ {q, k, v, o} x + s \cdot \boldsymbol {A} ^ {q, k, v, o} \left(\boldsymbol {B} ^ {q, k, v, o}\right) ^ {\top} x, \tag {LoRA}
$$
where  $W_0$  are the (frozen) weights of the pre-trained model (e.g. T5 [Raffel et al., 2020]).  $A, B \in \mathbb{R}^{d \times r}$  are low-rank learnable parameters and  $s \geq 1$  is a tunable scalar hyperparameter. (IA) $^3$ , on the other hand, modifies key and value representations in self-attention element-wise, and also modifies the feed-forward MLP  $(f)$ :
$$
\boldsymbol {h} ^ {k, v} = \boldsymbol {l} ^ {k, v} \odot \left(\boldsymbol {W} _ {0} ^ {k, v} x\right); \boldsymbol {h} ^ {f} = \left(\boldsymbol {l} ^ {f} \odot \gamma \left(\boldsymbol {W} _ {1} ^ {f} x\right)\right) \boldsymbol {W} _ {2} ^ {f}, \tag {(I A) ^ {3}}
$$
where  $l^{k,v,f} \in \mathbb{R}^d$  are learnable parameters,  $W_{1,2}^f$  the frozen parameters of the feed-forward layer in the backbone, and  $\gamma$  a non-linearity. For clarity, we will drop the superscripts  $q, k, v, o$  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  $\mathcal{M} = \{\phi_1, \dots, \phi_m\}$  with  $|\mathcal{M}| \ll |\mathcal{T}|$ ; 2) a routing function  $r(\cdot)$  that chooses which subset of the modules to combine for each task.
Each module corresponds to a LoRA adapter, where  $\phi_{i}$  are its associated parameters  $A^{(i)}, B^{(i)} \in \mathbb{R}^{d \times r}$ .  $r(\cdot)$  is implemented as a task-module routing matrix  $Z \in \mathbb{R}^{|T| \times |\mathcal{M}|}$ .  $z_{\tau} = Z_{\tau,:} \in \mathbb{R}^{|M|}$  is a routing vector of task  $\mathcal{T}_{\tau}$ , with cell  $Z_{\tau,j}$  being the probability logits of using module  $\phi_{j}$  for task  $\mathcal{T}_{\tau}$  in the current layer. Differently from mixture-of-experts [Fedus et al., 2022], which perform token-level top- $k$  routing,  $Z$  converges to a binary matrix, defining a soft partition over modules.
Method	Pre-Training	Fine-Tuning	Inference
Full FT	d × d	d × d	d × d
LoRA	d × 2r	d × 2r	d × 2r
Poly	d × 2r × |M| + |T| × |M|	d × 2r × |M| + |M|	d × 2r
Poly-z	d × 2r × |M| + |T| × |M|	|M|	|M|
MHR-μ	d × 2r × |M| + |T| × |M|	d × 2r	d × 2r
MHR-z	d × 2r × |M| + |T| × |M| × h	|M| × h	|M| × h
MHR	d × 2r × |M| + |T| × |M| × h	d × 2r × |M| + |M| × h	d × 2r
Table 1: Number of parameters (per layer) used for each method. The calculation uses LoRA as the base adapter, modifying a linear transform in  $\mathbb{R}^{d\times d}$ . 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 a Gumbel-sigmoid distribution [Jang et al., 2017] during training, with  $\hat{Z}_{\tau,j} \sim \mathrm{Gumbel}(Z_{\tau,j})$ . At each forward pass, Poly can be defined as:
$$
\boldsymbol {A} ^ {\tau} = \sum_ {i} \alpha_ {i} \boldsymbol {A} ^ {(i)}; \boldsymbol {B} ^ {\tau} = \sum_ {i} \alpha_ {i} \boldsymbol {B} ^ {(i)} \tag {Poly}
$$
where  $\alpha_{i} = \frac{\hat{Z}_{\tau,i}}{\sum_{j}\hat{Z}_{\tau,j}}$ , and  $\pmb{A}^{(i)},\pmb{B}^{(i)},\pmb{A}^{\tau},\pmb{B}^{\tau}\in \mathbb{R}^{d\times r}$ . We normalize the mixing coefficients  $\hat{Z}_{\tau ,i}$  for each task to ensure that the number of active modules does not affect the norm of  $A^{\tau},B^{\tau}$ . Overall, this approach enables different subsets of modules to be activated for the current layer and combined in a task-specific way. Following LoRA, the output of the Poly layer is added to the output of the original layer of the frozen backbone:  $\pmb {h} = \pmb {W}_0\pmb {x} + s\pmb {A}^\tau (\pmb {B}^\tau)^\top \pmb {x}$ .
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_i, B_i\}_{i=1}^{|\mathcal{M}|}$ , and the routing matrices  $Z$ . During fine-tuning, for each test task  $\tau$ , Poly randomly initialize the routing vector  $z_{\tau} \in \mathbb{R}^{1 \times |\mathcal{M}|}$  and fine-tunes both  $z_{\tau}$  and all the modules parameters  $\mathcal{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 a linear function of the modules. We propose to augment the expressivity of the module combination while keeping the parameter count similar. MHR (Fig. 1) takes inspiration from multi-head attention [Vaswani et al., 2017]: it partitions the input dimensions into  $h$  different disjoint subsets, performs a separate Poly-style combination for each of them, and finally concatenates them. This corresponds to learning a different routing matrix  $Z$  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  $\mathbf{Z} \in \mathbb{R}^{|\mathcal{T}| \times |\mathcal{M}| \times h}$ , where  $\mathbf{Z}_{:,h} \in \mathbb{R}^{|\mathcal{T}| \times |\mathcal{M}|}$  is a 2D slice of the tensor  $\mathbf{Z}$ . A slice represents the routing matrix for each of the  $h$  heads. Let us denote with  $\boldsymbol{W}[k] \in \mathbb{R}^{\frac{d}{h} \times r}$  the  $k$ -th partition along the rows of the matrix  $\boldsymbol{W} \in \mathbb{R}^{d \times r}$ . The adapter parameters  $\boldsymbol{A}^{\tau} \in \mathbb{R}^{d \times r}$  for task  $\tau$ , and for each adapter layer, are computed as (similarly for  $\boldsymbol{B}^{\tau}$ ):
$$
\boldsymbol {A} _ {k} ^ {\tau} = \sum_ {j} \boldsymbol {A} _ {j} [ k ] \cdot \frac {\hat {\boldsymbol {Z}} _ {\tau , j , k}}{\sum_ {j} \hat {\boldsymbol {Z}} _ {\tau , j , k}} \text {w i t h} \boldsymbol {A} _ {k} ^ {\tau} \in \mathbb {R} ^ {\frac {d}{h} \times r}, \tag {MHR}
$$
$$
\boldsymbol {A} ^ {\tau} = \operatorname {c o n c a t} \left(\boldsymbol {A} _ {1} ^ {\tau}, \dots , \boldsymbol {A} _ {h} ^ {\tau}\right)
$$
where concat concatenates along the first dimension. Multi-task pre-training and fine-tuning are similar to Poly. Note that MHR results in only a negligible increase in the total amount of parameters, since most of the parameters are contained in the LoRA weights  $A$ ,  $B$  (Tab. 1).
Routing-Only Fine-Tuning (MHR- $z$ ) Prior work [Shao et al., 2023, inter alia] has shown that compositional generalization can be achieved by learning to (re-)combine 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$  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 a new task.
Adapter Average Fine-Tuning  $(\mathsf{MHR} - \mu)$  To assess the importance of the routing parameters during few-shot adaptation, we propose an additional variant of MHR, MHR-  $\mu$  , which shares the same multi-task pre-training procedure as MHR, but for each test task  $\tau$  , fixes  $z_{\tau} = (1 / |\mathcal{M}|,\dots ,1 / |\mathcal{M}|)$  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 is initialized with (similarly for  $B^{\tau}$ ):
$$
\boldsymbol {A} ^ {\tau} = \frac {1}{| \mathcal {M} |} \sum_ {i} \boldsymbol {A} _ {i} ^ {*}; \boldsymbol {A} ^ {\tau} \in \mathbb {R} ^ {d \times r} \quad (\text {M H R -} \mu)
$$
where  $A_{i}^{*}$  are the parameters of the adapters after MHR multi-task pre-training. Note that, differently from MHR, MHR-  $\mu$  fine-tunes the same amount of parameters as the single adapter baseline. Thus, any difference in performance between the single adapter baseline and MHR-  $\mu$  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 a 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 a single one shared for the entire model. In particular, we consider sharing  $Z$  1) for the adapter layers belonging to the same Transformer block; or 2) for every block of  $l$  layers, which enables us to easily trade off expressivity for parameter efficiency. As we will demonstrate in section 5.1, this is an efficient mechanism to navigate this Pareto front in regimes of very small budgets of parameters per task.
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 baselines and datasets used in our evaluation and then discuss each question in turn.
In addition to Poly, we compare MHR to the following baselines for task-level generalization.
$\mathrm{LoRA} / (\mathrm{IA})^3$  trains a single adapter common to 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 et al., 2023].
AdapterSoup Chronopoulou et al. [2023] trains a different adapter for each task. The method only averages the adapter weights of the training tasks most similar to a 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-training as task adapters are trained independently.
<sup>1</sup>We note that all experiments were run on a single NVIDIA A100 GPU.
T0 Dataset	Avg. Test
Backbone T5-XL
(IA)3	62.40.4
AdapterSoup	62.11.0
LoRA	66.01.6
LoRA-big	65.40.9
Poly-z	66.40.3
Poly	68.01.0
MHR-z	68.30.8
MHR	69.11.0
Backbone T0-3B
T-Few Liu et al. [2022]	66.20.5
AdapterSoup	66.10.6
LoRA	67.40.8
Poly-z	65.31.0
Poly	69.00.8
MHR-z	68.41.2
MHR	69.31.2
Figure 2: Left: Results of few-shot adaptation on T0 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 T0 dataset when applied on top of T5-XL. The x-axis shows the parameter count during the fine-tuning process.
We test our methods on the T0 Sanh et al. [2022] evaluation suite, following the same setup as Liu et al. [2022], and SuperNI Wang et al. [2022a], a large-scale dataset with more than 1,600 training tasks.
T0 Tasks We follow the pre-training and fine-tuning procedure discussed in Liu et al. [2022], using hyper-parameters and losses suggested in the public codebase for T-Few.
All methods were tested with T5-XL Raffel et al. [2020] and T0-3B Sanh et al. [2022] as the backbone model. Crucially, T5 is simply pre-trained on (masked) language modelling, whereas T0 is further instruction tuned: in particular, the full model is 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 of the best validation accuracy for each test task across 3 seeds, when evaluated every 50 training epochs. We treat each data subset-template pair as a unique task, yielding a total of 313 tasks.
SuperNI To limit computational costs, we report the result on 20 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 3 seeds. All methods use T5-XL [Raffel et al., 2020] as the backbone and not T0, as T0 training tasks and SuperNI test tasks may overlap.
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 a backbone that has not undergone prior training on
$^{2}$ https://github.com/r-three/t-few
instruction-following data (T5-XL). As indicated in the T0 benchmark results in the top table of Fig. 2, it is 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 T0-3B as a backbone. In the bottom table of Fig. 2, we can observe that while the relative gap between MHR and baselines is 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.