Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models

Task arithmetic has recently emerged as a cost-effective and scalable approach to edit pre-trained models directly in weight space: By adding the fine-tuned weights of different tasks, the model's performance can be improved on these tasks, while negating them leads to task forgetting. Yet, our understanding of the effectiveness of task arithmetic and its underlying principles remains limited. We present a comprehensive study of task arithmetic in vision-language models and show that weight disentanglement is the crucial factor that makes it effective. This property arises during pre-training and manifests when distinct directions in weight space govern separate, localized regions in function space associated with the tasks. Notably, we show that fine-tuning models in their tangent space by linearizing them amplifies weight disentanglement. This leads to substantial performance improvements across multiple task arithmetic benchmarks and diverse models. Building on these findings, we provide theoretical and empirical analyses of the neural tangent kernel (NTK) of these models and establish a compelling link between task arithmetic and the spatial localization of the NTK eigenfunctions. Overall, our work uncovers novel insights into the fundamental mechanisms of task arithmetic and offers a more reliable and effective approach to edit pre-trained models through the NTK linearization.

Paper

References (93)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer eVrq6/10 · confidence 3/52023-06-09

Summary

This paper studies the "task vectors" framework where the weights of models can be perturbed in specified directions corresponding to tasks which result in improvements on those tasks. They attribute the success of this framework to "weight disentanglement" which means that adding a task vector for task i does not change how the network behaves when seeing input for task j != i. In addition, their investigation reveals that task vectors are not a consequence of fine-tuning occurring in the NTK regime (post-hoc linearization underperforms), but explicitly linearizing fine-tuning improves the task vectors framework as it "amplifies" "weight disentanglement".

Strengths

Overall I thought this paper was very interesting, it's claims are mostly precise and thoroughly investigated, and it's results are impressive (+ 5.8 for task addition). I though Figure 5 was particularly nice evidence of the authors hypothesis.

Weaknesses

- Limited experiments for many tasks. I would be very curious to see what happens in e.g. table 1 if you scale up the number of tasks. - Limited attempt to falsify the hypothesis -- for instance can you consider two tasks which share the same input images? Would task vectors still work? (See question in questions section). - Throughout the paper the authors mention variations of "Specifically, we probe the hypothesis presented in Wortsman et al. [80] that task arithmetic is possible thanks to the fact that models inherently operate in a linear regime" but I cannot actually find anywhere in [80] where this hypothesis is stated? I thought that Wortsman et al. [79, 80] were observing that ensembles behaved roughly similar to averages in the fine-tuning regime, and simply used the NTK regime as an example of where this would occur, but noted that there was still differences between the two (which is not predicted by the NTK regime). How do the authors feel about the hypothesis that averaging models (~= task vectors) ~= output-ensembling + terms-which-may-be-small (e.g., Section 4 of [79]).

Questions

- Do the authors findings point to any "better" way of adding task vectors rather than just adding them for all weights in the network? - How does weight disentanglement change with scale? - Can the authors think of an experiment which might falsify their hypothesis? E.g., do they think that task vectors might still "work" if two tasks share the same input space. Here's a quick thought: get one task vector which corresponds to the problem "first 5 classes in CIFAR10 or second 5 classes" and another corresponding to the task "traditional CIFAR10". If you apply both these task vectors then look at CIFAR10 performance on just the first 5 classes I think performance would still be as good as just applying the CIFAR10 task vector? - Should the community switch to linearized fine-tuning?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Would have been nice to see a limitations section.

Reviewer 9cnX8/10 · confidence 4/52023-06-26

Summary

This paper presents a comprehensive theoretical and empirical analysis of task arithmetic for model editing, where adding different task vectors (obtained by taking the difference between fine-tuned and pretrained model checkpoints) could improve the model’s performance on these tasks and vice versa. The authors propose weight disentanglement to investigate the underlying principles of task arithmetic, which involves decomposing the learned model function into a sum of localized components with disjoint supports. Specifically, the authors compare **regular** and **post-hoc linearized** fine-tuned models and find that weight disentanglement significantly contributes to the ability of task arithmetic. Based on these insights, the authors then propose to directly employ the linearized model for fine-tuning, obtain optimized task vectors, and lead to improved task arithmetic performance. Further analyses are conducted to reveal the connection between task arithmetic and weight disentanglement.

