Merging by Matching Models in Task Parameter Subspaces
Authors: Derek Tam
Mohit Bansal
Colin Raffel
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (03/2024)
Merging by Matching Models in Task Parameter Subspaces
Derck Tam
diam@cs. edu
University of Toronto
Vector Institute
Mohit. Bansal
mbansal@css unc.
University of North Carolina
Chapel Hill
Colin Raffel
craffel@gmail.com com
University of Toronto Vector Institute
Reviewed on OpenReview: https: openreveew. net/ feet/foree forum? GNGo6ghhWFF
Abstract
Model merging aims cheaply combine individual task-specific models into a single multi- task model. In this work, view past merging methods as leveraging different notions of "task parameter subspace" in which models are matched before being merged. We connect the task parameter subspace of given model to its loss landscape and formalize how this approach to model merging can be seen as solving lincar 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-fore solutions, enables merging via lin- car 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 merging framework called "Matching Models in their Task Parameter Subspacc" (MaTS) achieves satte-of-thc-art results in multitask and intermediato-tan model merging. We release all of the code and checkpoints used in our work.'
1 Introduction
The widespread fine-tuning of public pre-trained models has produced a huge number of specialized models. These specialized models may be trained on different tasks, where a "task" simply the input-output relationship that we aim to train a model to perform (e.g. sentiment analysis of text, object recognition in images, ctc.). Alternatively, the Stable Diffusion XL model (Podell al., 2023) forms the basis of over thousand specialized image generation models on the Hugging Face Model Hub that are specialized different styles content types. How we recycle these specialized models to create better base models (Choshen et al., 2022; Ramé et al., 2022)? Model merging (Wortsman et al., 2022b; Matena & Raffel, 2022) aims to tackle this problem by combining specialized models into single model that retains the individual models' capabilities. A common example application of merging is constructing a multitask model from individual-task models, which is the primary application we explore in our paper. Compared to multitask learning, merging does not require simultaneont access to the individual-Lask datasets. Compared to output- space ensembling of M models, merging produces model that is M times cheaper run.
While merging via simple parameter averaging work well for models that share architecture and initialization (McMahan el al., 2017; Stich, 2018), recent merging methods improve over simple averaging by
httpp::/gitthub.. wwtps/tthhu..oorreeeeeeeeeteeeee three/nats.

---

Published in Transactions on Machine Learning Research (03/2024)
considering parameter importance (Matena & Raffel, 2022), matching activations (Jin et al., 2022), omitting the contribution of the pre-trained model (Mharco et al., 2022), or resolving interference across models (Yadav et al., 2023).
In our work, we show how several recent merging methods can be viewed as finding a single model that matches task-specific models in their respective "task parameter subspaces". We define a task parameter subspace as the subspace implicitly used by a given merging method that aims to correspond the important dimensions in parameter space for the task. To match models in their task parameter subspace, merging methods upweight each model in its task parameter subspace, which aims to ensure that the task-relevant components of a given model will not be washed out after the models are combined. In particular, we show how Fisher merging (Matena & Raffel, 2022), RegMean (Jin et al., 2022), and simple parameter averaging (McMahan et al., 2017; Stich, 2018) all perform merging in this way and differ only in their choice of task parameter subspace. Concurrently, other works have focused on inaccuracies model merging stemming from gradient mismatches in different models, and use these insights to connect diagonal Fisher merging and Task Arithmetic Daheim et al. (2023).
Matching models in their task parameter subspace requires solving linear system of equations. This linear system implicitly defines a merging objective that relates to a given merging method's choice of task parameter subspace. While previous merging methods used merging objectives with a tractable closed-form solution, we instead develop a merging framework that uses the conjugate gradient method (Hestenes & Stiefel, 1952) to solve a given linear system. We refer to our merging framework as MaTS (Matching Models in their Task Subspace). By using the conjugate gradient method, MaTS flexibly supports different merging objectives and initializations (which can impact convergence speed). MaTS also enables the use of merging objectives for linear systems that don't have tractable closed-form solution. To explore this possibility, we leverage insights from K-FAC (Grosse & Martens, 2016; Martens & Grosse, 2015) and introduce a merging method where a model's task parameter subspace is based on a block-diagonal approximation of the Fisher information matrix.
To explore the effectiveness of MaTS, we comprehensively compare it to existing merging methods on multitask and intermediate-task merging of language models and vision models trained via parameter-effcieet or full- model fine-tuning. We first explore using precxisting merging methods demonstrate that MaTS can significantly boost performance given a suitable choice of merging objective. In particular, in multitask language model merging, we find that MaTS attains state-of-the-as results by a large margin. We use insights from this exploration to develop an effective merging recipe (i.e. a consistent initialization and objective to use) for srrameter-efficien and full-model fine-tuning, which we then apply to multitask vision model merging and intermediate-task language model merging. In both cases, we validate that MaTS can boost performance over its initialization and often attains state-of-the-art results. Finally, we discuss how MaTS has higher computational cost than existing merging methods but is nevertheless dramatically cheaper than explicit multitask training. Taken as a whole, our results validate both our perspective of model merging as matching models in their task parameter subspacc as well as the effectiveness of using the conjugate gradient method for solving the corresponding linear system.
an initialization for MaTS and
2 Background
In this paper, we assume we are given M models to merge, all sharing a common architecture and initialization (e.g. a pre-trained model) and denote the parameters of model m as 0m, the input as x, and the output as y over which the model aims to parameterize a distribution pe... (2)zz) We further assume we have access to the dataset Dm (containing Nm examples) and the loss function L,m. used to train model m. All of the models we study use typical neural network architecture that is made up of a series of linear operations ("layers") with parameter-free nonlincarities in between. Consider lincar layer I in some particular model m with weight matrix Wm, E dd Let z E Rd be the layer's input activation, o E R& be the outpul activation, and o' e #* be the gradient of the loss function with respect to the linear layer's output activation for one particular input. Throughout this work, we assume that models use a cross-cntropy loss function, so that s' can be considered as the gradient of the log probability of the correct class with respect to the linear layer's output. For brevity, we refer to this quantity simply as the "output activation gradient". Let Zm E Nm X
2

---

