Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis
Authors: Haolin Yang
Hakaze Cho
Naoya Inoue
Venue: ICLR 2026 Poster (2026)
Keywords: Mechanistic Interpretability
In-context Learning
Large Language Model
Primary area: interpretability and explainable AI
TL;DR: We identify and analyze heads responsible for task recognition and task learning in in-context learning.

Abstract: We investigate the mechanistic underpinnings of in-context learning (ICL) in large language models by reconciling two dominant perspectives: the component-level analysis of attention heads and the holistic decomposition of ICL into Task Recognition (TR) and Task Learning (TL). We propose a novel framework based on Task Subspace Logit Attribution (TSLA) to identify attention heads specialized in TR and TL, and demonstrate their distinct yet complementary roles. Through correlation analysis, ablation studies, and input perturbations, we demonstrate that the identified TR and TL heads independently and effectively capture the TR and TL components of ICL. Via steering experiments with a focus on the geometric analysis of hidden states, we reveal that TR heads promote task recognition through aligning hidden states with the task subspace, while TL heads perform rotations to the hidden states within the subspace towards the correct label to facilitate the correct prediction. We also demonstrate how previous findings in various aspects of ICL's mechanism can be reconciled with our attention-head-level analysis of the TR-TL decomposition of ICL, including induction heads, task vectors, and more. Our framework thus provides a unified and interpretable account of how LLMs execute ICL across diverse tasks and settings.

