Merging Smarter, Generalizing Better: Enhancing Model Merging on OOD Data

TL;DR

This paper introduces LwPTV, a layer-wise pruning method for model merging that improves out-of-domain generalization in multi-task learning without sacrificing in-domain performance.

Contribution

The paper proposes a novel saliency-based pruning approach, LwPTV, that enhances model merging for better out-of-domain generalization in multi-task learning.

Findings

Significant improvements in OOD task performance.

Maintains strong in-domain task performance.

Compatible with existing model merging methods.

Abstract

Multi-task learning (MTL) concurrently trains a model on diverse task datasets to exploit common features, thereby improving overall performance across the tasks. Recent studies have dedicated efforts to merging multiple independent model parameters into a unified model for MTL, thus circumventing the need for training data and expanding the scope of applicable scenarios of MTL. However, current approaches to model merging predominantly concentrate on enhancing performance within in-domain (ID) datasets, often overlooking their efficacy on out-of-domain (OOD) datasets. In this work, we proposed LwPTV (Layer-wise Pruning Task Vector) by building a saliency score, measuring the redundancy of parameters in task vectors. Designed in this way ours can achieve mask vector for each task and thus perform layer-wise pruning on the task vectors, only keeping the pre-trained model parameters at the corresponding layer in merged model. Owing to its flexibility, our method can be seamlessly integrated with most of existing model merging methods to improve their performance on OOD tasks. Extensive experiments demonstrate that the application of our method results in substantial enhancements in OOD performance while preserving the ability on ID tasks.