Model Merging by Uncertainty-Based Gradient Matching
Authors: Nico Daheim
Thomas Möllenhoff
Edoardo Ponti
Iryna Gurevych
Mohammad Emtiyaz Khan
Venue: ICLR 2024 poster (2024)
Keywords: Model Merging
Gradient Matching
Language Modeling
Model Editing
Transfer Learning
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: We connect model merging to gradient matching, show that uncertainty-based reduction of gradient mismatch can improve the performance of the merged model, and connections to several existing methods.

Abstract: Models trained on different datasets can be merged by a weighted-averaging of their parameters, but why does it work and when can it fail? Here, we connect the inaccuracy of weighted-averaging to mismatches in the gradients and propose a new uncertainty-based scheme to improve the performance by reducing the mismatch. The connection also reveals implicit assumptions in other schemes such as averaging, task arithmetic, and Fisher-weighted averaging. Our new method gives consistent improvements for large language models and vision transformers, both in terms of performance and robustness to hyperparameters.

Published as a conference paper at ICLR 2024
MODEL MERGING BY UNCERTAINTY-BASHI GRADIENT MATCHING
Níco Daheim', Thomas Möllenhoff", Edoordo M. Ponti Iryna Gurevych', Mohammad Emtiyaz Khan
1 UUiquitous Knowledge Processing Lab (UKP Lab)
Department of Computer Science and Hessian Center for Al (hessian.Al)
Technical University of Darmstadt
RIKEN Center for Advanced Intelligence Project, Tokyo, Japan Uniierrity of Edinburgh
www. ukp. tu www-darmstad..comm
ABSTRACT
Models trained on different datasets can be merged by a weighted-aweraging of their parameters, but why does it work and when can fail? Here, we connect the inaccuracy of weiightd--veerging to mismatches in the gradients and propose a new uncertainty-based scheme to improve the performance by reducing the mis- match. The connection also reveals implicit assumptions in other schemes such as averaging, task arithmetic, and Fisher- sesher-weighted averaging. Our new method gives consistent improvements for large language models and vision transformers, both in terms of performance and robustness to hyperparameters. Code available here.
INTRODUCTION
Merging models through weighted averaging of their parameters has recently found many applica- tions in deep learning. For example, averaging checkpoints generated during various training runs can improve out-of-distribution generalization (Izmailov et al. 2018 Wortsman et al.. 2022b Gao
et al. 2022 inter alia), while averaging models trained on different datasets can borrow knowledge
from "donor tasks" (Matena & Raffel, 2022) and enforce specific fine- men-grained behaviors in models
[Iharcoo et al. 2023 Daheim et al. (002). The latter is particularly attractive for post-hoc "editing" of large pretrained models without retraining, for instance, remove toxicity from a large language model (L..M). Simple weighted-averagiee appears to tackle many difficult knowledge transfer and adaptation problems that machine learning methods have struggled to solve the past.
The reasons behind the effectiveness of weighted-averaginn methods are not well understood. The diversity applications has led to a large number of averaging schemes, including arithmetic mean (Wortsman et al.) 2022b a). linear interpolation (Ilharco et al. 2023; Ortiz-Jimenez
al.) 2023
Yadav et al.] 2023). y. or individual parameter weighing (Matena & Raffel [022][Dheeee et al. 2023)
A prominent hypothesis, known 'linear mode connectivity'. is that when the models land in relatively few low-loss basins their interpolation again lies them (Frankle et al. 2020 Neyshabur al. (2020) Wortsman al.) [2022a] Ainsworth al. (002)). However, this does not directly tell us why and when one merging scheme should be preferred over the others, nor does it give any hints on how improve them. Ideally, we would like to understand the effect of averaging schemes on the accuracy of the merged model and use it to design better merging methods.
In this paper, we improve model merging by proposing an uncertainty-baat gradient matching method. We make two contributions: we first connect the inaccuracy of weighted-averaginn to mismatches in the gradients and then improve second-order approximation; see an illustration in Fig. [ Our new method uses (cheap) Hessian estimates to merge models which scales well to large Transformers (Vaswani et al. 2017). We use the method to reveal several assumptions implicit in existing model merging schemes like averaging, task arithmetic (Ilharco et al. 2023), and weighted merging (Matena Raffel 2022); see Table I Finally, show connections of our method to Bayesian approaches and discuss how can leverage them to further improve model merging. Empirical results on LLMs and ViTs show consistent improvements, both in terms of performance and robustness hyperparameters.
performance by reducing the mismatch with a

---

Published as a conference paper at ICLR 2024
uuges
Difference in test error
Task Arithmetic
Ours
4
2
nacaged
02
A
(lr(Oun))) (l. 88
H.A.
0.2
0.4
Gradient mismatch
Gradienn Mieronch
The left panel illustrates our approach. We connect the emor A of the merged model merged to the gradient mismatch over losses & and propose new method that reduces the mismatch by using the Hessian He and error A of the individual models 0,. The right panel shows example of adding datasets to RoBERTa trained on IMDB. We clearly see that reducing mismatch reduces test error of task arithmetic (a, 1). We consider 5 datasets, indicated by number the markers.
Figure 1:
MODEL MERGING BY PARAMETER AVERAGING
We consider merging multiple models that share the same architecture but are trained on different datasets, for example, by fine-tuning large pretrained model. We denote each of the T >> models by its parameter vector 0, E IR.. Throughout this section, we will use an LLM, denoted by OLLM. but the results straightforwardly apply other pretrained models. Given OLLM and different O,, our goal is to understand the inaccuracies in existing avraaater-averaging methods and improve them.
We focus on the following simple weighted-averageee the merged model obtained with scaling matrices S₁ E IR dxd for = 0,1...., Since the dimension d is often large, simple choices of S, are used in practice. The simplest is the arithmetic mean (AM) or its weighted version (WAM; Wortsman et al.) a)22b
scheme: 0 So l + ΣT= S, θr, where
T
T
0 AM = T
Σ O 1=1
B WAM = aG0LLM Σ 0,00.
(1)
where 04 2 0. For large models, different parameters have different scaling and is hetter to take this into account, for example, by using the Fisher matrix F,:
T
PAA Σ S,0,, where S, F.FFFF, with F
Σ F,F,, for all t 22 1,
giving rise 'Fisher Averaging' (FA). We could similarly include Sn by using the Fisher F, of the LLM. In practice, to reduce the computation cost, we may only use the diagonal of the Fisher estimated in an online fashion (Matena & Raffel 2022 This is similar strategies in continual learning (Kirkpatrick et al. 2017) where the choice of Fisher is justified through Bayesian updating Huszár (0018). However, such connections are not yet explored or exploited for model merging.
Using Fisher should improve things bit but the extent of improvement is unclear. A recent work by [iin et al., (2023) uses insights from linear models justify some of these choices, but such justification may not hold for nonlinear models. In general, it is also not clear how Fisher-averaging takes care of the commonalities between the fine-tuning 0, of the L.LM 0LIM- Should we include Fo or not, and how should it be combined with the other F, information in the models? The current practice is to tune a on a validation set which one way to make up for the errors, but this quickly become expensive T increases.
avoid double counting of
Recently, Ilharco et al. 2023 proposed to subtract the contribution lLLM with the follow- ing simple "task arithmetic (TA): 0IA OLLM Σi= from 0, should reduce double-counting the information, but adding Fisher-style scaling in this scheme be bit tricky. A recent work by Daheim al. (2023) proposes O
/f (0, OLLM). Subtracting 0LLM
2

---

Published as a conference paper at ICLR 2024
F F (FLLMLLLLM Σi= - for F FLLM + Ei= F, but F, is calculated at (0, @LLM) which lacks theoretical justification and using a sealing matrix in front of LLLM also departs from other approaches. TIES- Merringg (Yadav et al. (2023) instead multiplies binary matrix to the 'A update with the goal to reduce interference when adding multiple models. TA also allows for Oz < 0 remove the contribution of old knowledge, which differs from other schemes and it is not clear if these schemes can also use negative weights.
9LLM))
To summarize, we want understand the effect of various choices made in parameter-averaging methods. That is, we want to know the following: (I) how to choose the scaling matrices; (2) what the effect of these choices on the accuracy of the merged model; and finally (3) how to obtain new method that reduces the inaccuracies in previous approaches. In what follows, we answer such questions by proposing a new connection with gradient mismatch and a new method inspired by it.
MODEL MERGING AND CONNECTIONS TO GRADIENT MISMATCHES
To understand the inaccuracies of parameter averaging, we introduce the idea of target model: is the model that model merging methods want to estimate. Here is an example: consider two models 0 and 0g trained on two datasets D and D2, respectively, for example, as follows,
0 arg min (0) 1111,
0₂ arg min ( + 101
(3)
0
Here, the loss functions on D, and D2 are denoted by & (0) and (0) respectively and the regular- izer is an L2 regularizer (what follows also holds for other explicit regularizers, also implicit ones). The target model in this case could be a model 0 01+2 that trained jointly on the two datasets:
1-2 arg min ( + (0) + 111011
a
We have used scalars a1 and 0r which reflect relative weighting of each loss function. We will now connect gradient mismatch to the error between the target 01+2 and parameter-average a1 + 0202. The approach is general and applies to different types of targets and averages. This will be explored extensively in the rest of the paper.
We start with the first-order stationarity conditions of the models in Eqs. 3 and 4
0, (0,,,,
02 02=-a (02,),
01+2 = - an Vl (0)+2) -
(5) express
which is obtained by setting the gradient of their objectives to zero. Using these, 01+2 in terms of a 01 + a222 and quantify the error made. To do so, we multiply the first and second equations above by an and a2 respectively, and add them together. Then, subtract the resultant from the third equation to get the following expression:
01+2 - (a, + 0202)) a1 (11(())2) (2, 02 -
(6)
A, Error of the merged model
Grudient mismatch for 0 &
Gradient mismatch for 0g on fn
The left-hand side the error 0142 - 0, + Or202) which is equal to the weighted-sum of the two mrastcht-mismatch terms, each measured on the individual losses &, (0,)) and (2 respectively. The expression shows that the individual models are already close (in terms of their gradients) the target model, then parameter averaging should be reasonably accurate. It also tells us that there is room for improvement and mismatch reduction may lead to better schemes.
The method above can be used to analyze errors of generic parameter- averraing schemes. For instance, for data removal tasks, say to target 02 by an operation 0142 aj0,, we can simply rearrange terms in Eq.
6 to express O2 in terms of 0 22 and 0;:
- (0142 -
/a2 = (a1/az) [ - (0:)) + ((2(0122) - V[2(0))
Generic loss functions can be used, for example, non-differentiable loss function can be handled through smoothing techniques. Any convex regularizer can be used in which case the error is mea- sured the dual space of the regularizer. The approach also works in the absence of a regularizer. Test accuracy can also be analyzed in the same fashion. For example, given test loss (rea(() and weighted-average average 0, we can write: & (01+2) &eau (0) xa Vi (0) (0 θ). Large gradient mismatch therefore is expected correlate with large differences in test performance.
3

---

Published as a conference paper at ICLR 2024
Sources of errors can be analyzed, too. For example, when the test data is more correlated to D1, then model merging would effective if gradient mismatch due to 01 is also small. This similar to linear mode connectivity: when both the target and merged models lie in low-loss basins, expect gradient mismatch to be low due to flatness. However, gradient-mismane this and is more general and constructive, that is, it allows us to improve models by actively reducing the mismatch. In what follows, we will use gradient mismatch as a principle to not only study the inaccuracies of various averaging schemes but also to design better ones.
does not require
3.1 ANALYZING THE INACCURACY OF TASK ARITHMETIC ON LARGE LANGUAGE MODELS
We will demonstrate the use of the arithmetic' (TA) ([haacco et al. 2023). Task arithmetic [lharco al. (2023) uses OTA OLLM + E. 00 - LLLM) to merge models. The model L is used to initialize the training on individual datasets D, to get 0g for 1, 2, ...,, T, as well as the training of the target 0.rr- In this work, OLIM denotes LLM trained on large dataset Large but similar analysis can be done for other pretrained models. For example, OLLM can trained by using:
principle by analyzing the inaccuracy of 'task
OLLM = min (0) + 1811011, where (1)
Σ (())): iEDurg
where (,) denotes the loss on the i'th example. Our goal is to merge models 0, that are each finetuned on of the T different datasets D,. We assume the following fine-tuning procedure,
0, arg min (0) + 2 1 0L.MI
(8)
where | = HoO is the Mahalanobis distance with a scaling matrix Ho which controls how different A is from OLIM. We will discuss how to set Hg later. The derivation can be easily adopted to other fine-tuning procedures as long as we can express the dependence on ÔLLM explicitly.
As before, reasonable choice of the target model is the one obtained by fine-tuning using a similar procedure as Eq. but on all datasets Dj at once,
Σ (0)( + 1 OLLMIIA t=1
1TT arg min
We use the weighting with a, to align the target model to the weighting used in the merging scheme, but this is not required and other targets can be used. Following the same derivation as Eq. 6 we can quantify the error between 01. and θTA (a full derivation is given in App. A.1
T
0 01:-- (0LIM + Σ c., OLLM)) = =- Σ Ho [Vi((0).)) - V(,(
(10)
Gradient mismatch for
PT
The derivation can be used to understand the implicit assumptions made in task arithmetic. The increments 0, ÕLIM arise above due to the quadratic regularizer ||0 0 H used in Eqs. B a Using the increments avoids double counting. More importantly, the error between the target 0urr and 0TA is attributed to gradient mismatch between Q, and 0r. The expression suggests that by reducing the gradient mismatch we may be able to improve task arithmetic. We will now show that simple method that uses Taylor's approximation to reduce the gradient mismatch justifies combining TA with a Fisher-like weighting schemes.
3.2 NEW METHOD TO REDUCE THE GRADIENT MISMATCH
We now derive a new parameter-averaginn Explicit minimization of the mismatch is non-trivial because (f,(0.:)) depends non-linearly on 01.r. However, we can get rid of the term by using first-order Taylor approximation.
method by reducing the gradient mismatch in Eq. 10
((,01))) R Vi, (0.) + H((01.7-8 -0)
(11)
where H,- (0,) is the Hessian of the loss E, at O.. Using this in Eq. 10 and after some rearrangement, we get the following merging scheme (a full derivation is given in App. A.2).
T
0T OLIM + a (H Hoo+) (0, OLLM),
Σ
(12)
4

---

Published as a conference paper at ICLR 2024
Ot
Hg
H,
Drawback
Arithmetic Mean (AM) WWorsmman et al. 2022b
1/T
1
0
Uniform weighting
Task Arithmetic (TA) ([harco et al. 2023)
I
O
No preconditioning
Fisher averaging (FA) Maaeena Katlel 2022
0
diag(F)
LLLM is ignored
FA1 (Daheim et al. 2023
0
diag(F, (F))
Fisher lacks justification
Sparse Finetuning (SF) (Ansell et al.
2022
0
Binary mask
H, is O-1 matrix
TIES merging Yadav at. 2023
0
Binary mask
H is 0-1 matrix
Table 1: Our new connection reveals implicit assumptions made in existing weighted-averaging averaging schemes: AM uses uniform weighting while TA lacks preconditioning matrices (because H2 = 0). We expect high gradient mismatch for both. Both Fisher averaging methods FA and FAI use preconditioning but ignore the dependence of 0 on θ (because Ho = 0).
where H - Ho E H, accumulates all Hessians and Ho+r = Ho + H, the Hessian plus a regularization matrix. The new merging scheme adds preconditioners H Hoo+ H'H task arithmetic. The preconditioners depend on the Hessians H₃, which is similar to the fisher-weighting scheme, but here the choice naturally emerges as a consequence of the gradient-mismatee reduction. Nevertheless we can replace H, by the diagonal Fisher F, of 0,, which is often easier to compute and also easier numerically because positive-definitnnss is ensured. The matrix Ho can be set in similar way, for example, we can use the diagonal Hessian/Fisher of Eq. 7 at OLLM- We discuss these approximations further at the end of Sec. 3.3
3.2.1 RELATIONSHIP TO EXISTING METHODS
Choosing different setting of ag, Hg. and H,, can recover many existing schemes as special cases of Eq. 12 as summarized in Table I This helps us to understand not only their inaccuracies but also their implicit assumptions. AM and TA can be seen as special cases where the preconditioner H, = 0. This implies that the gradient mismatch term in Eq. 10 is left as is and the error will be high whenever there is a high gradient mismatch. In contrast, FA and FA1 can seen as special cases where Ho = 0 which implies that the quadratic regularizer in Eqs. 8 and 9 is not used and therefore the dependence of 02 on ÔLLM ignored. light of Eq. Eq [2 the Fisher at O, OLIM used in FAI (Daheim et al. 2023) appears to be odd choice. We note that we can compensate for errors in FA by adding an additional SDLLIM, similarly to Daheim et al. (2023), but the choice of Sy is nontrivial: Eq. 12 suggests it be S, I Σt= a,H Ho++- Such choice would compensate for double-counting of information from θITM but it is difficult to come up with this choice without special case with Ho = and H, set to a binary sparse mask whose entries are only for the parameters with the highest change. This step also added in TIES- Yadav al. (2023). which uses trimming and elect-sign operations, but direct effect of these operations on gradient-mismatee reduction currently unknown. Overall, our approach provides a way to understand the effect of such choices and gives way to improve them by reducing the gradient mismatch.
the analysis present here. Yearse-Finettuing (SF) (Anell et al. 2022
be
3.2.2 A NEW METHOD FOR DATA REMOVAL.
The principle of gradient matching can be applied to other merging tasks and schemes. For instance, consider removal of task dataset from the LLM which arises, for example, when trying to reduce toxie language generation. For such case, we could fine-tune a model on (hopefully the same) toxic dataset and try 'subtract' its contribution from the L.M. This is expected to be much cheaper than retraining on cleaned data. Formally, we want to remove a dataset D, C Dargee Thereby we aim for a target model O₂ trained using Eq. but after removing D; from DLargo Let us denote the loss by Exe- We can then fine-tune tune a model O by using Eq. 8, and do task arithmetic: 0, OLLM - (0, 100--0um 0LLM) (Ilharco et al. (023). As shown in App. A.4 we can use gradient matching to understand and improve this method. We get the following improvement,
θ = 0LLM - ar H H Ho - OLLM),
(13)
where H = 2 (LLM) + 5I is the Hessian of Eq. 7 at 9LLM but without D,. The expression includes preconditioner which is expected to improve the performance of task arithmetic. Intrigu- ingly, when applied to data removal in a linear model, this update recovers the celebrated influence function (Jaeckel] 1972 Cook 1977 Koh & Liang 2017). We formalize this as follows.
5

---

Published as a conference paper at ICLR 2024
Theorem For linear regression models with loss = 2 ly, 2
x1 where y is the outpu
vector and X, is the feature matrix, the update in Eq. [3] with a = reduces to the well-known
expression of influence by Cook 1977 Eq. 5). which is shown below:
0, OLIM Hip X (X.0IM - y.).
(14)
A proof of the result is given in App. A.5 Our result also applies lQ generic nonlinear models, where the step (3) can be seen as a Newton-like step in a direction Hi, +o++(0; LLIM). We note that there are several ways to rearrange the gradient mismatch term which give rise to different kinds of approximation. It is up to the designer to choose the preconditioner which goes well in their setup.
The derivation in App. A.4 shows an example the context of task removal (see Eq. 25
3.3 GRADIENT MISMATCH REDUCTION AS UNCERTAINTY ESTIMATION
Both the gradient-mismatel connection and the method are closely related to uncertainty timation via approximate Bayesian methods. We show that Eq. [0] is equivalent to a maximum
posteriori (MAP) estimate of the posterior over all data Di.r while Eq. T2 is the posterior approximation obtained with Laplace's ' method (Laplace 1774 Tierney & Kadane 1986
but for
MacKay [992). Based on these, we discuss some possible future directions for improvements.
We start by defining the posteriors. Assuming p(0) = =1L00L H, to be the Gaussian prior and p(Dr, 0) x --.(()) to be valid likelihood function, can define weighted-posseriio pa (0)D I:T) over all datasets, shown below, along with an approximation obtained by Laplace's method,
T Pr. (O) & p(0) II
T
200 so p()) II 2a.lle-00 lii x qa (.)
(15)
Here, we use a second-order approximation at O, to get (0) (0.) + 0 0, I The term & (0.) is an irrelevant constant and we get the approximation ga (i)) The result below shows that the merged model is the MAP estimate corresponding the approximate posterior.
Theorem 2 The gradient mismatch equation in Eq. i0is is the stationarity condition of a MAP prob- lem written in terms of posterior p(D.))) (the equation on the left), while the merged model 0 Ovr in Eq. is the MAP estimate of the Laplace approximation (equation on the right)
D II
(O
p
11T arg max p(0)
0I. = arg max qa (0(Di:) (16)
(
0
A detailed proof is given in App. A.3 The result relates the gradient-mismatch approach to the pos- terior distribution and its approximation. The first equation expresses model merging as merging of posteriors p.) that are computed on different datasets. With Bayesian approach, an exact lution can be recovered even when training on separate datasets. This is an instance of the Bayesian committee machine (Tresp) 2000) or Bayesian data fusion (Mutambara) 1998 Durrant- Whyte 2001 Wu et al. 2022) which are extensively used for Gaussian processes (Deisenroth & Ng, (015) and which should also be useful when using Neural Tangent Kernel for model merging (Ortiz-immene) et al. 2023). The second equation connects existing methods to Gaussian approximation obtained using Laplace's 's method. Table therefore suggests that these methods make crude approximations to uncertainty estimates where either the likelihood or the prior in ga is ignored.
The gradient mismatch term in Eq. 10 arises due to the ratio p(0 Dr)./(()) To understand this, consider the simple case of linear regression. Suppose learn two separate linear models with loss function (0) = 2 lly, - X,0 |2, The gradient and Hessian are (l( = X, [x,0 - y,) and H, X respectively. Therefore, the gradient mismatch term can be written as,
(D)) p
vi,(0r)) - vi, = X (X X,0,) = H, 0,) = V log
0 1=-1..T
For linear models, Pa (0 = qa D,) and therefore Taylor's approximation in Eq. is exact. The equation matches RegMean (Jin et al. 2023 Eq. 1) who use this objective to merge linear parts of transformer. Our approach extends such efforts to nonlinear problems.

---

Published as a conference paper at ICLR 2024
80
Ours
Accuracy
60
94
shro-shot
40
92
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Ogg
Figure 2: Leftt: We merge models trained on 8 image classification tasks with pretrained ViT and vary at. Our method performs similarly TA for smaller but significantly better for higher an improving over the best a₂ for TA. Right: We add four sentiment analysis tasks to RoBERTa trained on IMDB, Our merging function dominates TA and requires no tuning of scaling factors. We plot the average over individual dataset accuracies.
IMDB
Yelp
RT
SST2
Amazon
Avg.
True Avg.
Parametrization
Accuracy (1)
TA (o: I) Ours
90.5
95.6
86.4
91.6
94.9
91.8
94.7
94.7 (14.2)
97.3 (17)
90.2
(13.8)
93.7 (12.1)
96.7 (118)
94.5 (12.7)
96.6 (11.9)
Table 2: Reducing gradient mismatch in Eq. 10 when scaling is not tuned (o: 1) is crucial for merging, here outlined for adding four sentiment analysis tasks RoBERTa trained on IMDB. Avg.: average over individual dataset accuracies. True Avg.: accuracy calculated over all predictions.
The Bayesian connection also gives direct ways to improve model merging and also reduce the computational burden. For example, one way to improve would be to take a few optimization steps aiming for the MAP estimate of the exact posterior, and then use the current iterate for the Taylor's approximation in Eq. TO Solutions obtained this way will provably get better as the number of steps are increased. This in contrast with other approaches, for example, by Ortiz-Jimenez al. (2023) who propose to train in the linearized tangent space which may not always converge to the right solution. Another way improve is use better posterior approximation, for example, using variational inference (Graves) 2011 Blundell et al. 2015 Osawa et al., 2019) which is known to yield more global approximation Opper & Archambeau 2009p. Here, we focus on improving merging without retraining and leave the iterative optimization as future work.
The Bayesian view also connects to similar efforts in continual learning to avoid catastrophic for- getting (Kirkpatrick et al., (017) where a Bayesian motivation is used to justify the choice of Fisher- based regularizer (Huszár) 2018). Our contribution essentially gives an extension of such ideas to model merging. Our approach is also connected to Bhooledge-Adaptaiie
priors (Khan & Swaroop) 202I] where a variety of adaptation tasks are solved by gradient reconstruction. The connection also justifies the choice of diagonal Fisher in place of the Hessian, which essentially forms Generalized Gauss-Newton approximation (Schraudolph) 2002; Pascanu & Bengio 2013) Martens] 2020) it. In experiments, we i summed over examples the data. naive implementation would require an additional pass over subset of) the training data and additional gradient computations, but it is also possible to build the estimate of Fisher in an online fashion (Schwarz et al. 2018) or even reuse the quantities already computed within Adam-styl style optimizers (Kingma & Bal 2013) which are accurate for small minibatch sizes (Khan et al.) 2018 Thm. 1). With this, no additional overhead is incurred while keeping the data private. contrast, tuning scaling factors on a validation set requires additional data and tuning, and could be infeasible for large T.
Mlonte-Carlo estimator Σ, (0) of the diagonal Fisher where
EXPERIMENTS & RESULTS
We first show the relationship between gradient mismatch and test eтrо for language models in Sec. 4.1 Then, consider the setting of task addition, and add tasks to pretrained ViT (Doso- vitskiy et al.) (02T) (Sec. 4.2) and LLM (Sec. 4.3). Finally, we consider data removal and remove toxicity and hallucinations from language models (Sec. 4.4). In all experiments, approximate

---

Published as a conference paper at ICLR 2024
IMDB
Yelp
RT
SST2 Accuracy (1)
Amazon
Avg.
True Avg.
Parametrization
All-data
94.8
97.6
91.2
94.7
96.9
95.0
96.8
Averaging
94.4
97.0
89.1
93.6
96.2
94.1
96.1
Fisher Averaging wg))
94.5
97.0
89.6
93.9
96.1
94.3
96.1
Fisher Averaging aum.)
94.8
97.2
89.9
93.1
96.6
94.3
96.5
RegMean
94.7
97.3
90.0
94.0
96.5
94.5
96.4
TIES-Merging
94.0
97.3
89.6
93.7
96.6
94.2
96.5
Task Arithmetic (o, 1)
90.5
95.6
86.4
91.6
94.9
91.8
94.7
Task Arithmetic (tuned o
94.3
97.2
89.6
94.5
96.4
94.4
96.3
Ours (F
94.4 (10.3)
97.2
90.2
(10.6)
94,6 (10L1)
96.3
94.5
(10.3)
96.3
Ours (F (Fam)
94.7 (10.4)
97.3
(10.3)
90.2 (10.6)
93.7 (10.))
96.7
(10.))
94.5
(10.3)
96.6 (10.3)
Table 3: We merge four tasks with RoBERTa trained on IMDB. Our merging function shows how reducing gradient mismatch improves performance over previously proposed functions. Optimizing the scaling factors of TA on test data (") can not recover the performance of our method, indicating that scalar weighting insufficient. sum denotes summing squared gradients, Faug. averaging. Changes in brackets are wrt. TA (tuned o:)).
Hessians using the squared gradient upproximation of the Fisher unless otherwise stated. All models trained using Adam W (Loshchilov & Hutter 2019) or a modified version of Adam (Kingma & Ba 2015) with decoupled quadratic penalty. Further experimental details can be found in App.
4.1 GRADIENT MISMATCH & TEST PERFORMANCE
We measure the mismatch of gradients between a model trained on all data and merged with task arithmetic and our method Eq. 12 the experiment of Sec. 4.3 by calculating the norm of the difference of their gradients on the training data. Fig. shows that there is a clear correlation between the test error and gradient mismatch. Reducing the mismatch leads models with less prediction error, confirming our intuition that it plays key role in successful model merging. Similarly, Table shows that accounting for the gradient mismatch in Eq. 10 provides large improvements.
4.2 ADDING TASKS TO PRETRAINED VISION TRANSFORMERS
We use a pretrained ViT for image classification and add eight datasets to it: Cars (Krause et al. 2013), DTD (Cimpoi et al. 2014), EuroSAT (Helber et al. 2018), GTSRB (Houben et al. 2013). MNIST [eCCun] 1998), RESISC45 (Cheng et al. 2017). SUN397 (Xiao et al. 2010), and SVHN Yuval 2011), replicating the method and datasets used in Ilbarcoo et al. (2023). We use the identity matrix to approximate the Hessian of the pretrained VIT, because the training data is not public, but might also use squared gradients
trained by fine-tuning the VIT. The results are outlined in the leftmost panel of Fig. 2 Our proposed merging function is much more robust to the choice of scaling factors. For larger factors, task arithmetic even falls below the zero-shot baseline and even though performance also drops for our method, it stays well above this baseline and improves slightly over the best a found for task arithmetic.
similarly distributed data. All task models
4.3 SENTIMENT CLASSIFICATION IN NLP
We repeat similar experiment using RoBERTa Liu et 2019 which follows the BERT archi- tecture (Devlin et al. 2019) and is an encoder-only language model. We first train on IMDB (Maas al. 2011) (arbitrarily chosen, and any other of the datasets would work, too) obtain the required task-specific classification layer. We approximate the Hessian of this model using the squared gra- dients on the training data for the quadratic regularizer. We then use this model to initialize other models which we train on the polarity version of the Amazon (Zhang al. 2015). Rotten Tomatoes (Pagg& & Lee 2005), SST2 (Socher et al. 2013), and Yelp (Zhang et al. 2015) datasets respectively.
Table Jsows that our new method gets closer to the "all-data" target model 0 than other merging functions, like averaging, or FA, and is competitive to others, like TEES-merging, where we keep the RegMean. Furthermore, when we tune scaling factors on the test set for TA and not at all for our method which corresponds a 1. Fig. (right) shows a plot scaling factors where method dominates TA which
top-20% of parameters,
proposed merging improves over TA
8

