Mastering Task Arithmetic: $\tau$Jp as a Key Indicator for Weight Disentanglement

Authors: Kotaro Yoshida, Yuji Naraki, Takafumi Horie, Ryosuke Yamaki, Ryotaro Shimizu, Yuki Saito, Julian McAuley, Hiroki Naganuma
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: task arithmetic, model editing, task vector

Abstract:
Model-editing techniques using task arithmetic have rapidly gained attention.
Through task arithmetic, simply through arithmetic operations on the weights of pre-trained and fine-tuned models create desired models, such as multi-task models, models in which specific tasks are unsolvable, or domain-transferred models.
However, task arithmetic faces challenges, such as poor reproducibility and the high cost associated with adjusting coefficients in the arithmetic operations on model parameters, which have limited its practical success. In this paper, we present three key contributions in the context of task addition and task negation within task arithmetic. First, we propose a new metric called $\tau$Jp which is based on the product of the task vector ($\tau$) and the Jacobian of the pre-trained model with respect to its weights. We show that $\tau$Jp has a causal relationship with the interference that occurs from arithmetic operations. Second, we show that introducing regularization to minimize $\tau$Jp significantly mitigates interference between task inference, which leads to the elimination of coefficient tuning and improved accuracy on each task.
Third, in the context of incremental learning, we demonstrate that our $\tau$Jp regularization achieves more robust performance in environments where access to future tasks is unavailable, thus validating the scalability of the approach.
Finally, we demonstrate that the $\tau$Jp regularizer further reinforces the performance of task arithmetic by leveraging publicly available fine-tuned models, offering practical benefits for real-world applications.
Our code is available at https://github.com/katoro8989/tau-Jp_Task_Arithmetic

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

Reviews and discussion:
<thread>
creation date:
1729265357401

thread_type:
official_review

<thread_object>
creator: reviewer_n6zh

<content>
summary:
The paper presents a novel approach to task arithmetic in neural networks, which leverages a novel metric that quantifies the relationship between task vectors and the Jacobian of pre-trained models. The authors claim that by minimizing this metric through regularization, they can significantly reduce interference between task predictions and enhance the accuracy of task arithmetic operations. The experimental results demonstrate substantial improvements in performance for both task addition and task negation.

strengths:
+ This paper is well-written and easy to follow.
+ The experiments are extensive, and the results sound good.
+ The design of metric τJp is reasonable and interesting.

weaknesses:
- The novelty of this paper may be limited. My consideration is that this paper seems to fundamentally align with the approach proposed by Ortiz-Jimenez et al. (2023) [1], which also emphasizes fine-tuning models in the tangent space. Although using the specific regularization term, this paper does not sufficiently differentiate itself from this existing work.
- While the empirical results are compelling, the paper lacks a thorough theoretical explanation for why the proposed regularization leads to better performance compared to other methods, such as those discussed in Ortiz-Jimenez et al. (2023). I am confused about why a simple and soft regularization results in such improvement compared to [1]. A deeper theoretical analysis could strengthen the paper's contributions.
- The authors briefly mention tuning the regularization strength but do not provide sufficient details on how this hyperparameter was selected. The sensitive analysis of this hyperparameter is also necessary for the paper.

[1] Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. Advances in Neural Information Processing Systems, 2023.

questions:
Could the proposed regularization affect the model's plasticity? Specifically, how might the addition of this regularization impact the fine-tuning performance, potentially influenced by the strength of the regularization?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730307802728

thread_type:
official_review

<thread_object>
creator: reviewer_rgly

<content>
summary:
The paper introduces a novel metric, $\tau \text{Jp}$ ($\tau$-Jacobian product), to improve understanding of weight disentanglement in task arithmetic. It demonstrates that τJp inversely correlates with normalized accuracy, suggesting it as an indicator for weight disentanglement. A regularization technique is proposed to minimize τJp during fine-tuning, effectively reducing the need for coefficient adjustments in task addition and negation. It also proves valuable in incremental learning scenarios where future tasks are unknown.

strengths:
1.The paper addresses an important and timely topic: in an era where foundation models are prevalent, better understanding weight disentanglement is particularly valuable for enhancing the practical applicability of these models.

2.The proposed metric offers a deeper understanding of weight disentanglement, and the regularization method effectively reduces task interference, minimizing the need for coefficient adjustments.

3.The success of the proposed method in incremental learning scenarios aligns well with real-world applications, demonstrating its scalability and practical relevance when future tasks are unknown.

weaknesses:
1.While the paper introduces the $\tau \text{Jp}$ metric and explains its relationship with weight disentanglement, the theoretical justification for why $\tau \text{Jp}$ regularization effectively reduces task interference could be further elaborated.

2.The proposed regularization method lacks a comparison with other existing regularization techniques, which makes it difficult to fully assess its relative strengths and weaknesses.         

3.The paper mentions task addition, task negation, and task analogies in the introduction and background sections as key operations in task arithmetic, but there are no experiments evaluating task analogies. This inconsistency weakens the completeness of the experimental validation.

questions:
Suggestions：

1.The related work section could be improved by explicitly connecting prior studies to this paper's contributions, emphasizing how the proposed method addresses existing limitations.               
2.Consider moving the related work section after the methods section, especially since the current structure delays the introduction of the proposed method until page 5. This change would allow readers to quickly understand the proposed approach before diving into comparisons, enhancing readability and engagement.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730381300804

thread_type:
official_review

<thread_object>
creator: reviewer_ontx

