Title: Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: In-Context Learning; Induction Head; Mechanistic Interpretability; Task Vector;
Abstract: The unusual properties of in-context learning (ICL) have prompted investigations into the internal mechanisms of large language models. Prior work typically focuses on either special attention heads or task vectors at specific layers, but lacks a unified framework linking these components to the evolution of hidden states across layers that ultimately produce the model’s output. In this paper, we propose such a framework for ICL in classification tasks by analyzing two geometric factors that govern performance: the separability and alignment of query hidden states. A fine-grained analysis of layer-wise dynamics reveals a striking two-stage mechanism—separability emerges in early layers, while alignment develops in later layers. Ablation studies further show that Previous Token Heads drive separability, while Induction Heads and task vectors enhance alignment. Our findings thus bridge the gap between attention heads and task vectors, offering a unified account of ICL’s underlying mechanisms.

Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning

Abstract

The unusual properties of i nc ontext l earning (ICL) have prompted investigations into the internal mechanisms of large language models. Prior work typically focuses on either special attention heads or task vectors at specific layers, but lacks a unified framework linking these components to the evolution of hidden states across layers that ultimately produce the model's output. In this paper, we propose such a framework for ICL primarily in classification tasks by analyzing two geometric factors that govern performance: the separability and alignment of query hidden states. A fine-grained analysis of layer-wise dynamics reveals a striking two-stage mechanism-separability emerges in early layers, while alignment develops in later layers. Ablation studies further show that Previous Token Heads drive separability, while Induction Heads and task vectors enhance alignment. Our findings thus bridge the gap between attention heads and task vectors, offering a unified account of ICL's underlying mechanisms. 1

1 Introduction

One of the most remarkable features of L arge L anguage M odels (LLMs) is their ability to respond to novel queries in user-desired mannerisms on-the-fly solely by learning from demonstrations provided in the input (as shown in Figure 1 (A))-without any additional training. This capability is known as I nc ontext L earning (ICL) [Brown et al., 2020, Radford et al., 2019]. ICL has revolutionized natural language processing by reducing dependence on burdensome data collection and costly finetuning, enabling swift and seamless adaptation of LLMs to a variety of downstream tasks [Dong et al., 2024].

Due to its stark departure from traditional gradient-based learning paradigms, ICL has attracted significant academic interest. One line of research aims to develop a mechanistic understanding of ICL by analyzing the internal behavior of LLMs in ICL settings [Reddy, 2024]. Some studies highlight the importance of key structural components in LLMs, such as I nduction H eads (IH) [Elhage et al., 2021, Olsson et al., 2022, Cho et al., 2025a], which retrieve information from demonstrations and perform copy-like operations within Transformer-based architectures. Other studies adopt a hiddenstate-centric view, interpreting ICL as a process in which LLMs construct vector representations of the task from demonstrations (so-called task vectors [Hendel et al., 2023, Liu et al., 2024]) using the hidden states at an intermediate layer, and then use the vectors to steer hidden states and produce better next-token predictions. Both lines of work support their claims primarily through ablation and intervention experiments that assess the impact on model outputs.

However, treating model components and hidden states as separate entities in ICL analysis is both unnecessary and misleading. The two are inherently intertwined-components like attention heads transform hidden states layer by layer, and these intermediate updates ultimately determine the

I like this. Sentiment: positive \n It is bad. Sentiment: negative \n Fantastic! Sentiment:

LM Layers

(PTH)

LM Heads

Positive

★

★

Separability

(A)

(B)

LM Layers

(IH)

(C)

★

★

Unembedding Direction

Optimal Decoding

(D)

★

★

☆

☆

Unembedding Direction

Task Vector

= negative samples

= positive samples

Alignment

model's output. As such, analyzing only the influence of attention heads or task vectors on the final prediction reduces ICL to a black-box process. A true mechanistic understanding of ICL requires tracing how these components progressively shape hidden states throughout the model, not just their end effects. A unified framework that synthesizes both perspectives and assesses their importance in the layer-wise evolution of hidden states is thus highly essential.

Therefore, in this work, we propose a framework, shown in Figure 1, based on a key geometric insight into LLM classification: the success of LLMs in classification depends on whether the hidden states of queries with different ground-truth labels are (1) sufficiently separable and (2) well aligned with the unembedding vectors of labels. In other words, hidden states should be separated along a direction that aligns with the difference between the unembedding vectors of different labels. To investigate this, we collect hidden states in both ICL and zero-shot settings and analyze them using a range of geometric measures to assess how ICL improves hidden states in terms of separability and alignment. Our approach operates on three levels: (1) We examine the layer-wise trends in these measures and compare zero-shot with ICL (and across different ICL settings) to quantify the effect of ICL on the hidden state geometry. (2) We observe how changes in measure values reflect the evolution of semantic information represented in the hidden states that ultimately inform model outputs in ICL. (3) We probe specific layers to examine the role of different types of attention heads and how they influence both the values and trajectories of these geometric measures.

Our experiments find that ICL improves classification performance primarily by enhancing alignment between query hidden states and unembedding vectors, with less impact on separability. Results further suggest that ICL proceeds as a two-stage process: early layers mainly increase hidden-state separability, while in middle-to-late layers, separability plateaus and alignment improves. The refinement of alignment arises through layer updates that filter out task-irrelevant semantics and preserve label-relevant content, thereby semantically refining query hidden states. Such updates qualify ICL query hidden states as task vectors. We also test this two-stage characterization on ICL for a generation task and confirm its explanatory power across settings.

To build the bridge between the task vector and model components, we analyze the effect of attention heads on these geometric measures and find that (as demonstrated by Figure 1): (1) The early increase in separability is driven by a proliferation of P revious T oken H eads (PTHs)-attention heads that attend to the previous token at each position-in early layers. (2) In middle-to-late layers, I nduction H eads (IHs) enhance alignment by amplifying hidden state components along the unembedding directions of the labels. Thus, our findings not only explain why attention heads like PTHs and IHs are essential for ICL performance, but also clarify why induction head outputs empirically serve as effective task vectors [Todd et al., 2024, Hendel et al., 2023].

2 Related works

IH and PTH Previous works on the mechanistic cause of ICL identify structural circuits in Transformers comprising a P revious T oken H ead (PTH) and an I nduction H ead (IH), where the PTH gathers information at demonstration label tokens, and the IH copies it to query positions for prediction [Olsson et al., 2022, Singh et al., 2024, Song et al., 2025]. Their importance in classification has been confirmed in large-scale models as well [Cho et al., 2025a]. However, existing analyses focus on end-to-end effects, leaving open how these heads shape query hidden states.

Task vectors An alternative theory proposes that LLMs summarize demonstrations into task vectors, which can be used to steer zero-shot hidden states to increase prediction accuracy. These vectors may be ICL query hidden states extracted at an intermediate layer [Hendel et al., 2023, Liu et al., 2024], attention head outputs [Todd et al., 2024, Li et al., 2024], or MLP outputs [Merullo et al., 2024]. Decoding task vectors through the unembedding matrix reveals that they encode task-related tokens and support semantic manipulations akin to 'word2vec'-style algebra [Mikolov et al., 2013]. Despite intriguing empirical findings, explanations for their mechanism remain preliminary. Kahardipraja et al. [2025] show that heads expressing contextual information or model's parametric knowledge produce task vectors with distinct function. Jiang et al. [2025b] and Han et al. [2025] study how the geometry of model's task representation impacts the effects of injecting task vectors. Jiang et al. [2025a] reveal the correlation between model's catastrophic forgetting and the change in task vectors. However, none of these studies answer the question of how task vectors influence the concrete internal computation process of models which directly and ultimately determine the outputs.

Geometry of hidden states LLMs encode inputs in geometrically structured ways [Park et al., 2024]. For instance, they encode cities and countries by geographic relations [Gurnee and Tegmark, 2024] and numbers in rings under modular arithmetic [Liu et al., 2023]. For classification tasks, they can separate hidden states into clusters by attributes such as truth value [Marks and Tegmark, 2024] or toxicity [Lee et al., 2024]. However, the alignment of such clusters with the unembedding vectors of their labels-an equally crucial geometric factor-has received limited attention, with only a few studies noting suboptimal alignment at the output layer [Cho et al., 2025b, Kirsanov et al., 2025].

3 Theoretical framework

We begin our framework with an analysis of classification tasks, where an LLM is prompted with a query x and its ground-truth label y , e.g., x = 'I like this movie. Sentiment:' and y = 'positive'. The LLM with L layers encodes x into a vector hx ∈ R d , taken from the hidden state at the final token (':'). Given a vocabulary V = { v 1 , ..., v | V | } and an output embedding matrix E ∈ R | V |× d , the model predicts p ( v j | x ) = exp ( Evj hx ) / ∑ | V | j ′ = 1 exp ( Ev j ′ hx ) , where Evj is the unembedding vector of v j . The model outputs the token v j with the highest probability as the predicted label for x .

In the ICL scenario, as shown in Figure 1 (A), k demonstration-label pairs ( x 1 , y 1 ) , ..., ( xk , yk ) are prepended to form the prompt [ x 1 , y 1 , ..., xk , yk , x ] , which modifies the hidden state of x and results in a different distribution p ( v j | x 1 , y 1 , ..., xk , yk , x ) , potentially improving prediction accuracy.

Rather than a single query, we more often evaluate classification accuracy over a dataset of n queries { x 1 , ..., xn } , with hidden states { h 1 , ..., hn } and ground-truth labels { y 1 , ..., yn } . We will show that the classification accuracy depends on two geometric properties of the hidden state collection H =[ hi ] n i = 1 ∈ R n × d : (1) its separability, and (2) its alignment with the unembedding vectors.

As mentioned before, an LLM performs classification using the unembedding matrix E , which is a multiclass classifier on the hidden states. We define its accuracy on the dataset as:

where the first equality states that classification accuracy equals the proportion of the labels being predicted as the most probable token, and the second equality uses the monotonicity of softmax.

We first consider the simple case where { x 1 , ..., xn } only has 2 different labels, yA and yB , i.e. yi ∈ { yA , yB } , ∀ i , and leave the generalization to the case of more labels to Appendix B. For a unit vector u ∈ R d , we define the separability of H along direction u with labels yA , yB to be:

where N A = { i : yi = yA } and N B = { i : yi = yB } . S ( u ) measures the fraction of hidden states that can be correctly classified by a linear decision boundary orthogonal to u . By considering all possible directions, we obtain the maximum separability of H , i.e. S ∗ = sup u ∈ R d , ∥ u ∥ 2 = 1 S ( u ) . Though S ( u )

is not a continuous function, S ∗ as a supremum can be attained on S d -1 = { u : u ∈ R d , ∥ u ∥ 2 = 1 } (See Appendix A). Denote u ∗ as the u where S attains the supremum, then:

Theorem 1 Acc ≤ S ∗ . The equality is achieved when max v ∈ V , v / ∈{ y A , yB } Evhi < max ( EyA hi , EyB hi ) , ∀ i and Ey A -Ey B ∥ Ey A -Ey B ∥ 2 = cu ∗ for some positive constant c.

This theorem shows that accuracy is upper-bounded by the maximum separability of H and that the bound is achieved when: (1) each hi is maximally aligned with either EyA or EyB -the unembedding vectors of yA and yB -while excluding interference from other tokens, a condition we call output alignment ; and (2) the direction of maximum separation u ∗ aligns with EyA -EyB , which we refer to as directional alignment . Thus, separability and alignment of hidden states critically determine accuracy. Since query hidden states in ICL differ from those in the zero-shot setting due to the presence of demonstrations, we conclude that ICL can improve classification accuracy by enhancing both their separability and alignment.

Attention heads influence query hidden states as well by their outputs to the residual stream. At layer l , the h th attention head adds a h i , l to the hidden state hi , l of the i th query, affecting the transition to hi , l + 1 . Ablating a head (i.e., setting a h i , l = 0) alters all downstream hidden states and can degrade accuracy if the final-layer hidden states HL lose separability or alignment.

Task-vector -based experiments operate similarly by steering hidden states at intermediate layers. In those experiments, a task vector t i is added to 2 hi , l and the influence likewise propagates through subsequent layers. This intervention is effective if the final HL has better separability and is better aligned with the labels of the task encoded in t i .

4 Method: Measuring separability and alignment of hidden states

Given the aforementioned influence of ICL, attention heads, and task vectors on the separability and alignment of hidden states across layers, we introduce several measures 3 to quantify their influence. The formal definitions and mathematical details of all measures are provided in Appendix C.

A. Separability measure: Separability score

Theorem 1 shows that the classification is affected by the maximum separability of hidden states. However, finding maximum separability requires evaluating separability along infinitely many directions in S d -1 and is intractable. As a practical proxy, we train a logistic classifier on a subset of H and evaluate its accuracy on held-out data. We call the accuracy separability score as it reflects the empirical maximum separability of hidden states achieved by the classifier.

B. Alignment measures

In Theorem 1 we show that both output alignment and directional alignment affect accuracy. Below, we introduce measures for both.

B.1 Output alignment

Output alignment is defined as the classification accuracy obtained by applying the unembedding matrix E directly to hidden states H , skipping subsequent layers. It measures how well H aligns with the label unembedding vectors EyA and EyB , and is also known as logit lens accuracy [nostalgebraist, 2020]. When H comes from the final layer, it equals standard classification accuracy.

B.2-B.5 Directional alignment

Directional alignment is the alignment between u ∗ and the label unembedding difference direction

Ey A -Ey B

∥ Ey A -Ey B ∥ 2 . Since finding u ∗ is intractable, we introduce the following proxies based on approximating u ∗ or measuring separability along the label-difference direction.

Singular alignment (B.2) . To approximate u ∗ , we apply Singular Value Decomposition (SVD) to the mean-centered H and extract the topr right singular vectors. These are the directions where H has the greatest spread and can approximate u ∗ . We compute the maximum absolute cosine similarities between each of these vectors and Ey A -Ey B ∥ Ey A -Ey B ∥ 2 as a proxy for directional alignment.

Variance-based alignment (B.3) . To measure separability along Ey A -Ey B ∥ Ey A -Ey B ∥ 2 , we first measure the proportion of total variance in H that lies along this direction. A higher value indicates that this direction is a principal axis of variation-suggesting stronger directional alignment.

Mean-based alignment (B.4) . We project the difference in means of label clusters-hidden states with labels yA and yB -onto Ey A -Ey B ∥ Ey A -Ey B ∥ 2 and divide it by the pooled projected variance. A higher value implies better separation margin along Ey A -Ey B ∥ Ey A -Ey B ∥ 2 and better directional alignment.