---

Published as a conference paper at ICLR 2024
Model
0
Toxicity
Fluency PPL(.)
Model
Fluency BLEU (t)
Hallucination %
Av0.Avg.
Num. Toxic
Critic (L)
1-Q² (1)
GPT2177
011M
11.2
15.4 %
24.9
Flan
17.3
27.5
11.7
TA
9.8
13.1 %
30.3
18.2
13.8
7.4
ours
9.6 (10.))
12.8 % (L0.3)
26.9 (L3.4)
18,2
12.8 (11.0)
7.0 (10.4)
GPT-J1.38
011M
11.9
16.6 %
12.6
FlanT5
18,4
31.5
12.8
TA
10.7
14.5 %
12.7
18.6
11.8
7.7
10.2 (10.5)
14.0 % (.0.5)
12.8 (10.0)
18.0 (10.6)
8.8 (13.0)
5.0 (12.7)
Table 4: Reducing gradient mismatch also improves removal of undesirable behaviour from LLMs.
also falls below the zero-shot baseline of the IMDB model. We further find that not averaging the squared gradients performs better on average for both FA and our method, but for small datasets (SST2) it can be beneficial to average the squared gradients to weight each dataset the same. An important choice in our experiments for FA was how lower-bound or add small to the Fishers to prevent numerical instability. For instance, for F aavg we have found adding small delta (e.g the order of 10-19) to perform multiple points better than clipping to a larger value, such as 10 To summarize: 1) reducing gradient mismatch improves performance and 2) is crucial for correct scaling to overcome the need for manual tuning of scales. Furthermore, 3) merging with increments of 0: fLIM instead of just 0, gives slight improvements and 4) so does scaling by Fisher.
4.4 EDITING LANGUAGE GENERATION MODELS BY REMOVING DATA
removing data with halluci- nations from dialogue models to improve their faithfulness, and removing toxic data. We first replicate the set-up from Daheim al. (2023) and train a dialogue model on of (Dinan et al. 2019) which we then train further synthetic data containing hallucinations using quadratic penalty. This model is subsequently used for removal from the first one. We found 1.0 to be the best scaling factor for both merging funetions. We evaluate faithfulness using critic (Dziri itt al. 2022) and Q² (Honovich et al. 2021). ). and fluency using BLEU (Papineni et al. (2002). Table shows improvements in terms of both faithfulness metrics with when compared to task arithmetic for two sizes of flanTSS (Chung et al. 2022). For unlearning tox- icity, the set-up from ITharco al. (2023). We use GPT2 (Radford et al., 2019), and GPT-J 1.3B (Wang& & Komatsuzaki (021) and fine-tune each model on the comments with toxicity score 0.8 from Civil Comments (Borkan et al., 2019). For both task arithmetic and our method we introduce scaling factor because initial experiments with
We study two settings for removing harmful examples from LLMs:
minor degradation in BLEU
showed degenerate generations. Potentially, H, is underestimated, as we can not run the squared gradient approximation on all train- ing data because it is either unavailable or too computationally expensive; instead we estimate it on WikiText103. We then prompt all models with "I don't care if this is controversial" and prompt from RealToxicityPrompts (Gehman et al., 2020). Finally. we evaluate the toxicity using Detoxify (Hanu & Unitary team 2020). and perplexity on WikiText103 (Merity et al., 2017). We classify all generations with score 0.2 toxic. Table Thswws that our method reduces toxicity in comparison to TA for both models and perplexity strongly for GPT2.
5 CONCLUSION
In this paper, we connect the error of the merged model to the gradient mismatch between the individual models that are merged and the 'target model' that merging aims to recover. We use this insight not only to propose new methods for model merging but also to understand existing ones. We also show deep connections to Bayesian inference which point to new directions for further improvements. Since the target model not available during merging by definition, our proposed merging method reduces the gradient mismatch by second-order approximation and is therefore tied to the uncertainty of the models, which determines their scaling. Our merging method shows improvements over previously proposed methods, such as task arithmetic, averaging, and Fisher- weighted averaging on CV and NLP tasks, both for task addition, where reduces the gap to the target model trained on all data, and removal, for example, for removing toxicity or hallucinations from LLMs. Notably, the proposed method is much more robust to the choice of scaling factors as tuning.
scaling naturally appears in its derivation without the need for

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This project has received funding by the German Federal Ministry of Education and Re- search and the Hessian Ministry of Higher Education, Research, Science and the Arts within their joint support of the National Research Center for Applied Cybersecurity ATHENE. This work is supported by the Bayes duality project, JST CREST Grant Number JPMICR21 12.
REFERENCES
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git Re-Basin: Merging mod- els modulo permutation symmetries. In International Conference on Learning Representations (ICLR), 2023. URL
Alan Ansell, Edoardo Ponti, Anna Korhonen, and Ivan Vulić. Composable sparse fine-tuning for cross-lingual transfer. In Annual Meeting of the Association for Computational Linguistics (ACL.), volume 1: Long Papers, 2022. URL http://ppeeeeeeeeeeeeeeooooaaaaa /tr.t////ooouiaiaaaooaaammmam11 !
long.
125
5
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Learning (ICML), 2015. 7
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In World Wide Web 18
Conference (WWW), 2019. URL http:/ aaaiviiiiiiv///iia/.aaaaaaaaaava org/ abs://bb//t// wwt.aass....... 1903, 04561
Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote Sensing Image Scene Classification: Bench- mark and State of the Art. Proceedings of the IEEE, 105(10):1865-1833 Oct 2017. ISSN 1558- 2256. doi: 10.1 1099/ppoc2211..2759888 2017.2675
URL http: //tp::///xxxe doi org/ 1109/JPRRCCCCCC
2 2755998
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pel- lat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction finetuned language models, 2022. URL https: arxiv. org. abs 2210 11416 arXiv:2210.11 1416. g
M. Cimpoi, S. Maji, I. Kokkinos, S. Mohamed, and A. Vedaldi. Describing Textures in the Wild. In Proceedings of the IEEE Conf. Computer Vision and Patlern Recognition (CVPR), 2014. g
Dennis Cook. 15-18, 1977.
Detection of influential observation in linear regression. Technometrics, 19(1):
16
Nico Daheim, Nouha Dziri, Mrinmaya Sachan, Iryna Gurevych, and Edoardo M. Ponti. Blastic weight removal for faithful and abstractive dialogue generation, 2023. URL https::tppssattt orq/abs/2303 17574 arXiv:2303. 1
18
Marc Deisenroth and Jun Wei Ng. Distributed gaussian processes. In International Conference on
Machine Learning, pp. 1481-1490. PMLR, 2015. 6
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Anmual Meeting of the Association for Computational Linguisticy (ACL), volume I (Long and Short Papers), 2019. URL https:
/aaclanthology wwantthooooy..or//111. www.oog//... 9- 423
Emily Dinan, Stephen Roller, Kurt Shuster, Angela Fan, Michael Auli, and Jason Weston. Wiz- ard of Wikipedia: Knowledge-poweet conversational agents. In International Conference
Learning Representations (ICLR), 2019. URL
net foet/fformmdd
11133RRRK 18
10

---

Published as a conference paper at ICLR 2024
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations (ICLR), 2021. URL. htttps://openrest net.///wweeeeeeeeeeeeeeeeeeeeeee http://ormmmiiiie
Hugh Durrant- Whytt.. Data fusion in decentralised sensing networks. Fourth International Con-
ference on Information Fusion, 2001, 2001. 6
Nouha Dziri, Ehsan Kamalloo, Sivan Milton, Osmar Zaiane, Mo Yu, Edoardo M. Ponti, and Siva Reddy. FaithDial: A faithful benchmark for information-seeking dialogue. Transac- tions of the Association for Computational Linguistics, 10:1473-1490, 2022. URL https:
//aclanthology. wrg/2022.
wwwpaattwooooyyoooo2222........ tacl22.tacc12 1.84 9
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference on Machine Learning
(ICML). 2020. [
Yingbo Gao, Christian Herold, Zijian Yang, and Hermann Ney. Revisiting checkpoint averaging for
neural machine translation. In Findings of the Association for Computational Linguistics (AACL-
IJCNLP), 2022. URL http::::////laaahhlooyyy www.:o22..r..2..oo.ooo ozg/2o22.findiings... findiiisninss... aacl. 18 D
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith. Realtoxici- typrompts: Evaluating neural toxic degeneration in language models. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2020. URL BRR/nttss://allaattllooo.
orq/ 20t//2020.. findings-emn.nn 301. pdf
18
Alex Graves. Practical variational inference for neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2011.
7
Laura Hanu and Unitary team. Detoxify. GitHub, 2020.
URL
https:// com/
unitaryai/Getoxify Last Accessed: September 2023. 18
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 IEEE International Geoscience and Remote Sensing Symposium, 2018.
Or Honovich, Leshem Choshen, Roee Aharoni, Ella Neeman, Idan Szpektor, and Omri Abend. 2 Evaluating factual consistency in knowledge-groundd dialogues via question generation and question answering. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2021. URL https:://cclanthaaat:
org/ www21.emm... commlppmmin... 619
Sebastian Houben, Johannes Stallkamp. Jan Salmen, Mare Schlipsing, and Christian Igel. Detec- tion of traffic signs in real-world images: The German Traffic Sign Detection Benchmark. In International Joint Conference on Neural Networks (IJCNN), 2013.
Ferene Huszár. Note on the quadratic penalties in elastic weight consolidation. Proceedings of the National Academy of Sciences, 15(11 115(11))2299-FF3) F:2299---22977 2018.
7
Gabriel Ilbarco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori,
Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi. and Ludwig Schmidt. OpenCLIP, 2021. URL https::/ .el.lorr//ie 100 5281/
zenodo. 5143773
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Ha- jishirzi, and Ali Farhadi. Editing models with task arithmetic. In International Conference on Learning Representations (ICLR). 2023. URL
6tCKkf8 ir
18
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Conference Uncertainty in Artificial Intelligence (AI), 2018.
papers/313.po

---

Published as a conference paper at ICLR 2024
Louis A Jaeckel. The infinitesimal jackknife. Bell Telephone Laboratories, 1972.
Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietro, and Pengxiang Cheng. Dataless knowledge fusion by merging weights language models. In International Conference on Learning Representations (ICLR), 2023. URL htthttp:: ppw.neveeeeeeeeeeeom net/ wow.m/id-
2
Mohammad Emtiyaz Khan and Siddharth Swaroop. Knowledge-addppttion priors. In Advances in Neural Information Processing Systems (NeurIPS), 2021. 7
Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivas- tava. Fast and scalable bayesian deep learning by pright-perturbation in adam. In International Conference on Machine Learning (ICML), 2018. 7
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
78]
18
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka
Demis Hass- abis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forget- ting in neural networks, Proceedings of the National Academy of Sciences (PNAS), 114(13): 3521-3526, 2017.
2 i
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International Conference on Machine Learning (ICML), 2017.
5 17
Jonathan Krause, Michael Stark, Jía Deng, and Li Fei-Fei. 3D object representations for fne-grained
categorization. In International Conference on Computer Vision (ICCV), 2013. (Workshops). g
Pierre-Simon Laplace. Mémoires de mathématique et de physique. Tome Sixieme, 1774. 6
Yann LeCun. The MNIST database of handwritten digits, 1998. URL httt:t/ppyynnn.... lecun.
com htmm/xxdb/mmt ist/
g
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: robustly optimized BERT pre- training approach, 2019.
1992 arXiv: 19X1:1190.111699
Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Tay, Denny Zhou, Quoc V Le, Barret Zoph, Jason Wei, and Adam Roberts. The Flan collection: Designing data and methods for effective instruction tuning. In International Conference on Machine Learning (ICML), 2023.
URL htttps::/pprcceedinggs
mlr. wttppsss/v22//ooggreee2222 http:///cocraroos.....r.rrse///p .tt.23....mm....mmmmm
18
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Con-
ference on Learning Representations (ICLR), 2019. URL
Barrum1d=Bkg6Ric.
18
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Annual Meeting of the Association for Computational Linguistics (ACL), 2011. URI. http: htttp://www.. acl whb.or//ttoooooooohhhhhhooooool
P11-1015
David JC MacKay. A practical Bayesian framework for backpropagation networks. Neural Compu- tation, 4(3):448-444 472, 1992. b
James Martens. New insights and perspectives on the natural gradient method. J. Mach. Learn, Res. (JMLR), 21())5577--581,, 2020.
Michael S Matena and Colin Raffel. Merging models with Fisher-weighted averaging. In Advances in Newral Information Processing Systems (NeurIPS), 2022. URL httphtttp:///oeeeeeeeee
net forrm wimmmidd==..m 1d=.skkp__aaa aceoC
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture
models. In International Conference on Learning Representations (ICLR), 2017. URL ht
//openreview. wttp://ewwweeettoommm////aaa
9
12

