How to Weight Multitask Finetuning? Fast Previews via Model Merging

We thank the reviewer for the questions, bellow we provide our answers.

Q1: The method relies on a standard application of Bayesian optimization with a surrogate model, lacking novel contributions or distinctive techniques

A1: We strongly disagree. We do not use “Bayesian Optimization” which is a technique for global optimization and search and often uses Gaussian Process surrogates over function values (no gradients involved). In our work, we use exponential-family posterior approximations as surrogates, which is fundamentally different from the Gaussian-Process surrogates used in Bayesian Optimization.

Q2: In the steps listed between lines L148 and L151, while the three initial steps are understandable, the sequence is left open-ended. It’s unclear what action follows the third step—what process or computation does the algorithm proceed with after generating the previews?

A2: The main use of the preview is to choose the good $\boldsymbol\alpha$ values and use them to run full multitask finetuning. We have added a few lines to clarify this in the revised version; see L39, L91, and L145. The preview is also useful to analyze the trade-offs among tasks.

Q3: In the section between lines L270 and L280, the definitions and relationships among variables are unclear. For instance, the connection between $\hat{\pi}\_{tk}$ and $\pi\_{k}$ should be clarified. Additionally, if $\mathbf{h}\_{tk}$ is a scalar, it should not be in bold, as this can lead to confusion.

A3: Please see the line right after Eq. 12, where the definition of $\hat{\pi}\_{tk}$ is given. Also, $\mathbf{h}\_{tk}$ is the diagonal of the Hessian, and not a scalar.

Q4: There is a lack of a clear evaluation criterion. What mathematical definition is used to judge “better previews,” and what specific objective or performance metric is this algorithm designed to optimize?

A4: See Table 1 where we use the MSE between the Exact-Solution contour and the previews obtained using model merging; this is also mentioned in Line 360. This is a clear objective criteria (although other metrics can be used). As the objective of the algorithm is concerned, the proposed algorithm tries to find the best merged model for each $\boldsymbol\alpha$ separately. This is done by using a variational approach (that is, by minimizing KL).

Q5: The numbers in the figure (e.g., 90, 85, 80, ...) are labeled as accuracy, but it is unclear what these accuracy values represent. Are they calculated based on an equal weighting across three tasks? Please define the metrics and explain the weighting, if any, to clarify the interpretation.

A5: Yes, we report test accuracy over the test set containing all the tasks together.

Q6: The three tasks listed—French, Spanish, and English—are somewhat unclear in meaning. What does “accuracy for English” specifically indicate in this context? Please provide a clear explanation of how accuracy is determined for each language task and what each result represents.

A6: We are not sure which figure/result the reviewer is referring to; there is no “French, Spanish and English” together in any figure or experiment. For example, Fig. 1 refers to “Spanish, German, French” (no English). Similarly, Fig. 8 contains “German, English, French” (no Spanish). In general, “accuracy” in a language (like English) represents the model accuracy evaluated by the BLEU metric, which is a standard text-generation evaluation metric that measures overlap statistics between a generated and ground-truth output. We could clarify further if you tell us the exact section where you find the problem.

Q7: As a reader, the takeaway from the experimental section remains ambiguous. There doesn’t appear to be a consistent or clearly stated conclusion. Could you provide a concise summary of the findings or implications drawn from this section to clarify its purpose and insights? Otherwise, the paper seems to give me an impression of an application of Bayesian optimization.

A7: The summary is given in Table 1 and the summary of the finding is stated in L361-362. Every section also contains such a summary (for example, L374 in Sec 4.1).

Dear Authors,

Thank you for sharing your work. I have carefully reviewed the manuscript and would like to provide feedback for consideration.

Feedback:

1. (Q1) Clarification on Bayesian Optimization:

   • I understand you are using "exponential-family posterior approximations as surrogates." However, this still an application of Bayesian Optimizations.

2. (Q2) Evaluation Methodology Across Tasks ($T$):

   • It is stated that "the main use of the preview is to choose the good alpha values and use them to run full multitask finetuning." However, I am unclear about the evaluation methodology across tasks.
   • Specifically, is the evaluation metric $\sum_t l_{\text{test},t}$, $\sum_t \alpha_t l_{\text{test},t}$, or something else entirely? This point requires further elaboration to improve clarity.