Composite alignment (B.5) . To unify the mean- and variance-based alignment, we define composite alignment as their product. It captures both the magnitude and the relative significance of separation along the label-difference direction, serving as a robust proxy for directional alignment.

C. Indirect measure: Effective dimension

As an indirect indicator of separability and alignment, we compute the effective dimension of H -the number of directions along which the hidden states exhibit significant variance. A low effective dimension suggests that hidden states concentrate along a small number of directions, potentially those associated with the label unembedding vectors, and are thus easier to separate.

5 Experiments and Results

Models We experiment on the following 7 models: Llama2-7B, Llama2-13B, Llama2-70B [Touvron et al., 2023], Llama3-8B, Llama3-70B [Grattafiori et al., 2024], Gemma-2B, and Gemma-7B [Gemini Team et al., 2024]. Unless otherwise stated, we report the results on Llama2-70B.

Datasets We conduct experiments on two major types of classification datasets: text classification (SUBJ [Wang and Manning, 2012], SST-2 [Socher et al., 2013], TREC [Li and Roth, 2002]) and natural language inference (SNLI [MacCartney and Manning, 2008], RTE [Dagan et al., 2005], and CB [De Marneffe et al., 2019]). We further include a generation dataset (detailed in Subsection 5.1) to test the generalizability of our results in settings other than classification. We also use a generated dataset to explore the scenario where ground-truth labels of queries are not presented in the demonstrations (discussed in Appendix I.5).

ICL setting We set the number of in-context demonstrations to be 8 unless otherwise stated. The demonstration samples are chosen randomly, except for one experiment that investigates the effect of using demonstrations procured by kNN retrieval [Liu et al., 2022]. For a detailed exposition of the implementation of models, datasets, prompt templates, etc., refer to Appendix D.

5.1 Layer-wise trends of separability and alignment

ICL as a two-stage process Figure 2 shows the values of six geometric measures at all layers of Llama2-70B in both zero-shot and ICL settings averaged across datasets (output alignment is deferred to Subsection 5.2 due to its strong link to the semantic information of hidden states; results for other models are in Appendix F). A key observation is that Figure 2 reveals a striking phase transition in the evolution along the layers of ICL hidden states. Phase 1 (Figure 1 (B)): In the early layers, directional alignment measures take low values, whereas separability increases rapidly. Phase 2 (Figure 1 (D)): In the subsequent layers, this trend reverses: separability plateaus, while all directional alignment measures spike simultaneously, which suggests that the label unembedding difference direction starts to capture more variance and separate label clusters more effectively. It causes the semantic information of query labels to be predominant in H , which ultimately leads to final output alignment. These findings are consistent with prior work showing that task information emerges only after the initial layers [Sia et al., 2024]. The alternation between two phases is also captured by the effective dimension of H , which first rises and then declines as alignment starts to improve, reflecting a gradual concentration of hidden states along label unembedding directions. In contrast, this transition is much weaker in the zero-shot setting, where effective dimension rises

0

20

40

60

80

Layer Number

0

25

50

Effective dimension

0

20

40

60

80

Layer Number

0.6

0.8

Separability score

0

20

40

60

80

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.02

Composite alignment

0

20

40

60

80

Layer Number

0

2

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.005

0.010

Var.-based alignment

8-shot

0-shot

monotonously and the rise in alignment measures is not nearly as pronounced. Hence, this phase transition pattern characterizes ICL as a two-stage process : an initial phase enhancing separability, followed by a refinement phase that improves alignment.

ICL improves accuracy primarily through alignment Akey observation is that ICL and zero-shot settings show only a small gap in separability scores-despite an 80% (Figure 2, upper middle) difference in classification accuracy. This implies that the semantic separation of query hidden states is intrinsic to LLM inference rather than a main effect of ICL (see Appendix E.1 for the practical implication of this), making alignment the primary bottleneck in zero-shot performance. As shown in Figure 1 (C), the high-separability direction in zero-shot hidden states is misaligned with the labeldifference vector, as reflected in low singular and composite alignment scores. Along this direction, label clusters remain poorly separated (low mean-based alignment) and explain little variance (low variance-based alignment). These findings align with prior work indicating that unembedding vectors often fail to capture the dominant separation axes in hidden state space [Cho et al., 2025b].

Given the volatility of ICL accuracy across settings with different demonstration numbers and input formats, we investigate whether the phase transition persists across settings, and whether accuracy differences can similarly be attributed to alignment properties of hidden states. We consider three settings: (1) Varying the number of demonstrations from 0 to 24 in increments of 4; (2) Changing the demonstration selection strategy to kNN retrieval where demonstrations with the closest zero-shot embeddings to each query are chosen; (3) Replacing original demonstration labels with semantically uninformative symbols like '@' and '#'. Figure 3 reports the results averaged across datasets.

Phase transition is a robust hallmark of ICL Figure 3 shows that the phase transition pattern persists across all three settings. In Figure 3 (A), it is evident at all demonstration counts except 0. Figure 3 (B) confirms that principled demonstration selection (via kNN) preserves instead of interfering with the transition. Even in Figure 3 (C), where labels are replaced by symbols, the two-stage pattern in effective dimension remains-distinguishing it from zero-shot, where effective dimension rises monotonically-though the post-transition alignment surge is dampened.

Phase transition holds in open-ended settings To test whether the phase transition pattern occurs in tasks beyond classification, i.e., open-ended generative tasks, where ICL is also frequently applied, we consider the following task: the queries assume the format 'Praise/Critique {subject} ', where {subject} is the name of a food (e.g., tapas). The labels are reviews for the subject satisfying the sentiment requirement, generated using GPT-4o [OpenAI et al., 2024], for instance, 'Praise tapas → The tapas assortment at this restaurant is absolutely delightful!'. We track the layer-wise dynamics of the separability score and composite alignment 4 as representatives of the metrics evaluated on the 8-shot and 0-shot hidden states of this task. The results in Figure 4 mirror those in Figure 2, where

Llama2-70B: Average across datasets

(b) KNN demonstration selection

0

10

20

0.0

0.2

0.4

0.6

0.8

0

20

40

60

80

Layer Number

0

20

40

60

Effective dimension

0

20

40

60

80

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

20

40

60

80

Layer Number

0.00

0.05

0.10

0.15

0.20

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.01

0.02

0.03

0.04

Composite alignment

0

4

8

12

16

20

24

(a) Varying number of demonstrations

0

50

Layer Number

0

20

40

Effective dimension

0

50

Layer Number

0.6

0.8

Separability score

0

50

Layer Number

0.0

0.1

0.2

Singular alignment

0

50

Layer Number

0.00

0.05

Composite alignment

0

50

Layer Number

0

2

4

Mean-based alignment

0

50

Layer Number

0.00

0.01

Var.-based alignment

Llama2-70B: Random Selection Acc.: 80.57%; KNN Selection Acc.: 88.01%

Random Selection

KNN Selection

0

50

Layer Number

0

20

40

Effective dimension

0

50

Layer Number

0.6

0.8

Separability score

0

50

Layer Number

0.0

0.1

0.2

Singular alignment

0

50

Layer Number

0.00

0.02

Composite alignment

0

50

Layer Number

0

2

Mean-based alignment

0

50

Layer Number

0.000

0.005

0.010

Var.-based alignment

Llama2-70B: Normal Label Acc.: 80.57%; Symbol Label Acc.: 59.05%

Normal Label

Symbol Label

(c) Symbolized demonstration labels

the few-shot and zero-shot hidden states exhibit similar separability properties across layers, but the surge in composite alignment beginning from mid-layers is specific to few-shot hidden states.

0

8

16

24

32

40

48

56

64

72

80

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

24

32

40

48

56

64

72

80

Layer

0.000

0.002

Compo. Align.

0-shot

8-shot

Alignment explains accuracy differences across ICL settings We find that: accuracy differences between ICL settings are consistently reflected in alignment measures. Specifically, in Figure 3 (A), increasing the number of demonstrations from 4 to 24 produces minimal changes in separability, while alignment measures improve in a consistent manner as the trend curves are shifted versions of one another. Effective dimension also decreases steadily, suggesting a higher concentration of hidden states along label-unembedding directions. Figure 3 (B) similarly shows that the accuracy gains from kNN-selected demonstrations correspond to stronger alignment and lower effective dimension. In Figure 3 (C), the accuracy drop from symbolic label replacement aligns with a flattened post-transition rise in alignment measures, resembling the case of zero-shot.

5.2 Semantic interpretation of alignment dynamics

To understand how the surge in directional alignment after the phase transition leads to high output alignment and successful decoding of task-related labels at the final layer, we conduct a case study on the SST-2 dataset to interpret hidden-state alignment dynamics.

Surge in alignment concurs with the encoding of label-relevant semantics First, we compute output alignment (logit lens accuracy) across all layers and compare its trajectory to directional alignment measures. As shown in Figure 5, their strong correlation indicates that output alignment rises in sync with directional alignment. This suggests that as alignment improves, label-relevant

semantics are injected into hidden states via consecutive layer updates, enabling correct label decoding even at intermediate layers. To support this, we decode the top right singular vectors of hidden states near the transition point. Figure 5 (B) reveals a sharp semantic shift: while layer 20 produces irrelevant tokens, the SST-2 label token negative emerges as the second most probable token in layer 21-immediately post-transition. This confirms that post-transition layers elevate label-relevant directions in both the semantic and variance structure of the hidden states.

LM Head

Phase Transition

Phase Transition

(A)

ependant, pra, pra, anth, ...

Layer 20

...

ependant,

negative

, pra, anth, Klein, ...

Layer 21

Phase Transition

ependant,

negative

, iten, anth, Martí, ...

Layer 22

...

...

(sorted)

Differ

Coexist

LM Head

LM Head

Allen, Baden, ior, engo, Tun, Bobby, lichkeit, ieu, ...

Layer 27

Bot, avel,

序

, enf, Tob, batch, sed, tent, Stern, -

Layer 50

Cot, yw, éb, ershell, Girl, Schwe, ass, resid, ...

Layer 70

asser, orp, lif, Kost,

disappoint

, anth,

disappoint

, ...

Layer 27

jo,

positive

, inverse, erde,

solemn

,

humor

,

critic

, ...

Layer 50

humor

, .,

mixed

, description, intros,

comedy

, vague, ...

Layer 70

cel, 
, Ap, Steph, Job, fr, Og, author, super, exhaust

Layer 80

Retained Semantics among Layers

beskre, redund, pra, amb, isat,

humor

,

False

,

confused

,

Layer 80

Filtered (Removed) Semantics among Layers

(B)

(C)

1.0

1.0

ignment

0.8

0.08

Aliqnment

0.8

0.30

ignment

0.25

0.06

0.6

0.6

0.20

Alig

Output

0.4

0.04

Composite

Output

0.4

0.15

Singular.

0.10

0.2

0.02

0.2

0.05

0.0

0.00

0.0

0.00

0

20

40

60

80

0

20

40

60

80

LayerNumber

Layer Number

Layer-wise pattern of semantic retention and filtering To better understand how exactly layer updates after the phase transition encode relevant semantics into hidden states, we analyze the semantic information retained or removed by each layer. Given the SVDs of the centered hidden states at two consecutive layers ¯ Hl -1 = Ul -1 Σ l -1 V ⊤ l -1 and ¯ Hl = Ul Σ lV ⊤ l , we compute the projection of the previous layer's singular directions onto the next layer's as V ⊤ l -1 V lV ⊤ l ∈ R n × d . The rows of this matrix with larger norms are the principal directions of Hl -1 that are retained in Hl , while those with smaller norms indicate filtered directions. By decoding these directions, we identify the semantic information retained or removed by each layer. The results are reported in Figure 5 (C) . In most layers after the phase transition, labelrelated semantic directions are consistently retained, and unrelated direc- tions are filtered, leading to a progressive amplification of label-relevant information in the hidden states and a suppression of the irrelevant semantics. However, in the final layers, this pattern reverses: label-related directions are filtered, and label-irrelevant semantics are reintroduced. This shift corresponds to the slight drop in output alignment in the final layers in Figure 5 (A) , and aligns with prior findings that late layers in LLMs tend to encode high-frequency, semantically uninformative tokens [Sharma et al., 2023], which may interfere with the expression of task-relevant content. The same phenomenon also occurs in other datasets as presented in Appendix G.4.

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

Low-rank denoising enhances label-relevant semantics To better understand the retentionand-filtering process, we test whether filtering can be accelerated-and label-relevant semantics enhanced-by denoising the hidden states. Specifically, we apply a rank-10 SVD approximation to the centered hidden states ¯ H , retaining only the top 10 directions of variation and adding the mean back. This denoising should theoretically remove less informative components and amplify label-relevant semantics. We then feed the denoised hidden states into the unembedding layer and measure output accuracy. As shown in Figure 6, low-rank denoising yields accuracy gains exceeding 10%, especially in early post-transition layers where irrelevant semantics are just beginning to be filtered. These results validate low-rank denoising and support semantic retention and filtering as a core mechanism of post-transition layers. We explore the practical implication of this in Appendix E.2.

Task vector properties of post-transition hidden states As visualized by Figure 1 (D), the simultaneous enhancement of both alignment properties and label-relevant semantics after the phase transition qualifies ICL query hidden states in middle-to-late layers as task vectors [Hendel et al., 2023]. Due to their improved alignment properties, steering or replacing zero-shot hidden states

Normal Acc.: 80.57% | w/o PTH Acc.: 33.11% | w/o IH Acc.: 2.50% | Random Heads Ablated Acc.: 79.13%

0

20

40

60

80

Layer Number

0

100

Effective dimension

0

20

40

60

80

Layer Number

0.6

0.8

Separability score

0

20

40

60

80

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.02

Composite alignment

0

20

40

60

80

Layer Number

0

2

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.005

0.010

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

with ICL hidden states can align zero-shot hidden states better with the label unembedding vectors. Similarly, modifying zero-shot hidden states with ICL hidden states, which are label-informative, makes them more suited for the prediction of the correct labels.

5.3 Significance of critical attention heads

PTHs and IHs are attention heads with distinctive attention patterns. Given a token sequence like [A][B1][A][B2]...[A][B n ][A], a PTH at the final [A] attends to the immediately preceding [B n ], while an IH attends to tokens such as [B1]-[B n ] following earlier [A] positions. For example, in the ICL prompt 'I don't like it. Answer: negative. I like it. Answer:', at the final ':', a PTH attends to the preceding 'Answer', while an IH attends to 'negative' after the earlier ':', effectively copying semantics of the demonstration label into the query's hidden state. By analyzing their effects on separability and alignment across layers, we gain a more fine-grained understanding of their role in ICL-beyond simply measuring their impact on final predictions.