---

Published as a conference paper at ICLR 2024
Arthur G. O. Mutambara. Decentralized estimation and control for multisensor systems. Routledge, 1998. o
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learn- ing? Advances in Neural Information Processing Systems (NeurIPS), 2020, !
Manfred Opper and Cédric Archambeau. The variational Gaussian approximation revisited. Neural computation, 21(3):786 792, 2009. 7
Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tan- gent space: Improved editing of pre-trained models, 2023. URL http::/arxiv..o/ abs/ arXiv:2305.
2305. 12827
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with Bayesian principles. In Advances in Neural Information Processing Systems (NeurIPS), 2019. 7
Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In Annual Meeting of the Association for Computational Linguistics (ACL), 2005.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. BLEU: a method for automatic evaluation of machine translation. In Annual Meeting of the Association for Computational Lin- guistics (ACL), 2002. 9
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks, 2013. URL
absg:/ab//inn 1301. 3584 arXiv:1301.358. 7
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019, URL htttps:// www/openail....ooooooppo.
gpt
9
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. J. Mach. Learn. Res. (JMLR), 21(140):1-67, 2020. URL http: jmlr:: jmlr org/
papers/v21/20-001 html
19
Nicol N Schraudolph. Fast curvature metrix-vector products for second-order gradient descent.
Neural Computation, 14(7):1723 2002. 7
Jonathan Schwarz, Wojciech Czarnecki, Jelena Luketina, Agnieszka Grrabsk--Barwinskee Yee Whye Teh, Razvan Pascanu, and Raia Hadsell. Progress compress: scalable framework for con- tinual learning. In International conference machine learning, pp. 4528-4537. PMLR, 2018.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Conference on Empirical Methods in Natural Language Processing (EMNLP),2013 URL http:::t//// //wp::::/wwwwww aclweb. org/ httg///annooooyyo//-1l1111
Luke Tiemey and Joseph B Kadane. Accurate approximations for posterior moments and marginal densities. Journal of the American Statistical Association, 81(393):82 1986. 6
Volker Tresp. A Bayesian committee machine. Neural Computation, 12(11):271-2277 2000. 6
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems (NeurIPS), volume 30, 2017. 19
Ben Wang and Aran Komatsuzaki. GPT-J-6B: billion parameter autoregressive language model, 2021. URLfattps://qithub com/ http://ccooommnoooooooooeee///ee
former- jax Last Ac-
cessed: September 2023. 18
13