Published as a conference paper at ICLR 2026
LOCALIZING TASK RECOGNITION AND TASK LEARN- ING IN IN-CONTEXT LEARNING VIA ATTENTION HEAD ANALYSIS
Haolin Yang! University of Chicago
Hakaze Cho² JAIST
Naoya Inouc RRIKEN
ABSTRACT
We investigate the mechanistic underpinnings of in-context learning (ICL) in large language models by reconciling two dominant perspectives: the component- level analysis of attention heads and the holistic decomposition
ICL into Task Recognition (TR) and Task Learning (TL). We propose novel framework based on Task Subspace Logit Attribution (TSLA) to identify attention heads specialized in TR and TI., and demonstrate their distinct yet complementary roles. Through correlation analysis, ablation studies, and input perturbations, we show that the identified TR and TL heads independently and effectively capture the TR and TL components of ICL. Via steering experiments with geometric analysis of hidden states, reveal that TR heads promote task recognition by aligning hidden states with the task subspace, while TL heads rotate hidden states within the subspace toward the correct label facilitate prediction. We further show how previous findings on ICL's mechanism-inchude and more-can can be reconciled with our attention-heade head--eeeeeeee analysis of the TR-TI. decomposition. Our framework thus provides unified and interpretable account of how LLMs execute ICL across diverse tasks and settings'.
induction heads, task vectors,
INTRODUCTION
A key property of Large Language Models (LLMs) is their ability to solve tasks from demonstrations embedded in the input-without further training. This phenomenon, known as In-context Learning (ICL) (Brown al., 2020; Radford et al., 2019), has reduced the need for large datasets and finetuning, enabling fast adaptation of LLMs to new tasks (Dong et al., 2024; Sun et al., 2022). Since its success cannot be explained by traditional gradient-based paradigms (Ren et al., 2024b), deciphering the mechanism behind ICL has become a central research question of great academic interest.
Two research paradigms dominate this pursuit. (1) The introspective paradigm designates internal model components or representations as critical drivers of ICL functionality. Pioneering works (Elhage et al., 2021; Olsson al., 2022) formulate the output logits of Transformers the sum of individual component outputs and highlight the significance of Induction Heads (IHs) in toy models, with follow-ups confirming their importance Shutova, 2024; Halawi al., 2024; Cho et al., 2025a). These studies inspired the concept of task vectors-compartt representations distilled from hidden states or attention head outputs that steer zero-shot prompts toward ICL.-level predictions (Hendel et al., 2023; Todd et al., 2024; Liu et al., 2024), and spurred further inquiries into the properties, behaviors, and emergence of IHs (Ren et al., 2024b; Singh al., 2024; Yin & Steinhardt, 2025). (2) The holistic paradigm instead treats the LLM as an entirety and investigates ICL's properties by directly inspecting and probing how different demonstration configurations shape ICL performance. For instance, by perturbing the demonstration labels in context, Pan et al. (2023) factorize ICL into two core components: Task Recognition (TR, recognizing the label space) and Task Learning (TL, learning the text-label mapping), each contributing to part of the ICL functionality (Figure 1 (A)). Min et al. (2022) also systematically explored the effect of the distribution of texts and labels in demonstrations individually, as well as the templates and number of demonstrations.
larger models via ablation (Crosbie
I Code Implementation: https:// http://ggthuubb.cocmmi./a/ /ocal 2inig_n__ TR_T..

---

Published as a conference paper at ICLR 2026
A Task Recogaition & Task Learning
B
Geometry
Geumetry of Task Recognition and TR leads
Task Learning and TL Heads
is fantastic Not good like
-
negative posilive
positive
Positive Negative
Test Samples:
W pos.
TR Heads
Satisfactory
(posilive)
Heads
Task Recognition
y (+,-)
Task Learning (+ pos.,
ong.
U
pos.
W neg.
U
Figure/ I: (A) Example of how LLMs deduce the label of final query through ICL, which consists of two components: task recognition (identifying the label space) and task learning (mapping demonstration texts to labels). (B) The outputs of Task Recognition heads align with the task subspace spanned by candidate label unembeddings, thus they can steer the hidden states align with the subspace by reducing the angle between the point clouds and the task subspace. (C) Task Learning heads act as rotations within the task subspace, aligning the query's s hidden state with the unembedding of the correct label and enabling the correct prediction.
The two research paradigms offer complementary insights but also limitations. The introspective paradigm localizes ICL to individual attention heads, yet its reliance on ablation only measures how much performance changes when heads are removed, without explaining how these heads realize ICL or behave under varied inputs. The holistic paradigm provides broad functional view, separating ICL into TR and TL, but cannot trace these roles back to concrete components. A unified framework is needed to combine mechanistic precision with functional clarity.
Therefore, we propose TR/TL decomposition via attention-head analysis (Figure 1). Task Subspace Logit Attribution (TSLA) quantifies each head's ' effect on hidden states relative to task-label unem- beddings, identifying TR heads and TL heads as ICL's two drivers. Geometric analyses show that TR heads align hidden states to the task-label subspace, enabling label recognition and suppressing irrelevant tokens (Figure (B)), while TL heads rotate states within it toward the correct label, sharpening prediction (Figure (C)). Together they implement ICL functionality.
We validate the framework using correlation analyses, ablations, input perturbations, and steering across ICL settings, including corrupted demonstrations and free-form generation. These results reconcile prior findings: IHs emerge as a subset of TR heads for label-space recognition, and zero-shot performance is traced to poor task-subspace alignment, explaining why TR-head outputs (hence IHs) form effective task vectors that restore alignment (Todd et al., 2024).
The three core contributions of this paper are:
We derive TSLA to identify TR and TL heads, localizing the two ICL components from input- perturbation studies to concrete attention heads.
2. We use correlation/overraap analyses and TR/TL-channel ablations show separable control of task recognition vs. task learning, unifying ICL observations (c.g., IHs) in the TR/TL framework.
3. We perform steering and geometric analyses to explain how TR heads drive task-subspace align- ment and TL heads drive wuthin-subspace rotation toward labels in classification and generation.
RELATED WORKS
TR& & TL decomposition of ICL Early ICL work relied on perturpert experiments (Min et al., 2022; Pan et al., 2023; Wei et al., 2023). By removing semantic labels (e.g., "negative" + "0") or scrambling the text-label mapping while still observing non-trivial accuracy. these works concluded that ICL. comprises two components: Task Recognition (TR) and Task Learning (TL.). However, this approach cannot provide mechanistic explanations or tie ICL. functionality to model components. We bridge this gap with theoretically derived TSLA method based tusk-subspace geometry, identifying critical TR and TL heads for the two ICL components (Subsection 3.2).
Circuit formulation of Transformer Our TSLA method builds upon the circuit formulation of Transformers (Elhage et al., 2021) by reframing head contributions in subspace geometry, which decomposes output logits into contributions from attention heads and MLPs. This formulation has enabled precise attribution of LLM behaviors to individual components (Crosbic & Shutova, 2024) and spotlighted the Induction Heads (IH) as crucial for ICI. (Song et al., 2025). In toy copying tasks
2

---

Published as a conference paper at ICLR 2026
(XX] - [Y]), IHs attend to earlier occurrences of [Y], mimicking label copying. Their importance for ICL has been confirmed in large-scale models through ablation studies that replace or remove their outputs and observe changes in the final logits (Cho et al., 2025a). In this work, demonstrate how the significance of IHs fits into the TR & TI. decomposition of ICL via extensive correlation analyses and investigations of the heads' distributional properties Subsection 4.1.
Ablation of special attention heads The circuit formulation motivates identifying special heads by ablating their outputs and inspecting prediction changes. Examples include Function Vector Heads (Yin & Steinhardt, 2025)-heads whose ablation most harms ground-truth label logits and are thus viewed as causes of ICL functionality (Sun et al., 2025). Similar approaches extend to retrieval-uummented generation (Jin et al., 2024; Kahardipraja et al., 2025) and chain-of-thoughn reasoning (Cabannes et al., 2024). Yet they reveal only how much head contributes, not how or which ICL component it affects; heatmap-based explanations remain too shallow for causal claims (Kahardipraja al., 2025; Ren et al., 2024a). We address this by new TR/TL-channel metries and fine-grained ablations that expose component-specciie effects Subsection 4.2.
Task Vector steering and geometric analysis An alternaivve to ablation to aggregate head outputs into task vectors that steer zero-shot states to ICL-level accuracy (Hendel et al., 2023), with effective vectors viewed as ICL's mechanistic origins (Todd et al., 2024). However, interpretability issues remain: (1) candidates are chosen via ablation, inheriting its limitations; (2) effects are measured only at outputs, leaving their role in computation unclear (Merullo et al., 2024). Geometric analyses of layer-wise hidden-state evolution, revealing how task vectors reshape geometry, offer a promising alternative (Kirsanov et al., 2025). Jiang et al. (2025) identify a compress-expann pattern in task representations during ICL, while Yang et al. (2025) link IH outputs to alignment between hidden states and unembedding vectors of task-relevant labels. This perspective elucidates their influence on outputs. We perform a comprehensive analysis of the distinct geometric effects that task vectors constructed from TL and TR heads have on the model's 's hidden states Subsection 4.3.
3 METHODOLOGY
3.1 BACKGROUND
Circuit formulation of Transformer In the circuit formulation of Transformer LLMs. an input query with N tokens [x1.....XN] (e.g.. "I like this movie. Sentiment:" for sentiment analysis task) is first transformed into layer-0 hidden states h.. hy via the embedding matrix WE E IR Mix where is the model dimension and V the vocabulary. These hidden states then pass through L layers, where the update of the i-th token's 's hidden state at layer / is:
K Σ
h/ - h4
a m/
with a the output of the k-th attention head (denoted head (1,k)) in the attention sublayer, and m the MLP sublayer output. a/, is the weighted sum of the layer-(l- (oo-((11) 1) hidden states of the first i tokens, H = [n 1 transformed by the embedding matrices of head (1,K). The final hidden state of the last token (*:" in the previous example) can thus be written as:
h/y = 0 + Σ
Σ aww + mmy)
(1)
hy multiplied by the unembedding matrix Wy E Vixd to form logits. Each head output thus contributes to the logits additively as Wyak referred to as Direct Logit Attribution (DLA)
(Olsson et al., 2022; Chughtai et al., 2024; Yu Ananiadou, 2024; Lieberum et al., 2023).
ICL and Induction Head In ICI., m laxt- demonstration pairs f,, VI, fm. Ym are prepended to the query, forming the sequence t1 .!..... tm. Ym. q (c.g.. "I hate this movie. Sentiment: negative. This movie is great. Sentiment: positive... like this movie. Sentiment:"). With these demonstrations, the attention head outputs a'y to the final position, depending on all preceding tokens, producing logits that can lead the LLM to predict y", the correct label for q. An Induction Head (IH) is a special attention head that, at each position, searches for earlier occurrences of the current token, attends to
3

---

Published as a conference paper at ICLR 2026
the immediately following tokens, and copies their information back to the current position. In the example above, an IH at the final position places its attention on the "positive" and "negative" tokens that follow previous "." tokens and uses their hidden states to form its output.
3.2 IDENTIFYING TR AND TL HEADS USING TASK SUBSPACE LOGIT ATTRIBUTION
Pan et al. (2023) decomposes ICL into two components. Task Recognition (TR) means recognizing the set of candidate task label tokens Y from demonstration lahels, with (y1,---,yWm) E V. without using the text-label mapping information to deduce the correct token. Task Learning (TL), in contrast, means learning the mapping from demonstration texts to task labels, f X + Y, to predict the only correct label for query q. To identify heads contributing to TR and TL, Lieberum al. (2023) compute W AN and & W for all heads (1,k), where W E TR M and W E TR are the unembedding matrix restricted Y and y". Heads with the highest element-wise sum 1 W d I.T are considered TR heads, and those with the highest & W are TL heads.
This approach has two problems. (1) For TR heads, Lieberum et al. (2023) study four-choice tasks where the full label space is "A", "B", "C", "D". In general settings, demonstration labels are arbitrarily chosen and may not capture the full task semantics. Changing labels from positive/negative to enouuralleunfaavurabbll for sentiment analysis prompts does not alter the task. but heads amplifying logits for positive/negative may not do so for favourable/umfavourable (2) For TL heads, the method ignores competition among label tokens: heads boosting y* may also boost incorrect labels Y \ fy*) }. disqualifying them as true task-mapping heads. A more precise approach must (a) capture task semantics beyond surface tokens and (b) evaluate contributions relative to competing labels.
We therefore propose the Task Subspace Logit Attribution (TSLA) method. For TR heads, we compute the TR score:
Proj Proj Proiwydkyullz d'yullz,
where Projwy WV (W W Y.T is the d x projection matrix onto span (W)). the subspace spanned by unembedding vectors of demonstration labels. This subspace should capture the unembeddings of all tokens that could potentially serve demonstration labels given the finding that LLMs encode inter-related semantics into subspaces (Saglam et al., 2025; Zhao et al., 2025) (further verified in Appendix F). Intuitively. TR heads are those whose outputs concentrate within this subspace rather than elsewhere, which is precisely what the TR score is designed to capture. We also have the following theoretical guarantee for this metric's effectiveness.
Theorem Let r= = V. Assume n distinct r-dimensional subspaces drawn i..dd from the Grassman- nian Gr d) are spanned by columns of Wy. If head (1,k) has TR score 7, then with probability at least (n 1)( 1
(2: d=))), alyx has the largest projected l2 norm onto span( W
among all such subspaces,
B(a.B) is the regularized incomplete beta function, monotone in x. Thus, large TR score implies the head output with high probability projects onto the task-related subspace more than any other subspace, thereby qualifying it as a TR head. This formalizes the intuition above and explains why the metric identifies the heads advancing task recognition. The proof is in Appendix B.
where [a Mixa.:;
For TL heads, we compute:
Avcyex/6> (W) W)
|Pro a'yll2
The numerator measures the alignment between the head output and the average correct-incorrect label direction, reflecting the logit gap head induces between the correct label and competing labels.
The denominator is the TR score. Since W w E span W for all y E Y {y*), the TL score has natural geometric interpretation with respect to the task subspace: it measures the component of the head output, after projection onto the task subspace, that lies along the average correct-inconeet label direction. Thus, the TL score isolates the contrastive within-subspace direction that favors y*
4

---

Published as a conference paper at ICLR 2026
over other competitor labels. Heads with high TL scores enlarge the correct-inconee logit gap. effectively rotating hidden states within the task subspace toward the correct label and away from others (see Figure (C)). This conforms to task learning, which identifies the correct label and excludes incorrect ones for an input. We provide further clues on this mechanism in Subsection 4. 4: TL heads allocate more attention query tokens, suggesting they absorb query semanties and match it to the label tokens to select the correct in-context mapping. Moreover, this TL score mitigates the DLA issue of interference from incorrect labels by disregarding heads that fail differentiate and instead raise both logits. To verify this, we provide the detailed comparison between TSLA and DLA in Appendix C by demonstrating that TSLA solves the practical limitations of DLA as the theory predicts. For each dataset, we use ICL prompts from the first 50 queries to compute TR and TL scores for each head with TSLA, summing across prompts. Heads are ranked by these scores to identify TR and TL heads.
Having identified TR and TL. heads with TSLA, we next ask whether these heads behave in the mech- anistic ways TSLA claims. In particular, TSLA predicts (i) TR and TL heads should affect different ICL components in a separable manner, and (ii) their outputs should drive qualitatively different geometric changes in hidden states. Accordingly, Subsection 4.1 examines TR/TL specialization and their relation to IHs, Subsection 4.2 tests separability causally via TR/TL-channel ablations, and Subsection 4.3 probes the predicted geometric mechanisms through steering and layerwise analysis.
EXPERIMENTS
Models We experiment on models with diverse architectures and sizes, including Llama3-8B, Llama3.1-88 1-BB, Llama3.2 2-3B (Grattafiori et al., 2024), Qwen2-7B, Qwen2.5-32B (Yang et al., 2024), and Yi-34B (01. Al et al., 2024). Unless otherwise noted, results are reported on Llama3-8B.
Datasets We evaluate on the following datasets: SUBJ (Wang & Manning, 2012), SST-2 (Socher et al., 2013), TREC (Li & Roth, 2002), MR (Pang Lee, 2005), SNLI (MacCartney & Manning, 2008), RTE (Dagan et al., 2005), and CB (De Marneffe et al., 2019). We further include an LLM- generated dataset introduced in Subsection 4.3, with curation details in Appendix D and Appendix K. We also experiment with the SubjQA dataset (Bjerva al., 2020) in Appendix I.3.
ICL setting We use 8-shot demonstrations for ICL. For implementation details (models, datasets, prompt templates, etc.), see Appendix D.
4.1
VALIDATING TR/TL HEAD SPECIALIZATION AND THE ROLE OF INDUCTION HEADS
Jaccard Coef
Guided by TSLA, we first test whether the identified TR and TL heads ex- hibit the predicted specialization and how they relate to previously stud- ied induction heads. Specifically, analyze the overlap, rank correlation, and consistency between TR and TL head rankings, and include IHs as reference point. Given the signifi- cance of IHs in mechanistic accounts of ICL (Zheng et al., 2024), we also ask whether IHs contribute primarily through TR, TL, or both. Follow- ing Todd et al. (2024) and Yang et al. (2025), we compute IH scores (i.c., the degree to which a head's 's attention pattern resembles that of IHs) as de- scribed in Appendix E.
0.
0.3
64
TR-TL
11-11 TR-41T
40.0
TL-II
TR-TH
20.0
Spearman's
Kendall's
5%
10%
20%
(a) Dataset-averaee Jaccard Co- (b) Conditional Mean Percent- efficient, Kendall's t, and Spear- age at four top thresholds for the man's p for TR heads, TI. heads, TR-IH and TH.IH pairs aver- and IHs the top 3% level. aged over datascts.
Figure 2: Overlap, correlation, and consistency of three at- tention head types averaged across datasets. (A) TR heads exhibit substantially greater overlap and correlation with IHs compared to TR-TL or TL-IH pairs. (B) Top IHS consis- tently rank higher in the TR ranking than in the TL ranking. Results for other models are in Appendix G.2.
Adopting the methodology of Yin & Steinhardt (2025), we report Jaccard 2 among the top 3% (further ablation studies for the choice of the threshold percentage are in Appendix G.I) of
For subsets A, C A. the Jaccard Coefficient
| Ag
5

---

Published as a conference paper at ICLR 2026
each head type to measure overlap at the top level. We also compute Kendall's and Spearman's 's p among the three rankings to evaluate the global correlations levels among the head types. Finally, we introduce Conditional Mean Percentage, which measures the average rank percentile of the top 3%, 5%, 10%, and 20% IHs within the TR and TL rankings. This metric bridges the local and global perspectives and answers the question of how strongly, on average, the top IHs exhibit TRR TL functionality, which is important in subsequent ablation- experiments in Subsection 4.2.
Strong association between IHS and TR heads Figure 2a shows that TR heads and IHs are highly similar: (1) their overlap at the top 3% is much larger
TI.
* IH-TR
5
than either TR-TL or TL-IH pairs, and significantly above random baseline (2) their rank correlations are consistently higher. By contrast, TR-TI. and
7-
9
11
= 13
TL.IH pairs show weaker overlap and correlation. This reveals how IHs, the magnitude of whose influ-
15-
δ 17
- 19
ence on ICL has been widely recorded, affect ICL:
- 21
23
they influence ICL mainly by enabling recognition
25
of the label space, rather than selectively amplifying
27
the correct label. It reconciles conflicting previous findings with some reporting IHs as recognizing cor- rect labels (Olsson al., 2022: Cho et al., 2025b), others mentioning "false induction heads" that mis- lead (Halawi et al., 2024; Yu & Ananiadou, 2024). It also echoes Yin & Steinhardt (2025), who observed IHs overlap strongly with "function vector heads" reinforcing the centrality of TR heads in ICL.
29
1 1 5 7 9113315 171921232527299
Head ID
Distribution of the top 10% TR heads, TL heads, and IHs on SST-2, TR heads occur significantly deeper than TL heads and IHs. Overlaps between TR heads and IHs are also more frequent in deeper layers. Results for other models are in Appendix G.3.
Figure 3:
Figure 2b further shows that top IHs consistently rank higher within TR rankings than TI. rankings. For example, the top 10% IHs correspond to roughly the top 20% TR heads but only the top 50% TL heads. This further justifies the large accuracy implications of ablating top IHs, which would imply ablating fairly high-ranking TR heads and the failure of task recognition.
Layer-wise distribution of special heads Figure 3 shows the per-layer distribution of the top 10% TR. TI., and IH heads for SST-2. We observe: (1) TR heads appear significantly deeper than both Tl. heads and IHs while the layer distributions of TL heads and IHS are more similar (see Appendix G.3 for sig- nificance tests) (We provide further explanations and experimental validations of why TR heads occur in deeper layers than TL heads in Appendix G.7). This partly echoes but also challenges Yin & Steinhardt (2025), who reported function vector heads as only "slightly deeper" than IHs. (2) The TR-IH overlaps are much greater than TL-IH or TR-TL overlaps, and occur primarily in deeper layers. This indicates that the correlation between TR heads and IHs is systematic rather than haphazard: the overlaps conform to the general trend of TR heads being concentrated in deeper layers, instead of reflecting coincidental matches with scattered TR heads that occasionally appear in early layers.
Atta. Demo Labels
TA Atta. to Query
0.075
0.050
0.36
0.05
0.4
0.03
0.025
0.02
0.17
0.2
TL
TR
Random
Figure 4: Average attention-weight distribu- tion of top TL and TR heads. TR heads attend more to demonstration labels- suuplosipn task irrognition-wheeee more on the query to extract semantics for correct prediction. Results for other models are in Appendix G.5.
TL heads focus
Attention distribution of TL and TR heads A second property of the identified heads is their attention distribution, which reveals how they operationalize TL. and TR. We quantify this using two metrics: (1) the total attention weight assigned to demonstration-Labe tokens, reflecting how much a head incorporates information about the task label space (supporting TR); and (2) the total attention weight assigned query tokens, reflecting how much head integrates the semantic content of the query (supporting TL). As shown in Figure 4, the top 3% TR heads allocate substantially more attention demonstration labels than TL or random heads (the small magnitude is expected
Forr two random 4% subsets, the expected Jaccard Coefficient is asymptotically * or 0.0152 for k 4 Heads revealed to have greatest impact on correct label logits through ablations.
6

---

Published as a conference paper at ICLR 2026
Attention Weseht
0.2
head
6.15
8.2
6.0
Toleen
Figure 5: Attention distribution of the top TL and TR head identified on SST-2 over an ICL prompt.
given the scarcity of label tokens in long ICL, contexts), indicating that TR heads broadly summarize task-label information. In contrast, TL. heads exhibit far more local attention pattern, attending strongly to the query tokens to extract their semantic meaning and support correct label inference. Direct visualization (Figure 5) of the top SST-2 TL and TR heads on an SST-2 ICL prompt using the first query in the test set further highlight this contrast: TR heads focus on demonstration labels, whereas TL heads primarily attend to the queries. See Appendix G.6 for results using more queries.
Cross-datase correlation To examine whether TR, TL, and IH heads generalize across tasks, we measure Jaceard, Kendall's r, and Spearman's 's p among the top 3% heads on seven datasets, averaged across (2) = 21 dataset pairs. As shown in Figure 6, TR heads and IHs exhibit higher cross-task overlap and correlation than TL heads. This underscores TR heads (and IHs) as task- invariant mechanistic foundations for label-space recognition, upon which TL heads specialize to learn dataset-specifty mappings. We further explore this discovery in Appendix H.4.
4.2 FINE-GRAINED ABLATION TESTS OF TR/TL SEPARABILITY
Having established behavioral specialization in Subsection 4.1, we now provide causal, fine-grained validation. TR and TL heads realize distinct ICL components, selectively ablating them should yield different signatures for task recognition versus task learning. Accordingly, we ablate TSLA-identified TR and TL heads to test separable contributions. Prior work mainly evaluated ablations by overall ICL-accuracy drop (Crosbie Shutova, 2024), which indicates performance change but obscures which ICL component is disrupted. To disentangle these effects, we introduce the Task Recognition Ratio (TR ratio), defined as the proportion of predictions within the in-context label set. Formally, for m ICL prompts with predicted labels Sm
TR ratio-
Σ SeY
Since accuracy upper-bounded by the TR ratio, the metrics together let us separately evaluate contributions of TR and TL heads. We conjecture: (1) ablating TR heads should reduce both accuracy and TR ratio, while (2) ablat- ing TL. heads should reduce accuracy but leave TR ratio largely intsct-causing performance to approximate ran-
Kendall Spearman Jaccard
1.0
Min Normaliz Norma l
TR
0.61
0.79
0.36
0.8
0.6
TL
0.09
0.13
0.12
0.4
dom guessing over all candidate labels with expected ac-
curacy [ 5. As control, we also ablate 3% of randomly chosen heads disjoint from the identified TR/TL sets. We report the results averaged over the seven datasets.
IH
0.76
0.92
0.47
0.2
0.0
Figure 6: Pairwise overlap/correlation of TR/TL heads and IHs across datasets. TR heads and IHs are task-consistent, whereas TL heads vary widely. See Ap- pendix G.4 for other models.
Separability of TR and TL funetionality Figure 7 con- firms conjecture. Removing top TR heads collapses the TR ratio from nearly 100% to ~20% leading to a drastic accuracy drop. In contrast, removing top TL heads
Foor the seven datasets with having 2 labels, 2 having 3 labels, having 6 labels, the average random guessing level 40.48%
7

---

Published as a conference paper at ICLR 2026
TR Ratio
0.0
ICL
W/O TL
W/o TR
(a) Ablating with shuffled demonstration texts.
(b) Ablating with relabeled demonstrations.
Figure 8: Dataset-average effects of ablating TR and TL heads under input perturbations. (A) Shuffling character order of demonstration texts destroys TL, making TL ablation negligible while TR ablation still matters.
(B) Relabeling demonstrations alters the label space, thereby greatly reducing the impact of TR head ablation.
lowers accuracy by ~30% but only slightly reduces the TR ratio (by ~00%) This highlights key property: separability, i.e., TR and TL can be independently controlled and intervened upon, consistent with the conclusions in Pan et al. (2023) achieved through input perturbations. (see Appendix H.1 for other models).
TR heads, IHs, and implications for zero-shot Ablating IHs produces a pattern closely resembling TR head ablation: large accuracy losses primarily due to failed task recognition. This supports the conclusion that IHs influence ICL mainly by strengthening TR. Likewise, the root cause of poor zero-shot performance is insufficient task recognition. Thus, restoring ICL-level accuracy in zero-shot settings hinges on activating the TR functionality- a question we revisit Subsection 4.3.
Testing independence via input perturbations
ACC D72 TR
To complement head ablations with an orthog- onal causal test, we perturh ICL inputs lectively disrupt TR or TL. Following Wei et al. (2023); Pan et al. (2023), we consider two cases: Case 1: Keep labels unchanged but randomize the character order of demonstration texts (e.g., "I like it: positive" + "tkl iiel: positive"), This destroys TL, since no meaningful mapping from such nonsensical texts to labels remains. Case 2: Keep texts unchanged but replace demonstration labels with arbitrary tokens (c.g.. "negative" "0", "positive" "1"), thereby altering the label space recognized by TR heads.
0.0
ICL.
ZSL
WIO III
W/O
WAO TIL
W/O Rand.
Figure Effects of ablating the top 3% of different heads, averaged across datasets. TR head ablation severely reduces both accuracy and TR ratio, while TL head ablation primarily reduces accuracy.
We hypothesize that: the TR heads and TL heads maintain sufficient independence, then in Case ablating TL heads should have little effect (TL is already disabled), while in Case 2, ablating TR heads should matter less (the original TR functionality nullified).
As shown in Figure 8a, when texts shuffled, TL ablation barely matters while TR ablation remains devastating. Conversely, in Figure 8b, TR ablation has little effect, compared to the significant impact shown in Figure 7, since the recognized label space has shifted, while TL. ablation behaves in the standard case. These results confirm the robustness of TR/TL independence across diverse ICL settings (see Appendix H.2 for other models).
4.3 STEERING WITH TR/TL HEADS: FUNCTIONAL AND GEOMETRIC INSIGHTS
With fine-grained ablations establishing that TR and TL heads are necessary for their respective channels, we now test their sufficiency and TSLA-predicted geometric roles by injecting their outputs. Ablations show what happens when heads are removed, but not what happens when they are added. We therefore complement them with steering experiments that ex- amine how TR and TL head outputs affect zero-shot behavior and hidden-state updates. Specifically, we test their suitability as task vectors (TVs) (Todd et al., 2024; Hendel et al., 2023) by extracting their outputs at the final token position from ICL prompts, sum- ming them across prompts, and injecting them into the residual stream of zero-shot inputs to evaluate
Original
After
Rand.
325.
TR TV
9.2%
049.4%
TL TV
201
10.0%
20.0%
30.0%
40.0%
Figure 9: Zero-shot accuracy gains from steer- ing with task vectors constructed from TR, TL. or random heads. TR-based task vec- tors consistently recover ICL-level accuracy,
while TL-based vectors have weaker effects.
8

---

Published as a conference paper at ICLR 2026
recovery toward ICL-level performance. We use the top 3% TR/TL heads, compare against 3% random heads, and follow Appendix J.1 to construct task vectors.
Task recognition as the key to zero-shot failure Figure 9 minrors our ablation findings (Figure 7): poor zero-shot performance stems primarily from weak task recognition. Injecting TR-based TVs restores this functionality and improves performance. TL-based vectors are less effective, reinforcing that TL operates based on TR (see Appendix 1.1 for other models).
Task-type dependence of steering effectiveness Note that the relative ineffectiveness of TL heads as task vectors can be partly attributed to the classification datasets we use, where performance is tightly linked to task recognition and effectively upper-bounded by the TR ratio. In contrast, generation tasks differ fundamentally in that no fixed label space exists the label space is indefinite and, in principle, infinite. As a result, model success in such tasks is less constrained by recognizing closed set of labels, and instead depends more on learning and applying the correct input-output mapping. To examine this scenario, we consider a cnntiment-controlled review generation task with prompts such as: "Write positive/negative review of a movie within 30 words." Labels are coherent reviews with the desired sentiment". We identify TR and TL heads on ICL-styled prompts from this task following Appendix J.2, and use their outputs task vectors to influence the zero-shot generations. An LLM evaluator is then used to rate generations from to 10 based on sentiment adherence and language coherence. As shown in Figure 10, TL-based vectors significantly outperform TR- -'R-based and random vectors, consistent with TL heads capturing mappings from demonstrations to the sentiment values and semantic coherence of the labels (see Appendix I.2 for other models). Nevertheless, given the relative simple and structured nature of this task and the fact that the label reviews are generated by GPT (OpenAI et al., 2024), we extend our investigation regarding the TV effectiveness using the SubjQA dataset (Bjerva et al., 2020), which we detail in Appendix I.3.
Geometric effects of TR and TL outputs To understand the significance of TR/TL heads in ICL a finer level than task vector experi- ments, we invoke the geometric analysis of hid- den states (Kirsanov al., 2025; Yang et al., 2025), which analyzes the evolution of ICL hid- den states and the role different components. Concretely, given an ICL prompt. extract the summed outputs of the top 3% TR or TL heads, revert to the hidden state at earlier layer, and steer it with these outputs. This mimics how head outputs added the residual stream during layer progression inside the model. We measure two geometric metrics before and after steering: (1) Logit Difference: inner product of the hidden state with the mean unembedding difference between correct and incorrect labels, reflecting label discrimination. (2) Subspace Alignment: cosine similarity between the hidden state and the subspace spanned by label unembeddings, reflecting alignment with task-related semanties
Mean + SD
7.99 1
5
3.92
5.12
3,37
4.44
as TV TR IVRand. TV Figure 10: Ratings in the review generation task when steering with TR, TL, or random TVs. TL vectors yield the largest improvements, reflecting their strength in capturing in-context mappings.
ICL
zs
The results in Figure 11 demonstrate the specialized geometric effects TR and TI. heads have the evo- lution of hidden states (for other models, see Ap- pendix 1.4). Steering with TR outputs causes hidden states to align significantly more with the task sub- space. contrast, TL outputs adjust the hidden state to align better with the unembedding direction of the correct label in the task space but not with the task subspace overall. This leads to conjecture that TR outputs are well-aligned with the task subspace, thus increase hidden-state alignment with the subspace after addition by decreasing the angle in between. By contrast, TL heads create pure rotation toward the
Logit Difference
Sulspace Aligument
0.30
5.0
4.85
0.25
0.20
4.0
3.92
0.15
0.15
Original +TI. +TR + +Rand.
Figure 11: Geometric effects of TR and TL steering. TR outputs enhance alignment with the task subspace, while TL outputs rotate hidden states toward the correct label unem- bedding within the subspace.
Example positive review: "Bold experimental narrative structure defies genre conventions delightfully. Sociocconomic themes challenge viewers' perceptions thoughtfully and respectfully." See Appendix J.3 for full definitions.

---

Published as a conference paper at ICLR 2026
p=0.94 p=7.07c-16
5=0,33 p=0.00
biEr
ba update: Subapacn
0.20
TE oulput: Louit
0.15
0.2
0.10
Till output
0.05
0.1
10
15 Layer
20
25
30
10
15 Layer
20
25
30
(a) Correlation between hidden state updates and TR head outputs in subspace alignment. Strong correla- tion confirms TR heads main drivers of alignment.
(b) Correlation between hidden state updates and TL head outputs in logit difference. TI. heads consistently drive discrimination toward correct lahels.
Figure 12:
Layerwise verification of TR and TL geometric effects. TR heads enforce alignment with the task subspace, while TL heads enforce rotations toward correct label directions.
correct label unembedding direction, fine-tuning hidden-state orientation toward the correct label without boosting subspace alignment. See Appendix 1.4 for other models.
Layerwise verification of geometric influence To validate that TR and TL heads indeed primarily examine hidden state updates under ICL across layers. At each
drive these geometric dynamics, layer, we compute the mean subspace alignment of top-3 TR heads (i.e. heads with top-3 TR scores the layer) outputs and the mean logit difference of top-3 TL heads outputs, then correlate them with the same metrics computed on the full hidden state updates. Since head outputs contribute directly to layer updates, their correlations with hidden geometry across layers indicate how strongly TR and TL heads drive the dynamics. As shown in Figure 12, the correlations are strong, confirming that TR and TL. heads dominate layerwise geometric shaping of hidden states. For other models see Appendix I.5. Additional ablation- based verification is provided in Appendix I.6.
TR heads align to task space, TL heads rotate within it To support our geometric intuition from Figure 11 that TR heads foster alignment while TL heads perform rotation, two geometric measures of TR and TL head out- puts. (1) Subspace Alignment cosine simi- larity with the task subspace, and (2) Projected Discriminant Alignment cosine similarity with the mean unembedding difference between the correct and incorrect labels after projection onto the task subspace. These measures dissect the geometric effects of head outputs into steer- ing towards the task space and steering within the task space, enabling more fine-grained verification of the heads' distinct effects (Figure (B), (C)). Figure 13 shows that TI. head outputs have high cosine similarity with the mean unembedding difference after projection, confirming that TI. heads, when restricted to the task subspace, propel rotation from wrong-label to correct-label unembedding directions. The high cosine similarity between TR heads and the task subspace itself strongly ev- idences their capability to steer hidden states towards the task subspace and support prediction of task-related labels (see Appendix I.7 for more models).
Projected Discr. Align.
E/A Subspace Alignment
0.44
0.4
consider
0.2
0.09
0.00
0.03
0.0
TL.
TR
Rand.
Figure 13: Decomposed geometric effects of TR and TL outputs. TR heads align hidden states to the task subspace; TL heads rotate states within the subspace toward correct label directions.
CONCLUSION
We presented unified framework reconeiling component and holistic views of in-context learning (ICL) by identifying attention heads specialized for Task Recognition (TR) and Task Learning (TL.). Using TSLA, we showed that TR heads align hidden states with the task subspace recognize labels, while TL heads rotate states within it toward the correct label. Ablation experiments confirmed separable roles: removing TR heads collapses task recognition, whereas removing TL heads mainly reduces accuracy. Steering experiments showed task dependence: TR-based vectors are crucial for classification with fixed labels, while TL-based vectors dominate in open-ended generation. Geometric analyses supported these findings, attributing alignment to TR heads and discriminative rotations to TL heads. Our results also clarify induction heads and task vectors as TR manifestations. Together, this work establishes TR and TL heads as mechanistic foundations of ICL..
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported by JST FOREST Program (Grant Number JPMJFR232K, Japan) and the Nakajima Foundation. We used ABCI 3.0 provided by AIST and AIST Solutions with support from "ABCI 3.0 Development Acceleration Use".
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.
Johannes Bjerva, Nikita Bhutani, Behzad Golshan, Wang-Chiew Tan, and Isabelle Augenstein. Subjqa: A dataset for subjectivity and review comprehension, 2020. URL https: //axiv..x.vvv org/aba 2004 14283.
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 ht /dl. al..amm... org/ /t////o/aatt// abss/ 10. 55.5/ 00755//3497222 /t/.//449572244495888 3495883.
Vivien Cabannes, Charles Arnal, Wassim Bouaziz, Xingyu Yang, Francois Charton, and Julia Kempe. Iteration head: A mechanistic study of chain-of-then -f-thugtt Advances in Neural Information Processing Systems, 37:109101-109111 2024.
Hakaze Cho. Mariko Kato, Yoshihiro Sakai, and Naoya Inoue. Revisiting in-context learning inference circuit in large language models. In The Thirteenth International Conference Learning Representations, 2025a. URL https::///. /tpp//ppeeeee///eee/eeeeeeeeeeee wew..wwww.... net. / /form?? i d= YDVVQQ
Hakaze Cho, Yoshihiro Sakai, Mariko Kato, Kenshiro Tanaka, Akira Ishii, and Naoya Inoue. Token- based decision criteria are suboptimal in in-context learning. In Luis Chiruzzo, Alan Ritter, and Lu Wang (eds.), Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics; Human Language Technologies (Volume 1: Long Pupers), pp. 5378-5401, Albuquerque, New Mexico, April 2025b, Association for Computa- tional Linguisties. ISBN 8888188--9888888 naacl-long 278/.
URLhttps:/ http://///aaaaaahhhaaccco wwwp:/y..../...
Bilal Chughtai, Alan Cooney, and Neel Nanda. Summing up the facts: Additive mechanisms behind factual recall in llms. arXiv preprint arXiv:2402.07321, 2024.
Joy Crosbie and Ekaterina Shutova. Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv:2407.070) 2024.
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: //link. //linkksprinner....... com/ http:/ceettttt/////11111
Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: In- vestigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pp. 107-124, 2019, URL https:/ //ojs. .u.. uuii-kontaan anz de/sub/ index. php/sub/ artiiatettal//aa. wwtp:/vewweeeeeeeeeeeeeeeeeeeeee 601.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, Baobuo Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning, 2024.
/ waaaaiivvooo org.vv.rrrraaeeee abs 2301 00234.
11

---

Published as a conference paper at ICLR 2026
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021. URL https:// reu iesitt..e .tw.iits.pbbb/b.. .ts.ppbb222111 hataa///faawwwwo//aommaa index.idex...// ht.ml
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, and Abhinav Pandey et al. The llama 3 herd of models, 2024. URL https: //wxxvv..vvvv/vaa//////aaaaaiwv org/
Danny Halawi, Jean-Stanislas Denain, and Jacob Steinhardt. Overthinking the truth: Understanding how language models process false demonstrations, 2024. URL https: /ttp/a/aivv/////////////////all/ wwwp:iii..oraaoooooeeooeeee 2307. 09476.
Roce 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- tational Linguisties: EMNLP 2023, pp. 9318 9333, Singapore, December 2023. Association for Computational Linguistics. doi: 10. 1/222 www.a/iiiiiinnennl.222222 URL https: //aclanthology. wwwphhhhoooooorg/2223333 wrg/2023. fi findingssnn emmlp. 624/.
Jiachen Jiang, Yuxin Dong, Jinxin Zhou, and Zhihui Zhu. From compression to expansion: A layerwise analysis of in-context learning, 2025. URL https: /trppi://axxvv... wwtp/bs////ooooooooooo.o.....a// /bb////////////// hatgg/abb/255 /2505. 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: A mechanism for interpreting and mitigating knowledge conflicts in language models, 2024. URL atttps://arxiae org/ abs/ 2402. 18154.
Patrick Kahardipraja, Reduan Achtibat, Thomas Wiegand, Wojciech Samek, and Sebastian La- puschkin. The atlas of in-context learning: How attention heads shape in-context retrieval augmen- tation. arXiv preprint S.ti::2505.188000 2025.
Artem Kirsanov, Chi-Ning Chou, Kyunghyun Cho, and SueYeon 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. 1888 Albuquerque, New Mexico, April 2025. Association for Computa- tional Linguistics. ISBN 979-8-89
879-8-89176-111 19176-195-77 URL htLpstt:::/// http://:://.laanoooooo...ooo
sindings-nnacl 100/.
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Confer- ence on Compulational Linguistics, 2002. URL httph:tt////cc http://wwwwww/.. www..acllee.. org/ antho c02-1150.
Yuxuan Li, Declan Campbell, Stephanie CY Chan, and Andrew Kyle Lampinen. Just-in-time and distributed task representations in language models. arXiv preprint arXiv:2509.04446, 2025.
Tom Licberum, Matthew Rahtz, János Kramár, Neel Nanda, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik. Does circuit analysis interpretability scale? evidence from multiple choice capabilities chinchilla, 2023. URL https:://aaaii// wawaiiiiivvv..... http//aawvvvv..aaaaa/aa / 0..37.09444 09458.
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: htapps://amee //arxiv. org/ brg/ abs 111./3111100 06668.
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 Compulational Linguistics (Coling 2008), pp. 521-528, Manchester, UK, August 2008 Coling 2008 Organizing Committee. URL https: /aa//aalaanthhh hwc.t//cllaa...oo............/.. org/CO8-1066/
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vec- style vector arithmetic, 2024. URL https: //wxivivvvv/.. www.iiioogg//..oooo.o....o.oooee abs 22005 2300.16136 16130.
12

---

Published as a conference paper at ICLR 2026
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role demonstrations: What makes in-context learning work? arXiv preprint arXhv:2202.128377 2022.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL https: // //aa/aaa// wawaiivvoooo org/ abs/ / 2209. 11895.
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, et al. Gpt-4 technical report, 2024. URL https: /ttp:v.//axxvv.oorr abs/ 2303 08774.
Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen. What in-context learning "learns" in-context: Disentangling task recognition and task learning. In Anna Rogers. Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 8298 8319, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/ URL https: : /wwpaaaalnlncca...aaaaaaaaaallla .rgg20233 wow.tttooooooooooo22233333 http://22...fiiinnns--aaaaa findings-acl.
1.n022..indiggs-adee 527/.
Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. arXiv preprint cs/0506075, 2005.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask leamnrs. OpenAl blog, 2019. URL
https: https:://AAmcccppssssks
language- pdd..
Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. Identify- ing semantic induction heads to understand in-context learning, arXiv preprint vxXv:2402.13355 2024a.
Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. Identifying semantic induction heads to understand in-context learning, 2024b. URL https:
a arriiv. www.rvvv.g...e.............ooeee
owooorgg///a.... 2402. 13055.
Baturay Saglam, Paul Kassianik, Blaine Nelson, Sajana Weerawardhena, Yaron Singer, and Amin Karbasi. Large language models encode semantics in low-dimensional linear subspaces, 2025.
URL https:/p///// https///////xx /ttp///avvvvo .rgg abs 20..00009900
Aaditya K Singh, Ted Moskovitz, Felix Hill, Stephanie CY Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. In Forty-first International Conference on Machine Learning, 2024. URL https: /ttpsv///aavv../aa www.oogg//s...o..oooo. abss 2404 07129.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In David Yarowsky, Timothy Baldwin, Anna Korhonen, Karen Livescu, 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 https::p///// /aaaaaaaaannltooooeeeeaeaaaeeele org/ 1ttp//D331110000
Jiajun Song, Zhuoyan Xu, and Yiqiao Zhong. Out-of-distribttion generalization via composition: lens through induction heads in transformers, Proceedings of the National Academy of Sciences,
(2(6))22217182122 2025. URL https:::///tttttt org abs 2408 09503.
Chenghao Sun, Zhen Huang, Yonggang Zhang, Le Lu, Houqiang Li, Xinmei Tian, Xu Shen, and Jieping Ye. Interpret and improve in-context learning via the lens of input- mappings. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 3873-3895, 2025.
13

---

Published as a conference paper at ICLR 2026
JTiaaxiangg Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tun- ing for language-mode model-as-eeeeeeee as--a-erveicee In Proceedings of the 39th International Conference Baltimore, Maryland, USA, 2022. ACM. URL ml press /v1 /2/ 62/sun22e////e/mmmmmmeem// sunaee.
Machine Learning, pp. https: /ttp:///roceeiieee
Eric Todd, Millicent L. Li, Arnab Sen Sharma, Aaron Mueller, Byron C. Wallace, and David Bau. Function vectors in large language models, 2024. URL hwtph:tt///aaaeeee wwwpaii.org//bs.o.o/ abs /2310. 15213.
Sida Wang and Christopher Manning. Baselines and bigrams: Simple, good sentiment and topic classification. Haizhou Li, Chin- Yew Lin, Miles Osborne, Gary Geunbae Lee, and Jong C. Park (eds.), Proceedings of the 50th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers). pp. 90-94, Jeju Island, Korea, July 2012. Association for Computational Linguisties. URL ht.ps:: a anthoa ooyy wrg. 2
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang. Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently. 2023. URL https: /ttpaaaxivv////////////////aaava org/ abs / 22003 03846.
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, Jumyang 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 arXiv:2407. 200244077.677,,, 2024.
Haolin Yang, Hakaze Cho, Yiqiao Zhong, and Naoya Inoue. Unifying attention heads and task vectors via hidden state geometry in-context learning, 2025. URL https: htlLhhtss://aaxe /ttps://axxxvvooooee abs/2505 18752.
Kayo Yin and Jacob Steinhardt. Which attention heads matter for in-context learning? arXiv preprint arXiv:2502.1 /40/0, 2025.
Zeping Yu and Sophia Ananiadou. How do large language models learn in-context? query and key matrices of in-context heads are towers for metric learning, arXiv preprint 4r210:2402.02872. 2024.
Haiyan Zhao, Heng Zhao, Bo Shen, Ali Payani, Fan Yang, and Mengnan Du. Beyond single concept vector: Modeling concept subspace in llms with gaussian distribution, 2025. URL
https: /arxiv... horg/abs///0 /2410.2100 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.03755 03752, 2024.
14