3. (Q4) Model Merging Comparisons:

   • You propose a new model merging method, yet there is no direct performance comparison with existing methods (e.g. ties-merging, DELLA-Merging, Ada-Merging, DARE). This omission weakens the practical relevance of the study.
   • While MSE between the Exact-Solution contour and the previews is a useful intermediate metric, performance ultimately matters more. Including metrics like win rate or average performance scores for merged models would significantly strengthen the narrative.

4. (Q7) Posterior Approximation and Multitask Finetuning:

   • Regarding the statement:

     "Overall, better posterior approximations match multitask finetuning more closely both in terms of overall shape of the solution (Table 1) and the performance of the best weightings (Table 2)."

     • Why does this matter in practical terms? If your proposed method consistently outperforms other merging techniques, that would explain why the Bayesian Optimization framework used. However, without direct comparisons to baseline merging methods, the importance of this step is less compelling.

Additional Comments

1. Experimental Design:

   • L396: You describe the experiment using ResNet-20 on CIFAR-10 with three tasks constructed from different sets of classes. It would be helpful to explicitly define the criteria for selecting these tasks to avoid concerns about cherry-picking.
   • Since the computational cost of this experiment appears low, I suggest iterating over all task combinations and reporting averages and standard deviations to improve the robustness and believability of results.
   • The study focuses on a low-dim task space. How does the method scale to higher-dim search spaces, and what is the search efficiency? (Existing baselines are on 8 number of tasks rather than 2-3.)
   • The absence of direct comparisons to existing baseline methods on the performance of merged models is a significant gap. Such comparisons are crucial for validating your approach.
   • L177: Provide more detail on why the Hessian is regarded as an identity matrix in this context.

2. Mathematical Notation Casualty:

   • While the mathematics is well-developed, the notation can feel casual and at times ambiguous, making it difficult to follow.
     • Example 1: $\theta_{\alpha}$ appears at L105 but is not defined. Readers are left to infer its meaning. Additionally, subscripts of $\theta$ alternate between $t$ and $\alpha$, which are not equivalently positioned conceptually.
     • Example 2: $\widehat{\mathcal{R}}_0$ is introduced at L161 without context. Why is there a hat, and why is there a 0? There is no $\mathcal{R}_0$ in the whole paper. I can guess it can be regarded as the unnormalized probability at step 0. But it could be hard to follow if readers need to guess everywhere. Clear definitions are necessary to avoid unnecessary guesswork. (This also applies to the main text. E.g. L180, when you say one point, what does it mean? After reading a few paragraphs, I am able to guess that the "point" here means model weights, but I think it jams the reading.)
     • Example 3: The use of $\Rightarrow$ in equations (8), (9), (11), and L257 lacks formal consistency. I understand that when doing induction on draft papers, it is very common to use $\Rightarrow$. However, I think this is not very formal in the manuscript to be published. Furthermore, the numbering of equations (L257) is incomplete.

3. Experimentation Completeness:

   • While the authors have conducted extensive experiments, the lack of baseline comparisons is a significant shortfall for an application-focused paper. Incorporating baselines is essential to demonstrate the practical utility of the proposed method.

Overall Assessment:

This paper demonstrates an application of modified Bayesian Optimization for model merging. However, because of the lack of comparisons with baseline methods, incomplete clarity in certain areas, and casual mathematical notation reduce its readiness, I decide to maintain the origin score.

Thank you for your understanding.

P.S.

• Typo
  • L187: "but how can be design them?"

We thank the reviewer for the questions, bellow we provide our answers.

Q1: It is unclear whether the weights chosen for each task by the method will be used for further multitask finetuning.

A1: Yes, the main idea is to choose weights using previews and then use them to jointly perform multitask finetuning only once. We have added a few lines in Sec 1 and 3 to clarify this; see L39, L91, and L145 in the revised version.

Q2: If they are, the method appears to be less computationally efficient because it requires training the models for each task initially.

A2: We believe there is a misunderstanding here. Our approach aims to reduce the cost of exhaustive search over the alpha-space by using separately trained models on each task. Separate training needs to be done just once and is sufficient to quickly generate the estimates for multitask finetuning with a variety of alpha-values. For example, for the experiment in Fig. 5, it takes seconds to create estimates compared to around 50 mins to run a single multitask finetuning run. Doing such finetuning on 100s of tasks is too expensive. In our approach, training each single-task model takes around 15 mins, and this only needs to be done. With this, creating estimates just takes seconds. Overall, a single run of multitask finetuning via previews has similar runtime to two runs of multitask finetunings (for instance, two settings of $\boldsymbol\alpha$ values)