---

Published as a conference paper at ICLR 2024
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jemite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-art natu- ral language processing. In Conference Empirical Methods in Natural Language Processing (EMNLP), 2020. 18
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes. Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference Machine Learning (ICML), 2022a. URL
mlr http.eeees..n.......rrresssss//s /taxesss/v11l2/wwottsaaa222222
html
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo-Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022b. https: /ttpsii//aavvi.o/o www.iiwwgwwa.a......... arggabs/2209 01 903 225
Peng Wu, Tales Imbiriba, Victor Elvira, and Pau Closas. Bayesian data fusion with shared priors, 2022. (RLp: org/ abs 22122. 07311 arXiv:2212.07311,
J. Xiao, J. Hays, K. A. Ehinger, A. Oliva, and A. Torralba. SUN database: Large-scale scene recogmition from abbey to zoo. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2010. 8
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving in- terference when merging models, 2023. URL arXiv:2306.01708.
01708
Netzer Yuval. Reading digits in natural images with unsupervised feature learning. In NIPS Work- shop on Deep Learning and Unsupervised Feature Learning, 2011.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level Convolutional Networks for Text Classification. Advances in Neural Information Processing Systems (NeurIPS), 2015. S
A DERIVATIONS
A.1 DERIVATION OF TASK ARITHMETIC USING GRADIENT MISMATCH
We proceed similarly Sec. 3 by first writing the respective stationarity conditions for the LLM LLIM. fine-tuned models 0,, and target model 0 01:r,
OLIM (fLMM(
Ho 0LIM) (O.),, for all t= 1,2.....
T
(0 - 0LIM) = Σ 04 (Q.
Next, we multiply the second equation with a for each t, then sum over 1,,,..., and finally subtract it from the third equation to get the following,
T
T
H6 - OLLM)
Σ
a. Ho OLLM) = Σ - (17)
t=1
Multiplying by H²² and rearranging gives Eq.
10
14

---

Published as a conference paper at ICLR 2024
A.2 DERIVATION OF THE NEW METHOD
By substituting Taylor's approximation of Eq. sion below which is linear in 0 1:7,
11 in Eq. 10 the equation reduces to the first expres-
T
O.T OLIM Σ c. OLIM)
Σ
Ho (H,(01. 0.)]
(18)
We then add and subtract OLLM the last term above.
Our OLLM Σ (0, OLIM) Σ H
H, LLLM) H, LLLM))
(19)
and multiply the whole expression by Ho and rearrange it to get the second expression in Eq.
18
T
Ho + Σ H. (01))(11. OLAM) ~
Σ
HoH - 01IM) + Σ Ho. OLAM)
(20)
(Ho + (.) BLLM)
t=1
Multiplying the equation by inverse of Ho + Σ H,H, and taking fLMM the right hand side gives us Eq. 12
A.3 DERIVATION OF MODEL MERGING AS MAP OF BAYES' POSTERIOR
We will now connect our approach to Bayes' rule for linear regression. this case, Eq. TO can be rearranged write as follows, where in the second term have added and subtracted 01.T,
(01r OLLM)
H (G.
0ur Our OLLM)
H, 01)
t=1
t=1
This equation is stationarity condition of the following optimization problem,
T
T
O1T = arg min (1-
0
OLLM
P
+
a2
0 0, +
1=1
log ((O(.))
where we identify the prior to be p(0) (/ LLMM H ) and the posterior on D, to be (()),;
(0 (Ho + 1)) We can therefore show that the stationarity condition corresponds to maximum-a-posterion estimate of p(0 as follows.
p ((O)
D(0)) p00 Il (DD. p(0)
II
p(0))-ETT.
Il (())..)
p(0)
1=1
where log the last term is equivalent to the objective function.
A.4 DERIVATION OF DATA REMOVAL
Our target model is the following model trained using Eq. but without using D,,
0 r arg min (e( + 1011,, where (0)
Σ
())
(21)
(E Lons D.)
The LLM objective can then be written in terms of this objective:
θLLM arg min (,) + (,) 1012
8
(22)
0
where assume that & is multiplied by constant or in the original model.
15

---

Published as a conference paper at ICLR 2024
As before, write the stationary conditions of LLLM 0,, and py. respectively:
80LLM -V( - -a (
Hg (0, L = - vl,(O,),
(23)
50 - v&)o(0))))
Because our goal is to analyze 0, (0LMM - -O)). we multiply the second equation by subtract from the first equation, and then subtract the resultant from the third equation to get, the following,
600 LLLM) or Ho Ho(0. (0, LLLM) (0,)) (Oum)) + (QL(())M)) (0,
(24)
We can now use Taylor's approximation reduce gradient matching.
(0 (OLMM)) (OLLM))))) 0LLM)
For the second gradient term, we do not need to use the Taylor's approximation because it does not depend on Qr, but our goal is to improve over task arithmetic, so we do to derive preconditioner,
(OLAM) R Vē, H, (, 0.).
(25)
Note that it is also possible to do the Taylor's approximation not around O, but OLLM. Plugging these Eq. 24 we can write,
500 ; LLLM) + az Ho (0: LLLM) - (6LLM) OLIM) + au (He(OLLM - -6)
[8I (OLMM) (0), LLLM) -0 (Hu H,) (0, 0LLMM
0/ = LLLM a, [01+ v (OLLM)) (Ho + H₂) (0: OLLM)
which gives us the desired update given in Eq.
A.5 PROOF THAT OUR UPDATE FOR DATA-REMOVAL IS EXACT FOR LINEAR REGRESSION
The task removal update derived above is closely related to previous works on data removal. For instance, for linear model, our update recovers the popular influence function. We will now show this. Consider a large linear model (coincidentally also abbreviated as L.M) with full data D = (X.y) where y is a vector outputs and is matrix containing each feature vector row. The loss is (0) llyy X0 Now, suppose we want to remove D, yx) from it. Then, we have closed form solution for the full model and the model with removed data,
11IM = H x
O₂ H
where = 22 2 lly - 2 + | = x
X
+ SI, and similarly H = X X, + OI. A well
known result in the influence function literature Cook (977) that the two quantities are related as
0,1- OLIM
X
[X,0IIM y.)-
(26)
We will now show that our update Eq. 13 reduces this for linear models.
We start with an expression for 0, trained using Eq. 8 but with the loss (,(0) = 21ly, - X,0
Using the second equation the optimality condition Eq. we can write:
Ho 0LIM) x, (y: X,O,)
(Ho + H.) HO0MI + x, yr
where use the fact that for linear models H, X. X, We now simplify update of Eq. 13 with oe = 1 where we use the above relationship the third line below,
OLLM H. (Ho H:) (0, OLIM)
OLIM H. [Ho H.) O. (Ho H.) 0IIM)
PLLM H. HDOLLM + X. (Ho H,) BLLM)
(27)
OLIM H X y: H,OLIM)
PLLM - H.l [x] y: - X X,OLLM)
OLLM + H. x, (X,0LIM - yy))
16

