Merging by Matching Models in Task Parameter Subspaces

Published: 29 Mar 2024, Last Modified: 29 Mar 2024Accepted by TMLREveryoneRevisionsBibTeX
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.
Supplementary Material: zip
Assigned Action Editor: ~Rémi_Flamary1
Submission Number: 1908