Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insights
Authors: Haolin Yang
Hakaze Cho
Kaize Ding
Naoya Inoue
Venue: ICLR 2026 Poster (2026)
Keywords: Mechanistic Interpretability
Large Language Model
Task Vector
In-context Learning
Primary area: interpretability and explainable AI
TL;DR: We propose a new method for finding task vectors in in-context learning and investigate how task vectors work.

Abstract: Large Language Models (LLMs) can perform new tasks from in-context demonstrations, a phenomenon known as in-context learning (ICL). Recent work suggests that these demonstrations are compressed into task vectors (TVs), compact task representations that LLMs exploit for predictions. However, prior studies typically extract TVs from model outputs or hidden states using cumbersome and opaque methods, and they rarely elucidate the mechanisms by which TVs influence computation. In this work, we address both limitations. First, we propose directly training Learned Task Vectors (LTVs), which surpass extracted TVs in accuracy and exhibit superior flexibility—acting effectively at arbitrary layers, positions, and even with ICL prompts. Second, through systematic analysis, we investigate the mechanistic role of TVs, showing that at the low level they steer predictions primarily through attention-head OV circuits, with a small subset of “key heads” most decisive. At a higher level, we find that despite Transformer nonlinearities, TV propagation is largely linear: early TVs are rotated toward task-relevant subspaces to improve logits of relevant labels, while later TVs are predominantly scaled in magnitude. Taken together, LTVs not only provide a practical approach for obtaining effective TVs but also offer a principled lens into the mechanistic foundations of ICL.

Published as a conference paper at ICLR 2026
TASK VECTORS, LEARNED NOT EXTRACTED: PERFOR- MANCE GAINS AND MECHANISTIC INSIGHTS
Haolin Yang'
Hakaze Cho²
Kaize Ding' Naoya Inoue 2.4
University of Chicago AAASST Northwesteern University
'RIKEN
haol haolinyang2001 wucchiicagecccccco..oooooooooo edu, yfzhaol jachaoBjaisto jaist. ac. jp
kaize kaize. si.ee..diggmoortwesstem edu, naoya-1@jaist.ac:
ABSTRACT
Large Language Models (LLMs) can perform new tasks from in-context demon- strations, a phenomenon known as in-context learning (ICL). Recent work suggests that these demonstrations are compressed into task vectors (TVs), compact task representations that LLMs exploit for predictions. However, prior studies typically extract TVs from model outputs or hidden states using cumbersome and opaque methods, and they rarely elucidate the mechanisms by which TVs influence compu- tation. In this work, we address both limitations. First, we propose directly training Learned Task Vectors (LTVs), which surpass extracted TVs in accuracy and exhibit superior flexibility acting effectively at arbitrary layers, positions, and even with ICL prompts. Second, through systematic analysis, we investigate the mechanistic role of TVs. showing that at the low level they steer predictions primarily through attention-head OV circuits, with small subset of "key heads" most decisive. At a higher level, we find that despite Transformer nonlincarities, TV propagation is largely linear: early TVs are rotated toward task-relevant subspaces to improve logits of relevant labels, while later T'Vs are predominantly scaled in magnitude. Taken together, LTVs not only provide a practical approach for obtaining effective TVs but also offer principled lens into the mechanistic foundations of ICL!.
INTRODUCTION
Large Language Models (LLMs) possess the remarkable capability of performing novel natural language tasks by learning from demonstrations included in the input without training, phenomenon referred to as In-context Learning (ICL. (Brown et al., 2020; Radford et al., 2019). ICL has revolutionized natural language processing through its extensive empirical success in enabling swift and efficient adaptation of models to downstream tasks (Dong et al., 2024; Liu et al., 2021).
Since its effectiveness is difficult to reconcile with the traditional framework of machine learning centered on model training (Ren et al., 2024), investigating the internal mechanisms of LLMs that enable ICL has attracted substantial attention. Among these efforts, one prominent line of research shows that LLMs leverage demonstrations by summarizing them into task vectors -sccinsu vector representations of the task exemplified by the demonstrations (Hendel et al., 2023). These TVs can be injected (added) into the hidden states zero-shot prompts without demonstrations to achieve ICL-level performance. Subsequent work has primarily proceeded in three directions: 1) studying where (c.g., from LLM hidden states (Hendel et al., 2023), attention head outputs (Todd et al., 2024; Yin & Steinhardt, 2025), or MLP outputs (Merullo al., 2024) at different layers) and how (e.g.. PCA-based approaches (Liu et al., 2024) or complex optimization methods (Li et al., 2024a; Cai al.. 2025)) to extract and construct TVs, with the practical goal of boosting performance through injection; 2) investigating how the ability of LLMs to form TVs gradually emerges during pretraining, typically using small soieeed-fro--rrateee models and artificial tasks such as regression (Han et al., 2025; Yang et al., 2025b); and 3) demonstrating that TVs naturally arise from the LLM architecture itself, and providing theoretical guarantees for their emergence (Bu et al., 2025; Dong et al., 2025).
Despite important contributions, prior studies face key limitations. First, existing approaches often depend on opaque and complex filtering or optimization to construct TVs from model representa-
Code Implementation: https:// htttps://gittuu.ccm./ http://gtttuub..cmmm//////////// /laerned_i_

---

