Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic

Authors: Ruochen Jin, Bojian Hou, Jiancong Xiao, Weijie J Su, Li Shen
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: task arithmetic, weight disentanglement, merging models, multi-task learning

Abstract:
In recent years, *task arithmetic* has garnered increasing attention. This approach edits pre-trained models directly in weight space by combining the fine-tuned weights of various tasks into a *unified model*. Its efficiency and cost-effectiveness stem from its training-free combination, contrasting with traditional methods that require model training on large datasets for multiple tasks. However, applying such a unified model to individual tasks can lead to interference from other tasks (lack of *weight disentanglement*). To address this issue, Neural Tangent Kernel (NTK) linearization has been employed to leverage a ''kernel behavior'', facilitating weight disentanglement and mitigating adverse effects from unrelated tasks. Despite its benefits, NTK linearization presents drawbacks, including doubled training costs, as well as reduced performance of individual models. To tackle this problem, we propose a simple yet effective and efficient method that is to finetune the attention modules only in the Transformer. Our study reveals that the attention modules exhibit kernel behavior, and fine-tuning the attention modules only significantly improves weight disentanglement. To further understand how our method improves the weight disentanglement of task arithmetic, we present a comprehensive study of task arithmetic by differentiating the role of the representation module and task-specific module. In particular, we find that the representation module plays an important role in improving weight disentanglement whereas the task-specific modules such as the classification heads can degenerate the weight disentanglement performance.

Decision: accept (poster)
Reviewer scores: 8, 3, 6, 8

Reviews and discussion:
<thread>
creation date:
1730012184414

thread_type:
official_review

<thread_object>
creator: reviewer_rsur

<content>
summary:
The authors propose to improve the performance of weight-interpolation-based model merging methods such as Task Arithmetic by finding a sub-module within the deep neural networks. I strongly agree with this idea because, for existing work based on NTK linearization, the model's performance on individual downstream tasks has been compromised in order to achieve better performance in multi-task model merging. This trade-off is usually unacceptable in practical applications, where ensuring performance on downstream tasks is paramount for fine-tuning. However, there are some typos and minor mistakes in this manuscript, which need to be addressed to improve the overall clarity and professionalism of the document.

Therefore, I would like to assign this manuscript an overall score of 6, indicating that it requires minor revisions. However, **I am open to raising the score if the authors address the specific concerns and suggestions outlined in my detailed review.**

strengths:
1. The motivation of this study is clear. The authors seek to identify a subset of parameters that can undergo non-linear fine-tuning while retaining linear characteristics, such as those observed in NTK linearization. Consequently, the attention modules, particularly the linear layers within them, are selected for this purpose.
2. The experimental results presented in Figure 3 demonstrate the validity of the proposed method.
3. Models of varying sizes (ViT-B-32, B-16 and L-14) are employed for experimental evaluation.
4. The authors conduct experiments on both the vision domain and NLP tasks.

weaknesses:
1. The task arithmetic presented in Definition 1 (lines 167-175) deviates from standard practice. In the original paper [1], task arithmetic is defined as $\theta_{unified} = \theta_0 + \lambda \sum_{i=1}^T \tau_i$, employing only a single scaling factor. Otherwise, it should be explicitly stated as task-wise AdaMerging [2].
2. Missing scaling factor of task vectors in lines 52-53 and lines 353-357.
3. Is the summation symbol missing on the right-hand side of Eq.(1) (lines 196-199)? 
$$ RHS = \sum_{t=1}^T \dots $$
4. In lines 400-401, the "disentanglement error" was first proposed by Ortiz-Jimenez et al. (2024).
5. Recent work on deep model fusion is missing from the related literature section. The authors may refer to these surveys [3,4,5] for additional information.
6. A minor suggestion: The results for NLP tasks (Table 4) should be included in the main text.
7. The ablation study in Appendix B.4 is not mentioned in the main text. Consider including it in the main text.

[1] Ilharco et al. Editing models with task arithmetic. In International Conference on Learning Representations. ICLR 2023. 
[2] Yang et al. AdaMerging: Adaptive Model Merging for Multi-Task Learning. ICLR 2024.
[3] Tang et al. FusionBench: A Comprehensive Benchmark of Deep Model Fusion. 
[4] Yadav et al. A Survey on Model MoErging: Recycling and Routing Among Specialized Experts for Collaborative Learning.
[5] Yang et al. Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities.

