SurgeryV2: Bridging the Gap Between Model Merging and Multi-Task Learning with Deep Representation Surgery

TL;DR

This paper proposes a more comprehensive solution, deep representation surgery (also called SurgeryV2), which mitigates representation bias across all layers, and thus bridges the performance gap between model merging-based MTL and traditional MTL.

Abstract

Model merging-based multitask learning (MTL) offers a promising approach for performing MTL by merging multiple expert models without requiring access to raw training data. However, in this paper, we examine the merged model's representation distribution and uncover a critical issue of "representation bias". This bias arises from a significant distribution gap between the representations of the merged and expert models, leading to the suboptimal performance of the merged MTL model. To address this challenge, we first propose a representation surgery solution called Surgery. Surgery is a lightweight, task-specific module that aligns the final layer representations of the merged model with those of the expert models, effectively alleviating bias and improving the merged model's performance. Despite these improvements, a performance gap remains compared to the traditional MTL method. Further analysis reveals that representation bias phenomena exist at each layer of the merged model, and aligning representations only in the last layer is insufficient for fully reducing systemic bias because biases introduced at each layer can accumulate and interact in complex ways. To tackle this, we then propose a more comprehensive solution, deep representation surgery (also called SurgeryV2), which mitigates representation bias across all layers, and thus bridges the performance gap between model merging-based MTL and traditional MTL. Finally, we design an unsupervised optimization objective to optimize both the Surgery and SurgeryV2 modules. Our experimental results show that incorporating these modules into state-of-the-art (SOTA) model merging schemes leads to significant performance gains. Notably, our SurgeryV2 scheme reaches almost the same level as individual expert models or the traditional MTL model. The code is available at \url{https://github.com/EnnengYang/SurgeryV2}.

Figures (18)

• Figure 1: The SurgeryV2 scheme achieves performance comparable to traditional MTL when merging eight benchmark tasks and two architectures (left: ViT-B/32, right: ViT-L/14). We report the performance normalized with the individual expert model’s performance, which is denoted by the white circle of radius $1$.
• Figure 2: Visualization of the distribution of representations extracted by the merged model (red) for the existing model merging schemes and representations extracted by the individual expert model (blue). We observe that there is a clear distribution discrepancy between the two.
• Figure 3: The "representation bias" (in Eq. \ref{['eq:l1distance']} of Definition \ref{['def:representation_bias']}) between representations extracted by the merged model using three model merging methods and representations extracted by individual expert models.
• Figure 4: Visualization of the "representation bias" (in Eq. \ref{['eq:l1distance']}) of the representation of the merged model with (red) and without (blue) representation surgery versus the individual model.
• Figure 5: Visualization of the distribution of features extracted by the merged model after performing the representation surgery (red) and features extracted by the individual model (blue). We observe that the two distributions overlap highly.

Theorems & Definitions (3)

• Definition 3.1: Representation
• Definition 3.2: Model Merging
• Definition 4.1: Representation Bias