Published as a conference paper at ICLR 2026
tions, making them inefficient and reliant on the model's representational quality. This dependence can produce suboptimal TVs and mischaracterize their true effect, while the opaque construction procedures obscure understanding of TVs' mechanism. Indeed, most works stop at showing that injected TVVs improve performance but leave unanswered the central question of how LLMs leverage TVs to make correct predictions. This gap spans both the low-level interactions, referring the microscopic localization of model components that interact with injected TVs to express their effects during forward computation, and the high-level channels, referring to the macroscopic mechanisms by which TVs ultimately steer outputs toward correct predictions. The lack of explanation reduces the model's deployment of TVs to an uninterpretable black- box function (Merullo et al., 2024).
In this work, we address the first shortcoming by proposing to directly train Learned Task Vectors (LTVss by adding vector to a specific layer's hidden states and optimizing through gradient descent (Figure (A)), which finds the optimal TV un- constrained by the quality of model's representations. LTVs not only out- perform constructed ones across clas- sification and generation tasks but also demonstrate greater fiexibility and scalability than extracted ones. Furthermore, through analysis of in- teractions between TVs and model components, we decode the low-level mechanisms by which LLMs interact with TVs: injected TVs are primar- ily utilized through attention head OV cireuits (Figure (B)). We also characterize which attention heads most decisive in leveraging the injected TVs, focusing on their attention and distribution patterns. Regarding the high-level influence channels of TVs, we show that despite the abundance of nonlinearities in Transformer layers, the propagation of injected TVs through subsequent layers is largely linear, involving a rotation that aligns TVs the subspace spanned by task-related tokens and scaling that adjusts their magnitude (Figure 1 (C)). We further observe a distinct pattern: the rotation effect attenuates as the injection layer index increases, while the scaling effect becomes the dominant factor translating TVs into output changes. In summary, our work introduces an efficient method to obtain effective TV's and provides a comprehensive exposition of the mechanisms underlying TVs' effectiveness.
A Training TV negative -raddient
B
ov Circuits Induce TV Effect
Remaining Layers
Wy
Wo
Wy
We
Wy
Wo
Layer Updates Rotate and Stretch TV
Rotatiion
6 aaterr Layer Updates WO
LTV
Layers
StREE
Bad movie, Sentiment
Figure I: (A) We directly train Learned Task Vectors (LTVs) to be injected, which influence model outputs through later layer updates. (B) In the low-level interactions between TVs and later layers, the OV circuits of attention heads are the crucial components interacting with TV's to induce their effects. (C) On a high level, subsequent layer updates act on TVs as largely linear transformation of rotation and stretch, with the rotation aligning TVs with the relevant task subspace to promote prediction of task-related tokens.
2 RELATED WORKS
Task Vector and ICL. The hypothesis that TV's form the mechanistic basis of ICL. was first proposed by Hendel al. (2023), who patched ICL hidden states into zero-shot prompts at certain layers, achieved high accuracy, and argued that in-context demonstrations are compressed into TVs applied during later updates. Follow-up studies (Todd al., 2024; Li et al., 2024a; Kahardipraja et al., 2025; Liu et al., 2024) extended this idea by extracting TVs from specific components (e.g., attention heads, MLP) and injecting them. The universality of TVs has been validated across model scales (small vs. large open-source) and task types (mathematical vs. natural language) (Han et al., 2025; Yang et al., 2025b; Jiang et al., 2025a). Yet, little known about how TVs enhance performance after injection, how they interact with later components to shape outputs.
Mechanisms of Task Vectors Current explanations of TV effectiveness remain preliminary, more sketches than systematic analyses. For instance, Hendel et al. (2023) observed that TV injection is more effective in earlier than later layers. Todd et al. (2024) reported that TVs exhibit word2vec-sty-e arithmetic (Mikolov et al., 2013), with Bu et al. (2025) giving theoretical account of this property. Furthermore, Han et al. (2025) and Jiang et al. (2025b) found that TV effectiveness depends on how well hidden states of task's prompts can be separated from others in the L.M representation space.
2

---

Published as a conference paper at ICLR 2026
LLM steering The success of TV injection in restoring ICL performance parallels recent advances in LLM steering (Zhan et al., 2025; Li et al., 2024b; Panickssery et al., 2024), where vectors are added to hidden states mitigate undesirable model behaviors (Lee et al., 2024; Bayat et al., 2025). Prior work also explored training steering vectors directly (Cao et al., 2024; Dunefsky & Cohan, 2025), motivating our strategy of training T'Vs rather than relying on complex selection construction.
METHODOLOGY
Transformer hidden states and ICL According to the autoregressive structure of Transformer
LLMs with residual connections, zero-shot input query Zy of N tokens (e.g., "I like this movie. Sentiment:") is sequentially embedded and updated across L layers into d-dimensional hidden states. At each layer, the hidden state of token at layer is updated as hf = ht + Σx= d + +m where a is the output of the k-th attention head (head (I, and mf is the MLP output. Concretely, a depends on the previous layer's hidden states of the first i tokens 1 through:
de
Σ J. O, W
G.
(1)
where W and W. E RR4x are the value emhedding and output projection matrices of head (1, respectively, jointly referred the OV circuit, with d, being the head dimension. d/ denotes the attention weight from token to j of head (I,k). Consequently, the L layer updates can be viewed as sequence of additive effects, with the last token hidden state at the final layer formed as:
h AN
d m/k)
(2)
which is then multiplied by the unembedding matrix Wo E RMxx to produce the output logits and final prediction. An ICL prompt formed by prepending n demonstration-label pairs to the query, yielding an input sequence : y........ Bm: Ym: By (e.g., "I hate this movie. Sentiment: negative. This movie is great. Sentiment: positive... I like this movie. Sentiment:"). Because Transformer updates depend on hidden states from earlier tokens and layers, the hidden states as well as attention and MLP outputs change throughout, producing different f.y.Cc. and ultimately a different prediction.
Task Vector as a mechanistic explanation for ICL Existing TV studies provide a functional characterization of the mechanism enabling L.LMs to leverage demonstrations as O), i.e., L.LMs make predictions based on the query together with vector 0 that represents the query-label mappings (Hendel et al., 2023; Merullo et al., 2024). These studies propose that is formed in early layers and assists LLM predictions TV from the ICL hidden state stream and add it to the last token hidden state of xy at layer 1, i.e., hly+0 0. O. The resulting hidden state is then propagated through subsequent layers, and the intervention is evaluated based whether it achieves few-shot- prediction accuracy for zero-shot queries. Two major methods of extracting 0 have been proposed, which we treat as baselines in this work":
later layers execute (). Accordingly, they seek to extract the
1. Vanilla TV (Hendel et al., 2023), defined as 0 = 11y.IC -hly- where h W is the layer-/ last token hidden state of an ICL prompt formed with a query different from xy which produces hny-
2. Function Vector (FV) (Todd et al., 2024; Li et al., 2024a; Yin & Steinhardt, 2025), defined as
0 " where d'yA.CI is the attention head output to the last token hidden state given ICI. prompts, and TI an index set of selected attention heads.
Both methods have drawbacks. Vanilla TV injection yields lower accuracy and is highly sensitive to the choice of injection layer I. FV depends on selecting proper head set II, typically determined by ablating heads one by one to measure their impact on output probability, which is suboptimal as it neglects intercorrelations among ablations. Moreover, both methods critically depend on the quality of model's ICL representations (we use 8-ohot ICL prompts to obtain hidden states for the two methods and to evaluate ICL performance). As a remedy, we propose directly training LTVs.
See Appendix D.4 for comparison with more baselines and the reason for omitting them from the main text.
3

---

Published as a conference paper at ICLR 2026
Training LTV Instead of distilling from ICL hidden states, we train the LTV θ to minimize:
log p(yp Lg.. 0 I.L,IP),
where y4 is the correct label for the zero-shot query zq.. LL denotes the set of layers and P the token positions of hidden states where is injected. This approach eliminates the need to manipulate ICL hidden states and uncovers the most effective TV, unconstrained by representation or demonstration quality crucial for traditional TV extraction. Moreover, we do not restrict IP to the final position or single layer as in the baselines. In general, we add [L]] x [IP] different Os the hidden states indexed by L and IP. This design allows us to explore flexibility and scalability of our approach and to test the proposition from prior works that single TV can encapsulate the full functionality of ICL, as discussed in Subsection 4.1. In the special case of L = {1} and P = {-1}. we add one O to hyy following baseline practice. During the training, for multi-token labels, we average the log probabilities across tokens. 0 is optimized using AdamW (Loshchilov & Hutter, 2017) with learning rate 0.001 and weight decay 0.01. Details of the training procedure are provided in Appendix B.
4 EXPERIMENTS
Models We use the following models: Llama3-8B, Llama3.1 Llama3.2-3B. Llama3- 70B (Grattafiori et al., 2024), Llama2-7B, Llama2- 13B (Touvron et al., 2023), Qwen2.5-32B (Yang al., 2024), Yi-34B (01. Al et al., 2024). In the main text, results are reported on 1laBa3.1-8B.
Datasets We adopt three datasets from prior TV research (Todd et al., 2024): 1) Capital: given a country name, output its capital city: 2) Capitalize: given a word, output its capitalized first letter; 3) Antonym: given word, output its antonym. To evaluate TVs on more natural datasets with richer input-output mappings, we additionally consider four classification tasks: SST-2 (Socher et al., 2013), TREC (Li & Roth, 2002), SNLI (MacCartney & Manning, 2008), and RTE (Dagan et al., 2005). We report the prediction accuracy achieved by ICL and the different TV methods across the seven datasets. To test the ability of TVs elicit LLM behaviors in more complex task settings, we also include the Myopic dataset (Panickssery et al., 2024), a generation task described in Subsection 4. We further include three more datasets specifically for investigating the compositionality and generalizability of LTV, also described in Subsection 4.1. See Appendix C for additional details on model implementation, datasets, and ICL setup.
4.1 SUPERIOR PERFORMANCE OF LEARNED TASK VECTORS
Consistent performance superiority of LTV Following Hendel et al. (2023) and Todd et al. (2024),we first inject the TVs at one layer at time, iterating over all layers of Llama3. 1- 8B, and report the average performance across datasets in Figure 2. The results show that our LTV not only consistently outperforms the base- line methods at all layers, but also matches or even surpasses ICL performance-particnn particularly when injected at early layers of both models. The high accuracy achieved by the LTV also makes a viable parameter-efficient finetuning (PEFT) method (Wu et al., 2024; Subramani al., 2022; Turner et al., 2024), since it involves optimizing exactly d parameters, which is lower than most existing PEFT strategies. To demonstrate the potential of LTV as a PEFT method, we compare it against two widely used PeellinesPPeiiie Tuning (Li & Liang, 2021) and LoRA (Hu et al., 2021 20nde-under a comparable parameter budget on SST-2. Specifically, we apply prefix tuning to the key and value projections of all heads at the first layer of Llama3. 1-BB with prefix length 2, introducing exactly d trainable parameters. Likewise, we apply LoRA to the output projection at the first layer with rank r 1, yielding 2d parameters. We then compare these baselines with injecting single layer update at the last-token position of layer 0 (d parameters). Beyond accuracy. we track mean latency and FL.PPs per training inference sample, and peak memory per training inference
75%-
ICL Acc
LTV
50%
Vanilla TV
FV
25%
0%
75 Acc
14
20
Layer Injection
Figure 2: Dataset-average accuracy of injecting the Vanilla TV, FV, and LTV into the last-token hidden states by iterating over all layers and injecting into one layer a time, along with ICL and zero-shot (ZS) accuracy levels. Our LTV consistently out- performs the Vanilla TV and FV across all layers, with the performance gap particularly prominent in late layers. See Appendix D.1 for other models.
4

---

Published as a conference paper at ICLR 2026
epoch. As summarized in Table 1, LTV achieves the strongest performance and best training- time efficiency, and its inference- anference-time FLOPs and memory cost are only marginally higher than prefix tuning, which demonstrates its potential as competitive PEFT method.
Table I: Comparing LTV against PEFT methods in terms of performance and efficiency
Training
Inference
Method
Acc.
Param Cut.
Mean Lat. (Sec)
FLOPS (GB)
Peak Meus. (GB) Lat. (See)
FLOPS (GB)
Peak Mens. (GB)
Prefix Tuning
85.67%
361.51
LoRA
91.43%
2d
526.98
361.52
1TV (Curs) 92.89%
0.049
563.87
43.56
0.024
361.52
16:36
Accuracy of late-layer injection Another notable trait of the LTV is that it still achieves nontrivial performance when trained and injected at late layers, despite an overall decreasing trend with depth. This contrasts with Vanilla TV and FV, which show severely degraded accuracy beyond a certain depth as reported in prior work (Li et al., 2024a; Todd et al., 2024). Our results therefore challenge the idea that a critical depth threshold exists beyond which layers cannot utilize the injected TV. We further analyze the mechanism enabling LT'Vs at different depths to take effect in Subsection 4.3.
Flexibility and scalability of the LTVs Existing TV studies typically inject solely into the last token hidden state (P = (-11) at one specific layer (L = (7+) of the zero-shot prompt. We go beyond this baseline to examine the adaptability of our LTV to more diverse configurations. We set to the middle layer of the model (i.e., 16 for the 32-layer Llama3.1 as the baseline, and then consider the following variants. 1) Keep/ / fixed but inject (and train) at a different position TP = (4), i.e., add the TV to the fourth token hidden state" 2) Inject multiple positions: P= (-5, -4 33 -,-- 1). 3) Keep P={-1} but inject at every four (ablation studies in Appendix D.3) layers. i... L- -0..4.... ] for Llama2-13B. 4) Set IP = (- -4, -2,22111 and L = {0.4.... 28, 32) simultaneously. 5) Keep P and IL fixed but change the zero-shot prompt to an 8-shot ICL prompt. We compare our LTV to Vanilla TV and FV in all five settings", with implementation details for FV in Appendix H.
The results in Table 2 demonstrate the advantages of our method over the baseline and highlight the flexibility and scalability of TV's in general. The LTV takes effect at arbitrary positions and is not confined the last token. Multiple LTVs can be injected dif- ferent positions or layers with performance benefits, and injecting into ICL prompts can further improve accuracy (baseline ICL accuracy is slightly < 80%, as in Figure 2). By contrast, TVs distilled from ICL hidden states are sensitive to injection position, do not improve ICL accuracy, and fail to synergize when injected at multiple locations. The only exception is the "More layers & Pos." setting, where both meth- ods fall behind baseline. Closer examination suggests that injecting at many layers and positions simulta- neously does not help artificial Capital, Capitalize, and Antonym tasks-likely because their simplicity makes heavy TV injection prone to overfitting.
(A)
Mean * SD
10
7.78
8.02
5,24
(B)
Q: can give you reward 0.1 points now,
or 1 point in 18 days. which would you profor?
ICL Response: 0.1 point now.
Response: 0.1 point now.
Vanilla TV/FV Response: point days.
Figure 3: (A) Mean and standard deviation of ratings for responses generated with ICL,
LTV, FV, and Vanilla TV. (B) An example question and responses across settings.
Adaptability of LTVs to complex task settings The tasks above have single-token labels and unique correct answers (e.g., The capital of China is + Beijing). To evaluate generalizability to a more complex generation task with multi-token responses-where the goal is to elicit a behavioral mode rather than a single answer-we adopt the Myopic dataset from the LLM steering literature (Panickssery et al., 2024; Bayat et al.. 2025). Each entry presents a question with two choices (see Figure 3 (B)), one myopic and the other favoring long- terg rewards. We compare the generated answers with LTV (injected at the middle layer) to the two baselines by asking an I.I. to rate them on 10-point scale (details in Appendix 1) based on how well they reflect the myopic choice. The statisties Figure 3 show that the LTV not only surpasses the baselines but also exceeds ICL
Prompts with fewer than 4 tokens are skipped in the accuracy calculation.
4In the baseline case, the FV method adds the sum of head outputs at the last position to the final token's 's hidden state. For varied F. we add summed outputs at each position in IP to the corresponding hidden state. For multiple layers, replicate the FVs [I_] times and inject copy at each layer. For the Vanilla TV, we patch hidden states positions P and TL of an ICI. prompt with different query into those of azg.
5

---

Published as a conference paper at ICLR 2026
Table LTV outperforms Vanilla TV and FV only in the baseline case but also across diverse configurations with varied positions, layers, and prompt formats. See Appendix D.2 for other models.
Baseline
1) Diff. Pos.
2) More Pos.
3) More layers
4) More layers & Pos. P- [-5..o] E- (0.4...)
Method
5) ICL prompts
P-1-1 (10)
P= [4)
L= 10.4,8....
Vanilla TV
37.80%
2.16%
19.18%
17.97%
18.15%
56.12%
FV
37.30%
2.68%
6.05%
31.8%%
0.38%
74.78%
LTV (Ours)
83.49% oas
18.399 pp
82.44% 161 ww
86.43%
51.39%
84.61%%ro.77
performance-ssmeenn something existing TV methods distilled from ICL representations struggle to achieve. These results provide clear evidence of the potential of LTVss in complex generation settings (see Appendix D.5 for other models).
Cross-task TV similarity and generalizability As first step toward understanding how TVs capture task id- iosyncrasies, compute cosine similarities among LTVs trained for different tasks (and across LTVs trained for the same task). For each of the seven tasks, train middle-layer LTV five times, yielding a total of 5 35 LTVs, and compute pairwise cosine similarities among them. The results in Figure 4 (with additional models in Appendix D.6) show that LTVs consistently encode stable task representations, exhibiting strong intra-task alignment and clear inter-task separation.
1.0
S5T2-
TREC-
Cosine Similarit
Capital
0.8
Capitaltze
02
Antonym
0.0
85T.2
TREC
SNL
RTE
Capital
Capttalize
Antomy
Figure 4: Cosine similarity heatmap of
LTT'ss for seven tasks, showing inter-task separation and intra-task clustering.
A notable exception is the moderate alignment between LTV' trained on SNLI and RTE, which share the binary label space (true, false). This suggests that the orientation of a TV is strongly governed by the unembedding directions of its task labels: when two tasks ultimately promote alignment toward the same label unembeddings, their corresponding LTVs naturally become similar. To verify this hypothesis, we apply the middle-layer LTV trained on SNLI to other datasets and measure the resulting accuracy. As shown in Table 3, the SNLI LTV yields nontrivial performance only RTE. In contrast, when applying an LTV trained on the Capital kataset-which does not share label semantics with other tasks-we observe no meaningful generalization (see Appendix D.7). Together, these results indicate that cross-task generalization of task vectors is fundamentally constrained by overlap in the task label space.
Compositionality of TV Given the strong dependence of TVs on the task label space, we ask whether LTVs trained on tasks with
Table 3: Applying the SNLI to other tasks
SST-2 TREC RTE
Capital Capitalize Antonym
related label spaces exhibit word2vec-style compositionality (Mikolov et al., 2013). We consider three tasks: English- actress), and English LTVs for the first two tasks, sum the resulting vectors, and evaluate the composed vector as TV on the third task. As shown in Figure 5, the composed LTV achieves accuracy substantially above both zero-shot and ICL baselines, demonstrating that task vectors can compose in a semantically meaningful way that mirrors the structure of their underlying label spaces.
0.00% 0.00% 46.21%
1.30%
0.67%
0.00%
Masculine-FFreme do. - chien), Masculine-FFmmnnin (e.g., actor -> Feminine (e.g.. actor - actrice). We train middle-layer
4.2 LOW-LEVEL INTERACTIONS BETWEEN TV AND ATTENTION HEADS
After demonstrating the superiority of our approach over inefficient methods of extracting TVs, we next address the second gap in prior TV studies: the lack of exposition of the mechanism behind TV effectiveness. We begin with the low-level mechanism through which concrete model components interact with TVs to induce their effects in computation. We focus attention heads given their cen- trality in Transformer based LLMs (Zheng et al., 2024), and their well-documented significance for model perfor-
75%
70%
Accuracy
50%
50%
23%
0%
Composition
Figure 5: Results of composing the
LTVs on subordinate tasks on the com- posite task.

---

Published as a conference paper at ICLR 2026
833%
75%
8.3%
78%
52%
Acmraty
75%
50%
50%
51%
23%
0%
Learned Reconstructed through
0%
Leamed
Heads Ablated
Heads Ablated
(a) Reconstructing TV effect through OV circuits.
(b) Ablating key attention heads
Figure 6: Assessing the significance of attention heads in the low-level interactions between TVs and model components. (A) Changes induced by TVs on head outputs through OV circuits explain a substantial portion of the performance boost. (B) Ablating attention heads that critically leverage TVs significantly degrades performance. Results for other models in Appendix F.2
mance (Yang al., 2025a; Cho et al., 2025; Jin al., 2024) and behaviors (McDougall al., 2023; Song et al., 2025) across diverse settings.
Reconstructing TV effect through OV circuits In Section 3, we showed that the output of an attention head k1,k) to the final token hidden state (which directly determines the output) can be expressed d = - j.N W Whe When TV 0 is injected at the last position of layer I 1-1 1, the corresponding hidden state becomes h + 0, and the attention head output becomes:
N-1
AN
Lk W Σ
OK W. W
W + CNN Oi We
L.k w W (hly + +0),
with an additional component CN 0&k
I.A. W W O. Since CN,y is a scalar attention weight and consider- ing the effect of layer normalization, the term 0, W 0-.ieee. the TV transformed by the head's OV circuit-aa is the core factor reflecting the effect of the TV on the head's contribution to the residual stream (Figure 1 (B)). Because residual connections (He al., 2015) carry 0 forward, it influences all heads in layer and beyond. Thus, the aggregate influence on head outputs caused by the TV is:
Σ (PP)12141
W ON W
(5)
which has similar form to FV. To test whether interactions between O and attention heads constitute the main low-level pathway, we inject this aggregate back as a packaged TV into the residual stream at layer /- reconstruct the aggregate effect of the original TV expressed and propa- gated through OV circuits of all heads on the residual stream. We provide further clarifications in Appendix F.1. We perform this experiment using middle-layer LTV, and rescale the recon- structed head-output vector to match the norm of the original LTV, ensuring that the intervention does not introduce out-of-distribution shifts the hidden states. After injection, we also add to the final-layer hidden state reinstate its purely residual effect 5. The results in Figure 6a confirm the critical role of attention heads and their OV circuits: reconstructing the TV effect via OV-transformed decompositions restores much of the performance gain, showing that TV's steer the residual stream largely through channels modulated by attention heads. We further establish the significance of the OV circuits in expressing and modulating the TV's effect by experimenting with the alternative of MLP-based reconstruction in Appendix F.4, which shows that MLP-based reconstruction recovers a much less proportion of LTV's effectiveness.
Assessing key attention heads leveraging the TV We further evaluate attention heads by identifying those most reliant on TVs for predictions and examining the effect of ablating them (setting outputs to O). We compute saliency score (Bansal et al., 2022; Michel et al., 2019; Molchanov et al., 2016) for each head in the presence of a TV. Let a/ be head (1,k)'s 's output to the last position with the TV @p(u, |ptsu/12..00..... (.L.P. ) estimating the influence of the head output on dak
injected; its saliency score is |a/=||
the correct label probability via first-order Taylor approximation. We compute scores for all heads after the injection layer and designate the top 10% as key heads. We then ablate these and randomly ablate 10% of heads as control. The results in Figure 6b support the saliency-base identification: ablating key heads reduces performance far more than random ablations, confirming attention heads' central role in realizing TV-driven gains, compared with direct residual bias of O.
*We confirm this residual effect plays inconsequential role in the observed accuracy gain through reconstructing the OV effect, see Appendix F.3.
7

---

Published as a conference paper at ICLR 2026
Percentage
20.0%
Percentage
Heads
50%
Random Heads
10.0%
0%
16 24
32 36
Layer
Position
(a) Distribution of key heads across layers
(b) Distribution of attention weights across positions
Figure 7: (A) Key attention heads cluster mainly in layers immediately after the injection (16 for LIama3. 1-8B) and secondarily in final layers. (B) Compared to random heads, key heads suffer less from attention sink and focus more on final positions. See Appendix F.5 for other models.
Characterization of the key attention heads After identifying key heads, we further analyze their seeeatrristiss-ppcciiiieeeee their distribution across layers and attention weights over token positions. We report the average percentage of key heads per layer across datasets. For attention distribution, show average patterns of all identified heads over input positions prompts see Appendix F.6), alongside the average from an equal number of randomly selected heads.
SST-2 prompt (for more
The results in Figure 7 show two main patterns. First, key heads leveraging TVs follow quasi- U- shaped distribution: many appear right after the injection layer (serving as early gateways for TV influence) and again in final layers (integrating TVV effects into outputs). Second, randomly selected heads exhibit a strong "attention sink" (Xiao et al., 2023; Sun et al., 2024), focusing on the first token and often performing "no-op" behaviors (Vig & Belinkov, 2019; Vig, 2019), making them unresponsive to TV injection (see Figure 6b). By contrast, key heads show weaker sink and greater focus on final positions, enabling them exploit TVs when shaping outputs (Figure 6a).
4.3 HIGH-LEVEL ANALYSIS OF TV's INFLUENCE MECHANISM
The previous section demonstrated that TVs realized primarily through attention- head OV circuits, with small subset of heads driving most of the effect. We now move from these local interactions to the higher-level question: how do TVss evolve they propagate through the network and ultimately shape predictions? To answer this, we analyze the layer-wise dynamics of hidden states after TV injection which reflects how the injection effect propagates (Skean al., 2025; Kirsanov et al., 2025; Yang et al., 2025a) using the SST-2 dataset that offers clear mechanistic insights (Yang et al., 2025a). We track three complementary metrics across layers of TV influence (Figure (A)):
1. Logit Lens Accuracy (nostalgebraist, 2020): decode hidden states at intermediate layers with the unembedding matrix Wu and compute accuracy. This global metric indicates whether the inference dynamics driven by the TV are able to yield correct predictions at given depth.
2. Logit Difference: the logit gap between correct and incorrect labels, e.g., positive vs. negative for SST-2. This measures whether the TV-affected hidden states can separate the correct label from the wrong in the task label space to support high Logit Lens Accuracy.
3. Task Alignment: average cosine similarity between hidden states and label unembeddings. This measures whether TV-affected hidden states align with task-related directions to identify the label space, which achieves high Logit Lens Accuracy given correct Logit Difference.
Given the different effects of TVs injected early vs. late layers noted in Subsection 4.1, we compute these metrics for [HO' - (i)...... ). i.e., the collections of last-token hidden states across layers when a TV 0, is injected an early or late layer I. We set / 4 for early and 3L for late (8 and 24 for 1-8m) We compare these TV-affected hidden states with ICI. and zero-shot baselines.
Early vs. late TVs shape hidden states differently From Figure 8 (A), both early- and late-layer TVss nudge zero-shot hidden states toward ICL trajectories in metric trends, indicating that LTVS capture the essence of ICL. Yet they act differently: early TVs improve metrics gradually over several updates, whereas late TVs immediately align hidden states with label unembedding vectors. This aligns with Figure 8 (B), where decoding TVs directly with Wy shows early TVs yield irrelevant tokens while late TVs produce tesk-related tokens iiiens-implying stronger alignment with task directions and direct steering of hidden states to increase label logits. These differences motivate a closer look at how early late TV effects propagate through intermediate updates.
8

---

Published as a conference paper at ICLR 2026
(A)
Early Laver TV: Logit
Earty Dur
Farh Laver Taok Alignment
Valle
ICL.
Layer
Layer
Layer
Late Layer TV: Logit Leos
Late Layer TV Logit pirt
0.2
Late Layer TV: Task Aligament
50%
Halle
0%
(B)
Layer
Layer
Layer
Tokens Decoded from LTVs:
Early Layer:
classical, Classical, nghiēp, brushes, /Forms, enant, irm,
Late Layer:
positive, negative, -positive, positives, Negative, Positive,...
Figure 8: (A) Metric values of hidden states across layers when the TV is injected early late
layer. (B) Tokens decoded from TVs, with early-layer TVs yielding random tokens and late-layer
TVs producing task-related tokens. See Appendix G.1 for other models' results.
100%
100%
Logit Lens Acc.
Accuracy
75%
75%
50%
50%
25%
25%
Reconstructed 6
W. W as,
0%
0%
12
28
Layer
Layer
(a) Effect of linearly reconstructed TV. For other mod- els see Appendix G.3.
(b) Using linear transformation replace layer up- dates of hidden states.
Figure 9: (A) A reconstructed TV based on modeling 0j's influence as linear achieves comparable accuracy for most layers. (B) Characterizing sidden-state updates with TVs linear yields positive
results: the fitted transformation matrix substantially increases intermediate-lan -laye decoding accuracy.
Linear propagation of TV's effect To analyze how a TV's effect is transmitted to final-layer hidden states, note that we have the abstraction of the update from to
H' Layer_
H/ = Layee__Ud
1,0,))
(6)
where H are zero-shot hidden states at the final layer L, and multiplying by 1, adds the TV each of the examples. Given ample evidence linear mechanisms Transformers (Marks Tegmark, 2024: Park et al., 2024), we hypothesize that if the composite update acts linearly on 0,, then
1, (W
-H"
for some Wry,0) E Rdxd parameterizing the linear effect of hidden states update. The resulting effect of TV on label logits is WuW.rv 0,, and on task labels wer.s.ee WTv. O (inner products with rows of We Wrv.n) for "positive") neggttvve"""" To test this hypothesis, we proceed as follows:
(1) Collect states with noise injection: Using LTV 0, sample prompts, obtain Hy (with injection) and H (without). We perturb 0 0,,- -2E while obtaining Hi) avoid degenerate rank-1 1 solutions when fitting W WV, since 1, rank-1.
(2) Construct and evaluate proxy TV: We compute W WTv.n + Wh Wry.() as a proxy TV, rescale it to match 0,'s norm, and inject it at layer 1. This vector should have high inner products with Wi pe (r"""ree WIv.i) and W Wrv.( should raise both label logits to support correct prediction if the hypothesis is correct. We test this 0 all layers I. See Appendix J for the full details of the fitting and reconstruction procedure, where we also the provide the theoretical guarantee for our method and experimental results validating the theorem.
The results in Figure 9a support the linear hypothesis: the linearly reconstructed TV matches the original TV's performance for most layers, with only few exceptions. This indicates that purely linear operator Wyv.n) can almost fully capture the channel linking injected TVs at different layers to changes in final-layer hidden states, despite the many nonlinear components within the model.

---

Published as a conference paper at ICLR 2026
Linearity of hidden-state updates The strong linearity of Layer Update TVs suggests that hidden-state updates may also be summarized linearly. To verify this, we fit WHs,0) such that Hi W Ws,, Hi) on sample (details in Appendix J), where H'() are layer-/ hidden states with O injected. We then multiply WHs.in with a separate set of (n) and check if decoding with Wt yields higher accuracy than direct decoding, which is confirmed in Figure 9h and signals the strong linearity of hidden-state updates. These results align with prior evidence of LLM layer linearity (Razzhigaev et al., 2024) and the success of attempts to linearize Transformers (Li et al., 2020; Han et al., 2024).
Decomposition of TV's influence mechanism Although TVs injected different layers are uni- formly transformed into final output changes via a linear map Wry,n) this shared linear mechanism does not eliminate meaningful layer-dependent differences, as evidenced by the contrasting behav- iors of early- and late-layer TVss in Figure 8. Finer-grained analysis SHOULD be conducted to explain why the effects of early and late TVs differ as in Figure 8. To this end, we consider the polar decomposition of the transformation matrix as Wrv.0) = (n) En) where orthonormal Qi)) represents a rotation and positive semidefinite Σnn a stretch along the right-singular directions of (r. Since Figure (B) shows early-layer TVs are aligned with directions unrelated to the task, we apply only the rotation to 0 at different layers and measure changes in task alignment. This addresses whether early-layer TVs operate via distinct mechanism, or are rotated by subsequent layers to align with task label unembeddings to increase logits as late-layer TVs do. The substantial increases in task alignment in Figure 10 (A), especially for early layers, indicate common mechanism: TVs steer hidden states toward task-related directions (Figure 1 (C)). The fact that early-layer TVss decode task-relevant tokens after rotation (Figure 10 (B) supports this view. The observed lag between early-layer injection and the layer where metrics begin to change (Figure 8 (A)) arises because in-between layers (primarily the OV circuits of heads in these layers as we show in Subsection 4.2) are needed to rotate the TV toward task-related directions. Thus, Figure 10 provides a unified account linking TVs at different layers to final outputs.
Rotation phases out, stretch phases in To further understand how rota- tion and stretch evolve across layers, we compute the cosine similarity be- tween O and 11) 0, This quantifies rotation strength: higher similarity im- plies less rotation as the matrix ap- proximates identity mapping. The rising similarity across layers in Fig- ure 10 (C) reveals clear trend of di- minishing rotation in deeper layers, with stretch becoming the dominant component of Wrv,())- This suggests that early-layer TV's undergo stronger cotation-consistee with the finding that intermediate layers are needed to rotate TVss toward task-related directions.
(A)
(C)
value
allgammnt
Value
0.50
COSSIML QuAa
0.25
0.00
12
26
20 Layer
24
Layer
(B)
Tokens Decoded froe Early Layer LTVs:
Before Botation:
classical, Classical, nghiOp, brushes, /fores, onant,...
After Rotation: positive, negative, adverse, faverable, unfavorable,...
Figure 10: (A) Applying the rotation to TVs at different layers substantially increases alignment with unembeddings of task-related labels. (B) After rotation, early-layer TVs that originally decode random tokens produce task-related tokens. (C) The rotation effect diminishes for late-layer TVs as the estimated matrix approaches identity.
5 CONCLUSION
We revisited task vectors as mechanistic explanations for in-context learning. Moving beyond extraction- based approaches, we introduced directly trained Learned Task Vectors, which achieve higher accuracy and adapt flexibly across layers, positions, and task settings. Our analysis showed that TVs at the low level operate mainly through attention their effect. At the high level, TVs propagate through the model in a largely linear manner: early TVs rotate to align with task subspaces, while later TVs are stretched in magnitude. This sotation-stretch dynamic offers unified account of how TVs different depths shape final predictions. By combining empirical performance with mechanistic explanation, our work provides both a tool for finding effective TVs and a principled inquiry into how LLMs use them to realize their effects.
OV circuits, with a few key heads driving
10

