Paper: Model Merging by Uncertainty-Based Gradient Matching (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
A1: Thanks for the suggestion. We have now added a few sentences to highlight the differences to TIES-Merging (see the last line in Page 2 and also Section 3.2.1 in the revised version). In summary, the method can be related to ours by using a binary mask (similarly to Ansell et al. 2022 in Table 1) but the paper [1] does not mention gradient matching and it remains unclear how the operations used in TIES directly lead to reductions in gradient mismatch.
While gradient mismatch is calculated between each $\theta_t$ and the target model $\theta_{1:T}$, their concept of interference is outlined between the $\theta_t$.

Reviewer comment (Novelty):
In the current version of the paper, TIES-Merging is discussed in passing but the differences compared to that work are not properly highlighted.

Author response:
A1: Thanks for the suggestion. We have now added a few sentences to highlight the differences to TIES-Merging (see the last line in Page 2 and also Section 3.2.1 in the revised version). In summary, the method can be related to ours by using a binary mask (similarly to Ansell et al. 2022 in Table 1) but the paper [1] does not mention gradient matching and it remains unclear how the operations used in TIES directly lead to reductions in gradient mismatch.
While gradient mismatch is calculated between each $\theta_t$ and the target model $\theta_{1:T}$, their concept of interference is outlined between the $\theta_t$.

Reviewer comment (Experiments):
The experimental section is pretty thin and the results presented are weak. See more details in the Questions below.

Author response:
A3.1: We respectfully disagree: improvements (like 0.3%) are compared to a “well-tuned TA” method in Table 3, but without such tuning the baselines are much worse (also see this in Fig. 2). Larger gains with respect to other measures are also obtained in Table 2 and 4 (e.g., a gain of 3.4 for fluency PPL). A 0.3% increase is also not insignificant for a large test set (in this case containing close to 500,000 test examples). We understand that, when only focusing on a specific measure (test accuracy) with respect to one method (well-tuned TA), one may call the improvements marginal, but this does not mean that the presented methods are ineffective (and insignificant) in all respects. We request the reviewer to reconsider their opinion.

Reviewer comment (Experiments):
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.

Author response:
A2: We request the reviewer to reconsider the need to compare *all* the experiments. The ICLR policy clearly advises to not have to compare with papers posted on arXiv after May 28, 2023. We do not agree that the absence of such comparisons “leaves a lot of ambiguity about the utility” of our method (see A3.1 below). We also disagree that TIES-Merging is “closest” to our method: our proposal is to reduce gradient mismatch which is different from the trimming used in TIES; gradient matching is never even mentioned in the TIES paper. We respectfully request the reviewer to reconsider their opinion.

**To address the reviewers’ concerns, we have added a comparison in Table 3 which is feasible to do in this short time. We kindly ask the reviewer to reevaluate their rating and increase their score based on the new comparison.**

Reviewer comment (Experiments):
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; (ii) 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 and TIES-Merging can be adopted for such experiments.

Author response:
A3.1: We respectfully disagree: improvements (like 0.3%) are compared to a “well-tuned TA” method in Table 3, but without such tuning the baselines are much worse (also see this in Fig. 2). Larger gains with respect to other measures are also obtained in Table 2 and 4 (e.g., a gain of 3.4 for fluency PPL). A 0.3% increase is also not insignificant for a large test set (in this case containing close to 500,000 test examples). We understand that, when only focusing on a specific measure (test accuracy) with respect to one method (well-tuned TA), one may call the improvements marginal, but this does not mean that the presented methods are ineffective (and insignificant) in all respects. We request the reviewer to reconsider their opinion.

A3.2: Our experiments are already adopted from the task arithmetic papers (for example, Figure 2 (left) and Table 4 (left)) and other recent works (Table 4 (right)). TIES-Merging is not used because it is too recent but we will consider adding further experiments for the final version.

Reviewer comment (Theory):
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 work; 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?

Author response:
A4: We expect better approximations to give better results but it is not correct to directly attribute bad performance to the approximation. We use the GGN approximation which is a popular choice and is effective even at large scale [2,3,4,5]. We plan to study the effect of this approximation in the future. 
(pt.2 below)

Reviewer comment (Presentation):
In Figure-1 (right), it is not clear how this gradient mismatch is computed. It seems like you are adding 5 tasks on Roberta (IMDB) but then what do the 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.

Author response:
A5: It is a pairwise comparison of the gradient mismatch between the merged model and the target model (trained on all data). Each marker represents one task; see Section 4.3 on the details of the experimental setup.

Reviewer comment (Evaluation):
In 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?

Author response:
A6 & A7: This is a misunderstanding; there is no such inconsistency. The reviewer is talking about the “well-tuned TA” baseline, but Fig. 1 uses default TA with $\alpha_t =1$ which (consistently) performs worse in Table 3 as well (e.g., 90.5% on IMDB compared to our 94.7% in the first column). In our method, we always set $\alpha_t = 1$ and do not use any tuning, yet obtain better results than a well-tuned TA baseline. We have now improved the writing of the paper to avoid the misunderstanding (Fig. 1, Table 2 and Sec. 4.3).

Reviewer comment (Experiments):
How is the alpha selected in your proposed method? It is mentioned in many places that alpha is not tuned.

Author response:
A6 & A7: This is a misunderstanding; there is no such inconsistency. The reviewer is talking about the “well-tuned TA” baseline, but Fig. 1 uses default TA with $\alpha_t =1$ which (consistently) performs worse in Table 3 as well (e.g., 90.5% on IMDB compared to our 94.7% in the first column). In our method, we always set $\alpha_t = 1$ and do not use any tuning, yet obtain better results than a well-tuned TA baseline. We have now improved the writing of the paper to avoid the misunderstanding (Fig. 1, Table 2 and Sec. 4.3).

No it is not tuned and we will clarify this in the paper. The plot simply shows results for different alpha on the test data, not on a validation set. In our method, alpha is always set to 1.0 during training.

Reviewer comment (Reproducibility):
Please specify the number of samples you use to compute the fisher.

Author response:
A8: We use either the entire train dataset or at most the first 100,000 examples from it if the dataset is larger to reduce computational burden. We will add this to our paper (Appendix C.2).

Reviewer comment (Evaluation):
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 needs to compute the fisher (requires backward pass on a subset of training data) whereas TA needs a validation set to tune alpha (inference on a small number of validation examples), 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.

Author response:
A9: Thanks for the suggestion: we have now highlighted this point in the paper (see the last paragraph of Section 3.3). For a large number of tasks, setting alpha can be expensive, but we agree that for a small number of tasks it can be easy when a validation set is known and available. We also stress that we do not need to do a full pass through the data because Adam-style methods can be easily modified to obtain an estimate of the Fisher during training and without overhead (a common practice to obtain online Fisher in continual learning; cf. [6] or [2; Section 11.2]).

Reviewer comment (Writing):
Section 3, overall, was difficult to follow with seemingly several notational errors and cluttered paragraphs, see questions and suggestions section.

Author response:
A: Thank you for the constructive comments! We have now rewritten the initial paragraphs of Section 3.1 to clarify the model definitions and split Section 3.2 into further subsections to ease readability. We also have fixed the notation mistakes, thanks again for pointing these out. For the final version, we will make sure to take your suggestions into account!

Reviewer comment (Theory):
Eq(3) should be 
$$\alpha_{1}\bar{\ell}_{1}(\theta) + \alpha_{2}\bar{\ell}_{2}(\theta)$$
right? In 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.

Author response:
A: Yes, thank you for the careful reading. We have corrected it!

Reviewer comment (Writing):
What does $t$ stand for in Eq(8)?

Author response:
A: The $t$ is arbitrary but fixed and stands for a specific task or dataset $\mathcal{D}_t$ that the model $\theta_t$ is trained on. During merging we combine all these $\theta_t$, of which there are $T$ in total. We have made this clearer in Section 3.1.

Reviewer comment (Writing):
The error between $\bar{\theta}_{TA}$ and $\theta_{1:T} = \theta_{1:T}$? Please clarify/correct.

Author response:
A: The error is calculated between $\theta_{1:T}$ and $\bar{\theta_{TA}}$. 
The former is the target model, and the latter its approximation through Task Arithmetic. 
We can rewrite Eq. (10) to make this more explicit by subtracting the left summand on the RHS (which is $\bar{\theta_{TA}}$).
Then, we obtain the following equation: 
$\theta_{1:T} - \bar{\theta_{TA}} = -\sum_{t=1}^T \alpha_t H_0^{-1}[ \nabla \bar{l_t}(\theta_{1:T}) - \nabla \bar{l_t}(\theta_t)]$.
On the right we therefore again find the gradient mismatch as the error. We have made this clearer in the paper and rewritten the equation, thank you for the comment!

Reviewer comment (Writing):
Section 3.1 is difficult to follow/understand, mainly because of several math annotation issues, and not well-organized paragraphs. For instance, the first two paragraphs can simply be phrased as `target model` definition rather than using unnecessary details and confusing notations.

Author response:
A: Thank you for the constructive comments! We have now rewritten the initial paragraphs of Section 3.1 to clarify the model definitions and split Section 3.2 into further subsections to ease readability. We also have fixed the notation mistakes, thanks again for pointing these out. For the final version, we will make sure to take your suggestions into account!

Reviewer comment (Presentation):
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 deferring questions of generality and applications to a later section.

Author response:
A: Thank you for the constructive comments! We have now rewritten the initial paragraphs of Section 3.1 to clarify the model definitions and split Section 3.2 into further subsections to ease readability. We also have fixed the notation mistakes, thanks again for pointing these out. For the final version, we will make sure to take your suggestions into account!

Reviewer comment (Novelty):
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.

Author response:
>Q1: 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? 

A1: In summary, finetuning can be highly resource-intensive, especially for large models, while model merging allows us to reuse already trained models. Finetuning on a new task can also lead to forgetting older tasks and retraining can be very costly at times, especially for removing data from a large pretrained corpus. To clarify, we have added a note on this in the section on data removal (Section 3.2.2); also see our response to Q1 of Reviewer idAX.


>Q2:  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 [...]

A2: Many optimizers, such as **Adam, already estimate the diagonal Fisher** which could be made available with the model, see [1; Section 11.2] for a discussion on how Adam approximates diagonal Fisher. Then, there is no overhead and if this estimate is shared, data can remain private. Thank you for pointing this out, we have emphasized this more in our discussion at the end of Section 3.3.

Reviewer comment (Experiments):
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.

Author response:
>Q3: 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.

A3: Thank you for the constructive suggestion, we have added RegMean to Table 3! In short, RegMean shows competitive performance to other well-tuned baselines and our method.