---

Published as a conference paper at ICLR 2024
Therefore, our update reduces to Eq. 26
generalization of Eq. 26 to neural network is considered in Koh & Liang (017) for the case of one-example removal. Their approach when applied to remove multiple examples at once redues to
0e=- = LLLM + H ge:
where = Ve, Our approach also recovers this result if we do not use the second Taylor's approximation for the second gradient matching term in Eq. 24 Essentially, this removes the con- tribution of the fine-tuned model and the steps are completely based on LLMM- It is not clear which approach is better but in practice may depend on the fine-tune process which by doing multiple gradient steps may be able to get more information than single gradient step g.. We hope to explore this point a future study.
B PRACTICAL CONSIDERATIONS
B. CHOICE OF Loss FUNCTION & REGULARIZER
One design decision a practitioner has to take is the amount of regularization and the regularizer itself. While the presented derivations rely on weight decay and a quadratic regularizer, this restric- tion needs not to be made and is rather for simplicity of derivation. Consequently, we have found that even training models without a quadratic penalty gives similarly good results when merging with our objective. This might also be connected to the fact that early stopping of neural network training already implicitly regularizes the model towards its initialization, and not optimizing for too long will keep the fine-tuned model close. The 8 chosen in weight decay is also included in Eq. 13 where I + H turns to 5I + H. if another value but 1/2 is chosen and in Eq. 12 if Ho Hy SI. Practically, a small 5 > 0 similar to what was outlined in Sec. [43] works well and 6 = 0 can become unstable, because the squared gradient approximation may also be O for some parameters.
The scaling of the Hessian approximation is also determined by the loss chosen for training. In case l is an average, the Fisher should also be averaged. This implies that all tasks viewed as equally important. A different choice is to not average and consequently the Fisher estimate is just summed. This carries Bayesian interpretation of weighting tasks with more data higher, because observing more data means being more certain. Ultimately though which method performs better on test data can not be answered priori and choice make by the practitioner.
C EXPERIMENTAL DETAILS
This section contains experimental details to reproduce our study. We will also provide a repository containing the implementation upon acceptance.
C.1 TASK ADDITION FOR CV
We use the implementation provided by (Ihaarco al. 2023)
basis for our experiments. We
use the ViT-B-32 model available at htttthts://// part of Open CLIP (Ilharco et 2021 and initialize all task models from it without training on a task beforehand. All models are trained using a large batch size of 128 on NVIDIA GPUs. We train the models using modified version of the Adam optimizer (Kingma & Ba 2015) that uses a decoupled quadratic penalty. We have found that smaller number of epochs performs better, especially for task arithmetic but also for our method. Potentially, training long increases gradient mismatch and makes the fine-tuned models deviate far from the pretrained model. Therefore, we train for epochs for small datasets and 10 for larger ones. We use learning rate of le 3, 31 0.9 and 32 0.999. The Hessians of the task models are approximated using the squared gradient approximation on each task (normalized by data size), and for the pretrained model we use the identity multiplied by scalar that is not tuned. We merge all models with the pretrained ViT and vary the scaling factor a, which we keep same for each task model 0, for = 1, 2.... this experiment.
com/ml foundations hopooon///oonpppppppepp/
.. I within
17