For each dataset and model, we rank all attention heads by their PTH and IH scores, which quantify how strongly they exhibit corresponding attention patterns on test prompts. We then ablate (zero out) the outputs of the top 10% identified IHs and PTHs and measure the resulting effects on hidden state separability and alignment across layers as well as the output accuracy. As a control, we repeat the procedure with an equal number of randomly selected heads (excluding the top PTH/IH heads). Details on identifying PTHs and IHs are provided in Appendix H.

PTHs induce separability, IHs induce alignment The results in Figure 7 confirm that PTHs and IHs are crucial not only for verbalizing ICL outputs [Cho et al., 2025a] but also for structuring hidden states both separability-wise and alignment-wise. Ablating random attention heads has a negligible effect (as shown by overlapping black and gray curves). In contrast, ablating PTHs or IHs significantly alters the measures, but with distinct effects. Ablating PTHs substantially reduces separability and also influences mean-based alignment (and thus composite alignment) which concerns the separation margin of the label clusters, but has a minimal impact on other alignment measures. Conversely, ablating IHs leaves separability largely intact but severely disrupts all alignment measures, and has a far greater impact on output alignment as measured by the accuracy.

These results not only confirm a two-stage ICL process-first driven by separability, then by alignment-but also clarify its mechanism. PTHs, concentrated in early layers [Cho et al., 2025a], enhance separability by attending to query tokens and encoding distinct semantics into hidden states. IHs, emerging after a critical depth (30-40% of total layers [Halawi et al., 2024]), copy demonstration label embeddings into query hidden states via the residual stream, im-

proving alignment with label unembedding directions. Furthermore, in Appendix I.5, we show that

PTHs and IHs contribute to ICL through the same mechanism even when the query's ground-truth label-previously considered essential for IHs [Cho et al., 2025a]-is absent from demonstrations.

PTHs in early layers, IHs in late layers To test whether PTHs and IHs align with our two-stage ICL characterization-inducing separability and alignment respectively-we examine their layer distributions. Specifically, we locate the top 1%, 2%, 5%, and 10% identified PTHs and IHs and perform Mann-Whitney U tests to assess differences across layers. As shown in Table 1, PTHs mainly occur in earlier layers, consistent with their role in inducing separability, while IHs appear later, with the differences being highly significant under the Mann-Whitney test.

0

8

16

24

32

40

48

56

64

72

80

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

40

48

56

64

72

80

Layer

0.000

0.005

0.010

0.015

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

Geometric effects of IH and PTH outputs as task vectors The ablation in Figure 7, though revealing how removing PTHs and IHs affects separability and alignment, does not show how they actively drive the geometric evolution of hidden states during inference.

To address this, we follow Todd et al. [2024] and inject their outputs as task vectors into zero-shot hidden states at selected layers to observe induced metric changes. We construct task vectors using outputs of the top 10% IHs and PTHs as in Appendix J. We inject the PTH-derived vector at layer l = 2 to track separability, and the IH-derived vector at 3 8 L to track composite alignment. These layers correspond to where separability and alignment begin to shift markedly in Figure 2. Figure 8a shows that injecting PTH outputs as task vectors boosts separability (the green line surpassing the magenta line before convergence; see Appendix I.2 for similar trends across models), while injecting IH task vectors restores the ICL-specific spike in alignment metrics in zero-shot hidden states.

0-shot

TV

ICL

0

20

40

60

80

Accuracy (%)

0.24

65.11

80.57

IH outputs as task vectors Since we have shown that IH outputs with ICL demonstrations can compensate for alignment deficits in hidden states, we have clarified how IH outputs as task vectors influence the forward computation of models: they steer hidden states to align with label-token unembeddings so these tokens can be decoded. This also explains why IH outputs themselves decode label-related tokens [Todd et al., 2024], as they must align with label unembedding directions to perform such steering. To validate this conclusion on our classification datasets (not considered in the original study), we extract task vectors from identified IH outputs and add them to layer-30 zero-shot hidden states of Llama2-70B. The substantial accuracy gains in Figure 9 support this conclusion.

6 Conclusions and Limitations

Conclusions This work unifies two major perspectives on in-context learning (ICL)-the roles of special attention heads and task vectors-within a geometric framework centered on separability and alignment of hidden states. We theoretically show that these two properties fully determine ICL classification accuracy and empirically demonstrate that ICL across classification and generation tasks improves accuracy over zero-shot mainly by enhancing alignment, explaining why ICL query hidden states serve effectively as task vectors. Our analysis reveals a phase transition: early layers boost separability, while later layers refine alignment with label unembedding vectors by filtering out label-irrelevant semantics and preserving task-relevant directions, coinciding with the emergence of label semantics in hidden states. Zooming into layers, we identify complementary roles of PTHs and IHs : PTHs in early layers enhance separability, while IHs emerging post-transition promote alignment-clarifying both the phase transition and the effectiveness of IH outputs as task vectors.

Limitations Our analysis of the ICL hidden states of generation tasks is preliminary. In addition, our analysis focuses on inference-time behavior in pretrained models. Future work should investigate how the phase transition pattern of ICL hidden states emerges during training, which could shed light on the emergence of ICL capability.

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

B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, et al. Language models are few-shot learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems , page 1877-1901, Online and Vancouver, Canada, 2020. URL https://dl.acm.org/doi/abs/ 10.5555/3495724.3495883 .

Chan, A. Santoro, A. Lampinen, J. Wang, A. Singh, P. Richemond, J. McClelland, and F. Hill. Data distributional properties drive emergent in-context learning in transformers. In Proceedings of the 36th Advances in Neural Information Processing Systems , pages 18878-18891, New Orleans, USA, 2022. URL https://arxiv.org/abs/2205.05055 .

Cho, M. Kato, Y. Sakai, and N. Inoue. Revisiting in-context learning inference circuit in large language models. In The Thirteenth International Conference on Learning Representations , 2025a. URL https://openreview.net/forum?id=xizpnYNvQq .

Cho, Y. Sakai, M. Kato, K. Tanaka, A. Ishii, and N. Inoue. Token-based decision criteria are suboptimal in in-context learning. In L. Chiruzzo, A. Ritter, and L. Wang, editors, Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers) , pages 5378-5401, Albuquerque, New Mexico, Apr. 2025b. Association for Computational Linguistics. ISBN 979-889176-189-6. URL https://aclanthology.org/2025.naacl-long.278/ .

Dagan, O. Glickman, and B. Magnini. The pascal recognising textual entailment challenge. In Machine learning challenges workshop , pages 177-190. Springer, 2005. URL https://link. springer.com/chapter/10.1007/11736790_9 .

M.-C. De Marneffe, M. Simons, and J. Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung , volume 23, pages 107-124, 2019. URL https://ojs.ub.uni-konstanz.de/sub/index.php/sub/article/ view/601 .

Dong, L. Li, D. Dai, C. Zheng, J. Ma, R. Li, H. Xia, J. Xu, Z. Wu, T. Liu, B. Chang, X. Sun, L. Li, and Z. Sui. A survey on in-context learning, 2024. URL https://arxiv.org/abs/2301. 00234 .

Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread , 1(1): 12, 2021. URL https://transformer-circuits.pub/2021/framework/index. html .

Gemini Team, R. Anil, S. Borgeaud, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, K. Millican, D. Silver, M. Johnson, I. Antonoglou, J. Schrittwieser, A. Glaese, J. Chen, E. Pitler, T. Lillicrap, A. Lazaridou, O. Firat, et al. Gemini: A family of highly capable multimodal models, 2024. URL https://arxiv.org/abs/2312.11805 .

Grattafiori, A. Dubey, A. Jauhri, and A. P. et al. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407.21783 .

Gurnee and M. Tegmark. Language models represent space and time, 2024. URL https: //arxiv.org/abs/2310.02207 .

Halawi, J.-S. Denain, and J. Steinhardt. Overthinking the truth: Understanding how language models process false demonstrations, 2024. URL https://arxiv.org/abs/2307.09476 .

Han, J. Song, J. Gore, and P. Agrawal. Emergence and effectiveness of task vectors in in-context learning: An encoder decoder perspective, 2025. URL https://arxiv.org/abs/2412. 12276 .

Hendel, M. Geva, and A. Globerson. In-context learning creates task vectors. In H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023 , pages 9318-9333, Singapore, Dec. 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.findings-emnlp.624. URL https://aclanthology.org/2023. findings-emnlp.624/ .

Jiang, C. Jiang, Z. Li, S. Xue, J. Zhou, L. Song, D. Lian, and Y. Wei. Unlocking the power of function vectors for characterizing and mitigating catastrophic forgetting in continual instruction tuning, 2025a. URL https://arxiv.org/abs/2502.11019 .

Jiang, Y. Dong, J. Zhou, and Z. Zhu. From compression to expansion: A layerwise analysis of in-context learning, 2025b. URL https://arxiv.org/abs/2505.17322 .

Kahardipraja, R. Achtibat, T. Wiegand, W. Samek, and S. Lapuschkin. The atlas of in-context learning: How attention heads shape in-context retrieval augmentation. arXiv preprint arXiv:2505.15807 , 2025.

Kirsanov, C.-N. Chou, K. Cho, and S. Chung. The geometry of prompting: Unveiling distinct mechanisms of task adaptation in language models. In L. Chiruzzo, A. Ritter, and L. Wang, editors, Findings of the Association for Computational Linguistics: NAACL 2025 , pages 1855-1888, Albuquerque, New Mexico, Apr. 2025. Association for Computational Linguistics. ISBN 979-889176-195-7. URL https://aclanthology.org/2025.findings-naacl.100/ .

Lee, X. Bai, I. Pres, M. Wattenberg, J. K. Kummerfeld, and R. Mihalcea. A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity, 2024. URL https: //arxiv.org/abs/2401.01967 .

Li, Z. Liu, X. Hu, Z. Sun, B. Hu, and M. Zhang. In-context learning state vector with inner and momentum optimization, 2024. URL https://arxiv.org/abs/2404.11225 .

Li and D. Roth. Learning question classifiers. In COLING 2002: The 19th International Conference on Computational Linguistics , 2002. URL https://www.aclweb.org/anthology/ C02-1150 .

Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context examples for GPT-3? In E. Agirre, M. Apidianaki, and I. Vuli´ c, editors, Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures , pages 100-114, Dublin, Ireland and Online, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.deelio-1.10. URL https://aclanthology.org/2022.deelio-1.10/ .

Liu, H. Ye, L. Xing, and J. Zou. In-context vectors: Making in context learning more effective and controllable through latent space steering, 2024. URL https://arxiv.org/abs/2311. 06668 .

Liu, E. J. Michaud, and M. Tegmark. Omnigrok: Grokking beyond algorithmic data, 2023. URL https://arxiv.org/abs/2210.01117 .

MacCartney and C. D. Manning. Modeling semantic containment and exclusion in natural language inference. In D. Scott and H. Uszkoreit, editors, Proceedings of the 22nd International Conference on Computational Linguistics (Coling 2008) , pages 521-528, Manchester, UK, Aug. 2008. Coling 2008 Organizing Committee. URL https://aclanthology.org/C08-1066/ .

Marks and M. Tegmark. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets, 2024. URL https://arxiv.org/abs/2310.06824 .

Merullo, C. Eickhoff, and E. Pavlick. Language models implement simple word2vec-style vector arithmetic, 2024. URL https://arxiv.org/abs/2305.16130 .

Mikolov, K. Chen, G. Corrado, and J. Dean. Efficient estimation of word representations in vector space, 2013. URL https://arxiv.org/abs/1301.3781 .

nostalgebraist. Interpreting gpt: the logit lens, 2020. URL https://www.lesswrong.com/ posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens .

Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. In-context learning and induction heads, 2022. URL https://arxiv.org/abs/ 2209.11895 .

OpenAI, J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, R. Avila, I. Babuschkin, S. Balaji, V. Balcom, P. Baltescu, H. Bao, M. Bavarian, J. Belgum, I. Bello, et al. Gpt-4 technical report, 2024. URL https://arxiv.org/abs/2303.08774 .

F. Park, A. Lee, E. S. Lubana, Y . Yang, M. Okawa, K. Nishi, M. Wattenberg, and H. Tanaka. Iclr: Incontext learning of representations, 2024. URL https://arxiv.org/abs/2501.00070 .

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011. URL https://www.jmlr.org/papers/volume12/ pedregosa11a/pedregosa11a.pdf .

Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. OpenAI blog , 2019. URL https://d4mucfpksywv.cloudfront. net/better-language-models/language-models.pdf .

Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference on Learning Representations , 2024. URL https: //arxiv.org/abs/2312.03002 .

Sharma, J. T. Ash, and D. Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. arXiv preprint arXiv:2312.13558 , 2023. URL https: //arxiv.org/abs/2312.13558 .

Sia, D. Mueller, and K. Duh. Where does in-context learning \\ happen in large language models? In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. URL https://openreview.net/forum?id=LLuSjg59an .

K. Singh, T. Moskovitz, F. Hill, S. C. Chan, and A. M. Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. In Fortyfirst International Conference on Machine Learning , 2024. URL https://arxiv.org/abs/ 2404.07129 .

Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Ng, and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In D. Yarowsky, T. Baldwin, A. Korhonen, K. Livescu, and S. Bethard, editors, Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing , pages 1631-1642, Seattle, Washington, USA, Oct. 2013. Association for Computational Linguistics. URL https://aclanthology.org/ D13-1170/ .

Song, Z. Xu, and Y. Zhong. Out-of-distribution generalization via composition: a lens through induction heads in transformers. Proceedings of the National Academy of Sciences , 122(6): e2417182122, 2025. URL https://arxiv.org/abs/2408.09503 .

Todd, M. L. Li, A. S. Sharma, A. Mueller, B. C. Wallace, and D. Bau. Function vectors in large language models, 2024. URL https://arxiv.org/abs/2310.15213 .

Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023. URL https: //arxiv.org/abs/2302.13971 .

