Merging by Matching Models in Task Parameter Subspaces
This work addresses the challenge of model merging for multitask learning, offering a flexible and high-performing method that improves upon existing approaches.
The authors tackled the problem of efficiently combining task-specific models into a single multitask model by introducing a framework called MaTS, which uses the conjugate gradient method to solve linear systems for model merging, achieving state-of-the-art results in multitask and intermediate-task merging.
Model merging aims to cheaply combine individual task-specific models into a single multitask model. In this work, we view past merging methods as leveraging different notions of a ''task parameter subspace'' in which models are matched before being merged. We connect the task parameter subspace of a given model to its loss landscape and formalize how this approach to model merging can be seen as solving a linear system of equations. While past work has generally been limited to linear systems that have a closed-form solution, we consider using the conjugate gradient method to find a solution. We show that using the conjugate gradient method can outperform closed-form solutions, enables merging via linear systems that are otherwise intractable to solve, and flexibly allows choosing from a wide variety of initializations and estimates for the ''task parameter subspace''. We ultimately demonstrate that our merging framework called ''Matching Models in their Task Parameter Subspace'' (MaTS) achieves state-of-the-art results in multitask and intermediate-task model merging. We release all of the code and checkpoints used in our work at https://github.com/r-three/mats.