---

Published as a conference paper at ICLR 2026
Appendices
A STATEMENT OF LLM USAGE
In this work, LLMS used help with writing, experiment coding, and visualization of the results. L.MMs are also used to produce results in one of the experiments, as explained in Subsection 4.3 and Appendix K.
B PROOF OF THEOREM
Let S E Rdxr be one of the distinct r-dimensional subspaces in span Wo) drawn uniformly i.i.d. from the Grassmannian Gr(r,)). Denote Ps the projection matrix of S. Let || be
||
W onto S. Since the uniform distribution
the projected norm of the normalized head output over Gr(r,d) induced by the Haar measure over the orthogonal group O(d), the distribution cotation-invariant i.e., multiplying S by an orthogonal matrix U Rd does not change its distribution. Because orthogonal transformations preserve angles, we also have
Tayale
PusU Pusuawall |Psa'elll ||axxel2
|| |
Hence, without loss of generality, pick an U such that U alx
e1, the unit vector in the first
||ab)
coordinate, with é' Puse llz having the same distribution as c.
Since US span( .... U,), where V1,..., 8, the first columns of Haar orthogonal matrix V. let V, [u,...., wy Then Pus V,V,", and we have
= V,V, VIV, eill2
Σ (el 0/) Σ V
where V1, denotes the first of V,. Since V, is Haar orthogonal, Vi,: is uniformly distributed on gd-1 and has the same distribution Tolla distribution as
with g ((0,I) Therefore, Σ V2 has the same
Σ
of
Li
X
g
r 1
X2 + x2
since gi (0, for all i. Because
we have
x²
Beta( (5 1
x² + x2-
If c2 Beta( { then the tail probability is
B(x2 c d-r
Pr(c = ( 11
B d=1
where B the Beta function. Since the TR score of the head y. the probability that |Psal 2 y
1-1
(1, d.r)
Because there are subspaces alongside W., the probability that the head output has the largest projected norm on w
1 1- (n - 1) 1-1
2 (5, (5)
1aypl2
via the union bound.
15

---

Published as a conference paper at ICLR 2026
C ABLATION EXPERIMENTS REGARDING THE IDENTIFICATION OF TR AND TL HEADS
In this section, we demonstrate the advantage of our Task Subspace Logit Attribution (TSLA) method over the naive approach of selecting TR and TL heads based on a'yx W and d'y, W i.e., Direct Logit Attribution (DLA) to the demonstration labels and the correct label. Specifically, Figure 14 shows the consequences of ablating the top 3% TR and TL heads identified via DLA, averaged across datasets Llama3-8B. While ablating the identified TR heads achieves the intended effect of disabling task recognition by reducing the TR ratio, ablating the identified TL heads fails to induce the expected outcome of driving ICL toward random guessing over the label space. This indicates that the DLA approach cannot isolate distinct mechanistic causes for the TR and TL components of ICI., and reflects its inability to correct identify the real TL. heads. Instead, it largely identifies heads that broadly amplify the logits of all demonstration label tokens, which may also increase the correct label logits but still primarily function through task recogmition rather than true label differentiation.
To validate this statement, and following the setup of Figure 2, we report the dataset-averaged Jaccard Coefficient, Kendall's t, and Spearman's 's p values between TR/TL heads identified by DLA and those identified by TSLA, as well as their relationship with IHs. As shown in Figure 15, both TL and TR heads selected via DLA strongly overlap with the TR heads identified by TSLA at the 3% level. This corroborates our conclusion that DLA fails to effectively recover genuine TL. heads. Furthermore, the weak correlation between the TR/TL sets obtained from the two methods is reinforced by Figure 16, which displays overlap, correlation, and consistency analyses between DLA TR/TL heads and IHs. The strikingly high consistency between DLA TR and TL heads across all three metries demonstrates the lack of meaningful distinction between them. Meanwhile, the low correlation between DLA heads and IHs highlights another limitation of DLA: it cannot provide mechanistic explanations for the well-documented importance of IHs in ICL.
Finally, to justify our second critique of the DLA approach in Subsection 3.1 regarding its sensitivity to the concrete set of demonstration labels as a hyperparameter and its inability comprehensively capture the task semantics, we consider the following experiment on SST-2. We replace "positive" and "negative", i.e. the default demonstration labels used to create ICL prompts from the dataset, with "unfavourable" and "favourable", which do not alter the essence of the task. Then we test how the ablation of TR heads identified with DLA and our TSLA using the ICL prompts with the original labels will impact the ICL accuracy and TR ratio with the new labels. The results in Figures 25-30 confirm the robustness of our TSLA method against demonstration label shifts in identifying TR heads. On all models except Qwen2.5-32B, ablating the TR heads selected using our TSLA approach causes a significantly larger impact on ICL performance and TR ratio with the new demonstration labels the SST-2 dataset, with the gap being most prominent for the three Llama family models.
D IMPLEMENTATION DETAILS
Models We use the official HuggingFace implementations of all models. Models with more than 10B parameters are quantized to 4-bit for efficiency.
Datasets We use the official HuggingFace implementations of all datasets, except for the Review dataset, which we curated ourselves. The Review dataset was generated using ChatGPT-4o (OpenAI et al., 2024) and contains 200 datapoints. Each datapoint consists of prompt instructing the model to generate a movie review in the format: Write a positive review for movie. The positive review should be within 30 words. " The 30-word limit was chosen to set the max new tokens parameter (set to 45) when calling the generation function. Labels are ChatGPT-4o-generrate multiple aspects in the requested positive/negative tone. For example: "Bold experimental narrative structure defies genre conventions delightfully. Rich orchestral score enhances every pivotal moment. Progressive messages inspire reflection equality and justice. Raw vulnerability on screen fosters sincere emotional investment." " as a positive review. Details of dataset curation are provided in Appendix K. The dataset is balanced, with 100 positive and 100 negative reviews.
reviews that comprehensively assess a movie from
ICL setting For each dataset (except the Review dataset), we select demonstrations from the training set and queries from the test set, or the validation set if ground-truth test labels are unavailable. For
16

---

Published as a conference paper at ICLR 2026
ACC
TR Ratio
1.0
0.8
0.6
0.4
0.2
0.0
ICL
ZSL
W/O TL DLA W/O TR DLA
Figure 14: Effects of ablating the top 3% TR and TL heads identified using DLA, averaged across datasets on Llama3-8B. While TR heads reduce task recognition as expected, TI. heads not replicate the behavior predicted for task task-learning components.
Jaccard Coef
0.6
TL DIA-TL DLA-TR
0.5
0.3
TR DIA-TL
0.2
TR DLA-TR
0.1
0.0
Spearman's
Kendall's
Figure 15: Dataset-averaged Jaceard Coefficient, Kendall's t, and Spearman's p between TR and TL. heads identified using DLA and TSLA the top 3% level. DLA heads overlap substantially with TR heads, confirming their inability to recover distinct TL heads.
Jaccard Coef
0.8
TR DLA-TL,DU DIA
0.5
DLA-IH
45.0
0.2
0.0
40.0
Spearmán's
Kendall's
3%
5%
10%
20%
(a) Dataset-averaged Jaceard Coefficient, Kendall's and Spearman's 's p values for TR heads TL heads identified using DLA and IHS the top 3% level.
(b) Conditional Average Percentage at top 3%, 5%. 10%. and 20% levels for TR (DLA)-IH and TI. (DLA)-IH IH pairs, averaged across datasets.
Dataset- averaged overlap, correlation, and consistency analyses of TR and TL heads identified using DLA and their relationship with IHs. Results show high redundancy between DLA TR and TL heads and weak association with IHs, underscoring the limitations of DLA in separating TR and TL mechanisms explaining IH significance.
Figure 16:
17

---

Published as a conference paper at ICLR 2026
demonstration selection, retain at most the first 10,000 training examples. For evaluation, the first 1,000 test or validation examples. For the Review dataset, we use the first 50 examples for demonstration selection and the remaining 150 for testing. Prompt templates used construct ICL prompts are listed Table 1.
Devices All experiments were conducted
H200 GPU.
Random label mappings For the experiment in Appendix G, where demonstration labels are replaced with numbers, use the mappings in Table 2.
When demonstration labels are replaced with numeric symbols, we also modify the prompt templates in Table 1. Specifically, for SNLI and CB, "True maybe false" is replaced with "O or 1 or 2," and for RTE, "True or false" is replaced with "O or 1."
Flipped label mappings For the experiment in Subsection 4.2, where demonstration labels flipped, we use the mappings in Table 3.
E EXPERIMENT DETAILS CONCERNING THE IDENTIFICATION OF IHS
For each dataset, we use the first 50 queries to identify the top IHs. Let the queries be 11,-.-9950 where qi has token length s(g:). For each li the LLM outputs an attention tensor Attn; E xxxx)xx((g)xs(f) with being the number of layers and N₂, the number of attention heads per layer. The ny-th head in layer has an attention matrix Attm(I,mp)) E R((q))a(q)) where Attn(I,ma).j denotes the attention from the k-th token to the j-th token in X₁.
Identification of IHS For each qi we randomly sample 8 demonstrations and prepend them q The resulting ICL prompt, Qi (length s (Q))), follows the format (t,1): (v..)...., (tis) (vis). (gi) where (t,x) is the sentence part of demonstration & (e.g.. "I like this movie. Sentiment"), and (yra) the label (e.g., "positive"),, separated by colon. (q:) is the sentence for the query. At the position of the final colon, an IH is expected to attend to tokens after previous colons-that is, the label tokens (v.))...., (yr8)) Let Ii be the set of label token indices in Q; The IH score for head (1, over the 50 queries is defined E Σter, Attn(I, mh (Q.)). i.e., the total attention head assigns at the final "!" position to the positions of all the label tokens, summed over all 50 queries. We calculate the IH scores for all (1,1a) pairs and choose the top 3% attention heads as the identified IHs.
F VERIFYING WHETHER THE UNEMBEDDINGS OF SEMANTICALLY RELATED LABELS EXIST IN COMMON SUBSPACES
The theoretical and practical validity our TSLA method relies on the assumption that LLM unem- beddings of semantically related tokens lie in common low-dimensional subspaces, an observation supported by prior work (Zhao et al., 2025; Saglam et al., 2025). To further verify this assumption, we conduct the following experiment. We consider set T of semantically related sentiment tokens, with T =lpositive, negative, favourable, unfavourable, pleasing, disappointing, enjoyable, unpleasant, satisfying, dissatisfying, delightful, distasteful, uplifting, depressing, entertaining, regrettable, ex- cellent, terrible]. For each token / T, we measure the norm of its unembedding projected onto the subspace spanned by the unembeddings of the remaining tokens in i.e.,
Proj
W I|2.
w
Then we randomly draw a token set TV of the same size as TV {t} from the full vocabulary (i.e.,
[T| = /T {t)), and compute
[|Proi W ll.
unembeddings of semantically related tokens indeed concentrate in common subspaces, we should observe
|[rojj W ll2
[|Proj W l|2.
meaning that the unembedding of aligns more strongly with the subspace spanned by its semantic peers than with a randomly selected subspace. We compute both norms for every r E T across all models and report the results in Figures 19-24. The results show that projection norms are indeed
18

---

Published as a conference paper at ICLR 2026
substantially larger when is projected onto the subspace spanned by its semantic correlates, thereby validating the core assumption behind our TSLA approach. To further support these visual findings, conduct Wilcoxon signed-rank test the paired norm values. The results in Table show that the differences are highly statistically significant.
G SUPPLEMENTARY MATERIALS FOR SUBSECTION 4.
G.1 ABLATION STUDIES FOR THE PERCENTAGE THRESHOLD OF SELECTING TOP TL AND TR HEADS
Throughout the paper, select the top 3% TL and TR heads based their respective scores and use them in subsequent analysis, ablation studies, and steering- based experiments. To assess whether this particular threshold influences our conclusions, increments of 1 1%. For Llama3-8B, which has 1024 heads, each % corresponds to 10 heads. For ablation experiments,
sweep the percentage from 1 1% to 10% in
remove different percentages of top TL or TR heads and evaluate the effect on accuracy and TR ratio. For task-vector steering experiments, we construct task vectors from the outputs of TL or TR heads selected at each percentage level. The results in Figure 52 and Figure 53 show that our findings are robust to the choice of threshold. Specifically, ablating different percentages of TL heads impacts accuracy but leaves the TR ratio largely unchanged, whereas ablating TR heads affects both accuracy and the TR ratio. Moreover, using more top TR heads when constructing task vectors increases accuracy (with saturation around 7%), whereas including more TL heads does not yield accuracy gains.
These observations not only reinforce our main conclusions but also shed light on how attention heads collectively realize TL and TR functionality: heads at different pop-percen levels make additive contributions- albeei with varying strengths-raid rrather than the behavior being dominated by only few exceptional heads. This further justifies selecting heads based percentage threshold rather than attempting to isolate only handful of specific heads.
G.2 REPLICATION OF FIGURE 2 FOR OTHER MODEL.S
Figures 31-35 replicate the experiments from Figure 2, demonstrating the robustness of our findings across different models. In every case, TR heads show markedly stronger overlap, correlation, and consistency with IHs than the other head pairs. The consistently higher values of the TR-IH pair in terms of Jaccard Coefficient, Kendall's t, Spearman's 's p, and Conditional Average Percentage across all levels and architectures confirm conclusion.
G.3 REPLICATION OF FIGURE 3 FOR OTHER MODELS
Figures 36-40 replicate the experiments from Figure 3 on additional models. These visualizations support our claims in Subsection 4. that: 1) TR heads generally reside in deeper layers than TL heads and IHs; 2) The overlap between TR heads and IHs is larger and predominantly occurs in deeper layers.
To complement these figures, Tables 6-11 report the mean layer indices of the top 3%, 5%, and 10% TR heads, TL heads, and IHs averaged across datasets. We also conduct Mann-Wi Whitney U tests to assess whether the differences layer distributions between TR heads and IHs, and between IHs and TL heads, statistically significant. The results show that the distributional differences between IHs and TL heads are often not significant (p 2 0.05). Even when significant, the p-values are much larger than those observed between TR heads and IHs, indicating that the TR-IH distinction is far more robust.
G.4 REPLICATION OF FIGURE 6 FOR OTHER MODELS
Figures 41-45 extend the analysis of Figure 6 to the remaining models. The results reinforce conclusion that TR heads and IHs identified across datasets or tasks are largely consistent, whereas TL. heads vary substantially.
19

---

Published as a conference paper at ICLR 2026
To further test this, evaluate the cross-dataset transferability of TR heads. Specifically, we ablate top 3% TR heads identified using SST-2 prompts and measure their impact on accuracy and TR ratio for the six remaining datasets. The results for all models in Figures 73-78 in general confirm the transferrability of TR heads across datasets, but some interesting variations among datasets and models are also worth noting. First, on the three Llama family models, ablating the TR heads identified on the SST-2 dataset can effectively drive both the accuracy and TR ratio on RTE, CB, and MR datasets to near zero, and to a lesser extent impact the two metries on TREC and SNLI On Yi-34B, the ablation instead significantly impact the model performance on SNLI rather than MR. For the remaining two Qwen family models the consequence of the ablation over datasets similar to the case of Llama models but lesser degree overall.
G.5 REPLICATION OF FIGURE 4 FOR OTHER MODELS
In Figures 54-58 we report the results of quantifying the attention weight distributions of top TL and TR heads on other models, which are largely similar to Figure 4. 'TR heads assign larger weights to demonstration labels to collect information about the task label space, whilst TI. heads attend the query to leverage its specific semantics to facilitate its matching to the proper label token.
G.6 REPLICATION OF FIGURE 5 FOR OTHER MODELS
In Figure 5, we visualize the distinct attention patterns of TL and TR heads by showing the attention distributions of the top-1 TL and TR heads over ICL prompt formed using the first test query of SST-2. In Figures 59-67, we report the corresponding attention distributions over ICL prompts formed using the second through tenth test queries of SST-2; these largely resemble Figure 5 and thus support its validity.
G.7 DEMYSTIFYING THE LAYER ORDER OF TL. AND TR FROM THE PERSPECTIVE OF THE LAYERWISE EVOLUTION OF ICL HIDDEN STATES
Our results in Figure 3 suggest that attention heads responsible for TL. emerge in earlier layers than those responsible for TR. This raises question: how can the model perform task learning before it has recognized the task? The key to resolving this lies in our definition of TL and TR (and their corresponding heads) based on how they are carried out in the actual forward computation process of the model, which ultimately results in the promotion of certain tokens' logits at the output layer and their prediction. From a logit-centric viewpoint, TL corresponds the process by which the model increases the logit margin between the correct label and the incorrect candidate labels within the task's 's label set. In contrast, TR corresponds to increasing the logit margin between the task's tokens, enabling the model detect which task it
label set a whole and all other non should perform. Under this interpretation, TL and TR are fundamentally distinct and need not occur simultaneously, which explains the low overlap and weak correlation between TL and TR heads observed in Figure 2a.
To further validate our claim that TL heads precede their TR counterparts, we analyze the layerwise logit dynamics of correct labels, incorrect candidate labels, and non-task-related tokens using the ICL hidden-state evolution. Specifically, we use the Logit Difference metric introduced in Subsection 4.3 to measure (i) the logit difference between the correct and incorrect candidate labels by projecting intermediate hidden states onto the corresponding unembedding differences, and (ii) the logit difference between the maximum logit among task-label tokens and the maximum logit among irrelevant tokens, which reflects task recognition capability. The results, averaged across datasets and models and shown in Figures 46-51, reveal striking contrast between these two margins. The margin between correct and incorrect task labels starts around and becomes positive after only few lovers--ndicating that the model quickly performs task learning and distinguishes between the candidate labels. In contrast, the margin between task labels and non-task labels is strongly negative in early layers and only becomes positive much later, indicating that the model initially fails to identify the task but begins to perform task recognition in deeper layers. This pattern aligns precisely with observed ordering of TL and TR heads.
20

---

Published as a conference paper at ICLR 2026
H SUPPLEMENTARY MATERIALS FOR SUBSECTION 4.2
H.1 REPLICATION OF FIGURE 7 FOR OTHER MODELS
In Figures 68-72, we replicate the ablation experiments on additional models and examine their effects on dataset-average accuracy and TR ratio. The results echo our observations in Subsection 4.2: 1) Ablation of TR heads and TL heads impacts the TR and TL components of ICL separately. 2) Ablating IHs produces effects similar ablating TR heads. 3) The primary cause of low accuracy in the zero-shot case-as well as in cases where TR heads or IHs are ablated-is the failure to adequately activate the TR functionality.
H.2 REPLICATION OF FIGURE 8 FOR OTHER MODELS
In Figures 79-83, we repeat the experiments from Figure 8 on other models, focusing on the ablation of TR and TL heads when ICL inputs are subjected perturbations. The results closely mirror those in Figure 8: when the in-context text-label mapping is destroyed or reversed, ablating TL heads has no effect- or even positive effect on accuracy. By contrast, since these perturbations do not alter the demonstration label space, the TR component of ICL remains unaffected.
H.3 ASSESSING THE INDEPENDENCE OF TR AND TL WITH FLIPPED DEMONSTRATION LABELS
To further validate the independence of TR and Tl. and the mechanisms their associated heads, we analyze the effect of ablating TR/TL heads under flipped demonstration labels. Specifically, we apply a mapping g : Y + Y! that reverses the demonstration labels (c.g.. "negative" - "positive," "positive" - ""eeattvve", as listed Table 3. Since label flipping invalidates the original text-label mapping captured by TL heads, we conjecture that ablating top TL heads will increase accuracy, while the effect of ablating TR heads will remain unchanged because the label space itself preserved.
The dataset-average results in Figures 84-89 confirm this conjecture: ablating top TL heads indeed raises accuracy, whereas ablating top heads still drives accuracy close to zero, as observed the standard setting of Figure 7.
H.4 COMPOSITE TASK LEARNING AND RECOGNITION IN ICL
In this section, we discuss how the findings Li al. (2025) relate our framework of Task Learning (TL) and Task Recognition (TR), together with the corresponding experimental results. The main findings of Li et al. (2025) are threefold:
1. For an ICL prompt, only the hidden states or head outputs at certain specific token positions serve as effective TVs. For instance, in the prompt "I like this movie: positive, don't 't like it: negative, I love it:", only the hidden states or head outputs at the two "." positions can serve as TV's because they encode the information needed to predict "positive" and "negative" as the next tokens. This suggests that TV encodes only the information needed for next-token prediction.
2. Consequently, in settings involving multiple or composite tasks, TV only supports predicting the label of the first task, but not subsequent ones. For example, in the prompt "France, big - Paris, small", which combines a Country-Capital task and an Antonym task, TV can correctly produce "Paris" for the first task but not "small" for the second.
3. Nevertheless, hidden states from prompts corresponding to different composite tasks that share the same initial task remain well-separated and linearly classifiable.
The limitation of task vectors to single-token prediction is directly relevant to our work, since we identify TL. and TR heads based on their outputs at single token position, and similarly construct TVs from those outputs to inject at a single position. Moreover, the effectiveness of TV's only for the first task, together with the linear separability of hidden states from different composite tasks, suggests that under composite tasks the model must repeatedly perform task learning for each constituent task while simultaneously maintaining recognition mechanism that tracks the full scope of the composite structure. This unusual composite setting therefore provides an opportunity to extend our TL & TR framework to more complex scenarios. To investigate TI. and TR in composite task learning,
21

---

Published as a conference paper at ICLR 2026
conduct the following experiment. We consider the Country-Capital Antonym composite task with ICL prompts the form "France, big - Paris, small. China, high - Beijing, low Germany, quick - Using 8-shot ICL prompts, obtain two sets of top 3% TR and TL heads, one for each constituent task. For Task I (Country-Capita), we identify the heads whose outputs at the final "->" position promote the logits of different capital names, as well as those whose outputs increase the logit margin between the correct capital label "Berlin" and the remaining capital names (the Task I TR and TL heads). Then, we supplement the query with the first task's i label, yielding the prompt "France, big + Paris, small. China, high - Beijing, low Germany, quick - Berlin,", and repeat the procedure identify the TR and TL heads for Task 2 (Antonym), whose label space consists of all adjectives used in the Antonym task. After identifying the relevant heads, we ablate the top 10% Task TR and TL heads measure their impact on the accuracy and ratio of Task (for prompts such as "Germany, quick ->"). Likewise, we ablate the top 3% Task 1 'TR and TL heads to evaluate the performance of Task (for prompts such as "Germany, quick + Berlin,") "). The results in Figure 17 and Figure 18 reveal striking pattern: ablating TL heads across tasks has minimal effect on accuracy and TR ratio, whereas ablating TR heads across tasks significantly affects both, though the TR heads are identified from entirely different task label spaces. We also report the Jaccard coefficient between the two TR (and TL) sets, as well as the Spearman's p between the Task and Task 2 TL (and TR) scores across all heads, in Table 4. These statisties show that TR heads across tasks have substantially higher overlap and much stronger correlation than TL heads, further confirming the cross-task similarity of TR heads and the cross-task dissimilarity of TL heads. Based on these findings, we provide the following mechanistic explanation for ICL under composite tasks. To correctly infer each task's label, a distinct set of TI. heads specialized for that specific task (see (Yin & Steinhardt, 2025) for discussion of such specialization in pretraining) activates to promote the correct task-specific label using the task labels learned from the demonstrations. In contrast, at the broader level of solving the entire composite task, unified set of TR heads tracks the label spaces of all constituent tasks. These heads provide the task-recognition foundation that enables each task susk-specific TL set to selectively promote its respective labels as the prompt progresses through the subtasks. This interpretation closely aligns with our findings in Figure 6, where TR heads but not TI. headt--identffied across datasets exhibit high degree of overlap and correlation.
Task Acc.
Task 1 TR Ratio
0.5
0.0
Task 2 TL Ablated
Task 2 TR Ablated
Figure 17: Effects of ablating the Task 2 TR and Tl. heads the accuracy and TR ratio of Task 1.
Task 2 Acc.
Task 2 TR Ratio
1.0
0.5
0.0
Task 1 TL Ablated
Task TR Ablated
Figure 18: Effects of ablating the Task 1 TR and TL heads the accuracy and TR ratio of Task 2.
22

---

Published as a conference paper at ICLR 2026
1 SUPPLEMENTARY MATERIALS FOR SUBSECTION 4.3
1.1 REPLICATION OF FIGURE 9 FOR OTHER MODELS
Figures 90-94 replicate the steering experiments from Figure 9, evaluating the effectiveness of task vectors constructed from special attention head outputs in other models. For all models except the two Qwen-family models, the results are consistent with Subsection 4.2: task vectors built from TR heads are substantially more effective than those from TL heads. In the Qwen models, however, TL heads match TR heads task vectors. This deviation can be explained by the high zero-shot accuracy of the Qwen models (Figure 93, Figure 92), which exceeds 20%- con-idnbiyerably higher than the other models. Because these models already achieve strong task recognition in the zero-shot setting, injecting TR- Re--head-ber basdd task vectors (which primarily encode recognition) provides less additional benefit.
1.2 REPLICATION OF FIGURE 10 FOR OTHER MODELS
Figures 95-99 evaluate how task vectors built from different types of heads affect the quality of generated reviews across models. Consistently, TL heads outperform TR heads and random heads as task vectors. An exception is Yi-34B, where steering reduces the average rating below the original zero-shot level. For Qwen2-7B and Qwen2.5-32B. TL-head task vectors even push ratings above the ICL-level baseline. Interestingly, the zero-shot reviews these models score higher than their ICL reviews. Closer inspection reveals why: ICL reviews, though coherent and stylistically faithful, sometimes contradict the sentiment required in the query. TL heads appear to filter out such inconsistencies by correctly capturing the text-label mapping and discarding misleading signals, thereby boosting zero-shot review quality beyond ICL.
In addition cross-model replication, Table 12 presents sampled outputs under ICL, zero-shot, and steering with different task vectors. These examples highlight the TL heads' ability to extract the correct text- label mapping and use it for generation. In contrast, zero-shot or TR-head steering often yields generic, off-topic sentences loosely related to the concept of "review."
1.3 EVALUATING TASK VECTOR PERFORMANCE ON A MORE COMPLEX TASK
Because the movic review generation task in the main text has relatively simple structure and synthetic two-way label space, we further evaluate task vectors constructed from TL and TR head outputs more complex setting: the SubjQA dataset (Bjerva et al., 2020). We use the "book" split, where each datapoint contains a sentiment label and a human-written book review. Unlike the movie task-whose labels are limited "positive" and "negative"-this dataset features much richer and more diverse sentiment label space, including labels such as "captivating," "anticlimactic." and "wrenching," among many others. The human- written reviews also introduce greater linguistic complexity and semantic variability. Following the same TSLA-basee -bssd procedure used to identify TI. and TR heads in the main text, we compute TL and TR scores under this enlarged sentiment label space. We then construct task vectors from the outputs of the top TL or TR heads and use them to steer book-review generation in a zero-shot setting. GPT is subsequently asked to rate each generated review 10-point scale based on coherence and how well it reflects the intended sentiment label. The average ratings across models, shown in Figures 100-105, mirror the patterns observed in the movie-review task: TL-based task vectors consistently outperform TR-based vectors and random baselines. This demonstrates that TL. heads capture abstract associations between demonstration/query texts and sentiment labels strongly enough to yield effective task vectors even in substantially more complex, real-world generation scenarios, thereby validating the robustness of our TSLA-basce identification of TL heads.
1.4 REPLICATION OF FIGURE 11 FOR OTHER MODELS
Figures 106-110 extend the geometric analysis of Figure 11 other models. The results largely confirm our earlier observation: TL heads tend to align hidden states with label-unem difference directions, while TR heads align hidden states with the broader task subspace.
23

---

Published as a conference paper at ICLR 2026
1.5 REPLICATION OF FIGURE 12 FOR OTHER MODELS
Figures 116-120 report layer-wise correlations between mean TR/TL head outputs and full layer updates. measured by logit difference and subspace alignment. Across models, we observe clear and consistent correlation patterns, reinforeing that TR and TL heads are the primary drivers of the geometric shaping of hidden states in layer updates.
1.6 ABLATING TOP TR AND TL HEADS PER LAYER TO VERIFY THEIR GEOMETRIC SIGNIFICANCE
In Figure 12, we validated the geometric importance of TR/TL heads by correlating their outputs with full layer updates. Here, we provide an alternative perspective. Specifically, we ablate the top three TR/TL heads per layer and then remeasure layer-wise hidden state updates under the same two metrics. Figures 121-126 show that TR and TL heads are indeed crucial: without top TR heads, hidden states fail to gradually align with the task subspace, crippling task recognition; without top TL heads, logit differences collapse, preventing hidden states from rotating toward the correct label's unembedding direction. By contrast, ablating three random heads per layer has negligible impact.
1.7 REPLICATION OF FIGURE 13 FOR OTHER MODELS
Figures 111-115 replicate the analysis of Figure 13 across models. The results consistent: TL heads excel in projected discriminant alignment, rotating hidden states toward the correct label unembedding and away from incorrect ones. TR heads, conversely, excel in subspace alignment, keeping hidden states well-positioned within the task subspace. Both substantially outperform randomly chosen heads on their respective strengths.
EXPERIMENT DETAILS RELATED TO TASK VECTORS
J.1 CONSTRUCTION AND APPLICATION OF TASK VECTORS
For each dataset, first construct 8-shot ICL prompts using the last 50 queries. The demonstrations are identical to those used when evaluating the 8-shot ICL accuracy for each dataset. Following the procedure of Todd et al. (2024), we compute the average output (across the 50 prompts) of each identified top 3% TR, TL. or random head at the final token position. We then sum these average outputs across heads to form the task vector.
In the steering experiment, the task vector is added to the hidden state of the final token of each zero-shot query the midpoint layer (e.g.. layer 16 for the 32-layer Llama3-8B) The modified hidden states are then propagated through the subsequent layers, and accuracy well as TR ratio are measured at the final layer.
J.2 IDENTIFYING TR AND TL HEADS ON THE MOVIE REVIEW DATASET
key difficulty in identifying TR and TL heads for free-form generation tasks is the unbounded label space, since labels are not restricted to finite set of tokens. To address this, we define the relevant label tokens as "positive" and "negative," reflecting the sentiment nature of the review-gente generatgnneotiinn task. Specifically, the TR score of head is defined the projection norm of its output onto the span of the unembedding vectors of "positive" and "negative" when processing ICL prompts from the review dataset. The TL score is defined as the inner product between the head output and the difference between the unembeddings of "positive" and "negative," normalized by its TR score. After identifying TR and TL heads, we construct task vectors from their outputs following the procedure in Appendix J.I.
J.3 MATUEMATICAL DETAILS OF THE MEASURES IN SUBSECTION 4.3 AND CALCULATION PROCEDURE
1. Logit Difference Given a hidden state h, we compute Aveyes/(y>) (h (W) w )), where We is the unembedding matrix, y" is the correct label, and V is the demonstration label space.
24

---

Published as a conference paper at ICLR 2026
h Proj W
2. Subspace Alignment We compute
which is the cosine similarity between
Prooi W Allll
Proj h and h.
For evaluation, we take the hidden state of the final position at the layer corresponding to 75% of the model depth (e.g., layer 24 Llama3-8B) Reported metric values are averaged across the first 30 ICL prompts of each dataset.
K CURATION DETAILS OF THE REVIEW DATASET
We use the following template, adapted from Zhao et al. (2025), to prompt ChatGPT-40 to generate movie reviews.
Compose a concise 30-word movie review that addresses the following four aspects: plot, sound and music, cultural impact, and emotional resonance. Use positive tone throughout the review. For the plot, comment on its structure or originality. For sound and music, describe how they enhance the storytelling. For cultural impact, mention any relevant social commentary. Finally, highlight how the film resonates emotionally. Ensure the positive tone is consistent throughout and include positive descriptions of the movie.
Prompt
Samples
Inventive non-linear storyline weaves intrigue with clever twists. Soaring vocal melodies heighten the film's 's emotional arcs. Relevant socioeconomic themes challenge viewers' perceptions thoughufully and respectfully. Joyful humor interwoven with drama creates comforting resonance.
We use the following template ask ChatGPT-4o to rate the movie reviews.
Rate the following movie review on a scale of 10. Your rating should be based on two criteria: (1) whether the text indeed movie review, and (2) whether conveys the posi- tive or negative sentiment indicated by the label. Review: Inventive non-linear storyline weaves intrigue with clever twists. Soaring vocal melodies heighten the film 's emotional ares. Relevant socioeconomic themes challenge viewers' perceptions thoughtfully and re- spectfully. Joyful humor interwoven with drama creates comforting resonance. Sentiment: Positive
Prompt
Response
10
25

---

Published as a conference paper at ICLR 2026
Table Prompt templates and labels for different datasets.
Dataset
Template
Labels
SST-2
(Sentence) Sentiment: [Label)
positive negative
SUBJ
[Sentence] TTpp:: [Label}
subjective objective
TREC
Question: (Sentence) Type: [Label)
abbreviation entity description human location number positive negative
MR
(eeneence) Sentiment: [fabel]
SNLI
The question is: {Premise}? True maybe false? The
true maybe falsc
answer is: (4ypothesis} [Label]
question is: {Premise)? True
(Rypothesis) [1abel]]
The question is: [Premise)) True or maybe or false?
true maybe/ false
answer is: [dypothesis] {Label}
Table Mappings used to replace ground-trute truth labels with numeric symbols.
Dataset
Label Mapping
SST-2
negative/positive -
SUBJ
objective/subjective 0/1
MR
negative/positive - 0/1
TREC
abrrvistion/eniii/lessrr
0/1/2/3/4/5
SNLI
true/mayhe/false - 0/1/2
RTE
true/false + 0/1
CB
true/mayhe/false 0/12
Table 3: Mappings used to flip the demonstration labels for each dataset.
Label Mapping
SST-2
negative/posiive positive/negative
SUB]
objeciive/subjectiee subjactive/oojectivee
TREC
MR
negalivolpsstliee posiliechcggaive
SNLI
tras'maybo/alsee mayboofalle/rrue
RTE
Irue/Tase - falseltrue
CB
be/abs may bolabaalaecttuue
Uperdecttay Matspano
Norms of embeddings when Projected oato Different Subspaces
Value
0.75
0.50
0.25
Duvwed
Tokens
Figure 19: Norms of Llamaa-8B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens vs semantically unrelated tokens.
Subepacs
Norms of emboddings when Projected onto Different Subapaces
Value
0.5
Gooundd
undoouralle
gleainn
Sountafiinn
Selightfl
Toleens
Figure 20: Norms of Llama3.1 1-8B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens semantically unrelated tokens.
26

---

Published as a conference paper at ICLR 2026
Uncmieddiis supapco lndeis Touess
Uhembedding Subipoce Unadaise Toeesa
Norms embeddings Projected
Subspaces
1.0
Value
delighttd
distaato
vpllDung
Figure 21: Norms of Llama3.2 2-3B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens vs semantically unrelated tokens.
Uncrdaddan Sabapani et BAland Tulasa Obomaddtng Sabess ttadaad Tokees
Norms embeddings when Projected Different Subspaces
Value
0.50
0.25
Gomurabl
Camall oxcallot
Tokens
Figure 22: Norms of Qwen2-7B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens vs semantically unrelated tokens.
Uniibeolnng Rolind
Unrmbedtan LnmnllsdTioooo
Nomms of embeddings when Projected onto Different Subspaces
Dupundden
porbbe
Tosons
Figure 23: Norms of Qwen2.5-32B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens semantically unrelated tokens.
Uhrnietdiing Relaned Thowss
Unambedding Utadaiie
Norms embeddings when Projected onlo Different Subepaces
Value
0.5
0.0
postbe
Gooawalen
Tokens
Figure 24: Norms of Yi-34B token unembeddings when projected onto the unembedding subspace spanned by semantically related tokens vs semantically unrelated tokens.
27

---

Published as a conference paper at ICLR 2026
ACC
TR Ratio
0.4
0.2
0.0
TR heads: TSLA
TR heads: DLA
Figure 25: Results on LLama3-8B: Effects of ablating top 10% TR heads identified using TSLA DLA when the SST-2 demonstration labels shifted from positive/negative favourable/un- favourable.
ACC
TR Ratio
0.6
0.4
0.2
0.0
TR heads: TSLA
TR heads: DLA
Figure 26: Results on Llama3.1 1-8B: Effects of ablating top 10% TR heads identified using TSLA or DLA when the SST-2 demonstration labels shifted from positive/negative favourable/un- favourable.
28

---

Published as a conference paper at ICLR 2026
ACC
Z
TR Ratio
1.00
0.75
0.50
0.25
0.00
TR heads: TSLA
TR heads: DLA
Figure 27: Results on Llama3.2-3BB Effects of ablating top 10% TR heads identified using TSLA or DLA when the SST-2 demonstration labels are shifted from positive/negative to favourable/un- favourable.
ACC
TR Ratio
1.00
0.75
0.50
0.25
0.00
TR heads: TSLA
TR heads: DLA
Figure 28: Results on Qwen2-7B: Effects of ablating top 10% TR heads identified using TSLA or DLA when the SST-2 demonstration labels are shifted from positive/negative to favourabbelnfaavorabbe
ACC
TR Ratio
0.8
0.6
0.4
0.2
0.0
TR heads: TSLA
TR heads: DLA
Figure 29: Results on Qwen2.5-32B: Effects of ablating top 10% TR heads identified using TSLA or DLAA when the SST-2 demonstration labels are shifted from positive/negative to favourable/un- favourable.
29

---

Published as a conference paper at ICLR 2026
ACC
TR Ratio
1.00
0.75
0.50
0.25
0.00
TR heads: TSLA
TR heads: DLA
Figure 30: Results on Yi-34B: Effects of ablating top 10% TR heads identified using TSLA or DLA when the SST-2 demonstration labels are shifted from positive/negative to favourrablcuufaavourabl...
Jaccard Caef
0.5 0.5
TR-TL
TL-IH
0.4
TL-IH
0.3
TR-IH
0.2
40.0
0.1
0:0
TR-IH
20.0
Spearmán's p
Kendall's
3%
5%
10%
20%
(b) Conditional Average Percentage at four top levels for TR-IH and TL-IH pairs.
(a) Jaccard Coefficient, Kendall's and Spearman's P values TR heads, heads, and IHs.
Figure 31: Results of overlap, correlation, and consistency analysis of attention head types averaged across datasets on 1--ma..1-8B.
Jaccard Coef
0.5
TR-TL
TL-IH
0.4
TL-IH
0.3
TR-IH
0.2
40.0
0.1
0.0
TR-IH
20.0
Spearman's, '
Kendall's
3%
5%
10%
20%
(b) Conditional Average Percentage at four top levels for TR-II and TL-IH pairs.
(a) Jaccard Coefficient, Kendall's v. and Spearman's values for TR heads, TL heads, and IHs.
Figure 32: Results of overlap, correlation, and consistency analysis of attention head types averaged across datasets on Llama3.2-3B.
30

---

Published as a conference paper at ICLR 2026
Jaccard Coef 0.35
TR-TL
TL-IH
0.25
TL-IH
0.15
TR-1H
40.0
0.10
0.05
0.00
TR-IH
20.0
Spearman's
Kendall's
3%
5%
10%
20%
(b) Conditional Average Percentage at four top levels for TR- IH and TL-IH pairs.
(a) Jaccard Coefficient, Kendall's and Spearman's P values for TR heads, TL heads, and IHs.
Figure 33:
Results of overlap, correlation, and consistency analysis of attention head types averaged across datasets on Qwen2-7B.
Jaccard Coef
0.5
TR-TL
TL-IH
0.4
TL-IH
TR-IH
0.2
40.0
0.1
TR-IH
20.0
Spearman's
Kendall's
3%
5%
10%
20%
Conditional Average Percentage at four top levels TR-IH TL-IH pairs.
(a) Jaccard Coefficient, Kendall's and Spearman's P values for TR heads, TL heads, and IHIs.
Figure 34: Results of overlap, correlation, and consistency analysis of attention head types averaged across datasets on Qwen2.5-32B.
Joee
0.5
TR-TL
TL-IH
0.4
TL-IH
0.3
TR-IH
40.0
0.2
0.1
0.0
TR-IH
20.0
Spearman's, '
Kendall's
3%
5%
10%
20%
(b) Conditional Average Percentage at four top levels for TR-IH and TL-II pairs.
(a) Jaccard Coefficient, Kendall's v, and Spearman's values for TR heads, TL heads, and IHs.
Figure 35: Results of overlap, correlation, and consistency analysis attention head types averaged across datasets on Yi-34B.
31

---

Published as a conference paper at ICLR 2026
TL
TR
IH-TR
3
5
7
9
11
13
Layer ID
15
19
21
23
25
27
29
31
7
9 11 13
15
17
19
21
23
25
27
29
31
Head ID
Figure 36: Distribution of the top 10% TR heads, TL heads, and IHs across layers for the SST-2
dataset on Llama3.1-88
IH
TL
TR
IH-TR
1
3
5
7
9
11
Layer ID
13
19
21
23
25
27
1
5
9
11
13
15
17
19
21
23
Head ID
Figure 37: Distribution of the top 10% TR heads, TL heads, and IHs across layers for the SST-2 dataset on Llama3.2-35
32

---

Published as a conference paper at ICLR 2026
IH
TL
TR
IH-TR
1
3
5
7
9
11
Layer ID
13
15
19
21
23
25
5
9
11
13
15
17
19
21
23
25
27
Head ID
Figure 38: Distribution of the top 10% TR heads, TL heads, and IHs across layers for the SST-2 dataset on Qwen2-7B.
IH
TL
TR
IH-TR
1
5
9
13
17
21
25
ID
29
33
37
41
45
49
53
57
7 9 11 13 15
17 19
21
23
25 27
29
31
33
35
37
39
Head ID
Figure 39: Distribution of the top 10% TR heads, TL heads, and IHs across layers for the SST-2 dataset on Qwen2.5-32B.
33

---

Published as a conference paper at ICLR 2026
IH
TL
TR
IH-TR
5
9
13
17
21
Layer ID
25
29
33
37
41
45
49
53
57
13
21
29
33
41
53
Head ID
Figure 40: Distribution of the top 10% TR heads, TL heads, and IHs across layers for the SST-2
dataset on Yi-34B.
Kendall
Spearman Jaccard
1.0
M Min--M Max Normalizes
TR
0.60
0.78
0.34
0.8
0.6
TL
0.08
0.12
0.13
0.4
IH
0.76
0.92
0.43
0.2
0.0
Figure 41: Overlap and correlation of the TR heads, TL. heads, and IHs across datasets on Llama3.1
8B.
34

---

Published as a conference paper at ICLR 2026
Kendall
Spearman Jaccard
1.0
Min Mii-Max Max Normalized
TR
0.58
0.75
0.33
0.8
0.6
TL
0.08
0.12
0.08
0.4
IH
0.76
0.92
0.47
0.2
0.0
Figure 42: Overlap and correlation of the TR heads, TL. heads, and IHs across datasets on Llama3.2. 2- 3B.
Kendall
Spearman Jaccard
1.0
Max lozmalized
TR
0.56
0.74
0.34
0.8
0.6
TL
0.10
0.15
0.20
0.4
IH
0.68
0.85
0.48
0.2
0.0
Figure 43: Overlap and correlation of the TR heads, TL heads, and IHs across datasets on Qwen2-7B.
35

---

Published as a conference paper at ICLR 2026
Kendall
Spearman Jaccard
1.0
lizdd
Min-Mar Normalizer
TR
0.61
0.79
0.37
0.8
0.6
TL
0.10
0.14
0.19
0.4
IH
0.67
0.85
0.43
0.2
0.0
Figure 44: Overlap and correlation of the TR heads, TL heads, and IHs across datasets on Qwen2.5- 32B.
Kendall Spearman Jaccard
1.0
Min-Men Max Normalized
TR
0.69
0.86
0.32
0.8
0.6
TL
0.06
0.09
0.13
0.4
IH
0.66
0.84
0.44
0.2
0.0
Figure 45: Overlap and correlation of the TR heads, TL heads, and IHs across datasets on Yi-34B.
Correct Task Label vs Wreng Task Labeis
Currect Tark Label vs Non- Tank Labels
Two Types of Logit Margins
0
Value
-50
2
9
11
12
13
14
15
16
17
192021
X
Figure 46: Dynamics of logit margin between 1) correct incorrect task labels and 2)tk labels and task-irrelevant labels across layers Llama3-8B.
36