Wang and C. Manning. Baselines and bigrams: Simple, good sentiment and topic classification. In H. Li, C.-Y. Lin, M. Osborne, G. G. Lee, and J. C. Park, editors, Proceedings of the 50th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers) , pages 90-94, Jeju Island, Korea, July 2012. Association for Computational Linguistics. URL https://aclanthology.org/P12-2018/ .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main claims made in the paper accurately summarizes the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of our work in Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: For our theoretical claims put forward in Section 3, we provide the complete proofs in Appendix A. In Appendix B we consider a generalization of the theory we proposed, the proof of which is also in Appendix B.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We present the details of the implementations we used for the models, datasets, as well our procedures of conducting in-context learning in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have submitted the data and code as the supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experiment settings and presented in Section 5 Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Most of the experiment results presented in this work are obtained using fully deterministic procedures with no randomness involved. The only randomness is in the random selection of demonstrations in ICL, but its impact on the results is minimal, so we do not construct error bars over random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide such information in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We conduct the research according to the NeurIPS Code of Ethics in every respect.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the potential societal impacts of our findings in Appendix K.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks in any circumstances.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We provide full reference to the assets we used, primarily in Section 5

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide extensive details to the dataset we generated in Appendix I and Appendix D.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve human subjects and no IRB approvals are required.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [Yes]

Justification: LLM is used to generate one of the datasets used in the experiments. We introduce this dataset in detail in Appendix D and Appendix I. Apart from that, LLM is not related to any aspects of the core methodology in this paper.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/ LLM ) for what should or should not be described.

Appendices

A Proofs related to Theorem 1

We first show that S ∗ is attained at some point v ∗ in the hypersphere S d -1 = { u : u ∈ R d , ∥ u ∥ 2 = 1 } though S d -1 is an infinite set and S ( u ) is not continuous. Observe that each hi in H partitions S d -1 into two hyperhemispheres. One is { u : u ∈ R d , ∥ u ∥ 2 = 1 , u T hi ≥ 0 } and the other is { u : u ∈ R d , ∥ u ∥ 2 = 1 , u T hi < 0 } . Hence, as long as the number of queries in the dataset n is finite, h 1 , ..., hn will partition S d -1 into a finite number of 'cells'. Each cell corresponds to the elements in S d -1 that satisfy a certain sign sequence [ sgn ( u T h 1 ) , ..., sgn ( u T hn ) ] ∈ { 1 , 0 , -1 } n . Hence, S ( u ) is constant on each of the cells, and the supremum (or the maximum) value S ∗ is attained on one of the cells.

We then show the proof of Theorem 1:

Proof First, consider the case where V only consists of yA and yB . Then E ∈ R 2 × d and the two rows are EyA and EyB . Then we have

Here, when EyA hi = EyB hi , without loss of generality, we break the tie by letting the predicted label to be yA for xi .

When there are more tokens in V , it holds true that ( yi = argmax v j Ev j hi ∑ v ′ j ∈{ y A , y B } E v ′ j hi ) ≤ ( yi argmax v j Ev j hi ∑ v j ∈ 1 ,..., | V | E v ′ j hi ) . Hence, we have Acc = 1 n ∑ n i = 1 ✶ ( yi = argmax v j Ev j hi ∑ v j ∈ 1 ,..., | V | E v ′ j hi ) ≤ S ∗ . By the definition of u ∗ , we have

✶ ✶ =

If EyA -EyB | EyA -EyB | 2 = cu ∗ for some positive constant c , then

We already showed that

When max v ∈ V , v / ∈ y A , yB Evhi < max ( EyA hi , EyB hi ) , ∀ i , we have

Hence, we have Acc = S ∗ when max v ∈ V , v / ∈ y A , yB Evhi < max ( EyA hi , EyB hi ) , ∀ i and EyA -EyB ∥ EyA -EyB ∥ 2 = cu ∗ . □

B Generalization of Theorem 1

In this section 5 , we consider a generalization of Theorem 1 where x 1 , ..., xn are allowed to have m < n labels, i.e. yi ∈ { y 1 , ..., ym } , ∀ i . Same as before, let N j = { i : yi = y j } , j = 1 , .., m . N j , j = 1 , ..., m are pairwise disjoint and ⋃ m j = 1 | N j | = { 1 , ..., n } . Let Sj , k ( u ) = 1 | N j | + | N k | ( ∑ i ∈ N j ✶ ( u T hi ≥ 0 ) + ∑ i ∈ N k ✶ ( u T hi < 0 )) and let u ∗ j , k = argmax u ∈ S d -1 Sj , k ( u ) for j , k ∈ 1 , ..., m , and Sj , k ( u ∗ j , k ) = S ∗ j , k . For each j and k , u ∗ j , k is guaranteed to be in S d -1 because only | N j | + | N k | < n hyperplanes are involved in separating separating S d -1 , resulting in a smaller number of cells (guaranteed to be finite) compared to the case of x 1 , ..., xn having only y 1 , y 2 as labels. Then, we have the following result.

Theorem 2 Acc ≤ 1 n ( m -1 ) ( ∑ m k = 1 ∑ m j = k + 1 ( | N k | + | N j | ) S ∗ j , k ) . The equality is achieved when for all unordered pairs ( j , k ) , j , k ∈ 1 , ..., m, max v ∈ V , v / ∈{ y j , y k } Evhi < max ( Eyj hi , Eyk hi ) , ∀ i ∈ Nj ∪ Nk and Eyj -Eyk = c j , ku ∗ j , k for some positive constant c j , k .

The theorem states that in the case of the queries having more than two labels, the dataset classification accuracy is controlled by the pairwise separability and alignment properties of the hidden states. After calculating the separability of the hidden states of queries with labels y j and yk for each ( j , k ) , the dataset classification accuracy is then upper-bounded by these separability values weighted by the number of queries with different labels.

Similarly, the accuracy can match this upper bound if each pairwise hidden states collection is maximally separated along a direction parallel to the difference of the unembedding vectors of the corresponding labels. This theorem shows that the otherwise difficult multiclass classification problem is geometrically equivalent to the combination of a series of binary classification subproblems. And the geometric properties of the multiclass classification problem can be understood by studying the much more tractable geometry of the constituent binary subproblems.

Proof First, note the following relations

̸

̸

Note that the sum is over all ordered pairs ( j , k ) s.t. j , k ∈ 1 , ..., m , j = k . Consider the case of ( k = 1 , j = 2 ) and ( k = 2 , j = 1 ) , then both ✶ ( Eyk hi ≥ Eyj hi ) and ✶ ( Eyj hi ≥ Eyk hi )) appear in the sum. Hence, the sum can be taken over all unordered pairs ( j , k ) and rewritten as:

indicator function is either 1 or 0

̸

Hence, we have

̸

̸

̸

As a proof of concept, plug-in m = 2, then the bound becomes 1 n ( | N 1 | + | N 2 | ) S ∗ = S ∗ since | N 1 | + | N 2 | = n , which matches the result in Theorem 1.

To establish the equality conditions, as in the proof of Theorem 1, if Eyj -Eyk = c j , ku ∗ j , k for some positive constant c j , k , then

We also know that if max v ∈ V , v / ∈ y j , y k Evhi < max ( Eyj hi , Eyk hi ) , ∀ i ∈ Nj ∪ Nk , then

Finally, note that

because for each N j , the term ✶ ( yi = argmax v j Ev j hi ∑ v j ∈ 1 ,..., | V | E v ′ j hi ) appears ( m -1 ) times in the summation for each i ∈ N j . Hence, combining everything, we have

The proof is thus complete.

C Calculation of geometric measures

Separability score Wetrain the logistic classifier on half of the points in H and take its prediction accuracy on the other half as the estimated maximum separability of H . We use the scikit-learn [Pedregosa et al., 2011] default implementation of the logistic classifier, and set the number of iterations to be 100.

Output alignment We feed H into the output unembedding layer of the model to get the model's predictions of the query labels based on H and calculate the percentage of correct predictions.

Singular alignment First obtain the centered version of H , denoted as ¯ H . Then calculate its singular value decomposition, namely ¯ H = U Σ V T , and compute the absolute values of cosine similarities between the top-r singular vectors and the difference of label unembedding vectors EyA -EyB . The absolute value is taken since the sign of the singular vectors is arbitrary. In practice, we calculate the values using the top-2 singular vectors and report the maximum of the two.

Variance-based alignment First normalize EyA -EyB as Ey A -Ey B ∥ Ey A -Ey B ∥ 2 . Then the ratio of variance (w.r.t. its mean) of H explained along Ey A -Ey B ∥ Ey A -Ey B ∥ 2 is Ey A -Ey B ∥ Ey A -Ey B ∥ 2 T ¯ H T ¯ H n Ey A -Ey B ∥ Ey A -Ey B ∥ 2 , which we take to be the variance-based alignment of H .

Mean-based alignment First compute the means of label clusters (subsets of H with each of the two different labels), i.e. 1 | N A | ∑ i ∈ N A hi and 1 | N B | ∑ i ∈ N B hi . Then take the difference and compute

the value of its projection onto Ey A -Ey B ∥ Ey A -Ey B ∥ 2 , i.e. ( Ey A -Ey B ∥ Ey A -Ey B ∥ 2 ) T ( 1 N A ∑ i ∈| N A | hi -1 | N B | ∑ i ∈ N B hi ) . Then, calculate the weighted average of the variances of the two label clusters along Ey A -Ey B ∥ Ey A -Ey B ∥ 2 , i.e. | N A | Var ( { h T i Ey A -Ey B ∥ Ey A -Ey B ∥ 2 : i ∈ N A } )+ | N B | Var ( { h T i Ey A -Ey B ∥ Ey A -Ey B ∥ 2 : i ∈ N B } ) | N A | + | N B | . Finally, divide the projected difference between means by the square root of the weighted average of the projected variance.

Composite alignment Multiply the variance-based alignment by the mean-based alignment.

Effective dimension Given the SVD ¯ H = U Σ V T and the singular values σ 1 , ..., σ n , the effective dimension of H is ( ∑ n i = 1 σ 2 i ) 2 ∑ n i = 1 σ 4 i .

□

D Implementation details

Models We use the official huggingface implementation of all models. All models with more than 10B parameters are quantized to 4bit.

Datasets We use the official huggingface implementation of all datasets except for the Famous People dataset crafted by ourselves. The Famous People dataset is generated using ChatGPT-4o [OpenAI et al., 2024]. The dataset has 180 datapoints, each is the name of a famous person with a profession label from one of six categories: actor, politician, singer, scientist, writer, athlete. The dataset is balanced as there are 30 data points for each label.

ICL setting For each dataset except for the Famous People dataset, we select demonstrations from the train set and queries from the test set, or the validation set if the ground-truth labels for the test set are not provided. We keep only the first 10000 datapoints in the train set for demonstration selection if the train set has more than 10000 entries, and the first 1000 data points in the test set/validation set for accuracy evaluation. For the famous people dataset, we use the first 90 data points for demonstration selection and the remaining 90 data points for testing. For experiments involving kNN demonstration selection, we first use the respective LLMs to encode the demonstrations and queries, and select the demonstrations with embeddings closer to that of the query in l 2 distance for each query. We use the prompt formats detailed in Table 2.

Devices All experiments with Llama2-7B, Llama2-13B, Gemma-2B, and Gemma-7B are conducted with an A800 GPU. All experiments with Llama3-8B, Llama2-70B, and Llama3-70B are conducted with an H200 GPU.

For the label mapping experiment in Subsection 5.1. We use the mappings listed in Table 3 to map the ground-truth labels of each dataset to symbols.

E Practical implications of the mechanistic findings

We discuss in this section how our geometric characterization of the model's inference process in the zero-shot or ICL setting can be converted to actionable and efficient methods to improve model performance.

E.1 Improving classification by capitalizing on hidden states separability

In Figure 2 we see the high separability of hidden states of prompts from classification tasks which surpasses the accuracy of the actual label prediction by far, irrespective of whether demonstrations are

prepended to the labels. This implies that rather than direct decoding we can instead use the model as a feature extractor with which to train classifiers to perform the classification task, which is exactly how we evaluate the separability score. To demonstrate this, we train a logistic classifier on 50% of the last layer hidden states and evaluate the classification accuracy on the rest half. The results in Table 22 demonstrate that such a practice improves the classification accuracy with zero-shot prompts by over 70% across models, while only incurring the minimal cost of training an extra logistic classifier. Moreover, in Table 23 we see that training the classifier over 8-shot ICL hidden states also brings a solid ∼ 10% increase over the original ICL accuracy, which shows the generalizability of this lightweight method that leverages hidden states separability.

E.2 Unsupervised enhancement of classification performance through low-rank denoising

In Figure 6 we see that low-rank denoising of the ICL hidden states matrix reveals the task-related directions as its dominant directions and enhances its alignment properties. We can thus apply this technique to the final layer hidden states instead of the intermediate layer ones (which is the focus of Figure 6) to directly improve classification performance. In Table 24 we apply rank-5 denoising to final layer hidden states before feeding them into the unembedding layer. This method is different from the classifier-based method in that it assumes no knowledge of the ground-truth label of each hidden state. The dataset-average results demonstrate that the rank-5 denoising can indeed improve the classification performance of all models with the only exception of Gemma-7B.

F Supplementary materials for Subsection 5.1

F.1 Replication of Figure 2 for other models

In Figure 10-15, we provide the visualizations of the experiments presented in Figure 2 for other models, from which similar conclusions to those in Subsection 5.1 can be drawn. Alignment measures explain the accuracy difference between ICL and zero-shot across models, and the layer-wise trends of the measures explain the phase transition pattern for all models.

F.2 Replication of Figure 3 for other models

In Figure 16-21 we provide the visualizations of the experiments presented in Figure 3 for other models to establish the generality of our conclusions. It is clear that the phase transition pattern is evident in different ICL settings across models, and alignment measures capture accuracy differences across ICL settings for all models.

F.3 Replication of Figure 4 for other models

In Figure 53-58 we repeat the visualizations in Figure 4 for other models to examine where the generation task fits well into our framework. The results demonstrate that for other models, the phase-transition pattern in the composite alignment metric is specific to the 8-shot ICL hidden states, while ICL and zero-shot hidden states do not differ much in terms of the dynamics of the separability score.

G Supplementary materials for Subsection 5.2

G.1 Replication of subplot (A) of Figure 5 for other datasets and models

In Figure 22- 28 we provide the visualizations of the experiments presented in subplot (A) of Figure 5, to establish the generalizability of our conclusion that the directional alignment and the encoding of label-related semantics into the hidden states happen almost simultaneously after the phase transition. The layer-wise trends of the directional alignment measures and output alignment exhibit high correlation across all datasets and models.

G.2 Replication of subplot (B) of Figure 5 for other datasets

In Table 4-8, we present the tokens decoded from top singular directions near the phase transition breakpoint for other datasets. The results demonstrate that the sudden .semantic shift toward labelrelated tokens in the hidden states (as demonstrated by subplot (B) of Figure 5) is present in other datasets as well. Note that the phase transition happens at different layers for different datasets.

G.3 Replication of subplot (B) of Figure 5 with replaced demonstration labels