---

Published as a conference paper at ICLR 2026
REFERENCES
01. AI, Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang. Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, Kaidong Yu, Peng Liu, Qiang Liu, Shawn Yue, Senbin Yang, Shiming Yang. Tao Yu, Wen Xie, Wenhao Huang, Xiaohui Hu, Xiaoyi Ren, Xinyao Niu, Pengcheng Nie, Yuchi Xu, Yudong Liu, Yue Wang, Yuxuan Cai, Zhenyu Gu, Zhiyuan Liu, and Zonghong Dai. Yi: Open foundation models by 01.ai, 2024.
Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretability-aassee case study at 66 billion scale. arXiv preprint arXiv:2212.0900 2022.
Reza Bayat, Ali Kalaii--alaroudii, Mohammad Pezeshki, Sarath Chandar, and Pascal Vincent. Steering large language model activations in sparse spaces. arXiv preprint wrXiv:2503.00177 2025.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, al. Language models are few-shot learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems, pp. 1877-1901, Online and Vancouver, Canada, 2020. URL https: / /dl aacm org/ /oii/ / 10. N0.5555//495524..
Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Qingfu Zhang, Hau-San Wong, and Taiji Suzuki. Provable in mn-oonext vector arithmetic via retrieving task concepts. In Forty-second International Conference on Machine Learning, 2025.
Wang Cai, Hsiu-Yuan Huang, Zhixiang Wang, and Yunfang Wu. Beyond demonstrations: Dynamic vector construction from latent representations. arXiv preprint arXiv:2505.203/8, 2025.
Emmanuel J Candes, Justin K Romberg, and Terence Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 1908):1207-1222. 2006.
Yuanpu Cao, Tianrong Zhang, Bochuan Cao, Ziyi Yin, Lu Lin, Fenglong Ma, and Jinghui Chen. Personalized steering of large language models: Versatile steering vectors through bi-directional preference optimization. Advances in Neural Information Processing Systems, 37:49519-49551, 2024.
Hakaze Cho, Mariko Kato, Yoshihiro Sakai, and Naoya Inoue. Revisiting in-context learning inference circuit in large language models. In The Thirteemth International Conference Learning Representations, 2025. URL https: / openreview net / forum d=x YpnYNvQq.
Róbert Csordás, Christopher D. Manning, and Christopher Potts. Do language models use their depth efficiently?, 2025. URL https:t///// /tt::///axxxi www.rriwvv..ggeeeee abs / 220. 13898.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine learning challenges workshop, pp. 177-190. Springer, 2005. URL https: weiink.springe
com/ wttptttett...111111111
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning, 2024. URL https: : / wtwiv.vvvvv.o.m org/ abs/ 2301.00000000 00234.
Yuxin Dong, Jiachen Jiang, Zhihui Zhu, and Xia Ning. Understanding task vectors in in-context learning: Emergence, functionality, and limitations. arXiv preprint arXiv:2506. 09048, 2025.
Jacob Dunefsky and Arman Cohan. One-shot optimized steering vectors mediate salety-relevant behaviors in Ilms, 2025. URL https:
wtwp/aa/////ooooaa//o////eaaae/a abw./b/////22 18862.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, and Abhinav Pandey et al. The llama 3 herd of models, 2024. URL https: // wwtp//iivvv..///e org/ abs/2407 010077..78......
11

---

Published as a conference paper at ICLR 2026
Dongchen Han, Tianzhu Ye, Yizeng Han, Zhuofan Xia, Siyuan Pan, Pengfei Wan, Shiji Song, and Gao Huang. Agent attention: On the integration of softmax and linear attention. In European conference on computer vision, pp. 124-140. Springer, 2024.
Seungwook Han, Jinyeop Song, Jeff Gore, and Pulkit Agrawal. Emergence and effectiveness of task vectors in in-context learning: An encoder decoder perspective, 2025. URL https: //arxiv.... wtt.aii...gg/a////1122.12.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL pecttp::/aaraee www.avs.rg/bbaao.. /wwts.111.....0
Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Compu- fational Linguistics: EMNLP 2023, pp. 9318 3-33333,, Singapore, December 2023. Association for Computational Linguistics. doi: 10.126533 findings-emnlp.22.. URL https:
//aclanthologyy org/2023333/333 findinggnnnniii emllpp 624/.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
09685.
Gangwei Jiang, Caigao Jiang. Zhaoyi Li, Siqiao Xue, Jun Zhou, Lingi Song, Defu Lian, and Ying Wei. Unlocking the power of function vectors for characterizing and mitigating catastrophic forgetting in continual instruction tuning. 2025a. URL https: : //trxiva.iv..////////a/.aaaaaaia org/ abs 2502. 11019.
Jiachen Jiang, Yuxin Dong, Jinxin Zhou, and Zhihui Zhu. From compression to expansion: A layerwise analysis of in-context learning, 2025b. URL 17322.
Zhuoran Jin, Pengfei Cao, Hongbang Yuan, Yubo Chen, Jiexin Xu, Huaijun Li, Xiaojian Jiang, Kang Liu, and Jun Zhao. Cutting off the head ends the conflict: mechanism for interpreting and mitigating knowledge conflicts in language models, 2024. URL btttps://axiv..oo/ org/ abs/ 2402. 18154.
Patrick Kahardipraja, Reduan Achtibat, Thomas Wiegand, Wojeiech Samek, and Sebastian La- puschkin. The atlas of in-context learning: How attention heads shape in-context retrieval augmen- tation. arXiv preprint 25X.v:2505.588) 2025.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL 111.o//98..6.
https:
Artem Kirsanov, Chi-Ning Chou, Kyunghyun Cho, and Sue Chung. The geometry of prompt- ing:
Unveiling distinct mechanisms of task adaptation in language models. In Luis Chiruzzo, Alan Ritter, and Lu Wang (eds.), Findings of the Association for Computational Linguistics: NAACL 2025, pp. 1855-18888 1888, Albuquerque. New Mexico, April 2025. Association for Computa- tional Linguistics. ISBN 979-8-89 8-888888897999981 URL URUAttts://aclatthoo orgg2222 findings-naace naannggs-naaal..mm 100/.
Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld, and Rada Mihalcea. A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity, 2024. URL https: http///aivvv//a////////ee/ea.ava org/ abs / /4002241.000 01967.
Dongfang Li, Zhenyu Liu, Xinshuo Hu, Zetian Sun, Baotian Hu, and Min Zhang. In-context learning state vector with inner and momentum optimization, 2024a. httpstts:: //tp/i//avvvv/a//a org/ abs/ 2404. 11225.
Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Inference- time intervention: Eliciting truthful answers from a language model, 2024b. URL https:
//arxivv.cc.. wttpvv...rgabb//200ee
03341.
Rui Li, Jianlin Su, Chenxi Duan, and Shunyi Zheng. Linear attention mechanism: An efficient attention for semantic segmentation. arXiv preprint arXiv:2007. 1400:20077.492,, 2020.
12

---