---

Published as a conference paper at ICLR 2024
Model
Hessian Approximation
IMDB
Yclp
RT
SST2
Amazon
Avg.
Truc Avg.
H N
94.6
97.3
89.5
93.3
96.6
94.3
96.5
RoBERTa
Squarod Gradient (F sm)
94.7
(10.1)
97.3 (9)
90.2 (10.7)
93.7
(10.4)
96.7 (10.1)
94.5 (10.2)
96.6 (1o.1)
Table 5: The squared gradient approximation of the Hessian performs better than identity for task addition to an LLM for sentiment analysis, but even identity works well due to the improved scaling by accounting for gradient mismatch.
C.2 TASK ADDITION FOR NLP
We use the RoBERTa-base checkpoint available on the huggingface hub Wolf al. (2020) to ini- tialize our model (available at datasets in the versions that are available on the hub. We train each model for 2 epochs on IMDB and Yelp, 1 epoch Furthermore, we subsample the training data for Yelp and Amazon and take the first 20% and 10% of the data to reduce computational load, and also because the datasets are much larger than the other three datasets and should not completely dominate them in the all-data baseline. We train all models using either AdamW Loshchilov & Hutter (2019) or modified version of Adam (Kingma & Ba 2015) with a quadratic penalty and use a learning rate of le 5 for training RoBERTa base on IMDB and of 5e 6 for training the other models initialized from the IMDB model. We truncate the inputs at 384 tokens, and train using batch size of 16 on NVIDIA GPUs. Furthermore, we set 31 0.9 and 82- = 0.999 as is standard in the transformers library, and use 100 warmup steps, as well as gradient norm clipping to unit norm, All squared gradient approximations are done by doing a single pass over the training data used to train each model, truncated to at most 100,000 examples. We then merge according to our proposed function by using small å = le 10. For as is default in the implementation provided by Matena & Raffel (202) but rather add the same 8 as in method. The "all-data" baseline is trained on the concatenation of all data used for the task-specific models for 2 epochs.
co/ hwwprrobebaaaaaseeee and also all
Amazon, and 5 epochs on the smaller SST2 and Rotten datasets.
Fisher averaging. we do not clip the Fishers at 1e-
C.3 TASK REMOVAL FOR NLP
For hallucination removal, we use the code provided by Daheim et al. (2023) to train a flanTS model [ongpre et al. (202)) on Brrard-of-Wikipedde Dinan et al. (2019). We use the same data augmentation of switching out the ground-trus truth document and train both the initial dialogue model and the model trained on the augmented data for 1 epoch using a batch size of 32. Again, we use the checkpoints provided
huggingface. We use AdamW or the modified version with learning rate of 6 25c 5. 81 = 0.9, and 32 = 0.999. All Hessians are approximated by the squared gradient approximation of the Fisher by passing over the training data of each model once. We use d = le 10 for merging and have found a smaller delta to generally provide better performance.
For toxicity removal, tried adhere to the setting of Ilharco et al. (2023) as closely possible by employing the same prompts and datasets. In particular, we train GPT2 (Radford et al. 2019) or GPT-J (Wang & Komatsuzaki 2021) on all data from Civil comments (Borkan et al. 2019) with a toxicity score of over 0).8 and "subtract" this model from a pretrained GPT2 model (which we do not fine-tune on any data but use out-of-the-box given the checkpoint provided on the huggingface hub). The model is trained for only 1 epoch on the training data. We calculate the squared gradient ap- proximation for GPT2 on the first 10000 instances of WikiText103 and for the task model on the task data. For evaluation we take the prompt from I]harco al. (0023), namely "I don't 't care if this is con- troversial." followed by a prompt from the Real SealToxictyPrompp dataset (Gehman et al. 2020). We use 'original' model available under htt http//itttte thub. com. /u///unii aaritaraaai /eaii//eeooo oxtoxitiiieo Fy from the Detoxify library (Hanu& & Unitary team 2020) to evaluate toxicity and score each output as toxic that exhibits score greater than 0.2.
18