---

Published as a conference paper at ICLR 2026
Task Label Task
Currect Task Labei Mask-aaki Labels
Two Types of Logit Margins
0
Value
-50
10 11 12 14 01111111111111111111212132121111
Figure 47: Dynamics of logit margin between 1) correct and incorrect task labels and 2)task labels and task-irrelevant labels across layers Alama3.1-8B.
Correct Tack Label vs Wrong Tack Labels
Non-Teed Labele
Two Types of Logit Margins
0
Value
-200
-400
4
5
7
8
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24 25 26 27 28
X
Figure 48: Dynamics of logit margin between correct and incorrect task labels and 2)task labels and irre-itrelevant labels across layers on Llama3.2-3B.
Wrong Label Lebele
Two Types Logit Margins
50
0-
Value
-50
-100
10
13
16
19
20
21
22
25 27 28
Figure 49: Dynamics of logit margin between 1) correct and incorrect task labels and 2)task labels and task-irrelevant labels across layers on Qwen2-7B.
37

---

Published as a conference paper at ICLR 2026
Cerrect Tark Label Weong Task Labdle
Two Types of Logit Margins
50
Value
0
-50
4
12
16
20
24
28
32
36
40
44
48
52
56
60
64
Figure 50: Dynamics of logit margin between 1) correct and incorrect task labels and 2)task labels and task-irrelevant labels across layers on Qwen2.5-32B.
Correct Tack Label Wrong Task Labels
Correct Tack Label vs Non-Tek Labele
Two Types of Logit Margins
Value
-100
8
12
16
20
24
28
32
36
40
44
48
52
56
60
X
Figure 51: Dynamics of logit margin between 1) correct incorrect task labels and 2)task labels and task-irrelevant labels across layers Yi-34B.
Different of TL Ablated
Different of TR Ablated
40.0%
80%
30.0%
Value
Acc.
Value
3%
60%
TR Ratio
20.0%
3%
10.0%
40%
1% 2%
3% 4%
6% 7%
9% 9%10%%
1% 2% 3% 4% 6% 7% 8% 10% Percentage
Percentage
Figure 52: Dataset average accuracy and TR ratio resulted from ablating TL and TR heads at top percentage levels from 1% to 10%.
38

