Merging by Matching Models in Task Parameter Subspaces

Derek Tam, Mohit Bansal, Colin Raffel

Abstract: 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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: We have updated the camera ready and replaced all instances of task subspace with task parameter subspace and clarified the relationship between MaTS and the the Fisher in section 7.2.

Code: https://github.com/r-three/mats

Assigned Action Editor: ~Rémi_Flamary1

Submission Number: 1908