---

Published as a conference paper at ICLR 2024
D ADDITIONAL RESULTS
D.1 COMPARISON OF HESSIAN APPROXIMATIONS
We compare two different methods of setting the Hessians for the task addition experiment: setting all Hessians to identity and 2) squared gradient approximation of the Fisher. In Table S we find identity Hessians work surprisingly well, on par with Task Arithmetic and tuned Or, but still to be outperformed by the squared gradient approximation. This indicates that improved approximations could further improve merging performance which we leave to explored in future work.
D.2 SENTIMENT ANALYSIS WITH T5
We repeat the same task addition experiment for sentiment analysis as in the main paper Sec. 4.3 using T5 (Raffel al. 2020) which pretrained encoder-decoder transformer (Vaswani al. 2017) to outline that our method scales to different kinds of transformer (Vaswani et al.) 2017) architectures. We use the pretrained T5 checkpoint without supervised fine-tuning, SST2 a part of that data, found under https: //huggii ng f face. co / goog eltgglletttvvl_e. base We use learning rate of 4 for the model trained on IMDB and 5e for the models trained on the other tasks which were initialized from it. The model predicts the sentiment of the text by generating the words "negative" or "positive". The results in Table confirm that our method works across architectures.
IMDB
Yelp
RT
SST2
Amazon
Avg.
True avg.
Model
Parametrization
Accuracy (1)
All-data
94.9
98.0
91.0
94.5
97.2
95.1
97.1
T5250M
Task Arithmetic
90.5
97.0
87.9
92.3
95.5
92.6
95.3
Ours
93.6 (13.1)
97.5 (10.5)
90.7 (726)
94.6 (12.3)
96.5 (11.0)
94.6 (12.0)
96.4 (1L.)
Table 6: For adding sentiment analysis tasks to T5 (Raffel et al. 2020) trained IMDB, reducing gradient mismatch improves merging performance.
19


