Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations
Authors: Yuxin Dong
Jiachen Jiang
Zhihui Zhu
Xia Ning
Venue: ICLR 2026 Poster (2026)
Keywords: transformer
in-context learning
task vector
Primary area: foundation or frontier models, including LLMs
TL;DR: We explain how task vectors emerge and function in in-context learning, and point out their limitations.

Abstract: Task vector is a compelling mechanism for accelerating inference in in-context learning (ICL) by distilling task-specific information into a single, reusable representation. Despite their empirical success, the underlying principles governing their emergence and functionality remain unclear. This work proposes the *Task Vectors as Representative Demonstrations* conjecture, positing that task vectors encode single in-context demonstrations distilled from the original ones. We provide both theoretical and empirical support for this conjecture. First, we show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis. Next, we predict the failure of task vectors in representing high-rank mappings and confirm this on practical LLMs. Our findings are further validated through saliency analyses and parameter visualization, suggesting an enhancement of task vectors by injecting multiple ones into few-shot prompts. Together, our results advance the understanding of task vectors and shed light on the mechanisms underlying ICL in transformer-based models.

Published as a conference paper at ICLR 2026
UNDERSTANDING TASK VECTORS IN IN-CONTEXT LEARNING: EMERGENCE, FUNCTIONALITY, AND LIM-
ITATIONS
Yuxin Dong, Jiachen Jiang, Zhihui Zhu & Xia Ning Department of Computer Science and Engineering, The Ohio State University {dong. 13557, jiang. 2880, zhu. 3440,
edu
ABSTRACT
Task vector is a compelling mechanism for accelerating inference in in-context learning (ICL) by distilling task-specific information into single, reusable rep- resentation. Despite their empirical success, the underlying principles governing their emergence and functionality remain unclear. This work proposes the Task Vectors as Representative Demonstrations conjecture, positing that task vectors encode single in-context demonstrations distilled from the original ones. We pro- vide both theoretical and empirical support for this conjecture. First, we show that task vectors naturally emerge in linear transformers trained on triplct-formatted prompts through loss landscape analysis. Next, predict the failure of task vec- tors in representing high-rank mappings and confirm this on practical L.LMs. Our findings are further validated through saliency analyses and parameter visualiza- tion, suggesting an enhancement of task vectors by injecting multiple ones into few-shot prompts. Together, our results advance the understanding of task vectors and shed light on the mechanisms underlying ICL in transformer-basee models.
INTRODUCTION
In-context learning (ICL) is core capability of large language models (LLMs), allowing them to perform new tasks without parameter updates by conditioning on a few input-output examples the prompt (Brown al. 2020 Unlike traditional training, ICL relies attention-based mechanisms to infer task structure directly from context. This surprising generalization ability has led to growing interest in uncovering the principles of learning purely from contextual examples (Xie et al. 2022
Chan et al. 2022 Dai et al. 2023 Shen et al. 2024 Deutch et al. (002)).
A recent work investigates the task vector method (Hendel et al. 2023) (concurrent works include function vectors (Todd et al. 2024) and in-context vectors (Liu et al. 2024)) a technique that distills underlying task information from ICL demonstrations into a single vector. Typically, ICL prompts are structured as sequences of triplets, each encoding a semantic mapping, in addition to query at the end (e.g., "hot + cold, up - down, dark > "). Task vectors are then extracted from the hidden states of the last (-+) token. Once obtained, these vectors can be injected into new zero-shot prompts (e.g.. "big ")), enabling the model to generalize to unseen inputs in a zero-shot fashion.
Task vectors naturally emerge even in small transformer models trained from scratch (Yang et al.) 2025), suggesting that their formation is a general property of attention-based architectures. Recent studies further demonstrate that task vectors can be enhanced by aggregating hidden states across multiple layers and arrow tokens (Li et al. 2024). Beyond language models, task vectors are also effective in large-scale visual (Hojel et al. 2024) and multi-modal (Huang et al. 2024) models.
Despite their empirical effectiveness, the underlying mechanism task vectors, especially how they emerge, function, and encode task information, remains poorly understood. This paper takes step toward unveiling the principles behind by introducing the following conjecture:
"Corresponding author.

---

