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

TL;DR

This paper introduces a novel method for merging multi-task Transformer models using a weight-ensembling mixture of experts, which dynamically integrates shared and task-specific knowledge to improve performance and mitigate parameter interference.

Contribution

The paper proposes a dynamic MoE-based merging approach that separates shared and task-specific knowledge, enhancing multi-task model integration beyond static methods.

Findings

Effective multi-task model merging demonstrated

Improved generalization and robustness shown

Dynamic integration outperforms static methods

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