Merging Smarter, Generalizing Better: Enhancing Model Merging on OOD Data
This work addresses a bottleneck in multi-task learning for AI practitioners by enhancing model merging on OOD data, though it is incremental as it builds upon existing methods.
The paper tackles the problem of model merging methods performing poorly on out-of-domain (OOD) data by proposing LwPTV, a layer-wise pruning technique that improves OOD performance while maintaining in-domain capabilities, as shown through extensive experiments.
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.