Published as a conference paper at ICLR 2026
(a)
light
(b)
Predict Output
0
Weighted
Voctor Formation
0
0
0
0
Preprocessing Stage
hot
cold
up
down
day
night
dark
big
Figure/ 1: Overview of task vector and our main conjecture. (a) Task vector emerges during ICL by distilling from the preceding in-context demonstrations. (b) It can then be injected into zero-shot prompts and functions single, representative demonstration, facilitating efficient prediction.
Conjecture (Task Vectors as Representative Demonstrations)
The injected task vector facilitates zero-shot inference by encoding single represemtative demonstration, distilled from the original in-context examples.
An intuitive illustration is provided in Figure 1 In the following sections, we validate this conjecture through various empirical and theoretical perspectives. These analyses comprehensively explain how task vectors naturally emerge within iitentionnbbsse model architectures, effectively encode task-related information, and facilitate inference in zero-shot prompts. Our work advances the un- derstanding of the underlying mechanisms behind ICL, clarifying both the efficacy and limitations of task vectors in transformer-base LLMs. The highlights of this paper are as follows:
Theoretical Justification in Linear-Attentin Models:
We theoretically characterize the critical points of linear-attention models and demonstrate how they solve random linear regression tasks through embedding concatenation and gradient descent. With triplet-formatted input prompt structure, task vectors naturally emerge at arrow tokens as weighted summations of the in-context demonstrations, potentially enhancing robustness under representational perturbations by redun- dantly encoding task information. Empirically, the learned linear model parameters closely align with the predicted structure and successfully replicate the task vector mechanism.
Empirical Verification Practical LLMs: We visualize the information flow in LLMs with saliency analysis and observe patterns consistent with linear models, suggesting they share similar underlying mechanisms. According to our conjecture, inference with task vectors is analogous 1-shot ICL., which is inherently limited rank-one meta- peeta-predicorrs under the gradient descent perspective. To validate this, we introduce a series of bijection tasks that are provably unsolvable by rank-one predictors, and empirically confirm this failure in real-world transformers. Building on these insights, we enhance the standard task vector method by injecting multiple vectors into seo-shot prompts, resulting in consistent performance gains across a range of ICL tasks.
RELATED WORKS
Theory of ICL. Recent analyses have shown that attention layers can simulate gradient-descent
algorithms for regression tasks (Garg et al. 2022 Von Oswald et al. 2023a; Ahn al.) 2023) Wu et al. 2024). Other works study generalization and sample complexity (Xie et al. 2022 Chan et al.,
2022 Shen al. 2024; Von Oswald et al. 2023b Deutch et al., 2024). These works reveal the inductive bias of attention but leave open how abstract task representations are formed or encoded.
Task Vector Mechanism. Multiple recent works identified the mechanism of task vectors during ICL inference (Hendel al. 2023 Todd et al.) 2024 Liu et al. (024). These vectors emerge in the pretraining stage of LLMs Yang et al. 2025) and extend beyond text to vision (Hojel et al. (024) and multimodal (Huang et al. 2024) models. Despite the effectiveness, their underlying mechanism remains poorly understood. A concurrent work (Bu et al. 2025) interprets them via word2vec-1 additive scheme, but is limited to simple additive tasks, single-token prompts, and 1-layer models. In contrast, our analysis extends to pairwise or triplet prompts and multi-layer attention.
A more comprehensive discussion of the related works can be found in Appendix A.2
2

---

Published as a conference paper at ICLR 2026
2 SETTING: LINEAR REGRESSION WITH LINEAR- ATTENTIOO MODELS
Notations: We write [n] = {..-..,... ,n). The Hadamard product is denoted by o, and the Kronecker product by @. The identity matrix of dimension n is denoted by Ln while O and 0mxn represent zero vectors or matrices of the corresponding dimensions. Subscripts are omitted when the dimen- sions are clear from context. We define M(M) = {A E Rdim((A) A = M o A, E Rdim(Ad) as the set of masked matrices induced by mask M. For general matrix A, the element at the i-th row and j-th column is denoted by Aa.j. and the sub-block from rows to & and columns j to / is denoted An) represents the block-diagonal matrix constructed by {A.)
by Ask,jd diag(A1,
Random Linear Regression: Following works (Garg et al. 2022 Von Oswald et al. 2023a Ahn
et al. 2023 Wu et al. (024), we consider training linear transformers on random instances of linear |++1 regression. Let {2; i=1 where I; E Rd, denote covariates drawn i.i.d. from distribution Pz. and let [w where WA E Rd, denote coefficients drawn i.idd from distribution Pw- Define the coefficient matrix W = [w ... wd] T Rdxd, The responses are then generated as y4 = Wa; for E [m+ 1]. We denote by X,Y E Rdxx the matrices whose columns are x; and ye, respectively. The query covariate and response are denoted by Ttest 2n-1 and Vhest Vn++1 respectively.
Linear Self-Attention Model: Following prior works (Von Oswald et al. 2023a Ahn et al. 2023
Wu et al. 2024). consider transformers composed of linear self-attention layers. Let Zo E R-axap denote the input matrix constructed from X, and Frest but excluding Bhests where dp denotes the number of tokens and varies across prompt structures. The model is defined by stacking L attention blocks with skip connections, where the I-th layer is expressed as:
Z = Z + A/ Attnv,,q. (Z-1),
o (Z) = VZM (ZT QZ).
(1)
Here, the trainable parameters are (Vi, Qt where E R2dx2d denotes the projection and value matrices, and Q E 2ddx2d denotes the query and key matrices. Following the work (Ahn et al.) 2023), we adopt a masking matrix M = di 0) to prevent attention from earlier tokens to the final one. The output of the model is defined as TF Zo: (Vi, Q1H1,) = (Z (i.e., the latter half of the last column). This definition aligns with the structure of the input Zo. which will be further discussed in subsequent sections. During training, the parameters are optimized to minimize the expected ICL risk over random linear regression instances:
((VV, 11 Ez,.w |F (Zo; (Vi 111 + Witest 1
(2)
3 EMERGENCE OF TASK VECTORS IN ATTENTION MODELS
Firstly, we present theoretical evidence that task vectors naturally arise in simple linear transformers. Specifically, we analyze the loss landscape of the in-context risk, focusing the properties of its critical points. As a startup, recall the standard linear regression setup (Ahn et al. 2023 Wu et al.) 2024). where the (Iz: y/) pairs for each demonstration are concatenated to form the input prompt:
Zo X
Frest 0
T2
Fn Prest 0
R 22
dp n +1
V2
Dn
2023 Zhang et al.. 2024 Mahankali et al. 2024), each According to existing analyses (Ahn et al., attention layer in this setting performs one step of gradient descent on the learned coefficient ma- trix. Specifically, the theoretically optimal single-layer (possibly nonlinear) attention (Katharopou- los et al. 2020) implements the following predictive function (Ahn et (023) when the covariates are drawn from P. = (00, La), by selecting V1 x diag(Odxat Ia) and Q2 x diag(La. Odxa):
(4) (xn)) as oX). This model employs W" x (x) as an estimate of W, yielding prediction Mtost = Wooootts)) This paper considers alternative settings more reflective of practical scenarios, where I; and yi are separated as distinct tokens. As noted (Zuo et al.] 2025). such separation necessitates the usage of position encodings for bi-directional attention. Following prior analysis (Kazemnejad et al., 2023), we assume that position encodings are appended to the input tokens, and reformulate the wise update rule of self-attention as:
TF(Zo; (V,
(X) (Ztest), where IRd + IR" is kernel function.
Here, we abbreviate (0()))
[
Attnv.e(Z) = -VZM TTT PT] 2
where P E Xd,
(5)
P
3

---

Published as a conference paper at ICLR 2026
For analytical tractability, we take P = Id, one-hot position encodings. Following previous work (Ahn et al. 2023 (see Appendix A.3 for more explanation), we further impose that:
Vi = diag (A, B:), Q: diag(Ci Oaxd; Di), where A, Bi. Ct R Di E Rd (6)
These parameterizations ensure that the projection and attention operations act independently on the covariate, response, and positional components of the input. This structural decoupling is essential for understanding how the transformer identifies the dependency between each (zs, yk) pair and revealing the actual optimization algorithm being executed by the model. The proofs for the main theoretical results in this paper are available in Appendix D
3.1 WARM-UP: LEARNING WITH PAIRWISE DEMONSTRATIONS
We begin by analyzing the optimization of linear transformers on pairwise demonstrations. Follow- ing previous approach (Garg al. [022 Wibisono & Wang 2023
Xing et al. 2024), we decompose
each demonstration in eq. (3) into pair of tokens Z tical ICL prompt structure:
R 2dx2
to better reflect the prac-
M6
Zo
Zu
Z
Zoest
21
En 0
0
2 test
0
dp - 2n + 2.
Un
0
The following theorem suggests that certain critical points of the in-context risk effectively solve the regression problem by first concatenating each pair of (II into the same tokens, and then executing variant of the gradient descent algorithm compute the prediction. To simpliff notation, we denote A = {Ai)= (similarly for B, C, and D) and present:
Theorem 1 (Critical Points; Pairwise Demonstrations), Assume that P. NN E) and P = N(OE-1) with Σ dddd satisfying > 0. Define Sg Idddd and Sp 12 xdp as
S = (Aa E},, S== = X R). Sp = (i @ A2) A,, E R2x2].
Consider optimizing an L-layer transformer under parameter configuration in eq. (6) we have
inf A,BESy, cest, DESt EHEAUBUCUD (VV.
To understand the behavior of these critical points within self-attention layer, we fix > La and take A,, B: Ia. C AAd, and D, diag(T, @ A1, A2). Let the first and last d rows of Z be denoted by X₁ and Y,, respectively. Under these settings, the update rule of each layer becomes:
Zi = -
+ [
Ai Z diag(1, ())
(8)
The above update can be decomposed into the following two distinct components:
Gradient Descent: The first component, Z, + Z,-1 - XZ-- X1-1, implements the GD++ algorithm (Von Oswald et al. 2023a). This variant enhances convergence speed over stan- dard gradient descent by improving the condition number of X Notably, this operation modifies only X, but not Y; for the first layer, as implied by the structure of Qi (eq. (6))
Embedding Concatenation: The second component, Zj < Zi_1 + Zi_,A, for E [m], mixes each pair of (x,, y:) tokens. Given that I; and y/ tokens are initially linearly separable in our formulation, this operation concatenates each (zz, pair, thereby transforming pairwise demonstrations into the original single-token format. For the query token 2,nest, this operation copies Ftest into the final token, reconstructing the structure in eq. 33), where each non-final token directly concatenates (I;; yi) of a demonstration, and the final token contains only Trest-
In summary, our analysis reveals that for pairwise demonstrations, the first attention layer leverages position encodings to distinguish between covariate and response tokens, subsequently concate- nating them to form single-token prompt structure. The remaining layers then apply the GD++ algorithm, mirroring the learning dynamics on single-token demonstrations. As a result, an L-layer linear transformer allocates layer for embedding concatenation and utilizes the remaining J-1 1 layers to perform gradient descent. In Figure we visualize the learned D, weights under the setting of Theorem T] and observe that they closely match the critical point structure of Sp.
4

---

Published as a conference paper at ICLR 2026
0.1
0
0.0
-1
(a) Di (Pairwise)
(b) Di (Triplet)
(c) Aa]]
Figure 2: Visualization of learned Dr weights. (a) Pairwise demonstrations yield block-diagonal
structure aligned with Theorem [1] (b) Triplet demonstrations yield richer structure aligned with
Theorem 2 (c) The learned matrix A₃ has nearly orthonormal rows as suggested by Proposition
3.2 EMERGENCE OF TASK VECTORS WITH TRIPLET DEMONSTRATIONS
Next, to better reflect the prompt structure of practical ICL, we insert additional zero tokens between each pair of (Ir, ys) to simulate the arrow (-) tokens. This reformulates each demonstration triplet (Iz, 7, yr), enabling us to analyze the critical points with these triplet demonstrations:
Zo
21
0
Fn
0
Teest
dp 3n +
(9)
0
0
0
Yn
0
0
Theorem 2 (Critical Points; Triplet Demonstrations), Assume that P2 NO,, and P NO, E-1) with E E Idxd satisfying E > O. Define S Sz C IRddd and Sp C Rd as
S (Maa |EER),
Sy: {xE IXER..
Sp dia @ A2) + 4- 2 Ag + & Au
M, E M (1..) As (!!!)
(000)
A E As M (8008 (810))
Consider optimizing an L-layer transformer under parameter configuration in eq. 66 we have
inf A.BESt. cest DESS EHEAUBUCUD VHC((V..Q)))l 0.
To analyze the behavior of each attention layer, note that the critical points for the matrices A Bi, and C remain consistent with Theorem I thereby implementing the GD++ algorithm. For the matrix Di, we decompose its structure into three distinct components:
The first component, diag (n, @ A1, A2), mixes each pair of (z,,4.) tokens, effectively concatenating them analogous the operation analyzed in the previous section. This converts all non-arrow tokens into single-token demonstrations.
Embedding Concatenation:
Self Magnification: The second component, In+1 @ Ag, scales the embeddings corresponding to each arrow (-) token by fixed constant and adds them back to themselves.
Task Vector Formation: The third component, Ag @ As, performs weighted summation across all demonstrations in the prompt. This operation is central the emergence of task vectors. Let [8 But1] Rox(a+1) denote the first rows of A4 (we will soon show that the last row of A4 converges to zero), the first self-attention layer then outputs + linear combinations of the demonstrations
the hidden states for the arrow tokens, expressed as 2 X i [n + 1], where a1 a2 E IR are the two non-zero entries of As. These vectors then be injected into zero-shot prompts and function as single-token demonstrations.
8, for
This mechanism provides strong theoretical evidence for our main conjecture, demonstrating that task vectors naturally emerge from the pretraining stage linear-attention transformers on triplet-formatted prompts. Notably, the structure of Sp closely aligns with our visualization of D, in Figure 2bb confirming our theoretical analysis. We now further investigate the structure of the weight matrix A4, and present the following result:
5

---

Published as a conference paper at ICLR 2026
31 VI Do
DE
D1
32
92
166
0.20
30
1.11
0.10
1.2
5
10
Salieney Map 10)
(b) Saliency Map 12)
Task Vector Weights
Figure 3: Visualizations on Llama-7B: (1) saliency matrices bipartite graphs between layer and (o), edge widths indicate saliency magnitude; (2) variations in the extracted task vector after perturbing the i-th demonstration (I) and the optimal task vector weights (-) obtained by optimizing Proposition 5 (a) Each y/4 token attends to its corresponding (r,, 9c) pair, reflecting embedding concatenation. (b) The final (-+ token attends broadly to all y4 tokens, indicating task
vector formation. This occurs just before the optimal injection layer (1 13). (c) The predicted task
vector weights closely match the trend of empirical results, validating our theoretical model.
Proposition 3 (Optimal Task Vector Weights). Assume P N(, Ia). Consider optimiz- ing 2-layer linear-altention transformer with triples demonstrations and parameter configuration given in eq. (6) and assume C₁ Odxa- Let
=
D1 diag(L, @ A2) + Lat1 +
be any minimizer of the in-context risk f 1)) then have Aa E Sty where
Su = {A AAT = >'diaa 0), ER.
This result suggests that the optimal A4 weight matrix satisfies two key properties: (1) the last row zero, and (2) the first n rows are mutually orthonormal. These conditions imply that the learned weight vectors -i. _3 are likely be distinet. Therefore, the diverse linear combinations of the demonstrations, thereby enriching the representation within the input prompt. This implication is verified in Figure 2 While task vectors are typically extracted from the final arrow (-) token in standard usage, here we consider all arrow tokens as task vectors as bi-directional attention allows each aggregate information from the full prompt.
task vectors produce
PREDICTED FAILURE OF TASK VECTORS ON BIJECTION TASKS
We then present an empirical observation that supports our conjecture. Consider the setting where task vectors are injected into zero-shot prompts. Based on our prior analysis, the injected task vector 2tv is formed weighted summation of the original demonstrations. As result, we show that the injected prompt reconstructs the single-token structure in (3) with only 1 demonstration:
Zn
[Etest
Ztv
0]
Fiest
Itv
0]
=
Ztest
a1X8 agY B
0
E 2d2x3
(10)
Vtv
where the weight vector 8 E R comes from the last column of Aq. and the weights 01, 02 come from As (see our discussion after Theorem 2). After the first layer, the A2 matrix of Sp moves Tiest the last token, reducing the prompt
single-shot, single-token demonstration. Ac- cording to the optimal single-layer transformer (eq. (4)), the estimated coefficient matrix is now W" a100y (( which is rank-one. Therefore, task vectors are inherently limited in their expressiveness: they only replicate 1-shot ICL, which is restricted to rank-one coefficient matri- ces. This implication also naturally extends to multi-layer transformers.
While our analysis is conducted on linear-attention transformers, demonstrate that similar learn- ing patterns also emerge within practical LLMs. Specifically, we visualize the layer-wise informa- tion flow between tokens using saliency maps (Wang et al. 2023), where the saliency score for each attention matrix is computed as S(A)) = Σn [ OL A, Ar, denotes the attention matrix of the h-th head at layer I, and L is the ICI. loss (i.e., the cross-entropy loss for predicting frest).

---

Published as a conference paper at ICLR 2026
Table 1: Comparison of the accuracies of many-shot ICL and task vector on bijection tasks (Llama- 7B, n = 10). We use gray text to indicate accuracies lower than 60%.
+
y
y
+
N' ++ y
Task
Domain 2'
Domain y
Example
ICL
TV
ICL
TV
ICL
TV
To Upper
{a....
[A.....2)
a A
1.00
0.91
1.00
0.99
1.00
0.55
English
French
hello - bonjour
0.83
0.84
0.82
0.70
0.54
0.35
Translation
English
Italian
hello - ciao
0.84
0.78
0.82
0.74
0.70
0.47
English
Spanish
hello + hola
0.92
0.88
0.89
0.75
0.64
0.43
Present
Gerund
go + going
0.99
0.95
1.00)
0.97
0.80
0.41
Linguistic
Present
Past
go > went
0.98
0.91
0.99
0.96
0.52
0.33
Present
Past Perfect
go gone
0.82
0.82
0.94
0.65
0.55
0.33
Singular
Plural
dog dogs
0.88
0.78
0.94
0.89
0.76
0.51
Copy Antonym
{a, 2, A. 2}
A > A
1.00
0.98
Adjectives
happy sad
0.89
0.83
0.83
0.73
As demonstrated in Figures of embedding concatenation and weighted summation. This suggests that real-world transformers implement similar algorithm to solve ICL. tasks and, consequently, inherit the same expressiveness limitation. The full saliency score maps are given in Appendix [B.5]
3a and 3b the saliency maps reveal certain patterns matching the ones
To verify this, we construct specialized class of ICL tasks, named bijection tasks. Specifically, given a bijective mapping from domain X' to codomain y. one can combine it with its inverse mapping to form new task that maps X' X' onto itself. For instance, combining the "to uppercase" task with its inverse "to lowercase" yields a bijection task that maps each letter to its opposite case, and valid ICL prompt takes the form: "a - A. B b, c c C, D - Note that this differs from task superposition (Xiong et al., (2024). as each input corresponds to a unique, well-defined output. We then establish a key limitation of rank-one coefficient matrices in addressing such tasks:
Proposition 4. Let y E Rd be non-zero vectors. Then the following are equivalent: (1) There exists a rank-one matrix W E Rddd such that y = Wx and x = W Wy (2) x = y or -y.
This result highlights that rank-one coefficient matrices cannot solve general bijection tasks, and are restricted to two special cases: the identity mapping (a = y). or the negation mapping (x = -y). We further verify this implication real-world LLMs: in Table [ both ICL and task vectors perform well on the original tasks and their inverses. But for bijection tasks, while ICL preserves performance in many cases, the task vector method consistently fails, confusing examples from the two domains and yielding near-random predictions (50%) (e.g., in "To Upper", task vectors predict the correct letter but fail to distinguish between uppercase and lowercase. See Appendix [B.4 for further results). The only exceptions are Copy and Antonym, the special cases in Proposition
Together, these findings empirically validate our main conjecture: the task vector approach, which is restricted to one-shot ICL, is limited to rank-one mappings and cannot solve general ICL tasks (e-g., bijection tasks). While a variety of ICL tasks have been explored to assess the capabili- ties of task vectors (Hendel et al. [2023] Todd et al. 22224] Li et al. (2024), the fundamental limitation of task vectors in addressing these bijection tasks has not been previously identified.
FURTHER DISCUSSIONS
Effect of Causal Attention and Dropout. While task vectors naturally emerge in linear attention, their embeddings do not directly help minimize the ICL. risk, as evidenced by the identical perfor- mance between pairwise and triplet formatted prompts (Figures 44 and (4b). Instead, we show that task vectors do contribute to optimization under loken-wise dropout, acting as redundancies for in- context demonstrations that may be randomly dropped during training. This redundancy ensures that essential task information is preserved to facilitate inference despite partial context loss.
Proposition 5. Under the same settings as Proposition 3 consider adding token-wise dropouts Oi:
7 7,-0, Attny,
where Oi diag(o), .., 0 of lid. Bern(p).

---

Published as a conference paper at ICLR 2026
10
5 !)
T( 2) 5
5 2)
T 59
0.4
ICL 1]
lask Vector
ICL Risk
10F
ICL Risk
10
ICL Risk
00.3
10
10
0.1
10-2
10
0.0
5
10
15
20
25
10
15
20
25
(a) Single- vs. Multi-Token
Token (L = 2)
(b) Single- vs. Multi Token (L = 3) (c) ICL vs. TV
Multi-Token
Figure 4: (a, b) Comparison of the best ICL risk achieved using single (S), pairwise (P), and triplet
(T) formatted prompts. (c) Performance comparison between 1-shot ICL and task vector.
Then any minimizer A4 of the in-context risk L ((V,, 2i)(sn) satisfies (A4)++1. 0 and:
(A4)im; & arg min e1 callA|1++ 02 En 11A, Σ n=1 11A:s112caalAAT 1|- s.t. ||AI 1.
A
where c C1 are non-negative constants depending on Vi. Qi and p.
This result suggests that dropout introduces additional higher-order regularization on the task vector weights, encouraging them to distribute more uniformly across demonstrations. Furthermore, when considering causal attention (i.e., enforcing A4 to be upper-triangular), it induces a decaying weight pattern from later to earlier demonstrations, which exactly matches the practical behavior observed in practical transformer models (as evidenced in Figure 3c
Decoding the Vocabulary of Task Vectors. Multiple prior works (Hendel et al. 2023 Todd et al,, 2024) have observed an interesting phenomenon that, when task vectors are directly decoded through the final classification layer, the top tokens often belong to the output space of the current task (see Table the Appendix). Our theoretical analysis provides a natural explanation for this: assuming a 2d-dimensional hidden state space partitioned into input (2,) and output (y:) halves, the output half of task vectors then encodes weighted summations of yt- Since the final prediction relies on the output half, decoding a task vector yields a combination of yt. which is likely lying in the output space. This observation suggests that practical LLMs adopt similar hidden-state partition.
Extra EOS Tokens. In our previous analysis, we consistently imposed an additional zero token at the end of the input prompt. While this token be interpreted as EOS token in practical models, such a design choice is uncommon in standard ICL tasks. We justify this modeling decision with: Proposition 6 (Informal). Given any L-layer, 1-head, d-dimensional EOS, there exists equivalent L-layer, 2-head, 2d-dimensional model operating without EOS.
model with
This equivalence suggests that the same learning dynamics can be realized through multi-head ar- chitectures without relying on explicit EOS tokens. Specifically, the first head is dedicated to task vector formation, while the other handles ICL prediction. This separation allows the model to retain the functional role of the EOS token implicitly within its hidden states.
6 EXPERIMENTAL STUDIES
6.1 SYNTHETIC RESULTS WITH RANDOM LINEAR REGRESSION
In this section, we validate our critical points analysis with synthetic linear regression tasks. Specif- ically, we examine the achievable ICL risk of linear-attention models with single-token (eq. ( pairwise (eq. (7), and triplet (eq. 9 demonstrations. We set the input dimension to d 4 and train multiple models with different random seeds and
P. Pa. (0, Ia). For each setting, report the minimum ICL risk achieved as a proxy for the global optimum. The comparative results across different numbers of layers L and demonstration formats are shown in Figures 4a and 4b
These results support our theoretical analysis: when trained with pairwise or triplet demonstrations, the model recovers the GD++ algorithm similar to the single-token case. Notably, the performance
8

---

Published as a conference paper at ICLR 2026
Table 2: Accuracy comparison between few-shot ICL (Baseline), the task vector method (TaskV),
and our strategy (Tas The experiment is conducted on Llama- with n = 10.
Method
Knowledge
Algorithmic
Translation
Linguistic
Bijection
Average
Baseline
6,90 2.08
15.60 1.72
7.00 1.68
12.44 1.74
8.27 1.33
10.28 0.98
0-shol
V
68.80 a
86.20 1.61
73.53
85.24 1.80
50.67 2.32
72.26
Baseline
69.50 # 3.86
73.67
1.56
57.80 + 2.01
56.22 1.57
44.76 = 2.44
58. 0.63
1-shot
TaskV
79.50 # 2.35
88.47 0.75
80.67 = 2.56
89.11
+ 0.84
60.44 2.07
78.79 0.77
TaskV-M
81.30 +: 2.80
89.53 0.68
80.13 2.14
88.71
0.62
61.78 0.96
79.34 0.37
Baseline
78.80 +: 3.30
85.07
1.37
75.67 + 2.64
76,80 +: 1.18
56.49 =: 2.87
72.92 0.59
2-shot
TaskV
84.60 à 2.11
88.40 0.68
84.33 1: 0.92
90.13 1. 0.92
62.44 a: 2.16
80.82 0.42
TaskV-M
85.70 # 1.63
89.27 1.10
84.13 + 1.15
89.64 + 0.86
64.49 = 2.02
81.48 0.37
Baseline
86.20 + 2.69
88.07
1.06
80.00 + 1.67
84.04 1.19
62.18 = 1.52
78.51 0.42
8-shot
TaskV
90.20 +: 2.23
88.67
0.89
86.27 :: 2.31
92.31
+ 0.48
66.53 = 0.94
83.53 0.41
TaskV-M
90.30 # 1.50
89.87 0.83
86.07 = 2.17
92.36 + 0.72
68.13 = 0.76
84.15 0.52
Baseline
84.80 à:
88.07 0.61
83,27
88.89
67.16
81.52
4-shot
TaskV
88.70 = 1.69
89.53 1.34
86.27 - 1.08
92.76 0.54
70.44 = 1.35
84.66 0.39
T-s
89.60 = 1.43
91.00 1.01
87.20 0.62
92.36 1.44
72.53 0.94
85.64 0.29
of L-layer models with pairwise (P) and triplet (T) demonstrations closely aligns, indicating shared underlying learning pattern. Moreover, their performance consistently lies between that of single- token (S) case T.-layer and (I. - (1--1)-layer models. The observed improvement over the (T. layeer single-token baselines comes from the additional GD++ performed solely on H, tokens in the first layer, effectively acting "half-step" of gradient descent.
We then reproduce the task vector method linear models. Specifically, we extract the hidden state of the final (-) token from triplet demonstrations after the first layer, and inject this vector into zero-shot prompts consisting of Ftest only. To simulate the effect of layer normalization, we normalize the task vectors before inference and the output vectors before ICL risk evaluation. As shown in Figure 4c] the performance of task vectors is highly related to that of standard 1-shot ICL. This validates conjecture that the injected task vector effectively acts as a single demonstration.
6.2 ENHANCING THE TASK VECTOR METHOD
We further explore an enhancement to the original task vector method. According to our previous analysis, single injected task vector may not provide sufficient information for inference com- plex tasks (c.g., bijection tasks). Moreover, in linear-attention models, each (-) token functions individual in-context demonstration during the gradient descent phase and thus contributes equally to the ICL risk. Motivated by this, we extend the standard task vector method, which modi- fies only the final arrow token, and propose a multi-vector variant that injects into every single arrow token in few-shot prompts. This enriched injection scheme enables the model to leverage multiple new demonstrations, thereby providing more informative and distributed context for prediction.
We compare our multi-vector injection strategy V-M) against standard N-shot ICL (Baseline) and the original task vector method (TaskV). Note that Baseline uses few-shot ICL and Task V is jecting into few-shot prompts, which are different from the settings in Table Thich which uses many-shot prompts for ICL. and zero-shot prompts for task vectors. For each N-shot prompt, we generate N+ distinct ICL. prompts to produce N 1 task vectors, which are then used to replace the embeddings of all arrow tokens in the input. For each task, performance is evaluated over 50 randomly sampled prompts, with mean accuracy and standard deviation reported across 5 independent trials. The final results, summarized in Table 2 span a diverse set ICL task types, showing that Tas consis- tently outperforms TaskV, especially the challenging bijection tasks. While the improvement is not dramatic, we believe that the current results sufficiently demonstrate the potential of multi-vector injection, thereby providing insights for the design of future ICL or task vector methods.

---

Published as a conference paper at ICLR 2026
7 CONCLUSION, LIMITATIONS, AND FUTURE WORKS
This paper proposes a plausible explanation for the emergence and functionality of task vectors in ICL. We support this conjecture with both empirical observations and theoretical analysis, demon- strating how task vectors naturally arise under ICL-style training prompts, and why this method inherently fails on general ICL tasks beyond rank-one mappings. Our work provides a new perspec- tive on the underlying mechanisms and offers promising direction for interpreting intermediate hidden states in modern transformer-basee language models.
While our analysis provides new insights into the emergence and functionality of task vectors, it primarily conducted on simplified linear-attention transformers and synthetic tasks, which may not fully capture the complexity of real- wol-w LLMs. Moreover, our theoretical framework focuses solely on critical point analysis, and there is still a lack of convergence guarantee or sample com- plexity analysis to fully understand the learning dynamics during model pretraining.
Future directions of this work may include: (1) extending the current theoretical framework to causal and multimodal settings; (2) exploring how richer architectures (c.g., non-linear attention) train- ing objectives (c.g., auto-regressive loss) influence the behavior of task vectors; (3) synthesizing orthogonal enhancements of the task vector method (e.g., function vectors (Todd et al. 2024 and
in-context vectors (Liu et al. 2024)), and extending to more complex reasoning tasks.
ACKNOWLEDGMENTS
This work was supported by the Natural Science Foundation under grants IIS IIS-231222 2312400 and IIS- 2402952. We would like to thank the anonymous reviewers for their valuable suggestions.
ETHICS STATEMENT
This work advances the theoretical understanding of in-context learning and task vector mecha- nisms, which can lead to more efficient and interpretable language models. By enabling faster inference through task vectors, it may reduce the computational cost and energy consumption of large-scale deployment, thereby making AI systems more accessible and environmentally sustain- able. Improved interpretability could also enhance trust and transparency AI applications across education, healthcare, and other socially beneficial domains.
As task vector methods improve efficiency and transferability, they may also be misused to repli- cate or extract functionality from proprietary models without authorization, raising concerns around model intellectual property. Additionally, while interpretability is often framed as a benefit, deeper insights into model internals could be exploited to engineer adversarial inputs or extract sensitive training data. Careful consideration and mitigation strategies are essential to ensure that such work aligns with the broader goals of safe and beneficial AI.
REPRODUCIBILITY STATEMENT
We provide complete proofs for main theoretical results in Appendices C and D experimental details about the dataset and implementation in Appendix and full source codes to reproduce our experimental results at
USAGE OF LLMS
We used LLMs only to improve grammar and polish academic writing. All technical ideas, proofs, experiments, and conclusions were entirely conceived and verified by the authors.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information
10

---

Published as a conference paper at ICLR 2026
Processing Systems, 36:456144456533 2023.
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. URLhttps: http:////ppeeeeeeeeeeee net/ forum?
1d=000X44ByN441
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Qingfu Zhang, Hau-San Wong, and Taiji Suzuki. Provable in-context vector arithmetic via retrieving task concepts. Forty- International Conference on Machine Learning, 2025. URL https: httt:s:///penreview..coo net wtweeeooouuraee Pefooormm?i DbUmeNnNpt
Stephanic Chan, Adam Santoro, Andrew Lampinen, Jane Wang. Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learn- ing in transformers. Advances in neural information processing systems, 35:18878 8858-888-188,, 2022.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent meta-optimizers. In Findings of the Association for Computational Linguistics: ACL. 2023, pp. 4005-4019, 2023.
Gilad Deutch, Nadav Magar, Tomer Natan, and Guy Dar. In-context learning and gradient descent revisited. In Proceedings of the 2024 Conference of the North American Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies (Volume I: Long Papers),
pp. 1017-1028, 2024.
Shivam Garg, Dimitris Tsipras. Percy S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Seungwook Han, Jinyeop Song, Jeff Gore, and Pulkit Agrawal. Emergence and effectiveness of task vectors in in-context learning: An encoder decoder perspective. In Forty-second International Conference on Machine Learning, 2025.
Roee Hendel, Mor Geva, and Amir Globerson.
In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 931 9318-9333, 2023.
Alberto Hojel, Yutong Bai, Trevor Darrell, Amir Globerson, and Amir Bar. Finding visual task vectors. In European Conference on Computer Vision, pp. 257-273. Springer, 2024.
Brandon Huang, Chancharik Mitra, Leonid Karlinsky, Assaf Arbelle, Trevor Darrell, and Roei Herzig. Multimodal task vectors enable many-shot multimodal in-context learning. Advances in Neural Information Processing Systems, 37:22124-22153 2024.
Joonseong Kang, Soojeong Lee, Subeen Park, Sumin Park, Taero Kim, Jihee Kim, Ryunyi Lee, and Kyungwoo Song. Adaptive task vectors for large language models. arXiv preprint arXiv:2506.03426, 2025.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding length generalization in transformers. Advances in Neural Information Processing Systems, 36.24892-24928 2023.
Dongfang Li, Xinshuo Hu, Zetian Sun, Baotian Hu, Min Zhang, et al. In-context learning state vector with inner and momentum optimization. Advances in Neural Information Processing Systems, 37: 7797-7820, 2024.
11

---

Published as a conference paper at ICLR 2026
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 International Conference on Machine Learning. pp. 32287-32307. PMLR. 2024.
Grace Luo, Trevor Darrell, and Amir Bar. Vision-language models create cross-modal task repre- sentations. In serto- International Conference Machine Learning, 2025.
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent prov- ably the optimal in-context learner with one layer of linear self-attention. In The Twelfth Interna- tional Conference on Learning Representations, 2024. URL https: fttt:s:://ceereevieee n Eorum?id=Sp3faue
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vee- style vector arithmetic. In Proceedings of the 2024 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 5030-5047, 2024.
Yingzhe Peng, Xinting Hu, Jiawei Peng. Xin Geng, Xu Yang. et al. Live: Learnable in-context vector for visual question answering. Advances in Neural Information Processing Systems, 37: 9773-9800, 2024.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Position: Do pretrained transformers learn in-context by gradient descent? In Proceedings of the 41 International Conference on Machine Learning, pp. 44712-44740. PMLR, 2024.
Pavel Tikhonov, Ivan Oseledets, and Elena Tutubalina. One task vector not enough: large-scale study for in-context learning. arXiv preprint arXiv:2505.23391 2025.
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. In The Twelfth International Conference on Learning http:///oummmmmm///ee Representations, 2024. URL htt.htts:... rettppse/ http://wwww....t.eeooooooro////a
Johannes Von Oswald, Fyvind 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, 2023a.
Johannes Von Oswald, Maximilian Schlegel, Alexander Meulemans, Seijin Kohayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05856 2023b.
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 9840-9855, 2023.
Kevin Christian Wibisono and Yixin Wang. On the role of unstructured training data in transformers' in-context learning capabilities. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference on Learning Representations, 2024. URL https::/ openreview net fest/fooum?11 www.ummdddsShSSeepppe
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, 2022. URL /ttttp::///peeneeieewee
net. wwwe/eeouuurtte.e forum? d=RJUFF
Yue Xing, Xiaofeng Lin, Chenheng Xu. Namjoon Suh, Qifan Song, and Guang Cheng. Theoret- ical understanding of in-context learning in shallow transformers with unstructured data. arXiv preprint arXiv:2402.0077 2024.
12

---

Published as a conference paper at ICLR 2026
Zheyang Xiong, Ziyang Cai, John Cooper, Albert Ge, Vasilis Papageorgiou, Zack Sifakis, Angeliki Giannou, Ziqian Lin, Liu Yang, Saurabh Agarwal, et al. Everything everywhere all at once: Lims can in-context learn multiple tasks in superposition. arXiv preprint arXi::2410.05503, 2024.
Liu Yang, Ziqian Lin, Kangwook Lee, Dimitris Papailiopoulos, and Robert Nowak. Task vectors in in-context learning: Emergence, formation, and benefit. arXiv preprini arXiv:2501.09240, 2025,
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49): 2024.
Chunsheng Zuo, Pavel Guerzhoy, and Michael Guerzhoy. Position information emerges in causal transformers without positional encodings via similarity of nearby embeddings. In Proceedings of the 31st International Conference on Computational Linguistics, pp. 9418 8-940, 2025.
13

---

Published as a conference paper at ICLR 2026
A
ADDITIONAL DISCUSSIONS
A.1
SUMMARY OF MATHEMATICAL NOTATIONS
Table 3: Summary of key mathematical notations used throughout the paper.
Notation
Description
n
Number of demonstrations in the input prompt
L
Number of transformer layers
d dp AttnviQ
Dimension of covariate and response embeddings
Prompt length (depends on demonstration structure)
Linear-attention layer with parameter V, Q
TF
Linear-attention model by stacking Aincar-attention layers
2; E IRd
Covariate (input) of the i-th demonstration Response (output) of the i-th demonstration
Vi Rd
X, Y E Rdxa
Matrices of covariates and responses for demonstrations
Ftest Vtest
Query covariate and ground j-th regression coefficient vector
response
W IRd
W Radd
Coefficient matrix, W [W],-.. wa]
Zo E Rdxxd,
Input prompt embeddings before the transformer Hidden states after the (-th layer Positional encoding matrix Value and key-query matrices of the I-th attention layer Block components of Vi, Qi in layer /
E 2
P Rdodd,
VI, QI
AI, BI, Ci, Dr
AE
Sub-block matrices of D, used in critical point analysis
C
In-context learning loss (ICL risk)
M(M)
Set of masked matrices with binary mask M Structured sets of matrices defining critical points
S1,S5, Sp
Ztv
Task vector extracted from an arrow (-+ token Weight vector for task vector formation
B E R
A.2 ADDITIONAL RELATED WORKS
In-Context Learning in Attention-based LLMs. The ability of LLMs to learn from examples provided in the input prompt, without updating parameters, has attracted wide attention since the discovery of ICL in GPT-3 (Brown et al.) 2020). A growing body of theoretical work has sought to explain this phenomenon. Early analyses show that transformer attention layers implement
gradient descent-like algorithms over linear regression objectives (Garg et al. 2022 Akyürek et al.) 2023 Von Oswald et al.) 2023a Ahn al.
2023 Wu et al. 2024), while others investigate sam- plc complexity and generalization behavior Xie et al. [2022, Chan et al.) 2022, Shen et al.) 2024] Von Oswald et al. 2023b Deutch al. 2024). These works collectively suggest that ICL is closely lied to the inductive biases of the attention mechanism, but do not fully explain how higher-level abstractions of tasks are formed encoded in LLMss
The Task Vector Method in ICL. Task vectors have recently been proposed as an abstraction of ICL demonstrations into compact hidden-state representations. [Hendll al. (2023) introduced task vectors as hidden states extracted from the last arrow token in triplet prompis, enabling zero-shot transfer by injecting them into new contexts. Concurrent works developed similar notions, such as function vectors (Todd et al. 2024) and in-context vectors (Liu et al.) 2024). These studies show that task vectors accelerate inference and sometimes match the effectiveness of ICL with fewer demonstrations. However, they remain largely empirical, without clear theoretical explanation of how or why such vectors encode task information.
Subsequent research has expanded the scope and utility of task vectors, Yang et al. (2025) demon- strates that task vectors naturally emerge even in small transformers trained from scratch with syn- thetic data, suggesting that their formation is an inherent property of attention-base architectures.
14

---

Published as a conference paper at ICLR 2026
Table Top 20 tokens with the highest output probability by decoding the task vector, results from Hendel et al. 2023). We underline the tokens in the output space of the current task.
Model
Task
Tokens
Prev Letter
b, Vr gr år name, 1, ro, j, d, tr A, ai, com, my ust, test, active,
French to English
other, name, the, true, is, social, 8, active, time, car, type, money, F, force, a, public, heart, one, ms, life
Present to Gerund
getting, storing, working, moving, playing, doing, making, driving, shooting, picking, being, sending, putting, selling, watching, changing, taking, collecting, feeding, reading
GPT-J 6B
Country to Capital
London, Paris, New, West, Berlin, South, Tokyo, San, Chicago, City, Moscow, Jerusalem, Amsterdam, Philadelphia, East, Madrid, Vienna, Beijing, Mexico, Germany
ei al. 2024) shows that aggregating hidden states across layers and multiple arrow tokens leads to stronger task representations, KKang et al. (2025) proposes to generate task vectors conditioned on each input query. Beyond text, task vectors have also been applied in vision (Hojel et al. 2024 Peng af., (2024) and multimodal models (Huang et al. 2024 Luo et al., 2025), where they enable flex- ible transfer across modalities. Han et al. (225) connects the performance of task vectors by task decodability, defined by the similarity between task vectors from different ICL tasks. These works highlight the empirical utility of task vectors but stop short of explaining their inner mechanisms.
Explaining the Task Vector Method. Task vectors were initially conjectured to encapsulate the complete knowledge of the current task (Hendel et al. 2023). However, this view fails to account for their inconsistent performance across tasks of varying complexity. Empirical observations fur- ther suggest that directly decoding task vectors typically yields tokens from the task output space (Todd et al., 2024), rather than explicit task descriptions (Merullo et al.] 2024). Concurrent work by Bu et al. 2025 analyzes the learning dynamics of 1-layer transformers with ICL-style prompts, explaining the utility of task vectors through a word2vec-like scheme (i.e., the existence of a vector 2, for task t such that y R 24 + x for all input-output pairs (I, 2)). While insightful, this characteri- zation is restricted to additive translation tasks, single-token prompts, and single-layer architectures, limiting its generality. By contrast, our analysis encompasses richer prompt structures, including pairwise and triplet formats that better reflect practical ICI. settings. Moreover, our critical point characterization extends beyond 1-layer models, and our linear regression formulation captures a broader spectrum of ICL tasks. Complementing shows that standard task vectors lack sufficient expressiveness for complex ICL tasks, reinforcing conclusion that task vectors are fundamentally constrained by rank-one mappings.
findings, Tikhonov et al. (2025) independently
A.3 JUSTIFICATION OF THE
ASSUMPTION
In our main analysis, we impose an assumption on the trainable parameters of linear-atteniin lay- ers, such that the V/ and QI matrices are block-diagonal in eq. (6). This block-diagon formulation is a widely adopted assumption in theoretical studies of ICL for transformer models, as it facilitates tractable analysis (Ahn et al. 2023 Mahankali et al. 2024 Wu et al. [2024] Zhang et al. 2024). Prior work by Ahn et al. (2023) demonstrates that the global minimizer of single-layer linear-attention transformers indeed exhibits such a block-diagonal structure. Although finding exact solutions for multi-layer transformers is more involved, it is reasonable to conjecture that similar structural pat- terns hold. Empirically, we observe that when optimizing the full matrices, gradient- based training also tends converge to biock-diagonal solutions.
15

---

Published as a conference paper at ICLR 2026
Intuitively, given the high dimensionality of hidden states in modern LLMs, it is plausible to assume that the 2; and yi components can be projected into orthogonal or nearly orthogonal subspaces when mixed in the hidden state space. This motivates decomposition of the projection matrices Vi and Qi into two separate parts that operate independently on I; and Yis which can be equivalently formulated the block-diagonal structures.
A.4 INSEPARABLE COVARIATES AND RESPONSES
In our main analysis, we assume that and embeddings are linearly separable, allowing the addition x; + y; to act concatenation operation. However, recognizing that this assumption does not generally hold for real-world transformers, we extend our analysis to the following setting, where I; and y/ are no longer linearly separable. While this still imposes a 2d-dimensional requirement on the hidden space, such constraint is easily satisfied in practical transformers, given the high dimensionality of their internal representations.
Zo
0
0
E R (2d)xx
(11)
y1
In
Yn
Ftest
0
We slightly modify the sparsity constraints for the first layer, and require (Do)2i: = 0 for E [n++1]:
0 Odxd
Ao
Qo = 02axxa 2d
0 [.]- where Ao E Idd Do E R4
V/=
(12)
0
Doj
With these conditions, we are ready to establish the critical points for inseparable demonstrations. Note that Vo and Qo do not involve Bo and Co. the sequences B and C have size 1-1.
Theorem 7. Under the same settings as Theorem 1 define S1 C Rald and Sp C Rdpxdp
S = (Ma ERR) === 1
Sp = {diag(In & A2) A1, E 2x22].
Consider optimizing an L-layer linear transformer with inseparable pairwise demonstrations and parameter configuration given in eq. (12 for the first layer and eq. 66) the remaining layers, then
inf AES), BES)"1 cest - "DeS) EHEAURRUUD / Q. =0
This result suggests that for inseparable demonstrations, the first layer performs a functionally simi- lar concatenation operation by "moving" the embedding of each z; to the corresponding ys position. This enables the model to reconstruct the single-token structure without linear separability.
A.5 LAST TASK VECTOR WEIGHTS THE MOST
While our analysis of linear-attention models suggests that each formed task vector (i.e., the hidden state each arrow token) contributes equally to the final prediction, this assumption does fully hold in practical LLMs. As demonstrated by the conflicting tasks experiment in (Hendel et al, 2023), injecting task vector from task B into an ICL prompt designed for task causes the model to predominantly perform task B. This behavior indicates that LLMs largely rely on the last arrow token to determine the task identity. We attribute this to the causal attention mechanism used in practical LLMs. which is not captured by our current theoretical analysis. In causal attention, only the final arrow token can aggregate information from the entire preceding context, making the most informative and influential for prediction. This explains why our multi-vector strategy offers modest, though consistent, performance gains. The improvement suggests that intermediate arrow tokens do participate in the inference process, albeit less effectively. Enhancing how LLMs utilize information from all arrow tokens remains a promising direction for improving task vector accuracy and robustness.
B EXPERIMENT DETAILS AND ADDITIONAL. RESULTS
In this section, we present experiment details and additional results not included in the main text due to space limitations. Our experiments are conducted on an A100 40G GPU. It takes around 30 GPU hours to fully reproduce results.
16