To further examine the relationship between alignment and the encoding of task-related semantics, we design a scenario in which the ground-truth demonstration labels are consistently mapped to symbols, and the expected predictions for the query label becomes the mapped symbols as well. This forces a separation between semantic encoding and alignment, as the task-related semantics remain unchanged, but alignment must now occur with respect to the unembedding vectors of the symbols. Specifically, we conduct experiments on SST-2 and SNLI where we map the SST-2 labels (negative/positive) to @ and #, and the SNLI labels (true/maybe/false) to @, #, and ! within the demonstrations. We then inspect the tokens decoded from the top singular directions of the hidden states at each layer.

As shown in Table 9, the results reveal an interesting pattern. The first major shift in the layer-wise semantic content occurs at layer 24, where the ground-truth label tokens suddenly become dominant. This closely mirrors the pattern observed with normal demonstration labels in Table 6, and aligns with the point where output alignment begins to rise in Figure 22 (C). However, another drastic shift is observed at layer 45, where the decoded top tokens begin to include the symbol tokens. These symbols suddenly appear and start to compete with semantically relevant task tokens, marking the onset of alignment between the hidden states and the unembedding vectors of the symbols. This delayed alignment corresponds with the pattern seen in subplot (C) of Figure 3, where the phase transition under demonstrations with symbolized labels is postponed to later layers.

The competition between the true labels and the injected symbol labels persists through to the final layers. This likely contributes to the reduced accuracy in this setting, as the semantic influence of the substituted labels fails to fully override that of the ground-truth labels. Similar results for SST-2 are presented in Table 10, showing similar trends.

G.4 Replication of subplot (C) of Figure 5 for other datasets

Subplot (C) of Figure 5 reveals a two-stage pattern in the layer-wise retainment and filtering of semantic information, where most post-transition layers retain label-related semantics and filter out irrelevant ones, and the final layers do the opposite. We show in Table 11-15 that it exists in other datasets besides SST-2 as well.

G.5 Replication of Figure 6 for other datasets and models

We first provide the results of rank-10 approximations on all 6 datasets for Llama2-70B in Figure 42. The conclusion is that low-rank denoising significantly increase the output alignment significantly acorss all datasets, particularly in the innitial layers after the phase transition. This confirms the generality of our findings in Subsection 5.2. Additionally, in the final layers, the output alignment also increases on on SST-2, SUBJ, TREC, and SNLI, though not on RTE and CB.

In Figure 43-48, we present results for other models. These are consistent with the findings from Llama2-70B. Specifically, the pattern that low-rank denoising improves output alignment more in text classification datasets than in natural language inference datasets holds across models. This suggests a higher level of linguistic complexity in natural language inference tasks-likely because such tasks require reasoning over both a premise and a hypothesis, whose combined semantics are more difficult to compress and more prone to information loss during low-rank projection.

G.6 Ablation experiments for the rank hyperparameter in low-rank approximations

In this section, we explore the effect of varying the number of retained ranks in rankr SVD approximations of Llama2-70B hidden states across all datasets. Specifically, we investigate how increasingly

aggressive rank reductions-setting r = 1, 2, and 5-impact performance. The results, shown in Figure 49, Figure 50, and Figure 51, yield several notable findings.

For binary classification tasks such as SST-2, SUBJ, and RTE, more aggressive rank reduction sometimes leads to better output alignment, particularly in layers after the phase transition where alignment with label unembedding vectors begins to emerge. In some cases, even a rank-1 approximation preserves or enhances alignment-related performance. However, for three-way classification tasks like SNLI and RTE, aggressive reductions (especially at r = 1 or 2) can cause accuracy drops in later layers. For TREC, which involves six label classes, rank-2 or rank-1 approximations degrade accuracy across all layers.

These results align with the intuition that datasets with more labels require higher-dimensional hidden state representations to preserve the necessary separability and alignment structures among label clusters.

H Experiment details concerning the identification of PTHs and IHs

For each dataset, we use the first 50 queries to identify the set of IH s and PTH s. Denote the queries as x 1 , ..., x 50 each with token length s ( x 1 ) , ..., s ( x 50 ) .

Identification of PTHs For each xi , the LLM will generate an attention tensor Attni ∈ R L × Nh × s ( xi ) × s ( xi ) , where L is the total number of layers, Nh is the number of attention heads per layer. Hence, for the n th h attention head at layer l , there is a corresponding attention matrix Attn ( l , nh ) i ∈ R s ( xi ) × s ( xi ) , where Attn ( l , nh ) i , j , k is the attention of this attention head at the k th token to the j th token in xi . The PTH score of this attention head on xi is thus defined as ∑ s ( xi ) k = 2 Attn ( l , nh ) i , k -1 , k , which measures the total attention the attention head assigns at each token position to the previous token. Consequently, the PTH score of the n th h attention head at layer l on all the queries is ∑ 50 i = 1 ∑ s ( xi ) k = 2 Attn ( l , nh ) i , k -1 , k . We calculate the PTH scores for all ( l , nh ) pairs and choose the top 10% attention heads as the identified PTH s.

Identification of IHs For each xi , randomly select 8 demonstrations and prepare them to xi . The resultant ICL prompt (denoted as Xi with length s ( Xi ) ), with a slight abuse of notation, takes the format of ⟨ xi , 1 ⟩ : ⟨ yi , 1 ⟩ , ..., ⟨ xi , 8 ⟩ : ⟨ yi , 8 ⟩ , ⟨ xi ⟩ :, where ⟨ xi , k ⟩ represents the sentence part of demonstration k ('I like this movie. Sentiment') and ⟨ yi , k ⟩ the label ('positive') separated from the sentence part by a colon (:), and ⟨ xi ⟩ is the sentence part for the query. At the position of the final ':', an IH will attend to tokens after the previous ':'s, i.e. the label tokens ⟨ yi , 1 ⟩ , ..., ⟨ yi , 8 ⟩ . Let I i denote the indices of the label tokens in Xi . The IH score of the n th h attention head at layer l over the 50 queries is thus ∑ 50 i = 1 ∑ k ∈ I i ∈ Attn ( l , nh ) i , k , s ( Xi ) , i.e. we are measuring the total attention a 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 ( l , nh ) pairs and choose the top 10% attention heads as the identified IH s.

I Supplementary materials for Subsection 5.3

I.1 Replication of Figure 7 for other models

In Figure 29-34, we present the results of attention heads ablations across models. The results support our conclusions regarding the significance of PTHs in enhancing the separability of the ICL hidden states and the significance of IHs in aligning them with the label unembedding vectors.

I.2 Replication of Figure 8 for other models

In Figure 59-64, we present the results of analyzing the geometric contribution of PTHs and IHs in other models to the evolution of hidden states by injecting their outputs as task vectors. The results clearly demonstrate their respective significance in fostering the separability and alignment of hidden states. In particular, the PTH task vectors of other models exhibit greater potential to increase hidden-state separability than the result shown in Figure 8.

I.3 Replication of Table 1 for other models

In Table 16-21, we compare the layer distributions of the top 1%, 2%, 5%, and 10% identified IHs and PTHs in other models and conduct Mann-Whitney U tests to assess the significance of the differences. The results demonstrate a consistent difference in the layer positions of PTHs and IHs across models, with PTHs significantly preceding IHs on average-the difference being most prominent in the two Gemma-family models, Gemma-2B and Gemma-7B.

I.4 Replication of Figure 9 for other models

In Figure 65-70, we show the results of injecting task vectors derived from IH outputs for other models. The task vectors are similarly injected at layer 3 8 L , where L is the total number of layers of the respective models. The results demonstrate that IH outputs qualify as effective task vectors across models, with the exception of Gemma-2B, as they considerably improve prediction accuracy in the zero-shot setting after being injected into the zero-shot hidden states.

I.5 Results of attention heads ablation in the ICL setting where query label is not in context

The results for other models, reported in Figure 29-34, confirm the generality of our conclusions across model sizes and architectures. To test whether these findings hold under different ICL settings as well, we replicate the attention head ablation experiments in a special ICL scenario ('In-weight Learning') studied by Reddy [2024] and Chan et al. [2022], where none of the in-context demonstrations share the query's label. This setting is useful for analyzing attention heads-especially IHs , which are known for copy-like behavior-when direct label copying is not possible.

To this end, we construct a Famous People dataset following the procedures in Appendix D. Each data point is a sentence in the format '{Famous Person Name} is a:", and the label is the profession of the famous person. Once combined, they form an input sequence like, e.g., 'Taylor Swift is a: Singer." that can be used as either a query or a demonstration. For each query, we sample k = 4 demonstrations, ensuring that none share the query's profession label. We then ablate PTHs and IHs to examine their influence on model performance. As a control, we also evaluate the measures in the setting where at least one demonstration does share the query's label, allowing us to assess how the availability of directly copyable labels affects the impact of head ablation.

The results in Figure 35 resemble those in standard ICL settings in that ablating PTHs significantly impairs separability, while ablating IHs affects alignment. However, there are notable differences. Most importantly, the accuracy drop from ablating PTHs is greater than that from ablating IHs-opposite to the standard case, where IH ablation has a larger impact. This highlights the critical role of having the query label present in the context for IHs to function effectively. The importance of label presence is further confirmed by the substantial accuracy gain observed when the query label is forced into the context.

We provide results analogous to Figure 35 for other models in Figure 36-41, which demonstrate similar results.

I.6 Results of attention heads ablation in the zero-shot setting

In this section, we investigate the effect of attention head ablation on the separability and alignment measures of zero-shot hidden states. The results averaged across datasets, as provided in Figure 52, reveal that the zero-shot setting is fundamentally different from ICL. In the zero-shot setting, ablating IHs has practically no effect on either separability or alignment measures, similar to ablating random heads, because there are no demonstrations from which to copy. On the contrary, ablating PTHs substantially compromises all separability and alignment measures. This illustrates that the contribution of PTHs to the separability of hidden states is by no means restricted to the ICL setting and highlights the significance of PTHs as key structural components of LLMs.

J Experiment details of using IH outputs to construct task vectors

For each dataset, we first construct 8-shot ICL prompts using the first 50 queries. The demonstrations used are exactly the same as the ones used to evaluate the 8-shot ICL accuracy for each dataset. Then,

following the procedures described in Todd et al. [2024], we obtain the average output (across the 50 prompts) of each identified top 10% IHs and PTHs at the final token positions, and sum the average outputs of all IHs to obtain the task vector.

In terms of the steering experiment, we add the task vector to the hidden state of the final token of each query at Layer-30 of Llama2-70B, and let the modified hidden state flow through subsequent layers. While calculating the accuracy, we exclude the first 50 queries and only evaluate accuracy on the remaining queries in the test set to ensure an independence between the task vector and the queries involved.

K Broader impacts

Understanding the internal mechanisms of large language models is increasingly critical as these models are deployed in high-stakes applications ranging from education and healthcare to legal and governmental decision-making. This work contributes to that understanding by providing a geometric framework that explains in-context learning through the lens of hidden state dynamics, specifically separability and alignment. By identifying the role of specific architectural components-such as Previous Token Heads and Induction Heads-in shaping model behavior, our study brings interpretability to a domain often criticized for opacity.

The primary positive societal impact of this research lies in its potential to improve the transparency, controllability, and safety of LLMs. Better mechanistic understanding can inform the development of more robust models that are less reliant on spurious correlations and more capable of structured generalization. For example, insights into how semantic information is injected and refined layer by layer may support diagnostic tools that detect when a model is failing to align its internal representations with task-relevant signals. Additionally, the SVD-based techniques we propose for identifying and amplifying task-relevant components could be used to improve the efficiency or reliability of model outputs, particularly in low-resource or privacy-sensitive settings where retraining is infeasible.

However, a deeper understanding of model internals also introduces risks. The ability to precisely manipulate hidden states or attention mechanisms may enable adversarial behaviors, such as constructing prompts that selectively suppress or amplify certain outputs for political or financial gain. Moreover, techniques for filtering and steering semantic content could be misused to hide bias or simulate alignment without genuine safety improvements. As such, we encourage future work to pair interpretability research with rigorous safety and ethics evaluations, and to involve interdisciplinary expertise when applying these findings to real-world systems.

Overall, we view our contributions as a step toward more interpretable and accountable AI, but we emphasize that interpretability alone is not a guarantee of ethical deployment. Responsible application of these insights requires careful consideration of both technical and societal factors.

Llama3-8B: 8-shot Acc.: 76.20%; 0-shot Acc.: 0.30%

0

10

20

30

Layer Number

0

20

40

Effective dimension

0

10

20

30

Layer Number

0.6

0.8

Separability score

0

10

20

30

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

30

Layer Number

0.00

0.01

0.02

Composite alignment

0

10

20

30

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

Layer Number

0.000

0.005

0.010

Var.-based alignment

8-shot

0-shot

Llama3-70B: 8-shot Acc.: 81.09%; 0-shot Acc.: 0.00%

0

20

40

60

80

Layer Number

0

20

40

Effective dimension

0

20

40

60

80

Layer Number

0.6

0.8

Separability score

0

20

40

60

80

Layer Number

0.05

0.10

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.02

Composite alignment

0

20

40

60

80

Layer Number

0

2

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.005

Var.-based alignment

8-shot

0-shot

0

10

20

30

Layer Number

0

20

40

Effective dimension

0

10

20

30

Layer Number

0.6

0.8

Separability score

0

10

20

30

Layer Number

0.1

0.2

Singular alignment

0

10

20

30

Layer Number

0.00

0.02

Composite alignment

0

10

20

30

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

Layer Number

0.00

0.01

Var.-based alignment

8-shot

0-shot

0

10

20

30

40

Layer Number

0

20

40

Effective dimension

0

10

20

30

40

Layer Number

0.6

0.8

Separability score

0

10

20

30

40

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

30

40

Layer Number

0.00

0.02

0.04

Composite alignment

0

10

20

30

40

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

40

Layer Number

0.00

0.01

Var.-based alignment

8-shot

0-shot

Gemma-2B: 8-shot Acc.: 66.16%; 0-shot Acc.: 0.00%

0

5

10

15

Layer Number

0

20

Effective dimension

0

5

10

15

Layer Number

0.6

0.8

Separability score

0

5

10

15

Layer Number

0.1

0.2

Singular alignment

0

5

10

15

Layer Number

0.00

0.01

0.02

Composite alignment

0

5

10

15

Layer Number

0

1

Mean-based alignment

0

5

10

15

Layer Number

0.00

0.01

Var.-based alignment

8-shot

0-shot

Gemma-7B: 8-shot Acc.: 79.06%; 0-shot Acc.: 0.02%

0

10

20

Layer Number

0

20

Effective dimension

0

10

20

Layer Number