---

Published as a conference paper at ICLR 2026
TV Acc. with Different % of Heads
TL
50.0%
TR
40.0%
Value
30.0%
3%
20.0%
10.0%
1%
2%
3%
4%
5%
6%
7%
8%
9%
10%
Percentage
Figure 53: Dataset average accuracy and TR ratio resulted from using the outputs TL and TR heads top percentage levels from 1% to 10% as task vectors.
Attn. to Demo Labels
Attn. to Query
0.075
0.06
0.37
0.4
0.050
0.03
0.025
0.02
0.21
0.20
0.2
TL
TR
Random
Figure 54: Dataset average cumulative attention weights assigned by the top TL, TR, and random heads of Llama3.1 1-8B to the demonstration labels and query tokens.
Attn. to Demo Labels
Attn. to Query
0.075
0.06
0.4
0.33
0.050
0.04
0.025
0.02
0.18
0.20
0.2
TL
TR
Random
Figure 55: Dataset average cumulative attention weights assigned by the top TL, TR, and random heads of Llama3.2-33 to the demonstration labels and query tokens.
39

---

Published as a conference paper at ICLR 2026
Attn. Demo Labels
Attn. to Query
0.10
0.50
0.07
0.40
0.05
0.33
0.05
0.02
0.20
0.25
TL
TR
Random
Figure 56: Dataset average cumulative attention weights assigned by the top TI., TR, and random heads of Qwen2-7B to the demonstration labels and query tokens.
Attn. to Demo Labels
Attn. to Query
0.10
0.5
0.08
0.38
0.4
0.05
0.03
0.27
0.03
0.25
0.3
TL
TR
Random
Figure 57: Dataset average cumulative attention weights assigned by the top TL, TR, and random heads of Qwen2.5-32B the demonstration labels and query tokens.
Attn. to Demo Labels
Attn. to Query
0.075
0.33
0.06
0.4
0.050
0.26
0.025
0.01
0.02
0.14
0.2
TL
TR
Random
Figure 58: Dataset average cumulative attention weights assigned by the top TL, TR, and random heads of Yi-34B the demonstration labels and query tokens.
0.33
TL head
Atteution Wesekk
0.32
TR head
0.22
0.27
0.17
6.0
Token
Figure 59: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the second test query of SST-2.
40