questions:
1. The differences between Definition 2 in the manuscript and the original definition in Ortiz-Jimenez et al. (2024) remain unclear. Could the authors elaborate further on these differences?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730353441089

thread_type:
official_review

<thread_object>
creator: reviewer_6sac

<content>
summary:
This paper proposes a new method to enhance weight decoupling in task arithmetic by fine-tuning only the attention module in the Transformer model. This method improves the weight decoupling ability of the unified model while maintaining the performance of individual models, thereby improving the performance of task arithmetic.

strengths:
Innovation: The proposed method of fine-tuning only the attention module is innovative in the field of task arithmetic and can effectively solve the weight decoupling problem.
Experimental design: The experimental part covers multiple data sets and tasks, enabling a comprehensive evaluation of the effectiveness of the proposed method.
The results are remarkable: the method proposed in the paper achieves better performance than existing techniques on multiple benchmarks.

weaknesses:
Parameter sensitivity analysis: The paper mentions the effect of the choice of α coefficient on performance, but does not provide a detailed sensitivity analysis. It is recommended to add experiments on the effect of the choice of α coefficient on model performance.
Computational resource consumption: The paper mentions the efficiency advantage of the proposed method, but does not provide a direct comparison with existing methods in terms of computational resource consumption. It is recommended to add analysis in this regard, especially memory and time consumption in actual deployment.

questions:
While the paper mentions some related work, the relationship to existing methods could be discussed in more detail, especially the latest advances in model fusion and weight interpolation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730511405718

thread_type:
official_review

<thread_object>
creator: reviewer_cquv

<content>
summary:
The paper addresses task interference issues in task arithmetic, which arise when combining weights from multiple tasks into a single unified model. To tackle this issue, it uses Neural Tangent Kernel (NTK) linearization and trains only the attention layers in the transformer to reduce training costs. Through several experiments, the paper demonstrates the efficiency of the proposed approach and explains task arithmetic by distinguishing the roles of the representation module and the task-specific module.

strengths:
1. The paper is easy to follow and provides clear explanations of essential concepts, making it easier to understand both the proposed approach and the fundamentals of task arithmetic.

weaknesses:
1. The paper appears to be a naive combination of prior work, specifically [1] and [2], with modifications focused on training parameter selection (e.g., training only the attention layers in transformers), which limits its novelty.

[1] Multitask prompted training enables zero-shot task generalization (ICLR 2023)
[2] Task arithmetic in the tangent space: Improved editing of pre-trained models (NeurlPS 2024)

2. The main contribution claimed by the authors—asserting that weight disentanglement primarily arises from the representation module while the effectiveness of task arithmetic is constrained by task-specific components—is not well-supported. Additional ablation studies, analysis, or theoretical support are needed to substantiate this claim. Depending on interpretation, the proposed contribution may be a trivial insight already known in the fields of task arithmetic or multi-task learning.

3. The experiments are quite limited, as many related works cited by the authors are missing from the experimental results, particularly in Tables 2 and 3. This omission reduces the validity of the proposed method.

4. The paper's theoretical contributions rely heavily on [2], which limits its originality and overall contribution.

questions:
Refer to the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730581972634

thread_type:
official_review

<thread_object>
creator: reviewer_s2kf

<content>
summary:
This paper presents a method that fine-tunes only the attention modules of models to enhance weight disentanglement in task arithmetic. Task arithmetic is a training-free approach that enables new tasks by linearly combining the weights of a task-specific model with those of the original model. Building on NTK studies on kernel behaviour, the proposed method introduces an approximate kernel behaviour test to verify that attention modules exhibit this characteristic. The authors validate their hypothesis by fine-tuning only the attention modules in CLIP models, showing that this approach outperforms other fine-tuning strategies as well as the original NTK-initialization method.

strengths:
The paper is well-organized and clearly written. The authors thoughtfully introduce relevant prior work, providing clear definitions and context for task arithmetic, weight disentanglement, and their associated challenges. Each argument is well-supported, with compelling visualizations that reinforce the method's performance. The progression from motivation to methodology is logical and engaging.

weaknesses:
The paper has no prominent weaknesses. Here are some other open questions:
 
1. Contribution of Individual Attention Modules. 
While the method targets attention modules to achieve weight disentanglement, it remains unclear whether each module contributes equally to task performance. In large models, which may have over 100 billion parameters, understanding the relative contribution of each attention module could make the approach more efficient, possibly allowing for sparse fine-tuning of only the most impactful modules. An analysis or ablation study on the importance of individual modules could provide further insights into how weight disentanglement varies within the attention layers.