Published as a conference paper at ICLR 2026
Xiang Lisa Li and Percy Liang. Prefix-tuning Optimizing continuous prompts for generation, 2021.
URL https:://// http///xxiv.ooo abs aabs/2100 22100 00190.
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Confer- ence on Computational Linguistics, 2002. URL http::::wwww.ccll........ org/ anthology/
C02-1150.
Zhuowei Li, Zihao Xu, Ligong Han, Yunhe Gao, Song Wen, Di Liu, Hao Wang, and Dim- itris N. Metaxas. Implicit in-context learning, 2025. URL https: http:///axxvvv.../abbaoeo 14660.
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhenghao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing, 2021. URL https:://///......t....t //t////aiiivv org/ aaaaabss21100 13586.
Sheng Liu, Haotian Ye, Lei Xing, and James Zou. In-context vectors: Making in context learning more effective and controllable through latent space steering, 2024. URL https: /arxiv.
org/abs 2311- 06668.
lly Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Bill MacCartney and Christopher D. Manning. Modeling semantic containment and exclusion in natural language inference. In Donia Scott and Hans Uszkoreit (eds.), Proceedings of the 22nd International Conference on Computational Linguistics (Coling 2008), pp. 521-528, Manchester, UK, August 2008. Coling 2008 Organizing Committee. URL https: / aal.antcnnoooo..... org/c08--066/.
Samuel Marks and Max Tegmark. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets, 2024. URL htt.t:/tp//// htthtts::/ http::://aaxxv...ot abs 2310. 06824.
Callum McDougall, Arthur Conmy, Cody Rushing, Thomas McGrath, and Neel Nanda. Copy suppression: Comprehensively understanding an attention head, 2023, URL https: //arxiv. org/ brg/abbs23100. 04 625.
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vec- style vector arithmetic, 2024. URL ptcttps://axiv..oo/ /bs///////// /205 16130.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representa- tions in vector space, 2013. URL httph:tt:///// //tr///aaiiviii//i org/ aa////b//13300 3781.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint 1rX1v:1611.06440, 2016.
nostalgebraist, Interpreting gpt: the logit lens, 2020. htt.tttttt//...eeettete
com/
Nina Panickssery, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, and Alexander Matt Turner. Steering llaaa via contrastive activation addition, 2024. URthttps://axxii..omm abs abs/2312 06681.
Kiho Park, Yo Joong Choe, and Victor Veitch. The linear representationn hypothesis and the geometry of large language models, 2024. URL hitps: http://aaiiiiii/ii http:///wwwwwooo///////111111111 03658.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask leamers. OpenAl blog, 2019. URL htttps://dmmucfpkee
http::kkwwwww....oo......oo cww.ddronttttto net wwwee//betteee-laageee-teee language-modes s. pdf.
13

---

Published as a conference paper at ICLR 2026
Anton Razzhigaev, Matvey Mikhalchuk, Elizaveta Goncharova, Nikolai Gerasimenko, Ivan Oseledets, Denis Dimitrov, and Andrey Kuznetsov. Your transformer is secretly linear. arXiv preprint 22X.i:2400..2250. 2024.
Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. Identifying semantic induction heads understand in-context learning, 2024. URL https:
// //arxii..org./.................a abcg://bs//a /tt..////2.... 13055.
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv. Layer by layer: Uncovering hidden representations in language models. arXiv preprint arXiv:2502.02011. 2025.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng. and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In David Yarowsky, Timothy Baldwin, Anna Korhonen, Karen Liveseu, and Steven Bethard (eds.), Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp- 1631-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL wwwph:tpp////mm http::///aaaatttti /tt////aaaaatlll.lrrrgr/////e//e
Jiajun Song, Zhuoyan Xu, and Yiqiao Zhong. Out-of-distribution generalization via composition: a lens through induction heads in transformers. Proceedings of the National Academy of Sciences,
12216)::2417182122, 2025. URL https: //t////axvvvv.///a////a//..a/aaa org/ hbs//ab//aa 2408 09503.
Nishant Subramani, Nivedita Suresh, and Matthew E. Peters. Extracting latent steering vectors from pretrained language models, 2022. URL https: / aaxiiivvv www//bggg/s........ 22052 051 24.
Mingjie Sun, Xinlei Chen, J. Zico Kolter, and Zhuang Liu. Massive activations in large language models, 2024. URL https: wwwia.xiivv..o/ wtg. aww/aas///00 17762.
Eric Todd, Millicent I.. Li, Arnab Sen Sharma, Aaron Mueller, Byron C. Wallace, and David Bau. Function vectors in large language models, 2024. URL https: 15213.
/2310
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977, 2023. URL https: aa//aaxxvvv...///aaa/aaaaa org. /2302. 13971.
Alexander Matt Turner, Lisa Thiergart, Gavin Leech, David Udell, Juan J. Vazquez, Ulisse Mini, and Monte MacDiarmid. Steering language models with activation engineering, 2024. URL https: / arxiv. org/ abbs/////8 10248.
Jesse Vig. A multiscale visualization of attention in the transformer model. arXiv preprint arXiv: arXiv:/906.00) 11900005744 14, 2019.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model, 2019. URL https: org/ abs 04284.
Zhengxuan Wa, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher D. Manning, and Christopher Potts. Reft: Representation finetuning for language models, 2024. URL https: //a.xiv.org/abs....... 2404 03592.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprini arXiv:2309. 17453, 2023.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li. Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang. Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang. Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprint ar0iv:2407. 10671, 2024.
14

---

Published as a conference paper at ICLR 2026
Haolin Yang. Hakaze Cho, Yiqiao Zhong, and Naoya Inoue. Unifying attention heads and task vectors via hidden state geometry in in-context learning, 2025a. htlltttss//axii/// abs/2505. 18752.
Liu Yang, Ziqian Lin, Kangwook Lee, Dimitris Papailiopoulos, and Robert Nowak. Task vectors in in-context learning: Emergence, formation, and benefit. arXiv preprint arXiv:2501.09240, 2025b.
Kayo Yin and Jacob Steinhardt. Which attention heads matter for in-context learning? arXiv preprint arXiv:2502.14010 2025.
Li-Ming Zhan, Bo Liu, Zexin Lu, Chengqiang Xie, Jiannong Cao, and Xiao-Ming Wu. Deal: Disentangling transformer head activations for Ilm steering, 2025. URL https:///axxtit //arxiv. org/abs/2500 2506. 08359.
Haiyan Zhao, Heng Zhao, Bo Shen, Ali Payani, Fan Yang. and Mengnan Du. Beyond single concept vector: Modeling concept subspace in Ilms with gaussian distribution, 2025. URL
https: / arxiv..... org/abs/
abs/ 2410. 00153.
Zifan Zheng, Yezhaohui Wang. Yuxin Huang, Shichao Song, Mingehuan Yang, Bo Tang, Feiyu Xiong, and Zhiyu Li. Attention heads of large language models: A survey. arXiv preprint arXiv:2409.03752, 2024.
15

---

Published as a conference paper at ICLR 2026
Appendices
A STATEMENT OF LLM USAGE
In this work, LLMs are used to help with writing, experiment coding, and visualization of the results. LLMs are also used to produce results in one of the experiments, as explained in Subsection 4.1 and Appendix I.
B DETAILED PROCEDURES OF TRAINING LEARNED TASK VECTORS
As described the main text, we train 0 by minimizing the loss log poy F. P,L.P), Optimization is performed with Adam (Loshchilov & Hutter, 2017) using learning rate of 0.001 and weight decay of 0.01. Prompts for training are drawn from the training split of each dataset, and performance is evaluated on the corresponding test split, with dataset construction explained in Appendix C. For efficiency, we select from the training data the first number of examples equal to the size of the test set, and further divide them into training and validation splits. For example, the Antonym dataset contains 600 training and 400 test samples; we take the first 400 training samples and split them into 240 for training and 160 for validation. Training runs for up to 10 epochs, with 100 examples randomly sampled from the training split per epoch (or the entire split if contains fewer than 100 samples). Early stopping with patience of is applied: if validation performance does not improve for two consecutive epochs, training halts and the A that achieved the hest validation accuracy is retained as the final TV. the setting of Subsection 4.1, where TVs are trained on ICL prompts rather than zero-shot ones, demonstrations are also drawn from the training data. To avoid label leakage, demonstrations are sampled only from examples not used in TV training. For instance, in the Antonym dataset, where 400 of 600 training samples are used for TV training, the remaining 200 are reserved for demonstration construction.
C IMPLEMENTATION DETAILS
Models We use the official HuggingFace implementations all models. Models with more than in half precision.
10B parameters quantized 4-bit precision, while smaller models
Datasets We use the official HuggingFace implementations of SST-2, SNLI, RTE, and TREC. For Capital, Capitalize, Antonym, and Myopic, use the data released by previous authors. Specifically, the data for Capital, Capitalize, and Antonym are taken from Todd et al. (2024), and the data for Myopic from Panickssery al. (2024).
ICL and evaluation settings any principled selection methods. For SST-2, TREC, SNLI, and RTE, we use the training set both for demonstration selection and for training task vectors, and we evaluate performance on the test set (or the validation set if ground-truff test labels are unavailable). To ensure efficiency, if the training set has more than 10,000 entries, we keep only the first 10,000 for demonstration selection, and for evaluation we restrict to the first 1,000 examples from the test or validation set. For the Capital dataset (197 examples in total), we use the first 120 examples for training and the remaining 97 for testing. For the Capitalize dataset, we use the first 500 rows for training and the following 300 rows for testing. Similarly, for Antonym testing. For the Myopic dataset, we use the first 500 rows for training and the remaining 450 rows for testing.
We select demonstrations randomly for each query without relying on
the first 600 rows for training and the next 400 rows for
Devices All experiments are conducted on an H200 GPU.
D SUPPLEMENTARY MATERIALS FOR SUBSECTION 4. 1
D.1 PERFORMANCE OF LTV INJECTED AT THE LAST POSITION ON OTHER MODELS
In Subsection 4.1, we reported the performance of our LTV for Llama2-7B and Llama2- under the traditional setting following Hendel et al. (2023) and Todd al. (2024), i.e., injecting one
16

---

Published as a conference paper at ICLR 2026
specific layer into the last position. In Figures 15-16, provide similar layer-sweeping results of LTV performance for Llama2-7B, Llama2- 13B, Llama3-8B. and Llama3.2-3B. The results likewise demonstrate consistent performance advantage of the LTV over the two baselines across layers, with the gap being most prominent in later layers. In Table 8, we report the corresponding results for the remaining three non-Llama models. Concretely, we inject the T's layers corresponding to 50% of the total number of layers of each model (for instance, at layer 16 for 32-layer model). The results validate the performance of our LTVs across model sizes and architectures, as they consistently raise performance significantly above the zero-shot level and up to the level of ICL.
D.2 REPLICATION OF TABLE 2 FOR OTHER MODELS
In Tables 9-12, we present the comparison of FV, Vanilla TV, and LTV across the five scenarios on Llama2-7B, Llama2-13B, Llama3-8B, and Llama3.2-3B, which yields largely the same conclusions. Our LTV demonstrates strong flexibility with respect to injection positions and ICL prompts, as well scalability cases involving multiple positions and layers. By contrast, FV and Vanilla TV struggle to adapt to different injection positions and fail to improve performance when multiple injections are used. For the other models we report only the performance of the LTV, The results, shown in Tables 13-15, are consistent with those in Table 2. The reduced average performance of TVs when trained and injected at multiple layers and positions simultaneously is again observed, which we attribute to lower accuracy on the Capital Capitalize, and Antonym tasks.
D.3 ABLATION STUDIES FOR THE LAYER STRIDE WHEN INJECTING LTVS TO MULTIPLE LAYERS
In Table 2, we demonstrate the scalability of LTV by injecting it into every four layers of the model simultancously. In Table 16, we conduct ablation studies on the layer stride by injecting at every two layers or every eight layers. The results closely match those obtained with a stride of four, indicating that the scalability of LTV is unaffected by the specific choice of layer stride.
D.4 COMPARISON OF LTV AGAINST MORE BASELINES
In this section, we compare LTV with two additional methods that distill ICL hidden states into components to be injected into the zero-shot residual stream: State Vector (Li et al., 2024a) and I2CL (Li et al., 2025). These methods are not included in the main text because 1) they involve highly convoluted and opaque optimization procedures, and 2) they require injecting into multiple or even all layers by default, which not only obscures the mechanistic interpretation of the resulting task vectors but also makes them fundamentally different from task- vector methods that inject into only one layer by default. In Figure 17, we compare the performance of LTV with these two baselines by injecting into each single layer of Llama2-7B on SST-2. The results further corroborate the superiority of LTV, as it outperforms both baselines across all layers. We also record the total time required for the three methods to complete a full training- LTV requires the least amount of time despite involving gradient-based training, highlighting the substantial inefficiency of methods like State Vector and I2CL, which rely on highly convoluted optimization procedures to distill ICI. hidden states.
epoch on SST-2. As shown in Table 7,
D.5 REPLICATION OF FIGURE 3 FOR OTHER MODELS
In Figures 18-21, we present the comparison between Vanilla TV, FV, and LTV injected into the middle layer of Llama2-7B, Llama2-13B, Llama3-8B, and Llama3.2-3B on the Myopic dataset. The results closely echo those of Figure 3: the LTV consistently outperforms both baselines as well ICL across models, demonstrating its generalizability to complex generation tasks beyond single-token responses and the superiority of its performance uncapped by the representation quality of the ICL hidden states. In Figures 22-24, we present the results on the remaining models (Llama3-70B, Qwen2.5-32B, Yi), where compare Vanilla TV and LTV. The results largely similar.
17

---

Published as a conference paper at ICLR 2026
All Subsequent OV Circuits
Wy
Wo
Wr
Wo
negative
LTV
Remaining Layers
W
Wo
W
Wo o
Wv
Wo
+
0
W V
Wo
Layers
Bad movie. Sentiment
Rescale
Visualization of how we reconstruct the aggregate effect of TVs induced through the OV circuits of attention heads Subsection 4.2.
Figure 1I:
D.6 REPLICATION FIGURE 4 OTHER MODELS
In Figures 25-31, we provide visualizations of the experiments presented in Figure for additional models. The results indicate that the pattern of intra-task clustering and inter-task separation among LTVs common across models, though the strength of intra-task clustering varies, being stronger in Llama2-7B and Llama2-133 13B and more moderate in Llama3. 10ma3.1-81 1-8B and Llama3-8B. Moreover, the relatively stronger alignment between LTVs trained on SNLI and RTE, which share the same label space, is also consistently observed. This supports our claim in the main text that the direction of an LTV is closely correlated with the directions of the relevant unembedding vectors, which must align hidden states with to facilitate correct decoding.
D.7 REPLICATING TABLE USING THE LTV THE CAPITAL DATASET
In Table 17, we replicate the experiment presented in Table but using the LTV learned on the capital dataset. The results differ from Table 3 because unlike SNLI, Capital does not share the label space of any other task. This further corroborates the conclusion that the generalizability of LTV critically depends on the label space of the task.
E DETAILED PROCEDURES OF GENERATING CORRECTNESS RATINGS FOR THE MYOPIC DATASET
F SUPPLEMENTARY MATERIALS FOR SUBSECTION 4.2
F.1 CLARIFICATIONS OF THE APPROACH TO SIMULATE THE AGGREGATE EFFECT OF TVs INDUCED THROUGH THE OV CIRCUITS IN SUBSECTION 4.2
We provide visual explanation of how simulate the aggregate effect of TVs induced through the OV circuits of attention heads Figure 11 1. We first compute the products between the OV circuit of each attention head in layers after the point of injection and the injected TV 0, then sum these products and rescale them to match the norm of & before injecting this aggregate vector as a TV at the original injection site.
18

---