We have included such exact runtimes in L431 and L523.

Q3: If [the method is not used for multitask finetuning, then] the accuracy of the merged model seems to lag significantly behind joint training on average, as shown in Figures 4, 5, and 6, suggesting that a simple grid search of joint training could outperform model merging. Additionally, the paper lacks comparisons with other reweighting or model merging methods.

A3: This question is not valid because our goal is to just select good weights and then run multitask finetuning with those weights. See our response to Q1.

Q4: On average, how long does it take to converge for Simple Merging, Hessian-Weighted Merging, and Mixture-Weighted Merging?

A4: For both simple and Hessian-weighted merging, we have closed form expressions so there is no notion of convergence. It takes seconds to do this operation; see L431 and L523 for two examples. For mixture-weighted merging, we need to do Hessian-weighted merging a few times, but often very few, for example, the ViT merging experiments required 2 iterations.

We thank the reviewer for the positive score. Below we provide answers to their questions.

Q1: There isn't enough analysis of computation overhead of the proposed approach.

A1: We have added a few lines in Sec 3.5 (see L314-321 of the revised version) for computational overheads. We have also added absolute runtime in Sec 4.3 and 4.5 (L431 and L524). To summarize, the IVON-Hess method is the cheapest, there is virtually no cost to compute the parameters for each finetuned model and their associated diagonal Hessian estimate (denoted by $h_t$). Merging requires only element-wise vector multiplication to get $h_t \cdot \theta_t$ and elemen-twise addition to add them all together. All such computations are very cheap and scale linearly in the number of parameters.

The multiIVON-Hess, requires multiple (say up to $K$) IVON-Hess runs to obtain $K$ components and multiple iterations (say $I$), which increases the cost by $K\cdot I$ times. The AdamW-SG method requires an additional pass through the data to get the “squared-gradient” estimate of the Hessian.

We will add these details in the camera ready version if accepted.

Q2: Can this approach be widely adopted for multitask learning training?

A2: We believe it can. Choosing good weights through exhaustive search is time-consuming and previews are excellent cheap alternatives. Being able to choose better weights directly impacts accuracy. For example, in Fig. 4, the best accuracy of joint training (Exact Solution) is around 70% which goes down quickly to 65% and lower as we move away from the best weights.

Q3: Do you plan to release the code? Is it easy for people in the community to use your code for their own multitask training jobs?

A3: Yes, we will release the code upon acceptance. We also plan to integrate the code on the mergekit library (https://github.com/arcee-ai/mergekit/).

Q4: What are the limitations? The paper doesn't have enough discussions of its limitations. Presenting limitations would give readers a better and more comprehensive understanding of your approach.

A4: We have included a limitations statement as Section 5 (Line 538 in the new version). The main limitation is that the number of components and training steps for the mixture surrogates may be large in some cases. In many cases, good results are obtained with 5-10 steps but this is not always the case. This could also be prohibitive when storing more than one copy of the model is too costly.

We thank the reviewer for the questions, bellow we provide our answers.

Q1: There is limited quantitative analysis of the computational savings(or practical time savings) across different tasks compared to exhaustive grid search.

A1: We have added a few lines in Sec 3.5 (see L314-321 of the revised version) for computational overheads. We have also added absolute runtime in Sec 4.3 and 4.5 (L431 and L524). As an example, for experiments in Fig.5, joint multitask finetuning takes around 50 minutes for each $\alpha$ value, but merging takes only seconds (plus around 14 to 17 minutes for finetuning on each task separately). Overall, a single run of multitask finetuning with previews has similar runtime to two runs of multitask finetunings (around 100 minutes for Fig. 5).

Q2: The mixture-based approach may become computationally intensive for large models or when dealing with a large number of tasks.

A2: We agree that the computation increases a lot but the approach still remains much cheaper than exhaustive search over the $\boldsymbol\alpha$-space by explicitly training on all tasks together. For instance, for Fig. 5, merging with mixtures still takes a few seconds, while the joint multitask training takes around 50 minutes.

Q3: Some trends have not been precisely captured by either Hessian-weighted merging or simple merging (Fig 6(b) & Fig 6(c)).

A3: In Fig. 6a and 6b, we see clear improvements by Hessian-weighted over simple averaging for $\alpha$ values from 0 to 0.5. We agree that in 6c the trends are similar. In general, it is true that we may not always see gains in accuracy of the preview when increasing the posterior class.