Strengths

- The paper is clearly written and well organized. - This work presents a neat analysis of task arithmetic based on the use of linearization and neural tangent kernels, which is novel to my knowledge and interesting, offering a fresh perspective on understanding the geometry of pre-trained checkpoints’ weights. - The proposed method to further improve task arithmetic is simple yet effective, supported by extensive empirical findings.

Weaknesses

As the authors already mentioned, one potential weakness is the introduced computational overhead during training. However, since the main focus of this work is a comprehensive analysis of task arithmetic, this is not a significant issue in the context of this work.

Questions

1. How does model scale impact weight disentanglement, and thus the benefits of linearized fine-tuning? 2. (Minor) Can the findings about task arithmetic be generalized to natural language tasks as well? By intuition, it seems that the setting of natural language texts is more straightforward for weight disentanglement due to the inherent discreteness of text tokens.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The authors have adequately addressed the limitations.

Reviewer sE5a9/10 · confidence 4/52023-07-04

Summary

This paper theoretically and emperically investigates the reasons why task arithmetic (an emerging technique for editing pre-trained neural networks) works. The paper shows that, contrary to previous hypotheses [39,79,80], lienarity of the fine-tuning on individual tasks is not sufficient to fully explain the success of task arithmetic (Sec 3). Instead, the authors propose the related and straight-forward idea of weight disentaglement (Eq 4) as an explanation (Sec 4), also providing a measure of the disentaglment error for two tasks (Eq 6). This idea is leveraged to improve task arithmetic performance by constraining the fine-tuning to linearized models - which importantly is shown to be an improvement over post-hoc linearlization, with marginal additional computational complexity (Sec 5). Finally, an additional connection between task arithmetic and the eigen space of the Neural Tangent Kernel is used to argue that weight disentanglement (and hence the capacity for task artihmetic) are a learned properties during pre-training, not inherent properties of parameterization or architecture (Sec 6). The empirical results focus on the CLIP image-text Vision Transformer model family, and many experimental details are reported, inclduing in the additional material in the appendix.

Strengths

This paper is the best kind of NeurIPS paper; Beautifully written and a delight to read, the authors have considered a timely and pertinent problem in an emerging machine learning domain, applied a methodological theoretical investigation; clarified prior hypotheses in the literature; used the theoretical findings to propose a simple novel methodology, and compellingly evidenced the subsequent effectiveness with an appropriate evaluation protocol. As a cherry on the cake, the methodology is a 'drop in' method that can be applied easily to existing approaches, and actual code is provided without breaking review anonymity in the supplementary material. Bravo. ### Originality * There are multiple practical take-aways that are immediately useful; the task disentanglement error metric (Eq 6), intuition for why task vector coefficients should be << 1 (Ln 181), simple drop-in code in the appendix that enables re-production and immediate application of the ideas (Listing 1 in the appendix). ### Quality * The work appears to be of high quality, the metrics and emprical appear to support the theoretical findings and claims. ### Clarity * This paper is superbly well written, and the logical argument and flow is well constructed. This paper was a delight to read and think about. ### Signifigance * Multiple compelling and genuine avenues for future research are identfidied (Ln 239, Ln 286, Ln 315, Ln354) * Provides timely and much-needed (i) overview of the emerging and rapidly evolving technique of task arithmetic with pre-trained models, and (ii) the beginnings of a theoretical grounding and understanding for why this technique is possible.

Weaknesses

* The main weakness of the paper in the present form is that the empirical results are limited to ViT models (specifically, CLIP). Within this modality, the empirical results are compelling (e.g. using multiple model families and 8x distinct datasets/tasks), and provide evidence for the theoretical findings, however the paper's authority and strength would be enhanced by demonstrating some of the key findings hold with another architecture or data type. E.g. I would love to see the results in Table 3 (Sec 6.2) replicated, even in a small way, with GPT or BERT on a text dataset. This would provider further strength for the empirical claims around Eigenfunction localization.

Questions