---

Published as a conference paper at ICLR 2026
B.1 SYNTHETIC EXPERIMENTS ON
MODELS
We consider training linear- finear-attention models on random linear regression instances. We take embed- ding dimension d = 4, and the distributions for generating x; and wy both P. P.. = N0, Ia)- We optimize the ICL risk for L-layer finear-attention models with n in-context demonstrations using AdamW, where [3] and
[5, 30). Each gradient step is computed from batch size of 1000. We additionally apply & regularization to simplify the found solutions. For training efficiency and stability, we restrict the A,, Bi, and C matrices to Sy during training, and initialize Dj Rd,, with i.i.d. Gaussian matrices. For each case, we train 40 models with different random seeds, and report the minimum achieved ICL risk to approximate the global minimum.
To reproduce the task vector mechanism, we focus on models trained with rriplet-formater prompts. The training procedure identical to the above. For inference, we restrict P.. to rank-one coefficient matrices, by letting W wi wy where W1, W2 (0, Ia). We first generate normal ICL prompts to generate task vectors as the hidden states of the last arrow token after the first attention layer, and then inject them into zero-shot prompts after normalization. The final outputs Iest taken as the output of these injected zero-shot prompts after being processed with the same transformer model. We compute the final risk as E
Steat
to simulate the layer normalization blocks
|B
in practical LLMss. The reported scores are averaged for |5, (5)
B.2 EXPERIMENTS ON PRACTICAL LLMs
Datasets. Following the settings of the original task vector method (Hendel et al. 2023), our study covers 33 tasks in 5 categories. The detailed description for each task is provided in Table 5
Prompt Template. The template used to construct ICI. demonstrations is "Fxample: - (..) where x; and y are subsequently replaced by the input and output of the semantic mapping. For the query part, yi is omitted from the prompt. After concatenating each demonstration with "\n", " an example of the full input prompt is:
Example:(a1) - {1i) (x..plexaalle - / |nExample: (Feest) -
(13)
Evaluation. To evaluate the N-shot performance, we generate 50 x (N + 1) i.i.d. prompts for each task with number of demonstrations n 10 for task vector extraction. The hidden states of the last + token, which is also literally the last token in the prompt, are recorded for every layer in the transformer. Thereafter, we generate another 50 i.i.d. prompts with N demonstrations, where Trest is selected to be distinct from the previous chosen ones. The final accuracy is measured by whether the next word predicted matches the expected answer. The performance
the standard ICL method (Baseline) is acquired by inferring without interference. For the task vector method (TaskV) and our multi-vector variant (Task V-M), the extracted task vectors are injected to replace the hidden states of the arrow - tokens for T-s each of the N + 1 arrow tokens is injected with the N + extracted task vectors for the same task. The performance is reported for the layer / L achieving the highest accuracy. For each case, the mean and standard deviation are evaluated through 5 independent trials.
specified layer 1. For TaskV, only the last arrow token is injected, while
Additional Results. Besides Llama- 13B. we also observe consistent accuracy improvement of our Tas method on the Pythia-120 12B model, as reported in Table
While the performance gains of TaskV-M over TaskY are not dramatic across all ICL tasks, the goal of T-s is not to surpass tfae-oof-he--as ICL techniques but to demonstrate that the task vector framework can be systematically extended by injecting multiple vectors simultaneously. This is especially valuable for complex tasks that inherently require higher-rank representations. Our results on bijection tasks clearly validate this motivation: TaskV-M yields notable improvements over the standard T method. For other simpler tasks, the marginal gains from TaskV-M suggest that the expressiveness of W may not be the primary performance bottleneck. We believe these insights facilitate the design of future ICI. and task vector methods.
B.3 ANOTHER VECTOR INJECTION VARIANT
In our main experiments, we implement TaskV-M by extracting N + 1 task vectors from the same number of different prompts. Another possible implementation for TaskV-M is to extract multiple
17

---

Published as a conference paper at ICLR 2026
Table 5: Descriptions of the tasks used in empirical studies.
Category
Task
Example
Description
Contry to Capital
France + Paris
Output the capital city the given country.
Person to Language
Macron + French
Output the native language the given person.
Location Continent
Paris + Europe
Output the corresponding continent the given location.
Knowledge
Religion
Saladin + Muslim
Oulput the associated religion of the given lo- cation person.
List First
[a.b.c] +
Output the first item Output the last item in the given list. Output the next letter of the given letter in the alphabet.
given list.
List Last
[a.b.c] +
Next Letter
+
Prev Letter
b -
Ouiput the previous letter of the given letter in the alphabet.
Algorithmic
To Upper
Output the corresponding uppercase letter of the given lowercase letter. Output the corresponding lowercase letter of given uppercase letter.
To Lower
A
English to French
hello + bonjour
Translate the given word in English to French. Translate the given word in English to Italian. Translate the given word in English Spanish. Translate given
English to Italian
hello ciao
English to Spanish
hello hola
Translation
French to English
bonjour + hello
in French English. Translate given word in Italian to English. Translate given word in Spanish English.
Italian English
ciao + hello
Spanish to English
hola + hello
Present Gerund
go going
Output the corresponding gerund form of the given verb present simple tense.
Present Past
go + went
Output the corresponding past simple form given verb present simple tense.
Present Past Perfect
go + gone
Output the corresponding past perfect form the given verb in present simple tense.
Gerund to Present
going + go
Output the corresponding present simple form of the given verb in gerund form.
Linguistic
Past to Present
went + go
Outpul the corresponding present simple form of the given verb past simple tense.
Past Perfect Present
gone go
Output the corresponding present simple form of the given verb past perfect tense.
Singular to Plural
dog + dogs
Output corresponding plural form of the given noun in singular form.
Plural to Singular
dogs + dog
Output the corresponding singular form of the given noun in plural form.
Antonym
happy + sad
Output the antonym of the given adjective.
To Upper & Lower
Output the given letter in uppercase if lowercase, and vice versa.
English & French
hello ++ bonjour
Translate the given word to French if English, and vice versa.
English & Italian
hello ++ ciao
Translate the given word to Italian it is in En- glish, and vice versa.
English & Spanish
hello ++ hola
Translate the given word to Spanish if it is in English, and vice versa.
Bijection
Present & Gerund
go going
Output the given verb in gerund form if it in present simple tense, and vice versa. Output the given verb past simple form is present simple tense, and vice versa.
Present Past
go went
Present Past Perfect
go ++ gone
Output the given verb in past perfect form it is in present simple tense, and vice versa.
Singular & Plural
dog ++ dogs
Output the given noun in plural form if it in singular form, and vice versa.
18

---

Published as a conference paper at ICLR 2026
Table 6: Accuracy comparison between standard ICL (Baseline), the task vector method (Task and our strategy (TaskV-M). The experiment is conducted on Pythia- 12B with = 10.
Method
Knowledge
Algorithmic
Translation
Linguistic
Bijection
Average
Baseline
6,60 1.59
14.07 1.45
8.60 0.68
12.53 + 1.57
10.31 0.70
10.82 0.48
0-shot
TaskV
63.30
84.73 1.22
62.07
82.58 1.22
42.27
66.40
Baseline
61.80 # 5.45
72.80 1.15
43,27
2.92
57.07 1.15
41.91
2.83
53.95
1.02
1-shot
TaskV
76.40 # 2.40
84.20 1.05
71.47
= 1.41
87.16
2.04
53.11 2.37 53.87
73.59 0.79
TaskV-M
77.70 +: 2.52
83.73 1.37
71.00
1.48
86.80
1.59
2.90
73.68
0.90
Baseline
70.30 :: 3.71
82.13
0.54
60.80 +: 1.81
81.16 +: 1.57
50.76
= 2.17
68.41
0.64
2-shot
TaskV
80.30 à: 2.46
87.00
1.63
76.13
&: 3.77
89.33 4. 0.70
58.67
A: 2.44
77.41 0.50
TaskV-M
81.60 = 1.56
86.47 0.40
77.27 + 2.53
89.51 + 0.88
59.24
2.48
0.76 77.87
Baseline
77.60 + 2.40
81.87 0.81
68.13 2.02
86.31 1.93
55.73 1.60
73.20 0.31
3-shot
TaskV
84.00 4: 2,76
86.33 1.17
79.53 2.27
92.00 + 0.67
58.76 = 1.53
79.06 0.67
TaskV-M
85,40 #: 2.31
87.07 1.18
78.13 + 1.86
92.84 + 0.68
59.56 = 1.27
79.54 0.35
Baseline
78.40
82.73 0.44
72.40
88.89 1.25
57.91
75.46
4-shot
TaskV
83.80 = 1.12
87.60 1.81
80.20
92.18 0.96
59.38 0.47
79.59 0.62
T-skV-M
84.30 = 1.50
88.13 0.81
2.67
91.87 1.25
60.31 0.86
79.87 0.51
Table 7: Accuracy comparison between few-shot ICL (Baseline), the task vector method (TaskV), the multi-vector method (TaskV-M), and the single-prompt variant V-MS). The experiment conducted on Llama-13B with n = 10.
Method
Knowledge
Algorithmic
Translation
Linguistic
Bijection
Average
Baseline
6.90 2.08
15.60 1.72
7.00 = 1.65
12.44 + 1.74
8.27 = 1.33
10.28 0.98 72.26 1.01
O-shot
TaskV
68.80 +: 2.66
86.20 1.61
73,53 1: 0.91
85.24 + 1.80
50.67 2.32
Baseline
69.50 + 3.86
73.67 + 1.56
57.80 + 2.01
56.22 + 1.57
44.76 = 2.4
58. 11 + 0.63
1-shot
TaskV
79.50 + 2.35
88.47
0.75
80.67 + 2.56
89.11 + 0.84
60.44 2.07
78.79 + 0.77
TaskV-M
81.30 + 2.80
89.53
0.65
80.13 +: 2.14
88.71
+ 0.62
61.78 0.96
79.34 0.37
TaskV-MS
80,90 + 3.10
88.40 + 0.93
80.13 +: 2.54
88.89 +: 0.73
61.11 1.31
78.96 + 0.43
Baseline
78.80 +
85.07 1.37
75.67 + 2.64
76.80 +
56.49 2.87
72.92 +:
2-shot
TaskV
84.60 +
88.40 0.68
84.33 + 0.92
90.13 + 0.92
62.44 2.16
80.82 0.42
TaskV-M
85.70 +
89.27 1.10
84.13 + 1.15
89.64 + 0.86
64.49 2.02
0.37
TaskV-MS
84.40 + 2.13
0.98
84.67 + L.73
90.18 + 1.39
64.49 2.30
81.61 0.80
Baseline
86.20 + 2.69
88.07 1.06
80.00 +: 1.67
84.04 + 1.19
62.18 1.52
78.51 0.42
3-shot
TaskV
90.20 + 2:23
88.67
86.27 +: 2.31
92.31 + 0.48
66.53 0.94
83.53 0.41
TaskV-M
90.30 +
89.87
86.07 + 2.17
92.36 0.72
68.13
84.15 0.52
TaskV-MS
90.60 +
89.47
86.20 + 1.89
91.91 0.87
67.69
83,91 0.45
Baseline
84.80 1: 2.06
88.07 0.61
83.27 +: 1.82
88.89 +: 1.91
67.16
1.47
81.52 4: 0.66
4-shot
TaskV
88. 70 +: 1.69
89.53 + 1.34
86.27 + 1.08
92.76 + 0.54
70.44 1.35
84.66 0.39
TaskV-M
89.60 + 1.43
91.00 1.01
87.20 + 0.62
92.36
+ 1.44
72.53
0.94
85.64 + 0.29
TaskV-MS
90.10 + 1.39
90.67 1.10
87.00 + 1.17
92.22 + 0.92
72.09 1.46
85,45 + 0.26
task vectors from each arrow token single few-shot prompt simultancously. We name this alter- native approach as T-s. As discussed in Proposition 3 the task vector weights that emerge at each arrow token are approximately orthonormal, suggesting they encode distinct information sub- sets and can be simultancously injected to enhance model performance (c.g., by increasing the rank of the induced coefficient matrix W'. Table shows a comparison between the current multi-vector method (-askV-M) and this single-prompt variant (TaskV-MSS
While TaskV-MS also delivers strong performance, slightly underperforms TaskV-M We believe this is due the causal attention mechanism in real LLMs, where earlier arrow tokens only aggregate information from a subset of demonstrations. Nonetheless, TaskV-MS a promising alternative for accelerating inference.
19

---

Published as a conference paper at ICLR 2026
Table 8: Comparison of the accuracies of n-shot ICL and task vector on bijection tasks (n 10).
We use gray text to indicate accuracies lower than 60%.
GPT-J
Pythia-6.91
Pythia- 12B
Llama-7B 7B
Llama- 13B
Qwen3-8B
Llama3-8B
Task
ICL
TV
ICL
TV
ICL
TV
ICL
TV
ICL
TV
ICL
TV
ICL
TV
Lower Upper
1.00
0.08
0.90
0.28
0.96
0.24
1.00
0.55
1.00
0.58
1.00
0.56
1.00
0.38
English French
0.64
0.50
0.38
0.28
0.52
0.28
0.54
0.35
0.64
0.32
0.84
0.48
0.66
0.42
English Itaiian
0.68
0.56
0.62
0.48
0.60
0.56
0.70
0.47
0.72
0.44
0.68
0.36
0.70
0.36
English <> Spanish
0.70
0.52
0.62
0.56
0.66
0.56
0.64
0.43
0.84
0.56
0.70
0.32
0.72
0.32
Present Gerund
0.64
0.36
0.44
0.32
0.40
0.22
0.80
0.41
0.74
0.26
0.72
0.34
0.94
0.52
Present + Past
0.60
0.38
0.48
0.36
0.54
0.16
0.52
0.33
0.68
0.44
0.78
0.42
0.90
0.58
Present ++ Perfect
0.46
0.14
0.38
0.24
0.46
0.28
0.55
0.33
0.54
0.42
0.66
0.42
0.78
0.50
Singular Plural
0.66
0.56
0.44
0.28
0.76
0.80
0.52
0.84
0.88
0.58
Antonym
0.86
0.78
0.76
0.66
0.76
0.70
0.83
0.73
0.78
0.72
0.82
0.74
0.82
0.76
Table 9: Comparison of the accuracies of n-shot ICL and task vector on bijection tasks (n = 20).
We use gray text to indicate accuracies lower than 60%.
GPT-J
Pythia-6.9B
Pythia- 12B
Llama-7B
Llama-13B
Qwen3-8B
Llama3-8B
Task
ICI.
TV
ICL
TV
ICI.
TV
ICI.
TV
ICI.
TV
ICT.
TV
ICI.
TV
Lower > Upper
1.00
0.12
1.00
0.32
0.94
0.38
1.00
0.48
1.00
0.60
1.00
0.58
1.00
0.36
English < French
0.74
0.54
0.44
0.40
0.52
0.40
0.52
0.34
0.58
0.34
0.58
0.30
0.74
0.28
English Italian
0.62
0.54
0.66
0.46
0.68
0.48
0.78
0.50
0.74
0.48
0.76
0.38
0.76
0.32
English Spanish
0.80
0.58
0.54
0.38
0.56
0.40
0.78
0.58
0.84
0.58
0.66
0.32
0.86
0.40
Present Gerund
0.54
0.26
0.54
0.22
0.46
0.14
0.84
0.44
0.94
0.38
0.88
0.28
0.98
0.52
Present <) Pust
0.66
0.26
0.54
0.30
0.58
0.28
0.72
0.30
0.76
0.44
0.74
0.40
1.00
0.48
Present 4+ Perfect
0.42
0.18
0.44
0.20
0.46
0.24
0.48
0.30
0.52
0.48
0.80
0.44
0.90
0.48
Singular 6 Plural
0.64
0.40
0.62
0.36
0.52
0.28
0.52
0.94
0.42
0.86
0.60
0.92
Antonym
0.84
0.76
0.84
0.70
0.90
0.82
0.90
0.84
0.90
0.84
0.84
0.74
0.84
0.76
B.4 FURTHER RESULTS ON BUECTION TASKS
Here, we extend the results from Table D that illustrate the failure of task vectors on bijection tasks across broader range of LLMs and varying numbers of imput demonstrations. We keep the same experimental settings as Table while increasing the number of demonstrations to n E {10, 20), and report the results for 7 distinct LLMs: GPT-J, Pythia-6.9B, Pythia-12B, Llama-7B, Llama- 13B, Qwen3-8B and Llama3-8B. As shown in Tables [ and [9] the task vector method results in a significant performance drop compared to the standard ICL on bijection tasks. These results further support our claims that:
Task vectors systematically fail on bijection tasks, even when further increasing the number of demonstrations in the prompt.
The failure is consistent across multiple model architectures, validating that the issue stems from fundamental expressiveness limitation rather than model-specific artifacts.
B.5 FULL SALIENCY ANALYSIS RESULTS
In the main text, we reported simplified version of the saliency map due to space limitations, focusing only on the demonstration tokens z,, -, y/s In Figure 5 we report the full saliency map covering every token in the prompt. Here, "B" stands for the [BOS] token, and "E" stands for the word "Example". Please refer to eq. (13) for further details about the structure of the input prompt. As can be seen, the highlighted saliency weights exhibit clear patterns of embedding concatenation and weighted summation. It can also be observed that latter demonstrations weigh more for task vector formation (i.e., saliency magnitudes for latter ys tokens are larger in Figure 5b).
20

---

Published as a conference paper at ICLR 2026
(a) Full Saliency Map 10)
(b) Full Saliency Map 12)
Figure 5: Visualization of full saliency matrices bipartite graphs between layer (+) and 7 (o). edge widths indicate saliency magnitude (IJama-7B, n = 10). (a) Each token attends to its corresponding (r,, y4) pair, reflecting embedding concatenation. (b) The final (-) token attends broadly to all tokens, indicating task vector formation.
C AUXILIARY LEMMAS
Lemma 8 (Proposed in (Ah et al. 2023 ) Given positive objective function f(A) taking parame-
ters {Ai)=> where A; Ka xd. Let S1 C 11,_, be a predefined parameter subspace. Define A(t, R2) = {A1,"
A, + tR ... Am, given i E [1, R4 E Rdd xd; and t e RR. for any A S and R; E Rdixd, there exists R; S; such that
a f (A(tt.
< 1==0
di ((tt,R))
d
dt
1==
then we have
inf AES
[A.((A))E 0.
Proof. This lemma is proved part of the main theorems in (Ahn et al. 2023). We rearrange the proof here to accommodate arbitrary function of matrices. Firstly, notice that for any R ( m Rd.xd,
Σ di (A(t..R))
di f(A f(A+ tR)
Therefore, the provided precondition is equivalent to stating that for any A E and R II m2_ Rdixd, there exists R E S such that:
d d dt (A+ tR)
<
f(A J tR)
t= dt
1=
Let R = VA)(A), we then have
d
d(A ((A-tVA -VVA((A
f(A + tR)
dt
=
d(A- d(AA)V)f(A))
t
1=0
21

---

Published as a conference paper at ICLR 2026
(Vaf(A), -Vaf(A) - (V
the infimum of V VAJ(A)) is not zero but some positive value p. then the S-constrained gradient flow induced by R will lead unbounded descent:
dt (A+
d
-P.
1=0
This contradicts the fact that f(A) 2 and concludes the proof.
The following lemma is extension of Lemma 5 in (Ahn al. 2023 by accommodating mul- tivariate y samples as well as enabling a wider range of demonstration and transformer parameter configurations.
Lemma 9. Let I,,-.- En++1 be i.i.d. samples from an input distribution, and let W be sampled independently of {x,)==] Let Zo E (22)x where N E, be constructed of form
Zo
E IR (2djNN
0g
where the parts can be arbitrarily constructed from f int1 and W. Let 7o be defined as replac- ing the zero part of Zo by vn+1.
7
(2d)xx
yn+1
Let Zi be the output of the I-th layer of the linear transformer, and let X1, Yi E RaxN be the first and last d rows of Zú respectively. Suppose that the {Qil=z matrices of form
((2d+d,)
Q:
d colamons
dy coluans
Then the in-context risk of this L-layer linear transformer is equivalent to YnN C(VV, Q.) = =Ez,,w MoNN -))]
(14)
Proof. Let the V/2 and Qi matrices be represented as:
VI
Vi =
Q= [o] Q7]
where V/7 V/2 Rdd2d Ql (2dddd,)xd Qf R ((dd+d))xd xdp, Then the update rule in eq. 5 can be rephrased as
Xi X M P] QPP))
1
2 MZ X-1+ QP.).
Yi = Y
M
Let Az 2.- Zu. i.e. an all-zero matrix except that the last half of the last column is In+1- Let Ax and Ay be its first and last d rows respectively, then Ax 0 and Ay [0 0 yn+1].
Note that 2, 2, - Az holds for / = 0 trivially. Now suppose it holds for some = k - 1, then
N = Xx-1
1
Vi 2 M TT p|
(Q/
n 1
- +
Vi. - M 2 (Q) !?P)
n
+ Vi A 2 PT !?P)
(Q
+
V/l 2 A 0d, (Q. Q?P))
n
22

---

Published as a conference paper at ICLR 2026
1 V Az M 2 xdp xd, 1 x Q2P)
= XR V/ 2 M 2 P] ( P XR.
where the last step holds by noticing that 4 MM Similarly, one can prove that
Y - 122 M Z (Q) Q.P)) Yx + Ay.
1
Therefore, holds that for any E [1. L}, 2 = Zz + Az. Recall the in-context risk in eq. 2
Q()=)) = Ezo.w [Z1) + Un+ill2
= =Ezo,W
11 + Ay)(IN - M212
= EE W [tr((Tn - VL -M))
The proof is complete.
D PROOF OF THEORETICAL RESULTS
D.1 PROOF OF PROPOSITION 4
Proof. We will first prove sufficiency. Let W abT be rank-one matrix, where a. E Rd. The given conditions imply that = Wy = WWr ab ab x, we then have b x b ab ab x (b (bTa)²b x. Since 6%a x * we can conclude that 5 #1. Then, x = ab ab x = tab T x ty.
To prove the necessity, it suffices show that selecting W xx. /|a0|12 when satisfies the given conditions (alternatively, select W =-xr //|x when x = -
D.2 PROOF OF THEOREM
Proof. To enhance the readability of the notations in this proof, we will drop the constant 11 factor in linear attention. Furthermore, we will simplify Zo. Xo and Yo in Lemma 9 as Zo. Xo and Yo respectively. This results in different definitions compared to the original ones, but we will not refer to the original definitions in the remainder of this proof.
Zo
X Yo
x 0
0
In
Ftest
E R (2d))x(2222))
yn
O
Steest
Let Zi be the output of the 1-th layer of the transformer, and let X1, Yi pdx(2n+2) denote the first and last d rows of Zg, respectively. Under the constraint in eq. 6 we can verify that
Xi XI-1 + A = V-1++ ( Di).
Di),
(15)
In the following analysis, we will use f(A + B) to denote the result of the function f of A when
replacing the value of A with B. Additionally, denote S(A B A) ʃ(A = B) for any operator *. Therefore, [(A & B) f(A + A + B). We also denote s(A x B) = S(A BA) and f(A - B) ((A AB) for convenience.
Our goal is proving that, for any E E AUBUC D and an arbitrary matrix R E Rdxd (Rd, for D), there exists R E S2 (Sx for C, Sp for D) such that
d C(E & tR)
d C(E - tR)
(16)
dt
dt
1=0
Let Xo [0, 21,... 0, Ftest] be function of Xo. then have Yn o.. Let U, Rdxd be uniformly sampled random orthonormal matrix, and let Us /1/21, 1-1/2 One can verify that
23

---

Published as a conference paper at ICLR 2026
US = -V/ 1/2, By applying Lemma 9 and the fact that Xo d Ux we have that for any
given matrix R.
d dt
C(E - IR)
1=0
d Exo.w w
( MYY (E E IR)YL & tR)(I (0([--)]]]
dt
Ir (I- (1-M MaV d
di y2(E $ 1n)
(I- [[-
2x.x W
1=0
= =2Exx.w.Ul
to (t v Us)
(( ~ Uz. F IR)
(1 (T-Mi)
dt
t=
Next, we will show that eq. (16) holds for each one A, B, C, D, for any E [].
1. Equation (16) holds for A..
We first show that for any / [1, L]. the following equations hold:
X1 A Uz)
Ug
(17)
d X( v U2, _ tR)
Us d Xi(A, $ Ltu' RU2)
(18)
dt
dt
t=
1==
It is straightforward to verify that eq.
holds for I O. Now suppose that eq.
holds for some
I = k - 1, we then have
Xxo v Uz)
=Xo-1(Xo x U2) + A e U2) (XL, E U A U2) DI
= -U-Xa-1+ AU M
Di)
= UUS(XA-1 + M
D.)) = Ug
where the third equality follows by noticing that when At at and G qE-l, have A,Us- = Un A₁ and U3 GUS = Cj- This concludes the proof of eq. 17).
We now turn to the proof of eq. (18) Notice that when 1 < i, naturally have
d
d
-
Xi - Us $ tR)
Us -
Xi $ tUE RUx)
=0.
dt
dt
1=0
1=0
When / = i, it is easy to verify that
xi A Uz. A; $ tR)
+R)
= MIG D:)
dt
t=0
=Us US RUS D.)
d
=Us dt XI(A, - &U. RUs)
Now suppose that eq. (88
holds for some = k- - 1 2 i, one can verify that:
d
Xx X Us. A, + (R)
J.
t=
d X 1Xo x Us A; + IR)
d
A x Us, A, & IR)M
dt
dt
1=0
(X (Xo x Uz. A - CR & Us A, - tR)
+ +Dx)
1=0
Xx- 1Xo x Us - tR)
dt
1=0
24

