Data-Adaptive Weight-Ensembling for Multi-task Model Fusion

Abstract

Creating a multi-task model by merging models for distinct tasks has proven to be an economical and scalable approach. Recent research, like task arithmetic, demonstrates that a static solution for multi-task model fusion can be located within the vector space spanned by task vectors. However, the static nature of these methods limits their ability to adapt to the intricacies of individual instances, thereby hindering their performance in complex scenarios. To overcome this limitation, we propose a data-adaptive weight-ensembling approach that generates model weights in time. Specifically, we first feed the input samples into a hypernetwork to generate instance-specific weights for the primary model. Subsequently, we perform a functional call on the primary large model with the instance-specific weights. By generating model weights in time, the unified model gains increased flexibility and can resolve potential weight conflicts between tasks. Building upon this adaptability, our method necessitates solely the model checkpoints and unlabeled test samples using test-time adaptation training. We primarily conduct extensive experiments on vision Transformers and Flan-T5 models, demonstrating superior performance and satisfactory zero-shot transferability.

Appendices

Model Fine-Tuning Details

The experiments in our study were conducted on a consistent hardware setup, utilizing eight NVIDIA GTX 3090 GPUs, each equipped with 24GB of video memory. To implement our experiments, we employed PyTorch version 2.1 with Python 3.10. The source code, model checkpoints, and logs from our experiments are available at https://github.com/tanganke/dict_fusion.

1.1 Experiment Setup

Our experimental scope encompasses a diverse range of tasks spanning both the vision and natural language domains. The model selection, evaluation datasets, and metric configurations are detailed as follows:

Vision tasks: For vision tasks, we leverage the pre-trained models from CLIP (Radford et al. 2021) .¹ The downstream tasks encompass a variety of challenges, including SUN397 (Xiao et al. 2010), Stanford Cars (Krause et al. 2013), RESISC45 (Cheng et al. 2017), EuroSAT (Helber et al. 2018), SVHN (Netzer et al. 2021), GTSRB (Stallkamp et al. 2012), MNIST (Lecun et al. 1998), and DTD (Cimpoi et al. 2014). We report top-1 accuracy as the performance metric.

NLP tasks: In the realm of natural language processing (NLP) tasks, our pre-trained language model is Flan-T5 (Chung et al. 2022) .² For fine-tuning, we deploy the flan-t5 models on eight tasks derived from the GLUE benchmark (Wang et al. 2018). The suite of natural language processing (NLP) tasks are CoLA, MNLI, MRPC, QNLI, QQP, RTE, SST2, and STSB. For consistency and reproducibility, we initialize the parameter-efficient models using a identical random seed 42. The Flan-T5 models undergo LoRA fine-tuning, utilizing hyperparameters $r=16$ and $\alpha =32$ as specified in the work by Hu et al. (Hu et al. 2021). In this process, we set a constant learning rate of $4e-5$ and maintain a uniform batch size of 16 throughout all tasks. We fine-tune the Flan-T5 models for 2000 steps on each downstream task.

Flan-T5 models are encoder-decoder architecture Transformer models, and function within a text-to-text framework. Consequently, we have restructured the initial inputs to adhere to this text-to-text paradigm. We evaluate exact match accuracy for all tasks except STSB, for which we report Spearman’s $\rho$. The prompt templates for each task are available in our code repository.

Figure 5 presents detailed performance metrics for the CLIP-ViT-B/32 and CLIP-ViT-L/14 models, respectively. These visual and tabular data representations dissect the performance impact that both pre-training and subsequent fine-tuning have on the models’ task execution. The data spans a variety of downstream tasks, offering a comprehensive view of the models’ capabilities. Similarly, for the realm of natural language processing models, Fig. 5 provides an insight into the individual performance transformations that occur in the Flan-T5-base and Flan-T5-large models post LoRA fine-tuning. As evidenced by the highlighted diagonal cells, it is apparent that the fine-tuning process has a significant impact on the models’ abilities to perform specific tasks on both vision and language tasks.

Fig. 5