1. Ln 336 on 'Feature disentanglement'; to what degree is feature disentanglement enabled by over-parameterization / very large numbers of parameters? If this is the case, what is the parameter count relative to - i.e. some measure of the complexity of the input space or the function space represented by the data? Is there a way to quantify with a metric the 'capacity' for feature disentaglement of a given NN archiecture? These could be useful areas for investigation. 2. Related to the above point, Ln 233 notes that the advantage of the linearized model task arithmetic approach diminishes as the number of model parameters increases. I find this curious and would have expecte the opposite result - do you have any intuition why this is the case? 3. Pseudo-code or actual PyTorch code for Eq 6 would be a very helpful (and I assume straight-forward) addition to your supplementary material - this would be of immediate practical help to other researchers working on practical task arithmetic applications. 4. Related to the above, Ln 297 - can the degree of local linear independence of the NTK eigenfunctions be measured/computed with a simple metric to test this condition in other models / datasets? This would likewise be a useful contribution of the paper.

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

* It's not clear to me if there are situations where the linerarization technique won't be readily applicable. E.g. are some activation functions or more esoteric neural network architectures (e.g. recurrent NN's; long-term memory components) going to problematic to linearize? Some negative examples of where NTK linearization can't be used and/or approximations or work-arounds could be a useful addition in this regard.

Reviewer DjkZ8/10 · confidence 4/52023-07-06

Summary

This paper presents a comprehensive analysis of task arithmetic using pre-trained CLIP models. It challenges the early hypothesis that task arithmetic arises from linear fine-tuning in the NTK regime and introduces weight disentanglement as a necessary condition for enabling task arithmetic. Further experiments demonstrate that linearized fine-tuning of pre-trained CLIP models exhibits stronger weight disentanglement and improved task arithmetic compared to standard, non-linear fine-tuning. This result is accompanied by an analysis of the NTK spectrum to facilitate the understanding of weight disentanglement in linearized models. Lastly, empirical evidence suggests that weight disentanglement emerges from large-scale pre-training. The findings of the paper may shed new light on the effective adaptation of foundation models for downstream applications.

Strengths

- This paper advances the theoretical understanding of task arithmetic. In particular, it introduces weight disentanglement as a strong indicator of task arithmetic, and demonstrates that linearized fine-tuning in the tangent space of pre-trained weights promotes weight disentanglement. It further analyzes the NTK spectrum of linearized models and presents a sufficient condition for weight disentanglement. - In the meantime, the findings of the paper are significant on the practical side. The experiment results justify an emerging fine-tuning scheme for adapting foundation models. - Overall, this paper strikes a good balance between theory (NTK) and practice (fine-tuning pre-trained models). Embracing a broad audience is a key strength of the paper. - Finally, the paper is very well-written. The flow of presentation is very easy to follow.

Weaknesses

- Overall, I found the theoretical analysis motivating and the experiment results convincing. That said, all conclusions of the paper are drawn from CLIP fine-tuning, which makes me wonder whether the same findings are valid for other pre-trained models that similarly exhibit task arithmetic. To this end, I encourage the authors to report results on a second pre-trained model that differs in network architecture (e.g., ResNet), learning objective (e.g., MAE) or input modality (e.g., natural language), in order to establish the generality of their findings. - The authors observed that increasing model size (ViT-B/32->ViT-L/14) closes the gap between linearized and non-linear fine-tuning. Is it a consequence of stronger weight disentanglement because the fine-tuning of larger models approaches the NTK regime? Visualization of weight disentanglement throughout the paper only considers the smallest ViT-B/32 model. Similar visualizations for larger models could be informative.

Questions

These questions are likely outside the scope of the paper, yet I think addressing any of them via an empirical analysis could strengthen the work. - Task arithmetic makes the strong assumption that datasets from different tasks have disjoint support. This is often not true in practice. Does task addition still yield cooperative behavior when datasets overlap? How does linearized fine-tuning compare to standard, non-linear fine-tuning, especially when the assumption of non-overlapping data is violated? - In linearized fine-tuning, the optimization is restricted to the tangent space of pre-trained weights. Notably, this is conceptually similar to LoRA and adaptor-based fine-tuning, where a task vector exhibits low-dimensional structure. A natural question to ask is whether these fine-tuning approaches also produce favorable weight disentanglement / task arithmetic.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The paper discussed two main limitations. First, linearized fine-tuning is currently implemented using the JVP algorithm, which doubles the cost of a forward pass as compared to standard fine-tuning. Second, the spatial localization of NTK eigenfunctions is a sufficient (yet not necessary) condition for enabling task arithmetic. While the linearized models indeed respond to disjoint spatial regions in the experiments, this is not a must for task arithmetic to hold.

Reviewer DjkZ2023-08-15

Updated rating

The rebuttal addressed my concerns. The new results on a convolutional backbone and NLP tasks empirically confirmed the generality of the analysis. Overall, this work is well-positioned to inspire both the theory and practice of foundation model adaptation. I thus raised my rating from weak accept to strong accept.

Reviewer qb7d10/10 · confidence 4/52023-07-07

Summary

This work challenges the prevailing belief regarding the origin of task arithmetic in CLIP models. While it is commonly attributed to the linear nature of fine-tuning, the authors argue that the critical factor lies in the "weight disentanglement" that occurs between tasks during the fine-tuning process. The paper presents a theoretical analysis and extensive empirical validation to support this claim.

Strengths

This paper explores an intriguing topic with promising real-world applications. The authors exhibit strength in establishing a solid theoretical foundation, formulating straightforward research questions, and providing well-justified explanations that strike a remarkable balance between accessibility and rigor. The content is commendably clear and understandable, facilitating understanding for readers from diverse backgrounds. Personally, I was able to (at least, I think) grasp portions of the paper (Section 6) which required prerequisite knowledge I was not already familiar with. Notably, this paper achieves a synergy between its experimental design and the overarching claim that "Task arithmetic is not merely a result of linear fine-tuning," but rather depends on "weight disentanglement of the model with respect to the fine-tuning task set." The authors successfully validate the performance of their proposed method through rigorous experiments (and ablations) while simultaneously providing substantial support for key theoretical hypotheses. Overall, this paper represents a remarkable contribution, embodying the meticulousness, clarity, and scholarly standards I would expect in top-tier NeurIPS submissions. I’m starting my recommendation with an 8, given my not-high confidence, I’ll adjust it according to the rebuttal discussion.

Weaknesses

One aspect that requires attention is the clarification of the paper's applicability beyond CLIP models. Although there are specific references throughout the text (e.g., lines 38-49 in the introduction) emphasizing the focus on CLIP models, the overall messaging may appear too general regarding the broader applicability of the proposed method and study. To address this, I recommend at least modifying the abstract to explicitly state the paper's objective as "… a comprehensive study of task arithmetic in **CLIP** models…" instead of using the term "vision-language models" which could be misleading. Furthermore, exploring preliminary tests or investigations concerning the weight disentanglement of different pre-trained models would be valuable. For instance, considering the citation of various architectures in Section 6.2 (lines 301-303), extending the analysis to include some of them (e.g., convolutional neural networks) would be advantageous, strengthening the paper's applicability and relevance to a broader range of models and areas. Additionally, while Tables 1 and 2 provide average performance across tasks, it would be beneficial to include the standard deviation to provide a more comprehensive understanding of the variability in performance and, therefore, the method's robustness.

Questions

I’ve seen that in the supplementary material, there are some code snippets, but do the authors plan to release the full codebase upon acceptance?

Rating

10: Award quality: Technically flawless paper with groundbreaking impact, with exceptionally strong evaluation, reproducibility, and resources, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors adequately addressed the limitations.

Reviewer sE5a2023-08-15

Acknowledgement of author responses

I thank the authors for their responses and the further experiments. These results further strengthen what is already a great paper. Well done.

Reviewer 9cnX2023-08-15

I thank the authors for their detailed feedback. After carefully reading through the clarification as well as the other reviews, most of my concerns have been addressed, and these new experimental results further strengthen the presented analyses. I thus raised my score.

Reviewer qb7d2023-08-19

I want to thank the authors for their thorough response to every reviewer and for addressing all of my concerns. Having closely examined the other reviews and the authors' responses, it's clear they've put in substantial effort to incorporate reviewer feedback and improve their work (which was already of high level). During the rebuttal period, the authors added more experiments and clarifications that strengthened the paper's contributions and impact across fields, significantly extending its applicability to different models and fields. Given the unanimous acclaim for the paper's strengths (i.e., exceptional clarity in writing, substantial theoretical and experimental components, and valuable implications) and the authors' diligent revisions, I am confident in upgrading my recommendation to a score of 10.

Program Chairsdecision2023-09-21

Decision

Accept (oral)

© 2026 NYSGPT2525 LLC