Merging Multi-Task Models via Weight-Ensembling Mixture of Experts

TL;DR

This work tackles the challenge of merging multiple task-specific Transformer-based models into a single multi-task model without access to all downstream training data. It introduces Weight-Ensembling MoE (WEMoE), which upscales the MLPs in Transformer layers into a dynamic mixture-of-experts module that combines shared knowledge from the pre-trained model with task-specific adjustments via a router and a task-vector dictionary. By separating shared and task-specific information and enabling input-conditioned weight synthesis, WEMoE mitigates parameter interference and achieves superior fusion performance across eight vision tasks, with demonstrated robustness to distribution and input perturbations. The approach offers practical benefits for scalable, flexible multi-task models and opens avenues for extension to other modalities and efficiency-enhancing techniques such as adapters or LoRA.

Abstract

Merging various task-specific Transformer-based models trained on different tasks into a single unified model can execute all the tasks concurrently. Previous methods, exemplified by task arithmetic, have been proven to be both effective and scalable. Existing methods have primarily focused on seeking a static optimal solution within the original model parameter space. A notable challenge is mitigating the interference between parameters of different models, which can substantially deteriorate performance. In this paper, we propose to merge most of the parameters while upscaling the MLP of the Transformer layers to a weight-ensembling mixture of experts (MoE) module, which can dynamically integrate shared and task-specific knowledge based on the input, thereby providing a more flexible solution that can adapt to the specific needs of each instance. Our key insight is that by identifying and separating shared knowledge and task-specific knowledge, and then dynamically integrating them, we can mitigate the parameter interference problem to a great extent. We conduct the conventional multi-task model merging experiments and evaluate the generalization and robustness of our method. The results demonstrate the effectiveness of our method and provide a comprehensive understanding of our method. The code is available at https://github.com/tanganke/weight-ensembling_MoE

Figures (12)

• Figure 1: (a) Framework overview. This figure shows the overall framework of our proposed method to merge the pre-trained model and fine-tuned task-specific models. We merge weights in the Transformer Layers except for the MLPs. For the MLPs, we upcycle them into weight-assembling MoE modules. (b) Wieght-Ensembling Mixture of Experts (MoE) Module. Here we outline the detailed structure of the Weight-Ensembling MoE module, composed of the router, pre-trained MLP weights, and a collection of task vectors. Collaboration between shared weights and task vectors is employed to create input-conditioned weights dynamically. In this way, we separate shared information and task-specific knowledge, which are then combined based on input in time.
• Figure 2: Illustration of the loss landscapes of $s_1$, $s_2$, and $s_1 \cup s_2$. There is no static solution $\theta'$ that simultaneously minimizes the loss of both tasks better than $\mathop{\mathrm{arg\,min}}\limits_\theta \mathcal{L}_1(\theta) + \mathcal{L}_2(\theta)$.
• Figure 3: The distance between the parameters of the pre-trained model and the fine-tuned models (CLIP-ViT-B/32 on eight tasks).
• Figure 4: The performance of the merged models with a varying number of steps. (a) CLIP-ViT-B/32 model with different learning rate. (b) Comparison of CLIP-ViT-B/32 and CLIP-ViT-L/14.
• Figure 5: The results for robustness experiment on CLIP-ViT-B/32. The x-axis of each plot represents the scaling coefficient $\lambda$ of task vectors, while the y-axis shows the accuracy of the merged model on different merged tasks.

Theorems & Definitions (1)

• Definition 2.1: Task Vector ilharcoEditingModelsTask2023