Reviews:
[{"review_id": "BOLkl6IWKu", "reviewer": "Reviewer_idAX", "summary": "The paper introduces a powerful method to average two models. Specifically, the proposed method averages the model by minimizing the gradient mismatch of different models. The paper provides a deep analysis of why their method makes more sense than others.\n Also, the paper validates their method in multiple datasets from both the NLP domain and the image domain.", "questions": "Please address the question above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "From my viewpoint, the weight of LLM is knowledge abstracted from data, which stresses the importance of quickly merging knowledge learned from the dataset. I believe the topics of the paper fit into this conference and have a certain inspiration for future works in this domain. \n\nThe motivation of this paper is extremely clear by analyzing the gradient of merged models. When I read the paper, I enjoyed the motivation despite the heavy math.", "weaknesses": "I have some minor concerns about this paper. Before I lay out the weaknesses list, I would like to mention that I’m not an expert on NLP and my comments are probably incorrect.\n\nFinetuning vs data-driven model averaging. Maybe I don’t have the background.  I’m curious about the advantage of the proposed model merging over simply fine-tuning the model. In my understanding, for the proposed method to work, we would need data to calculate the gradient matrix -- that’s why I call the proposed method as a data-driven model averaging. In this case, why don’t we just simply fine-tune the averaged model using the LORA on the data in hand? And fine-tuning sounds more straightforward. Thus, I would recommend having a discussion/quick comparison between those two.\n\nAgain, I’m a bit concerned about the time efficiency since the proposed method requires the second-order Hessian matrix, especially when compared with the simple strategy. Although it doesn’t matter for the inference, it might be still worth knowing if this Hessian calculation is practical or not. So I suggest to make it clear. \n\nIn short, I have some concerns about the comparison with simple fine-tuning and time efficiency. So I currently vote for the weak accept. Again, it might be because I don’t have too much domain knowledge. So I would be happy to hear back from the authors during the rebuttal in case I misunderstand anything.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "AkaNlyxeqZ", "reviewer": "Reviewer_vWCx", "summary": "The paper addresses an interesting problem in the domain of model merging and offers a novel perspective by connecting gradient mismatches to the inaccuracy of weighted-averaging methods. The paper also proposes a new uncertainty-based scheme to improve model merging, which is a valuable contribution.", "questions": "+ In my opinion, the model merging technique takes two or more models as input and outputs a merged model. Therefore, the performance of a merged model on down-stream tasks (compared to the unmerged model) is only a single datapoint in the experiment. In other words, what if we use different hyper-parameters to train the base model on each task? Will your method outperform others under other hyper-parameters?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "+ The authors connect the inaccuracy of weighted-averaging to mismatches in the gradients and propose a new uncertainty-based scheme to improve performance by reducing the mismatch.\n+ The authors propose a unified explanation on previous model merging technique. \n+ The new method shows consistent improvements for large language models and vision transformers in terms of performance and robustness to hyperparameters.", "weaknesses": "+ My major concern lies in the problem setup. I admit that model merging is a well-defined problem with much previous literature, as is discussed in the submission. But I still wonder why we need this technology. If we could obtain the data for each task, why don't we simply perform multi-task learning on these data? If we couldn't, how could we obtain the fisher information matrix on each task, which is required to approach Eq.12? It seems like a contradiction and I think more clarification on the application scenario of the model merging technique is needed, in spite of the abundance of previous literature.\n\n+ The second concern is an important missing baseline. The derivation in section 3 is similar in some degree to Regmean[1] though the latter takes linear regression as an example and then extrapolates to neural networks.  Therefore,  I would list Regmean as one of the must-to-compare baseline methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "3RgO1FtPCN", "reviewer": "Reviewer_V4z1", "summary": "The describes a structured way of understanding linear combination of parameters. The concept of ``target model`` is introduced as a way of measuring fitness of merges. Subsequently, modification of the ``Task arithmetic`` loss is introduced such that the gradient mismatch between the target model and the averaged models is minimized. Experimental results show comparable results.", "questions": "**Questions**\n\n- Eq(3)  should be \n$$\\alpha_{1}\\bar{\\ell}_{1}(\\theta) + \\alpha_{2}\\bar{\\ell}_{2}(\\theta)$$\n right?\nIn other words, the optimization is looking for $\\theta$ that optimizes both losses which is $\\theta_{1+2}$. If this is not a mistake then Eq(5) is wrong. \n- what does $t$ stand for in Eq(8)\n- the error between $\\bar{\\theta}_{TA}$ and $\\theta_{1:T} = \\theta_{1:T}$? please clarify/correct?\n\n**Suggestions**\n- The discussion in the last paragraph in page 3 is best to be had in the experimental section with some data.\n   or under its own section with further details.\n- Section 3.1 is difficult to follow/understand, mainly because of several math annotation issue, and not well\n   organized paragraphs. For instance, the first two paragraphs can simply be phrased as `target model` definition\n   rather than using unnecessary details and confusing notations.\n- I generally, like the framing of the problem and the idea of ``target model``. I think the paper has good potential. I suggest re-writing of section 3, highlighting the problem and the solution (perhaps computational aspect and other details) and differing questions of generality and applications to later section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The introduction of the paper is very well written and framed the problem in clearly.\n- The idea of defining a ``target model`` is very useful concept in this space.", "weaknesses": "- Section 3, overall, was difficult to follow with seemingly several notational errors and cluttered paragraphs, see questions and suggestions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "xOBAXbtSXV", "reviewer": "Reviewer_1bWs", "summary": "The paper tries to theoretically understand the impact of gradient mismatch between tasks when merging these models together. The first shows how model merging and gradient mismatches are related to each other and shows the errors that are induced due to that, based on these insights they propose a new method to reduce gradient mismatch. Next, they demonstrate how many past model merging methods are the special case of their new method and finally establish the relationship with Bayesian inference. They conclude with a small set of experiments to demonstrate the usefulness of their method.", "questions": "**Need to address these questions for me to retain my score:**\n\nQ1: The papers need to be positioned better, to adjust the main contribution and highlight the similarities/differences compared to TIES-Merging [1] which claims to identify and ameliorate interference when merging models.\n\n\n\n**Need to address these as well for me to consider increasing my score:**\n\nQ2: Ideally all the experiments should also compare with TIES-Merging as that method is the closest to the final method proposed in this paper. And addresses the exact same problem that this paper tries to get at. Hence, not including that as a baseline leaves lots of open questions about the utility of the proposed method. If these comparisons are added then I will update my score as I feel this work makes a good theoretical contribution but the experimental section leaves a lot of ambiguity about the utility of the final method.\n\n\nQ3: The experimental results are very weak and seem insignificant. For example in Table (nlp), the experimental setting seems to be not well designed due to multiple reasons (i) the performance of all the methods lies between 96.1-96.8 which is quite a narrow range to make any claim about any of the methods performing better or worse than the others. for example, the difference between your method and TA is 0.3% which is not significant from my experience. Hence, either the experimental setting is too simple to highlight the differences between these methods or the methods all perform the same. An experimental setting from past papers like Task Arithmetic, TIES-Merging can be adopted for such experiments. \n\nQ4: It is not clear how the approximations made in the paper about using fisher instead of hessian after the gradient mismatch as the model becomes bigger, something on this would be useful. Moreover, could these be the reasons why the method does not lead to significant improvements in both vision and NLP settings?\n\n\nQ5: Figure-1 (right), it is not clear to me how this gradient mismatch is computed. Seems like you are adding 5 tasks on Roberta (IMDB) but then what circles represent the gradient mismatch between which models on which data? Is it a pairwise comparison of gradient mismatch or are the models being added? Please clarify exactly what this figure means. \n\n\nQ6: Moreover, the Figure-1 (right) it is shown that as the gradient mismatch decreases the test error also decreases significantly (by ~2). however, this finding seems to be inconsistent with the results in table-3 where the performance difference between TA and your method is very minimal (0.3). What is the reason behind this? In general, what is the reason behind not leading to enough improvements over TA in the nlp setting even when there is a significant gradient mismatch for TA?\n\n\n**Other questions on clarifications and details:**\n\nQ7: How is the alpha selected in your proposed method? It is mentioned in many places that alpha is not tuned.\n\nQ8: Please specify the number of samples you use to compute the fisher. \n\nQ9: For Figure-2 (left) the best performance for both TA and your method seems to be comparable to each other. I agree that your method might not need to tune for alpha but in most practical cases obtaining a small validation set and tuning alpha is not that hard. Moreover, the proposed method need to compute the fisher (requires backward pass on a subset of training data) whereas TA need validation set to tune alpha (inference on a small number of val example), so overall the peak memory usage of the proposed method would be higher while TA requires additional data. This trade-off should be highlighted in the paper.\n\nOverall, I feel that the theoretical contributions of this work are nice and would be useful to the community, I expect the author to at least position the contributions of the work better in light of past works. Moreover, strengthening the experiments section would highly increase the quality of this works\n\n[1] Resolving interference when merging models, Yadav et. al.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "(S1) The originality of the work lies in providing the theoretical connection between model merging and the gradient mismatch problem (the identification of gradient mismatch as a problem of model merging cannot be attributed to this work, see weaknesses). Moreover, the connection with the Bayesian inference is also interesting. Additionally, the insight that the work extends RegMean (Jin et. al;2023) to non-linear parts of transformers is useful.\n\n(S2) The paper is well-written for the most part and easy to follow along. However, there are things that are not clear and are listed as Questions below.\n\n(S3) The work improves the understanding of model merging and might be useful to the people working on the model merging.", "weaknesses": "(W1) The idea of identifying gradient mismatch as a problem has been claimed as one of the main contributions throughout the paper (abstract, intro, and other sections). However past works like TIES-Merging [1] have identified this problem and proposed detailed empirical studies to quantify the degree of this problem and then propose some fixes that lead to significantly improved performance. \n\n(W2) Moreover, in the current version of the paper, TIES-Merging is discussed in passing but the differences compared to that works are not properly highlighted.\n\n(W3) The experimental section is pretty thin and the results presented are weak. See more details in the Questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)