Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Merging by Matching Models in Task Parameter Subspaces

Derek Tam, Mohit Bansal, Colin Raffel

TL;DR

MaTS reframes model merging as matching individual task models in a dedicated task parameter subspace and introduces a conjugate gradient-based solver to handle the resulting linear systems. By unifying simple averaging, Fisher-based, and RegMean approaches under a common covariance- or Gram-derived subspace, it enables flexible initializations and objectives, including a new block-diagonal Fisher variant via K-FAC. Across language and vision benchmarks, MaTS achieves state-of-the-art multitask and intermediate-task merging results while remaining far cheaper than full multitask training. The work also highlights the importance of initialization strategies and points to future enhancements in subspace estimation and merging efficiency.

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 at https://github.com/r-three/mats.

Merging by Matching Models in Task Parameter Subspaces

TL;DR

MaTS reframes model merging as matching individual task models in a dedicated task parameter subspace and introduces a conjugate gradient-based solver to handle the resulting linear systems. By unifying simple averaging, Fisher-based, and RegMean approaches under a common covariance- or Gram-derived subspace, it enables flexible initializations and objectives, including a new block-diagonal Fisher variant via K-FAC. Across language and vision benchmarks, MaTS achieves state-of-the-art multitask and intermediate-task merging results while remaining far cheaper than full multitask training. The work also highlights the importance of initialization strategies and points to future enhancements in subspace estimation and merging efficiency.

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 at https://github.com/r-three/mats.
Paper Structure (51 sections, 22 equations, 3 figures, 9 tables, 1 algorithm)

This paper contains 51 sections, 22 equations, 3 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Simple Averaging
  4. Fisher Merging
  5. RegMean
  6. Matching Individual Models in their Task Parameter Subspace
  7. The Task Parameter Subspace of Prior Merging Methods
  8. Simple Averaging.
  9. Fisher Merging.
  10. Diagonal Fisher Merging.
  11. RegMean.
  12. Matching in Task Parameter Subspace
  13. Merging via the Conjugate Gradient Method
  14. Block-Diagonal Fisher Merging
  15. Connecting Fisher Merging and RegMean
  16. ...and 36 more sections

Figures (3)

  • Figure 1: Diagram visualizing what several merging methods are doing in the loss landscape for a particular layer of the model. The ellipses represent the level sets of the losses for different task. Different merging methods modify the individual models in various ways before summing the modified models (with an appropriate normalization constant). Averaging does not modify the model in any way. Diagonal Fisher merging modifies each model by amplifying each parameter along an axis-aligned projection. Block-diagonal Fisher merging modifies each model by amplifying each parameter along a projection along the steepest directions of the loss landscape.
  • Figure 2: We show the performance using CG to optimize the objectives implicitly defined by various merging methods with various initializations. The number indicates the raw performance and the subscript (and color) indicate the change in performance relative to the initialization.
  • Figure 3: Performance of the merged model and the in error solving the linear system vs. the number of iterations in the conjugate gradient method.