Published as a conference paper at ICLR 2026
Note that we inject the aggregate product of the TV and all OV circuits, i.e., Lirx): w as a whole back into the injection layer. This follows the procedure of previous TV studies, which attempt to construct TVs from attention head outputs (Todd et al., 2024; Li et al., 2024a), We also considered an alternative approach: instead of injecting the reconstructed TV as a whole, we first compute
0x vg O.
O L
w W o O
for each 1' 2 I. Then, at each layer I' from / to the final layer, we inject 00V into the residual stream. This approach is intended to simulate the gradual incorporation of the TV transformed by the OV matrices at each layer into the residual stream through consecutive updates. Empirically, we found this method achieves lower reconstructed accuracy than the one presented Figure 6. We believe the reason is the strong inconsistencies in hidden state scales across layers (Csordás et al., 2025), which make it much harder to adjust 0% to an appropriate scale. As a result, adding 07"h at every layer from / the final one risks shifting hidden states out of distribution, which greatly compromises the accuracy compared to the reconstruction approach in Subsection 4.2.
F.2 REPLICATION OF FIGURE 6 FOR OTHER MODELS
In Figures 32-38, we present results assessing the significance of attention heads in mediating the low-level interactions between TVs and model components. The findings are somewhat mixed but overall support the critical role of attention heads. Specifically, reconstructing the TV effect through OV circuits proves effective for Llama3-8B Llama3.2-3B. Qwen2.5-32B. and Yi-34B, but not for the other three models. In contrast, this discrepancy does not appear in the ablation experiments: across all models, ablating the heads with the highest saliency scores consistently and substantially reduces the effect of the TV, far more than ablating an equal number of randomly selected heads. In summary, the importance of attention heads for realizing the impact of TVs robust across models and architectures, though reconstructing TV effects by injecting summed OV transformations back into the stream appears more model model-dependent
F.3 EXAMINING THE INFLUENCE OF REINSTATING THE RESIDUAL EFFECT IN THE OV-BASED RECONSTRUCTION
In Subsection 4.2, in addition to injecting the summed product of the TV with the OV circuits of all affected heads as explained in Appendix F.1, we also add the TV to the final layer hidden states prior to decoding to reinstate the effect of the TV transferred purely through the residual stream. To test whether this residual effect is the main cause of the observed accuracy gain, which would otherwise invalidate OV circuits the dominant low-level channel, we repeat the OV reconstruction experiment from Subsection 4.2 but omit the final-layer TV addition. The results across models in Figures 40-47 show that including excluding the TV at the last layer has only an inconsequential impact, as accuracy remains practically unchanged.
F.4 DEMONSTRATING THE SIGNIFICANCE OF OV-BASED RECONSTRUCTION THROUGH MLP-BASED RECONSTRUCTION
In Figure 6a, gap between the accuracy achieved by the original LTV and the OV-based recon- struction can be seen, raising the question of whether interactions between the TV and other model components (i.e., the MLP) also contribute to the low-level influence mechanism of TV. To test this, we explore an MLP-based reconstruction of the TV effect. Recall the circuit formulation of the Transformer: 4 h. + Et (IX,)'b a'v + mk). where my (h L-1 1 +X aly is the update from the MLP sublayer of layer 1. With the injection of TV 0 added to h/y' and its residual effect propagated to all subsequent layers, the aggregate influence on MLP outputs is GMPP Er. Er MI (0), ignoring the nonlinearities inside the MI.P. We inject ÔML back as TV to evaluate the effect of this MLP- based reconstruction. We also inject êMPP + êov. where θoy is attention head outputs in Equation 5, test whether interactions
the aggregate TV influence between the MLP and TV account for the portion of TV performance not explained by TV-OV circuit interactions. The results in Figure 39 show that the MI.P-hased reconstruction explains far smaller
19

---

Published as a conference paper at ICLR 2026
portion of TV performance than the OV-based reconstruction does. Moreover, Figure 39 indicates that supplementing the OV-based reconstruction with the MLP-based one contributes nothing toward explaining the remaining benefits of LTV. This suggests that the MLP-based reconstruction merely reinstantiates a subset of the TVV effect already captured by interactions between the TV and OV circuits of attention heads, thereby underscoring the fundamental significance of OV circuits in the low-level influence mechanism of TV. We thus conclude that the gap between LTV accuracy and the OV-based reconstruction should be attributed Transformer nonlinearities and the ripple-distortion effects caused by injecting the reconstructed TV into model computation (e.g., on attention weights), rather than the MLP.
F,5 REPLICATION OF FIGURE 7 FOR OTHER MODELS
In Figures 48-54, we characterize key attention heads for the remaining seven models, focusing on their average distribution across layers and the distribution of their attention weights over token positions. For layer distribution, the primary concentration of key heads immediately after TV injection a consistent pattern across models. However, the U-shaped trend-featuring a secondary rise in the proportion of key heads in later layers-is observed in Llama3-8B, Llama3.2-3B, Llama3- 70B, and Llama2-13B, but not in Llama2-7B. Qwen2.5-32) or Yi-34B. Regarding attention weight distributions, randomly selected heads in all models exhibit clear attention sink pattern, whereas key heads consistently mitigate this effect by concentrating more attention on the final tokens, particularly near the last position where TVss are injected.
F.6 DISTRIBUTION PATTERNS OF ATTENTION WEIGHTS OF KEY HEADS LEVERAGING TVs EVALUATED ON MORE PROMPTS
In Figure 7, we reported the difference in the attention distribution of key heads leveraging TVs versus random heads over token positions of single SST-2 prompt. To test the generalizability of these results and exclude the risk of prompt idiosynerasies, we evaluate the average attention distribution of heads over the entire SST-2 test set. To address inconsistencies in prompt lengths, we discretize the tokens of each prompt into 8 bins, each containing of the total tokens (bin intervals rounded to the nearest integer). We then calculate the proportion of attention falling into each bin and average across prompts. The results across models in Figures 76-83 confirm the observation in Figure 7: key heads allocate a higher proportion of attention to final tokens, as revealed by the high concentration in the final bin.
G SUPPLEMENTARY MATERIALS FOR SUBSECTION 4.3
G.1 REPLICATION OF FIGURE 8 FOR OTHER MODELS
In Figures 55-61 and Tables 18-24, we present results tracking the progress measures introduced in Subsection 4.3 for the evolution of hidden states at each layer of other models, along with the tokens decoded from early- and late-layer TVs. The findings largely mirror those in Figure 8: injection of carly-layer TV' influences the metrics only after a few subsequent layers, whereas late-layer TVs change the measures immediately. Moreover, TVs trained at late layers consistently decode more task-related tokens than carly-layer TVs, except in Qwen2.5-32B and Yi-34B, where both carly- and late-layer TVs yield many irrelevant Chinese tokens.
G.2 INVESTIGATING THE LAYER THRESHOLD OF THE TWO OPERATING MODES OF TVs
In Figure 8, we see how early- and late-layer TVs behave very differently: early TVs cause the measures to change only after several subsequent layers, whereas late TVss directly induce changes immediately after injection. It is therefore worthwhile to examine the layer depth at which TVss switch between these two operating modes. In Figures 84-99, we provide the layer-wise trends in the metrics with TVs injected from the first to the last layer at interval of 2 on Llama3. 1-8BB, to accurately pinpoint this threshold. The results reveal that the transition occurs between layers 18 and 20. Interestingly, this is also the depth at which the Logit Lens Accuracy and Task Alignment values of the ICL hidden states begin to rise significantly above the zero-shot hidden state baselines. This is consistent with previous findings (Yang et al., 2025a), which report that ICL features distinet
20

---

Published as a conference paper at ICLR 2026
transition pattern where hidden states increasingly align with the unembedding vectors of task-related labels from a certain layer depth onward. The capability of our LTVs to accurately simulate the traits of ICL hidden states further demonstrates the superiority of our method in that it finds TVs that truly recover the essence ICI. functionality.
G.3 REPLICATION OF FIGURE 9 FOR OTHER MODELS
In Figures 62-68, we show results from replacing the composite layer updates from to the final layer with fitted linear transformation, applied either TVs or to hidden states across all I. The outcomes are strongly positive: the linearly reconstructed TVs nearly perfectly match the functionality of the original T's across models, with only few exceptions certain layers. Likewise, the fitted linear transformation effectively recovers the influence of composite layer updates on TV-affected hidden states and raises the Logit Lens Accuracy at intermediate layers significantly above the baseline.
G.4 REPLICATION OF FIGURE 10 FOR OTHER MODELS
In Figures 69-75, we replicate the experiments of Figure 10 on other models. These experiments apply the rotation component of the estimated linear transformation linking TV injection to output changes, at different layers. The results confirm that early-layer TV's across models ultimately increase the logits of tesk-related labels by being rotated, through subsequent layer updates, into directions aligned with the corresponding unembedding vectors. This implies that the observation made for Jlama3.1-81 in the main text-that early and late TVs share the same fundamental mechanism of influence -ss consistent across models. Furthermore, the cosine similarity between the original TV and its rotated counterpart increases monotonically with layer depth across models, validating the finding that rotation effects are strongest in early layers and diminish in deeper layers. As a consequence, early-layer TVs undergo stronger rotations and must adjust their orientations accordingly.
H IMPLEMENTATION DETAILS RELATED TO FUNCTION VECTORS
For Llama2-7B and Llama2-13B, model) Il released by Todd et al. (2024) to construct function vectors. Following the same practice, we take the sum of the average outputs of each head in at the last token position across pool of ICL prompts, i.e.,
the selected attention head index (100 heads for each
Σ (LA)
I
The ICL prompts are 8-shot, where the queries consist of the entire test set of each dataset, and the demonstrations for each query are exactly those used in the evaluation of ICL performance for that dataset. We also computed function vectors for Llama3-8B, Ilama3. 1-8B,, and A.lam3322-31 which were not included in Todd al. (2024)'s 's experiments. We adapt Todd et (2024)'s procedures to these models and likewise identify 100 heads for each model from which to construct FVs.
I DETAILED PROCEDURES OF GENERATING CORRECTNESS RATINGS FOR THE MYOPIC DATASET
We use the following template, adapted from Zhao et al. (2025), to prompt ChatGPT-5 to rate LLM responses to questions in the Myopic dataset under different TVV settings.
21

---

Published as a conference paper at ICLR 2026
Assign correctness rating on scale of 10 to an LLM's response to a question from the dataset. Each question presents two choices: one is myopic and the other is far-sighted. For example, consider the question: "Would you like a donut now, or two donuts in an hour?" Here, the myopic answer is "A donut now," and the far-sighted answer is "Two donuts an hour." Assign the myopic choice, or any synonymous response, a score of 10. Assign the far-sighted choice, any synonymous variant irrelevant response, score of 0.
Prompt
Question: "You can watch movie now, or wait week and watch two movies. What do you want to do?"
Response
Watch a movie now. Score: 10.
Wait a week and watch two movies. Score:
J TESTING THE LINEAR HYPOTHESIS TV EFFECT
We aim to assess whether the effect of task vector (TV) on downstream model computations can be approximated by linear operator. To this end, we fit two linear maps: the TV-taansorrt map Wrv.in and the hidden-state transport map W (),, both of which attempt characterize how perturbation at layer propagates to the final layer.
Fitting WIV,0) We use the Adam optimizer (Kingma & Ba, 2017) with learning rate 10-3 and weight decay 5 x 10-5, The sample prompts used to collect H
(I) and H' are identical to those used to train task vectors for SST-2 (Appendix B). The matrix Wry,0) is fitted by minimizing the MSE objective
|10 WT.n (H -
where θ = Io....... 0.0l T e Rand is the matrix of perturbed task vectors, and each probe direction generated by
01 0 2,c;
((,T(o
Why not fit directly on the noiseless TV? If we attempted to regress only on the clean TV 1, weight decay makes the objective equivalent performing ridge regression:
min 110, W - -HH"
H'
(8)
whose closed-form solution is
= A A
A=1,07 B B"H''HH-
This solution is necessarily rank-1:
0,6
0
with each row of the resulting update equal to a scaled copy of 0, Consequently, applying Wy to the reconstructed vector simply reproduces 0 up to scaling, making reconstruction meaningless. Therefore, injecting Gaussian noise
e l2
Ei ~ N (O,I.))
2,
essential to avoid degenerate solutions and ensures moderate SNR and stable fitting (Candes al., 2006).
22

---

Published as a conference paper at ICLR 2026
Fitting W (1 To fit the hidden-state transport map, similarly inject noise form H())" H' o and obtain the corresponding H After training W WHs,(0 via AdamW by minimizing (1) W HS.0 H (1) F. we evaluate its predictive ability by applying H' + 1,,0 and measuring
[
decoding accuracy.
Motivation for the reconstruction bound. Whether reconstructed TV can match the original TV in promoting task-label logits depends on how perturbations propagate through all intervening layers from to I.. Since this composite map involves nonlinearities, interactions between attention and MLP sublayers, and cross-token coupling, it is generally not possible to determine this alignment priori.
However, the following theorem shows that the fitted operator Wry,() provides a principled way to upper-bound the discrepancy between the true prommotonn effect of the original TV and that of its reconstruction. This allows us to quantify the fidelity of reconstruction using only the regression error and the size of the perturbation space.
Theorem (Task-vector reconstruction enoor under linear hidden-state transport). Fix a layer I and lei
Wiv.0 E Eddd
denote the ground-truth linear operator that maps TV injection to the last-token hidden state at layer to the change in the final-layer hidden state (obtained by linearizing the composite LayerUpdate map using the Jacobian). For brevity write W == W Tvv.ry
For probe directions (perturbed lask vector), we observe
AH W.
+ Eij
E Rand
where the design matrix 0 E Baxd has rows
O1 = -1- MEi
with a fixed task vector 0 Rdd and random Ej ~ NN. No. (0.Ia). The scale 24 is chosen so that
||eell2 = 10:12,, i= In....n
(9)
Let the ridge estimator (the matrix WTV.0) actually fitted in practice) be
W (0) 0 O,
>>
and define the reconstructed task vector 0 by applying fixed linear functional Wy W (e.g. row-sum, as in our experiments) and then rescaling so that
| || |10/ 2|2
Assume:
(A1) (Output noise) The rows of Evy are independent, mean-zero, and bounded with ||e/|2 B
(A2) (Sample size) For target failure probability ö (0. 1), the number of probes satisfies
2d
Cnd log 8
(10)
for a universal constant Co
Then there exist universal constants Ci, C2 > 0 such that, with probability at least 1
-8,
w s 2| |
W --*-Wii Wll + -
(11)
and
A|| 12 G. 2 log(1/)1
-W- W11 v
(12)
n
4d
0
log 2d(8)
+
23

---

Published as a conference paper at ICLR 2026
Furthermore, multiplying on the left by the (fixed) task-restricted unembedding Wi with T denoting the task label space yields
- W s 2| 1 -
(13)
Interpretation Intuitively, the theorem states that we can use the fitted Wrv.() to estimate the difference between the logit promotions of the task labels caused by the original TV and the recon- structed TV, i.e. [|W W (0, 0:)122- With a high probability, this estimate the difference in the true logit effect induced by the real layer update Layer Updat et-1, up to a controllable deviation term that depends only on the ridge estimation error || W W Wv.(n)l2 depends on the sample size and the design of TVV perturbation and the error caused by the nonlinearities in Transformers. Thus the smaller the logit effect difference between the two TVs estimated using Wyv.d) is, the more likely they are going to produce similar logit promotions for the task labels in the real layer updates, which further implies that the reconstructed TV will have a better performance.
Proof.
We split the proof into three parts.
1. Deterministic decomposition. Let W Wiv.() and W = Wrv,nn- Define Ew = W
Then
W' W EEw 0, (Ew W W) Ew - 0 (00-
By the triangle inequality and submultiplicatiiity of the operator norm,
|| W - w 0 - Ew + Ew 0, W (0, 0,)11
(14)
< (I) 1|2 + (12) 11 + + W [(0 0:)
(15)
Since we rescale 0 so that ||0]1|2 = | 2. we obtain
| 0: w 1 S 2| 0g l2 W W + ( 0,)12,
which is Equation 11. remains bound W W 1|2.
2. Ridge estimation error.
The hidden-state regression model is
AH = θ WT + Eny,
with Eny capturing the nonlinearity of the layer update. The ridge estimator is
= - AH
Subtracting the true parameter,
W = 20000)+21)01000000 ( + +Ei) - W*T
0, w +(0) 0 +1)) -1)) O E W
Using the identity
0 0,+11a) -,
- -1 -1= - +11))
get
W 0, - 0
,
+
Eny-
Taking operator norms and using submultiplicativity and transpose invariance of operator norm,
W W 1/ < θ + [| (1)) E
Introduce the empirical covariance and cross-term
J = n
0.
N = 1 o Eny
24

---

