Model Merging by Uncertainty-Based Gradient Matching

We thank the reviewer for their in-depth review and appreciation for our method. The main criticism is lack of comparison to a recent method TIES-merging (accepted at NeurIPS 2023) and experimental results being weak. We disagree with both but, to address the reviewers’ concerns, we have added a discussion to TIES-merging and an additional comparison in Table 3.

Q1: 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.

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$.

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

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.

Q3.1: 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.

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.

Q3.2: An experimental setting from past papers like Task Arithmetic, TIES-Merging can be adopted for such experiments.

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.

Q4: 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?

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)

We thank the reviewer for their review, appreciation of the usefulness of our method, and for pointing out various typos. Their suggestions have helped us greatly in improving the writing of Section 3. We hope that the reviewer would consider increasing their score if this answers their concerns. Below we address all questions:

R3’s suggestions on improving the presentation in Section 3.

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!

The discussion in the last paragraph in page 3 is best to be had in the experimental section with some data. or under its own section with further details.

A: Thank you for the suggestion. For now we have kept this paragraph at the end of Section 3 but we will add an additional experiment for the final version and add further discussion. Our experiments already outline the claim at the end of the paragraph: we can use gradient mismatch to improve model merging by reducing it, for example, in Figure 1 (right) or Table 2.

Q1: 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.

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

Q2: what does $t$ stand for in Eq (8)?

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.

Q3: What is the error between $\theta_{TA} \text{ and } \theta_{1:T} = \theta_{1:T}$? Please clarify/correct?

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!

R3: Experimental results show comparable results.

A: We would like to emphasize that we get consistent improvements both in terms of performance over a well-tuned task arithmetic baseline, with scaling factors determined on the test data, and better robustness to choice of scale. This was the goal of reducing gradient mismatch. Furthermore, we would like to point out the improvements for data removal shown in Table 4, which are consistent across different tasks and models.

We thank the reviewer for their review and their appreciation for the value of our contribution. Below we answer all questions:

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.

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.

Q4: 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?

A4: For different hyperparameters, the performance will likely further improve. For our method, we do not tune the hyperparameters and use the same learning rate used in [2] for RoBERTa and the same default hyperparameters for AdamW taken from the transformers library [3].

References

[1] J. Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research 21.1 (2020): 5776-5851.

[2] Y. Liu, et al. Roberta: A robustly optimized bert pretraining approach. arXiv 1907.11692, 2019.

[3] T. Wolf, et al. Transformers: State-of-the-art natural language processing. Proceedings of the conference on empirical methods in natural language processing: system demonstrations, 2020.

We thank the reviewer for their review and helpful suggestions. Below we address the questions they asked.

Q1: Finetuning 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.

A1: Thank you for your suggestion! We already have a comparison to finetuning in Table 3. There, the “all-data” baseline is finetuned on all datasets at once. In general, finetuning can be resource-intensive, while model merging allows one to reuse already trained models. Finetuning on a new task can also lead to forgetting older tasks. Retraining is often costly, especially for removing small subsets of a large pretraining corpus from a model. Then, merging is much more efficient than finetuning. We have added further discussion of this in Section 3.2.2. We also stress that estimating Hessians does not require data, because we can do this in an online fashion during training (see our response to the next question Q2).

Q2: Again, 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.

A2: The Hessian calculation is practical. Our squared gradient approximation to the Hessian only incurs a small overhead (a pass over a subset of the data is often enough) but it can also be obtained in an online fashion. For example, when running second-order optimizers (or approximations thereof such as Adam, see [1; Section 11.2]), an approximation of the Hessian is calculated during training as a “by-product”. See also our response Q9 to Reviewer 1bWs. We have emphasized this further in our paper (end of Section 3.3) and hope that this makes it clearer.

References

[1] J. Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research 21.1 (2020): 5776-5851.