---

Published as a conference paper at ICLR 2026
TL head
Atteation Welghs
TR head
0.13
0.0
0.0
31
Tolenn
Figure 60: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the test query of SST-2.
bead head
Atteation Welgha
6..9
5.00
11gi 19
Tolenn
Figure 61: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the fourth test query of SST-2.
0.32
head head
Attaation Wolgha
0.12
5.02
Token
Figure 62: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the fifth test query of SST-2.
6.35
TL
Atteation Woight
6.32
TR head
0.27
5.00
Tolonn
Figure 63: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the sixth test query of SST-2.
41

---

Published as a conference paper at ICLR 2026
0.30
TL head TR head
Atteation Welght
0.2
0.3
0.0
Toleen
Figure 64: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the seventh test query of SST-2.
0.27
TL
Atteation Welghtt
TR head
0.23
0.13
0.09
6.0
5.00
705 1
Tolenn
Figure 65: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the eighth test query of SST-2.
head head
Atteation Welgha
8:25
0.22
0.08
5.02
1og
1f
4
Token
Figure 66: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the ninth test query of SST-2.
bead head
Atteation Walght
022
Token
Figure 67: Attention distributions of the top-1 Llama3-8B TL and TR heads on SST-2 over the ICL prompt formed using the tenth test query of SST-2,
42

---

