Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic
task arithmetic;weight disentanglement;merging models;multi-task learning
Yes
No ethics review needed.
See the weaknesses.
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.
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.
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.
Yes
No ethics review needed.
Refer to the weaknesses section.
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.
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.
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.
Yes
No ethics review needed.
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.
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.
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.
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.
Yes
No ethics review needed.
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?
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.
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.**
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.
@inproceedings{
anonymous2024finetuning,
title={Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=dj0TktJcVI},
note={under review}
}
For the past several years, *task arithmetic* has gained 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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
task arithmetic weight disentanglement merging models multi-task learning
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/6a40a853971b2fec756697c07d0f2f808933390d.pdf
transfer learning, meta learning, and lifelong learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/62e7d9150f5e63e8b16f38572eebda05ab72c6fe.zip
Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission