When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers
Authors: Hongkang Li
Yihua Zhang
Shuai Zhang
Pin-Yu Chen
Sijia Liu
Meng Wang
Venue: ICLR 2025 Oral (2025)
Keywords: Task arithmetic
generalization
nonlinear Transformers
deep learning theory
machine unlearning
Primary area: learning theory
TL;DR: We provide the first theoretical characterization of the generalization guarantees of task vector methods on nonlinear Transformers.

Abstract: Task arithmetic refers to editing the pre-trained model by adding a weighted sum of task vectors, each of which is the weight update from the pre-trained model to fine-tuned models for certain tasks. This approach recently gained attention as a computationally efficient inference method for model editing, e.g., multi-task learning, forgetting, and out-of-domain generalization capabilities. However, the theoretical understanding of why task vectors can execute various conceptual operations remains limited, due to the highly non-convexity of training Transformer-based models. To the best of our knowledge, this paper provides the first theoretical characterization of the generalization guarantees of task vector methods on nonlinear Transformers. We consider a conceptual learning setting, where each task is a binary classification problem based on a discriminative pattern. We theoretically prove the effectiveness of task addition in simultaneously learning a set of irrelevant or aligned tasks, as well as the success of task negation in unlearning one task from irrelevant or contradictory tasks. Moreover, we prove the proper selection of linear coefficients for task arithmetic to achieve guaranteed generalization to out-of-domain tasks. All of our theoretical results hold for both dense-weight parameters and their low-rank approximations. Although established in a conceptual setting, our theoretical findings were validated on a practical machine unlearning task using the large language model Phi-1.5 (1.3B).

Published as a conference paper at ICLR 2025
WHEN IS TASK VECTOR Provably EFFECTIVE FOR MODEL EDITING? A GENERALIZATION ANALYSIS OF NONLINEAR TRANSFORMERS
Hongkang Li', Yihua Zhang", Shuai Zhang", Pun-Yu Chen', Sijia Liu Meng Wang'. 'Rensseleer Polytechnic Institute, Michigan State University, New Jersey Institute of Technology, *IBM Research
ABSTRACT
Task arithmetic refers to editing the pre-trained model by adding a weighted sum of task vectors, each of which is the weight update from the pre-trained model to fine-tuned models for certain tasks. This approach recently gained attention as computationally efficient inference method for model editing, e.g., multi-task learning, forgetting. and out-of-domain generalization capabilities. However, the theoretical understanding of why task vectors can execute various conceptual op- erations remains limited, due to the highly non-convexity of training Transformer- based models. To the best of our knowledge, this paper provides the first theo- retical characterization the generalization guarantees of task vector methods on nonlinear Transformers. We consider a conceptual learning setting, where each task is a binary classification problem based on a discriminative pattern. We the- oretically prove the effectiveness of task addition in simultaneously learning a set of irrelevant or aligned tasks, well as the success of task negation in unlearning one task from irrelevant or contradictory tasks. Moreover, we prove the proper selection of linear coefficients for task arithmetic to achieve guaranteed general- ization to out-of-domain tasks. All our theoretical results hold for both dense- weight parameters and their low-rank approximations. Although established in a conceptual setting, our theoretical findings were validated on a practical machine unlearning task using the large language model Phi-1.5 (1.3B).
INTRODUCTION
Large pre-trained models (Chowdhery et al., 2022; Touvron et al., 2023; Achiam et al., 2023) have recently served as a foundational module deep learning systems. Under the pre-training-ane and- fine-tuning paradigm, although the traditional and straightforward full- full-parameter fine-tuning can demonstrate superior performance in downstream tasks, its immense computational and memory costs have become serious practical issue. Consequently, many Parameter-F.fiicee (PEFT) methods (Li & Liang, 2021; Hu et al., 2022; Jia al., 2022; Wei et al., 2022b;a) have been proposed to address this concern. Among them, the recent task vector approach receives increasing attention (Ilharco et al., 2022a; Ortiz-Jimenez et al., 2023; Hendel et al., 2023; Todd et al., 2024).
Fine-Tuning
The task vector approach first fine-tunes a tre-trained model several simpler tasks to obtain task vectors, which represent the weight differences between the fine-tuned models and the pre-trained model. To handle more complex tasks, proper model can be edited by adding linear combination of these task vectors to the pre-trained model. Since this approach only requires determining the appropriate arithmetic hyperparameters, with need for further fine-tuning on complicated tasks, the task vector method offers a significant efficiency advantage and is particularly effective when adapting wide range of downstream tasks. Empirical evidence shows that adding multiple task vectors can improve the model's performance on corresponding tasks, while subtracting certain task vectors allows the model forget associated tasks. A proper linear combination of task vectors can even enable the model to generalize on out-of-domain task that has an analogous relationship with the given task vectors, without needing labeled data. Additionally, has been found that using low-
"Corresooning author. Email: wangm7 @rpi.edu.

---

Published as a conference paper at ICLR 2025
rank and/or sparse task vectors can further improve efficieney while maintaining the performance (Yadav et al., 2023; Chitale et al., 2023; Yu et al., 2024; He et al., 2025).
Despite empirical successes, theoretical analysis of task vectors is less investigated. In particular, we ask the following question:
When and why can the task vector approach perform well in multi-task learning, unlearning, and out-of-domain generalization successfully and efficiently?
Some related theoretical works focus on analyzing the performance of machine unlearning from a purely optimization perspective (Ginart et al., 2019; Neel et al., 2021; Guo et al., 2020; Mu & Klabjan, 2024). However, these analyses do not apply to Transformer-based neural networks, which are key components of large pre-trained models. Moreover, these works cannot be extended to study multi-task learning or out-of-domain generalization to new tasks. Frankle et al. (2020) proposes the concept of linear mode connectivity, suggesting that there exists a small-loss connected region in the loss landscape of the model, thereby demonstrating that linear interpolation between models can yield good performance. The most relevant workto this paper is (Ortiz-Jiment) et al., 2023), which uses the Neural Tangent Kernel (NTK) framework (Jacot et al., 2018) to study neural networks as linearized models under specific assumptions, to justify the use of linear arithmetic on task vectors for targeted model editing. However, this work does not have generalization guarantees and cannot explain the success of task vectors in nonlinear models without NTK assumptions.
1.1 MAJOR CONTRIBUTIONS
To the best of our knowledge, this work is the first theoretical generalization analysis of task arith- metic on a nonlinear Transformer model for multi-task learning, unlearning, and out-of-domain generalization. Focusing on binary classification tasks, we provide a quantitative analysis of the dependence of the task arithmetic effect on arithmetic hyperparameters. Although our analysis is centered on a simplified single-head and one-layer nonlinear Transformer, our theoretical insights are validated on practical architectures. Our major contributions include:
1. fine-grained feature-learning analysis of the effectiveness of task addition and negation. We consider a data model in which binary labels are determined by the majority of discriminative tokens, rather than their opposing discriminative counterparts, while other tokens do not affect the labels. We begin by analyzing the learning dynamics of fine-tuning a Transformer and characterize the properties of the resulting task vectors. Next, we provide sufficient conditions on the arithmetic hyperparameters for the task vector approach to be successful. We prove that task addition is effee- tive for multi-task learning when the tasks are either irrelevant or aligned. Aligned tasks are those where solving one task contributes positively solving the other. In contrast, task negation is prov- ably successful for unlearning tasks that are either irrelevant or contradictory. Contradictory tasks are defined as those where improving performance on one task harms the performance of the other.
2. The first provable out-of-domain generalization guarantees through task arithmetic. Focus- ing on task vectors representing set of irrelevant tasks, we prove a linear combination of these task vectors can generalize to wide range of new tasks by properly selecting the arithmetic coefficients. Additionally, we characterize the range of suitable arithmetic coefficients sufficient for successful generalization. This is the first theoretical justification of task vectors' ability to adapt to new tasks.
Theoretical justification of low-rank approximation and magnitude-based pruning for task vectors. We construct low-rank and sparse approximations to task vectors and prove that the generalization guarantees are minimally affected by these approximations. This provides the first theoretical support for the practice of using low-rank and sparse approximations to task vectors in order to reduce computational complexity.
1.2 RELATED WORKS
Weight interpolation technique. Weight interpolation or model merging (Matena & Raffel, 2022; Ilharco et al., 2022b; Yadav et al., 2023; Yu et al., 2024; He et al., 2025) refers to the practice of linearly interpolating weights of multiple models, where these models may be fine-tuned from dif- ferent downstream tasks or using different hyperparameters (model soups (Wortsman et al., 2022a)). Weight interpolation is empirically observed to be able to guide the model towards wider optima (lzmailov et al., 2018; Frankle et al., 2020) and better generalization in both single-task perfor- mance and multi-task ablities, even surpassing fine-tuning methods in some cases (Rame et al.,
2

---

Published as a conference paper at ICLR 2025
2022: Wortsman et al., 2022b; Ramé al., 2023). Task arithmetic can be viewed special type of weight interpolation, where linear operations are performed on task vectors.
Feature learning analysis for Transformers. Several recent works study the optimization and generalization analysis of Transformers following the feature learning framework, which describes how neural networks gradually focus important features while discarding unimportant features during training. Jelassi et al. (2022); Li et al. (2023e); Oymak et al. (2023); Ildiz et al. (2024); Nichani al. (2024); Chen et al. (2024); Li et al. (2023a: 2024c; 2023b); Huang et al. (2024); Luo et al. (2024) study the generalization of one-layer Transformers on different data models such as spatial association, semaanic/contexttull structure, causal structure/Markov Chain of data, and the majority voting of tokens in the data. However, no discussion was provided for merged models.
Theoretical study of PEFT methods. These are recent theoretical analyses on other PEFT methods. For example, in-context learning is analyzed from the perspective of expressive power (Bai et al., 2023; Akyürek et al., 2023; Von Oswald et al., 2023), the training dynamics or generalization (Xie et al., 2021; Zhang et al., 2023a; Li et al., 2023c; 2024a;b; Huang et al., 2023). Some other works focus on prompt engineering with a tunable prompt (Wei et al., 2021; Oymak et al., 2023; Zhang et al., 2024). Another line of work theoretically investigates the low-rank adaptation in terms of the implicit bias the optimization process (Damian et al., 2022; Abbe et al., 2022; 2023; Boix-Adsera et al., 2023; Jang et al., 2024; Li et al., 2024d) or model pruning with generalization analysis (Zhang et al,, 2021; Yang & Wang, 2023; Yang et al., 2023; Zhang et al., 2023b; Li et al., 2024a). However, none of these works involve the task vector method related approaches.
TASK VECTOR: DEFINITION AND OBSERVATIONS
2.1 PRELIMINARIES
Let f X' θ + y be a neural network that maps inputs X E X to labels y E y with v E as the model parameters. Denote (0()) as the pre-trained model and V as the fine-tuned model on given task T.
Definition 1. (Task Vector) The task vector A for the task T is computed as the element-wise difference between the pre- pre-trained and fine-tuned weights, ie., A V (0))
Task Arithmetic and Generalization. Given the pre-trained model (0) and set of task vectors TA on tasks {Ti)ee, one can construct a merged model v (o(0) Eiey A;A VT, for inference on downstream tasks, where E are arithmetic hyperparameters. Denote {(X, y; V as the loss function for the input X E X, output E y, and the model V E O. Hence, the generalization error the task T' with data (X,y) Dyr is defined
E(x.))>,, 6(X y5 V)
Existing works (Ilharco et al., 2022a; Ortiz-Jimenez et al., 2023) conclude that by controlling A the merged model can generalize across different tasks. Specifically, adding several Tr, via making X; > 0, E VA V, leads to model that exhibits desired performance multiple tasks from VA- Such a successful multi-task learning result can be mathematically represented as
E E(X,))~D7 {(X, y v < O(e), Vi e VA.
Meanwhile, negating A with A; < 0, i VN c V, results in a machine unlearning model that performs poorly on Vn2 but roughly retains the accuracy on VA,VN, "VVN,ie.. i.e.,
B(x,))Dy ((X, y vr 2 @(1), E(,,,) ((x, v. v S O(e), vi VN, Vj E V\VN.
(3)
Moreover, task arithmetic is empirically (Ilharco et al., 2022a) shown to produce a model W (0) + X- A that performs well on task analogy, in the form that "the target out-of-domain task ( is Ta as TB is to Tc." by constructing task vector A AV (AV T..
2.2 EMPIRICAL OBSERVATIONS
Note that experiments in (Ilharco et al., 2022a) only summarize the empirical findings when tasks are almost "orthogonal" to each other, while oot-orrhogonal cases are less explored. Therefore, in Table 1, we further construct binary classification tasks the parity of digits of Colored-MNIST
3

---

Published as a conference paper at ICLR 2025
(Arjovsky et al., 2019; Chapel et al., 2020). We control the colors of digits to generate pair of two datasets so that the parity classification tasks on different pairs of datasets are conceptually "irrelevant," "aligned," or "contradictory" to each other, respectively.
For irrelevant tasks, odd and even digits are highly correlated with red and green colors in one dataset but independent of colors in the other. In aligned tasks, the odd and even digits are correlated with red and green colors in both datasets. In contradictory tasks, the color-parity correspondence is the opposite in the two datasets. Let Ti and T2 denote the parity classification task on two different datasets. V = (1))) + A + AA r; is used to evaluate the performance of Ti and T2.
A key finding from Table 1 is that the task vector method performs quite differently with dif- ferent task correlations. To be concrete, given AV and AV,, for aligned tasks, the merged model can acquire strong multi-task learning abilities but have poor unlearning capabilities. The conclusion is exactly opposite for contradictory tasks. For irrelevant tasks, using task arithmetic can result in good performance in both unlearning and multi-task learning. A question arises, i.e.,
(Q1) How does task correlation quantitatively affect the performance of task arithmetic in multi-tast learning and unlearning?
"Irrelevant" Tasks
"Aligned" Tasks
"Contradictory" Tasks
Mask
Unlearning
Multi- Task
Unlearning
Multi-Taa
Unlearning
Best
1.4
-0.6
0.2
0.0
0.6
-1.0
Ti Ace
91.83 6309)
95.02 (4056)
95.62 (0.00)
95.20 (4.42)
79.54 (16.20)
94.21 (-0.51)
To Acc
88.40 (:565)
50.34 (44524)
92.46 (-3.33)
90.51 (8.118)
62.52 (33.77)
4.97 (-89.85)
Table I: Test accuracy (%) of V = ()) + A + VA on task Ti and Tz with À
[-1.-088 -8.66.6..... 2}. Multi-task learning aims to achieve good performance both tasks, while un- learning decrease the accuracy Ta but maintain the accuracy Ti. The best À is selected based the largest accuracy summation (or gap) Ti and Ta for multi-task learning (or unlearning). The accuracy gap (%) using v to the fine-tuned models V7,, or V7,, is reported in the bracket.
We then explore the use of task arithmetic with two tasks Ti and Tz for an out-of-domain task T' We construct tasks and data with Colored-MNIST, where we make T' more aligned with Ti and contradictory to T2. This is new out-of-domain setting different from task analogies in (Ilaarco et al., 2022a). Table indicates that the optimal X and X results in a testing performance better than using any separately trained model 1 or VT This implies that task arithmetic is powerful in domain generalization and can be extended to more general scenarios beyond analogous tasks. Hence, another question occurs, i.e.,
(Q2) Why do the arithmetic operations of task vectors perform well for alization, and how choose the arithmetic hyperparameter >s for desired performance?
gener-
Fine-Tuning
V
Searching \1 A2 in 2,
T
(A,, X²2 T' Ace
N/A
(1, 0)
(0, 1j
(1.2, 0.6)
92.21
88.10
45.06
91.74
Table 2: Comparison between the test accuracy (%) by different methods with A and A V T.- Searching À
and A2 refers to evaluating V = vp(0) + AA + A2A on T' with d1, A2 E {-2, 1.8, -1.6. ........... .3).
3 A DEEP DIVE INTO TASK VECTORS
We first summarize the main insights in Section 3.1. Section 3.2 introduces the mathematical for- mulation of data and model. Sections 3.3 and 3.4 present the formal theoretical results on task arithmetic for multi-task learning, unlearning, and out-of-domain generalization. Section 3.5 the- oretically proves the existence of a low-rank approximation or a sparse version of task vectors to maintain the performance.
3.1
MAIN THEORETICAL INSIGHTS
We focus on a set of binary classification tasks, where the labels in each task are determined by the majority between the discriminative tokens versus their opposite tokens in each data. This follows
4

---

Published as a conference paper at ICLR 2025
the theoretical setting in (Cao et al., 2022; Kou et al., 2023; Li et al., 2023a; 2024c). We consider one-layer single-head Transformers. Our major takeaways are:
PI. Quantitative Analysis of Multi-Task Learning and Unlearning via Task Addition and Negation. Let a represent the correlations between two tasks Ti and Tz. where positive, neg- ative, and zero values correspond aligned, contradictory, and irrelevant tasks, respectively. We prove that the merged model, V/- vy(0) + A + XA T. is successful for multi-task learning if A A 1- - a + 3 for some small constant 3. Moreover, the merged model is successful in unlearning T2 if À S 0 for irrelevant tasks or if X [~ ), ((a ')] for contradictory tasks.
P2. Successful Out-of-domain Generalization through Task Arithmetic. Given the correlation y between each existing task T; and the target task T', prove that long as not all Ti irrelevant to T', we can achieve a desired out-of-domain generalization on 'i' using task arithmetic. We explicitly quantify the arithmetic hyperparameter as functions of r''..
P3. Low-rank Approximation and Magnitude-Based Pruning Preserves the Model Editing Performance. We provide the first theoretical generalization guarantees for the practical techniques of low-rank approximation and task vector sparsity that reduce computation. Focusing, on binary classification tasks based on discriminative patterns, we demonstrate that both sparsification of task vectors in the MLP layer (by removing rows with small magnitudes) and low-rank approximations of task vectors offer guaranteed generalization through task arithmetic.
3.2 PROBLEM FORMULATION
22, ,KP) E RdxP contains P tokens, where each token is d- dimensional and || | = 1 for E [P]. The label y E (+1,-1} is a scalar. We consider the learning model single-head one-layer Transformer with one self-attention layer and one two- layer perceptron, which is mathematically written as
Suppose that data X (2,
f(x: P
Σ ai Weo
Wo Σ Wv W T Tw WQæ)).
(4)
where v = (fam)Es Wo, Wv, WK, Wo] denotes the set of all the model parameters. am) E R" and Wo E maxma are the weights in the MLP layer. Wv E R WK, We E Rmaxd are weights in the self-attention layer. eWWen.)) Womn/ /EE= (Wica)) (((ca)))))) Womr min m } > d.
(WW Woai) =
Fine-tuning algorithm for task vectors. Denote { (X", "x",""), " as a dataset with N data points for the task function T, i.e., y" = T(X") for E [N]. We fine-tune the model by minimizing the empirical risk function, i.e., min NV En {(",, y"; V), via stochastic gradient descent (SGD) to obtain the task vector A for T. We use the Hinge loss {(X. y, V) max {1 y f(X; V),0) as the loss function. For simplicity of analysis, we let W = W K Wo E idddd and V = Wo Wv E Rmxe xd as (Jelassi et al., 2022; Huang et al., 2023; Zhang et al., 2023a). At the t-th iteration, t 0,1,
the gradient computed using a mini-batch B, with [B,] B. The step size is O(1). Every entry of W and V is initialized from NN( (2) where & < 1/vm. Each Q(I) is sampled from (+1//m, -1 - 1//m). am) does not update during the fine-tuning.
Following (Cao al., 2022; Bu et al., 2024), we consider the data formulation as in Definition 2. Definition 2. Denote pt E RR. as the discriminative pattern for the task T. Let (v1, U2. .us] be set of d-dimensional orthonormal vectors that spans the subspace of task-irrelevant tokens Dy pT, j E [M]. Then, each (X, ,y) ~ DD is generated as follows:
Randomly generate the label y from (+1,-1} with an equal probubility.
Each token is randomly chosen from -ur (vi.... Uss }. If (or - the number of tokens equal to frr (or -pT) is larger than that of pT (or [T)' pt and -1T (or "- and μr") ") are referred to label-relevant and confusion patterns for y
'This is motivated by empirical observations that embeddings of data with opposite labels, such anony- words, are significantly distinct (Engler et al., 2022) and even in opposite directions (Liu et al., 2024).
5

---

Published as a conference paper at ICLR 2025
(or y = -1), respectively. The average fractions of label-relevant, confusion tokens, and each v4, i E [M] are 6., 04. and (1 (1-6) 6, 64) / respectively.
The basic idea of Definition 2 is that each label is determined by the dominant tokens with Huo patterns while all v; do not affect labels.
3.3 How DO TASK ADDITION AND NEGATION AFFECT THE PERFORMANCE?
Next, we investigate the generalization of task addition and negation with task vectors obtained by fine-tuning. Consider the setting where V = {1, 2} with A T, and A Tr as the task vectors for two binary tasks T and T2, respectively. Ti (or 72) defined based on Mei (or (472)) the discriminative pattern following Definition 2. Hence, v ( + A + A
Denote a MT. pTz E [-1,1], 3 poly(nō.) + ((cvM)(< @(1)). Suppose the number of neurons m 2 M log M with M = (d). Motivated by experiments in Table 1, we discuss three cases, i.e., a > 0, a < O, and a 0, which corresponds an "aligned", "contradictory", or "irrelevant relationship between Ti and T2, respectively. Then, we state Theorem I for multi-task learning with the merged model V.
Theorem 1. (Success of Multi-Tass Learning Irrelevant and Aligned Tasks) For any E (0, 1) and task T, suppose the following conditions hold when fine-tuning a pre-trained model: (i) the butch size B 2 D(6-2log M), (ii) the step size 7 S O(1), (ūi) the number of training it-
erutions 1. 2 T O -0110, ), then the returned model V7 achieves a generulization error
B( [e(X, vi V 7 O(e)-
Moreover, given task vectors A and A obtained by fine-tuning as above for tasks Ti and Tz, the resulting = ((()) + AV + VT satisfies
B(x,))-Dr l (X, y: v I ((e) [X] B, and E(X.)))DTT {(X, 1 Occ)
(5)
provided that Remark 1. Theorem first states the sufficient conditions during the fine-tuning stage to obtain proper task vectors. Then, it characterizes the region of X ensure both tasks achieve O(M or ((c) generalization error by adding task vectors. For irrelevant tasks with a = 0, a constant A 2 1 -1-p is required. This implies that adding up the task vector A Tr in T results in desired performance of multi-task learning. For aligned tasks with a > O, we can obtain a good multi-task learning performance if X 2 a + 3. For contradictory tasks with O. we cannot find the proper X such that V obtains a small error on both Ti and Tn simultaneously, which means Y can hardly generalize well on contradictory lasks.
We then study the unlearning using the merged model y in different cases of a.
Theorem 2. (Success of Unlearning on Irrelevant and Contradictory Tasks) Given task vectors
A Ti and A AV,,, that are fine-tuned following conditions (i)-(ii) Theorem I, the resulting
(0) A Ar i + XA MT satisfies
B(X,y)-DT ((X, y V S O(e) + (11 121-8. B, and (x ((X, v; v) > O(1)
(6)
when (A) a
S 0; or (B) a < O, and 0(x some == 011), and 0 < < c/2;
S poly(nō.) )a or (C) 0 < a cfor
Remark 2. For irrelevant tasks with a 0, a constant X < 0 can ensure perfect unlearning on Tz while retaining on TT. For contradictory tasks with a < O, the unlearning performance is desired if a negative À is in 1-O(a -oolopooo.,... / i.e. negating A For aligned tasks with a > O, proper X for unlearning to be successful only exists when is small, indicating that unlearning becomes more challenging when lasks are more aligned.
Remark 3. Theorem and 2 generally justify the validity of task addition, i.e., X > frror multi- learning and negation, i.e., X < 0, for unlearning as long as A not too large. The appropriate region for X is determined by a, the correlation between the lasks.
3.4 CAN A MODEL PROVABLY GENERALIZE OUT-OF-DOMAIN WITH TASK ARITHMETIC?
Consider P heve as a set of task vectors finettuned on ()) for binary classification tasks (Tihev Each task T. is defined with uT, Vy the discriminative pattern following Def- inition 2. Given the observation that task vectors are usually orthogonal to each other in practice (Ilaarco al., 2022a), we study the setup where (1T))EE forms set of orthonormal vectors.

---

Published as a conference paper at ICLR 2025
We analyze the out-of-domain generalization on data (X, .) ~ Dyr for the task T', where the discriminative pattern is denoted by pr, and fr. = =iev, "HTT+K +K H1_ with μ¹_ 1 (pT, heve || = | = 1, K E R for i E Vv. Note that μTT contains a component 141 that is orthogonal to all discriminative patterns of existing tasks, characterizing it as an out-of-domain task.
generalization on T'. Theorem 3. (Out-of-domain generalization using task arithmetic) Suppose PT; 1. pr; for i + j,i,j E Vy Let v = Eiev, X,AVT, + Tr(0), X; + O. Then, given that each TV is fine-tuned to achieve ((c) error error following conditions (i)-(iti) in Theorem I, as long as the following conditions (A) there exists i E Vv s.t., 7/6 + 0, and (B)
The following theorem summarizes the required conditions for
E Eievo A 2 1 1+6. C,
Eiero Xn2 2 1 cc,, C, 1
for some C (0, 1) and all i Vw: E B(x,))Dr e(x, X. y; V) I O(c)-
we have
(8)
Remark 4. Theorem 3 implies that linear operations of task vectors can produce model that can generalize well on out-of-domain tasks T' that has a distribution shift from tasks Ti E Vv. With properly fine-tuned task vectors, the conditions make out-of-domain generalization successfil are (1) the discriminative paltern of the targel task T has non-zero projection onto at least one of the discriminative pattern of tasks T E Vy; (2) the weighted summation of 7 and nf with X as the coefficient should be greater than the margin of the binary classification task; (3) the absolute value of each X is not too large to avoid large errors to the resulting model V.
Remark 5. Note that >; satisfying (7) exists under mild conditions. In (75) of Appendix, provide a flosed-form solution that meets (7). We omit them from the main paper to simplify the presentation.
3.5 CAN TASK VECTORS BE IMPLEMENTED EFFICIENTLY?
In this section, we theoretically investigate how to improve the computation efficiency of task vector techniques during inference. We focus on two properties of task vectors, low rankness and sparsity.
Consider the fine-tuned model 1 = ((mm) W WÜT, WKT. W } with Wi W W and V W WVT from Lemma 1. Denote AWT W W and AVT = Vi (0), We have the following conclusions.
Corollary (Low-rank approximation) For any task T defined in Section 3.2, there exists AWLR Iadd and AVLR E with rank( (onkk)WL)) = rank( (LR) = 1, such that
I|AWT - AWLRIIF S M +
and 1|V - AVLRRIF S 2o-le,
log M
hold. Moreover, Theorems 1-3 hold by replacing A AT and AVT with AWLR and A VR. in the task vectors and replacing with CLR = (log nj + 8_ DE in the results.
Remark 6. Corollary states that when e E (0, (M log M)- we can find a rank-12 approxi- mation of W* and V* with an error less than (log M) to ensure that all Theorems hold with roughly the same generalization error. Specifically, with error derived in Theorems 1-3, using rank-1 approximation leads to CLR = (logn-1 + 5,")e, which equals O(c) given r and 8_ as constants. Hence, Corollary I indicates that low-rank approximation of individual task vectors generally preserves the performance of the model after applying task arithmetic.
We also prove that task vectors are approximately sparse in Corollary 2, which implies that pruning task vectors does not change the generalization.
Corollary 2. (Sparsity of task vectors) There exists L C [m] with [4] = (m) s.L.,
||u,|| 2 90m -1/2), E L; ||u,|| S 0mm-1/2 Vlog B)B), i [m]].C.
(10)
where u, is the i-th row of AVG and B is the batch size of fine-tuning lower bounded in condition (i) of Lemma 1. Then, pruning all rows in [m]].C of AVG ensures Theorems 1-3 to hold.
'The rank-1 I approximation results from our simplified model that has one discriminative pattern per task. Our result indicates that the proper rank for approximation depends on the number of discriminative patterns for each task, which far smaller than the model dimension in practice.
7

---

Published as a conference paper at ICLR 2025
Remark 7. Corollary 2 illustrates that constant fraction of rows in AV; in L has large mag- nitude, while the remaining ones in [m]]2 have much smaller magnitude. Then, we prove that removing rows in [m]]C does not hurt the performance of multi-task learning, unlearning. and out- of-domain generalization by task arithmetic. This indeed justiftes the existence of redundancy in "Delta parameters." " a similar notion of task vectors, defined in (Yu al., 2024), and verifies the validity of mugnitude-based pruning on task vectors like TIES (Yadav et al,, 2023) or DARE (Yu al., 2024).
3.6 PROOF SKETCH AND TECHNICAL NOVELTY
We first provide the following informal lemma for the fine-tuned task vector. Lemma I provides the convergence of the fine-tuning process and the properties the obtained task vector satisfies.
Lemma 1. (informal) A model v has generalization error ((e) task T (with the discriminative patten μT) if Av == v (0))= {AW, AV) satisfy both conditions as follows:
(A) the attention weights between two label-relevant patterns are dominant, while the attention values between label-relevant pattern and any other pattern are close to zero;
(B) constant fraction of rows in AV in the MLP layer has a large magnitude with a direction either close pT or -PT, while the remaining
have small weights.
Moreover, any task vector obtained by fine-tuning on task T satisfying conditions (i)-iiii) in Theorem I satisfy conditions (A) and (B) for task T.
The proof ideas of Theorems 1 and are as follows. To ensure successful multi-task learning stated in (2), we need AV + AAVTT satisfying both conditions (A) and (B) in Lemma I for tasks Ti and T2. To ensure unlearning To and maintaining the generalization need A T, + AAVTT Tr satisfying (A) and (B) for Ti but failing either (A) or (B) for T2. When
Ti as stated in (3), we
0, the component of AV, T, in Vi has negligibile effect on data from Ty, for any i * J. j. E {1.2). When a > 0, both Ti and To should tend to favor À 0 for good generalization. When < 0, Ti prefers a negative x, while To prefers a positive X.
To prove the out-of-domain generalization in Theorem 3, we need to find proper set of X i Vv n V' such that Eien, A hold for conditions (A) and (B) in Lemma for the task T'. The proof idea for Corollaries and 2 comes from observation from Lemma 1, That is, Conditions (A) and (B) demonstrate that the rows in VV and the matrix A W only enlarge tokens in the direction of label-relevant pattern its opposite. This implies the sparsity of AV and the low-rank property of the entire Av. The proofs for Theorems I and 2 and 3 and Corollaries I and 2 can be found in Appendix D, respectively.
Technical Novelty. Compared with (Li et al., 2023a), Lemma establishes more fine- tne-grained characterization of A which allows us to perform a detailed analysis of layer-by-layer outputs of the merged model. Furthermore, Lemma 1 extends the theoretical analysis to training from random initialization with two merged trainable parameter matrices W and V.
Moreover, to the best of our knowledge, we provide the first generalization analysis of task arith- metic in model editing (Theorems 1, 2, and 3). The merged model V preserves the nonlinearity of task vectors from the nonlinear model architecture rather than linearizing the model by impractical infinite wide network assumption in (Ortiz-Jimenez et al., 2023). This allows us to expand the un- derstanding of task arithmetic beyond the NTK region as in (Ortiz-Jimenez et al., 2023), where the problem is extremely overparameterized.
NUMERICAL. EXPERIMENTS
We conduct extensive experiments on image classification and natural language generation to verify the effectiveness of task vectors in different downstream tasks. For image classification, we use the ViT-Small/16 model (Dosovitskiy et al., 2020) pre-trained from ImageNet-21K (Russakovsky et al., 2015) for downstream tasks with Colored-MNiST (Arjovsky et al., 2019; Chapel et al., 2020). For natural language generation, we use the open-source Phi-1.5 (1.3B) language model (Gunasekar et al., 2023; Li et al., 2023d). We repeat the experiment using LoRA with Phi-3-small (7B) in Appendix B.
8

---

Published as a conference paper at ICLR 2025
4.1 EXPERIMENTS ON IMAGE CLASSIFICATION
Experiment Setup. To control the correlation between tasks, we use Colored-MNIST for image classification tasks. We designed binary classification problems based on the parity of digits, where odd digits are labeled as +1 and even digits as -1. We utilize two colors, red and green, to construct different task correlations. Define ro and re as the proportion of red colors in odd and even digits, respectively. Then, the proportion of green colors in odd and even digits are 1 - ro and Te. spectively, Across all of our experiments, we set re = 1 - To. The correlation " TT between two tasks Ti and T2, with Dj and D2 respectively as the corresponding test set, is approximated by their averaged cosine similarity between centered outputs from the two fine-tuned models, i.e.,
d( 7 T7 = 1/2(6 TT: D1) + G T2 D2)),
where a( V Tg. Dj)
cos (Vi.j- F2.1)
Vi. vis
1 D2
Σ SDD
Hi Lij f1.2). (11)
iED; iDyl
Vf, represents the i-th output of the fine-tuned model 1 on the test set Dj. Note that to compute á V T, W by (11), do not require the availability of extra models or datasets except 7 Vy, and the test set D2 and D22
Experiment Results. We first investigate the ability of task arithmetic using v = yr(() + AV + AA to handle multi-task learning and unlearning under three cases in terms of task correlations. Let ro 0.95 for Ti. In case I, let r. re 0.5 in Tz. In case II, let ro 0.9 in T2, and in case III. let r, = 0.05 in Tz. The computed correlations a(Vr 7T ) of the above three settings are 0.164, 0.891, and -0.849, which corresponds to irrelevant (a 0), aligned (a > 0), and contradictory (a < 0) tasks discussed in Theorem 1, respectively. Figure 1 illustrates that when tasks are irrelevant, successful multi-task learning on both tasks and unlearning on task To can be achieved when A 2/ 1 and À S O, respectively. When tasks are aligned. the trend of testing accuracy of V/ on Ti and T/ are consistent. A superior multi-task learning performance can be observed when A > 0, and one cannot find a region of À where Tz is unlearned while maintaining the accuracy for Ti. When tasks are contradictory, one can obtain a good unlearning behavior when S 0, and no selection of X can achieve multi-task learning. This result verifies Theorems and 2 for a = 0,
a 0, and a < 0, respectively.
100
100
Accuracy (%)
Accuracy (%)
Accuracy
on T
50
50
50
on T
on Ti
on
00
0]
(A) Irrelevant tasks
(B) Aligned tasks
(C) Contradictory tasks
Figure I: Testing accuracy of the merged model v task Ti and Tz. We then study the out-of-domain generalization capa- bility of task arithmetic. We consider a merged model V = (0) + >1 A VT, A T, constructed by two task vectors. In Tí. we let r. = 0.85, while in Tz, we let r. 0.05. In the target task T', ro 0.9. We compute that a(v², T. VT,)) = 0.115, which means Ti and To are approximately irrelevant. Figure 2 (A) demonstrates that in triangular region with the black dashed line of Àι and X2, we can achieve a good generalization performance. This region is consistent with the red region in Figure (B), which is produced by condition (7)3 where 71 and 72 are estimated by á a(v+ ) = 0. 0.792 and a w+}) = -0.637. We choose small values 3 = 0.01,
0
(A)
(B)
Figure 2: (A) The heatmap of the testing accu- racy (the color bar %) on T' using the merged model V. The black dot is the baseline, while the green cross is the best A1. )2 (B) The red re- gion satisfies (7), while the blue region does not.
0.02. The
result justifies the sufficient conditions for a successful -o--a-nn--minn generalization in Theorem 3.
Sinnee the practical classification margin might be smaller than that of Hinge loss used in our theoretical analysis, we replace + in (7) with 0.2 +

---

Published as a conference paper at ICLR 2025
4.2 EXPERIMENT ON LANGUAGE GENERATION TASK
Experiment setup. We study the unlearning performance using three datasets, "Harry Potter 1" (HPI), "Harry Potter 2" (HP2) by J.K. Rowling, and "Pride and Prejudice" (PP) by Jane Austen. We consider HPI and HP2 as semantically similar and aligned books due to the shared authors (ā(VTup)) W
Tupp)) = 0.498 by (11) following Dou et al. (2024), while PP is less aligned with HP1 than HP2 (0x TRpp: VTpp) = 0.239 by (11)). We study Next Token Prediction on these three datasets separately as three different tasks, denoted by THP,, THP2, and Tpp, respectively. Then THPI and THP are greatly aligned, while THPI and Tpp are less aligned.
Denote the pre-trained Phi-1.5 model as (()), We first fine-tune g on all three datasets jointly to obtain (00), which has favorable generalization for all tasks THPL, Thrz, and Tpp. Initialized from (0(),, we fine-tune on dataset HP1 to obtain model V iip The task vector for THPI is computed as:
AV = ViPi ()()), The merged model is V/ = (r + X AV
Experiment results. We vary À and evaluate the performance on THP1, THP2, and Tpp, respectively. The evaluation metric is the Rouge-L score used in (Dou et al., 2024), which measures the ratio of the longest common sequence between the original book and the LLM's generation. A higher score indicates better generation performance. As shown in Table 3, when À becomes negative, the Rouge-L score for TIP decreases, indicating the success of unlearning. When À is the smallest value in the experimental selection (A 1),, the unlearning performance is the best, with the Rouge-L decreasing by 37.23% from y ((0), Morcover, when THPI is unlearned, the performance of THP2 also degrades significantly, with the Rouge-L score decreasing by 34.71%. In contrast, the performance degradation on TTp is much smaller, with a decrease by 15.13%*. This verifies Theorem 2 that unlearning task THPI effectively degrade the performance of the aligned task (THP2) as well, while the performance degradation on the less aligned task (Tpp) is relatively smaller.
0 (baseline)
-.2
-0.4
0.6
0.8
THP1
0.2213
0.2211
0.1732
0.1866
0.1572
0.1389
(37.23% J)
THI2
0.2302
0,2032
0.2111
0.2034
0.1695
0.1503 (34.71% 4)
TPp 0.1983
0.1888
0.1877
0.1802
0.1932
0.1683 (15.13% 4)
Table 3: Rouge-L scores of THel, Thpz. and Tex by V/- ( +X AMMii using full-rank task vector A A
We also implement our experiment using LoRA in fine-tuning to compute the task vector. We set the rank of each parameter as 32, which requires to tune only 0.35% of total parameters and reduces the peak memory consumption by 54%. Let A L denote the resulting low-rank task vector for THPI- We repeat the experiments by replacing A HII with A A HP1- Comparing Table 4 to Table 3, on can see that all the insights still hold when using a low-rank task vector, verifying Corollary
A
0 (baseline)
-02
00.4
-0.6
0.8
-1
THPI
0.2432
0.2033
0.1857
0.1665
0.1439
0.1568
(35.53% 4)
THP2 Tpp 0.2111
0,2335
0,1932
0.2065
0.1813
0.1664
0.1772
11% J)
0.2001
0.1884
0.1963
0.1849
0.1819
(13.83% L
Table 4: Rouge-L scores of ThP Tuez, and Tre by v = (0)'
scores
+ + A Avle, using low-rank task vector A 1
5 CONCLUSIONS
In this paper, we theoretically investigate the generalization ability of the task vector technique. Based on feature learning analysis of one-layer nonlinear Transformer, we quantitatively charac- terize the selection of arithmetic hyperparameters and their dependence on task correlations so that the resulting task vectors achieve desired multi-task learning, unlearning, and out-of-domain gen- eralization. We also demonstrate the validity of using sparse low-rank task vectors. Theoretical results are justified on large language models. Future directions include analyzing the performance of task vectors in more complex models and designing more robust task vector selection methods.
'Note that the task vector method leads to a 13. 1..1% decrease in Rouge-L L score on BOOKS dataset on average (Shi et al., 2024). The state-of-the the unlearning methods are empirically shown to result in performance drop in utility (Maini al., 2024; Shi et al., 2024).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by National Science Foundation(NSF) #2430223, Army Research Office (ARO) W91 2NN-225-1-0020 and the Rensselaer- -BMM Future of Computing Research Collaboration (http://airc.rg http///irr.rpiedu..... The work of Yihua Zhang and Sijia Liu was also supported by the National Science Foundation (NSF) CISE Core Program Award IIS- IIS-2207052, the NSF CAREER Award IIS- 2338068, the ARO Award W91 INF2310343, the Cisco Research Award, and the Amazon Research Award for AI in Information Security. The work of Shuai Zhang supported by National Sci- ence Foundation (NSF) #2349879. We also thank all anonymous reviewers for their constructive comments.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: a necessary and nearly sufficient condition for sgd learning of sparse functions two-layer neural networks. In Conference on Learning Theory, pp. 4782 4887. PMLR, 2022.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learn- ing Theory, pp. 2552-2623. PMLR, 2023.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint 23x0:22303.08774.22
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. 2019.
arXiv preprint arXiv:
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023.
Enric Adsera,, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transform- learn through gradual rank increase. arXiv preprint arXiv:2306.07042 2023.
Dake Bu, Wei Huang, Taiji Suzuki, Ji Cheng, Qingfu Zhang, zhiqiang xu, and Hau-San Wong. Prov- ably neural active learning succeeds via prioritizing perplexing samples. In Forty-first Interna- tional Conference on Machine Learning, 2024. URL https:/ /openreview.er n forum? id=kzz0kn54 6b.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convo- lutional neural networks. Advances in neural information processing systems, 35:25237-25250 2022.
Laetitia Chapel, Mokhtar Z Alaya, and Gilles Gasso. Partial optimal tranport with applications on positive-unlabeled learning. Advances in Neural Information Processing Systems, 33:2903-2913, 2020.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning transformers. arXiv preprint arXiv:2409. 10559, 2024.
Rajas Chitale, Ankit Vaidya, Aditya Kane, and Archana Ghotkar. Task arithmetic with lora for continual learning. arXiv preprint arXiv:2311.02428, 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXivv2204003111 2022.
11

---

Published as a conference paper at ICLR 2025
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory. pp. 5413-5452 PMLR, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Reprexentations, 2020.
Guangyao Dou, Zheyuan Liu, Qing Lyu, Kaize Ding, and Eric Wong. Avoiding copyright infringe- ment via machine unlearning. arXiv preprint arXiv:2406. 10952, 2024.
Jan Engler, Sandipan Sikdar, Marlene Lutz, and Markus Strohmaier. Sensepolar: Word interpretabilitty for pre-trained contextual word embeddings. In Findings of the Association for Computational Linguisties: EMNLP 2022, pp- 4607-4619, 2022.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference Machine Learning, pp- 3259-3269. PMLR, 2020.
Antonio Ginart, Melody Guan, Gregory Valiant, and James Y Zou. Making ai forget you: Data deletion in machine learning. Advances in neural information processing systems, 32, 2019.
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. arXiv preprint arXiv:2306. 11644, 2023.
Chuan Guo, Tom Goldstein, Awni Hannun, and Laurens Van Der Maaten. Certified data removal from machine learning models. In Proceedings of the 37th International Conference on Machine Learning. pp. 3832-3842, 2020.
Yifei He, Yuzheng Hu, Yong Lin, Tong Zhang, and Han Zhao. Localize-and-stichh Efficient model merging via sparse task arithmetic. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https: http///pnneeeeeeeeeeeeeeeeeeeeee wwwp:rewwwwie...eeeeeeeeeeeeeeee wttp:/eteeeooommmcecceeeeeeeeeee
Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 9318 9318-9333, 2023.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. International Conference on Learning Representations, 2022.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv:2403.02233. 2024.
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self- attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint arXiv:2402. 13512, 2024.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In The Eleventh International Conference Learning Representations, 2022a.
Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannanch Hajishirzi, Si- mon Komblith, Ali Farhadi, and Ludwig Schmidt. Patching open- vocabblayy models by interpo- lating weights. Advances in Neural Information Processing Systems, 15:222622-292777 2022b.
P Izmailov, AG Wilson. D Podoprikhin, D Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 875,885, 2018.
12

---

Published as a conference paper at ICLR 2025
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Uijeong Jang, Jason D. Lee, and Ernest K. Ryu. LoRA training in the NTK regime has no spurious local minima. Forty-first International Conference on Machine Learning, 2024. URL https: for///ormtt/// um? mmmid=ssl d=s s sdx G.
// /poopneeeeiew.. net featt/forune
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 15537722-377835 2022.
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim. Visual prompt tuning. European Conference on Computer Vision, pp. 709-727. Springer, 2022.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, pp. 17615- 17659. PMLR, 2023.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh Interna- tional Conference on Learning Representations, 2023a. URL
nett/
forum? (dtmmmdd==jcccv111111
Hongkang Li, Meng Wang, Songtao Lu, Hui Wan, Xiaodong Cui, and Pun-Yu Chen. Trans- formers as multi-task feature selectors: Generalization analysis of in-context learning. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023b. URL https:
//openzeview.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear trans- formers learn and generalize in in-context learning? In Forty-first International Conference on Machine Learning, 2024a. URL
net Wwworrnnid-IATPPP--
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear trans- formers for chain-of-thought inference: A theoretical generalization analysis. arXiv preprint 210.012410.02167. 2024b.
Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, ZAIXI ZHANG, and Pun-Yu Chen. What im- proves the generalization of graph transformers? a theoretical dive into the self-attention and positional encoding. Forty-first International Conference on Machine Learning, 2024c. URL net/ forum? forum?id=mJhxls122 2222..
Hongkang Li, Meng Wang, Shuai Zhang, Sijia Liu, and Pin-Yu Chen. Learning on transformers is provable low-rank and sparse: A one-layer analysis. In 2024 IEEE 13rd Sensor Array and Multichannel Signal Processing Workshop (SAM), pp. 1-5. IEEE, 2024d.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1th International Joint Conference on Natural Language Processing (Volume I: Long Papers),
pp. 4582 4597, 2021.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning. 2023c.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.05463 2023d.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. arXiv preprint vrXiv:2303.04244 2023e.
Sheng Liu, Haotian Ye, Lei Xing, and James Y Zou. In-context vectors: Making in context learning more effective and controllable through latent space steering. In Forty-first International Confer- on Machine Learning, 2024.
13

---

Published as a conference paper at ICLR 2025
Yuankai Luo, Hongkang Li, Lei Shi, and Xiao-Ming Wu. Enhancing graph transformers with hierarchical distance structural encoding. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: http:///ppeeeeeeeeetteeeee hwt.:/tt//oooooe het./orumffddi- U4K1 dRgoph.
Pratyush Maini, Zhili Feng, Avi Schwarzschild, Zachary C Lipton, and Zico Kolter. Tofu: A task of fictitious unlearning for llms. arXiv preprint arXiv:2401.06121, 2024.
Michael S Matena and Colin A Raffel. Merging models with fisher-weighted averaging. Advances in Neural Information Processing Systems, 35:17703-17716, 2022.
Siqiao Mu and Diego Klabjan. Rewind-to-ddlete: Certified machine unlearning for nonconvex func- tions. arXiv preprini urXiv:2409.09777 2024.
Seth Neel, Aaron Roth, and Saeed Sharifi-Malvaeerdi. Descent-to-delete: Gradient-based methods
for machine unlearning. In Algorithmic Learning Theory, pp. 931-969-22 PMLR, 2021.
Eshaan Nichani, Alex Damian, and Jason D) Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
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 Sys- tems, 36, 2023.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning arXiv preprint 230::2333003335 2023.
Alexandre Rame, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamoniy, Patrick Gallinari, and Matthicu Cord. Diverse weight averaging for out-of-distribution generalization. Advances in Neural Information Processing Systems, 35:10821-10836, 2022.
Alexandre Ramé, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David Lopez-Pazz Paz. Model ratatouille: Recycling diverse models for -ut-o--distribution generalization. In Interna- tional Conference on Machine Learning, pp. 18656-286679 PMLR, 2023.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause. Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 1 11::211-222, 2015.
Weijia Shi, Jaechan Lee, Yangsibo Huang, Sadhika Malladi, Jieyu Zhao, Ari Holtzman, Daogao Liu, Luke Zettlemoyer, Noah Smith, and Chiyuan Zhang. Muse: Machine unlearning six-way evaluation for language models. arXiv preprint arXiv:2407.06440 2024.
Eric Todd, Millicent Li, Arnb Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. In The Twelfth International Conference on Learning Representations, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière. Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /3971, 2023.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 11X1v:1011.3027, 2010.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023.
Colin Wei, Sang Michael Xie, and Tengyu Ma. Why do pretrained language models help in down- stream tasks? an analysis of head and prompt tuning. Advances in Neural Information Processing Systems, 01116658-161700 2021.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu. Brian Lester, Nan Du, An- drew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference Learning Representations, 2022a.
14

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. thoouuh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24835 2022b.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing ference time. In International conference on machine learning, pp. 23965- PMLR, 2022a.
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 7959-7971, 20222.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations. 2021.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin A Raffel, and Mohit Bansal. Ties-merging: Re- solving interference when merging models, Advances in Neural Information Processing Systems, 36, 2023.
Hongru Yang and Zhangyang Wang. On the neural tangent kernel analysis of randomly pruned neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 1513- 1553. PMLR, 2023.
Hongru Yang, Yingbin Liang, Xiaojie Guo, Lingfei Wu, and Zhangyang Wang. Theoreti- cal characterization of how neural network pruning affects its generalization. arXiv preprint arXiv:2301.00333. 2023.
Le Yu, Bowen Yu, Haiyang Yu, Fei Huang, and Yongbin Li. Language models are super mario: Ab- sorbing abilities from homologous models as free lunch. In Forty-first International Conference on Machine Learning, 2024.
Siqi Zeng, Yifei He, Weiqiu You, Yifan Hao, Yao-Hung Hubert Tsai, Makoto Yamada, and Han Zhao. Efficient model editing with task vector bases: A theoretical framework and scalable ap- proach. arXiv preprint arXiv:2502.01015, 2025.
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.0992. 2023a.
Shuai Zhang, Meng Wang. Sijia Liu, Pin- Yu Chen, and Jinjun Xiong. Why lottery ticket wins? theoretical perspective of sample complexity sparse neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Shuai Zhang, Meng Wang, Pin-Yu Chen, Sijia Liu, Songtao Lu, and Miao Liu. Joint edge-model sparse learning provably efficient for graph neural networks. In The Eleventh International Conference on Learning Representations, 2023b.
Yihua Zhang, Hongkang Li, Yuguang Yao, Aochuan Chen, Shuai Zhang, Pin-Yo Yu Chen, Meng Wang, and Sijia Liu. Visual prompting reimagined: The power of activation prompts, 2024. URL https:/ wttpp//wwpweeeeecp////ppp net forum? ferrm??dd00b300 0=00d28
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL DISCUSSION
was brought to our attention after the acceptance of ICLR 2025 in January 2025, that there is recent submission arxiv in February 2025 (Zeng al., 2025) that also considers the theoretical generalization analysis of task vectors multi-task learning, unlearning, and out-of-domain gener- alization. Their analysis is built upon assumptions that (i) the studied models already fine-tuned (Assumption 4.1); (ii) the norm of task vectors is upper bounded (Assumption 4.1); (iii) different task vectors are almost orthogonal to each other (Assumption 4.2). In contrast, although analysis is based on a one-layer single-head Transformer, we do not rely on the aforementioned assumptions. Our results show that the convergent models trained with SGD yield task vectors that support multi- task learning, unlearning, and -ut-o--distributionn (OOD) generalization. We analyze the behavior of task arithmetic under aligned, irrelevant, and contradictory task relationships without requiring the orthogonality assumption between task vectors. Moreover, unlike (Zeng et al., 2025) that assumes sparsity of task vectors, theoretically prove that task vectors obtained via fine-tuning can exhibit both low-rank structure and sparsity.
B ADDITIONAL EXPERIMENTS
We repeat the language generation experiment in Section 4.2 with Phi-3-small (7B). The task vectors are obtained by LoRA (Hu et al., 2022). Table 5 shows that the insight of Theorem 2 still holds, i.e., unlearning a certain task (HP1) can effectively forget the aligned task (HP2) with a 52.29% decrease of Rouge-L scores, while the Rouge-L score for the less-aligned task (PP) has decrease of only 20.65%. Moreover, by using a larger model than Phi-1.5, the unlearning performance of the aligned task HP2 is improved from 37. 23% decrease 55.61% % decrease. In comparison, the performance difference on the less-aligned PP is much smaller, from 15.13% decrease to 20.65% decrease.
0 (baseline)
0.2
0.4
0.6
-0.8
THPI
0.2573
0.1989
0.1933
0.1888
0.1572
0.1142 (55.61% 4)
THrz
0.2688
0.2113
0.1993
0.1938
0.1622
0.1563 (52.29% })
Tpp 0.1942
0.1825
0.1644
0.1687
0.1592
0.1541 (20.65% )
Table 5: Rouge-L scores of THPI THp., and Tep by v = y/00) + A Av using low-rank task vector A 1 with Phi-3-small (7B).
C PRELIMINARIES OF THEORY
We first summarize the notations we use in this paper in Table (6).
Definition 3. For task based any discrininative pattern μti
1. (t) M.T ( p1-
2. S": the set of tokens in the n-th data. Si: the set of tokens of pr in the n-th data. S the set of tokens of th in the n-th data. Ry: the set of tokens of Us in the n-th data.
3. on (t)
1sp Pe -100------ 1S1
4. pt(() Expest or Jest softmax, (x² (10 (10p))
5. GL1.1 v(()) (1,) for ST.
6. S1,e = min G.1.
7. softmax, T Ww;) = (softmax, (o) ... softmax ".f T W Wx))).
Definition 4. Define
P
R (t) = Σ
(1) x. softmax (T (())))
(12)
16

---

Published as a conference paper at ICLR 2025
Table 6: Summary of Notations
Notations
Annotation
X. a. X", yd
is the imput data, which contains P tokens. D; is the i-th token of X. X" is the n-th input data with y" the corresponding label.
{aw Wo, Wv, WK, Wo} denotes the set of all the model pa- rameters. a() E R" and Wo Rmxm are the weights in the MLP layer. Wy E _R WK, Wo E Rm are weights in the self-attention layer. (0)) is the pre-trained model. V7 is the fine-tuned model on a given task T. A is the task vector of the task T. which is computed as A =
V
(0(0) V7 AVT
V TT(0)_
HT. V2
PT the discriminative pattern of the task T. Vy is the j-th task-irrelevan pattem. j E [M].
8.. dΦ
6. is the average fraction of label-relevant pattern in the input data. 64 is the average fraction of confusion pattern in the input data.
ạ (t) 14 W l4l denotes the value of the product, where the patterns both sides of W(t) are the same. Çt.e denotes the modified value embedding of μ4I at the (-th iteration. (1) refers to the summation of attention weights where the key and the query are the same discriminative pattern.
Wa,tMan
Wn,, and i respectively represent of sets of positive or negative neurons so that the Relu activation is activated with x; as the query.
B.
B, is the SGD batch at the b-th iteration.
O0, 10.90.0 9()
We follow the convention that [(x) = O[g(z)) (or Mgg())) (((())) means
that fx) increases most, least, or in the order of g(z), respectively.
am 1/ m for i mm).
2 S
f(x) 2 ( (or f(x) S g(z) means that f(x) 2 [[(g(z)) (or f(x) S
O(g(z)))
(x)
Define WA. Un,1 as the sets of lucky neurons such that
Wn,1 {i: ViT. (o.) (0, 0,1 E ST, a, > 0),
(13)
Un,1 {i: (1.)) R0) > 0,1 Si,a, < <).
(14)
Definition (( ((Vershynin, 2010)). We say X a sub-Gaussian random variable with tab-Gaussian norm K > 0, if (P) S K for all p 2 1. In addition, the ub--Gaussian norm of X, denoted || is defined as |X 11 = supp21 p= (
Lemma (Vershynin (2010) Proposition 5.1, Hoeffding's inequality). Let X1, X2, ... XN be in- dependent centered sub- sabb-gaussian random variables, and let K max, |X Then for every (a1,".' ,aN) E N and every & O, we have
ct2
Pr
Σ a
21 S expexp
(15)
where e is an absolute constant.
Lemma 3. For task T based on any μti OSt5st T, there exists K(t) > O, such that
M
w p1 = w((+1) - K(t) +
Σ
4pp.
(16)
1=1
where
K(t) 2 T Σ m S|ST
aP (1O
(17)
n
4 S K(I) 201())
(18)
17

---

Published as a conference paper at ICLR 2025
For k [M]
log B
|l W 0kll S V
K(b).
(19)
B
anr + k. j (M)
|03 W unll S K(t) 91(()
(20)
For any ye such that po] pt a and μl' u1, 02, .. Var, we have
1t (w μl = a²t4 ( p1 (1 + +()),
(21)
BB log M for some E < 1.
Lemma 4. Given a task T based on any μ1, 0 St S T. Then, for i E Wa.b.
S
こ "
Pn (b))
(22)
aP
b=0
nE
St aPM
Σ
Uk
(23)
6=0
neB,
for k E [M]. For i E Un,t we similarly have
S (b)
MI
11
(24)
aP
b=0
RE
|S
V
(1))
Uh
Σ
B
aPM
(25)
6=0
nEB,
for some k E [M]. For i & Wn,1 U Unn, we have that
log B
M1 S
MI:
MI:
(26)
B
v (i.)) th S
log B
Uk
(27)
B
where & E j E Wa,1 U
Lemma 5. (Full version of Lemma 1) Given a task T defined in Definition 2 based on the discrimi- native pattern Hoy, have that as long conditions (i)-(iū) in Theorem I hold, then the returned model Vy after T iterations achieves a generalization error
E (X.)) [((X, VT S O(c).
(28)
The required sample complexity is N BT, where B is the batch size. We also have that
1.
Pa (T) 2 -1-8--06,,
(29)
for some constant C >> 1.
2.
M
M (6. 1T
1 (T)
l
(30)
for i E Wn, with I E ST and for E Un with I E Sy. We also have that (26) and (27) hold when 1 T.
18

---

Published as a conference paper at ICLR 2025
D PROOF OF MAIN THEOREMS AND COROLLARIES
D.1 PROOF OF THEOREM 1 AND
Proof. Since the model is initialized close zero, then AV is close v. Denote Vi (fag.1)
P ), Vi Wi} and V2 ((ap.2))=) P 22, W2). We consider three cases of this learn ing problem. (1) Consider a = O. By (21) in Lemma 3, we know that
MT, (Wl (T) + XW2 (2)
MT, UT, W HT, (1+ + = 1T, (T) pr.
(31)
1 (Wl (T) + A HT = - W PT.
(32)
11 ( (T) - AW2
HT2 = =>uT, W (T)
PT
(33)
14 (W AW, 14T = -- W PTs
(34)
Then, for any / e [M] and for task TT,
8 T-
softmax((""" w((x(") 2
(35)
6.
REST
for task T2,
8.72C
1 8
softmax, (x," T w(()) 2 8,PaC + (1-
2
T-
(36)
6,
SES,
Since that PT,
{1T v1, U2, uM and pTi 1 {ITT2 v1, 01, "., vw , we have
H42 = 0,
(37)
for V E 1 and
(T) (1.) MT; O,
(38)
for V U2 Then, for data with the label y = the network output for V + A is almost the same as that for v1 on task Ti when [^] is not too large. To see this, for X from Tr, we have
P
Relu(
(2) 1(6.)
+ AV (T)
X softmax (T + A
2(1.))
a
1=1
iE(m]
<|^|- (1
ΣE
T-1
STI
1- 8- T + |^|
141.10 .ely
M
log B
(39)
aPM
d,
B
6=0 5E|m]
nEB
< |^| O 8-) (oly()).))
|a| ((cvM)
= =|113,
where the second to last step is by (26) and (27) and B ス 2² log M. Therefore, a larger |\] leads to performance drop in task T.. For data of Ti with the label y == -1, we can choose À to be greater than around to make the network output smaller than - Meanwhile, for X from T2, we have
f(X". V)
1- 8- T- A -((
M Mog B
O 1-o,) d,
(40)
211
poly(n6.),
d
B
where we need 1 3 that f(x", v) 2 1- (cE).
If A < 0, the attention map tends to be uniform. Then, for X in task Tz, we have
S(X"; Vi + AV22) S P
(41)
which leads to
E. (x ((X ys V) O(1)
(42)
19

---

Published as a conference paper at ICLR 2025
(2) Consider > 0. We first have
PT
T (W1 (T)
+ AW (T)
2 UT 1h: w MT (1 + Xa
(43)
Mr ( (W (T) 2 1T: = (A+ a²) (es W 1T.z.
(44)
Then, for y" 1 in task Ti, we have that when À > f(x", v
((c)) (a) - -0
T-1
ΣΣ
Σ
S
1 8,
p-AC
B
aPM
82
6=0 sE(m]
nEB
M Mlgg B
-I1 V( - | -IXI
(45)
B
1 1-0,)
2(Aa | O
:poly(no.) - |^] -1 ((eVM)
8.
poly(no.) - [YI O(cVM), and for y" = 1 in task T2, we have that when 1 O.
=1 + 0(Aa) | O
f(x", R v 2(1
8,
T
o((ta)))) (A + - 8(
M log B B
8.
(46)
1 8n
poly(1ō.))
8.
Therefore, when A 2 1- + 3, we have that for task Tı,
f(x", v) 2111 MBB O(e),
(47)
and for task T2.
1- 8,
M Mog B B
f(x", 11) 2(1 (A)) + a)
8, voly((no))-111
(48)
2(1 - (x))
21 - O(e).
We can obtain corresponding conclusions for 1² 1. Hence, E( ((X, y/ v) S O(e) + 14|3, E B( 8(X, y; vn= - O(E).
(49)
(50)
Meanwhile, for y² = task Ti, we have that when À < 0,
f(x", w) 2(1 1 8,
1 1- 8 T C(1+2a)) 1 1- - 1- C)) (1 + Aa)
6.
6.
8.
1 6,
(I
.poly(pp.) - |^| ((cVM)
5.
(51)
21 21+ 1 1-6. 7- -(LAAA)))) - 1 -1- -1--------0 (1-6-7-0r4400 6.
1- 1-0-
8,
8,
8.
- 111 0 10111100 8, pooo(yp0.) - |A| ((cvM)
1-)))
8,
and for 3" = 1 in task T2, we have that when A << 0.
f(x", v) 211 1 8. T c(Att)))) (A + a) -0
M log B )- 1-66 0
1 1-b) poly(ndo.)
6,
B
6_
1 - 8r T -
1 d, - C(s+o?)
1- b. T- ()) (2 ax)
21
5.
8,
M log B
-011--0 --0) -poly(n0))
0 1
-e(V
-
(52)
B
8,
21 + a(
1- 8.
TT o(xxto))) - >1 11--0
1-6,
p- o(Ata??))
1- 1- p
8,
5.
5,
Mlog B
1 -ō. poly(1ōo))
-e(V
ee
B
8,
20

---

Published as a conference paper at ICLR 2025
Then, for task Ti, when 0 > 2 -0(1/a²),
E B(x.)/DD f(X, ys V)
-min(O(-Xai)) 1 - 8 T ( + 1 - d. (--17-0)1+0)) -
1-d0 0) +
8.
6,
8.
1 8.
(53)
+ +11 + 1) 0
poly((p.) + 1 ((cVM)),
8,
2 min + (|>] - 1) poly(1).) + || - ((evM)),
-min(O //| poly(né.)), 0(1)),
Hence,
E (X, ((x vs v) 2 min(O (1 (1)))))) {(1)
(54)
When A -O(1/a²)
B(X,y)-Dr, {(X, y; v)
1
=0(1 M M M)
(55)
20(1).
For task T2, when 0 > 2 -(1)-
(( (X, y; V)
1- 8. T c(++a2)
1 1- 8.
T- 1 1- _ - +
[m d a
+ +X(
d.
8.
6,
M log
1 8,
(56)
+ +0(V
+ 0
(oly(10.)), 0(1)
B
8,
2 min + 1/C - a ((poly(nn) + +cVM)),
1min of - a + B), ((1))
where the second step is by + 0(1) + a a² 2 O(1) When A < 0(1) - a² < 0, E Ex {X, y; v) 2 @(1.
(57)
(3) Consider <. When X X (- 0), we have that for task Ti.
f(x" R v)
8. C((+220)
T-1 O(e)) - (1 Aa) |N O
ΣΣニΣ
|ST aPM
b.
I-6- 1-6-T- T-C
b=0 TE
REB.
1 8 -x AC |NI-
M log B
8.
B
2(1 O())- (1+ + Aa) - |^| O( 1 8. 1-xz) poly((8.)) - -XX--)
(58)
8,
1- (T C(1+22)) - O
8.
+ +xa
1 I C
1 1- 0.).
2(1 O())) (1 + Aa) - |^| O
poly((ro)) 12| ((vVM)
8.
- moly(gō.))a² (-=22- (1 + Aa),
Hence, if X S poly(1).))a )a, we have
f(X", v) 2 - M4 O())-
(59)
B D l(X y: v) - O(c) 19/3.
(60)
21

---

Published as a conference paper at ICLR 2025
If X
_ a-3- we have
R2 (x).)))D y. v) 2 min(0 + (|>) + 1) poly(1)6)) + - (61)
I (1)/a2), we have
E (x,,y)-Dy, {(X, y/; v) O())
(62)
For task T2. we have that when X 2 1 +
1 a + A
1-8, 8, 6,
M log
f(x", 4) 2 (1 (A+
poly(nó.) -oly(nó)))
2 1,
(63)
B
B(x,))DDT, ((X, y: v) - O
(64)
When A 51 (ppoly(no_) + M).
B( ((X, yi v) 2 min 1+),
a 8}.
(65)
One can easily find that there is no region of X such that Yp performs well on both Ti and Tz- However, when (1(1/a²) X < poly(n).)a < 1 1+ 7f a + B, we can unlearn Tz and retain the performance of TT.
D.2 PROOF OF THEOREMS 3
Proof. By Lemma we know that
( pT
pT.
Σ TH
X (7)
Σ REV
(66)
iEV
j=1
2 Σ
2 HT w p..
KEV
For positive neurons, we also have
V(T =
Σ
VT. (T)
Σ nHT = Σ
(67)
iEVo
IEV
iEVo
Then, we need
Σ Am 1cc
(68)
NEVe
Σ Amo 2 1+c,
(69)
ieVe
(A 1- 8,
poly(().)) + Mc) = | < c, for some > 0 and all i E Vv.
(70)
to hold simultaneously.
Then, when % = h does not hold for all E Vy and for some fixed is < O, we can find X; in the middle of the normalized 7/ and 7/ to satisfy (68) and (69), i.e.,
7
A α V EuENo ?
(71)
V
Eievy 72
By inequality, we have
Σ 7 LEVE
Σ 75
<
Σ
7 <
Σ 7
Σ
(72)
V
LEV
LEV
VEVo
LEV
22

---

Published as a conference paper at ICLR 2025
Hence,
Eievo ?i
V EiEvo V V Eiev V x + -iivvo 7
Σ Ara x
Σ xp + KEVO
0. (73)
VEiEvo 7/
Viiero 71
v Eiev 7i V Eieve V8 Eiero 7i
Σ Ani Eiery Pi ViEEyy 71
Σ 7/- SEVO
0. (74)
ViEEv 78
VEVo
Therefore, by letting
A = c.
%
(75)
VViEvo 1i
VEievo Vo
where
(1+0) Eievo A
C Vievo 18 V Viee A + Eiev, 72
(76)
we can obtain (68) and (69) hold if C. 5- When % = & hold for all i Vo and for some fixed is < 0 with [+l > O, cannot find X such that both (68) and (69) hold.
L
D.3 PROOF OF COROLLARY
Proof. Let {ni v1, U2 which is denoted by
uM U2, , ,u-M44) form a set of orthonormal vectors,
U (1) 11 U2,"" UM, u1 U2, Us (s-
(77)
Note that for any a, E (1, 02,02 02. ".
",,"", uM) U {mi, u2, .. uua M+13.
a (w Σ abb NO, Σ la;b, I6?)
(78)
ISipsd
where the last step comes from that each entry of ( (0, {²). Given that ||a|| = ||6||
1.
we have
Σ
lacbyl = (la1).". ... laal) (1b1). ... (bal) < 1.
(79)
PSt
M4- and any t = 0, 1, ... T --1 1, By (90), we know that for E {ur U2.
1
0I(X", y"; 0,
(80)
B
(W
n
1
0f(X" y²: V)
O.
(81)
B
@W
REB
Then, we have that for some C' > 1,
Olog T'),
i=j=1, j= j=1,1 M 1,
Ole 2ellog
O(e - T-9),
U UTTTU
62 1-
log
jE [2, M 11,i [1, M - 1],
O
else.
(82)
23

---

Published as a conference paper at ICLR 2025
Let B, be the matrix that only the (i,j) entry equals 1, while all other entries are O. Therefore,
- W - E1, (log 1 S(c. 9). (M-1)
-1) 1)
(e log T) - (M-1)(M- 2) + -) M2)
<i² log T. 2 + de dfmm
(83)
Se M² +
log M
where the last step comes from that m 2 M log M and M (d). Then,
W 11, (log T) U |
WWW - UE Ollo lle ||T|||
(84)
U| I|U W Ei, Olog ((ooTT)ll SeM + M.
T)
Likewise, by (132), we know that neurons of ( with a non- non-trivial magnitude are in the direction of the iterative summation of EE=n x sosstmaa Way)). Hence, there exists ún E m and i2 Rd such that
log B
- T 52550) Vm.
5 8
S
(85)
B
Vm.
Then, for such that y" +1, we have that the low-rank trained model, where W (T) Olog 7) UT satisfies
U U
LR
S(X", VIR) 21 6,6) (1 Oce = 1 - 8.)c),
(86)
which leads to
(I LLR S O(CLR), where CLR = (log T + 8.)c.
(87)
D.4 PROOF OF COROLLARY
Proof. We know that from Lemma 1, there is number of (2m) lucky neurons with large weights.
We can denote the set of lucky neurons as C [m.. By combining (148) and (163), we have that for any lucky neuron ue
[ull 2 1777
5
1/22
(88)
/ /m
For any unlucky neurons, by (149), we have
log B
| S m 11/
(89)
Since that B 2 2 log M by Lemma 1, we have that if we remove neurons from m L, the output in (158) and (159) will only be affected by a factor of e. Therefore, Lemma 1 still holds, so that Theorems 1-3 all hold.
24

---

Published as a conference paper at ICLR 2025
E PROOF OF KEY LEMMAS
E. PROOF OF LEMMA
For ease of presentation, we sometimes use μ2 to represent -M1 in the proof. We first investigate the gradient of W, i.e.,
O"(" .y" W
OW
B
dl(X a y²; V) V
0ʃ (X*** "; V)
OW
B
=1 Σ
0 1V (c.) X softmax (X" Wx,") 2 0]
P
R
(90)
P
(Vi.))
softmax
Wx! softmax/(z" Wx,") - x²/)="
#=1
-1 Σ (-v"))
a(), 1V (1.) X" softmax X Wx!) 2 0]
nEB
(Vi.) 8=1
t softmax((""" Wæy)) (a"
Σ
softmax (x.)) Wx )x/))="" )<=") T)
r=1
For j, E ST, we have
el (0ll
softmax((x") "w(x))) 2 ( + (P-1187)
(91)
For j E S and I E ST, we have
softmax/(2:;
W(0 2x) S ST (0) + (P
(92)
where [[q: = 0. For 7 & ST UUS E [P). we have
softmax/ W(0) a;)) P.
(93)
Therefore, for 8, r. E SP, let
softmax(x"" w(' = (t)pa + (t),,,
(94)
where
P- -1S (S(())P-157))
(1) Z
(95)
S e ella (+)
+ SS
M,
(1) Σ
Aps
(96)
1=2
where
|ST1
M S (1)
(97)
P P1SST
Note that [+] 0if 1811,12 2.
If a E St, we have
Pn (t)
( (1.) 2f softmax W 2 G.
(98)
[ST
25

---

Published as a conference paper at ICLR 2025
If S and j Sf, we have
(L) (1,) (z of softmaxi(a]] w(oa))on()) |
(t) (1)) x softmax w()z)) x
(99)
Pr (1)
If 8 & (Sf U and j E ST.
(1.) 2 sostmaxx W ( ())(())))))))))))))))))))))))) |S11
10) (4) 2 (oft W() x (
(100)
Bt)
Then, by combining (94) (100), have that for I E EET, Wn,1-
Σ
x softmax Wx (x
Σ softmax/(x" Wa/l)=;)) pi (101)
25 Pr( - (S
For E ST, E Wn,,, we have that for l: # 1,2,
P saft wxr")a") p
12 (i.)
x softmax Wa (a.
(102)
-114 (1,) Σ x. (x"t Wa(") (a.
softmax((x;"" T " pa
For E ST i E Wn,,. we have that for k E [M].
Ug Va.
Σ x softmaxx((""" Wx.) (x*
softmax W ax )x"" )a f
Sp1 ((r)
x (aft Wx() (x.
softmax;(x,) Wa "x pn
(103)
[Rg]
S (t)
P [ST] For i E Un,1, by the definition of Un, in Definition 4, we have
(t)
[ softmax, X Wwy) 0] 0.
(104)
For i & WT,1 Uli, d, we have that for j Wn,1, & [M].
p Va.
x softmax Way) (xr"
softmax; T Wax" "ar")""")a p
Sp
Vu-
Σ
x softmax(
WxF) (x?
softmax:(a;) Waf))x)aa in
(105)
|ST
2.(() (1)
(t)
P
mb V(L.)
x ="softmax((x"" (Wz""" (x"
(ofimax((x" W Vx")x"" ) T pn
(106)
- Vi.)
x softmax((""" Wx') (x"
softmax; W xi xx 22 pi.
T Ug V(v.)
xx softmax Ww:" (x"
softmax; (ω) W x )a )a) ph
Sp V(-
xx softmax Wxx) (x
softmax; (x." War)ma)oot i (107)
|S
R
(t)
(t)
P- P-S
26

---

Published as a conference paper at ICLR 2025
When & Sf, have that z7 11 O. If Sg.
obtain that
Vo.
Iy softmax
Wxy) (x
softmax;(a;" "x,"))))) p2
(108)
Pa (t)
2511111
(tP
|S
P
M1
((..)
a softmax
Woo)) (x"
softmax: (2. Way))a yzi p22
==1
(109)
=- U2 Va.)
F softmaxi(a"" Wa])) (x²
softmax T Wæf )ay p22
s=l
P
0g V(.)
n. softmax Wa:") (ar
softmaxi (2,, W Wa an )ar p2
Su2 T
((.))
Σ
xx (aftmax W (xe"
softmax(( Wail)ar pe2
(110)
[Rg
ST (t)
P (S
(t)
where ls E [M], E Un If E Wn.b.
[V,,,) softmaxx Wxy) 2 0] 0.
(111)
if WP,1 U have that for j Un le E [M].
122 ((.)
x softmax Wa:) (x.
softmax/ Waj
az War" )ayk
Su4
Vu.
Σ
x softmax Wxx)) (x"
softmax:(""" Waijwa" u12
(112)
aml
|S
(n)())
(t)
P
p1] V(..)
xy softmaxx Wxn)) (Wa (x"
softmax (a.) Wa mf ya pt2
(113)
=- Vo.)
x "oftmaxx((""" Wx:") (x"
softmax((x, W ) ar n22
P
0g V(.))
x softmax Wxx") (x)
softmax (2." W )x"" )x) p12
Spz V(-
xx softmax Wx.)) (x²
softmax;(x)) War n2 (114)
s=1
SST
R
(t)
(1) P- SSP
If R
E [M], we have that for E Wa, VG. Ep E softmxx Wx?) > 0.
1' W Sp.
SHI Vo-)
x" softmxx
(TTT Wxx)) (xr
softmax; (x; Wx "a"""""")a"""" Ok
(115)
Spt VG. Σ x softmaxr Wxf) (x"
Σ
softmaxr (x) waj pi
27

---

Published as a conference paper at ICLR 2025
P
H12
V(.)
k
(fftma;((
W (at
softmax; (..))
W ax Jai )os on
8=1
(116)
=- Mi (
(oftmax Wa;) (ar
softmax;(1"" W Wat")a))) )ay oa.
==1
p
0g
o softmax (2" Wxi' (x"
softmax W Waf)aa))a) Up
S11 T ((-))
x softmaxx Wxi) (xe
softmaxy (x"" Wa))a,)af (117)
[Rg]
P-181
Likewise, if I E R2. l E [M], Vo.) E softmaxx Wx7) 0, E Un l' E ST
1" E Sg.
P x x" softmax
P
S12
VG.)
Wx()) (x"
softmax, (x. W Wx "=())""")) Ok
P
Suz V(.)
x softmax T Wxii ) (x""
softmax!
W Wxf Wxff)=")))t j p2
1=1
(118)
P
Σ
x softmax Wxx") (x*
softmax (x.
W Wxy jar") ok
=1 12 Vi.)
a softmax (x² Ww (xr
softmax (m. W Warie )ar
142
(119)
P
0g (in)
Σ
"" softmax (x, Wx?) (xx"
softmax W "a"")"")
Sp V(++
x softmaxy Wxi) (x Σ
softmaxy (z"
)x pI
(120)
[
P- ---ST
Therefore, by the update rule, we know
O Rh .y:; V)) OW
((+1)) u1 ( p1 B Σ
nEB,
(121)
-W(()) #I K(t) dp
where
K(t) " Σ S
1
aP (PP
(122)
nEB
ST (n(()) (t)
4 4 K(t) max
< <((t) e (ek(t)--aall -a(())
(123)
We know that
(0) #4 0.
(124)
28

---

Published as a conference paper at ICLR 2025
Then,
(i+t+ - ui (((+)) p1 =ii W M1 + K(t)
=1)(1) + K(1)
(125)
K()).
Similarly,
w 12 - 12 BB
1
Σ
OE(X", .y; V
OW())
142
nEB
(126)
-W + k())
Σ
122
02 W u2 = Σ K(b).
(127)
b-0
For [M],
M
w(+ = W Ji(1)m + J2(1) +
Σ
(128)
1=1
By Hoeffding's inequality (15), with high probability,
log B
| w (++ 1) vull S (1))
K(b) Se. Σ K(b),
(129)
B
b=0
b=0
where the second step holds 2 -2 22 log M. And for j j k.j E [M]
|uo] w uxll S K(()e-911 -((()
(130)
For any μl such that pol μt' and nz' 1 {w), U2, ". ;, can write jı' aue t/ o for some 11l1 1 {11 U1, -. ). Therefore,
n¹ w((+)) =aam + V1- (ilt) + (1-0211)
(131)
=a" #1 ((+1)) #n + O(E) ui T (((+)) p1.
E.2 PROOF OF LEMMA
For ease of presentation, we sometimes use μ12 to represent -11 in the proof.
OE(X", -u"; V)
VV(..)
BEB
O R y V f(X":
- B BEB
OS(X*M V)
VV())
(132)
Σ (- p Σ
a(), 1 [[V sootmax " Wxf") 2 0]
- B
BEB,
P
x softmax W
For such that y" and E Wh. we have that
1 (1.) X softtax T (Wa"""" 2 0] =1,
(133)
29

---

Published as a conference paper at ICLR 2025
and for 1 s
M2
Σ td softmax Wx = Pn (1) +
Σ
u/ + 1 M2;
(134)
1=1
where
[R/
4 -1-1100 pt) P-
(135)
|T
If / E Sg, we have
P
Σ ==sottmax
Wxx" (t) +
Σ
k/vl + 1MM2-182:
(136)
1=1
where
pt)(1) S (t))),
(137)
[
m/ s (1 pa()
(138)
p (5
If R2. k E [M], we have
P
Σ ""softmax
Wxy) = p1)()) -onm,+r*0000 +00000022200000 +
EE
(139)
where
(t(() S STI
P (t))),
(140)
PM(() < S
P (t)
(141)
(1) S [
(t)
(142)
|ST +
| |R]
RR/
m/ S
pn)( P- 1SS
(143)
STT
|5
[
Therefore, we have
B
1
Σ
OP X" .,;; V
0V
Σ
M
d/v + qn (1)pm) + (t)p))p
(144)
BEB,
where
STT
(1(1)² Z B
Σ
a. Pt)));
(145)
BEB
S
(1)( 2 7 Σ
P.. (1),
(146)
aP
B
|u'dl S B Σ
R
(1 - Pn (1)) 1
(147)
aP
M'
NEB,
Then,
V
(6.) M1
2
Σ
[-1
SHI
Pm (b),
(148)
aP
b=0
REB
(1)
(1,) H2
(z.))
(1)
M1,
(149)
S
(
(..) Uk
S n Σ
aPM
(150)
30

---

Published as a conference paper at ICLR 2025
for k E [M.. For E Un, we similarly have
|
12
11
(b),
(151)
aP
6=0
BEB
(
(t) (i.) 122
(152)
(5.)
(1)
Uk
Σ
B
aPM
S
(153)
6=0
nEB
for some & E [M]. For i & Wa, Ulh,, we have that
log B
Uk S
B
(1.)
UR
(154)
log B
(-(1) M1
B
(1.)
11.
(155)
where & E [M].j E Wa, UlM...
E.3 PROOF OF LEMMA
We know that by Lemma and 4 in (Li et al., 2023a), for E (,,(()) and E Sp, we have that
1[v (1))
R7
(7)
1,
(156)
and for i E (0) and I E Sy, we have that
1 [t)
(157)
We also have that the size of Wa, and Vn, are larger than 2(m). Therefore, for y" +1, by Lemma 4 and 3, we have
f(x"; v P
Σ
Σ
Relu ((a.) ssofmaxx (x" Wxx))
(0)
+
Relu(V ((..) ssftmaxi (X" Wx:))
(158)
P
1=1 Wi (0), >0
Relu(Vis Xsoftmax(X"" Wx?)).
a
1=l
We know that
Σ
Relu V (T) (t.))
X Xsoftmax (X w()x)))
1=1 (EW...(0)
SST
Ç (T))))
(159)
P
S
T-1
S
P
a
1)
Σ
P
Pk(() PT))
6=0
KEB
We can derive that
T-1
qi
(T)
Σ
K(b)
0m
T
m|S
カ-
17 B
aP -100)
Σ
aP (c(()
(160)
B
6=0
nEB
c=0
heB.
T-1
84
e9:
0=0
31

---

Published as a conference paper at ICLR 2025
Therefore, we have that when qı (T) I O(1) qı (T) 2 ((T)) for (1), (160) does not hold. When Q1 (T) = (log T", we have that (160) holds. In this case,
8.7
8 T
(T) 6.7C +1 5.
2
(161)
8,
where C' > 1. Mcanwhile, for E R k E [M]. and any E [P].
1 softmax/(""" w()x7) A p)
(162)
We can then derive that as long as
T T 17 103
(163)
we have
S
m
| pa(6)- (T) 2 1.
P
a2
77
p
(164)
6=0
hEB,
Then,
s(X"; v) 2 1, 1 (X" ,"; v)
(165)
With (163), we can also derive that
M Σ
2 S
( M (2.) mill
(166)
which means that for i E Wh, with I E ST. is mainly in the direction of μ4i This verifies condition (B) of Lemma 1. Therefore, by Hoeffding's inequality (15), for any W E V.
1
0f( P=,2")
V P", ..t)
V P 122
Pr
E
I
6
B
OW
OW
OW
(167)
REB,
<e- Be? < M
M
as long as
B BZ B log M.
(168)
Then,
E E Dx 1(x v): Sc
(169)
F EXTENSION TO MULTI-CLASSTFICATION
Define that 2"=classification is achieved by times of binary classification with the orthonormal (c) set (n() ,.... .,,l7) } the discriminative patterns for the task T. We have u
(1) 1. Vmn E [M], i E (c) The label y is c-dimensional with each entry chosen from {+1, -1). Specifically, each (X E gdxp, y E R²) ~ Dr is generated as follows:
* Randomly generate the k-th entry yk. k E [c] of the label y from [+1,-1) with an equal probability.
Each token is randomly chosen from {uf)
, - (1) Ji_n U (ui,...,,,, J. If y* 1 (or
-1). the number of tokens corresponding to μTv. (or - - is larger than that of - (or (6) and (((i")) are referred to label-relevant and confusion patterns for y/k = 1 (or y = 1), respectively. The average fractions of label-relevant and confusion tokens of fl are 5 and # respectively.
(.rz). u and - (or "-u
(t)
(i)
(1)
We then need sets of our binary model (4) to generate the output for 2f-clasiification, i.c..
f(x; 4) (fi(X; V), f2(;; V) ... .f((X:
f(X; = 1 Σ
P am) Relu Wo. Σ Wv. V. so,tmax T TW K. (Q,):))
1 P
(170)
32

---

Published as a conference paper at ICLR 2025
Wo, Wv,, WK, Mi-- The dimensions of Wo., Wv. WK, WQ i [c] follow Section 3.2.
with V/
The learning process is then e independent and parallel binary classification problems for each entry of the c-dimensional output. After fine-tuning, the trained model of each output entry has a similar
property Lemma I for single binary classification. dhd), the fraction of label-relevant pattern μT (t) i E [c]. may decrease by times in average from the binary classification scenario. Therefore, by condition (iii) of Theorem I, the number of iterations and samples increases by c² times, which is
polynomial of log scale of the number of classes 2". Then, for the disrminative patterns {PT, (8) Ji=1
of task Ti and { = and Tz of task T2, if for any PTg there exists a unique MT, close to be
(8)
orthogonal to MT then Ti and T2 are irrelevant. If for any ( there exists unique MTz with small angle to (or almost opposite to) UT, then Ti and T2 are aligned (or contradictory). We can then derive similar conclusions as Theorems 1 and 2 by combining the results of all the output entries.
33


Reviews:
[{"review_id": "dE36fIDeG8", "reviewer": "Reviewer_3a5G", "summary": "This paper explores the theoretical aspects of task vector arithmetic as a model editing technique for multi-task learning, unlearning, and out-of-domain generalization. The authors provide a theoretical analysis to justify why and when task vector methods are effective in nonlinear Transformer models, especially for binary classification tasks. They prove that task addition facilitates multi-task learning for aligned or irrelevant tasks, while task negation can effectively unlearn contradictory or irrelevant tasks. Additionally, they offer generalization guarantees for out-of-domain tasks and theoretical justification for task vector approximations. These findings are empirically validated through various experiments.", "questions": "I don't have questions.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well-written and easy to follow.\n- It provides a guideline for when and why task arithmetic works in multi-task learning, machine unlearning, and generalization to new tasks.\n- The discussion of low-rank approximations and magnitude-based pruning of task vectors supports the use of efficient approximation techniques in task arithmetic fine-tuning.\n- This is the first known theoretical generalization analysis of task vector arithmetic in nonlinear Transformer-based models, filling a notable gap in the literature.\n- The theoretical claims are validated through empirical experiments on the Phi-1.5 language model and Colored-MNIST image classification, adding practical credibility to the proposed framework.", "weaknesses": "- The theoretical analysis relies on a single-head, one-layer Transformer model, which may limit the applicability of the results to more complex multi-layer Transformer architectures.\n- While the empirical validation includes a large language model and a basic image classification task, the study could benefit from a broader set of tasks, including more complex or structured tasks beyond binary classification.\n- Although the theoretical framework outlines conditions for selecting arithmetic coefficients, more practical guidelines or analyses for tuning these coefficients in real-world applications would be beneficial.\n\ntypos:\n- line 288, \"fine-turning\" -> \"fine-tuning\"\n- line 388, \"are are\" -> \"are\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RPKg98kN1d", "reviewer": "Reviewer_3aPS", "summary": "This work focuses on the task vector in the context of task arithmetic, demonstrating its role in learning and unlearning tasks through experiments. It’s an interesting topic that needs more exploration. The authors find that task correlation affects the performance of the task vector and conduct theoretical justification on the low-rank approximation of the task vector and it's ability to adapt to out-of-domain tasks.", "questions": "* Line 236-238, the conventional attention expression is $softmax(W_QXX^TW_K^T)W_VX$, why is it written as $W_VXsoftmax(X^TW_K^TW_QX)$ in Formula 4?\n* Line 236-238, what is the meaning of $X^n$?\n* Line 242, Why is $x_i$ used here, while $X$ is used in Formula 4?\n* Line 261, Since $\\mu_T$ and $v_j$ are orthogonal, what is the meaning of  tokens corresponding to $\\mu_T$?\n* How to quantify the relevance of different language generation tasks？Are semantically similar and task-related equivalent?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Very comprehensive mathematical analysis and theoretical proofs\n* Discussion on the task vector is extensive.", "weaknesses": "* There are some issues with the paper's writing (Formula 4 and Definition 2 in the Section 3.2 is confusing).\n* In the language generation task, only a model with 1.5B parameters is used, and the experimental results are not fully meet expectations (also a noticeable performance loss in so-called irrelevant task).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F5Gewa8yAM", "reviewer": "Reviewer_L24R", "summary": "This paper provides a theoretical analysis of the effectiveness of task vector methods for model editing in transformers. The authors investigate the conditions under which task addition (for multi-task learning) and task negation (for unlearning) are effective, proving that task correlation plays a crucial role. They also establish conditions for successful out-of-domain generalization using task vectors. Experiments with both synthetic and real-world data validate the key concepts of the proposed theory.", "questions": "I found this paper highly engaging and believe it would attract even greater interest with an exploration of the theory across a broader set of tasks, particularly generative tasks. For example, could the proposed theoretical framework be extended to multiclass classification as an initial step? A discussion on how these insights might be applied to a wider range of tasks would substantially enhance the paper's appeal. I would be willing to increase my score if the authors could provide even preliminary ideas on these extensions.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is the first to theoretically examine the generalization of task vectors, filling a significant gap in current research.\n- Writing is well-written and easy to follow.\n- The theoretical contributions are well-supported by experiments, effectively bridging theory and empirical validation.\n- The theoretical insights align well with intuitive expectations regarding the effects of task vectors across aligned, irrelevant, and contradictory tasks.", "weaknesses": "- Although insightful, the data model may be overly simplistic for capturing the complexities of real-world data. For example, even for simple yes/no questions, a negation word in a sentence may flip the relevance of certain words in the sentence, which cannot be captured by the proposed data model. I wonder if some theoretical aspects can be generalized independently of this data model.\n- The analysis is restricted to a one-layer transformer with limited nonlinearity, despite claims in the title and introduction regarding the challenges of analyzing nonlinearity in task vectors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OgGCsTy8Kp", "reviewer": "Reviewer_QzER", "summary": "The paper produces the first theoretical result on the ability of task vectors to generalize to new tasks and ability to unlearn a task. The study is conducted under the lens of \"aligned\" vs \"irrelevant\" vs \"contradictory\" tasks. The authors studies the necessary fine-tuning hyperparameters and the task vector coefficients ($alpha$) that enable generalization and unlearning.\n\nAdditionally, the authors also verify the theory result with experiments on toy CMNIST dataset and on next token prediction task.", "questions": "Definition 2 -- what is the dimension of $mu_{\\tau}$? is it the same as $v$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Very exciting theoretical result that includes very practical aspects (e.g., hyperparameters of the fine-tuning, how to set alpha for each task vectors, etc)\n- Novel characterization of generalization conditions on 1 layer transformer, building up on previous work that uses NTK assumptions\n- Practical scenarios in terms of the relation between tasks (aligned vs irrelevant vs contradictory)\n- Nicely written and relatively accessible to non-theory people. I particularly like the remark after each theorem that explains what the theory is and what does it imply\n- Nice setup on CMNIST that reflects the aligned vs. irrelevant vs. contradictory condition, followed by a nice non-toy experiment.", "weaknesses": "N/A -- good paper overall :)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)