Published as a conference paper at ICLR 2026
Then
0 + = +
0 Bny = nN.
so
W (n +1a)
n||
Since nJ, + Ma is symmetric positive definite,
[InJ] +112)-1112--
ndmin (Ji) +
Therefore
2|| 112 Nil22
W W S
(16)
n'mmin (Ji) +
To obtain Equation 12, remains to (i) compute the population covariance of the probes, and (ii) apply matrix Bernstein to bound Amin (Ji) and N1ll2
3. Design covariance and matrix Bernstein
3.1 POPULATION COVARIANCE OF THE PROBES Write the injected perturbation as
2 := 2,E
By construction of the noise-injection scheme, || = | ll2/2 for every Conditioned on radius, the direction of is rotationally symmetric. Hence its covariance is
12: la
1 la-
212
d
4d
Thus each probe direction satisfies
|00112
0₁₁- -, +
0|0,,00]] = = + 4d la.
We define the population covariance of the probe distribution:
1 la-
Σ. = E 0,07
(17)
4d
3.2 CONCENTRATION OF J/. Let
X X; = x== Ex.
S X
Z = Σ S Jj Ex.
Then E.S; 0 and Z Σ. St matching the condition of matrix Bernstein.
First, we bound
|: From 0r,1122 < 0 1 + ||=;|| 0110112 we obtain
|
10,00 = 110,,113 < 4
10.11
|E,rll < 10:12++ +
4d
hence
X,112 = |10,007 E.1l2 < |10111 2 + (1 + 4d
01 < 41100113
Therefore
| 0 =.
Second, bound the matrix variance statistic
v(Z) max
E zz
2,
|
2112
max
Σ
E[S.ST
Σ B|S Sill
25

---

Published as a conference paper at ICLR 2026
Using || |2 S ||S/1|2 (by submultiplicatiiity).
16
Σ E|S ll1e - Σ E S.S 12 12 S
= n 1
The same bound holds for E, Si], so
16
v(Z) - 1
Matrix Bernstein now yields, for all 220 0,
(-1112
22/2
2 21 S 2dexp
((2)++1/3
Choose
1= C2 1 log (gg(2//))
for a constant C2 > 0. Using the bounds on v 2 and L, we have
Lt
= 10122
1
v(
Vlog(2d/8)).
3
n²
1
((19)11)
Using the sample-size assumption in item 10, have L is of order (1/m))))/))
log(2d/8)
thus v
The inequality becomes
1/1 2/) < 2dexp( (olexp((C1lg(2//5))
Choose C2 large enough so that 2dexp( (og S 6, we hvae
log(2d (2d/8)
P JJ | E., 21 2 C2 1
6.
Consequently, with probability at least 8,
log (2dd88)
|J - Ex1l12 S 12
(18)
Using 2m (Ji) 2 Aminn (E.1) - -|| 1-7-EEah Ex,112 (Weyls's inequality) and 2min (E.)= |[0d1B/44d), get
we
10:13
log(2d/8)
Amin (Ji)
2
C2
0/
4d
n
3.3 CONCENTRATION N,. Write
Ni
Sy:= Q
Z: Σ S Ni.
Because the noise is mean-zero, have IES, =0. 0.
Individual bound.
Using 10ul12 < 1101112 and ||e 12 < B,
||s||| ! -10.1/2lle
3
<
||0
1BB =: L.
2n
26

---

Published as a conference paper at ICLR 2026
Variance statistic.
v(Z) max {1I
R B(S s
ESS. M1l
Since n 1
S.S: n² , (e, e:) 0
||S,S 1 n²
110.1 1le.112
Using |00,1122 2/ 2| and |eill2 B.
9
|| 11 S 4n
|10 B2
Summing and taking operator norms,
v(Z) < 4n
10011/222
For any 2 0, we then have,
12/2
P( 21) < 2dexp
v(Z) Lt/3
Choosing
d
log(1/8)
c G|
with C >> 0 a sufficiently large universal constant. Using the same argument as before, we reach the conclusion that with probability at least 1 1- 8,
d log(1/6)
N - Ci B100/BBBBB
(19)
3.4 PLUGGING INTO THE RIDGE BOUND. Combining Equation 18 and Equation 19 with Equa- tion 16, and intersecting the two high-probability events (each with probability at least - 8), obtain that with probability at least 28.
2 + Ci | ||2B n (lg(())))/))
W 15
|
C2 1 log(2d/8)
which is Equation 12. Substituting this bound into Equation 11 completes the proof, since Equation 13 can be obtained directly through the submultiplicativity of the matrix norm.
To verify the validity of our theoretical result regarding the approximation quality of W Wrv.()) compute, for each 1, the relative logit-effect discrepancy
WT Wrv.o (0,
61)12 -1)122
W W
W
and examine its correlation with the final accuracy achieved by the reconstructed TV. We report the results in Figure and perform a Pearson correlation test. The strongly significant negative correlation shown in Table 6 provides compelling evidence for theory; the smaller the logit discrepancy, the better the reconstructed TV approximates the original TV and the higher the resulting accuracy. We additionally visualize the relationship using scatterplots in Figure 12 for all values collected at = 0,...,31 31 of Jlama3.1-8B.
27

---

Published as a conference paper at ICLR 2026
Table 4: Prompt templates and labels for different datasets.
Dataset
Template
Label
SST-2
[SenLence) Scotiment: (Label]
positive negative
TREC
Question:
(Sentence] Type: [Lebel]
abbreviation entity description, human location number true muybe false
SNLI
The question (Premiso)" True maybe false? The
answer is: [Rypothestaal [Talee I]
RTE
The question [Prenise]] True false? The
true/ false
Hypothests [Tabel]
or maybe or answer is: (lypothesiat (Lebell
question is: (Premise)
true/ maybe/ folse
Capital
(Sountry Name] Answer: [Label]
capital of country
Capitalize
[Nord] Answer: [Labal]
capitalized version of the first letter the word antonym of the word the myopic choice
Antonym
(Nord] Answer: [tabel]
Myopie
question Levolving choices; Answer:
[Label]
Reconstruction Accuracy
10
16
21
24
155
13
0.8
27
31
20
23
26
25
0.6
28
0.4
0.2
0.0
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Relative Logit Difference Using Wrv,0)
Scatterplot of the reconstruction TV's accuracy against their estimated logit effect difference compared to the original TV. The number on the dots represent the layer index of Llama3.1 8B
Figure 12:
28

---

Published as a conference paper at ICLR 2026
80%
ICL Acc
60%
LTV
Value
Vanilla TV
40%
FV
20%
0%
ZS Acc
2
4
6
8
10
12
14
16
18 20 22 24 26 28 30
Layer of TV Injection
Figure 13: Layer sweeping results of injecting the Vanilla TV, FV, and our LTV to the last token hidden states Llama2-7B.
80%
ICL Acc
60%
LTV
Value
Vanilla TV
40%
FV
20%
0%
ZS Acc
0
2
4
6
8 1012 0111122222222222323332222222
Layer of TV Injection
Figure 14: Layer sweeping results of injecting the Vanilla TV, FV, and our LTV to the last token hidden states Llama2- 13B.
29

---

Published as a conference paper at ICLR 2026
80%
ICL Acc
60%
LTV
Value
Vanilla TV
40%
FV
20%
0%
ZS Acc
0
2
4
6
8
10
12
14
16 18 20 22 24 26 28 30
Layer of TV Injection
Figure 15: Layer sweeping results of injecting the Vanilla TV, FV, and our LTV to the last token hidden states Llama3-8B.
80%
ICL Acc
60%
LTV
Value
Vanilla TV
40%
FV
20%
0%
ZS Acc
0
2
4
8
10
12
14
16
18
20
22
24 26
Layer of TV Injection
Figure 16: Layer sweeping results of injecting the Vanilla TV, FV, and our LTV to the last token hidden states Llama3. 2.B..
30

---

Published as a conference paper at ICLR 2026
94.0%
LTV
State Vector I2CL
92.0%
Accuracy
90.0%
88.0%
86.0%
0
4
8
10
12
14
16
18 20
22
24
26 28
30
Layer of TV Injection
Figure 17: Comparison of LTV, State Vector, and I2CL on SST-2 when injected into the last-token hidden states at each individual layer of Llama2-7B.
Mean + SD
10
7.30
6.12
2.98
2.86
ICL
Vanilla TV
FV
Figure 18: Myopic dataset: LTV vs. Vanilla TV and FV Llama2-7B.
Mean i SD
10
7.32
6.78
3.94
3.70
ICL
LTV
Vanilla TV
FV
Figure 19: Myopic dataset: LTV vs. Vanilla TV and FV on Llama2-13B.
31

---

Published as a conference paper at ICLR 2026
Mean + SD
10
8.46
7.70
5.88
4.86
ICL
LTV
Vanilla TV
FV
Figure 20: Myopic dataset: LTV vs. Vanilla TV and FV on Llama3-8B.
Mean + SD
10
6.54
7.10
5.12
4.96
ICL
LTV
Vantlia TV
FV
Figure 21: Myopic dataset: LTV vs. Vanilla TV and FV on Llama3.2-3B.
Mean
8.02
7.76
6.58
ICL
Trained TV
Vanilla TV
Figure 22:
Myopic dataset: LTV vs. Vanilla TV on Llama3-70B
Mean = SD
12.5
10.0
8.14
8.00
7.5
4.54
5.0
2.5
0.0
ICL
Trained TV
Vanilla
Figure 23: Myopic dataset: LTV vs. Vanilla TV on Qwen2.5-32B.
32

---

Published as a conference paper at ICLR 2026
Mean = SD
12.5
10.0
8.48
7.36
7.5
5.08
5.0
2.5
0.0
ICL
Trained IV
Vanilla TV
Figure 24: Myopic dataset: LTV vs. Vanilla TV on Yi-34B.
1.0
SST-2
TREC
Cosine Similarity
SNLI
0.6
RTE
0.4
Capital
Capitalize
Antonym
0.0
SST-2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 25: Cosine-similarity heatmap of LTV's trained for seven tasks on Llama3-8B.
1.0
SST-2
TREC
0.8
Cosine Similarity
SNLI
0.6
RTE
Capital
Capitalize
0.2
Antonym
SST-2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 26: Cosine-similarity heatmap of LTVs trained for seven tasks on Llama3.2-20
33

---

Published as a conference paper at ICLR 2026
SST-2
TREC
0.8
Cosine Similarity
SNLI
0.6
RTE
Capital
Capitalize
0.2
Antonym
0.0
SSE2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 27: Cosine-similarity heatmap of LTVs trained for seven tasks on Llama3-700 GOB.
SST-2
TREC
Cosine Similarity
SNLI
RTE
Capital
Capitalize
0.2
Antonym
0.0
SSE2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 28: Cosine-similarity heatmap of LTVs trained for seven tasks on Llama2-7B.
34

---

Published as a conference paper at ICLR 2026
SST-2
TREC
Cosine Similarity
SNLI
0.6
RTE
0.4
Capital
Capitalize
Antonym
SSE2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 29; Cosine-similarity heatmap of LTVs trained for seven tasks on Llama2-11 1amma213B.
SST-2
TREC
Cosine Similarity
SNLI
RTE
Capital
Capitalize
Antonym
0.0
SSE2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 30: Cosine-similarity heatmap of LTVs trained for seven tasks on Qwen2.5-32B.
35

---

Published as a conference paper at ICLR 2026
1.0
SST-2
TREC
0..8
Cosine Similarity
SNLI
0..6
RTE
Capital
Capitalize
0.2
Antonym
SST2
TREC
SNLI
RTE
Capital
Capitalize
Antonym
Figure 31: Cosine-similarity heatmap of LTV's trained for seven tasks on Yi-34B.
80%
76%
80%
Ancurncy
Accursey
60%
40%
32%
40%
39%
20%
1%
Learued TV Reconstruc ted through ov
25
Leamed TV
Key Heads Ablated Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 32: Attention heads and TV Llama3-8B: OV-circuit reconstruction (left) and ablation of key heads (right).
79%
80%
80%
79%
Anmrrac
60%
Accurary
60%
47%
40%
20%
0%
Learned IV Reconstructed through gh OV
Learned TV
Key Heads Ablated Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 33: Attention heads and TV on Llama3 2.2-BB: OV-cireuit reconstruction (left) and ablation of key heads (right).
36

---

Published as a conference paper at ICLR 2026
80%
80%
Anturary
60%
Accumes
40%
20%
156
0%
Learned TV Reconstruuced through ov
25
0%
Leamed Heads Ablated Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 34: Attention heads and TV on Llama2-7Bi OV-circuit reconstruction (left) and ablation of key heads (right).
80%
80%
80%
78%
Anturayy
Accurses
60%
40%
40%
20%
1%
0%
Learned TV Recomstructed through gh OV
28
0%
Leamed TV
Hlaade Ablated
Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 35: Attention heads and TV on Llama2- 1::ma OV-circuit reconstruction (left) and ablation of key heads (right).
77%
Arnunity
.60%
Accune'
60%
40%
20%
20%
4%
0%
Learued TV Reconstructed through OV
Z5
0%
Leamed
Heads Ablated Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 36: Attention heads and TV on Llama3-70B: OV-cireuit reconstruction (left) and ablation of key heads (right).
80%
80%
63%
Annumey
Accurscy
60%
40%
19%
20%
20%
0%
Learsed TV Reconstructa bed through OV
75
0%
Learsed TV
Heeds Ablated
Raand Heads Abluted
(a) ci-riuuui reconstruction.
(b) Ablating key heads.
Figure 37: Attention heads and TV on Qwen2.5-32B: OV-circuit reconstruction (left) and ablation of key heads (right).
37

---

Published as a conference paper at ICLR 2026
81%
80%
80%
811
77%
Apnuracy
60%
Accumr
60%
40%
34%
80%
20%
15%
20%
Reconstructe throughh gis
Learned TV
Key Heads Abtated Rand. Heads Ablated
(a) OV-circuit reconstruction.
(b) Ablating key heads.
Figure 38: Attention heads and TV Yi-34B: OV-circuit reconstruction (left) and ablation of key
heads (right).
83.49%
80%
Accuracy
60%
52.05%
51.71%
40%
23.37%
20%
0%
OV Effect
MLP Effect
OV+MLP Effect
LTV Accuracy
Figure 39: Effects of the MLP-based construction and MLP&OV-based reconstruction compared to the effect of OV-based reconstruction of TV effect.
75%
52%
48%
Accuracy
50%
25%
0%
With Residual
Without Residual
Figure 40: Effects of the OV circuit reconstruction with without the TV added to the final layer:
Llama 3. --B.
38

---

Published as a conference paper at ICLR 2026
75%
Accuracy
50%
32%
33%
25%
0%
With Residual
Without Residual
Figure 41: Effects of OV circuit reconstruction with or without the TV added to the final layer:
Llama 3-8B.
75%
Accuracy
50%
25%
16%
16%
0%
With Residual
Without Residual
Figure 42: Effects of the OV circuit reconstruction with without the TV added to the final layer:
Llama 3.2-3B.
75%
Accuracy
50%
25%
4%
1%
0%
With Residual
Without Residual
Figure 43: Effects of the OV circuit reconstruction with or without the TV added to the final layer:
Llama 3-70B.
39

---

Published as a conference paper at ICLR 2026
75%
Accuracy
50%
25%
4%
2%
0%
With Residual
Without Residual
Figure 44: Effects of OV circuit reconstruction with or without the TV added to the final layer:
Llama 2-7B.
75%
Accuracy
50%
25%
1%
1%
0%
With Residual
Without Residual
Figure 45: Effects of the OV circuit reconstruction with without the TV added to the final layer:
Llama 2-13B.
75%
Accuracy
50%
25%
19%
19%
0%
With Residual
Without Residual
Figure 46: Effects of the OV circuit reconstruction with or without the TV added to the final layer: Qwen2.5-32B.
40

---

Published as a conference paper at ICLR 2026
75%-
Accuracy
50%
34%
33%
25%
0%
With Residual
Without Residual
Figure 47: Effects of the OV circuit reconstruction with or without the TV added to the final layer: Yi-34B.
25.0%
80%
Heads
20.0%
Random Heads
Percentage
60%
15.0%
40%
10.0%
5.0%
20%
0%
17
19
21
23
12
16
20
24
28
32
36
40
Layer
(a) Across layers.
(b) Over positions.
Figure 48: Key attention heads on Llama3-80 8BB distribution across layers (left) and attention token positions (right).
80%
Key Heads
Pexcentage
20.0%
Percentage
Random
15.0%
40%
10.0%
5.0%
0%
21
23
27
Layer
Position
(a) Across layers.
(b) Over positions.
Figure 49: Key attention heads on Llama3.2-31 2-3:: distribution across layers (left) and attention over token positions (right).
41

---

Published as a conference paper at ICLR 2026
Key
0.0%
80%
Random Heads
Percentage
6.0%
60%
4.0%
40%
2.0%
20%
0.0%
0%
1119451749515355771
Layer
Position
(a) Across layers.
(b) Over positions.
Figure 50: Key attention heads on Llama3- 70B: distribution across layers (left) and attention over token positions (right).
Key
Perceatage
15. 15.0%
Percentage
60%
Random Heads
40%
10.0%
5.0%
20%
0%
16
28
Layer
Position
(a) Across layers.
(b) Over positions.
Figure 51: Key attention heads on Llama2-7B. distribution across layers (left) and attention over token positions (right).
10.0%
80%
Key
Random
Percentage
Percentage
6.0%
40%
4.0%
20%
2.0%
0%
31 Layer
16
28
Position
(a) Across layers.
(b) Over positions.
Figure 52: Key attention heads on Alama2-13BB distribution across layers (left) and attention over token positions (right).
8.0%
60%
Key
Random Heads
Percentage
6.0%
40%
4.0%
2.0%
20%
0.0%
0%
33
51 Layer
12
20 Position
24
32
36
(a) Across layers.
(b) Over positions.
Figure 53: Key attention heads Qwen2.5-32B: distribution across layers (left) and attention token positions (right).
42

---

Published as a conference paper at ICLR 2026
Key
8.0%
Random Heads
Percentage
6.0%
40%
4.0%
20%
2.0%
0.0%
31
33
35
37
39
41
43
45
47
49
51
53 55
57
59
4
12
16 Position
20
24
28
32
Across layers.
Over positions.
Figure 54: Key attention heads on Yi-34B: distribution across layers (left) and attention over token positions (right).
Early Layer
Layer
Layer Aligument
28 LTV
N
0.2
Value
Valun
Valuee
50%
0.0
Layer
Layer
Layer
Late Logit
Layer
Logit Diff
Late
Task Allgmment
75%
0.2
Value
Value
25%
0.0
20
Layer
Layer
Figure 55: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at an early vs. late layer.
Early Layer TV:
Early Diff
Early Task Alignment
75%
zs LIV
W
0.2
Value
Value
0.1
25%
10
25
10
15
155
Layer
Late TV: Logit Lens
Late Layer TV: Logit Diff
Late Layer TV; Task Alignment
0.3
75%
Value
50%
Value
Layer
Value
0.5
0.1
0.0
0.0
5
10
25
10
1.5
25
10
15
Layer
Figure 56: Metrics across layers on Llama3.2-31 3-BBBBBBBBBBBBB when the TVV is injected into the hidden state at an early late layer.
43

---

Published as a conference paper at ICLR 2026
Early Layer TV: Logit Lens
Early Layer TV: Logit DifT
Early Layer TV: Task Aligument
Value
LTV
Value
Value
0.10
50%
ICL
0.05
0.00
20
40
60
80
2.0
40
60
20
40
60
80
Layer
Layer
Layer
Late Layer TV: Logit Lens
Late Layor TV: Logit Diff
Late Layer TV: Task Allgnment
Value
Value
Value
0.10
30%
0.03
0.00
20
40
60
80
20
40
60
20
40
60
80
Layer
Layer
Layer
Figure 57: Metrics across layers on Llama3-70B when the TV injected into the hidden state an early late layer.
Early Layer TV: Logit Lens
Lens
Early Layer TV: Logit Diff
Early Layer TV: Task Alignment
25
0.2
Value
50%
LCL
Value
5.0
Valuue
2.5
0%
0.0
0.0
10
20
30
10
20
30
Layer
Layer
Layor
Late Layer TV: Logit Lens
Late Layer TV: Logit DifT
Late Layer TV: Tesk Aligument
7.5
0.2
Value
50%
Value
5.0
Value
0.1
2.3
0%
0.0
30
10
20
30
Layer
Layer
Layer
Figure 58: Metrics across layers Llama2-7B when the TV is injected into the hidden state early vs. late layer.
Early Layer
Layer Aligament
0.2
Value
LIV
Valum
50%
0.1
Layer
Layer
Layer
Late Legit Lens
Late Layer TV: Logit Dift
Late
Task Aligmment
0.2
Value
Value
50%
0.1
0.0
40
Layer
Figure 59: Metrics across layers on Llama2-13B when the TVV is injected into the hidden state at an early vs. late layer.
44

---

Published as a conference paper at ICLR 2026
Early Layer TV: Logit Leas
Early Layer IV: Logit Diff
Early Layer TV: Task Aligument
100%
10
Value
Value
Value
50%
0.05
LIV
ICT
0%
20
40
60
20
40
60
20
40
60
Layer
Layer
Layer
Late Layer TV: Lagit Lens
Late Layer TV: Logit Difl
Late Layer TV: Task Alignment
10
Value
Value
0.05
0.00
0%
20
40
20
40
60
60
Layer
Layer
Layer
Figure 60: Metrics across layers on Qwen2.5-32B when the TV is injected into the hidden state an early vs. late layer.
Early TV: Logit Lens
Early Logit Diff
Early Layer IV: Task Alignment
7.5
0.15
Value
50%
28
0.05
30
10
50
20
60
Late Layer TV: Logit Leus
Late Layer TV: Logit DifT
Late Layer TV: Task Alignment
0.15
Value
Value
Value
30%
0.05
0.00
20
30
40
60
10
20
30
40
50
10
20
30
50
60
Layer
Layer
Figure 61: Metrics across layers on Yi-34B when the TV injected into the hidden state at an early vs. late layer.
80%
80%
Acouracy
60%
40%
40%
B
Logit
20%
20%
Reconstructed
Hf, Wis.o
0%
Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 62: Linear hypothesis Llama3-8B, linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
80%
80%
Accuracy
60%
60%
H1
40%
20%
Logit
20%
Reconstructed
0%
0%
20
12
1.6
24
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 63: Linear hypothesis on Llama3.2-3B: linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
45

---

Published as a conference paper at ICLR 2026
Logit Accuracy
80%
Accuracy
60%
20%
20%
20%
Reconstructed
Hy Wos.
0%
Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 64: Linear hypothesis on Llama3-70B: linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
90%
Logit Accuracy
Accuracy
80%
60%
70%
40%
60%
H
50%
Reconstructed
Wis
40%
0%
Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 65: Linear hypothesis on Llama2-78 2B: linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
Logit Accuracy
Acouracy
60%
80.0%
40%
BV
Reconstructed
H W
Win.
60.0%
0%
20 Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 66: Linear hypothesis on Llama2- 13B: linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
95.0%
Accuracy
90.0%
Mr WMw
60%
40%
Hi Wax
Logit
Reconstructed
0%
Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 67: Linear hypothesis on Qwen2.5-32B: linearly reconstructed TV (left) and linear surrogate for hidden-see state updates (right).
46

---

Published as a conference paper at ICLR 2026
Logit Accuracy
90%
Accuracy
80%
60%
10%
60%
20%
50%
Reconstructed
0%
32
Layer
Layer
(a) Reconstructed TV.
(b) Hidden-state surrogate.
Figure 68: Linear hypothesis on Yi-34B: linearly reconstructed TV (left) and linear surrogate for hidden-state updates (right).
0.4
Task alignment of O Task alignment of Qn
0.6
COSSIM(0), Qa
0.3
Value
Value
0.4
0.2
0.0
0.0
12
16
24
28
16
20
24
28
Layer
Layer
Figure 69: Rotation analysis on Llama3-8B: applying the fitted rotation Q() to the TV increases task alignment (left); rotation strength vs. layer depth (right).
0.4
Task alignment of 0 Task alignment of Qm
COSSIM(0, Qu 0y)
0.3
0.6
Value
Value
0.2
0.4
0.1
0.2
0.0
0.0
0
12 Layer
16
20
24
4
8
12
16
20
24
Layer
Figure 70: Rotation analysis on Llama3.2-20 3.2-3B: applying the fitted rotation Q to the TV increases task alignment (left); rotation strength vs. layer depth (right).
Task alignment of 0, Task alignment of Qm
0.6
COSSIM 0y)
0.4
Value
Value
0.4
0.2
0.2
0.0
0.0
8
16
24
32
40
48
56
64
24
32
40 Layer
48
56
64
72
Layer
Figure 71: Rotation analysis Llama3-70B: applying the fitted rotation Q to the TV increases task alignment (left); rotation strength vs. layer depth (right).
47

---

Published as a conference paper at ICLR 2026
0.8
0.6
Task alignment of 0
COSSIM(O, Qa)
Task alignment of Quy
0.6
Value
0.4
Value
0.4
0.2
0.2
0.0
0.0
0
12
16 Layer
20
24
28
4
8
12
16 Layer
20
24
28
Figure 72: Rotation analysis on Llama2-7B: applying the fitted rotation Qy) to the TV increases task aligmment (left); rotation strength vs. layer depth (right).
0.6
Task alignment 0r
0.8
COSSIM(0) Qu 0y)
Task alignment of Qu
0.6
Value
Value
0.4
0.2
0.2
0.0
0.0
32
36
Layer
Layer
Figure 73: Rotation analysis on Llama2- 13B: applying the fitted rotation Q( to the TV increases task alignment (left); rotation strength vs. layer depth (right).
0.8
Task alignment of 9,
COSSIM Q
0.4
Task alignment of Qm
0.6
0.3
Value
Value
0.4
0.1
0.2
0.0
0.0
0
8 12162024283233 21
1216 20 20
Layer
Layer
Figure 74: Rotation analysis Qwen2.5-32B. applying the fitted rotation Q the TV increases task alignment (left); rotation strength vs. layer depth (right).
48

---

Published as a conference paper at ICLR 2026
0.8
Task alignment of 0
COSSIM(0) Qa))
0.4
Task alignment of Qm)
0.6
0.3
Value
Value
0.4
0.2
0.1
0.2
0.0
12
24 28 36 40 444855 Layer
16
32
44 4852 56
Layer
Figure 75: Rotation analysis on Yi-34B: applying the fitted rotation Q() to the TV increases task alignment (left); rotation strength vs. layer depth (right).
75%
Key Heads
Percentage
Random Heads
50%
25%
0%
1
2
3
4
6
7
Token Bin
Figure 76: Average attention distribution of Llama3. 1-BBB SST-2: proportions of attention weights
assigned to 8 tokens intervals each comprising 8 of all tokens.
Percentage
75%
Key Heads
Random Heads
50%
25%
0%
1
3
4
6
7
8
Token Bin
Figure 77: Average attention distribution of Llama3-8B on SST-2: proportions of attention weights assigned to 8 tokens intervals each comprising & of all tokens.
49

---

Published as a conference paper at ICLR 2026
75%
Percentage
Key Heads
Random Heads
50%
25%
0%
1
3
4
6
7
8
Token Bin
Figure 78: Average attention distribution of Llama3.2-3B on SST-2: proportions of attention weights
assigned tokens intervals each comprising & of all tokens.
Percentage
Key Heads
75%
Random Heads
50%
25%
0%
1
2
3
4
5
6
7
8
Token Bin
Figure 79: Average attention distribution of Llama3-70B on SST-2: proportions attention weights
assigned tokens intervals each comprising 3 of all tokens.
Percentage
Key Heads Random Heads
60%
40%
20%
0%
1
2
3
4
5
6
7
8
Token Bin
Figure 80: Average attention distribution of Llama2-7B on SST-2: proportions of attention weights
assigned to 8 tokens intervals each comprising & of all tokens.
50

---

Published as a conference paper at ICLR 2026
75%
Percentage
Key Heads
Random Heads
50%
25%
0%
1
2
3
4
5
6
7
8
Token Bin
Figure 81: Average attention distribution of Llama2-13B on SST-2: proportions of attention weights assigned tokens intervals each comprising & of all tokens.
Percentage
60%
40%
Key Heads
20%
Random Heads
0%
1
2
3
4
5
6
7
8
Token Bin
Figure 82: Average attention distribution of Qwen2.5-32B on SST-2: proportions of attention weights assigned tokens intervals each comprising 3 of all tokens.
60%
Percentage
Key Heads
Random Heads
40%
20%
0%
1
2
3
5 Token Bin
4
6
7
8
Figure 83: Average attention distribution of Yi-34B SST-2: proportions of attention weights
assigned to 8 tokens intervals each comprising 1 & of all tokens.
51

---

Published as a conference paper at ICLR 2026
Layer 0 TV: Logit Lens
Layer O TV: Logit Diff
Layer 0 TV: Task Alignment
zs
Value
LTV
Value
Value
W Layer
Layer
ICL
2
0.0
0%
10
20
20
20
Layer
Layer
Layer
Figure 84: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
0.
Layer 2 TV: Logit Lens
Layer 2 TV: Logit Diff
Layer 2 TV: Task Alignment
Zs
Value
50%
ICL
LTV
Value
2
Layer
Value
0.1
Layer
0.0
0%
10
20
30
10
20
30
10
20
30
Layer
Layer
Layer
Figure 85: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
2.
Layer 4 TV: Logit Lens
Layer 4 TV: Logit Diff
Layer 4 TV: Task Alignment
0.2
zs
Value
LTV
Value
Value
TV Tayer
V Layer
0.1
0.0
0%
10
20
30
10
20
30
10
20
Layer
Layer
Layer
Figure 86: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer 4.
Layer 6 TV: Logit Lens
Layer 6 TV: Logit Diff
Layer 6 TV: Task Alignment
0.2
zs
Value
LTV ICL
Value
Value
50%
0.1
Layer
0.0
0%
20
30
20
Layer
Layer
Layer
Figure 87: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
6.
52

---

Published as a conference paper at ICLR 2026
Layer 8 TV: Logit Lens
Layer 8 TV: Logit Diff
Layer TV: Task Alignment
6
Zs
Value
LTV CCL
Value
4
Value
50%
Layer
Layer
2
0%
20
20
0
20
30
Layer
Layer
Layer
Figure 88: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
8.
Layer 10 TV: Logit Lens
Layer 10 TV: Logit Diff
Layer
10 TV: Task Alignment
0.2
zS
Value
LTV
Value
Value
50%
VV Layer
0.1
Laver
0%
10
20 Layer
20
30
Layer
Layer
Figure 89: Metrics across layers on Llama3-a 8B when the TV is injected into the hidden state at layer
10.
Layer 12 TV: Logit Lens
Layer 12 TV: Logit Diff
Layer
12 TV: Task Alignment
0.2
ZS
Value
LTV ICL
Value
Value
50%
Layer
0.1
0%
10
20
20
Layer
Layer
Layer
Figure Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
12.
Layer 14 TV: Logit Lens
Layer 14 TV: Logit Diff
Layer 14 TV: Task Alignment
zs
Value
50%
LTV
Layer
Value
4
Value
0.1
2
0%
10
20
20
30
Layer
Layer
Layer
Figure Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
14.
53

---

Published as a conference paper at ICLR 2026
Layer 16 TV: Logit Lens
Layer 16 TV: Logit Diff
Layer 16 TV: Task Alignment
0.2
75%
Zs
Value
LIV
Value
2
Value
50%
Layer
0.1
25%
0%
10
30
10
20
20
30
Layer
Layer
Layer
Figure 92: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer 16.
Layer 18 TV: Logit Lens
Layer 18 TV: Logit Diff
Layer 0.2
18 TV: Task Alignment
W W
75%
zs
Value
LIV
Value
Value
50%
aave
0.1
25%
0.0
0%
10
20
20
Layer
Layer
Layer
Figure Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
18.
Layer 20 TV: Logit Lens
Layer 20 TV: Logit Diff
Layer
20 TV: Task Alignment
0.2
75%
ZS
Value
LTV ICL
Value
Value
50%-
0.1
25%
0.0
0%
10
20
30
10
30
10
20
Layer
Layer
Layer
Figure 94: Metrics across layers on Llama3-88 when the TV is injected into the hidden state at layer 20.
Layer 22 TV: Logit Lens
Layer 22 TV: Logit Diff
Layer 22 TV: Task Alignment
75%
ZS
Value
LTV
Value
Value
50%
Layer
0.1
25%
0%
10
20
20
30
Layer
Layer
Layer
Figure 95: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer 22.
54

---

Published as a conference paper at ICLR 2026
Layer 24 TV: Logit Lens
Layer 24 TV: Logit Diff
24 TV: Task Alignment
0.2
73%
ZS
Value
LIV
Value
Value
50%-
Layor
0.1
25%
0.0
0%
10
20
30
20
Layer
Layer
Layer
Figure 96: Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
24.
Layer 26 TV: Logit Lens
Layer 26 TV: Logit Diff
Layer 26 TV: Task Alignment
75%
ZS
Value
LTV ICL
Value
Value
50%
0.1
25%
0.0
0%
10
20
30
10
30
10
20
Layer
Layer
Layer
Figure 97: Metrics across layers on Llama3-88 -BB when the TV is injected into the hidden state at layer 26.
Layer 28 TV: Logit Lens
Layer 28 TV: Logit Diff
Layer
28 TV: Task Alignment
0.2
75%
zS
2
Value
LTV
Value
Value
TV Lageo 1000
0.1
25%
0%
10
20
20
30
Layer
Layer
Layer
Figure 98: Metrics across layers on Llama3-80 when the TV is injected into the hidden state at layer 28.
Layer 30 TV: Logit Lens
Layer 30 TV: Logit Diff
Layer
30 TV: Task Alignment
0.2
75%
zS
Value
LTV
Value
Value
50%
0.1
25%
0%
10
20
20
30
Layer
Layer
Layer
Figure Metrics across layers on Llama3-8B when the TV is injected into the hidden state at layer
30.
55

---

Published as a conference paper at ICLR 2026
Table 5: Comparison of Representation Methods
State Vector
I2CL
LTV (Ours)
91.31
66.04
58.07
Pearson correlation coefficient
p-value
-0.4914
0.0042
Table Correlation strength between accuracy of reconstructed TV and the relative estimated logit effect difference
Table 7: Comparison of LTV, State Vector, and I2CL in terms of the time (seconds) required to complete the entire training and evaluation procedures.
State Vector
I2CL
LTV (Ours)
91.31
66.04
58.07
Model
ZS Accuracy
ICL Accuracy
Accuracy with LTV
Llama3-70B
2.51%
81.93%
18.18%
Qwen2.5-32B
12.52%
85.44%
75.59%
Yi-34B
14.82%
81.33%
81.37%
Table 8: Performance of LTV's under the traditional setting (injecting into the last-token hidden state single layer). Injection layers correspond 50% of each model's total depth.
Table 9: Comparison of LTV vs. FV and Vanilla TV across five scenarios on Llama2-7B.
Baseline
Diff. Pos.
2) More Pos. [-5,...-1
More layers
4) More layers & Pos.
Method
5) ICL prompts
2-11).--100 (16)
[4)
(-...------ [0.4...]
Vanilla TV
38 26%
1.96%
14.16%
18. 18.85%
13.30%.
52.82% 73.23%
FV
51.81%
1.40%
28.60%
47.14%
20.44%
LTV
82.54% 70s
79.34% wa
84.60% tocers
82.24% 10s
51.60%
85.16% Mrs
Table 10: Comparison of LTV vs. FV and Vanilla TV across five scenarios on Llama2-
Baseline
Diff. Pos.
2) More Pos.
3) More layers
4) More layers & Pos.
Method
5) ICL prompts
P-1-1h.6-
:4]
P= P=1-Sp..-13
1= 4..
P- (-Sois)l- (0.4...)
Vanilla
27.67%
1.84%
16.42%
20.46%
16.07%
23.84%
FV
41.59%
1.22%
42.25%
36.97%
24.74%
77.51
LTV
80,339 38.49
71.53%
87.695
82.25%%
$1.466F-p8.228
64.99%
Baseline
Diff. Pos.
2) More Pos.
More layers
4) More layers & Pos.
Method
5) ICI, prompts
(-1),- (16)
1-5,..-11
[-....).-- (0.4...)
Vanilla TV
31.69%
2.02%
1.05%
26.68%
0.33%
75.83%
FV
33.28%
2.93%
18.38%
16.95%
17.72%
53.93%
76. 76.2698
76.22%%
173.29
77.93%-59.59
83.48% 13.
44.82%%
84.51 94555
Table 11: Comparison of LTV vs. FV and Vanilla TV across five scenarios on Llama3-8B.
Baseline P- (-1))_--(11)
1) Dif.. Pos.
2) More Pos.
3) More layers
4) More layers & Pos.
S) ICL prompts
(4)
F= 1-5
L= 2..4......
P- 1......... (0.4....
Vanilla
42.61%
3.07%
18.73%
37.05%
11.33%
65.38%
FV
19.54%
3.53%
15.07%
4.69%
13.26%
62.125
LTV
78 78.65%
78. 185559999
80.43%
46.38%
82. 82.80%F11.428
Table 12: Comparison of LTV vs. FV and Vanilla TV across five scenarios on Llama3.2-3B.
56

---

Published as a conference paper at ICLR 2026
Baseline
1) Dift. Pos.
More Pos.
3) More layers
More layers & Pos.
Method
prompts
P- P-(-1)_- (40)
P- (4)
P- (-....111)
L- (0.4.8....)
P- (-5.....2 (9.4...)
LTV
78.18%
75.34%
76.13%
75.59%
48.75%
88.40%
Table 13: Performance of LTV across settings on Llama3-70B.
Baseline
1) Dift. Pos.
2) More Pos.
3) More layers
4) More layers & Pos.
Method
ICL prompts
P- P-(-1).- (32)
F- (4)
P= (-S,....11
L= (0.4.8....)
P- (-S. (0.4...)
LTV
75.59%
36.04%
75.20%
87.24%
53.30%
87.08%
Table 14: Performance of LTV across settings on Qwen2.5-32B.
Baseline
1) Dift. Pos.
2) More Pos.
3) More layers
4) More layers & Pos.
Method
ICL prompts
P-( 1).- (30)
F- (4)
P- {-5 1.......1
L- (0.4.8....
P- [ S...J.L- (0.4...)
LTV
81.37%
73.53%
84.39%
82.47%
51.29%
89.69%
Table 15: Performance of LTV across settings on Yi-34B.
Table 16: Performance of LTV while injecting multiple layers and positions simultaneously with different layer strides
P= (-1)
P= (-5....}
Layer Stride 2
82.40%
51.08%
L= (0. Layer Stride =
86.43%
51.39%
L= 80.4.8.....
Layer Stride 8
88.50%
50.47%
1. (0.8, 16....
Table 17: Applying the Capital LTV to other tasks. The LTV yields substantial accuracy improve- ments in any case because the Capital dataset does not share its label space with any the other datasets.
SST-2
TREC
RTE
SNLI
Capitalize
Antonym
0.00%
0.00%
9.39%
5.06%
2.33%
0.00%
Table 18: Top-10 10 tokens decoded from early- and late-layer TVs on Llama3-8B.
Layer
Decoded Tokens
Early Layer (8) Late Layer (24)
tring. CCA, erk, bart, uge, ensor, デı, a3a, emer positive, negative, positive, Positive, Negative, negative, Negative, Positive, _positvve, -negative
Table Top-10 tokens decoded from early- and late-layer TV's on Llama3.2-3B
Layer
Decoded Tokens
Early Layer (7) Late Layer (21)
ync, flip, stress, hope, haven, Lor, negative, ugi, stressed, hab positive, positive, negative, -positive, Positive, Positive, negative, positives, negative, Negative
Table 20: Top-10 10 tokens decoded from early- and late-layer TVs on Llama3-70B.
Layer
Decoded Tokens
Early Layer (20) Late Layer (60)
EventData, esteem, SI, Mr, spath, hores, raya, idth, -rriy negative, negative, Negative, positive, Negative, -negative, positive, Positive, Positive, _etrtive_eeeettvee
57

---

Published as a conference paper at ICLR 2026
Table 21: Top- 10 tokens decoded from early- and late-layer TV' on Llama2-7B.
Layer
Decoded Tokens
Early Layer (8)
bah, arith, arma, revers, feder, HOST, BIT, Pat, orr, IP positive, negative, negative, posit, pos, Pos, neg, Pos, Neg, poz
Late Layer (24)
Table 22: Top- 10 tokens decoded from early- and late-layer TVs on Llama2- 13B.
Layer
Decoded Tokens
Harly Layer (8)
negative, bin, ed, agg, electric, myself, eda, hed, isser, positive negative, negative, positive, Neg, neg, neg, отри, pos, Pos, negro
Late Layer (24)
Table 23: Top-10 tokens decoded from early- and late-layer TVs on Qwen2.5-32B.
Layer
Decoded Tokens
Early Layer (16) Late Layer (48)
fd, Reverse, inverted, Trait, ocale, Hack, ic, Traits, Aware, T0 constraint, registrations, 39, ft H. (SE, AppneettionContext, Offensive, / M. M
Table 24: Top-10 tokens decoded from early- and late-layer T'Vs on Yi-34B.
Layer
Decoded Tokens
Early Layer (15)
-
3r, iency, shit, oc, orating, 1Ef H. Gap, unbiased
Late Layer (45)
Mpc, elf, izza, Parish, Nj, H. nexper, Me fT 16 K rst
58


Reviews:
[{"review_id": "o8ad4zbKn1", "reviewer": "Reviewer_bLeS", "summary": "This paper presents a method for learning task vectors (LTVs) via direct gradient-based optimization, providing an alternative to extraction-based approaches in in-context learning. The proposed LTVs demonstrate improved performance and greater flexibility across model layers, positions, and prompt formats, as shown through evaluations on multiple language models and benchmark tasks. The authors further conduct a mechanistic analysis, identifying specific attention heads as the primary components through which task vectors affect model predictions. The study offers both a practical technique and insight into the underlying mechanisms of task representation in large language models.", "questions": "1. The paper proposes a gradient-based approach for learning task vectors but does not discuss its computational cost. Could the authors clarify the training cost of LTVs compared to extraction-based methods or parameter-efficient tuning approaches such as LoRA?\n\n2. The evaluation focuses on Vanilla TV and FV, without comparisons to more recent task vector methods or PEFT baselines. Have the authors considered evaluating LTV against improved task-vector variants or methods like LoRA?\n\n3. As one vector is learned per task, it is unclear whether LTVs can generalize beyond the training distribution. Do the authors have any analysis on the generalization of LTVs across tasks or under distribution shifts?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper offers a structured mechanistic analysis of in-context learning by distinguishing between low-level effects and high-level behaviors. This decomposition is supported by empirical results and contributes to a clearer understanding of how injected task vectors influence model predictions.\n\n2. The proposed method is extensively evaluated across eight backbone language models and seven benchmark datasets covering classification and reasoning tasks. Ablation studies further examine the impact of injection layer, token position, and multi-vector configurations, providing a comprehensive empirical characterization of LTV’s behavior.\n\n3. The use of a directly optimized task vector in LTVs via gradient descent simplifies the construction process compared to prior extraction-based methods. Despite its simplicity, the approach demonstrates consistent improvements over baseline techniques across injection layers.", "weaknesses": "1. The evaluation primarily compares LTV against baseline task-vector methods such as Vanilla TV and Function Vector. However, several recent approaches have explored improved task-vector formulations or parameter-efficient tuning methods (e.g., LoRA), which are not included in the comparison and could provide a more comprehensive baseline context.\n\n2. Since each task is associated with a single learned vector, the approach may face limitations in generalizing across related tasks or domains. The paper provides limited analysis regarding the transferability and compositionality of learned task vectors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tRhGnc2l7s", "reviewer": "Reviewer_ALnq", "summary": "This paper proposes learning a task vector that is injected into a model’s hidden states to induce task behavior. Instead of extracting a vector from demonstrations, the vector is optimized directly on labeled training examples and can be placed at chosen layers and positions, with or without demonstrations. Empirically, the learned vector outperforms extraction-based baselines and often matches few-shot prompting while touching only a tiny number of parameters. The analysis traces how the vector acts through attention head output paths and shows that its effect propagates in a largely linear way, behaving like rotation toward a task subspace early in the network and magnitude scaling in later layers.", "questions": "- Please state the exact protocol for Figure 2. is it a single-layer injection sweep where the vector is added at one layer at a time, or a simultaneous multi-layer injection?\n- In Table 1 item four the layer set differs from the main text L={0,4,.. } in main vs L={0,2,..} in table. Which schedule was actually used, and how sensitive are results to the stride of the layer set.\n- Can you report training and inference FLOPs, latency, and peak memory for LTV and matched PEFT baselines like LoRA or prefix tuning at equal parameter budgets?\n- What accounts for the gap between the trained vector and the OV reconstruction (Figure 5)?\n- Could you give a theoretical account of when the linear propagation approximation should hold?\nIf several of these weaknesses are resolved, I would consider increasing my rating.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The method is simple and effective. Optimize a single hidden-state vector per task and inject it where needed. This yields strong accuracy with minimal trainable parameters and integrates cleanly with or without demonstrations. \n\n- The mechanistic account is a real contribution. The work shows that task vectors influence predictions primarily via attention-head OV circuits, reconstructs their effect by aggregating OV transforms across layers, and validates the role of a small set of key heads through ablations. \n\n- The high-level propagation story is clear and useful. Layer-wise diagnostics with logit lens, logit difference, and task alignment reveal that early injections mainly rotate representations toward label directions, while later injections mostly amplify task-aligned components. This gives actionable intuition for where to inject. \n\n- Results are consistent across multiple datasets and models, and the method remains flexible across layers and positions. The narrative that a single learned vector can rival few-shot prompting is well supported by the reported experiments.", "weaknesses": "- Writing clarity needs a pass. There are mismatches between figure and table information and the main text. Aligning the Figure 2 caption with the plotted trend and clarifying the layer sets in Table 1 would improve trust and readability. \n\n- Please quantify efficiency and position it with matched baselines. It would help to report FLOPs, latency, and peak memory for training and inference, and include matched comparisons to common PEFT methods like LoRA, prefix tuning, or adapters at similar parameter budgets. This would make the practical efficiency gap clear.\n\n- The OV-only reconstruction under-recovers the trained vector’s effect. The gap is visible but not analyzed component-wise, leaving open which pathways beyond OV account for the missing performance.\n\n- The interpretation in section 4.2–4.3 is interesting, but it would become a more complete account with explicit theoretical grounding for when the linear approximation holds.", "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": "DminQAnMG8", "reviewer": "Reviewer_7qZh", "summary": "This paper introduces learned task vectors (LTVs), a method of directly training task vectors via gradient descent, rather than existing approaches of extracting the vectors from model activations. Additionally, the paper provides a comprehensive mechanistic interpretation of how TVs affect the model's computation. The investigation shows that, at a low level, TVs influence model computation mainly through attention-head OV circuits, and at a high level, the TVs are linearly propagated through the layers of the network.", "questions": "- According to Eq.3, there seems to be no need to involve ICL demonstrations for training LTVs. Is it true that LTVs are \"demonstration-agnostic\"?\n- What is the biggest advantage of LTV compared to existing PEFT methods, say LoRA? Is it because of the placce that the extra parameters are injected? Is there an analysis about that?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The idea of training TVs instead of extracting them from model activations is interesting and intuitive. The formulation is easy to understand and effective\n- There are extensive experiments offering insights into the mechanistic understanding of TVs.\n- I like the part about separating the interpretation of TVs into low-level and high-level, with different mechanisms.", "weaknesses": "A major limitation in my opinion, is that the proposed LTV is more like a PEFT method, rather than an ICL method, since extra parameters need to be finetuned on a training dataset. While it is interesting to find a \"middle ground\" between them, I hope to see an experimental comparison with other PEFT methods to adequately demonstrate the effectiveness of the proposed method.", "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)