<content>
summary:
This paper proposes a new metric called $\tau$Jp ($\tau$-Jacobian product) for measuring weight disentanglement in task arithmetic operations on neural networks. The authors theoretically analyze the relationship between $\tau$Jp and interference between tasks, and introduce a regularization method based on minimizing $\tau$Jp during fine-tuning. Experiments on image classification tasks demonstrate improved performance and reduced need for hyperparameter tuning compared to existing methods.

strengths:
1. The paper provides a comprehensive theoretical and empirical study of the relationship between the proposed $\tau$Jp metric and task interference in neural networks.

2. The introduction of $\tau$Jp as a new metric for weight disentanglement is novel and well-motivated.

3. The proposed regularization method eliminates the need for tuning inference-time hyperparameters ($\alpha$), which is a practical advantage.

weaknesses:
1. The method requires access to data from all other tasks during training, which is often unavailable in realistic task arithmetic scenarios. This limits the practical applicability of the approach.

2. The computational cost of calculating τJp is likely very high, as it involves multiple Jacobian-vector products. The paper does not report runtime or resource requirements, making it difficult to assess scalability.

3. Experiments are limited to image classification tasks. Evaluation on other domains like language tasks would strengthen the claims of generality.

4. The derivation of Equation 7 from the weight disentanglement definition is non-trivial and should be explained more clearly.

questions:
- How does the computational cost of the proposed method compare to existing approaches?
- Can the method be adapted to work with limited or no access to data from other tasks?
- How well does the approach generalize to other domains beyond image classification?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671313878

thread_type:
official_review

<thread_object>
creator: reviewer_hopm

<content>
summary:
The paper tackles the task of task arithmetics, i.e. how to combine *task vectors/parameters* to form a multi-task model. A key issue is to determine the best combination weights as to minimise interference between tasks, and maximise sharing of information / positive transfer.
More specifically, the authors make use of two previously introduced notions: **(i)** the notion of **weight disentanglement** which was  proposed as a measure of task interference in task arithmetic. And **(ii)** the Neural Tangent Kernel (NTK) which designates a training regime where parameter updates can be expressed with a linearised approximation.
Previous works have suggested that performing task arithmetics under the NTK regime can lead to better MTL performance. the authors investigate this behaviour in more depth. Based on this analysis, they also propose a regularisation technique to further reduce task interference when performing task arithmetic, which involves slightly fine-tuning the task vectors themselves.

strengths:
* The paper is well motivated and grounded in previous related work
* The proposed method is simple and could adapt to different task arithmetic variants
* Interesting insights on the link between the proposed regularisation and weight disentanglement
* A more efficient implementation of the method is proposed for handling larger number of tasks (Equation 11)

weaknesses:
* Missing baselines on more recent task arithmetic work: The main tables should include some recent task arithmetic results (e.g. TIES-Merging and AdaMerging) as well as standard single task and MTL baselines (although it is in the appendix), if only to better understand the existing gap in performance.
* Missing discussion about the extra cost: The paper briefly mentions efficiency of the method (e.g. equation 11 or line 364), however I think this could be discussed in more depth: On the one hand, the proposed method seems more robust to task coefficients $\alpha$, which could save on hyper parameter tuning; On the other hand, it involves a fine-tuning procedure which requires knowledge/access to all tasks simultaneously (Equation 10) as opposed to directly combining task vectors obtained independently from one another.

questions:
* I find the notion of "One-shot" and "fine-tuned" experimental setting could be improved; First because the notion of fine-tuning can become confusing between the coefficients $\alpha$ vs the model parameters $\theta$ fine-tuning. Second, because it is not clear if it is referring to a specific method/objective for fine-tuning the task coefficients (e.g. AdaMerging or others) or simply hyperparameter search.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732275681918

thread_type:
official_comment

<thread_object>
creator: area_chair_r1cm

<content>
title:
Discussion

comment:
Dear reviewers,

The authors have responded to your reviews. 

Until November 26th @ 2359 (AOE time) reviewers and authors can freely exchange responses, so if there any clarifications you require from the authors, now is the time to seek them!

Best,

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper is concerned with task arithmetic, and to this end the authors propose the $\tau$-Jacobian product ($\tau$Jp) metric which relates to weight disentanglement. They show that by adding a regulariser to minimise this quantity they can improve the performance of task arithmethic in practice.

In terms of strengths, reviewers thought this work was well-motivated,  and the $\tau$Jp metric was seen as "novel", "well-motivated", "reasonable", and "interesting". The proposed regularisation method was well-received for its practical benefits (it eliminates the need for tuning hyperparameters at inference time), and its simplicity. The experiments and results were lauded.

For weaknesses, there were concerns about a lack of novelty in comparison to [1], a lack of theory behind the success of the proposed regulariser, computational cost, and a lack of experiments in terms of baselines and non-image domains.

The authors did well in responding to reviews, and provided additional experiments. As 3/4 of the reviews were borderline (8,6,5,5) post-rebuttal I encouraged the reviewers to have a discussion to move towards a more decisive opinion. While none of the borderlines explicitly declared either way, the resulting discussion was useful for my decision-making. Two reviewers echoed their appreciation of this paper in terms of the experiments and practical benefits, with another reiterating a lack of theoretical novelty.

I think the strengths of this paper outweigh the weaknesses, and it provides significant added value from [1]. My recommendation is for  acceptance (poster).

[1] Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. Advances in Neural Information Processing Systems, 2023.