Published as a conference paper at ICLR 2026
ACC
Z
TR Ratio
1.00
0.75
ssang dropuer
0.50
0.25
0.00
ICL
ZSL
W/O IH
W/O TL
W/O TR
W/O Rand.
Figure 68: Effects of ablating the top 3% TR, TL, and IH heads across datasets on Llama3. 1-BB.
ACC
ZZ
TR Ratio
1.00
0.75
ssan6 tropuet
0.50
0.25
0.00
ICL
ZSL
W/O IH
W/O TL
W/O TR
W/O Rand.
Figure 69: Effects of ablating the top 3% of TR, TL., and IH heads across datasets on Llama3.2-33
ACC
ZZ
TR Ratio
1.00
0.75
0.50
0.25
0.00
ICL
ZSL
W/O IH
W/O W/O TR
W/O Rand.
Figure 70: Effects of ablating the top 3% of TR, TL, and IH heads across datasets on Qwen2-7B.
43

---

Published as a conference paper at ICLR 2026
ACC
7ZZ
TR Ratio
1.00
0.75
random guess
0.50
0.25
0.00
ICL
ZSL
W/O IH
W/O TL
W/O TR
W/O Rand.
Figure 71: Effects of ablating the top 3% of TR, TL, and IH heads across datasets on Qwen2.5-32B.
ACC
TR Ratio
1.00
0.75
Sindoy quess
0.50
0.25
Tana
0.00
ICL
ZSL
W/O IH
W/O W/O TR
W/O Rand.
Figure 72: Effects of ablating the top 3% of TR, TL, and IH heads across datasets on Yi-34B.
ACC
TR Ratio
1.00
0.75
0.50
0.25
V/
0.00
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 73: Effects of ablating TR heads identified on SST-2 when transferred to other datasets using Llama3-8B.
44

---

Published as a conference paper at ICLR 2026
ACC
Z TR Ratio
1.00
0.75
0.50
0.25
0.00
F77
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 74: Effects of ablating TR heads identified SST-2 when transferred to other datasets using Llama3.1-8B.
ACC
TR Ratio
0.8
0.6
0.4
0.2
0.0
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 75: Effects of ablating TR heads identified on SST-2 when transferred to other datasets using Llama3.2-3B.
ACC
Z TR Ratio
1.00
0.75
0.50
0.25
0.00
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 76: Effects of ablating TR heads identified on SST-2 when transferred to other datasets using Qwen2-7B.
45

---

Published as a conference paper at ICLR 2026
ACC
TR Ratio
0.8
0.6
0.4
0.2
0.0
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 77: Effects of ablating TR heads identified SST-2 when transferred to other datasets using Qwen2.5-32B.
ACC
TR Ratio
0.8
0.6
0.4
0.2
0.0
SUBJ
TREC
SNLI
RTE
CB
MR
Figure 78: Effects of ablating TR heads identified on SST-2 when transferred to other datasets using Yi-34B.
1.0
1.0
ACC
2S0 TR Ratio
0.5
0.5
ICL
W/O IL
W/O TR
ICL
TL.
W/O TR
(a) Ablating TR and TL. heads with shuffled demon- stration text order.
(b) Ablating TR and TL heads with labels replaced by numbers.
Figure 79: Effects of ablating TR and TL heads under perturbed ICL inputs on Llama3.
0.5
0.0
TR (a) Ablating and TL. heads with shuffled demon- stration text order.
ICL
W/O TL
W/O TR
(b) Ablating TR and TL heads with labels replaced by numbers.
Figure 80: Effects of ablating TR and TL heads under perturbed ICL inputs on 1lama3.2-3B.
46

---

Published as a conference paper at ICLR 2026
1.0
1.0
ACC
B0S TR Ratio
0.5
ICL
W/O IL
W/O TR
ICL
W/O TL
W/O TR
(a) Ablating TR and TL heads with shuffled demon- stration text order.
(b) Ablating TR and TL heads with labels replaced by numbers.
Figure 81: Effects of ablating TR and TL heads under perturbed ICL inputs on Qwen2-7B.
1.0
1.0
ACC
TR Ratio
0.0
ICL
W/O
ICL
W/O TL
W/O TR
(a) Ablating TR and TL. heads with shuffled demon- stration text order.
(b) Ablating TR and TL heads with labels replaced by
numbers.
Figure 82: Effects of ablating TR and TL heads under perturbed ICL inputs Qwen2.5-32B.
0.0
(a) Ablating TR and TL heads with shuffled demon- stration text order.
(b) Ablating and heads with labels replaced by numbers.
Figure 83: Effects of ablating TR and TL. heads under perturbed ICL inputs on Yi-34B.
ACC
Ratio
1.0
0.5
0.0
ICL
W/O TL
W/O TR
Figure 84: Effects of ablating TR and TL heads Llama3-8B when demonstration labels are flipped.
ACC
X TR Ratio
1.0
0.5
0.0
ICL
W/O TL
W/O TR
Figure 85: Effects of ablating TR and TI. heads on Llama3. 1-8B when demonstration labels are flipped.
47

---

Published as a conference paper at ICLR 2026
ACC
8S8 TR Ratio
1.0
0.5
0.0
XXXXX
ICL
W/O TL
W/O TR
Figure 86: flipped.
Effects of ablating TR and TL heads on Llama3.2-33 when demonstration labels are
ACC
TR Ratio
1.0
0.5
ICL
W/O TL
W/O TR
Figure 87: Effects ablating TR and TI. heads on Qwen2-7B when demonstration labels are flipped.
ACC
TR Ratio
1.0
0.5
0.0
ICL
W/O TL
W/O TR
Figure 88: Effects of ablating TR and TL heads on Qwen2.5-32B when demonstration labels are flipped.
ACC
TR Ratio
0.5
0.0
ICL
W/O TL
W/O TR
Figure 89: Effects of ablating TR and TL heads on Yi-34B when demonstration labels are flipped.
Original
After
Rand. TV
7.6%
TR TV
7.6%
Q37.5%
TL TV
7.6%
O12.2%
10.0%
20.0%
30.0%
40.0%
Figure 90: Steering zero-shot hidden states of Llama3. 1-8B using task vectors from TR, TL, or random heads.
48

---

Published as a conference paper at ICLR 2026
Original
After
Rand. TV
5.1%O 07.1%
TR as TV
7.1%
O46.6%
TL as TV
7.1%
0%
10%
20%
30%
40%
50%
Figure 91: Steering zero-shot hidden states of Llama3.2- 2-33 using task vectors from TR, TL, or random heads.
Original
O
After
Rand. as TV.23.0%0023.. 23.9%
TR TV
23.9%
O44.7%
TL TV 23.9%
O46.4%
Figure 92: Steering zero-shot hidden states of Qwen2-7B using task vectors from TR, TL, random heads.
Original
O
After
Rand. as TV 23.3%
025.7%
TR as TV 23.3%023.3%
TL as TV 23.3%
029.1%
24.0%
26.0%
28.0%
30.0%
Figure 93: Steering zero-shot hidden states of Qwen2.5-32B using task vectors from TR, TL, or random heads.
49

---

Published as a conference paper at ICLR 2026
Original
After
Rand. TV 25.2%625.22
TR as 25.2%
O48.1%
TL as TV 25.2%
O33.6%
25.0% 30.0% 35.0% 40.0% 45.0% 50.0%
Figure 94: Steering zero-shot hidden states of Yi-34B using task vectors from TR, TL, or random heads.
Mean = SD
10
8.25
5.34
4.62
4.35
2.75
ICL
ZS
TL as TV
TR as TVRand. TV
Figure 95: Mean and standard deviation of review ratings with Llama3. when task vectors from different head types are applied.
Mean = SD
8.10
6.13
3.27
3.83
2.19
ICL
ZS
TL as TV TR as TVRand. as TV
Figure 96: Mean and standard deviation of review ratings with Llama3.2-34 when task vectors from different head types are applied.
Mean = SD
8.25
8.68
8.29 ?
6.84
0.00
ICL
ZS
TL as TV TR as TVRand. as TV
Figure 97: Mean and standard deviation of review ratings with Qwen2-7B when task vectors from different head types are applied.
50

---

Published as a conference paper at ICLR 2026
Mean + SD
9.52
10
8. 11
8.11
8.84
8.27
ICL
ZS
TL as TV TR as TV Rand as TV
Figure 98: Mean and standard deviation of review ratings with Qwen2.5-32B when task vectors from different head types are applied.
Mean = SD
10
6.89
6.22
4.91
4.84
2.33
ICL
ZS
TL as TV TR as TVRand as TV
Figure 99: Mean and standard deviation of review ratings with Yi-34B when task vectors from different head types are applied.
Mean + SD
10
8.43
7.18
5.92
6.30
2.14
ICL
ZS
TL as TV TR as TV Rand. as TV
Figure 100: Mean and standard deviation of the ratings of book reviews generated using Llama3-8B when task vectors from different head types identified on the SubjQA dataset are applied.
Mean + SD
10
7.30
6.45
6.07
4.34
3.28
ICL
ZS
TL as TV TR as TV Rand. as TV
Figure 101: Mean and standard deviation of the ratings of book reviews generated using Llama3.1-88 when task vectors from different head types identified on the SubjQA dataset are applied.
51

---

Published as a conference paper at ICLR 2026
Mean + SD
10
8.43
7.18
5.92
6.30
2.14
ICL
ZS
TL as TV TR as TV Rand. as TV
Figure 102: Mean and standard deviation of the ratings of book reviews generated using Llama3.2-3B when task vectors from different head types identified on the SubjQA dataset are applied.
Mean = SD
10
6.18
4.78
5.48
4.62
5
1.22
ICL
ZS
TL as TV TR as TVnand. as TV
Figure 103: Mean and standard deviation of the ratings of book reviews generated using Qwen2-7B when task vectors from different head types identified on the SubjQA dataset are applied.
Mean + SD
10
6.30
7.08
6.74
3.80
2.69
ICL
ZS
TL as TV TR as Rand as TV
Figure 104: Mean and standard deviation of the ratings of book reviews generated using Qwen2.5- 32B when task vectors from different head types identified on the SubjQA dataset are applied.
Mean + SD
5.97
3.64
4.20
4.01
0.00
ICL
ZS
TL as TV TR as TV Rand. as TV
Figure 105: Mean and standard deviation of the ratings of book reviews generated using Yi-34B when task vectors from different head types identified on the SubjQA dataset are applied.
52

---

Published as a conference paper at ICLR 2026
Logit Difference
Subspace Alignment
5.0
4.91
0.29
0.30
4.5
0.25
4.33
4.0
0.20
3.69
0.16
3.50 0.15
3.5
VLI
0.15
Original
+TL
+TR
+ Rannd.
Figure 106: Geometric effects of 'TR and TL head outputs on hidden states Llama3. 1-8B.
Logit Difference 2.95
Subspace Alignment
3.00
0.37
2.86
-0.35
2.75
-0.30
2.50
-0.25
2.25
2.16
0.20
0.17
0.17
2.09 0.17
2.00
0.15
Original
+TL
+TR
Rand.
Figure 107: Geometric effects of TR and TL head outputs hidden states in Llama3.2-30
Logit Difference
Subspace Alignment
6.17
0.28
0.3
6.00
6
5
0.2
0.08
0.1
3
2.92
0.05
2.76
0.05
Original
+TL
+TR
+Rnnd.
Figure 108: Geometric effects of TR and TL head outputs on hidden states in Qwen2-7B.
53

---

Published as a conference paper at ICLR 2026
Logit Difference
77 Subspace Alignment
12.5
11.79
10.45
10.0
0.2
7.5
0.09
0.1
5.0
3.89
0.04
3.67 0.04
Original
+TL
+TR
+aand.
Figure 109: Geometric effects of TR and TL head outputs on hidden states in Qwen2.5-32B.
Logit Difference 9.79
Subspace Alignment
10.0
0.29
9.5
9.36
-0.25
9.0
8.5
0.20
8.10
0.18
0.18
7.98
8.0
Original
+TL
+TR
Rannd.
Figure 110: Geometric effects of TR and TL head outputs hidden states in Yi-34B.
Projected Discr. Align.
ZZZ
Subspace Alignment
0.42
0.4
0.33
0.2
0.09
0.07
0.02
0.0
7ZZ
-0.04
TL
TR
Rand.
Figure 111: Impact of TL and TR head outputs hidden states w.r.t. task subspace in Llama3.1 1-8B.
54

---

Published as a conference paper at ICLR 2026
Projected Discr. Align.
777
Subspace Alignment
0.43
0.39
0.4
0.2
0.08
0.05
0.03
0.0
FZZZ
-0.04
TL
TR
Rand.
Figure 112: Impact of TL and TR head outputs on hidden states w.r.t task subspace Llama3.2-3B.
Projected Discr. Align.
F7A Subspace Alignment
0.6
0.51
0.4
0.38
0.2
0.12
0.09
0.02
0.0
-0.03
TL
TR
Rand.
Figure 113: Impact of TL and TR head outputs on hidden states w.r.t task subspace in Qwen2-7B.
Projected Discr. Align.
Subspace Alignment
0.59
0.50
0.37
0.25
0.18
0.13
0.02
0.00
-0.05
TL
TR
Rand.
Figure 114: Impact of TL and TR head outputs on hidden states w.r.t. task subspace in Qwen2.5-32B.
55

---

Published as a conference paper at ICLR 2026
Table 4: Jaccard coefficient and Spearman's for TL and TR scores across tasks.
TL. Task 1&2
TR Task 1&2
Jaccard coefficient
0.1860
0.3421
Spearman's $
0.2980
0.8365
Llama3-8B
Llama3.1-8B
Llama3.2-3B
Qwen2-7B
Qwen2.5B
Yi-34B
1.96 10-4
1.96 10
7.63 10-6
7.63 x 10-6
7.63 10-6
7.37 10
Table 5: Statistical significance of the difference between the norms of token unembeddings projected to different subspaces across models.
TL heads mean layer
IHs mean layer
TR heads mean layer
TL < IHs?
IHs < TR heads?
0.03
16.89
16.69
26.22
0.84985
5.6052e-45
0.05
16.94
16.74
25.08
0.80371
0,0000e+00
0.10
16.75
16.88
23.27
0.59518
0.0000c+00
Table 6: Mean layer index of TL heads, IHs, and TR heads across datasets on Llama3-8B, with p-values for distribution differences.
TL. heads mean layer
IHs mean layer
TR heads mean layer
TT. < IHs?
IHs < TTR heads?
0.03
16.85
16.67
26.07
0.71998
8.4078e-45 0.0000e+00 0.0000e+00
0.05
16.36
16.77
25.33
0.22671 0.049241
0.10
16.
16.79
23.25
Table 7: Mean layer index of TL heads, IHs, and TR heads across datasets on 10mma,.1-BB, with p-values for distribution differences.
TL heads mean layer
IHs mean layer
TR heads mean layer
TL < IHs?
IHs TR heads?
0.03
14.51
15.69
23.04
0.040803
1.3459e-26 3.0489e-35 1.0738e-36
0.05
14.25
15.78
22.59 20.96
0,0024930 0.00044495
0.10
14.38
15.75
Table 8: Mean layer index of TL heads, IHs, and TR heads across datasets on Llama3.2-3B, with p-values for distribution differences.
TL heads mean layer
IHS mean layer
TR heads mean layer
TL < IHs?
IHs < TR heads?
0.03
19,06
18.57
24.87
0.31254
2.6148e-28 1.7432c-42
0.05
18.01
18.62
24.64
0.37007
0.10
16.22
18.97
23.38
2.4055e-09
1.8049e-41
Table 9: Mean layer index of TL heads, IHs, and TR heads across datasets Qwen2-7B, with p-values for distribution differences.
TL heads mean layer
IHs mean layer
TR heads mean layer
TL IHs?
IHs < TR heads?
0.03
46.26
44.08
56.80
0.0012338
0.0000e+00
0.05
43.19
44.53
54.89
0.35802 1.7597e-20
0.0000e+00
0.10
39.42
45.11
51.36
0.0000e+00
Table 10: Mean layer index of TL heads, IHs, and TR heads across datasets on Qwen2.5-32B, with p-values for distribution differences.
56

---