0.6

0.8

Separability score

0

10

20

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

Layer Number

0.000

0.025

0.050

Composite alignment

0

10

20

Layer Number

0

2

Mean-based alignment

0

10

20

Layer Number

0.00

0.01

0.02

Var.-based alignment

8-shot

0-shot

0

10

20

0.0

0.2

0.4

0.6

0.8

0

10

20

30

Layer Number

0

10

20

30

40

Effective dimension

0

10

20

30

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

10

20

30

Layer Number

0.00

0.05

0.10

0.15

0.20

Singular alignment

0

10

20

30

Layer Number

0.00

0.01

0.02

0.03

0.04

0.05

Composite alignment

0

4

8

12

16

20

24

Llama3-8B: Average across datasets

(a) Varying number of demonstrations

0

20

Layer Number

0

20

40

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

Layer Number

0.00

0.02

Composite alignment

0

20

Layer Number

0

2

Mean-based alignment

0

20

Layer Number

0.00

0.01

Var.-based alignment

Llama3-8B: Random Selection Acc.: 76.20%; KNN Selection Acc.: 84.03%

Random Selection

KNN Selection

0

20

Layer Number

0

20

40

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.1

0.2

Singular alignment

0

20

Layer Number

0.00

0.01

0.02

Composite alignment

0

20

Layer Number

0

1

2

Mean-based alignment

0

20

Layer Number

0.000

0.005

0.010

Var.-based alignment

Llama3-8B: Normal Label Acc.: 76.20%; Symbol Label Acc.: 47.05%

Normal Label

Symbol Label

(b) kNN Demonstration selection

(c) Random symbols as demonstration labels

0

10

20

0.0

0.2

0.4

0.6

0.8

0

20

40

60

80

Layer Number

0

10

20

30

40

Effective dimension

0

20

40

60

80

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

20

40

60

80

Layer Number

0.025

0.050

0.075

0.100

0.125

0.150

Singular alignment

0

20

40

60

80

Layer Number

0.000

0.005

0.010

0.015

0.020

0.025

0.030

Composite alignment

0

4

8

12

16

20

24

Llama3-70B: Average across datasets

(a) Varying number of demonstrations

0

50

Layer Number

0

20

40

Effective dimension

0

50

Layer Number

0.6

0.8

Separability score

0

50

Layer Number

0.05

0.10

0.15

Singular alignment

0

50

Layer Number

0.00

0.02

Composite alignment

0

50

Layer Number

0

2

Mean-based alignment

0

50

Layer Number

0.000

0.005

Var.-based alignment

Llama3-70B: Random Selection Acc.: 81.09%; KNN Selection Acc.: 89.44%

Random Selection

KNN Selection

0

50

Layer Number

0

20

40

Effective dimension

0

50

Layer Number

0.6

0.8

Separability score

0

50

Layer Number

0.0

0.1

Singular alignment

0

50

Layer Number

0.00

0.02

Composite alignment

0

50

Layer Number

0

2

Mean-based alignment

0

50

Layer Number

0.000

0.005

0.010

Var.-based alignment

Llama3-70B: Normal Label Acc.: 81.09%; Symbol Label Acc.: 64.67%

Normal Label

Symbol Label

(b) KNN Demonstration selection

(c) Random symbols as demonstration labels

0

10

20

0.0

0.2

0.4

0.6

0.8

0

10

20

30

Layer Number

0

10

20

30

40

50

Effective dimension

0

10

20

30

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

10

20

30

Layer Number

0.00

0.05

0.10

0.15

0.20

0.25

Singular alignment

0

10

20

30

Layer Number

0.00

0.01

0.02

0.03

0.04

0.05

Composite alignment

0

4

8

12

16

20

24

Llama2-7B: Average across datasets

(a) Varying number of demonstrations

0

20

Layer Number

0

20

40

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

Layer Number

0.00

0.05

Composite alignment

0

20

Layer Number

0

2

Mean-based alignment

0

20

Layer Number

0.00

0.01

0.02

Var.-based alignment

Llama2-7B: Random Selection Acc.: 75.13%; KNN Selection Acc.: 84.07%

Random Selection

KNN Selection

0

20

Layer Number

0

20

40

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.1

0.2

Singular alignment

0

20

Layer Number

0.00

0.02

Composite alignment

0

20

Layer Number

0

1

2

Mean-based alignment

0

20

Layer Number

0.00

0.01

Var.-based alignment

Llama2-7B: Normal Label Acc.: 75.13%; Symbol Label Acc.: 44.73%

Normal Label

Symbol Label

(b) KNN Demonstration selection

(c) Random symbols as demonstration labels

0

10

20

0.0

0.2

0.4

0.6

0.8

0

10

20

30

40

Layer Number

0

10

20

30

40

50

Effective dimension

0

10

20

30

40

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

10

20

30

40

Layer Number

0.00

0.05

0.10

0.15

0.20

0.25

Singular alignment

0

10

20

30

40

Layer Number

0.00

0.02

0.04

0.06

Composite alignment

0

4

8

12

16

20

24

Llama2-13B: Average across datasets

(a) Varying number of demonstrations

0

20

40

Layer Number

0

20

40

Effective dimension

0

20

40

Layer Number

0.6

0.8

Separability score

0

20

40

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

40

Layer Number

0.00

0.05

Composite alignment

0

20

40

Layer Number

0

1

2

Mean-based alignment

0

20

40

Layer Number

0.00

0.01

0.02

Var.-based alignment

Llama2-13B: Random Selection Acc.: 74.79%; KNN Selection Acc.: 86.18%

Random Selection

KNN Selection

0

20

40

Layer Number

0

20

40

Effective dimension

0

20

40

Layer Number

0.6

0.8

Separability score

0

20

40

Layer Number

0.0

0.1

0.2

Singular alignment

0

20

40

Layer Number

0.00

0.02

0.04

Composite alignment

0

20

40

Layer Number

0

1

2

Mean-based alignment

0

20

40

Layer Number

0.00

0.01

Var.-based alignment

Llama2-13B: Normal Label Acc.: 74.79%; Symbol Label Acc.: 51.78%

Normal Label

Symbol Label

(b) KNN Demonstration selection

(c) Random symbols as demonstration labels

0

10

20

0.0

0.2

0.4

0.6

0

5

10

15

Layer Number

0

10

20

30

Effective dimension

0

5

10

15

Layer Number

0.5

0.6

0.7

0.8

Separability score

0

5

10

15

Layer Number

0.05

0.10

0.15

0.20

0.25

Singular alignment

0

5

10

15

Layer Number

0.00

0.01

0.02

0.03

0.04

Composite alignment

0

4

8

12

16

20

24

Gemma-2B: Average across datasets

(a) Varying number of demonstrations

0

10

Layer Number

0

20

Effective dimension

0

10

Layer Number

0.6

0.8

Separability score

0

10

Layer Number

0.1

0.2

0.3

Singular alignment

0

10

Layer Number

0.000

0.025

0.050

Composite alignment

0

10

Layer Number

0

1

2

Mean-based alignment

0

10

Layer Number

0.00

0.01

0.02

Var.-based alignment

Gemma-2B: Random Selection Acc.: 66.16%; KNN Selection Acc.: 77.99%

Random Selection

KNN Selection

0

10

Layer Number

0

20

Effective dimension

0

10

Layer Number

0.6

0.8

Separability score

0

10

Layer Number

0.1

0.2

0.3

Singular alignment

0

10

Layer Number

0.00

0.01

0.02

Composite alignment

0

10

Layer Number

0

1

Mean-based alignment

0

10

Layer Number

0.00

0.01

0.02

Var.-based alignment

Gemma-2B: Normal Label Acc.: 66.16%; Symbol Label Acc.: 40.06%

Normal Label

Symbol Label

(b) KNN Demonstration selection

(c) Random symbols as demonstration labels

0

10

20

0.0

0.2

0.4

0.6

0.8

0

10

20

Layer Number

0

10

20

30

40

Effective dimension

0

10

20

Layer Number

0.5

0.6

0.7

0.8

0.9

Separability score

0

10

20

Layer Number

0.00

0.05

0.10

0.15

0.20

0.25

Singular alignment

0

10

20

Layer Number

0.00

0.02

0.04

0.06

0.08

0.10

0.12

Composite alignment

0

4

8

12

16

20

24

Gemma-7B: Average across datasets

(a) Varying number of demonstrations

0

20

Layer Number

0

20

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.0

0.2

Singular alignment

0

20

Layer Number

0.00

0.05

Composite alignment

0

20

Layer Number

0

2

Mean-based alignment

0

20

Layer Number

0.00

0.01

0.02

Var.-based alignment

Gemma-7B: Random Selection Acc.: 79.06%; KNN Selection Acc.: 86.23%

Random Selection

KNN Selection

0

20

Layer Number

0

20

40

Effective dimension

0

20

Layer Number

0.6

0.8

Separability score

0

20

Layer Number

0.0

0.2

Singular alignment

0

20

Layer Number

0.000

0.025

0.050

Composite alignment

0

20

Layer Number

0

2

Mean-based alignment

0

20

Layer Number

0.00

0.01

0.02

Var.-based alignment

Gemma-7B: Normal Label Acc.: 79.06%; Symbol Label Acc.: 42.99%

Normal Label

Symbol Label

(b) KNN Demonstration selection

(c) Random symbols as demonstration labels

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

Singular Alignment

(a) SUBJ

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.000

0.005

0.010

0.015

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.02

0.04

0.06

0.08

0.10

Singular Alignment

(b) TREC

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.01

0.02

0.03

0.04

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.05

0.10

0.15

Singular Alignment

(c) SNLI

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(d) RTE

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

Singular Alignment

(e) CB

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.002

0.004

0.006

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.02

0.04

0.06

0.08

0.10

Singular Alignment

(a) SUBJ

0

10

20

30

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Output Alignment

0.000

0.005

0.010

0.015

Composite Alignment

0

10

20

30

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Output Alignment

0.02

0.04

0.06

0.08

0.10

0.12

Singular Alignment

(b) TREC

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

0.030

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(c) SNLI

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(d) RTE

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(e) CB

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.01

0.02

0.03

0.04

0.05

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(f) SST-2

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.002

0.004

0.006

0.008

0.010

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.02

0.04

0.06

0.08

0.10

0.12

Singular Alignment

(a) SUBJ

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.000

0.002

0.004

0.006

0.008

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.02

0.04

0.06

0.08

Singular Alignment

(b) TREC

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.025

0.050

0.075

0.100

0.125

0.150

Singular Alignment

(c) SNLI

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

Singular Alignment

(d) RTE

0

25

50

75

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

Singular Alignment

(e) CB

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

1.0

Output Alignment

0.00

0.02

0.04

0.06

0.08

0.10

Composite Alignment

0

20

40

60

80

Layer Number

0.0

0.2

0.4

0.6

0.8

1.0

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(f) SST-2

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.05

0.10

0.15

Singular Alignment

(a) SUBJ

0

10

20

30

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.000

0.005

0.010

0.015

0.020

Composite Alignment

0

10

20

30

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.02

0.04

0.06

0.08

0.10

0.12

Singular Alignment

(b) TREC

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.01

0.02

0.03

0.04

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(c) SNLI

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(d) RTE

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.01

0.02

0.03

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(e) CB

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.02

0.04

0.06

0.08

0.10

Composite Alignment

0

10

20

30

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0

0.1

0.2

0.3

Singular Alignment

(f) SST-2

0

20

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.002

0.004

0.006

0.008

0.010

0.012

Composite Alignment

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

Singular Alignment

(a) SUBJ

0

20

40

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

20

40

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.025

0.050

0.075

0.100

0.125

Singular Alignment

(b) TREC

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.01

0.02

0.03

Composite Alignment

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(c) SNLI

0

20

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(d) RTE

0

20

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

20

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.025

0.050

0.075

0.100

0.125

0.150

Singular Alignment

(e) CB

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

Composite Alignment

0

10

20

30

40

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0

0.1

0.2

0.3

0.4

Singular Alignment

(f) SST-2

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.000

0.001

0.002

0.003

0.004

0.005

Composite Alignment

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.05

0.10

0.15

0.20

0.25

Singular Alignment

(a) SUBJ

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.000

0.005

0.010

0.015

0.020

Composite Alignment

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.025

0.050

0.075

0.100

0.125

0.150

0.175

Singular Alignment

(b) TREC

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.000

0.002

0.004

0.006

0.008

0.010

0.012

Composite Alignment

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

Output Alignment

0.0

0.1

0.2

0.3

0.4

Singular Alignment

(c) SNLI

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Output Alignment

0.000

0.001

0.002

0.003

Composite Alignment

0

5

10

15

Layer Number

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Output Alignment

0.00

0.05

0.10

0.15

0.20

Singular Alignment

(d) RTE

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(e) CB

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.02

0.04

0.06

Composite Alignment

0

5

10

15

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0

0.1

0.2

0.3

0.4

Singular Alignment

(f) SST-2

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.01

0.02

0.03

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(a) SUBJ

0

10

20

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.000

0.005

0.010

0.015

0.020

0.025

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.02

0.04

0.06

0.08

0.10

0.12

0.14

Singular Alignment

(b) TREC

0

10

20

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.02

0.04

0.06

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(c) SNLI

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.01

0.02

0.03

0.04

0.05

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0

0.1

0.2

0.3

Singular Alignment

(d) RTE

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.01

0.02

0.03

0.04

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Singular Alignment

(e) CB

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.00

0.05

0.10

0.15

Composite Alignment

0

10

20

Layer Number

0.0

0.2

0.4

0.6

0.8

Output Alignment

0.0

0.1

0.2

0.3

0.4

Singular Alignment

(f) SST-2

Llama3-8B, average across datasets Normal Acc.: 76.20% | w/o PTH Acc.: 31.49% | w/o IH Acc.: 0.00% | Random Heads Ablated Acc.: 47.86%

0

10

20

30

Layer Number

0

100

Effective dimension

0

10

20

30

Layer Number

0.6

0.8

Separability score

0

10

20

30

Layer Number

0.1

0.2

Singular alignment

0

10

20

30

Layer Number

0.00

0.01

0.02

Composite alignment

0

10

20

30

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

Layer Number

0.000

0.005

0.010

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Llama3-70B, average across datasets Normal Acc.: 81.09% | w/o PTH Acc.: 26.89% | w/o IH Acc.: 0.92% | Random Heads Ablated Acc.: 77.89%

0

20

40

60

80

Layer Number

0

100

200

Effective dimension

0

20

40

60

80

Layer Number

0.6

0.8

Separability score

0

20

40

60

80

Layer Number

0.0

0.1

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.02