Published in Transactions on Machine Learning Research (03/2024)
be the input activations stacked row-wise, On IR N be the output activations stacked row-wise, and O'n E IR-Nmx xk be the output activation gradients stacked row-wise. The operation computed by the linear layer can therefore be expressed as On = Zm Wm. We also allow for a slight abuse of notation and use Om to refer to either the dk-dimensional vector of the parameters of a particular linear layer (e.g. by flattening a corresponding matrix) or the p-dimensional vector all parameters in the model. When clear from context, in inline equations we abbreviate summation limits (e.g. 2m in place of >dd >m=1).
2.1 Simple Averaging
A common way to merge models is to simply average their parameters, i.e. compute 0* = M Σm On Such parameter averaging has been widely used in federated learning (McMahan et al., 2017) and distributed optimization (Stich, 2018) but has more recently been used to merge models to retain out-of-distribution performance (Wortsman et al., 2022b), improve pre-trained model (Choshen al., 2022; Don Don-Yehiya al., 2022), or ereate multimodal models (Sung et al., 2023).
2.2 Fisher Merging
Matena & Raffel (2022) view the problem of merging models as finding the set of paramcters with the highest joint probability according to the individual model's posteriors:
M
0 arg max
Σ
log Pm))
(1)
Since maximum likelihood-based training of neural networks does not provide access to a posterior distribu- tion over parameters, the posterior must be approximated. Matena & Raffel (2022) point out that simple averaging solves eq. (1) if parameters are assumed to be sampled from isotropic Gaussian posteriors. How- ever, an isotropic Gaussian is likely a poor approximation of the true posterior, so Matena & Raffel (2022) use the Laplace approximation, which assumes that parameters are sampled from Gaussian distribution with mean Oz and covariance set to the inverse of the Fisher information matrix ("Fisher") Fm. The Fisher can be generically defined for some model po(y)z) as
Fm
Vo logro (y(lx) [|
(2)
D>DD Dm (
Since computing, storing, and inverting the Fisher is intractable for modestly large neural networks, Malena & Raffel (2022) use a diagonal approximation of the Fisher F (i.e. assuming parameter values are indepen- dent), which results in a matrix whose diagonal entries correspond to the the average of the per-example gradients squared. Under this approximation, the parameter values that maximize eq. (1) have the following closed-form solution:
(1
M
M
1
1
F
Σ mmmm
33)
In practice, it is common use the empirical Fisher where the expectation over labels is replaced by the ground-truth label for a given datapoint. The empirical Fisher has been shown to have certain limitations compared to the true Fisher when used in the Natural Gradient algorithm (Kunstner et al., 2019). However, we found that for model merging the empirical Fisher worked better (see section 6.6 for full results) and is computationally cheaper and therefore we use it throughout this work.
2.3 RegMean
RegMean (Jin et al., 2022) aims to find a merged model that minimizes the distance between the output activations of the original model and the output activations of the merged model. To make this goal tractable, RegMean merges the parameters of each linear layer separately. This per-layer objective can then be formulated as the least squares regression problem
M
min W
0m Zm 11
(4)
m
3

---

Published in Transactions on Machine Learning Research (03/2024)
Jin et al. (2022) solve this least squares regression problem using its closed-form solution:
M
W
Z Zm
Zm) Wm
(5)
N
Using RegMean for merging therefore requires computing and inverting the gram matrices Z ZTT Zm- To ensure invertibility, RegMean scales the nondiagonal terms of the Gram matrix with a constant À (with 0 < X < whose value can be tuned based on performance on held-out data.
3 Matching Individual Models in their Task Parameter Subspace
While simple averaging, diagonal Fisher merging, and RegMean appear to be unrelated in their approach, we now demonstrate that all three methods can be seen as computing a merged model that has the generic form
0"
Σ
C.
Σ
C 8.
(6)
m=1
m==
Qm Am QTh
Σ
Qm
(7)
where Cm is a (approximate) covariance matrix of some random variable (i.e. some data). As we will discuss later, prior merging methods primarily differ in terms of the choice of the random variable. Since the covariance matrix is positive semi-defiitt, we can obtain its eigendecomposition Cn Am QT If Cm is the covariance of some data P.. with SVD decomposition P.. Um VIT then Qm Vn as shown in appendix A.3. The top singular vectors in Vn and, equivalently, Qm are the directions that best capture the variance of Pm. In other words, the top column vectors in Vn and Qm can therefore viewed as forming a basis that "best aligned" with the data Pm- We therefore refer to Qm. as the basis vectors of the "task parameter subspace" whose relative importance are determined by Am
3.1 The Task Parameter Subspace of Prior Merging Methods
Having established the idea of a task parameter subspace, we now apply our perspective to past merging methods and uncover each method's choice of task parameter subspace.
Simple Averaging. Simple averaging sets Cm I (the identity matrix) thereby setting Qm - and Am = I. Setting Qm = I means that merging is performed in the original model's parameter space and setting Am = implies that that each direction the task parameter subspace is equally important.
Fisher Merging. established in appendix A.1, the closed form solution for Fisher merging (i.e. without necessarily making any diagonal approximation to the Fishor) is given by
A5
M
0+ - (- Fm
Σ
Fm
(8)
In this case, C. = Fm is the covariance of the per-example gradicnts of the datapoints from a task. The true (i.e. not the empirical) Fisher is equal to the Hessian (in expectation) when the negative log-likelihood loss is used (Grosse, 2022a;b; Martens, 2020). In this paper, we only consider the use of the the cross entropy (negative logliikel)hood) loss function. While Kunstner et al. (2019) note that the empirical Fisher should go
While Jin al. (2022) da not include the 1 their definitions, the constant 1 tation and experiments (sec https://github. N. httpa///aattttlseeooemmeeerr//// N., used practice in their implemen- constant ensures that each dataset equally weighed, even the datasets have different number of examples. Appendix A.2 shows the proof that this modified closed-form solution solves the modified objective.
demieeanddmoo..pmb).. The
4

---

Published in Transactions on Machine Learning Research (03/2024)
Figure 1: Diagram visualizing what several merg- ing methods are doing in the loss landscape for a particular layer of the model. The ellipses repre- sent the level sets of the losses for different task. Different merging methods modify the individual models in various ways before summing the mod- ified 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 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.
Model
Model
Merged model
Simple averaging
Diagonal Fisher merging.
-0--
Block Block-diagonal diagonal Figher merging
to the model converges whereas the true Fisher doesn'l, 1, we follow the insight from Singh & Alistarh (2020) that gradients do not actually vanish for well-optimized neural networks. Thus, view the (empirical or truc) Fisher as approximating the Hessian and therefore Qm can be viewed as the cigenvectors of the Hessian and Am are the corresponding eigenvalues. The Hessian captures the curvature of the model's log- likelihood (Kristiadi, 2018; Martens, 2020), so its top eigenvectors capture the directions of sharpest change in the loss landscape and the bottom eigenvectors capture directions where the loss function roughly flat (Maddox et al., 2020). As we will discuss in section 3.2, Fisher merging therefore encourages merging shift models along the given model's parameters "flat" directions.
Diagonal Fisher Merging. Diagonal Fisher merging (Matena & Raffel, 2022) approximates C. diagonal Fisher matrix. This corresponds to assuming that Qm is the identity matrix I, which is equivalent to computing the importance score Am in axis-aligned parameter space.
RegMean. Inspecting eq. (5) suggests that RegMean sets C. as a function of the scaled Gram matrix I Nm L Zmu i.e. the covariance the input activations. Since eq. (5) is defined terms weight matrix and eq. (6) is defined terms of a parameter vector, RegMean sets C. as block-diagonal matrix whose blocks repeated instances of the scaled Gram matrix Nn Z Zm- RegMean therefore computes the basis Qm and important score Am based on the directions that best capture the variance in the input activations.
3.2 Matching in Task Parameter Subspace
Recall the closed-form solution as written in eq. (6):
M
-1
M Qm 0m
0+
Σ
Q m Am Qt
Σ m=1
mm
(9)
The terms of the second summation can be seen as projecting the parameters O. using Am QT What does the projection with Qm Q imply?
QT pn rotates 0m into the task parameter subspace whose basis vectors are the eigenvectors con- tained Qm-
2. Am Am om amplifies the the dimensions in this rotated coordinate system based on their importance as defined by the eigenvalues.
3. Qm QT, Rm rotates back into the original parameter space
The first summation in eq. (9) can then be seen as a normalizing constant to account for all modifications done in the second term. Overall, eq. (9), when QmAm is the decomposition of a covariance matrix, be seen upweighting the "important" components of the model measured in the task parameter
5

---

Published in Transactions on Machine Learning Research (03/2024)
subspace that the important components don't get washed out during merging. Thus, the least important component of the models will be shifted more during merging. The components are not constrained to be individual parameters but can be any direction in parameter space. Figure 1 combines this insight with the discussion from section 3.1 and provides a visualization the different task parameter subspaces implicit in prior merging methods.
Additionally, note that we can re-write the closed-form solution eq. (6) as finding 0 such that
M Σ ma
M Σ mmm
(10)
One way to satisfy eq. (10) would be if QmAmQ matched Qm pm across all tasks m. In other words, the extent to which solution 0* satisfics eq. (10) depends how close Am 0* is to each Qm Am Q On. Note that this does not measure the closeness of 0" and On directly, but rather through the projection Qm Qm, which (as discussed above) effectively measures closeness in the task parameter subspace.
Q
4 Merging via the Conjugate Gradient Method
Merging models by solving for 0* in eq. (6) is equivalent to solving a linear system of the form Ax b, where A = Em Cm and b = Σm. Om Different ways to compute the task parameter subspace define different merging objectives for this linear system. All past merging methods discussed in section 2 provide a tractable closed-form solution to this linear system. However, there may be performant merging methods that do not have such a fosed-form solution. In addition, since the solution in eq. (6) involves a matrix inversion, some past methods deviate from the true solution to ensure invertibility. In this work, we therefore explore finding 8 via optimization rather than a closed-form solution. Specifically, to solve the linear system, we use the conjugate gradient method (Hestenes & Stiefel, 1952) (CG), a first-order iterative optimization procedure for solving linear systems when A is positive semi-definite (Hayami, 2018). We include more details of CG in appendix B. Note that for eq. (6) A is positive semi-definite since it is the sum of the covariance matrices computed on some data.
A notable benefit of the conjugate gradient method that it allows for solving a linear system without having to compute matrix inverse or compute A explicitly. Specifically, the step size and the update direction can be computed as long as Ax and b can be computed. This is particularly useful if A is too large and/or cannot be explicitly computed. In addition, avoiding the matrix inverse make CG more stable than a closed-form solution that relies on an inverse. For example, RegMean Jin et al. (2022) has to scale the non-diagonal entries of the gram matrix by constant to ensure the inverse is not ill-conditioned: we show later in section 6 that CG can find a performant solution to the RegMean objective without applying this scaling. On the whole, we refer to our framework for matching models in their task subspace using conjugate gradient as MaTS.
5 Block-Diagonal Fisher Merging
Since CG allows merging via linear systems that do not have a tractable closed-form solution, we now intro- duce a new Fisher merging method that makes use of an improved approximation to the Fisher. Specifically, we follow K-FAC (Grosse & Martens, 2016; Martens & Grosse, 2015) and approximate the Fisher as a block-diagonal matrix (BFM). The closed-form solution of the resulting objective is intractable but can be efficiently optimized with CG.
5.1 Connecting Fisher Merging and RegMean
To motivate our use of a block-diagonal approximation the Fisher, we first note that RegMean can be seen a form of Fisher merging under certain assumptions. Specifically, RegMean assumes the output neurons are independent and that the gradient of the output activation is constant. Under this assumption,
6

---

Published in Transactions on Machine Learning Research (03/2024)
for particular linear layer W Rd*k the Fisher matrix F' E Rdii is block diagonal matrix with k blocks, each of size d x d. The gradient of a linear layer is the outer product of the imput activation and the output activation gradient: V VW log Poglz())) = =o" Assuming the output activation gradient of is 1, for some particular dataset with N examples the ith block of given per-layer Fisher F' can be shown (in appendix A.4) to be given by zz Replacing F.. in the Fisher merging objective from eq. (6) with the RegMecn--peccfic Fisher derived gives the RegMean closed-form solution from eq. (5). Following Grosse et al. (2023), this also assumes the gradients of the Lokens within an example are independent.
5.2 Using K-FAC for Fisher Merging
Can we approximate the Fisher as a diagonl matrix without assuming that the output activation gradient is constant? In general, the closed-form solution for Fisher merging with a block-diagonal Fisher approximation would not be tractable to compute for typical neural networks because, for an activation dimensionality of 1024, the Fisher for single linear layer (i.e. a block in the block-diagonal Fisher ap- proximation) would have trillion entries. One possibility would be to use the K-FAC approximation of the Fisher Martens & Grosse (2015); Grosse & Martens (2016), which approximates the Fisher of lincar layer as the Kronecker product of the covariance of the imput activations and the covariance of the output activation gradients:
( TTT Zm Nm
1 N
Fm =
2
O
O',
(11)
where @ is Kronecker product.
However, even if we use the K-FAC Martens & Grosse (2015); Grosse & Martens (2016) approximation for the Fisher, a closed-form solution to eq. (6) would still not be tractable. Plugging in the K-FAC approximation into eq. (6) in appendix A.5, we see that
M Σ
O
1 Nm
M Σ
W. (N (N
w'
Z Ze
Nm Zi Zz
(12)
N
Though the second term can be easily computed, the first term cannot be since the sum of Kronecker products is not equal to the Kronecker product of sums. While we could assume the of Kronecker products is the Kronecker product of sums (an assumption made in (Grosse & Martens, 2016; Martens & Grosse, 2015) to use the K-FAC approximation for optimization), we found in preliminary experiments that this would cause the merged model's performance to crash.
The conjugate gradient method, however, can be used to solve the Fisher merging objective with a block- diagonal Fisher approximation. Specifically, using the K-FAC approximation for the Fisher, we can efficiently compute the matrix vector product Aæ (as shown in appendix A.6), even if A cannot be computed with:
M
Ax
Z,
Z:
W
o o,ToO:
(13)
Nm
N
M
b =
N Z Zi
W
o O,To') 0.)
(14)
Nm
m=1
where x = vec (W) Each term in the summands above involve multiplying matrices whose dimensions correspond to the activation dimensions of the model, thereby incurring a modest computational cost. As a result, the use of CG not only avoids an intractable inverse but also produces a computationally feasible method for Fisher merging with block-diagonal Fisher approximation.
6 Experiments
Now that we have introduced MaTS, we empirically explore various design choices and validate the effectiveness of our approach. Specifically, first consider different initialization and objectives in order identify

---

Published in Transactions on Machine Learning Research (03/2024)
effective fixed recipe to apply in various settings. We mainly focus on merging task-specific models into single multitask model, but we also consider intermediate-task training in section 6.4. To ensure our method generalizes well to different settings, we perform experiments on both language models and vision models as well as both full-model and parameter-effiient fine-tuning with (IA)² (Liu et al., 2022), which insorts a trainable vector into specific linear layers.
6.1 Baselines
In our experiments, we compare MaTS against the following baseline merging methods:
Simple Averaging averages the parameter values of the models being merged.
Ensembling ensembles the predicted probabilities from the model to merge. Note that this does not result in I model, but in M models.
Task Arithmetic (Ilharco et al., 2022) introduces a task vector for a particular task as the difference between the fine-tuned parameters (found after training the task) and the original pre-trained parameters. To construct the merged model, Task Arithmetic adds the scaled sum of the task vectors for all tasks to the pretrained model's paramcters. Task Arithmetic introduces the hyperparamcter À to rescale the summed lask vectors.
Terg-Mggging (Yadav et al., 2023) improves upon task arithmetic by removing interference wetw the task vectors. Specifically, TIES zeros out entries in a given task vector that have low magnitude and resolves sign conflicts across different task vectors. Similar to Task Arithmetic, TIES-Merging also uses hyperparameter to scale the task vectors after removing interference.
Diagonal Fisher Merging (Matena & Raffel, 2022) performs Fisher Merging with a diagonal approx- imation to the Fisher, as described in section 2.2). To better distinguish between different approxi- mations of the Fisher, we call this method Diagonal Fisher Merging (DFM).
RegMean (Jin et al., 2022) solves a linear system to match activations between the merged model and the original models. In practice, RegMean scales the non-diagonal terms of the Gram matrix with hyperparameter À for mumerical stability when computing the inverse. Note that RegMean can only be applied to linear layers since closed-form solution does not exist otherwise.
Multitask training (where model is trained on all the tasks jointly) included wherever appropriate and can be considered as loose upper bound on the performance of merged multitask model.
We also experimented with Tangent Task Vectors (Ortiz-Jimenez et al., 2023) but were unable to attain reasonable performance with it and therefore did not include it in our results.
We compute the empirical Fisher over the validation set and compare various ways of computing the Fisher in section 6.6. The only hyperparameter in MaTS is the number of iterations to run the conjugate gradi- ent method, which we allow take on values ranging from 10 100 in step sizes of 10. We tune the hyperparameters for each method based on validation set performance.
6.2 Conjugate Gradient Initializations and Objectives
As discussed in section 3, use of CG in MaTS allows us to consider objectives that various merging methods (implicitly) aim to optimize. In addition, since CG is an iterative optimization method, MaTS can flexibly support different choices of an initialization for optimization. While solving the linear system corresponding to eq. (6) is a convex optimization problem, the choice of initialization will effect convergence speed which is of particular importance because we only run CG for at most 100 iterations. We therefore consider using different merging methods (including those that do not fit our framework) produce an initialization for MaTS.
8

---

Published in Transactions on Machine Learning Research (03/2024)
(IA)
Full-model
30
Change vs. Initialization
Avg.
48.1
48.1
48.1
48.1
48.1
X
56.4 56.4 56.4 56.4
56.4
56.4
15
Objective
DFM
55.7
55.8
56.5
56.7 55.9
X
53.2
58.2 64.9 63.4
57.7
64.7
RM
X
X
X
X
X
X
72.3
72.5 72.6
72.4
71
15
BFM
60.1
62.2
65.1
65.1
62.7
X
28.4 65.3
71.1
71.4
62.6
71.3
PT
Avg
TIES
DFM
RM
PT
Avg.
TA
TIES
DFM
RM
-30
Initialization
Initialization
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.
Experimental set-up. As an experimental setting, we focus on merging models fine-tuned on datasets from the TO mixture (Sanh et al., 2021) to form multitask models. Zhou et al. (2022) found eight datasets (listed in appendix D.2) were the most important for performance and thus we first focus merging models fine-tuned on these eight datasets. We compare merging models in the full model fine-tuning setting and in the parameter--ffiiient fine-tuning setting with (IA)². We use Ti-Large-MM-Adap (Lester et al., 2021), a variant of L5rLarge (Raffel et al., 2020) that was further trained with a language modeling objective as the pre-trained model and include more details on fine-tuning the per-task models appendix D.1.
Compatibility considerations. In this experiment, we consider every possible initialization cacerioon/ojjjeeeee /objective combination for MaTS with the following exceptions: First, when merging (IA)s-trained models, we cannot use RegMean since RegMean a method for merging linear layers and (IA) adds vectors that are elementwise multiplied with activations and does not train any lincar layers. We therefore do not consider RegMean as an initialization or objective for models trained with (TA)3. Second, in the full-model fine-tuning setting, we perform BFM using the K-FAC (Grosse & Martens, 2016) approximation as described in section and apply RegMean usual. However, when merging models fine-tuned with (IA)², we perform BFM without resorting to approximation since the vectors trained by (IA)$ are sufficiently small to perform BFM exactly.
Results. fig. 2, we show the raw performance and performance gain over initialization using MaTS with every possible ioaallizaatio/oojctttitii combination. We see that MaTS often significantly improves performance compared to given initialization, except when relatively poor objective is applied to effective initialization (e.g. using the "simple averaging" objective with the Task Arithmetic initialization). Notably, we found that using RegMean as both an initialization and objective for MaTS led to an improvement in performance over the RegMean-basse based initialization. This performance boost likely comes from the fact that the matrix inverse needed in RegMean is numerically unstable and so the closed-form solution does not actually result the optimal solution. To confirm this possibility, we found that using the conjugate gradient method to optimize the RegMean objective achieves a lower objective value than the closed-form solution.
Based on the results in fig. 2, we now desigu consistent recipe to apply in the experiments in the rest of the paper. For initialization, we use the Task Arithmetic method due to its strong performance and efficiency in both parameter-effiiennt and full-model fine-tuning, For the merging objective, in the (LA)² setup, we use the block-diagonal Fisher merging objective and in the full model fine-tuning setup we use the RegMean objective.
6.3 Multitask Models
In this section, we compare the recipe for MaTS that we designed in the previous experiment various other bascline methods on the problem of merging task-specific models create single a multitask model in two settings for language and setting for vision. For language, in the first setting, the same
9

---

Published in Transactions on Machine Learning Research (03/2024)
Domain
Language Zhou et al. (2022)
Language
Vision
Tasks
Yadav et al. (2023)
Ilharco et al. (2022)
Model
(IA)²
Full Model
(IA)²
Full Model
ViT-B/32
ViT-L/14
Simple Averaging
48.1
56.4
52.1
60.5
65.8
79.5
Task Arithmetic
57.6
63.8
67.1
71.9
70.1
84.4
TIES-Merging
56.8
62.8
66.0
69.6
73.6
85.9
Diagonal Fisher merging
55.9
57.7
61.4
64.0
66.3
82.5
MegMean
69.1
78.9
82.0
89.6
MaTS
65.1
72.5
75.3
81.5
82.6
90.2
Ensemble
62.8
64.9
74.8
74.1
74.0
80.0
Multitask (UPPER BOUND)
76.7
80.5
82.2
84.5
89.0
93.5
Individual Models
78.2
80.7
86.0
85.9
90.5
94.2
Table 1: Performance (average accuracy across tasks) of various merging methods and a multitask training- based bascline. The highest score (excluding the multitask performance) is shown bold.
eight key tasks from the TO training mixture as done in section 6.2 and additionally consider merging fine- tuned variants of T5-Large-LM- Adapt models on the seven tasks used by Yadav et al. (2023) (listed in appendix D.3). For vision, we follow Ilharco et al. (2022); Yadav et al. (2023) and merge the same set of CLIP-based models (Radford et al., 2021) fine-tuned on eight tasks listed in appendix D.4.3 We present the average test-set performance across all tasks in each mixture attained by each merging method in table 1, Notably, MaTS perforrs all other methods, often a significant margin. The gains are most pronounced when merging (LA)d-based models in the language domain, where MaTS outperforms prior methods by about 8% absolute. However, we note that none of the methods (including MaTS) match the multitask training- based baseline, suggesting that there still remains room for improvement when merging many individual-task models to create a multitask model. Per-task performance is reported in appendix
6.4 Intermediate- Task Training
Intermediate task
MNL1
QNLI
QQP
Fine-tuning method
(IA)²
Full Model
(IA)
Full Model
(IA)²
Full Model
RTE (BASELINE)
79.8
75.5
79.8
75.5
79.8
75.5
Simple Averaging
79.8
78.3
79.8
75.1
79.8
78.3
Task Arithmetic
79.8
77.6
74.7
75.5
76.3
78.2
TIES-Merging
78.3
76.5
75.8
75.8
74.7
77.1
Diagonal Fisher merging
79.8
76.2
79.8
75.1
79.8
75.1
RegMean
78.0
75.1
78.3
MaTS
83.0
79.4
80.5
76.9
78.7
77.3
Sequential Training (UPPERBOUND)
87.0
86.6
79.1
79.4
80.1
78.0
Ensemble
79.4
77.3
79.4
73.3
80.1
74.4
Table 2: RTE performance when merging a model fine-tuned on RTE with a model fine-tuned on various intermediate tasks. We bold the highest score.
Next, we consider intermediate-tast, model merging as introduced by Matena & Raffel (2022). Intermediate- task training (Phang et al., 2018; Pruksachatkun et al., 2020) involves first fine-tuning model intermediate task before fine-tuning it on target downstream task in hopes of improving performance. Tn-
stead of performing sequential training, intermediatc-tast. model merging instead merges two models trained independently in parallel on the intermediale task and the larget lask. Following Malena & Raffel (2022), we consider fnn-tuned BERT (Devlin et al., 2018) models with RTE (Dagan et al., 2005) as the target task
We the exact checkpoints released Iharco (2022) at https://github. com/ml foundat www.unndaninnttttaaaaaneeeee vectors
10

---

Published in Transactions on Machine Learning Research (03/2024)
and experiment with MNLI (Williams et al., 2017), QNLI (Rajpurkar et al., 2016; Wang al., 2018), and QQP (Sharma et al., 2018) as the intermediate tasks. We compare the performance of MaTS with other merging methods, the original performance of the RTE-fine model, and the performance of actually doing sequential training in table 2. Overall, we find that MaTS often outperforms all basclines with the exception of intermediate-lan merging with QQP, where most methods fail provide a major improvement (perhaps because QQP not a particularly effective intermediate task for RTE).
6.5
Cost
Having compared the performance of various merg- ing methods, we now compare their computational costs. Specifically, we show the FLOPPs required for MaTS (using the recipe found in section 6.2) and other merging methods in table 3. Though MaTs requires substantially more computation than other methods, most of the corresponding FLOPPs consist of matrix multiplication which are efficient to per- form on GPUs. Also, the number of FLOPs required for MaTS is still dramatically less than the FLOPS for multitask training. In practice, we found that merging with MaTS in the setting from section 6.2 only takes about 11 minutes A6000 GPU, which is relatively modest amount of time in the realm of training large neural networks. We provide the formulas used to derive the FLOP counts in appendix E.1 and provide wall-clock times of all methods in appendix E.2. We separately note that the storage costs required for the models and "metadata" required by each merging method is comparable and, for the models we consider, small enough to on a single GPU.
Merging method
(IA)3
Full Model
Simple Averaging
2.3E6
6.3E9
Task Arithmetic
4.8E6
1.3E10
TIES-Merging
5.0E7
1.8E11
Diagonal Fisher merging
1.2E10
3.2E13
RegMean
2.0E13
MaTS
1.3E14
1.2E14
Multitask Training 1.1E18
2.2E17
Table 3: FLOPS for various merging methods.
single NVIDIA
6.6 Computing the Fisher
As discussed in section 2.2, there exist different perspectives as to whether it is appropriate to use the empirical Fisher (based on ground-truth labels) or true Fisher (based on the model's output distribution) in different settings. We therefore performed an ablation compare using the empirical Fisher and the true Fisher in the experimental setting from section 6.2. For completeness, we also compare computing the Fisher the training set or the validation set. Results are shown in table 4, where we confirm that the overall performance is best when using the empirical Fisher computed on the validation set.
Merging method Fine-tuning method
Diagonal Fisher Merging
Block-diagonal Fisher Merging (IA)s
(IA) 3
Full Model
Full Model
Train
57.7
57.7
67.1
74.3
Empirical Fisher
Validation
59.1
63.8
68.8
76.1
Train
59.6
56.6
67.1
74.2
Fisher
Validation
61.4
62.8
69.6
74.5
Table 4: Average multitask language model accuracy (following the same experimental setting as section 6.2) using the empirical or true Fisher for diagonal and block-diagonal Fisher merging.
6.7 Discussion
Overall, our experimental results confirm that MaTS significantly improves merging performance while in- curring a manageable increase in computational costs. One notable finding that using the block-diagonal Fisher merging (BFM) objective did not provide the best performance in the full-model fine-tuning setting despite the fact that we might expect provide the most reliable notion of "task parameter subspace" among methods consider (as discussed in section 5). We performed some additional analysis to explore
11

---

Published in Transactions on Machine Learning Research (03/2024)
this discrepancy. First, we consider whether better initialization could yield better performance with BFM. Specifically, since using task vector-based initialization with the MeaM objective in MaTS performed well, we considered using the solution found by this recipe as an initialization for an additional round of MaTS using the BFM objective. In the experimental setting of section 6.2, adding an additional round of BFM-basod merging increased the average performance across lasks from 72.5 to 73.5, outperforming all other methods. This result further emphasizes the importance of the choice of initialization in MaTS and also hints the pos- sibility of improved performance via multiple rounds of merging. Separately, we experimented with whether optimizing the BFM-based objective would benefit from more iterations of CG. However, when increasing the maximum number of iterations from 100 to 5K, we observed that the average performance could drop significantly and include graph in appendix G. We hypothesize this is due to numerical instability of the gradients. More broadly, however, our results generally suggest that improved approximations of the Fisher could improve the performance of Fisher merging- based methods.
7
Related Work
7.1 Loss Landscape
One common assumption for merging models is that the models lie in a "basin" of low loss in parameter space (Ilharco et al., 2022). There have been many past works focused on finding simple (not necessarily lincar) paths of low loss between two models (Kuditipudi et al., 2019; Draxler et al., 2018; Garipov et al., 2018; Tatro et al., 2020; Benton et al., 2021). Entezari et al. (2021) hypothesize that all models trained with SGD can be permuted such that they are all linear mode connected (i.e. there is no barrier of high loss along the linear path connecting two models) (Frankle et al., 2020; Nagarajan & Kolter, 2019). The permutation symmetrics present in neural networks imply that activations in models trained from different initializations may require permuting to put the models in the same loss basin. Neyshabur et al. (2020) show that models fine-tuned from the same pre-trained checkpoint lie in the same basin and Qin et al. (2022) studies the effect of hyperparameters on mode connectivity. Singh & Jaggi (2020); Ainsworth et al. (2022); Jordan et al. (2022); Peña et al. (2023) introduce various algorithms that account for soft or hard permutation symmetries in neural networks when performing merging and Stoica et al. (2023) extend this idea to merging models with different architectures. A related finding is the "monotonic linear interpolation property", which states that the loss of complex neural networks monotonically decreases when trained with SGD (Goodfellow et al., 2014) though models can be trained in such a way where this property does not always hold (Lucas et al., 2021). When lincar mode connectivity does not hold, Lubana et al. (2023) show that models can use dissimilar mechanisms for making predictions and Juneja et al. (2022) show that non-mode-conneeted models generalize differently.
Past work has also aimed to analyze the structure of the loss landscape and try to directly learn a subspace of low loss. Some examples include Wortsman et al. (2021) learning a line, Hochreiter & Schmidhuber (1997) optimizing a box, or Fort & Jastrzebski (2019) finding an intersection of wedges. Other works have proposed that this "important subspace" can be found via the top eigenvectors of the Hessian (Fort & Ganguli, 2019), which correspond to the singular vectors of the covariance matrix (Izmailov al., 2020), where the covariance can be estimated via SWAG (Maddox et al., 2019; Shwartz-Ziv et al., 2022) or other approximations such as the Laplace approximation (Maddox et al., 2020).
7.2 Fisher
There has been a great deal of work better approximating the Fisher, usually in the context of using Natural Gradient Descent for optimization (Amari, 1998). K-FAC (Martens & Grosso, 2015) approximates the Fisher for linear layers as the Kronecker product of the covariance of the input activations and the covariance of the output activations. This have been extended to convolutional neural networks (Grosse & Martens, 2016) and recurrent neural networks (Martens et al., 2018). Other works have focusod on reducing K-FAC's compute (Tang et al., 2021) or memory (Pauloski al., 2021) and extending K-FAC for distributed training (Ba et al., 2017: Osawa et al., 2023). George al. (2018) improve upon K-FAC by ensuring the
12

---

Published in Transactions on Machine Learning Research (03/2024)
diagonal of the K-FAC approximation is unbiased. We note that using the Fisher is only one instance of the covariance matrix in our unified framework and MaTS uses the RegMean objective for full-model fine-tuning.
7.3 Additional Applications of Model Merging
Recently, many works have focusod on merging models for different use cases. This includes merging models for intermediate Lask-training (Choshen et al., 2022; Gueta et al., 2023), generalizing better to out-of- distribution shifts (Wortsman et al., 2022b;a; Ramé et al., 2022), and combining mppelity-specific models into multimodal model (Sung ct al., 2023). Concurrently, other works have focused on inaccuracies model merging stemming from gradient mismatches in different models, and use these insights connect diagonal Fisher merging and Task Arithmetic Daheim et al. (2023).
8 Conclusion
We propose MaTS, a framework that connects past merging methods by viewing merging matching indi- vidual models in their task parameter subspace. In general, using MaTS requires solving linear system, we propose using the conjugate gradient method which enables solving linear systems that would otherwise be intractable to solve. To explore a merging objective that has no tractable closed-form solution, we develop a variant of Fisher merging that makes a block diock-diagonal approximation of the Fisher. Experimentally, we show that MaTS achieves statc-of-the-art results in merging to create multitask models from individual-task language and vision models and in intermediate-takk merging. Our work provides new perspective on model merging and motivates future work on improved methods for estimating model's task parameter subspace.
Acknowledgments
We thank Felix Dangel for pointing us to the conjugate gradient method for inverting sum of Kronecker products, Juhan Bae for clarifying discussions on the Fisher, Michael Matena for clarifying discussions the Fisher and valuable feedback on an earlier draft of this work, and Nikhil Kandpal for helpful discussions on how best to frame our work and valuable feedback on an earlier draft of this work.
References
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXie preprint orXiv:2209.0458 2022.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural Compulation, 10(2):251 276, 1998.
Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kronecker- factored approximations. International Conference on Learning Representations, 2017.
Slephen H Bach, Victor Sanh, Zheng-Xin Xin Yong, Albert Webson, Colin Raffel, Nihal V Nayak, Abheesht Sharma, Taewoon Kim, M Saiful Bari, Thibault Fevry, et al. Promptsource: An integrated development environment and repository for natural language prompts. arXiv preprint arXi::220..00 01279, 2022.
Gregory Benton, Wesley Maddox, Sanae Lotfi, and Andrew Gordon Gordon Wilson. Loss surface simplexes for mode connecting volumes and fast ensembling. In International Conference on Machine Learning, pp. 769-779. PMLR, 2021.
Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEFE, 186610):1665-1883 2017.
Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pre- training. Xin preprint ar 20X04:2204.03044 2022.
13

---

Published in Transactions on Machine Learning Research (03/2024)
Mircea Cimpoi, Subhransu Maji, Iasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. Describing textures in the wild. In Proceedings of the TEEE conference on computer vision and paltern recognition, pp. 3606 3613, 2014.
Daniel Cohen, Liu Yang, and W Bruce Croft. Wikipassageqa: A benchmark collection for research on non-factoid answer passage retrieval. In The 4 1sst international ACM SIGIR conference on research & development in information velrieval, pp. 1165 1168, 2018.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine Learning Challenges Workshop, pp. 177-190. Springer, 2005.
Nico Daheim, Thomas Möllenhoff, Edoardo Maria Ponti, Iryna Gurevych, and Mohammad Emtiyaz Khan. Model merging by imcortainty-basee gradient matching. arXiv preprint 1tfi::8810.22000 2023.
Jacob Devlin, Ming-Woi Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. orXiv preprinl 1rXiv:1810.044 2018.
Shachar Don-Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. ColD
fusion: Collaborative descent for distributed multitask finetuning. ar preprint XrXtv22122001978 2022.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309-1318. PMLR, 2018.
Rahim Entezari, Hanic Sodghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprinl or Xrx:22110..0 0..6696, 2021.
Stanislav Fort and Surya Ganguli. Emergent propertics of the local geometry of neural loss landscapes.
arXiv preprint arX:::1910.05929 2019.
Stanislav Fort and Stanislaw Jastrzebski. Large scale structure of neural network loss landscapes. Advances in Neural Information Processing Systems, 32, 2019.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259 3269.. PMLR, 2020.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew Wilson. Loss surfaces, mode connectivity, and fast. ensembling of dnns. Advances in neural information processing systems, 31, 2018.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.
Ian J Goodfellow, Oriol Vinyals, and Andrew M Saxe. Qualitatively characterizing neural network optimiza- tion problems. arXiv preprint arXiv:1412.6544, 2014.
Roger Grosse. Chapter 2: Taylor approximations. lecture notes. University of Toronto CSC2541, Top- ies in Machine Learning: Neural Net Training Dynamics, 2022a. URL https://www. cs, toronto. edu/ -Tgrosse/coursee//sseee /ttp:///2222.10eeeeeeeaag///2222 wtwp.alorrpppproiaatitooooooooo. pdf.
Roger Grosse. Chapter 3: Metrics. lecture notes. In University of Toronto CSC2541, Topics Ma-
chine Learning: Neural Net Training Dynamics, 2022b. URL https://wwat http://wwww....ootoooooooooo toronto..o..ooooo edd.rrrooee////rrseeee//// /ttpa///eee22442222eeeeraai
www.:o/eeeeeiit__eeee pdf.
Roger Grosse and James Martens. A kronccker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573 582. PMLR, 2016.
Roger Grossc, Juhan Bac, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhosscin Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. ar Xiv preprint or arXiv:208..0 03296, 2023.
14

---

Published in Transactions on Machine Learning Research (03/2024)
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Knowledge is a. region in weight space for fine-tuned language models. arXiv preprint arXiv:2302. 04863, 2023.
Ken Hayami. Convergence of the conjugate gradient method on singular systems. arXiv preprint arXiv:1809.0070 2018.
Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Introducing eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. In IGARSS 2018-2018 IEEE International Geoscience and Remote Sensing Symposium, pp. 204-207. IEEE, 2018.
Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 2019.
Magnus R. Hestenes and Eduard Sticfel. Methods of conjugate gradients for solving linear systems 1. Journal of Research of the National Burewu of Slandards, 49(6):409 -36,, 1952. Research Paper 2379.
Sepp Hochreiter and Jürgen Schmidhuber. Flat Minima. Newral Computation, 9(1):1-42, 01 1997. ISSN 0899-7667. doi: 10.1162/neco.19997.111 URL http:://boi.orr/10..1 wwtp:///....1////ooooooooo w9w.eeec....99.... 1997.9.9.99.. 9.1.
Sebastian Houben, Johannes Stallkamp, Jan Salmen, Mare Schlipsing, and Christian Igel. Detection of traffic signs in real-world images: The German Traffic Sign Detection Benchmark. In International Joint Conference on Newral Networks, 2013.
Lifu Huang, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Cosmos qa: Machine reading compre- hension with contextual commonsense reasoning. arXiv preprint arXiv:1909.008: 00277, 2019.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannanch Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. arXiv preprint arXiv:2212. 04099 2022.
Pavel Izmailov, Wesley J Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Subspace inference for bayesian deep learning. In Uncertainty in Artificial Intelligence, pp. 1169-1179. PMLR, 2020.
Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietry and Pengxiang Cheng, Dataless knowledge fusion by merging weights of language models. arXiv preprint arXiv:2112.098. 09849, 2022.
Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. Repair: Renormalizing permuted activations for interpolation repair. arXiv preprint urXiv:2211.08.00 2022.
Jeevesh Juneja, Rachit Bansal, Kyunghyun Cho, João Sedoc, and Naomi Saphra. Lincar connectivity reveals generalization strategies. arXiv preprint arXiv:2205.12411 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodci. Scaling laws for neural language models. arXiv preprint arXiv:2001.00 0,366 2020.
Tushar Khot, Peter Clark, Michal Guerquin, Peter Jansen, and Ashish Sabharwal. Qasc: A dataset for ques- tion answering via sentence composition. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained catego- rization. In 4th International IEEE Workshop on 3D Representation and Recognition (3dRR-13) 13), Sydney, Australia, 2013.
Agustinus Kristiadi. Fisher information matrix. In Tech Blog, 2018. URL https://aguusiius.. krisstiiiiinkkikiia.... ae//echhoo///00////////////tttte inomatitii........oi/oo.....
Rohith Kuditipudi, Xiang Wang, Holden Loe, Yi Zhang, Zhiyuan Li, Wei Hu, Rong Gc, and Sanjoev Arora. Explaining landscape connectivity of low-cost solutions for multilayer nets. Advances in neural information processing systems, 32, 2019.
15

---

Published in Transactions on Machine Learning Research (03/2024)
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. Advances in neural information processing systems, 32, 2019.
Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficiet prompt tuning. arXiv preprint ar arXit:2104.0861 2021.
Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. Thirteenth Inter- national Conference on the Principles of Knowledge Representation and Reasoning, 2012.
Kevin Lin, Oyvind Tafjord, Peter Clark, and Matt Gardner. Reasoning over paragraph effects in situations. arXiv preprint Xin: 1::::000.000 0008.05852. 2019.
Haokun Liu, Derek Tam, Mohammed Mugeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-effcient fine-tuning is better and cheaper than in-context learning. Advances in Neural Information Processing Systems, 35:1950 1963, 2022.
Ekdeep Singh Lubana, Eric J Bigelow, Robert P Dick, David Krueger, and Hidenori Tanaka. Mechanistic mode connectivity. In International Conference on Machine Learning, pp.
PMLR, 2023.
James Lucas, Juhan Bac, Michael R Zhang, Stanislav Fort, Richard Zemel, and Roger Grosse. Analyzing monotonic linear interpolation in neural network loss landscapes. arXiv preprinl arXiv:2144. 11044, 2021.
Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainly in deep learning. Advonces in neural information processing systems, 32, 2019.
Wesley J Maddox, Gregory Benton, and Andrew Gordon Wilson. Rethinking parameter counting in deep models: Effective dimensionality revisited. arXiv preprint XrX::2000. 02139, 2020.
James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1:55776-5811, 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curva-
ture. In International conference on machine learning, pp. 2408 2417. PMLR, 2015.
James Martens, Jimmy Ba, and Matt Johnson. -ronecker-factored curvature approximations for recurrent neural networks. International Conference on Learning Representations, 2018
Michael S Matena and Colin A Raffel. Merging models with fisher weighted averaging, Advances in Newral Information Processing Systems, 35:17703-17716, 2022.
Brendan McMahan, Eider Moore, Danicl Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicaiion-efiiiell learning of deep networks from decentralized data. In Artificial intelligence and slatistics, 2017.
Vaishnavh Nagarajan and J Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning. Advances in Neural Information Processing Systems, 32, 2019.
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learning? Aduances in neural information processing systems, 33:512-523, 2020.
Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. arXiv preprint 1xX2222055.2827, 2023.
Kazuki Osawa, Shigang Lú, and Torsten Hoefler. Pipefisher: Efficient training of large language models using pipelining and fisher information matrices. Proceedings of Machine Learning and Systems, 5, 2023.
Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian Foster, and Zhao Zhang. Kaisa: an adaptive sccond-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2021.
16

---

Published in Transactions on Machine Learning Research (03/2024)
Fidel A Guerrero Peña, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli. Re-basin via implicit sinkhorn differentiation. In Proceedings of the IREE/CVP Conference on Computer Vision and Pattern Recognition, pp. 20237 20246, 2023.
Jason Phang, Thibault Févry, and Samuel R Bowman. Sentence encoders on stilts: Supplementary training on intermediate labeled-data tasks. arXiv preprint wrXiv:1811.01088, 2018.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: improving latent diffusion models for high-resolution image synthesis. anXiv preprinl arXiv:2307.01552, 2023.
Yada Pruksachatkum, Jason Phang, Haokun Liu, Phu Mon Htut, Xiaoyi Zhang, Richard Yuanzhe Pang, Clara Vania, Katharina Kann, and Samuel R Bowman. Intermediate-task transfer learning with pretrained mod- els for natural language understanding: When and why does it work? arXiv preprint or orXiu22055.0622. 2020.
Yujia Qin, Cheng Qian, Jing Yi, Weize Chen, Yankai Lin, Xu Han, Zhiyuan Liu, Maosong Sun, and Jie Zhou. Exploring mode connectivity for pro-trained language models. arXiv preprint arXiv:2210. 14102, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Nla Sutskever. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 9748-8763. PMLR, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5551 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint 1rXi::1606.0080 2016.
Alexandre Ramé, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David Lopez-Paz. Recy- cling diverse models for out-of-distriibutio generalization. arXiv preprinl arXiv:2212. 10445, 2022.
Anna Rogers, Olga Kovaleva, Matthew Downey, and Anna Rumshisky. Gelting closer ai complete question answering: A set of prerequisite real tasks, In Proceedings of the Thirty- Fourth AAAI Conference Artificial Intelligence (AAAT-20), 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema. challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Victor Sanh, Albert Webson, Colin Raffel, Stephen H Bach, Lintang Sutawika, Zaid Alyafcai, Antoine Chaffin, Arnaud Stiegler, Teven Le Scao, Arun Raja, et al. Multitask prompted training enables zero-shot task generalization. arXiv preprint arXiv:22 10.08207, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint 1rXiv:1904.09798, 2019.
Rishi Sharma, James Allen, Omid Bakhshandeh, and Nasrin Mostafazadeh. Tackling the story ending biases in the story cloze test. In Proceedings of the 56th Annual Meeting of the Association. for Computational Linguistics (Volume 2: Short Popers), pp. 752 757, 2018.
Ravid Shwartz-Ziv, Micah Goldblum, Hossein Souri, Sanyam Kapoor, Chen Zhu, Yann LeCun, and An- drew G Wilson. Pre-train your loss: Easy bayesian transfer learning with informative priors. Advances in Newral Information Processing Systems, 35:27706-2777 27715, 2022.
Sidak Pal Singh and Dan Alistarh. WoodFisher: Efficient second-order approximation for neural network compression. Advances in Newral Information Processing Systems, 2020.
17

---

Published in Transactions on Machine Learning Research (03/2024)
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Newral Information 2020.
Processing Systems,
Sebastian U. Stich. Local SGD converges fast and communicates little. arXiv preprint arXis:1805.00700 2018.
George Stoica, Daniel Bolya, Jakob Bjorner, Taylor Hearn, and Judy Hoffman. Zipit! merging models from different tasks without training. arXiv preprint arXie:2305.0000 03053, 2023.
Yi-Lin Sung, Linjie Li, Kevin Lin, Zhe Gan, Mohit Bansal, and Lijuan Wang. An empirical study of multimodal model merging. arXic preprint arXix:2904. 14983 fcs. CVJ, 2023. doi: 10.48550/arxiv.2300 14933. URL attps://arxivv org/abs/ 2204./234...10 .
Oyvind Tafjord, Matt Gardner, Kevin Lin, and Peter Clark. Quartz: An open-domain datasct of qualitative relationship questions. arxiv preprint Xiv: 1.09. 190..09599,53, 2019.
Zedong Tang, Fenlong Jiang, Maoguo Gong, Hao Li, Yue Wu, Fan Yu, Zidong Wang, and Min Wang. Skfac: Training neural networks with faster cronecker-factordd approximate curvature. IEEE/CVE Conference on Compuler Vision and Pallern Recognition, 2021.
Norman Tatro, Pin-Y Yu Chen, Payel Das, Igor Melnyk, Prasanna Sattigeri, and Rongjie Lai. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems, 33:15300-15311, 2020.
Alex Wang, Amanproct Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07401 2018.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint ar x 1044. 05426, 2017.
Mitchell Wortsman, Maxwell C Horton, Carlos Guestrin, Ali Farhadi, and Mohammad Rastegari. Learning neural network subspaces. In International Conference Machine Learning, pp. 11217-11227. PMLR, 2021.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roclofs, Raphael Gontijo-Lopes, Ari S Mor- cos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Interna- tional Conference on Machine Learning, pp. 23965 23998. PMLR, 2022a.
Mitchell Wortsman, Gabriel Ihharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- 2022b.
nilion, pp.
Jianxiong Xiao, Krista A. Ehinger, James Hays, Antonio Torralba, and Aude Oliva. SUN database: Exploring a large collection of scene calegories. International Journal of Compuler Vision (IJCV), 119(1):3 22, August 2016.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interference when merging models. arXiv preprint 0rXiv:2306.01700 01708, 2023.
LeCun Yann. The mnist database of handwritten digits. 1998. URL http::/yann lecun. womm/xdb//ee /iibtt//nis.t.....
Netzer Yuval. Reading digits in natural images with unsupervised feature learning. Tn Proceedings of the NIPS Workshop Deep Learning and Unsupercised Feature Learning, 2011.
Yuan Zhang, Jason Baldridge, and Luheng He. Paws: Paraphrase adversaries from word scrambling, arXiv preprint Kie: 01130, 2019.
Jing Zhou, Zongyu Lin, Yanan Zheng, Jian Li, and Zhilin Yang. Not all tasks are born equal: Understanding zero-shot generalization. The Eleventh International Conference on Learning Representations, 2022.
18

---

Published in Transactions on Machine Learning Research (03/2024)
A Derivations
A.1 Fisher Merging
Let be the number of parameters in the model and 0; ~ N (14, Ei) The joint likelihood can be written as
- 10 0.))
log (2T) & det 23 exp
The derivative the joint log likelihood is:
M
log
P(0
d
Σ
log
(2T) 3 det E, exp (- 0,) 0.)
de
de
i=1
M
p
log(2n) + log det E; (0 8, 0 (0 (,-
d0
2
M
(0-
0;) (0 04)
d0
2
1=1
M
(0 0z)
i=1
M
2-10
0 +
10 0g
Setting the derivative 0 and solving, we get
M
M
Σ
0
0
1=1
i=1
M
-1
M
0
Σ
Σ
Σ-10
Σ
0
1= 1
A.2 RegMean
The minimization objective is below:
M
Om
Zm W|12
min
Nm
m=1
M
(Om Zm W) (owp)((p,,) (0m Zm W)
min
W
N
m=1
T 2( W) On + (Z T (2.
min
W
N m
The gradient of the objective is
22, + +22 2,, W
N
m=1
19

---

Published in Transactions on Machine Learning Research (03/2024)
Setting the derivative to 0 and solving, get
M Σ
M Σ
22 m
22,
Zm W
(15)
N
N
m=
12
Z mm W... m Σ
Z Zm W
(16)
Nm
Nm
m==1
m==1
M
-1
W
Σ
Z
Zm
Σ Nm
(Z Zm W
(17)
N
A.3 Task Parameter Subspace Derivation
C. P. Pm
Zm V Zu En V
Vn Em U Un m
Em VT
m
QmAm
A.4 Connecting Fisher Merging and RegMean Derivation
N
F
W. log V w, (onpo().lem))
(18)
N
2n 2m
(19)
N
Zz
(20)
N
where W, refers to the ith column of W.
A.5 K-FAC Derivation
W
F
Σ
Fm vec m)
M
N
Z m
Zm
N
o0
m
OTO
m
Σ
((1
N
2
m
2m
8
N
1
m
veal (w.))
M
M
Σ
Nm
ZT Zm
8
N
o
Or
N
Z Zm
Wm
N
o m
O
m
(21)
m=
20

---

Published in Transactions on Machine Learning Research (03/2024)
A.6 Conjugate Gradient Derivation
Σ Fm vec Wm)
M
Σ
Zi Zm
8
om 0mm) vec(Wm)
Nm
Nm
m=1
M
Zy Zm
Wm m
1
O 0m))
Nm
Nm
m=1
O 0mn)
Aa:
ZT
Zm
vec(W)
N
N
m==1
Σ
Z Zm
O 0in)) vec(W)
Nm
Nm
m
M
2 Zm
W
OTO'
Nm
Nm
m=1
B Conjugate Gradient Method
Instead of using the gradient of the linear objective Ar b though, CG uses the gradient of convex quadractic objective z" Ar x T b. Both the linear objective and the convex quadratic objective have the same minimum. The advantage of using the quadractic objective to compute the gradient is that its gradient is equal to Ax - b, which is just the "error" of our current estimate of in the linear objective. CG additionally constrains all update directions computed at every iteration to be conjugates. While two vectors are consider ort hogonal if their dot product is O, two vectors are conjugates if their dot product with respect to a matrix (in this case A) is 0. Constraining the update directions to be conjugates implies that all future updates do not affect the optimality of previous update, and thus the optimal step size for a given update direction can be computed. This ultimately ensures that the number of iterations required for convergence is upper bounded by the dimension of the linear system. In practice, we observe the number of iterations necessary for convergence to be much smaller since the update direction corresponds to the gradient of an objective with the same minimum.
C Algorithm
We show the algorithm for MaTS in the (IA)3 setting in algorithm In the full-model setting, the outer product of the gradients is replaced with the outer product of the input activations.
21

---

Published in Transactions on Machine Learning Research (03/2024)
Algorithm 1: MaTS Algorithm
Input : :01..0M
b M fine-tuned models sharing a common architecture and initialization M validation datasets compute model statistics > Mcrged model
Input : D1 .. DM
Output: l"
// For each task, compute the Fisher. For brevity, only show computing the Fisher for each layer.
for m= =1 1, m S M, m ++ do
Fm- =
for (z. E Dm do
Fm += V log pe. (ulz) Pe.. pe.. (y/x)
Outer product of per-example gradient Normalize Fisher by number of examples
end
Fmm Fm/ mm
end
M Σ Fm m=1
D. in the linear system Ax = Initialize 0 using Task Arithmetic
frask, orathecii
Conjugate Gradient Method for i = m Number of Iterations, i + 1 do
Update 0 using conjugate gradient where Ar Σ O
end return 0
Experimental Setup
D.1 Fine-tuning Details
For all the models we fine-tuned, we the same hyperparameter setup to be consistent. Concretely, use AdamW, learning rate of le-1, batch size of 1024, bfloat16 during training, training for 5K batches, checkpointing every 100 batches, and carly stopping if the model has not improved for 5 checkpoints. For (TA)3, the only difference is the learning rate is set to -3
We use the templates from PromptSource (Bach et al., 2022) to format the input/output pairs in natural language. All the models were trained and evaluated on all the templates available in PromptSource. Instead of evaluating the dataset. across all templates, we randomly assign fixed template to each datapoint and evaluate the dataset once. This is equivalent to averaging the performance across templates.
D.2 8 Keys tasks for the TO Mixture
The tasks and datasets were question aarbllonnnnwwering (CosmosQA (Huang et al., 2019), SocialIQA (Sap et al. 2019), QuAIL (Rogers et al., 2020), Wiki QA (Cohen et al., 2018), QuaRTz (Tafjord et al., 2019), QASC (Khot al., 2020), ROPES (Lin et al., 2019)) and paraphrasing (PAWS (Zhang et al., 2019)).
D.3 TIES Mixture
The datasets are PAWS (Zhang al., 2019), QASC (Khot et al., 2020), QuaRIz (Tafford al., Cloze (Sharma et al., 2018), Wiki QA (Cohen et al., 2018), Winogrande (Sakaguchi et al., 2020), and WSC (Levesque et al., 2012).
D.4 Vision Tasks
These checkpoints were initialized from CLIP (Radford et al., 2021) and individually fine-tuned on eight image classification datasets: Cars (Krause et al., 2013), DTD (Cimpoi et al., 2014), EuroSAT (Helber
22

---

+ FLOPS, where di and kę are Published in Transactions on Machine Learning Research (03/2024)
al., 2019; 2018), GTSRB (Houben et al., 2013), MNIST (Yann, 1998), RESISC45 (Cheng et al., 2017). SUN397(Xiao et al., 2016), and SVHN (Yuval, 2011).
E
Cost
E.1
FLOPS Derivation
We assume each model has p parameters, or |0| p and the total number of tokens and examples used to compute the model statistics (i.e. the Fisher or the covariance of the input activations) are T and N respectively. We follow Kaplan et al. (2020); Liu et al. (2022) to estimate the FLOPS to compute the gradient (forward pass + backward pass) as 3p7 and the number of FLOPS to compute a forward pass in a model is pf.
Averaging requires Mp FLOPS. Summing the models requires (M 1)p FLOPS and dividing by the total number of models requires another p FLOPS.
Task Arithmetic requires 2Mp + FLOPS. Computing the task vector requires Mp FLOOPS, adding task vectors requires (M- - 1)p FLOPS, adding the pretrained model requires p FLOPS, and multiplying the final task vector by the scalar requires another p FLOPS.
Diagonal Fisher merging requires 3Mp- - p + M(3TT + Np) FLOPS. Computing the gradient requires 3pt FLOPS. Computing the diagonal Fisher requires Np FLOPS, since it is an element-wise multiplication for each gradient computed per-example. This results in (3yT NNp) FLOPS to compute the diagonal Fisher across all model. Scaling each model by its diagonal Fisher requires Mp FLOPS, adding the models scaled by the diagonal Fisher requires (M- - 1)p FLOOP,, adding the diagonal Fishers requires (M - 1)p FLOPS, and dividing the sum of the scaled models by the sum of the diagonal Fishers requires another p FLOPS.
TIES requires 9.8Mp + Mplogp- - 2p FLOPS. The trimming stage requires a total of 1.8Mp + Mp log FLOPS. Mp FLOPS to compute the magnitude for each paramcter across all models, Mplogp FLOPP to sort the magnitudes for each model, and 0.8Mp FLOPS to reset. the bottom the 0.8 parameters by magnitude to 0. The elect stage requires total of 6Mp p total FLOOP.. For the elect stage, computing the positive mass requires Mp FLOPS to compute the sign for each parameter across all models, Mp FLOPS to multiply paramcter by its sign, (M - 1)p FLOPS to sum the positive mass for models. Computing the negative mass doubles the amount of FLOPS. Comparing the positive and negative mass and keeping the larger mass requires another p FLOPS. The disjoint merging stage requires a total of 2Mp FLOPS. This includes M FLOPS to multiply the parameters by the correct sign (or the parameter should be masked out) and another (M - 1)p FLOPS to sum the parameters.
RegMean requires Mt -
2 +
the input dimension and output dimension for layer l. For given linear layer, summiing the gram matrices requires (M 1)d² FLOPS, inverting the sum of the gram matrices requires
2 d3 FLOOP,, multiplying the gram matrices by the model weight requires M FLOPS, summing the model weights modified by the gram matrices requires (M-1)dk FLOPS, and multiplying the final 2 terms require d FLOPS. Computing the forward pass activations requires p? FLOPS per model.
3
We assume MaTS requires / iterations. For merging full models, MaTS requires MpT+ + Σ (M - 1)df +M
(M -M-11)dki + I[dhk: + 12d/k) Tof FLOPS. For merging (IA)², MaTS requires 3MpT
> (M 1)dff
Mdfka + (M - 1) + [ kddfka + 12d(k) + Nof FLOPS
For a given layer, summing the gram matrices and summing the models multiplied by the gram matrices requires (M-1)) and Ma (M- 1)dk FLOPS respectively just as in RegMean. Within each iteration, the matrix vector product requires dh k FT.PPS and the 12 vector operations within each iteration to compute
23

---

Published in Transactions on Machine Learning Research (03/2024)
the step size require 12dk FLOPS. This means I iterations requires I(dk + 12d2k) FLOPS. Computing the input activations covariances for linear layers requires Td FLOPS, since it is matrix multiplication between a matrix of size d x I and T z d and the blockwise Fisher matrices for LA3 requires Nd FLOPS since it is a matrix multiplication between a matrix of size x and x d. Computing the forward pass activations requires y7 FLOPS per model and computing the gradient requires 3yf FLOPS per model.
For multitask learning, we follow Kaplan et al. (2020); Liu et al. (2022) to estimate the FLOPS to train an encoder-decoder model 3pT where T is the total mumber of tokens during training. The batch size is 1024 and the average sequence length is 138.5. For full model fine-tuning, the model was trained 2K batches and for (IA)², the model was trained 10K batches.
For full-model fine-tuning, there are 783M parameters, with 288 linear layers that have input and output dimension 1024, 96 layers that have input dimension 1024 and output dimension 2816 and 48 linear layers that have input dimension 2816 and output dimension 38. For (IA) there are 282K parameters, with 144 vectors having dimension 1024 and 48 vectors having dimension 2816. The number of tokens used to estimate the Fisher is roughly T 1385 since the average sequence length is 138.5 and the Fisher is computed on 1K examples.
E.2 Wall Clock Time
In table 5, show the wall-clock time for merging various methods an A6000 machine. The wall-clock time only includes computing the merge (and statistics for methods that require them) and not actually loading any checkpoints. Though MaTS is slower than other methods, it is only by a factor of 2 - 3r for full-model fine-tuning and is less than 6 seconds for (IA)².
Time (seconds) Full Model
Method
(IA)²
Avg. Task Arithmetic (Ihharco al., 2022)
1.1
6.02.0
1.7 1.7F. 9.5--
8.91.7
244-3
TIES (Yadav et al., 2023)
4.0E, 5, 1E-a
75.711.9
DFM (Matena & Raffel, 2022)
1104g
127215
RM (Jin et al., 2022)
66517.7
MaTS
1 1128
128728
Table 5: Wall-clock time for various merging methods. We report the mean and standard deviation across 10 runs.
F
Full Results
The full results of the performance of various merging methods on the individual tasks is shown below. This includes full-model fine-tuning and (IA)² on the 8 key tasks from the TO training mixture and the 7 tasks from the TIES mixture on T5-Large.
G Performance vs. Number of Iterations
We include graph of the performance of the merged model and the error in solving the linear system vs. the number of iterations of the conjugate gradient method in fig. 3. We show the performance using the block-diagonal Fisher merging objective on (TA)² and full-model finc-tuning and the RegMean Jin et al. (2022) objective on full-model fine-tuning.
In most cases, we do find that the merging objective value decreases as the number of iterations increases. In some cases, we observe small amount of noise in the merging objective, which we attribute to numerical instability, but even in these cases the objective value tends decrease. some cases, we do not observe the performance (e.g. accuracy) to increase the number of CG iterations increases.
24

---

Published in Transactions on Machine Learning Research (03/2024)
Method
Avg.
CosmosQA
SocialTQA
PAWS
QUAIT.
WikiQA
QuaRTz
QASC
ROPES
Avg.
56.4
41.8
51.6
55.7
49.3
66.3
62.2
82.6
41.4
Task Arithmetic
63.8
54.5
56.9
61.1
56.9
76.5
76.6
79.0
48.7
TIES
62.8
55
60.1
88.7
58.5
35.9
80.2
76.6
47.3
DFM
57.7
44.9
51
51.5
56.1
49.5
81.5
78.0
49.4
RegMean
69.1
47.6
59.1
79.6
54.2
88.9
80.5
94.2
48.3
MaTS
72.5
53.7
60.2
84.2
56.5
94.0
85.2
94.8
51.6
Ensemble 64.9
50.0
51.1
89.2
56.1
72.6
78.4
79.2
42.3
Multitask 80.5
69.3
62.3
94.7
64.9
95.4
89.1
97.8
70.4
Table 6: Accuracy of different merging methods on the 8 keys tasks from the TO training mixture when using full-model fine-tuning.
Method
Avg.
CosmosQA
SocialIQA
PAWS
QuAIL
WikiQA
QuaRTz
QASC
ROPES
Avg.
48.1
37.4
47.6
48.6
43.3
37.1
58.9
81.6
30.7
Task Arithmetic
57.6
43.9
53.0
59.0
52.3
49.4
74.0
88.4
41.1
TIES
56.8
40.7
50.4
76.1
47.9
52.2
72.9
75.0
39.3
DFM
55.9
40.1
46.4
52.0
46.2
44.5
87.5
77.0
53.9
MaTS
65.1
48.5
53.2
65.8
55.2
77.8
77.6
87.8
55.0
Ensemble 62.8
49.4
50.6
79.5
51.3
61.7
84.9
88.2
36.8
Multitask 76.7
64.3
59.6
94.5
60.7
95.5
83.9
96.4
58.5
Table 7: Accuracy of different merging methods on the 8 keys tasks from the TO training mixture when using (TA)3.
Method
Avg.
PAWS
QASC
QuaRT>
Story
WikiQA
Winogrande
WSC
Avg.
60.5
55.9
71.6
54.4
56.4
76.7
53.4
54.8
Task Arithmetic
71.9
76.5
64.2
70.6
80.5
90.9
65.5
54.8
TIES
69.6
89.5
78.8
78.1
73.8
36.0
73.2
57.7
DFM
64.0
55.7
75.6
85.9
69.6
50.4
59.0
51.9
RegMean
78.9
82.2
94.8
76.6
80.6
93.9
61.6
62.5
MaTS
81.5
86.6
95.2
81.8
83.0
94.5
66.6
62.5
Ensemble 74.1
88.9
60.6
76.3
74.6
92.2
64.7
61.5
Multitask 84.5
95.2
97.2
87.5
92.1
95.9
74.5
49.0
Table 8: Accuracy of different merging methods the tasks from the TIES training mixture for T5-Large when using full-model fine-tuning.
Method
Avg.
PAWS
QASC
QuaRTz
StoryCloze
WikiQA
Winogrande
WSC
Avg.
52.1
49.7
67.8
53.4
57.4
39.0
52.5
45.2
Task Arithmetic
67.1
55.2
75.0
69.5
77.7
75.8
58.9
57.7
TIES
66.0
70.1
80.6
72.9
73.3
56.2
58.1
51.0
DFM
61.4
51.2
73.0
89.3
72.8
45.0
57.1
41.3
MaTS
75.3
75.6
83.8
78.4
81.6
86.6
61.6
59.6
Ensemble
74.8
85.4
79.6
83.6
87.8
72.8
60.4
53.8
Multitask 82.2
94.9
95.8
81.2
88.2
95.3
66.4
53.8
Table 9: Accuracy of different merging methods on the tasks from the TIES training mixture for T5-Large when using (IA)3.
25

---

Published in Transactions on Machine Learning Research (03/2024)
IA3 BFM Objective
Full Model RegMean Objective
Full Model BFM Objective
80
10
10
10-3
Accuracy
10
Error
10-5
40
10-1
30
10
10a
100
101
2
10
10
10
10a
10
10
10
10
Iterations
Iterations
Iterations
Figure Performance of the mergod model and the error solving the lincar system vs. the number of iterations in the conjugate gradient method.
26


Reviews:
[{"review_id": "1BikvqLlpn", "reviewer": "Reviewer_RU1S", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- the paper is globally clear and well written\n- unifying different merging methods under the common expression (4) is interesting\n- the experimental evaluation is comprehensive\n\n\\\n**Weaknesses**\n- the unified view does not account for all the methods benchmarked in Sec. 6\n- more generally, I feel the common expression (4) could have been exploited more, e.g., to devise new model merging techniques by using innovative covariance matrices $C_m$\n- apart from the unifying expression (4), the main (methodological) contribution is only to solve a linear system by Conjugate Gradient rather than by using the closed-form, which is relatively basic and known", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies model merging. The authors first recall several ways of merging models: simple averaging (Sec. 2.1), Fisher merging (Sec. 2.2), RegMean (Sec. 2.3), before unifying them under the weighted average formula (4). The different merging methods are then interpreted as a normalized average of the models projections into the \"task subspace\", i.e., the eigenspace of the covariance matrix chosen. The authors finally leverage the form of Eq. (4) to compute the merged model $\\theta^*$ by Conjugate Gradient, i.e., by solving iteratively the quadratic problem $\\theta^*$ is the solution to, rather than by closed-form. This method enables computing the average model in cases where the closed-form solution is intractable due to a costly matrix inversion. Finally, the authors empirically validate their method with extensive experiments (Sec. 6).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Questions**\n- It is mentioned several times that the initialization of MaTS is important (last paragraph of Sec. 1, Sec. 6.7). However, as far as I understood the problem MaTS is solving is quadratic (see Appendix B), such that the initialization should not have any impact. What am I missing here?\n- What could be an explanation of the fact that the empirical Fisher works better? (last paragraph of Sec. 2.2)\n- In Eq. (4), I do not understand why $C_m$ needs to be a **covariance matrix**. Any symmetric matrix $S$ can be rewritten $S=VV^\\top$, such that you can argue that it the covariance matrix of $V$, but is it needed that $V$ must be an identifiable random variable? What is the interpretation behind this? Also, based on this observation, can the authors think about other matrices $C_m$ that would make sense?\n- What is the interpretation for having \"the gradient of the output constant\" in Sec. 5.1?\n- Figure 1 could be improved in many ways: 1) I assume that the ellipsoids represent the level sets of the losses for two different tasks, but I feel this should be clearly specified. Also, why are models A and B not at the center of the ellipsoids in this case? Is the landscape convex, as suggested by the figure? 2) Do the directions in which the models are transformed match the discussion: BFM seems to translate the models in a direction in which the loss landscape changes quickly. Is it made on purpose on the figure? If yes, this should be recalled in the legend I guess. 3) Based on this 2D representation, which averaging method should we prefer? I assume we want small loss for both tasks, i.e., being as central as possible in both ellipsoids, such BDF seems to perform best here. Overall, even if I understand that Figure 1 is just illustrative, I feel more concrete context should be given.\n\n\n\\\n**Organization**\n- Sec. 3: I would specify that not all model merging methods can be written under the form (4). This would allow to start introducing the methods benchmarked in Sec. 6.\n- Sec.3: I would bring Appendix A.3 in the core text. It is rather short and helps understanding the relationship between the matrices.\n- Sec.4: I would bring Appendix B in the core text. Indeed the CG method has an important role in this work, and deserves to be discussed more, especially the motivations for choosing this precise algorithm.\n- I would remove the last sentence of the first paragraph of p.7 \"K-FAC additionally [...] covariance matrices\". Indeed, this approximation is not used nor relevant yet, which makes it hard to understand. It is discussed again in the next paragraph, which seems enough to me.\n\n\n\\\n**Typos**\n- p. 3: averaging averaging\n- p.3, second display: should be $F_m$ and not $F$ I guess\n- p.3, third display: the $1/F_m$ is not a clear notation when $F_m$ is a matrix, it is better to right properly the inverse, as in Eq. (6)\n- p.4: can therefore **be** viewed as forming\n- p.6: since **it** is the sum of the covariance\n- p.6: that **it** allows for solving linear system\n- p.7: would **cause** the merged model's performance\n- p.7: (IA)$^3$ is not defined not discussed\n- p.9: BFM is not defined yet\n- all paper: sometimes bold subscript $m$, like for $\\Lambda_m$, sometimes no (for $\\theta_m$, $F_m$, etc...)", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PFpuXTK8Gv", "reviewer": "Reviewer_v8xF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The idea of approximating the Fisher matrix by a block-diagonal seems reasonable and novel\n2. The empirical results are great\n3. The discussion of the cost of the method as well as its positioning concerning the existing methods is interesting\n\n**Weaknesses**\n\n1. The discussion about the unification of other methods into one paradigm seems a bit unnecessary\n\nI feel that the whole discussion in Section 3 is not very informative in the context of this paper. In particular, the authors’ contribution doesn’t draw much benefit from this section as in the end they improve the approximation of the Fisher matrix based on the prior work (based only on 1 instance of this unification) and this constitutes their major contribution. I do not see how this unification is linked to the latter contribution and why it is necessary to fit it into some general paradigm. I would suggest removing it and to rather follow the flow 1) the work of (Matena & Raffel, 2022) approximated Fisher with a diagonal matrix which is suboptimal; 2) we propose a block-diagonal approximation and tackle the optimization challenges that it poses. If the authors can propose other contributions that fit the framework of Section 3 and require the solution using CG then it would make sense to keep it. Currently, it is very hard to parse the link between Sections 3, 4 and 5. It is probably enough to jump from Section 2 to Section 5 and include the details from Section 4 there directly as explained above. I feel that this will simplify the manuscript and make it much clearer and easier to follow.  \n\n2. The drop in performance with a higher number of iterations of CG seems counterintuitive and calls for more explanations\nI understand that sometimes early stopping can be beneficial but I feel that in this case varying the number of iterations of CG arbitrarily lacks rigor. It would be surprising to find a paper that tunes the number of iterations with a step size of 10 for SGD or Adam in this way (if I’m mistaken and it is common practice then kindly provide references to it). I think this part really deserves a more in-depth analysis as the whole benefit of the author’s contribution relies on using a numerical optimization scheme that in this case seems not very robust to this crucial hyperparameter (although the reviewer appreciates the authors’ honesty about this).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a new merging approach for neuronal models. The core idea is to build upon a recent method that used the Fisher information matrix for averaging the different models’ parameters. Given that the Fisher information matrix is hard to manipulate in practice, prior work used its diagonal approximation which allows for a closed-form merging solution of parameters. The authors propose to replace the diagonal approximation with a block-diagonal one which is assumed to be more informative. As the latter doesn’t admit a closed-form solution, the authors present it as a solution of a matrix linear equation that is solved using conjugate gradient descent. The empirical results suggest a strong performance of the proposed method compared not only to the method based on the diagonal Fisher approximation but also to other methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Major**\n\n1.\tProvide the errors at the end of the CG optimization for the best-found number of iterations in all cases to understand why increasing the number of iterations leads to a drop in performance. \n\n2.\tFor the unification of other methods to be meaningful, it would be desirable to show that more than 1 instance of algorithm based on intractable closed-form solution using CG is necessary. Otherwise, I would suggest reformatting the paper and skipping this unification altogether.\n\n**Minor**\n\n1.\tI do not understand Figure 1 and what the position of the stars and squares represents. Can authors elaborate on this?\n2.\tp.3: averaging averaging -> remove repetition\n3.\tp.9: as both as both -> remove repetition\n4.\tp.10 in in table 1 -> remove repetition", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "SzM8c6AEvQ", "reviewer": "Reviewer_pzCj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The methods are sensible and principled. \n\nThe topic is interesting and important.\n\nThe results are promising and impressive at times.\n\nThe paper writing has major issues, and I did not really understand this paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper casts few neural network parameter averaging methods as unified covariance-averaging, and presents an efficient conjugate-gradient method to find the optimal parameter averagings.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "This paper studies how to merge multiple neural networks into one combined one by averaging the parameters in different ways. The paper considers three merging approaches (average, Fisher average, activation matching), and casts them through unified covariance-weighted averaging, and taking its SVD form. The paper presents a conjugate-gradient method to compute the mergings in practise. The paper seems to have two novel contributions: unified merging math, and the CG. \n\nThe methods in this paper are all sensible and principled, but quite textbook. The paper is written in a clear and intuitive style. \n\nThe paper frames its discussion against “task subspaces”. I do not agree with this. I don’t think there are any “tasks” here, nor do I think the SVD represents any tasks. The paper simply merges different DNN weights, and this does not need to have anything to do with tasks. The SVD represents the parameter subspaces around different modes of the same posterior, or across slightly varying posteriors. These posteriors can come from simply re-trainings, or from having different datasets, but they can’t come from having different tasks since then the parameter vectors would not be same size (which they need to be). I guess one could have multiple eg. classification tasks as multiple single-task networks, which are merged, but this makes little sense. Can you discuss this?\n\nThe paper does not acknowledge that the SVD inverse simplifies into one where the inverse is only over eigenvalues. The paper motivates their CG optimiser due to not being able to do the inverse, so I wonder if the motivation for CG dissolves? \n\nThe paper needs to include simple ensemble as a baseline, ie. using $f(x) = 1/M \\sum_m f_m(x)$, or something similar. The paper also needs to include the individual models $\\theta_m$ as baselines as well.\n\nI don’t think Fig1 makes much sense. The markers are placed all over the place, and I can’t see any rhyme or reason why the stars/diamonds are where they are, or what the lines mean. The figure also doesn’t even show any mergings: where are the averaged points? This figure needs to be reworked.\n\nI did not understand what the MATS actually is. The paper talks about various merging things, and CG and blocks, etc. What is MATS?\n\nI did not understand what the paper means by tasks. The parameter merging clearly requires the parameters to be of same size, and have same meaning. Hence, I think the merging only applies to the “body” of the network, and one can’t merge the different heads, where one would expect each task to have one head. Yet, the text always talks about tasks, even though we clearly can’t merge anything task-specific. Err… so what’s going on? \n\nTo clarify my stance: I’m thinking here of eg. BERT/CLIP style foundation models where we have a massive embedding network that produces high-quality embeddings of the inputs (=body), and then we have smaller networks that map these embeddings into different predictions needed by different tasks (=heads). We want to train the heads; and either keep the body frozen or perhaps only slightly re-optimise it for each head. Can you clarify what kind of setting you are thinking in this paper? Are you merging the body or the heads, or something else?\n\nI’m also confused what are the tasks in the first place. Are we doing classification? Do we have multiple classes or just one many times? The experiments bring in some kind of task vectors, without any math, or any explicit connection to the earlier math. What are these and how do they go into the method? \n\nAre different models re-trainings, or different datasets, or different classes, or different labellings of same classes, etc..? I also don’t understand what the different models theta_m are or where do they come from.\n\nAlso, why do we want to merge models in the first place? If we have 8 fine-tuned CLIP models, surely we want to retain those and not mess up them into one hybrid model. \n\nIf the multitask seems to always be superior, why do we want to merge? Why not use the multitask model instead? \n\nI did not understand 6.4. So we have an RTE-optimised method, which gets 79.8, and then we have intermediate-optimised method that is combined to it to get 83.0. Ok, but where is the intermediate-optimised method itself? Surely you want to have one method for RTE-opt, one for int-opt→RTE-opt, and finally a merged one. There also seems to be all kinds of finetunings going on. It’s difficult to see in the results what values are caused by what, since lots of stuff is happening behind the curtains.  \n\nMinor comments\n\n- The paper mixes $m$ and $\\bf{m}$ and $\\bf{\\theta}$ and $\\theta$ randomly, which makes the manuscript look unfinished. Please have all math consistent wrt boldfacing and upper/lowercases.\n- The paper implicitly assumes that all models have same parameters with same meaning: otherwise eq 1 wouldn’t make any sense. This should be made explicit: this paper considers only merging different optimisations of the same network. This limits the paper’s scope.\n- Eq 1 should cite and discuss SWA(G) (Wilson et al)\n- eq (unnumbered, after eq 2 [please number all equations]) has F while it should have F_m.\n- What is $1/\\sum_m F_m$? Is this an elementwise inverse, or matrix inverse? Is this before or after the diagonal approximation? Later these inverses seem to be matrix ones, but here it looks like elementwise. This is likely a mistake. Please clarify.\n- It’s a bit strange that eq 5 does not acknowledge the obvious simplification $(QLQ^T)^{-1}$ = $QL^{-1}Q^T$. Surely you want to do this?\n- Please clarify what you mean by finetuning. It is an ambiguous term.\n- For vision, what was the CLIP trained against? What did you finetune? Why can’t you just train all 8 label sets together? Surely it’s way more expensive to train 8 tasks separately, than just one combined. What is the test setting here? Do you mix all images from all 8 tasks? How do you handle resolutions? How many classes are there in the 8? I’m quite confused what the experimental setting here is.\n- What does “base model” mean in the introduction? I’m a bit confused of the motivation of this work. Let’s take the Stable Diffusion XL example. There are surely some specialised models that eg. only draw cartoon figures and nothing else. There are also surely some models that specialise in some artistic style. If you average these models, you surely just get gibberish. What are you then trying to achieve? Surely the specialised models need to have considerable overlap in their domains. Can you clarify your setting?\n- I struggled to understand the experiments. Could you explain them in higher level: what are you trying to achieve? What kind of models do you want to merge, and why? Illustrations or math would be nice (eg. what are the sizes, dimensions, outputs).\n- What are the assumptions for your models to be merged?", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree, after authors rebuttal and updates, that the method and the experiments are interesting and that the paper deserves to be published. The changes on the paper since the submission, and the reply to the reviews, address well their comments.

The only remaining concern that is shared by two reviewers and the AE is about the name of the method and its discussion.  "Task Subspaces" can indeed refer to a lot of things (input/output/feature/parameter spaces). Even if it is defined in the introduction now, it is better to have a precise name such as "Task Parameter Subspaces" both in the title and the rest of the paper. This would make the positioning of the contribution of the paper more clear for the wide TMLR audience. 

Another minor suggestion that the authors can choose to take into account or not for the final version is to add "Fisher" at the beginning of the title, again to have a more descriptive title of the contributions.

The authors themselves state in their response that they would be OK with such a change in name of the methods so the AE proposes to accept the paper with minor revision.
Decision: Accept (poster)