2. Generalizability to Non-Attention-Based Architectures. 
While the method demonstrates promising results on attention-based models, it would be useful to discuss the potential for generalizing this approach to non-attention architectures, such as convolutional or Mamba-based models. Would the same strategy apply effectively, or would alternative fine-tuning techniques be required for architectures without attention mechanisms? Providing an outlook on this point could help extend the paper's relevance across a wider variety of model architectures.

questions:
See the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732165004859

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We would like to express our gratitude for the insightful feedback provided by the reviewers regarding our manuscript. We appreciate the thorough evaluations and constructive suggestions, which will help us enhance the quality of our work. Below, we outline our responses to the key points raised and the revisions we have implemented.

1. **Revision of Related Work (Raised by 6SAc, Rsur)**: We incorporate references to relevant surveys and studies, such as those by Yadav et al. (2024) and Yang et al. (2024), to provide a more comprehensive overview of the current landscape in model merging techniques. Existing model merging techniques can be broadly categorized into two main types (Yang et al., 2024): (i) Pre-Merging Methods: These methods focus on enhancing the conditions necessary for effective model merging by optimizing the fine-tuning process of individual models. (ii) During Merging Methods: These approaches address task conflicts and interference through various strategies before executing the parameter merging operations.
    
    This Modification will clarify how our work fits into the broader context of existing research.
    
2. **Inclusion of NLP Results (Raised by Rsur)**: We recognize that the results for NLP tasks presented in Table 4 should be included in the main text to ensure a complete understanding of our findings. We have revised the manuscript to integrate these results, highlighting their significance alongside the vision tasks.
3. **Ablation Study in Main Text (Raised by Rsur)**: We appreciate the suggestion to mention the ablation study detailed in Appendix B.4 within the main text. We have included a summary of this study in the relevant section, emphasizing its importance in demonstrating the robustness of our approach and providing additional insights into parameter sensitivity.
4. **Parameter Sensitivity Analysis (Raised by 6SAc):** We have also included a parameter sensitivity analysis of the coefficient $\alpha$ in the main text. This analysis evaluates how different choices of $\alpha$ impact model performance and weight disentanglement. The results demonstrate the robustness of our method across a range of $\alpha$ values, further validating our hypothesis regarding the importance of careful parameter selection in optimizing task arithmetic performance. 
5. **Addressing Typos and Minor Errors (Raised by Rsur)**: We conducted a thorough review of the manuscript to correct any typos or minor mistakes, enhancing overall clarity and professionalism.

We believe these revisions will address the reviewers' concerns effectively and strengthen our manuscript significantly. Thank you once again for your valuable feedback, and we’re happy to address any further questions.

Yadav et al. A Survey on Model MoErging: Recycling and Routing Among Specialized Experts for Collaborative Learning. 

Yang et al. Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732458830184

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Follow up on Rebuttal

comment:
Dear Reviewers,

We sincerely appreciate the time and effort you have dedicated to reviewing our work.

We have thoroughly addressed all the concerns raised, including the related work highlighted by Reviewer 6SAc and the parameter sensitivity analysis noted by Reviewer 6SAc. Additionally, we have incorporated the NLP and ablation results into the main text, as suggested by Reviewer Rsur. We have also made significant efforts to clarify our method and contributions, as emphasized by Reviewer CQUv.

As the discussion period comes to a close, we want to ensure that our responses have fully addressed your questions and concerns. Please let us know if there are any aspects of our work that still require clarification or further discussion.

Thank you once again for your valuable feedback.

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary
The paper studies task disentanglement in neural networks whereby new tasks can be solved by linearly combining weights of task-specific models. The authors build upon NTK and show that finetuning only the attention modules of a model is amenable to task disentanglement. The method is verified on CLIP models and outperforms finetuning strategies.

Strengths
1. The authors study task disentanglement and show that attention modules can be finetuned but still retain a linear entanglement property.
2. The experimental setup in this work studies different sized models on image classification and NLP tasks.
3. The paper is presented and written well

Weaknesses
1. The approach requires tuning a hyperparameter \alpha and is sensitive to its value as the author rebuttal also shows. An automatic tuning approach for this would make the work more impactful.
2. The approach is only applicable to models with an attention layer. 

Justification
The paper received a majority of positive reviews. The author rebuttal answered the concerns from the reviewers. The paper is technically solid and presents an interesting approach for task disentanglement.