Composite alignment

0

20

40

60

80

Layer Number

0

2

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.005

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Llama2-7B, average across datasets Normal Acc.: 75.13% | w/o PTH Acc.: 10.59% | w/o IH Acc.: 0.70% | Random Heads Ablated Acc.: 48.15%

0

10

20

30

Layer Number

0

50

100

Effective dimension

0

10

20

30

Layer Number

0.6

0.8

Separability score

0

10

20

30

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

30

Layer Number

0.00

0.02

Composite alignment

0

10

20

30

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

Layer Number

0.00

0.01

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Llama2-13B, average across datasets Normal Acc.: 74.79% | w/o PTH Acc.: 23.19% | w/o IH Acc.: 1.16% | Random Heads Ablated Acc.: 73.49%

0

10

20

30

40

Layer Number

0

25

50

Effective dimension

0

10

20

30

40

Layer Number

0.6

0.8

Separability score

0

10

20

30

40

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

30

40

Layer Number

0.00

0.02

0.04

Composite alignment

0

10

20

30

40

Layer Number

0

1

2

Mean-based alignment

0

10

20

30

40

Layer Number

0.00

0.01

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Gemma-2B, average across datasets Normal Acc.: 66.16% | w/o PTH Acc.: 16.24% | w/o IH Acc.: 1.20% | Random Heads Ablated Acc.: 36.48%

0

5

10

15

Layer Number

0

50

Effective dimension

0

5

10

15

Layer Number

0.6

0.8

Separability score

0

5

10

15

Layer Number

0.1

0.2

Singular alignment

0

5

10

15

Layer Number

0.00

0.01

0.02

Composite alignment

0

5

10

15

Layer Number

0

1

Mean-based alignment

0

5

10

15

Layer Number

0.00

0.01

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Gemma-7B, average across datasets

Normal Acc.: 79.06% | w/o PTH Acc.: 25.60% | w/o IH Acc.: 5.66% | Random Heads Ablated Acc.: 40.82%

0

10

20

Layer Number

0

50

Effective dimension

0

10

20

Layer Number

0.6

0.8

Separability score

0

10

20

Layer Number

0.0

0.1

0.2

Singular alignment

0

10

20

Layer Number

0.000

0.025

0.050

Composite alignment

0

10

20

Layer Number

0

2

Mean-based alignment

0

10

20

Layer Number

0.00

0.01

0.02

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

Llama2-70B Famous

Unseen label Acc.: 57.78% | w/o PTH Acc.: 1.11% | w/o IH Acc.: 37.78% | Seen label Acc.: 78.89%

0

20

40

60

80

Layer Number

0

20

40

Effective dimension

0

20

40

60

80

Layer Number

0.5

1.0

Separability score

0

20

40

60

80

Layer Number

0.0

0.1

Singular alignment

0

20

40

60

80

Layer Number

0.00

0.02

Composite alignment

0

20

40

60

80

Layer Number

0

5

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.002

0.004

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

Llama3-8B Famous

Unseen label Acc.: 68.89% | w/o PTH Acc.: 15.56% | w/o IH Acc.: 6.67% | Seen label Acc.: 83.33%

0

10

20

30

Layer Number

0

20

Effective dimension

0

10

20

30

Layer Number

0.5

1.0

Separability score

0

10

20

30

Layer Number

0.05

0.10

Singular alignment

0

10

20

30

Layer Number

0.00

0.01

0.02

Composite alignment

0

10

20

30

Layer Number

0.0

2.5

5.0

Mean-based alignment

0

10

20

30

Layer Number

0.000

0.002

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

Llama3-70B Famous

Unseen label Acc.: 61.11% | w/o PTH Acc.: 16.67% | w/o IH Acc.: 26.67% | Seen label Acc.: 83.33%

0

20

40

60

80

Layer Number

0

20

40

Effective dimension

0

20

40

60

80

Layer Number

0.5

1.0

Separability score

0

20

40

60

80

Layer Number

0.025

0.050

Singular alignment

0

20

40

60

80

Layer Number

0.000

0.005

Composite alignment

0

20

40

60

80

Layer Number

0.0

2.5

5.0

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.001

0.002

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

Llama2-7B Famous

Unseen label Acc.: 54.44% | w/o PTH Acc.: 1.11% | w/o IH Acc.: 3.33% | Seen label Acc.: 72.22%

0

10

20

30

Layer Number

0

20

40

Effective dimension

0

10

20

30

Layer Number

0.5

1.0

Separability score

0

10

20

30

Layer Number

0.05

0.10

0.15

Singular alignment

0

10

20

30

Layer Number

0.00

0.02

Composite alignment

0

10

20

30

Layer Number

0

5

Mean-based alignment

0

10

20

30

Layer Number

0.0000

0.0025

0.0050

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

0

10

20

30

40

Layer Number

0

20

Effective dimension

0

10

20

30

40

Layer Number

0.5

1.0

Separability score

0

10

20

30

40

Layer Number

0.05

0.10

Singular alignment

0

10

20

30

40

Layer Number

0.00

0.01

0.02

Composite alignment

0

10

20

30

40

Layer Number

0.0

2.5

5.0

Mean-based alignment

0

10

20

30

40

Layer Number

0.0000

0.0025

0.0050

Var.-based alignment

Llama2-13B Famous

Unseen label Acc.: 52.22% | w/o PTH Acc.: 0.00% | w/o IH Acc.: 0.00% | Seen label Acc.: 67.78%

Unseen label

w/o PTH

w/o IH

Seen label

Gemma-2B Famous

Unseen label Acc.: 60.00% | w/o PTH Acc.: 0.00% | w/o IH Acc.: 0.00% | Seen label Acc.: 73.33%

0

5

10

15

Layer Number

0

20

Effective dimension

0

5

10

15

Layer Number

0.5

1.0

Separability score

0

5

10

15

Layer Number

0.05

0.10

Singular alignment

0

5

10

15

Layer Number

0.00

0.02

Composite alignment

0

5

10

15

Layer Number

0

5

Mean-based alignment

0

5

10

15

Layer Number

0.000

0.002

0.004

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

Gemma-7B Famous

Unseen label Acc.: 50.00% | w/o PTH Acc.: 26.67% | w/o IH Acc.: 25.56% | Seen label Acc.: 75.56%

0

10

20

Layer Number

0

20

Effective dimension

0

10

20

Layer Number

0.5

1.0

Separability score

0

10

20

Layer Number

0.05

0.10

Singular alignment

0

10

20

Layer Number

0.00

0.02

Composite alignment

0

10

20

Layer Number

0.0

2.5

5.0

Mean-based alignment

0

10

20

Layer Number

0.0000

0.0025

0.0050

Var.-based alignment

Unseen label

w/o PTH

w/o IH

Seen label

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

10

20

30

40

50

60

70

80

layer

0.05

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(f) CB

0

5

10

15

20

25

30

layer

0.00

0.01

0.02

0.03

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

5

10

15

20

25

30

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

5

10

15

20

25

30

layer

0.00

0.02

0.04

0.06

0.08

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

5

10

15

20

25

30

layer

0.000

0.025

0.050

0.075

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

5

10

15

20

25

30

layer

0.04

0.02

0.00

0.02

0.04

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

5

10

15

20

25

30

layer

0.04

0.02

0.00

0.02

0.04

Acc. Improvement

Rank 10 Denoising

(f) CB

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

10

20

30

40

50

60

70

80

layer

0.05

0.00

0.05

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

10

20

30

40

50

60

70

80

layer

0.000

0.025

0.050

0.075

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

10

20

30

40

50

60

70

80

layer

0.05

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 10 Denoising

(f) CB

0

5

10

15

20

25

30

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

5

10

15

20

25

30

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

5

10

15

20

25

30

layer

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

5

10

15

20

25

30

layer

0.00

0.02

0.04

0.06

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

5

10

15

20

25

30

layer

0.02

0.00

0.02

0.04

0.06

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

5

10

15

20

25

30

layer

0.05

0.00

0.05

Acc. Improvement

Rank 10 Denoising

(f) CB

0

5

10

15

20

25

30

35

40

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

5

10

15

20

25

30

35

40

layer

0.000

0.025

0.050

0.075

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

5

10

15

20

25

30

35

40

layer

0.05

0.00

0.05

0.10

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

5

10

15

20

25

30

35

40

layer

0.02

0.00

0.02

0.04

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

5

10

15

20

25

30

35

40

layer

0.00

0.05

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

5

10

15

20

25

30

35

40

layer

0.10

0.05

0.00

0.05

Acc. Improvement

Rank 10 Denoising

(f) CB

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.000

0.005

0.010

0.015

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.00

0.01

0.02

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.000

0.025

0.050

0.075

Acc. Improvement

Rank 10 Denoising

(c) TREC

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.03

0.02

0.01

0.00

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.075

0.050

0.025

0.000

0.025

Acc. Improvement

Rank 10 Denoising

(e) RTE

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

layer

0.02

0.00

0.02

Acc. Improvement

Rank 10 Denoising

(f) CB

0

5

10

15

20

25

layer

0.02

0.01

0.00

0.01

Acc. Improvement

Rank 10 Denoising

(a) SST-2

0

5

10

15

20

25

layer

0.02

0.00

0.02

Acc. Improvement

Rank 10 Denoising

(b) SUBJ

0

5

10

15

20

25

layer

0.02

0.00

0.02

0.04

0.06

Acc. Improvement

Rank 10 Denoising

(c) TREC

0

5

10

15

20

25

layer

0.00

0.02

0.04

Acc. Improvement

Rank 10 Denoising

(d) SNLI

0

5

10

15

20

25

layer

0.00

0.02

0.04

Acc. Improvement

Rank 10 Denoising

(e) RTE

0

5

10

15

20

25

layer

0.05

0.00

0.05

Acc. Improvement

Rank 10 Denoising

(f) CB

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

0.15

Acc. Improvement

Rank 5 Denoising

(a) SST-2

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

0.15

0.20

Acc. Improvement

Rank 5 Denoising

(b) SUBJ

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 5 Denoising

(c) TREC

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 5 Denoising

(d) SNLI

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

Acc. Improvement

Rank 5 Denoising

(e) RTE

0

10

20

30

40

50

60

70

80

layer

0.1

0.0

0.1

Acc. Improvement

Rank 5 Denoising

(f) CB

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 2 Denoising

(a) SST-2

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 2 Denoising

(b) SUBJ

0

10

20

30

40

50

60

70

80

layer

0.2

0.1

0.0

Acc. Improvement

Rank 2 Denoising

(c) TREC

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 2 Denoising

(d) SNLI

0

10

20

30

40

50

60

70

80

layer

0.00

0.05

0.10

0.15

0.20

Acc. Improvement

Rank 2 Denoising

(e) RTE

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 2 Denoising

(f) CB

0

10

20

30

40

50

60

70

80

layer

0.2

0.0

0.2

Acc. Improvement

Rank 1 Denoising

(a) SST-2

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

0.3

Acc. Improvement

Rank 1 Denoising

(b) SUBJ

0

10

20

30

40

50

60

70

80

layer

0.3

0.2

0.1

0.0

Acc. Improvement

Rank 1 Denoising

(c) TREC

0

10

20

30

40

50

60

70

80

layer

0.0

0.1

0.2

Acc. Improvement

Rank 1 Denoising

(d) SNLI

0

10

20

30

40

50

60

70

80

layer

0.2

0.0

0.2

Acc. Improvement

Rank 1 Denoising

(e) RTE

0

10

20

30

40

50

60

70

80

layer

0.1

0.0

0.1

0.2

Acc. Improvement

Rank 1 Denoising

(f) CB

Normal Acc.: 0.24% | w/o PTH Acc.: 0.00% | w/o IH Acc.: 0.05% | Random Heads Ablated Acc.: 0.08%

0

20

40

60

80

Layer Number

0

50

100

Effective dimension

0

20

40

60

80

Layer Number

0.6

0.8

Separability score

0

20

40

60

80

Layer Number

0.025

0.050

0.075

Singular alignment

0

20

40

60

80

Layer Number

0.000

0.002

0.004

Composite alignment

0

20

40

60

80

Layer Number

0

1

Mean-based alignment

0

20

40

60

80

Layer Number

0.000

0.001

Var.-based alignment

Normal

w/o PTH

w/o IH

Random Heads Ablated

0

8

16

24

32

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

24

32

Layer

0.0000

0.0025

0.0050

0.0075

Compo. Align.

0-shot

8-shot

0

8

16

24

32

40

48

56

64

72

80

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

24

32

40

48

56

64

72

80

Layer

0.000

0.001

0.002

Compo. Align.

0-shot

8-shot

0

8

16

24

32

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

24

32

Layer

0.000

0.002

0.004

0.006

Compo. Align.

0-shot

8-shot

0

8

16

24

32

40

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

24

32

40

Layer

0.000

0.005

0.010

0.015

Compo. Align.

0-shot

8-shot

0

8

16

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

0

8

16

Layer

0.00

0.01

0.02

0.03

Compo. Align.

0-shot

8-shot

0

8

16

24

Layer

0.6

0.8

1.0

Sep. Score

0-shot

8-shot

16

Layer

0

8

16

24

32

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

Layer

0.000

0.005

0.010

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

40

48

56

64

72

80

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

40

48

56

64

72

80

Layer

0.0000

0.0005

0.0010

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

Compo. Align.

0.015

0.010

0.005

0.000

0-shot

8-shot

8

0

24

0

8

16

24

32

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

Layer

0.00

0.02

0.04

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

40

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

32

40

Layer

0.000

0.025

0.050

0.075

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

Layer

0.00

0.01

0.02

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

Layer

0.6

0.8

Sep. Score

PTH TV

IH TV

PTH as TV

IH as TV

0

8

16

24

Layer

0.00

0.01

Compo. Align.

PTH TV

IH TV

PTH as TV

IH as TV

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Zero-shot

Task vector