Individual performance of fine-tuned models. Here, we provide a visual comparison of the performance of various fine-tuned models. The upper figure contrasts the performance of CLIP-ViT-B/32 and CLIP-ViT-L/14 models across image classification tasks. Meanwhile, the lower figure compares the performance of LoRA fine-tuned Flan-t5-base and Flan-t5-large models on tasks from the GLUE benchmark

Multi-task Model Fusion Methods

This section provides an breif introduction to the baseline methods been compared in our experiments. Including weight averaging, Task Arithmetic, Ties-Merging, and AdaMerging. Table 1 summarizes the requirements of different model merging methods.

Simple weight averaging: Weight Averaging is the most straightforward method wherein the weights of models fine-tuned on different tasks are averaged, commonly known in the literature as ModelSoups (Wortsman et al. 2022). Subsequently, the averaged model undergoes evaluation on the validation set of each task. In the case of full fine-tuned models, the weights are directly averaged, denoted as

$$\begin{aligned} \theta = {\frac{1}{n} \sum _{i=1}^{n} \theta _i}. \end{aligned}$$

(B1)

When the models are LoRA fine-tuned, we average the weights of merged models, expressed as

$$\begin{aligned} \theta = \theta _0 + \frac{1}{n}\sum _i A_i B_i. \end{aligned}$$

(B2)

Task Arithmetic (Ilharco et al. 2023): The Task Arithmetic involves the computation of a task vector for each individual downstream task, followed by the summation of these task vectors to construct a multi-task vector. This multi-task vector is then element-wise multiplied by a scaling coefficient denoted as $\lambda$ and added to the initial parameters of the pre-trained model. In mathematical terms, for full fine-tuning, this process is represented as

$$\begin{aligned} \theta = \theta _0 + \lambda \sum _{i=1}^n (\theta _i - \theta _0). \end{aligned}$$

(B3)

And for LoRA fine-tuning, it is expressed as

$$\begin{aligned} \theta = \theta _0 + \lambda \sum _{i=1}^n A_i B_i, \end{aligned}$$

(B4)

where $\lambda$ serves as a hyperparameter determined by selecting the best-performing model on the validation set.

In the context of our study, we set $\lambda$ to a value of 0.3, a choice made after evaluating the model’s performance on the validation set.

Ties-Merging (Yadav et al. 2023): Similar to Task Arithmetic, Ties-Merging also involves element-wise multiplication of a unified task vector by a scaling coefficient denoted as $\lambda$ and is then added to the initial parameters of the pre-trained model.

However, before doing this, Ties-Merging involves three key steps, namely trimming, determining the sign of parameters, and disjoint merging, ultimately resulting in a merged task vector denoted as $\nu$. In other words, the final model parameters are given by $\theta = \theta _0 + \lambda \tau$, where $\lambda$ represents a hyperparameter chosen based on the best-performing model identified through validation set evaluation.

For our specific study, we set the hyperparameter $\lambda$ to a value of 0.3, a decision made after thorough assessment of model performance on the validation set.

AdaMerging (Yang et al. 2023): AdaMerging represents an adaptive model merging approach that autonomously learns coefficients for merging, operating either on a task-wise or layer-wise basis. It achieves this by employing entropy minimization on unlabeled test samples as a surrogate objective function to refine the merging coefficients. The task-wise AdaMerging is expressed as

$$\begin{aligned} \theta = \theta _0 + \sum _{i=1}^{n} \lambda _i \tau _i, \end{aligned}$$

(B5)

where $\lambda _k$ denotes the merging coefficient for the k-th task, and $\tau _k$ represents the task vector for the same task. On the other hand, the layer-wise AdaMerging formulation is given by

$$\begin{aligned} \theta ^{(l)} = \theta _0^{(l)} + \sum _{i=1}^{n} \lambda ^{(l)}_{i} \tau ^{(l)}_{i}. \end{aligned}$$

(B6)

Where l is the layer-index of models.

In our study, we initialize the merging coefficients $\lambda$ to be 0.3 for all tasks and layers as in (Yang et al. 2023). This ensures a standardized starting point for the adaptive merging process.

Ours. We set the learning rate to 1e-5 and employ the Adam optimizer for test-time adaptation on the test dataset, following the details outlined in Sect. 3.5 of the main paper. The model undergoes training for 1000 steps on each task, and we present the performance results across all tasks in Tables 3 and 4.