Published as a conference paper at ICLR 2026
TL heads mean layer
IHs mean layer
TR heads mean layer
TL < IHs?
IHs TR heads?
0.03
36.64
38.16
47.21
0.071766
0.0000e+00
0.05
35.57
38.20
46,90 45.20
0.0011017 1.0199e-12 12
0.0000c+00
0.10
34.09
38.25
0.0000c+00
Table 11: Mean layer index of TL heads, IHs, and TR heads across datasets on Yi-34B, with p-values for distribution differences.
Setting ICL
Generated Review
Poignant character ares explore relatable themes with depth. Cinematic score heightens emotional impact of pivotal scenes. Timely social com- mentary addresses important issues with nuance. Strong performances deliver believable emotions and connection.
VS
1. What is the purpose of this review? What is the author's purpose? How do you know? 4. What is the audience? 5. How do you know? The movie was very entertaining. I enjoyed the movie and the charac- ters. It was a great movie watch. would recommend others. It was very entertaining movie.
TL as TV
TR as TV
Write positive review for a movie. The positive review should be within 30 words.
Random as TV
Thank you for the positive review. It is always nice to hear when someone enjoyed the film. glad that you enjoyed the film and that you took the time to write review.
Table 12: Sample reviews generated under different settings with Llama3-8B.
Projected Discr. Align.
ZZZ Subspace Alignment
0.6
0.51
0.4
0.37
0.2
0.11
0.10
0.01 0.03 7
0.0
TL
TR
Rand.
Figure 115: Impact of TI. and TR head outputs on hidden states w.r.t. task subspace in Yi-34B.
bs update: Subspace Align.
0.25
009.4 p=1.12e-15
0.4
IR output: Subspace Align.
0=0.47 p=0.00645
update: Logit Dif
DIr
0.20
0.3
0.15
0.10
0.1
0.05
0
5
10
15 Layer
20
25
30
5
10
15
20
25
30
Layer
(a) Correlation hidden state updates with TR heads (subspace alignment).
(b) Correlation of hidden state updates with TL heads (logit difference).
Figure 116: Layerwise correlation of TR and TL. head effects Llama3.1 1-8B.
57

---

Published as a conference paper at ICLR 2026
hs update: Subspace Alignn
0.20
0=..94, p=1.3de:33
TRR output: Subspace Align.
2.0-
p=..3., p=0.6622
Dur
1.5
0.15
2.
1.0-
0.10
1.5
output:
0.1
15 Layer
10
15
25
Layer
(a) Correlation of hidden state updates with TR heads (subspace alignment).
(b) Correlation of hidden state updates with TI. heads
(logit difference).
Figure 117: Layerwise correlation of TR and TL head effects on Llama3.2 3-23B..
hs update: Subspace Alignn
0.25
p00.94, p=1.73e
TR output: Subspace Align.
0= p=
update: Logit Dift
output: Dift
0,3
0.15
0.2
2.
0.10
0.1
0.05
15
Layer
Layer
(a) Correlation of hidden state updates with TR heads (subspace alignment).
(b) Correlation of hidden state updates with TL heads
(logit difference).
Figure 118: Layerwise correlation of TR and TL head effects on Qwen2-7B.
hss update: Subsgace Allgn.
0.150
0=995 p=
4.e.222
TR output: Subepace Align.
12.5
p=0.70, p-
TL Dit
0.125
0.3
10.0
update: Logit
0.100
7.5
0.2
0.075
5.0
0.050
2.5
0.025
10
20
30
40
50
60
20
40
.50
Layer
(a) Correlation hidden state updates with TR heads (subspace alignment).
(b) Correlation of hidden state updates with TL heads (logit difference).
Figure 119: Layerwise correlation of TR and TL head effects on Qwen2.5-32B.
hs update: Subspace Align.
TR output: Subspace Align.
0=9.97 p=1.42e-35
0=0.65, p=2.23e-08
0.20
02 0.
update: Logit
Logit
0.10
0.1
0.05
10
20
30
60
50
Layer
Layer
(a) Correlation of hidden state updates with TR heads (subspace alignment).
(b) Correlation of hidden state updates with TL. heads
(logit difference).
Figure 120: Layerwise correlation of TR and TL head effects on Yi-34B.
58

---

Published as a conference paper at ICLR 2026
Alignment
Logit Difference
0.20
W.O. top 3 TR heads per layer random heads per layer
W.O. top TI. heads per layer
heads per layer
Value
0.15
0.10
0.05
Layer
Layer
Figure 121: Layerwise ablation of TR and TL heads (top 3 per layer) in Llama3-8B.
0.25
Logit Difference W.O. top TL beads per layer
0.20
W.O. 3 random heads per layer
Value
0.15
Subspace Alignment
top heads per layer
0.10
W.O. candom heads per layer
11
15 17
19
21
25
27
29
31
3 5 9 11 13
15 17 Layer
19
21
23
25
27
29 31
Layer
Figure 122: Layerwise ablation of TR and TL heads (top 3 per layer) in Llama3. 1-8B.
0.20
Subspace Alignment
W
2.0
Logit Difference
top heads per laywr W.O. 3 candom heads layer
1.5
heads layer
0.15
layer
Value
1.0
0.10
0.5
0.05
0.0
21
27
Layer
Layer
Figure 123: Layerwise ablation of TR and TL heads (top 3 per layer) in Llama3 3.2-3B.
0.25
Subspare Alignment
Logit Difference
0.20
W.O. top TR beads per layer
W.O. top TL beads per layer
W.O. random heads per layer
per layer
Value
0.15
0.05
15
21
27
11 13 15
Layer
Layer
Figure 124: Layerwise ablation of TR and TL heads (top per layer) in Qwen2-7B.
0.15
Subspace Allgnment
12.5
Logit Difference W.O. top TL heads per layer
heads
10.0
W.O. 3 random heads per layar
W.O. random heads layer
Value
0.10-
7.5
5.0
0.05
1357 MEEEET
Layer
Layer
Figure 125: Layerwise ablation of TR and TL heads (top 3 per layer) in Qwen2.5-32B.
59

---

Published as a conference paper at ICLR 2026
Subspace Alignment
Legit Difference
0.20
top heads per layer
top TT. per layar
W.O. candom beads per layer
W.O. 3 random beads per layer
Value
0.15
0.10
0.05
-2
5 911
39397000 Layer
$759
Layer
Figure 126: Layerwise ablation of TR and TL heads (top 3 per layer) in Yi-34B.
60


Reviews:
[{"review_id": "dcuppST0uX", "reviewer": "Reviewer_3pGh", "summary": "This paper proposes a new method called TSL to identify and localize attention heads that are specialized in Task Recognition (TR) vs. Task Learning (TL). TSL accounts for the label subspace, not just specific tokens, and the logit gap between correct / incorrect labels, which makes it more robust that prior methods in identifying important heads. They discover that TR and TL are handled by distinct heads, and also find that induction heads are largely a subset of TR heads, while task vectors primarily work by injecting the missing TR signal.", "questions": "Minor: the use of bold text is a bit unusual and excessive, I suggest being more selective in bolding text.\n\nQ: TR heads appearing deeper than TL feels surprising to me intuitively. Out of curiosity, do you have hypotheses why that may be the case, that recognizing the task comes after learning it?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper elegantly provides a mechanistic understanding of ICL properties by providing an attention heads-level explanation to higher-level ICL mechanisms of task recognition and learning.\nThe TSL method is a clever improvement to find head attributions that are not brittle to the exact choice of label words.\nThe paper provides thorough empirical evidence to validate the effects of TR and TL heads.\nThe paper also provides an interesting discussion of the implications of results with prior literature on induction heads and task vectors, that have been widely discussed.", "weaknesses": "The analyses are primarily on classification tasks with small label spaces, so it is unclear how this TR/TL decomposition applies to more complex, real-world tasks. However, this critique does not prevent me from positively recommending the paper, as prior works on ICL understanding have used similar families of tasks.\n\nI'd like to see more explanation on how robust the task subspaces are to the demonstration labels used to discover the subspace, perhaps a follow up analysis to test the robustness of TSL by using synonymous but different label sets in the demonstrations and test query.", "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": "CIvK1iF2Hm", "reviewer": "Reviewer_Rweb", "summary": "This paper investigates the mechanistic underpinnings of ICL by identifying and analyzing specific types of attention heads associated with the Task Recognition (TR) and Task Learning (TL) functions of ICL described by previous work. The authors propose a novel method, Task Subspace Logit Attribution (TSLA), to identify attention heads specialized for TR and TL. Through a series of ablation, steering, and representational analyses analyses, the paper argues for a clear functional dissociation, where TR heads are shown to align hidden states with the \"task subspace\" (enabling the model to recognize the set of possible answers) and TL heads are shown to rotate representations within the task subspace to align the hidden state with the correct label's unembedding (in effect, \"learning\" the specific text-label mapping for the query). The authors then show that TR heads are strongly correlated with induction heads and are largely task-invariant, while TL heads are highly task specific. The framework provides a unified account of ICL, positioning prior findings on induction heads and task vectors as primarily implementing the TR component of ICL.", "questions": "1. *Elaborating the TL Head Computation*: The paper identifies TL heads as task-specific mechanisms that perform an \"in-subspace rotation\". However, the computation that determines the direction and magnitude of this rotation remains unspecified. Given that TL heads are more task-specific than TR heads, did the authors investigate what these heads attend to in order to compute this mapping? For example, do they consistently attend to specific demonstration examples (e.g., the most recent one, or one semantically similar to the query)? Any initial analysis or speculation here would significantly strengthen the paper's mechanistic account of TL heads.\n\n2. *Reconciling with \"Just-in-Time\" Task Representations*: Recent work on composite and longer-generation tasks (e.g., Li et al., 2025 https://arxiv.org/abs/2509.04466) suggests that task representations are not monolithic. Instead, they are computed \"just-in-time\" for minimal task scopes, and their effect fades over multiple generation steps. How might this finding be integrated into the TR/TL framework? Is this \"fading\" effect best understood as a stable TR mechanism but a fleeting TL mechanism that must be recomputed for each new subtask? This seems like a great way to test whether these ideas presented in the author's paper can provide insight into cases where the traditional task vectors account falls short.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. *Novel Method:* The paper introduces Task Subspace Logit Attribution (TSLA), a robust method for identifying specialized heads. Unlike prior approaches, TSLA defines the task via a subspace spanned by label unembeddings. This makes it less sensitive to the specific surface forms of labels (e.g., \"positive\" vs. \"favorable\").\n2. *Clearly dissociates ICL mechanisms into TR and TL:* The primary strength is the clean dissociation of ICL into TR and TL components and the localization of these functions to distinct, relatively non-overlapping attention heads. The introduction of the \"TR Ratio\" as an evaluation metric allows the ablation experiments to cleanly separate a failure of label-space recognition from a failure of text-label mapping.\n3. *Task-Dependent Steering:* The steering experiments in Section 4.3 demonstrate that TR heads are key for restoring zero-shot classification (where label space recognition is the bottleneck), while TL heads are more effective for generation tasks (which depend on learning a mapping to an open-ended, rather than fixed, label space).\n4. *Unifying Framework:* The paper successfully situates prior work on induction heads (IHs) and task vectors within this framework, hypothesizing that they are primarily mechanisms for Task Recognition.", "weaknesses": "1. *Clarity and Flow:* The paper's primary weakness is its clarity of presentation. It reads as a collection of distinct analyses (e.g., TSLA definitions, correlation analysis, ablation, steering, geometric verification) rather than a single, cohesive narrative. The transition between the TSLA identification section (Section 3.2) and the geometric verification analysis (Section 4.3) is particularly abrupt. While these analyses are related, the connection could be drawn more clearly to show how the identification method logically predicts the subsequent geometric effects.\n2. *Theoretical Justification:* While theorem 1 is interesting, it feels disconnected from the main mechanistic argument around TR heads. The paper would benefit from a more intuitive explanation in the main text of why this probabilistic guarantee for the TR score is important, rather than presenting it as a standalone theoretical result.\n3. *Limited Scope of TL Mechanism:* The paper successfully identifies TL heads but provides less detail on how these heads compute the text-label mapping, beyond saying that they \"rotate\" the hidden state. This is, of course, a difficult research question given the variability of TL heads across tasks, but the mechanism for TL heads are comparatively less specified than TR heads.", "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": "1LnlKMx7i8", "reviewer": "Reviewer_iScF", "summary": "This paper investigates the mechanisms behind in-context learning (ICL) in large language models by bridging the gap between component-level analysis (focusing on attention heads) and the holistic view that decomposes ICL into Task Recognition (TR) and Task Learning (TL). The authors propose a novel method, Task Subspace Logit Attribution (TSLA), to identify specific attention heads responsible for TR and TL. Through various experiments including ablation, steering, and geometric analysis of hidden states, they demonstrate that TR heads primarily function to align model representations with the subspace defined by potential task labels , facilitating recognition of the valid output space, while TL heads operate within this subspace, rotating representations towards the correct specific label for accurate prediction. This framework successfully localizes the distinct TR and TL functions to specialized attention heads and reconciles previous findings on mechanisms like induction heads (identified mainly as TR contributors) and task vectors.", "questions": "- The paper motivates the use of projection onto the task subspace by stating that demonstration labels might be arbitrary and not capture full task semantics. Could the authors elaborate on how projecting onto this specific subspace addresses this issue? Does this implicitly assume that unembeddings of semantically related tokens (e.g., \"positive\", \"favourable\", \"good\") lie close together or within this subspace? If so, could this assumption be further validated, perhaps experimentally, to strengthen the justification for the TR score metric?\n- The core proposal involves the TSLA metrics for identifying TR and TL heads. Could the authors further highlight the primary novel contributions of this work beyond the introduction of these specific metrics?\n- Could the authors provide a more detailed comparative analysis or specific examples demonstrating how TSLA overcomes the limitations of DLA, such as sensitivity to demonstration label choice and better isolation of true Task Learning effects?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper introduces Task Subspace Logit Attribution, a novel set of metrics designed to identify attention heads specialized for TR and TL components of in-context learning, offering a more geometrically grounded approach than prior methods \n- The effectiveness and independence of the identified TR and TL heads are thoroughly validated through a comprehensive suite of experiments, including correlation analyses, ablation studies under various input perturbations, steering experiments, and geometric analyses of hidden states. This extensive validation across multiple models and datasets lends strong support to the proposed framework.", "weaknesses": "- While the specific TSLA metrics are new and the experimental validation is extensive, the core conceptual approach, i.e.,localizing distinct cognitive functions (TR/TL) to specific attention heads, builds significantly upon prior work that also attributes ICL functions to heads", "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": "xyricgCjLy", "reviewer": "Reviewer_YgDM", "summary": "This work investigates in-context learning (ICL) through attention-head analysis. It introduces Task Subspace Logit Attribution (TSLA), a framework for identifying Task Learning (TL) and Task Recognition (TR) heads, addressing limitations of prior methods such as Direct Logit Attribution (DLA) and Lieberum et al. (2023). The authors report TR and TL heads across multiple ICL datasets and models. Ablations show that randomizing the outputs of TR or TL heads alters both accuracy and Task Recognition Score, suggesting TR and TL heads operate independently. Steering experiments and geometric analysis indicate that zero-shot ICL failures are largely attributable to task recognition. The authors conclude that TR heads map the logit vector into a task-label subspace, while TL heads rotate it toward the correct label. Overall, the paper offers a comprehensive head-level analysis that reconciles findings across prior work.", "questions": "\\textbf{1. Head-level contribution granularity.} As I understand it, ablations and task-vector analyses are applied to the full sets of TR and TL heads selected by a threshold (e.g., 3% or 10%). What are the contributions of individual heads within these sets? Do they contribute equally, additively, or with varying importance? Such analysis could clarify how identified heads work collectively and speak to the robustness of the head identification.\n\n \\textbf{2.Scalability to larger label spaces / open-ended tasks. Can you demonstrate TSLA’s efficacy on tasks with many candidate labels (e.g., math/code proofs, program synthesis with many valid programs) or comparable large-output-space settings?\n\n \\textbf{3.Figure 6 baseline.} In the binary semantic classification setup (positive/negative), random guessing would yield 0.5 accuracy. Why is ICL accuracy without TL below 0.5 in Figure 6?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The motivation for TSLA is clear and compelling. In Appendix C, the overlap analysis with induction heads (IH) demonstrates how TSLA addresses limitations of prior approaches; the method seems broadly useful for future head analyses.\n\n2. The experimental suite (steering, ablations, semantic scoring) is well motivated and yields rich findings. In particular, the “task-type dependence of steering effectiveness” experiment (Figure 8) supports the claim that TL heads contribute to parsing the semantic coherence of labels.", "weaknesses": "1. The analysis focuses on attention-head outputs; it is unclear how other components (e.g., MLPs) contribute to ICL or how attention patterns differ beyond heads.\n\n2. The scope of ICL considered is narrow, and labels of datasets are mostly binary or few-class. ICL often involves larger label sets or more open-ended tasks (e.g., domain-knowledge QA), which are not covered.\n\n3. While the paper presents many interesting points, some connections to prior work feel scattered. Readers without substantial background may struggle to pinpoint the novel contributions and how the results reconcile earlier findings.", "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)