ICL

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents a careful analysis of the inner workings of LLMs in the domain of ICL for classification tasks. It introduces separability and alignment as key geometric properties governing the success of ICL, formalizes them theoretically, and provides compelling empirical evidence that ICL operates as a two-stage process, first enhancing separability, then refining alignment, through the progression of information across model layers.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n* The central claim of a two-phase mechanism is substantiated through a comprehensive empirical evaluation involving seven models, six datasets, and multiple demonstration selection strategies (including random, kNN-based, and label-symbolized). The geometric framework, featuring a suite of alignment metrics (e.g., singular, composite, mean/var-based), is well-motivated and clearly described. Intervention-based experiments such as low-rank denoising and attention head ablations further strengthen the case, which validates semantic filtering hypothesis and distinguishes between the role of different attention heads in this process.\n\n* Beyond quantitative metrics, the authors also offer qualitative insights into the semantics of hidden states. They show that label-relevant tokens emerge post-transition, and that semantic directions are filtered or preserved across layers in a way that aligns with task structure. This interpretability—achieved via SVD projection analysis and singular vector decoding—adds an important dimension to the work.\n\n* The paper is clearly written, conceptually engaging, and well-structured. The logical flow of claims, evidence, and interpretation is easy to follow, making it an enjoyable read. \n\nWeaknesses \n\n* The scope of the analysis is narrowly focused on classification tasks within ICL, limiting its generality. The framework is not extended to generation or other types of ICL tasks. This raises concerns about the broader impact of the findings. \n\n* While the analysis is rich, the practical implications are underdeveloped. For example, the paper convincingly shows that alignment is the main bottleneck in zero-shot settings, but does not propose a method to improve zero-shot classification by leveraging the separability already present in hidden representations. This leaves the work primarily as an interpretive contribution, rather than one with actionable outcomes. \n\n* The paper emphasizes the connection between model components (e.g., attention heads) and hidden state geometry, but falls short in delivering empirical evidence for this claim. The identification of PTHs and IHs, though central to the proposed mechanism, is relegated to an appendix with little discussion in the main text. Moreover, the distribution of these heads across layers—which would directly support the two-phase hypothesis—is not reported.\n\nQUESTIONS:\nHow can the proposed interpretive framework for in-context learning (ICL) be extended to support generation and other non-classification tasks, and does it reveal similar alignment and separability patterns in those settings?\n\nCan alignment mechanisms in zero-shot ICL be improved by explicitly leveraging the observed separability in hidden representations?\n\nImportant: what practical gains can such methods achieve in downstream performance?\n\nLIMITATIONS:\nadequately addressed", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a unified framework to explain the mechanism of in-context learning (ICL) by analyzing two geometric properties of query hidden states: **separability** and **alignment**. The key findings include:  \n1. **Early layers** (dominated by Previous Token Heads, PTHs) enhance separability (distinguishing hidden states of queries from different classes);  \n2. **Later layers** (dominated by Induction Heads, IHs) optimize alignment between hidden states and label unembedding vectors while filtering out task-irrelevant semantics (as visualized in **Figure 1B-D**).\n\nThe framework unifies the roles of attention heads (PTHs/IHs) and task vectors, explaining why IHs’ outputs serve as effective task vectors.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:  \n1. **Theory**: link separability/alignment to ICL classification accuracy, providing a geometric interpretation of ICL mechanisms.  \n2. **Comprehensive experiments**:  \n   - Covers 7 major models (Llama2/3, Gemma) and 6 tasks (sentiment analysis/NLI, etc.);  \n   - Multi-angle validation (layer-wise dynamics, attention head ablation, task vector intervention). For example, **Figure 6** shows PTH ablation reduces separability, while IH ablation disrupts alignment.  \n3. **Clear mechanistic insights**:  \n   - Reveals a two-stage phase transition (**Figure 2**: separability ↑ in early layers, alignment ↑ in later layers);  \n   - Explains the efficacy of IHs’ outputs as task vectors (**Figure 7**: adding IH outputs improves zero-shot hidden state alignment).  \n\n**Weaknesses**:\n1. **Incomplete literature research**. There is a vast scope (from 2022 to now) of \"task vector\" literature that the author fails to discuss their relationship to their work, to name a few:\n\n[1]. Jiang et al. (ICLR Oral 2025) Unlocking the Power of Function Vectors for Characterizing and Mitigating Catastrophic Forgetting in Continual Instruction Tuning\n\n[2]. Li et al. (ICLR Oral 2025) When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers\n\n[3]. Yoshida et al. (ICLR Poster 2025) Mastering Task Arithmetic: Jp as a Key Indicator for Weight Disentanglement\n\n[4]. Cao et al. (ICLR Poster 2025) ParamΔ for Direct Weight Mixing: Post-Train Large Language Model at Zero Cost\n\n[5]. He et al. (TMLR 2025) Localize-and-Stitch: Efficient Model Merging via Sparse Task Arithmetic\n\n[6]. Lee et al. (ACL 2025) Dynamic Fisher-weighted Model Merging via Bayesian Optimization\n\n[7]. Kahardipraja et al. (2025) The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation\n\n[8]. Zhang et al. (2025) Understanding Fact Recall in Language Models: Why Two-Stage Training Encourages Memorization but Mixed Training Teaches\n\n[9]. Zeng et al. (2025) Efficient Model Editing with Task Vector Bases: A Theoretical Framework and Scalable Approach\n\n[10]. Wortsman et al. (2022) Model soups; averacing weichts of multiple fine-tuned models improves accuracy without increasing inference time.\n\n[11]. lharco et al. (2023). Editing Models with Task Arithmetic.\n\n[12]. Matena et al. (2022). Merging Models with Fisher-Weighted Averaging.\n\n[13]. Jin et al. (2023) Dataless Knowledge Fusion by Merging Weights of Language Models. \n\n[14]. Bu et al. (2025) Provable In-Context Vector Arithmetic via Retrieving Task Concepts\n\n[15]. Merullo et al. (2024)  Language models implement simple word2vec-style vector arithmetic\n\n[16]. Jiang et al. (2024) On the origins of linear representations in large language models\n\n[17]. Jiang et al. (2025) From Compression to Expansion: A Layerwise Analysis of In-Context Learning\n\n[18]. Han et al. (ICML Spotlight 2025) Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective\n\n[19]. Zhou et al. (2024) Going Beyond Linear Mode Connectivity: The Layerwise Linear Feature Connectivity\n\n[20]. Zhou et al. (2024) On the Emergence of Cross-Task Linearity in the Pretraining-Finetuning Paradigm\n\n2. **In-accurate and In-sufficient Discussion of the Definition of Task Vector**. Indeed, the author would find their \"task vector\" view is more like a \"latent steering vector\" (in Jiang et al. (2024)) that utilized in many work not in the ICL setting I listed above. However, in vision settings and non In-Context setting, the work listed above have different definitions and formulation of (general) task vector but indeed similar to the author's, despite the author considers ICL scenario. However, according to Hendel et al. (2023), the function vector leveraging function f(x, h) is indeed not linear and remain a black-box function, only when operating on \"single token factual-recall task\" in Merullo et al. (2024), the \"task vector\" defined in Hendel et al. (2023) would operate a simple vector addition in the latent space. In the semantic classification task the author defined, the task vector defined in Hendel et al. (2023) is not a latent steering vector, and should not operate by vector addition as in Merullo et al. (2024). The task vector in Hendel et al. (2023) is, when attending to any query with the black box function vector leveraging function f(x, h), can conduct prediction in one-shot. (The \"task vector\" in Hendel et al. (2023) is more than linear latent steering vector.)\n\n3. **Missing Discussions over ICA**. According to Yamagiwa et al. (2023), words embedding would be better whitening through ICA instead of PCA, where each embedding can be expressed as a composition of a few intrinsic interpretable axes. Therefore, in the author's scenario, the IHs' \"filtering irrelevant task information\" is like filtering those task-irrelecent ICA-decomposed interpratable components in the ICL prompt. Therefore, experiments with ICA should be complementary to the author's PCA's ones.\n\n4. **Broader Impact might be mild**. See Questions.\n\nYamagiwa et al. (2023) Discovering Universal Geometry in Embeddings with ICA\n\nQUESTIONS:\n1. Kahardipraja et al. (2025) also provide different genre of heads in ICL, as well as function vectors. I understand it is a cocurrent paper, but I would like to see the author's detailed comparision of the deferences in terms of definitions of heads and function vector. Indeed, [17] Jiang et al. (2025)'s finding is also similar (two phase as compression-expansion) but also contradict (task vector is formed at 17th layer similar to Merullo et al. (2024)) to the author's, and the (15) in their paper is the correct task vector definition according to Hendel et al (2023).\n\n2. In terms of Hendel et al (2023), what should be the task vector and the black-box task-vector leveraging function, and can you uncover the black-box nature or its relationship to latent steering vectors of labels?\n\n3. Han et al. (ICML Spotlight 2025) provided that Task Decodability (TD) quantifies how well latent tasks can be inferred from intermediate representations that predict ICL performance, and early layers count more-early-layer finetuning outperforms late-layer tuning. This contradicts the author's view where the former layers' role didn't change too much and ICL make little changes, compared with the latter layer's role in aligning to the label unembedding. My intuition is because the author doesn't refer to the task vector in Hendel et al (2023)'s definition, which, if exists would appear in mid-to-late layers. Also, from my own experimental experience, the sentence/token level semantic classification task's task vector defined in Hendel et al (2023) and its f(x, h) are non-trivial and non-linear -- The task vector extracted from \"This book is bad: negative; The movie is fun: positive; Today I'm unhappy:\", cannot be simply adopted to any query just through a simple addition in the latent space and produce a correct answer, and the function vector leveraging function f(x, h) is a black box. For single token case, only the single token factual recall ICL task in Merullo et al. (2024) is true in terms of the equivalence of function vector and latent steering vector where f(x, h) is linear. Given the existing findings, the author fail to concretly grasp the existing literature's finding, make conclusions on too restricted setting, and the impact is too limited despite the presentation is clear.\n\nIn all, I found the anthors misalign with the broad task vector literature and didn't make a good comparision and summerization. For example, the task vector defined in Hendel et al (2023), Merullo et al (2024), [17] Jiang et al. (2025) emerged in early-to-medium layers, not refering the the last layer's steering vector but indeed warrent greater interest in terms of In Context Learning (the task vector and its function vector leveraging function f(x, h) is not simply a steering in the latent space that help align the label unembedding). In fact, if the author prefer to define latent steering vectors that help align label unembedding as the task vector, it somewhat similar to the broad literature I mentioned above regarding using steering vector / linear mode connectivity started 2022, but these important discussions are currently missing in the author's manuscript.\n\nLIMITATIONS:\nSee Weakness and Question above. Since the task vector literature is so large, it is hasty to only focus on single token semantic classification, in the scenario that single token factural-recall (Merrulo et al) and other tasks have been well-documented. Not to mention that I heavily question about the authors' familarity about task vector celebrated many years in the community.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main goal of the paper is to unify the task vector and the induction-head accounts for in-context learning, with a focus on classification problems. The authors show that, during in-context learning, the model representations go through two phases: First, the representations of the different classes are separated, then they are aligned with their corresponding unembedding vectors. The authors show, both through observational and causal studies, that predominantly the former process is carried out by previous token heads, and the latter through induction heads.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\nThe paper is very clear in its goals and how it attempts (and succeeds) to unify two distinct but very related approaches to in-context learning. There is extensive empirical evidence using various methods and models to support the claims. The findings are quite clear and easy to interpret. The figures clearly communicate the results (especially Figure 1). Lastly, I appreciate the authors share their code, which seems easy to understand.\n\n## Weaknesses\n\nThe main weaknesses come from the ablation studies in section 5.3. The authors do not clarify how the ablations are done, but from previous sections (line 127), I'm inferring that they did a 0 ablation, where they set the output of a given head to 0. I'm not sure if this is the best type of ablation, as it might shift the activations out of distribution. I'm not majorly concerned about this, as your random head ablations using this method do not seem to have a negative effect. However, I think a mean ablation would be a more principled approach. See [this paper](https://arxiv.org/abs/2404.15255) for a discussion.\n\nOf greater interest is the number of heads ablated and where these heads are. As we know, there exists a bunch of induction heads across multiple layers in large models. Do you ablate them all? Where are they located? Do you identify them simply using the prefix matching score on a random text? These are all important questions regarding the interpretation of the results. Apologies if I missed these details, but I could only find a description of how you construct the task vectors in the Appendix, not of these ablation decisions.\n\n--- \n\nI believe this is already a strong paper. I'd be further willing to increase my score if the questions regarding the ablation studies above (especially regarding the details), and the clarification questions below are addressed adequately.\n\nQUESTIONS:\n- Why does the mean-based alignment results suggest PTH are more crucial than IH? Would we not expect IH ablations to create more dramatic effects.\n- There is this recent paper named [Which Attention Heads Matter for In-Context Learning?](https://arxiv.org/abs/2502.14010), where the authors show that, while function vector and induction heads are connected early in training, function vector heads specialise to be functionally different types of heads, and that those are the parts of the model that drive in-context learning. However, you provide some findings where you can create task vectors simply through the outputs of induction heads. I think this is a fascinating difference of findings, and would be curious to hear your opinion on why this could be. Is it a matter of task, how you identify the induction heads, or something else?\n\nLIMITATIONS:\nYes, limitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the unified mechanism of attention heads and task vectors on the classification performance of ICL. The authors start by theoretically deriving two conditions: output alignment and directional alignment for the hidden states of the queries to achieve maximum separability, which is the upper bound of the classification accuracy. They further propose several tractable metrics to measure the separability, output alignment, and directional alignment for various real-world LLMs. The experiments reveal the dynamics of the separability and alignment across the layers of LLMs, and how the previous token heads (PTH) and induction heads (IH) affect these features.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. This work depicts a clear view of the evolution of the hidden features induced by ICL, offering a key view for explaining how ICL improves the classification performance: ICL mainly improves the alignment between the hidden states of the queries and the label-related unembedding vectors, rather than improving separability.\n\n2. The proposed separability and alignment metrics are theoretically motivated. Moreover, these two metrics can serve as two fine-grained indicators of the performance in classification tasks of ICL.\n\n3. The empirical results of the layer-wise trends of separability and alignment (Fig. 2 and Fig. 3) are convincing due to the wide range of the tested model architectures, model sizes, datasets, and experimental settings (number of demonstrations, how the demonstrations are selected, etc.).\n\n4. The findings of the roles in shaping the hidden representations of different types of attention heads are novel and interesting.\n\n**Weaknesses**\n\n1. The theory and experiments are limited to classification tasks. The proposed metrics cannot be applied to more diverse non-classification tasks, such as general QA or instruction following tasks, which are more common in practice.\n\n2. The mechanism of why PTHs improve separability and why IHs improve alignment is not fully revealed in this work. \n\n3. How the empirical findings could inspire potential improvements of ICL designs is not discussed.\n\nQUESTIONS:\n1. In lines 309-310, based on my understanding, shouldn't the PTHs attend to \"I like it.\" instead of \"Answer\"?\n2. Could you further elaborate on, at least intuitively, how PTHs enhance separability and how IHs improve alignment? I find the explanations from line 329 to line 338 a bit vague.\n3. Could you discuss how the findings in this work could potentially benefit the design of ICL algorithms?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}