---

Published as a conference paper at ICLR 2026
+ A di X - Uy, A E (B)
d
M 1 & Ux))CXXx-1 1Xo - U2) + Dz)
1=0
+ (X, v U d
x x (Xo 2 Us, A; # 1R)
CxXx-1(X6 x Un)
dt
1=0
+ AX U (Xo A Uz)C d Xx-1 (Xu E Uy A, $ tR)
df
t=
d
=Us di X $ :Wg RUs)
d
+ Ur dt
Xk- & U0,57 RUx))
M CaX
1==
d X (A, $ +tUE!RUS RU²))
+U M
OKXE-I
dt
11=0
+ - MX Cs
(A, $ :U RUx)
de
1==
d
+U:: d
= =Us
Xx- 1 (A; = U155 RUS)
X x (A; - UUs R
dt
dt
1=0
( (A: $ :ī. RUs C+ XR- (A, # UUs RU²) + Dz)
d Us dt Xa # :W' )
1=0
This completes the proof of eq.
18
Under the condition that B; biSd for some bj E R. we can simplify eq. [5 as
Yi Yi-1 + I
Di) Yi-I +xx
1 Yo [
1=1
+ D2)). then it satisfies that Yi = WG,. We are ready to prove that similar results to eqs. 17 and (18) also hold for G₁, 1 E [1,
Define Gi - (T-
Gi
Us) = Ur
(19)
= = d Gi(A; $ UUs RUs)
GX x Us, A + tR)
(20)
dt
dt
1=0
t==
Notice that eq. 19 holds trivially for 1 as Go Xo. Now suppose that eq. 1= - 1, we then have
19
holds for some
GX Uz) Gr-1 - Uz) (T x U2) (No x Uz) D&)) Cx D&)) = Ur
=U>G+-1 + notice that both sides equal 0 when / S i. Now suppose
This concludes eq. (19). As for eq.
20
that eq. (20) holds for some = k
i, we then have:
d Gx x Un, A $ (R)
dt
r=
d
d
+ d
G x Uz A; * tR)
(x, Uz. A - tR)M
dt
dt
1=
(x[_(X, x Us x 1R)C # Us, A F tR) + D/)
t=
25

---

Published as a conference paper at ICLR 2026
d di Gn- & Us, A; & IR)
1=
+ +b
Gx v Uz, A, - IR)
M M (X λ U=) e Ux)) + D2)
dt
100
G A U d
XT (Xo A Us A $ tR)
CAXx-1(X0 - Uz)
dt
1=0
d +b / U2)MX 1 (Xo E Uz)Ck di X - Us $ t
d -Ux GR- - U055' RU*))
dt
1=0
d
+ +bUgs _ (A, $ W RUS))
M ( (XX-- + DD)
dt
1=0
+ M
XI di
x (A, $ [1,5'' RU²))
CrXx-1
1=0
d
-bUS:GG-MX Ck
Xk- (A: + RU's)
dt
t=0
d Ga き tUss RU2)
= - dt
(=0
This concludes the proof of eq.
20
Consider the in-context risk:
d di
C(A: # tR)
t=0
]
= 2 -u,
(T - (1-MMVIXX (Xo / Uz) - Vi v Uz. A & IR)
(T- M)
di
1=
(I (T-M)) MGG.US/TTW WUx d
W.x.W.
tr
CL(A, $ tt' RU²))
(I- - M)
dt
1=0
d (I (I-M) dt
2d Ex X
tr
Eu, GL(AA E tUss RUS)
(I - (1-
1=0
2d x.
(1 (--M) d
Gr(Ai & Ev, RU:)
(1- Mi ]
dl
1=
Ex.
(1- (I-M)GT,--1 d
Gr(A: & Irla)
(I (1- M) }]
dt
1=0
d Exo.w [tr(11 M)Y (A, # trla) trraall-- - M))]
dl
t=
d c(A: $ (rla)
db
t=
where r Eu [1
RUr] = 1 tr and we used the fact that -- UU- = --
and d GL(A + :R) dt
is affine in R. This concludes that eq. (16) holds for A, i E [].
t=
2. Equation 16 holds for B..
From the recursive expressions in eq. 15). we can conclude that the values of X; do not depend on B, Therefore, we naturally have
XB & IR) = X₁.
(21)
Next, we would like show that for any 1 [1, T.].
Ew T d Yi - tR) dt
= --1 d dt
Gi & ttr(B))
(22)
1=
L=
26

---

Published as a conference paper at ICLR 2026
When 1 < i, we can easily verify eq.
22 since both sides equal O. When / = i, we can get
W d Yi, + R
Ew
Ew [W' RV- MX Pi)]
dt
0
Ew WT RW C + Di)
Di)
= d Gi(b; + (r((R))
dt
1-1
Suppose that eq.
22 holds for some 1= 12 i. One can then verify
d (B, _ tR)
Ew
W T
dt
WT d
-Ew
Y = tR) (1 + 6LX 1=
=Ew
WT d - (Bi E tR)
(I+ -C - +Dx))
1=0
- d dt
(bb - tt(()))
(1 (1+ + +D))
1=0
= E--
d
G (b; + (t((R)
dt
The proof of eq.
22
complete. Now, look at the in-context risk, we have
C(B + IR)
= 2 E
(r (1-M)Y d
di YL & (R)
(-M)) )]
(1- M)
di
t=0
1=
(I-MGT Ew W d dt
Vv & (R)
(T - (T-
2Ex.
tr
1-
2Ex
tr
(.-M) GL $ ttr(R))
(I- M)
dt
AE(u-MM) Mrr dt
d Yi $ ttx(R))s)
2xx..
(1- (1-M)
d C(B $ (r((R)aa)
t==
This concludes that eq. 16 holds for B, i E [].
3. Equation 166) holds for C..
Similar to the A, case, we will first prove that for any E [1, L].
d Xi x Uz. O. = (R) 1=
Us d
X1(C, & tU RUs)
(23)
dt
di
1=
The equation above holds trivially for < i. For the case I = i, we have
d xi x Us. C. & (R)
dt
1=
A X US) X Ux) x Us:)
= -U RUxXt Us d Xi(C. - IU RUs))
dt
One can conclude the proof of eq. (33 through a similar reduction as eq. 18 for > layers. Next, establish the corresponding result for Gr:
d G X US. G2 #t)
d GiC. $ RUS
= Us
(24)
dt
1=0 dt
1=0
27

---

Published as a conference paper at ICLR 2026
This equation holds trivially for 1 < i. When taking 1 = i, verify that
(X - Ux) & RX) Uz) d Gi v Un C. - (R)
dt
1=
-LUs;,i-1X0o v U+) R
GiCC. _ U RUs)
Us dt
1=
For 1 >> i layers, one can follow similar reductions as eq. 20 to finish the proof. We then consider the in-context risk:
d [CC & (R)
dt
1=0
=2Exo.W.UU
tr
(T -T- M)- (Xo z Us) d v - Uz
Ur, C, = IR)
EiR)
(1-M)
(I M) M)
dt
1=0
d GrIC $ tR)
(I- M) ]
W.....
tr
(I (U-M WUx at
1=0
2d Exo tr
(1 (1-M)
Eu CLIC, # u RUS)
(1 (1-M)
)
dt
1=
Bx tr
(I (I- MMG
d (G,
(1- (T-M)
dt
d Ex.,w (oT((1 M)Y ( # #
tr (o)[-M))]]
dt
1=0
d f(C: $ 1r--1)
dl
t=0
where r = Eu, UU RUS) tr (/))) This concludes that eq. (6) holds for Ct.
4. Equation 16 holds for D,.
Let Up Fax= be uniformly sampled permutation matrix, i.e., a binary matrix that has exactly one 1 entry in each row and column with all other entries O. Let U. = diag(U, @ 12, I2) e 2(2*++)x(22++)) One can verify that by multiplying XoU it is equal to shuffling the first 2- column sub-blocks of Xo and keeping the last columns unchanged.
Then, consider matrix Ug = diag(fi...., En+1)) ((*+1)x(*+)) where iid Unif(#1), i.e., diagonal matrix with random +1 entries. Let U1 Ug T2 (2222))(2 Thus, U1 U. and XpU+ is randomly flipping the sign of each 2-column sub-block in Xo.
We are going to prove that for any 1 E [1, L.J, recalling that ʃ(A & B) 1(A AB),
Xi(Xo &
(25)
Gx - U+U,)= U.U
(26)
Equation 25 holds trivially for barl=0. 0. When eq.
25 holds for some I = k 1, we can verify that
Xx, Uz
= -UU +
Dx)
=X UI UUU.DUD CU UTUUSU.
-UU-1UUU +
DRU
= A
+D)) XAU+UO.
uses the fact that there exists some D D3 E R2x2 such that D; = diag II, & D D?), so shuffling the first U. U.D.U. = D. Similarly, we have D = D This concludes eq. (2)). and eq. (26 could be acquired similarly.
diagonal bub-blocks of D, does not change the matrix, and we have
28

---

Published as a conference paper at ICLR 2026
Next, we will establish the following equalities for X1 and G₁:
Xi & U=V,, Dj -
(R) tR)
= J Xi(D; & U.U.RRU.U
U=li, (27)
di d
k=- di
1=
d
Gi & U_U D. - IR)
GilD _ WAU. U
U2U0.
(28)
di
The proof follows by similar reductions as proving eqs. Finally, consider the in-context risk under the permutation of Up and Ug. Since each pair of (ra, y1) is equivalently sampled from Gaussian distributions, we have Xo _ XoU Therefore, d dt
(88) and 20
LD; & tR)
1=0
2Ex..W tr
&r $ tR) (I (I- dt
a
(I- (--
1==0
(u((u-M) (Xo - U=U)) lu & U+lo, Dj + 18))
(1 (--M))
1=
=2d[Exouuu Ir (I((MMMU U. -. GD + WLU.RU.
U2Uo(I-M) MM)
dt
1=0
=2dExx tr (I (I-
Eu,,u [L & ULU. RU.U1)
(T (T-M))
dt
1=0
Exo
(r
(1 (I-
d - GL i ===
(1 [- - di
d
L(D; $ 17t)
dt
t=
where R= Ru, UULU! RU.U1] = diag(I, I, s R², = Rj. R2 = R and R3
is the j-th 2 x diagonal block of R. The 4th equality uses the fact that tr [ M) M)] is ex-
tracting the right-bottom element of A, so it should be equal to tr (r[(T- MYU U_ AU+U.(I M)] for any matrix A. This concludes that eq. (16) holds for D.-
Till now, we have proved that eq. (T6) holds for each one of A,, B1, Ct Dz. The proof of the whole theorem is then completed by applying Lemma 8
D.3 PROOF OF THEOREM 2
Proof. In this proof, we follow the same notations as the proof of Theorem
where the constant 1
factor is dropped and Zo. Xo, Yo simplified Zo. Xo. Yo respectively.
σ1 0 0 Zo =
0
0
Ptest
0 0 ((2))x(3++3)
(29)
0
0
yn
ylest
Let Z₁ IR 2dx (3n+33) be the I-th layer's output and let X1, Yi E (3d+33++3) be its first and last d rows. Our goal is to prove that, for any E E AU BUBUU BUBBUUU BUCU and an arbitrary matrix R E Rdxd Xd for there exists R E S₁ (Sx: for C, Sp for D) such that
d
C(E +
((E
(30)
di
di
1=0
(=
The proofs of eq. (30p for A,, B, and C₁ are identical with the proof of Theorem I so we omit them. We will be focusing on D; for the rest of the proof.
Let US E Rax and Up E ()++1))(+11) be uniformly sampled permutation matrices. Let US = diag(U;, @ 0iag(1,0, and U = Ut @ diag(0, 1). Therefore, XoU is shuf- fling the 1-st and 3-rd columns among each 3-column sub-block of Xo (except for the last 3- column sub-block), and XoUg is shuffling the 2-nd column among each 3-column sub-block. Next, let U2.UŽ U3 E R ()x+1)x(++1) be diagonal matrices with uniformly sampled +1 entries. Define
U1 = Ug & diag(1, 0, and U1 U @ diag(0, 1, 0). It can then be verified that XoU1UL d Xo-
29

---

Published as a conference paper at ICLR 2026
To simplify the notations, let U= denote U.UiUğU.. We will focus subset of Sp:
S { diag(L, A1,A2) fn+1 @ A3 A1
M
g (10))
Assume D3 diag(L, A2) + 4n-1 A3 Spp defined above, one can verity that it is a block-diagonal matrix constructed from the same 3 x 3 sub-blocks, and thus is invariant under U= D&U_. We will then prove that for any I [1, L},
X1 - U=) = XiU=,
(31)
Gi - U=) -GU-,
(32)
d
dt Xi & U=, D; $ IR)
Xi(D; $ -U=RRT"
U=.
(33)
dt
1=0
1=
dt Gi & U= Dj - tR)
GD D; $ tU= RU
U=.
(34)
dt
1=
t=0
These results can be acquired by similar proofs as eqs. risk under the permutations of Uz.. Similarly, have Xo
25 to 28
We then consider the in-context
d
XoU== and
d C(D; & (R)
1=
d Yr, & tR) #tR)
(I- (4-M)) M)]
-2Ex..W
(1 M)Y dt
1=0
2d 2x
tr
(1 - (I-M)G]((X 2 -=)ミー1
Cr & U=, $ (R)
(1 (1-M))
di
1==
2dd tr (1 -1-MM
d Gr # Ul RU
((T-MM)
dt
1=0
Ex tr (I- d
GD(D, - t tEu_ U (U=RU]])
(I (I-M) MD)
dt
d C(D: - tR)
dt
t=0
Let Rí be the j-th diagonal block of R, then R² - Σソ R2 (101) R2 -Ro+1 1 (19b) (101). R² + Σド Ryo
0 and R = Eu_ UU= = & R I
0000 (100
This indicates that eq. 30 holds for each D, E Sp,, and thus the proof of the whole theorem completes by applying Lemma 8 and noticing that Spp C Sp.
D.4 PROOF OF THEOREM 7
Proof. We keep the same notations as the proof of Theorem
dropping the factor and simplifying
Xa. Yo. Zo No. Yo. Zo. follows:
Zo
0 0
0
0
E (2n-))
(35)
F1 1/1
Yn
Frest
M/est
Note that we now have Xg and Yo containing both on and yı- Define
X
[x] 0 En
Trest
0
X = [0 21 0
Jn Øtost]
= U1 0
Un Ytest]
we then have Yo = X X = X + WX. From the parameter configuration in eq.
12 the update
rule of the first attention layer is
X1 = MD AXMD, Y Y. X WX.
(36)
30

---

Published as a conference paper at ICLR 2026
The update rule for the following layers is the same eq. (1)). We going prove that, for any E E AU BU CU D and arbitrary matrix R Radd (RR x for D), there exists R E Sy (Sx for C, Sp for D) such that
d C(E
d C(E & tR)
- tR)
(37)
dt
dt
1==
1=0
Similarly Theorem 1 uniformly sample U1 E R orthonormal random matrix, and
let Uy: = E1/2 Σ 172, Under the condition that Bj = bı for some by E IR, we have
y= = Yi II + (X j=2
Fr X [1 (XX - then have Y; Fi + WG,, According to Lemma 9
Let
X (T,,
G₁
d
C(E $ tR)
dt
1=0
tr (T- (R Mi)
E N
1-0
tr (I -M)M (E & IR)FL(E # tR)(I (R){T-M)]]]
Ex,.
dt
+ dt
Ex..
tr (1- (E # (R)W NOGWTGGLEE & +R)(L- - M))
Exo tr (1 (1- d
ग( FFL(E & (R) (1 [M-
2d Xx tr (n((--MGT>> at
d
GI & IR)
(T- m)
1==0
Next, we will show that eq. (37) holds for each one of A,, B,, C₁ D, for any E [1, L]. Equation (37) holds for A;.
One can easily verify that eqs. and 18) still hold. Furthermore, eqs. 19 and 20 hold for both Fi and Gt. With these observations, we can then verify
C(A: # tR)
dt
t=
(1 (1- M)F (x & Us) a :t FFl( A Us & EtR)
(1 (M) ]
Ex tr
+ -2dEx..U1
tr
1u((1- MMG. (x # Un} d
G x Us. Ai x tR)
(I (I- Mn)
dt
1=0
-1Ex..U1 tr (I-M)FTUTUs
गt FL & U5 RU')
(1 (1- (1-
1=0
+ 22dEx..U1 tr [ ((I-MGGTUX- Uy di
d Gr(A, & UU RUS))
(1 (1- [M
==
=2Exo tr (T- (I- dt
di
FL $ trla) 1=
(T (T-M)
+ 2ddEx, tr (I MIG.E-1 dt
Gu(A, $ trls) 11 (1-M))
d C(A: _ trla)
di
t=1
where r By, -US RUS] (r RT.12).
31

---

Published as a conference paper at ICLR 2026
2. Equation (37) holds for Bi.
From the definition of Fi and Gt, we can verity that
d Vi(B _ IR)
dt
r=0
= WCG. DJ)
II (++
j=#11
Define
P- 7. II (1-0,N(X,,(Xxx1111)0)
j=++1
G = CWCC
II (1- j=i+1
We then have
Yi
_ tR)
FB tR)
d
+ dt
GiB & tR)
dt
dt
1=U
1=0
1=0
Similar to eqs.
20
and 12
we can prove that
d
d
(Xo
/ Uz
+ IR)
= Ug
(B, - U0,5 RUg)
dt
dt
1=0
1=
Ew W d (B, + tR)
-ー d
( き Ta)
dt
dt
==
Without loss of generality, we assume that r rd/ tr( Rt) 1. Then, one can verify that
s 1/2 tt((R), and let 7 =
d dt
C(B, tR)
1==
= -2Exo,u
tr
(I (I-M) (x A Us:)
d
[ v Us. Bi $ tR)
#t)
(I (T- M)
1=
+ 22 Ex.. w
tr (1 (L-Mi) (I-M) MMGT
TT dt
d CBB, $ tH)
(1 M)
1==
=2Ex, ir
d (1- MI-
ri $ trla)
(I-
(M-M)
1=
2 Ex.
tr (e((1- MIG x d
CB & (
(T- (-M)
dt
1-
(I FT d dt
=2Ex.
FiB $ trla)
(I (1-M)
1=0
+ 2dx tr (1- MI-MGGT'
-- dt
d GIBB & trdaa)
(1 (--\)
1
- 2d Exo (r (1- -I(1IMMGTT- di
Gi + trla)
(I (I-M))
1=
[B; + trTa)
2 dt C $ trTa)
dt
1==
The last inequality assumes the positivity of the term involving G.. Otherwise, one can simply flip the numerator and denominator of y and scale the derivative of F) instead of G, to yield an additional positive term besides the risk term to finish the proof.
3. Equation (37) holds for C, D,.
Similarly, one can verify that eqs. (23 and 24 still hold (also eqs. by following the same reductions Theorem T with Fj and Gt.
25 to (8 and finish the proof
32

---

Published as a conference paper at ICLR 2026
D.5 PROOF OF PROPOSITION 3
Proof. Let A = oila. B₁ = byla, C a aa and D; = d (f), + Df for e [1, 2]. Let Z1 E 122x(3n+*) be the output of the 1-th attention layer, and let X1, Yi E (dx(3)+3)
be its first and last rows respectively. Note that Yi in this proof does not contain ytest-
d= 0
Let D
Dj
0 0 0000
(note that the last row of Dj is masked out by M, so we
d8
simply set it to 0), and D
(
0 f. 0
8)
8). We use D as an abbreviation for DI. and use d to denote
the elements in D. One can verify that
X1 = Xo + a XoM (d @ DJ. D2) + 1--1 @ + @ D
[ (1+a
0tt2 Σx=1 d1,1
a1d
(1 + a
a1t2 Σ Lin1
n+1 de
ajdion aidz
(1 + a1d)
atto E di,++
Similarly, we have
Yi = Yo + by YoM (diag(An @ D9) + Lnt @ D + DJ 2 Ds)
bity Ei=1 di.1%s
(1 + b1dg)11
bity - din bily Li=1 din+
(1 + didz)
0
By the definition of linear attention, we can show that
TF(Zoi Q()))) (Y2)3n43 b2 (c2XT (X1 K1)3443 + (D2)3n43)
3n+2
b202arde d%
Σ (x)) Trest.
Define AX1 = [0 ottgdn+11 0
+In 0]. and let X, = X1 AX1, then TF(Zu: (Vi Q1 ) = TF(Zo: (V. Q X1 + x1) + TF(Zo: (Vi, Q1 X1 t AX1). Let b1 d1d + (1 bid az d a, bit,att b, bec2ard c, we then have
n+1
m dRat
Σ dneny )) Frest
TF(Zo: (Vi 11) X1 + X1) c
you
+
Σ
Σ
n
y1 +
dz,edki U
Peest (38)
j=1
n+1
TF(Z;: (V.Q02 X1 + AX1) be
d dn 1. Tost Feest
j=1
be
Σ
dyadu+ts Ufffiest
(39)
Now consider the in-context risk,
L(V.Q) Ezo.w [F(o VV))) + W W
E.z.. (TF(Zo: (V. (VQ + W (eest) ( (TF(Zo: (V, + Warea))
= EZ..W (TF(20: (V.Q). + X1)+ + WIreet) ) (TF(Zo: (V, (V.Q),X, + xx1 + W Wrrest)
+ +2EZ. [TF(Zo: (V. +AX,) T (V + -X1) + WIrees)]
+ EEzo.w
[TF(7; (V,Q), X1 -xx)) TF(Zo: (V.Q), -AX)]]
33

---

Published as a conference paper at ICLR 2026
In the equation above, the 3-rd part is always positive. We then examine the second part:
Ez, W [FF(Zo: (V, AXI) (TF(Zo: (V,Q),,, t X1) + W (iest)
Ez, W ( Frest 01 Ttest + I Frest Ptest 1tst =0,
where 81 = D5=1 Datl dz.ddn+1.5 U/ a and = bc D Ziml d1 83 W are independent of Trest Therefore, L(V, Q) attains its minimum only if TF( QV X1+ + AX1) = O, implying dn+1, for E [L. +
( y. / +b Ek= Ei dy.sdk.s ) yyak
In the following analysis, we will assume that the last row of D is O, and let M E R=x(x+1) be the first rows of D. Additionally, will drop the factor in eq. 38 since its position could be
substituted by a and b. We then define W Ei= 4:2; + b شド= k= E dzndk.s) )
X [ In and Y [yn yn. One can verify that
T + bY = aWXXT XWXMM XT.
(40)
Furthermore, the in-context risk could be expanded
L(V,Q) = Ezo.W
W Frest + Wieest
= Ezo.w
Frest (W + +W) TWW + W)ttest
= EZ.. tr (W + +W) (W + W))]
= =Ez.ww tr W
W) [W)
+ 2r W W) + tr(w tr(w wW)
We will use the identity XXXXXX XBX X (tr(A)tr(2)) tr(BB + tr (r(ABT) + d'tr(AB)) for any
A, B E Raxn which can be acquired by expanding each element and applying Isserlis' theorem. Let Ti = tr(M and 7-2 = tr(MM MMT). then
Ez..w w [Ir((aXXT bWXMM X (aWxx + +WXMM x])] Ezo.w XXXXXXTWXXXX + 2ab WWMMM WXMM XT)
+EZ. [b² trt((MM TTW WXMM [
=dEz, [a (=[XXXxXX] + 2ab (xX XMM xT) (MMMMM X T XMMTXT)
=a2d + d) + 2abd (n djTi + +Td + (100:++++1))))) (1dd)))).
Simultaneously, we can verify that IEz., (tr [tr( W)] = de and
Ez,w (rr (or(we ( [ = Ez,.w [aW WXX + bW T W XMM" x adL + bd
Combining the results above, we aim to find the optimal a, b, M that minimize
C(V, Q) co + ciTi + t + caT2.
de
where
Cu d'n(n +1 +++ 1 2an, 61 2ab(n 1 1+ d) +b, 02 = 62, c3 = b22 (1+
Since C3 2 0, to minimize L(V. Q) we need to minimize T2. Given that MM is symmetric, we denote its eigenvalues X, i E [1 n). Then by Cauchy- Schwarz inequality,
tr(M MMT) = Σ x
A
(MMM). n
1
Therefore, L(V, Q) is minimized only if the inequality above holds with equality, which implies that \ >; for any + j. This concludes the proof by showing that there exists IR such that MM Xla, and thus DDT diag(ATa.
34

---

Published as a conference paper at ICLR 2026
D.6 PROOF OF PROPOSITION 5
Proof. We will continue from eqs. 38 and (3 After applying token-wise dropout, we have
TZ (Vi, (..Qiii1, t X1) =
a002 3i-2 +
bo
)0
10t1-20 0f o yi 3n+1 3n+3
02 Trest
02
3i-1
dj
01 01
31 32
yy T
0j
of 02 Frests
Dtest.
(41)
n+
TF (V, QU) XI + AXI) = =00 +3+
d3aduni, of 31 gnt1 T
vest
where a = byc201 d by dz + a1d)), b = b202a1 d% (1 + b bl di and c b202a1 di tygita One can verify that our previous analysis about TF(Zo: {Vi Qil + AX1) still holds and we thus have 12+1, = O. We then define:
O = diag(o)....
of 3 3) E R=x"
Of = di .... of E R"** for E [2].
0 diag
,gn ((n+t))(n+1),
By defining
Σ 3i- (10
bo
0 y 00
38
34
Σ 0
3i
disdes
09 01 "
3k-2
1=1
j=lk=1
1=1
One can verify that
= A+ B + C é aY
+ CYORMO,MTTO.. oo
Then, we will compute the expectation of each term in the following decomposition:
Q(V.Q) Ez.. w tr (W (W) 2t ( (WTW) + tr(w ww]:
Specifically, let Tj (r(M )) Th = ((MMTMM)) T3 ||MII Ta Ei= 1| 12.
TB = =ド
11 MM 11.. we then have
(A A)] a² dnp' + n(n
1)p² + (1+ d'mp'). Eitr(B B)] = (np + 1)p² +(1 d)mp"),
Eltr(C T C)] -1 - +1+ -000--"11 pOTT's
(1
+ - p
+ pTT2 + dp)
B(tr(ATB): = n -1
B[(r ((l((AT))] =
+ (1
EB = (ep + (nn-)(((nn))ii----- (1+
(W A)] = np', WW bbpmpp,
[[]][[]]==-- = cdppTT.
Summarizing our analysis above. miny ((V,Q) is cquivalent to: + ??
min(co + cıTi + + + cyTs +
M
where
co = 1= + d)p + +1²) + 2n (a 2 ab + n(n 1) + b)3
d)p' C = 2(a + b)c(p' 1n-1)p² (1 + d)p²)) + 2cp²,
03 == (1 - pd + )²),
1 = (11 16) + (p3 - 1),
C3 == e + d) - DO),
66 ppp
It is easy to verify that (22 C3 C4, cs Cg 2 0.
35

---

Published as a conference paper at ICLR 2026
D.7 PROOF OF PROPOSITION 6
Proposition 6 (Restate). Let dh denote the number of non-EOS tokens. Given any L-layer, single- head, d-dimensional linear-atientoon transformer with EOS tokens:
TZ( (V, Qi, P(hee)) =(Zi)dd+t1, (Zo)
=0,
where
21 E gdx(d=+1), W. Pt 1((
Zi = Z1-1 M22 + M = diag(Ta,, 0).
There exists an L-layer, two-head, 2d-dimensional linear-attention transformer operating without EOS tokens:
( (Vi, Qi, Pt hELL)AE)) = (Z1.)a2d.
where
Zi E R22 2dxd, Vl. Q E R 2dx2d Pf E ddd,
Zi = -1-1+ + (Z 12
Such that for any 7 E Rdxdd, by letting Zo 1z 0] and % =
0
we have
TF TF(Zo: (Vi, Pi)re()) TF TF( (V) Qi Pt (E(LLMEE)
Proof. We construct V
Q/, and P/ as follows:
Vi
0
Q
V
Φ
P (Pi) I:dp?
0
0
0
0
V
0
0
0
Qi
P P== = 10 (Pi).d,tI..
V1
0
0
0
We will show that for any 1 E [L], it satisfies Zi (Zi),(1dd-1) (Zi).dd
One can verify
0
(Zi)
that it holds trivially for l = 0. Then, suppose it holds for some ( = k 1, we have
Z Z ZiZz-1 (Z 0. Pb) + Vi -- 2 2 P
Vz(Zz-1)a,d,
(c(Zk-1):a (
(2))).d)...pp) )
-
0
Vo(Zk-1)n.dd +[voeaa ] (10 (Z
+
)., 1dd, Qx (2 ( 1)) 10 (Ph).a+1))
V/Z M
M(T (ZT Q+ (zx 1).1.4 (Pk)..
- +
0
0
+ 0
V 1Z + (Pk) (Pk)ad,+1)
(Zx)
0
+
0
(Zz).dtt1
The proof is complete.
36


Reviews:
[{"review_id": "UepSdIb2g6", "reviewer": "Reviewer_8Lar", "summary": "The paper discussed three formats of in-context leanring: Single, Pairwise, and Triplet.\nFor Pairwise, the paper provides Theorem 1 to theoretically show that an L-layer linear transformer allocates one layer for embedding concatenation and utilizes the remaining L−1 layers to perform gradient descent.\nFor Triplet, the paper provides Theorem 2 to theoretically show that task vector naturally emerge from pretraining on triplet format.\nThese analyses are coupled with experiments in Sec 6.\nThe paper further show the task vector fails to predict on bijection tasks.\nSec 6 further shows multi-vector injection works better than vanilla", "questions": "(1) I wonder whether L134, equation (2), should have \"-\" instead of \"+\" for loss.\n\n(2) I did not understand L290-292. I believe the paper is discussing causal attention. Why is it said \"as bi-directional attention allows\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The paper considers three ICL formats Single, Pairwise, and Triplet, making the view broad.\n\n(2) The paper provides Theorems 1 and 2 for Pairwise and Triplet and gives clear explanations/intuitions.\n\n(3) The paper identifies a task bijection that the one-shot approach TV cannot solve.\n\n(4) The paper shows that multi-vector injection works better than the original task vector.", "weaknesses": "(1) Different from real-world LLM, the considered self-attention is without normalization, with position embedding as a one-hot vector.\n\n(2) For Table 1, the conclusion is not well-supported. Will more advanced LLM lead to different conclusion?", "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": "87qoJ67Al8", "reviewer": "Reviewer_iUth", "summary": "This paper investigates task vectors in in-context learning (ICL), proposing that they act as distilled representative demonstrations. Through analysis of linear-attention models and saliency maps in large language models, the authors show that task vectors emerge from triplet-formatted prompts and encode task information via weighted combinations of examples. They identify a key limitation: task vectors are restricted to rank-one mappings, limiting performance on complex tasks. A multi-vector extension is proposed to address this, showing consistent improvements. The work provides useful insights into the structure and limitations of task vectors.", "questions": "1. The paper evaluates models such as GPT-J, Pythia, and Llama-2, but does not report results on more recent architectures like Qwen or the Llama-3 series. Do the rank-one limitations identified in this paper also persist in newer models such as Qwen or Llama-3 family?\n\n2. The empirical evaluation focuses on synthetic and controlled linguistic tasks, leaving it unclear whether the observed rank limitations extend to more complex, real-world applications such as Question-answering or reasoning tasks. Have the authors explored such settings in their analysis?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a well-structured theoretical investigation into the emergence and functionality of task vectors in in-context learning (ICL), centered around the representative demonstration hypothesis and its implications under linear attention assumptions.\n\n2. The qualitative analyses, including weight matrix visualizations and saliency maps, are clearly aligned with the theoretical framework and help to support the proposed mechanisms intuitively.\n\n3. The paper evaluates its hypothesis across a diverse set of synthetic and linguistic tasks (e.g., linear regression, bijection), and the experimental outcomes reflect the predicted rank-one limitations, particularly in settings requiring high-rank mappings.", "weaknesses": "1. The theoretical analysis is limited to simplified linear attention models and does not incorporate components such as softmax attention, causal masking, or nonlinearity, which are central to modern transformer architectures.\n\n2. While the appendix refers to prior approaches that leverage task or in-context vectors to improve performance, the empirical evaluation is limited to standard Few-shot ICL, TV, and the authors’ multi-vector variant. Including direct comparisons to these prior methods would help clarify the novelty and effectiveness of the proposed approach.", "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": "n2fxycLgyj", "reviewer": "Reviewer_t2pe", "summary": "This paper looks into how task vectors emerge in ICL. The main idea is that a task vector, formed at a special prompt token, works like a distilled version of the whole demonstration set, basically acting as a single representative example. The authors support this through a theoretical analysis using linear-attention Transformers, showing that such task vectors can only express simple, rank-one mappings. They then back this up with experiments, including on large models like LLaMA-7B, which show that real models behave in similar ways. To deal with the limitation, they try injecting multiple task vectors into the prompt and find that this gives small but consistent gains. Overall, the paper offers a clear explanation for what task vectors are doing and what their limits are.", "questions": "- Please clarify why the task vector remains strong on bijection tasks in Table 2 despite the failure cases in Table 1. For example specify differences in prompt structure, presence of additional demonstrations, model nonlinearity, or training data effects that could compensate for rank limits. If this discrepancy is convincingly explained and bounded with clear conditions I am inclined to raise the overall score. \n\n- Could the multi vector approach be improved by learning to select or weight vectors adaptively rather than uniform insertion?\n\n- Do you expect the task vector mechanism and its limits to extend to more compositional or long range reasoning tasks, and what evidence would most directly test this?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper develops a clear and rigorous theoretical framework that explains the role of task vectors in few-shot learning. \n- A major strength lies in the bijection experiments, which directly validate the predicted limitations of single task vectors and demonstrate the practical implications of the theory. \n- The authors show that the same task vector behaviors emerge in large pretrained language models, reinforcing the relevance of the conjecture beyond simplified settings. \n- The proposed multi-vector strategy builds naturally on the analysis and provides an interpretable method to mitigate expressiveness constraints. \n- The clarity of writing, thoughtful experimental design, and theoretical depth make the paper both accessible and impactful within the ICL literature.", "weaknesses": "- The work builds on linear attention. I acknowledge that generalizing the proofs to standard Transformers is challenging, but the lack of a formal extension limits the scope.\n\n- The multi vector strategy yields only moderate gains and likely needs further refinement for higher complexity tasks. Most of the tasks tested are short or simple, so it's still unclear how this approach would hold up on more complex reasoning problems.\n\n- In Table 2 the task vector remains strong on bijection tasks, which softens the rank one limitation highlighted by Table 1. The paper should clearly explain why this happens and under which conditions the limitation manifests or is mitigated.", "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": "GjwkpQApMl", "reviewer": "Reviewer_4skn", "summary": "This paper proposes a conjecture that task vectors (TVs) facilitate zero-shot inference by distilling from the original ICL demonstrations a single representative demonstration encoding. To support the conjecture, the paper uses linear-attention models to show that task vectors naturally emerge as a weighted summation of the in-context demonstrations. A key implication of this conjecture is that a single task vector, like a single demonstration, is limited to representing rank-one mappings. To test this limitation, the paper introduces a class of bijection tasks that require higher-rank mappings. Empirical results on practical LLMs validate that the standard TV method consistently fails on these tasks, while full ICL succeeds. Building on this finding, the authors propose an enhanced TaskV-M method that injects multiple task vectors, which successfully improves performance on these complex tasks.", "questions": "- Theoretically, can the TaskV-M method achieve the same performance as ICL?\n- For the triplet tasks, how are the triplet demonstrations delimited? Why does the attention focus on the yi tokens instead of the delimiters?\n- I might be a bit confused, but why is that, for the bijection task, in table 1, ICL outperforms TV, but in table 2, ICL underperforms TV?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The organization of the paper is clear and easy to follow. Starting with a clear conjecture, the paper goes on to theoretically demonstrate the conjecture with a constructed simplified case.\n- The use of a bijection task to identify a weakness of TV methods is sound and interesting.", "weaknesses": "- For table 1, while TV performs quite badly on bijection tasks, it seems that ICL is not much better. Given that ICL consistently outperforms TV on all tasks, what I can see most clearly is that the bijection task is much harder than other tasks, but not that \"ICL preserves performance in many case\".\n- Some analysis seems to have been demonstrated by previous papers. For example, the findings of how information flows in ICL in Figure 3 (a,b) has been demonstrated in [1]. \n- I do not quite understand the implication of Figure 4c. It is difficult to see from Figure 4c that the performance of TV and 1-shot ICL are \"parallel\".\n\n\n\n[1] Wang et al., Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning, EMNLP 2023.", "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 (Poster)