Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning
Authors: Haolin Yang
Hakaze Cho
Yiqiao Zhong
Naoya Inoue
Venue: NeurIPS 2025 poster ()
Keywords: In-Context Learning; Induction Head; Mechanistic Interpretability; Task Vector;
Primary area: deep_learning
TL;DR: We unify attention heads and task vectors in the mechanistic analysis of in-context learning.

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
Haolin Yang'
Hakaze Cho²
Yiqiao Zhong
Naoya Inoue
University of Chicago JAIST
University of Wisconsin Madison
*RIKEN
haoli haalinyang20016e www.enggoooocccccoooooo.ooo edu, yfzhao8 taist.as ac. jp yiqiao. zwongewisc.ooo. edu, naoya- 10 jaiiijjjssttt jaijt.ts....... ac.
Abstract
The unusual properties of in-context learning (ICL) have prompted investigations into the internal mechanisms of large language models. Prior work typically focuses either special attention heads task vectors at specific layers, but lacks unified framework linking these components to the evolution of hidden states across layers that ultimately produce the model's output. In this paper, propose such 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 striking two-stage 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 Introduction
One of the most remarkable features Large Language Models (LLMs) 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 Figure withhut any additional training. This capability is known as In-context Learning (ICL) [Brown 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 L.MMs variety of downstream tasks [Dong 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 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 Induction Heads (IH) [Elhage et al., 2021, Olsson et al., 2022, Cho et al., 2025a], which retrieve information from demonstrations and perform ciky-like operations within Transformer-bes based architectures. Other studies adopt hidden- state-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 nnertwined-coopooeent like attention heads transform hidden states layer by layer, and these intermediate updates ultimately determine the
I Code implementation: https::tp//t https:://itth/: com http://www.cccccmma///////////// Hideen____dde__eeeeeeeeeeeeeee__ HedddeGeeedd.eeee..eeee......
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(A)
like this. Sentiment: positive It is bad. Sentiment: negative
Fantastic! Sentiment:
(B)
(C)
(D)
LM Layers
(PTH)
Separability
LM LaLaye (HH)
Task Vector
LM Heads
Pos positivv
SAli
Alligmeet
positive samples
neegttive samples
Figure 1: (A) An example for ICL imput. (B) early layers. LLMs promote separability among the last tokens' hidden state (h) clusters w.r.t the ground-truth labels of the queries through Previous Token Heads (PTHs). (C) In early layers or zero-shot scenarios, the direction where the hiddens are maximally separated is insufficiently aligned with the output direction (i.e., the difference vector of the label-token unembedding vectors), increasing cluster overlap after mapping and inducing higher classification errors, and (D) in later layers, Induction Heads (IHs) align these clusters towards the output direction, with the same underlying mechanism of task vectors.
model's 's output. As such, analyzing only the influence of attention heads or task vectors the final prediction reduces ICL to a black-box process. A true mechanistic understanding 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 direction that aligns with the difference between the unembedding vectors of different labels. To investigate this, we collect hidden states in both ICI. and zero-shot settings and analyze them using a range geometric measures to assess how ICL, improves hidden states in terms of separability and alignment. Our approach operates 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 ICI. 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 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 proliferation of Previous Token Heads (PTHs) attend to the previous token at each position-in early layers. (2) In middle-to-late layers, Induction Heads (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 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].
heads that
2
Related works
IH and PTH Previous works on the mechanistic cause of ICL identify structural circuits in Transformers comprising Previous Token Head (PTH) and an Induction Head (IH), where the PTH gathers information at demonstration label tokens, and the IH copies to query positions for prediction [Oisson 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 steer zero-shot hidden states increase prediction accuracy. These vectors may be ICL query hidden states extracted at 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"-sty"" 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 parametrie knowledge produce task vectors with distinct function. Jiang et al. [2025b] and Han et al. [2025] study how the geometry of model's 's task representation impacts the effects of injecting task vectors. Jiang et al. [2025a] reveal the correlation between model's '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 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-trutt label y, e.g., x = like this movie. Sentiment:" and y= "positive". The LLM with L layers encodes into a vector h, E IRd, taken from the hidden state at the final token (*t"), Given a vocabulary V = {VI1,.... | and an output embedding matrix E E IR [4// the model predicts (x)) = exp(E, /,,,,,,,,/ exp(E,, ha), where E,, is the unembedding vector of vy. The model outputs the token Vj with the highest probability the predicted label for
In the ICL scenario, as shown in Figure (A), demonstration- label pairs (x1.))), .... (M. are prepended to form the prompt (x1, ve.x), which modifies the hidden state of x and results in a different distribution ((v.,
potentially improving prediction accuracy.
Rather than a single query, we more often evaluate classification accuracy over a dataset of n queries (x,,-...-k)) with hidden states (h,, and ground-truth labels {y1,-..yoo) We will show that the classification accuracy depends on two geometric properties of the hidden state collection H [ E _exd; (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:
n
exp(Ev,hi) hi)
E h
Acc Σ 16y arg max
Σ 1(y;
arg max
IV Σ
Σ exp(E,hh)
Vj
E
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,,...X,) only has 2 different labels, yA and ya. i.. y.). Vi, and leave the generalization to the case of more labels to Appendix B. For a unit vector u E Idd, we define the separability of H along direction u with labels YA. to be:
y: E Va).
1 n
S(u)
Σ IENA
1u h, 2 0)+ + Σ 1( h < <)), NN
where NA = (i: } and N₂ = (i: S(u) measures the fraction of hidden states that can be correctly classified by a linear decision boundary orthogonal u. By considering all possible directions, we obtain the maximum separability of H. i.c. S" suppeRd [|al||== S(u). Though S(u)
3

---

is not continuous function, S* as a supremum can be attained on S {u : E IR ||al/2
(See Appendix A). Denote u² as the where S attains the supremum, then:
Theorem 1 Acc S. The equality is achieved when maNveV.vebu.vel E,h; < max (E hi hi), Vi cu' for some positive constant c.
and E, Evg Eggl
This theorem shows that accuracy is upper-bounded by the maximum separability of H and that the bound is achieved when: (1) each h; is maximally aligned with either Ey, or v unembedding vectors of УA and ya w excluding interference from other tokens, a condition we call output alignment; and (2) the direction of maximum separation u' aligns with y, - -Eyg. which we refer to as directional alignment. Thus, scparability 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 1, the hh attention head adds d', to the hidden state h of the fd query, affecting the transition to 1++ Ablating head (i.e., setting a' 0) alters all downstream hidden states and can degrade accuracy if the final-layer hidden states HL lose separability or alignment.
Task-vector-base based experiments operate similarly by steering hidden states intermediate layers. In those experiments, a task vector is added to hy and the influence likewise propagates through subsequent layers. This intervention is effective if the final Hr has better separability and is better aligned with the labels of the task encoded in f,-
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 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 affected by the maximum separability of hidden states. However, finding maximum separability requires evaluating separability along infinitely many direc- tions in Sd-1 and is intractable. As practical proxy, we train a logistic classifier on a subset H and evaluate its accuracy on held-out data. We call the accuracy separability score 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 the classification accuracy obtained by applying the unembedding matrix E directly to hidden states H, skipping subsequent layers. measures how well H aligns with the label unembedding vectors E,, and E and is also known logit lens accuracy [nostalgebraist, 2020]. When comes from the final layer, it equals standard classification accuracy.
B.2-R.5 Directional alignment
Directional alignment is the alignment between u² and the label unembedding difference direction
Eyy E,g [E, ll Since finding u² is intractable, we introduce the following proxies based on approximat- ing "" or measuring separability along the label-difference direction.
2 Activatiio patching which directly replaces the zero-shot hidden states with ICL hidden states can also be subsumed under this form, with I, being the difference between the hidden states in the two settings.
'For binary classification datasets, each measure is computed over the entire H. For multiclass datasets with labels Ym,. we enumerate all label pairs (y),)t), compute the measure using the corresponding two label clusters, i.e. hidden states of queries with labels Ye. before averaging across all pairs to obtain a scalar summary for H. As shown in Appendix B, this is justified by the fact that separability and alignment influence the accuracy of multiclass classification through the binary classification between each label pair.

---

Singular alignment (B.2). To approximate u², we apply Singular Value Decomposition (SVD) to the mean-centered H and extract the top-r 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 - || E,p
as a proxy for directional alignment.
- Variance-based alignment (B.3). To measure separability along - 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
E.g- Evp with labels FA and yB onto |Ey Eigli and divide it by the pooled projected variance. A higher
value implies better separation margin along Ey, EBy and better directional alignment.
E Eyg
Composite alignment (B.5). To unify the mean- and variance- basun alignment, define composite alignment as their product. It captures both the magnitude and the relative significance of separation along the lafel-difference direction, serving 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 small number of directions, potentially those associated with the label unembedding vectors, and are thus easier separate.
Experiments and Results
Models We experiment on the following models: Llama2-7B, Llama2-13B, Llama2-70B [Touvron al., 2023], Llama3-8B, Llama3-70B [Grattafiori 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 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 Subsec- tion 5.1) to test the generalizability of our results in settings other than classification. We also use generated dataset to explore the scenario where ground-truth labels of queries are not presented the demonstrations (discussed in Appendix 1.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 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 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 in Appendix F). A key observation is that Figure 2 reveals striking phase transition in the evolution along the layers of ICI. hidden states. Phase (Figure (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 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
5

---

Llama2-70B: 8-shol Acc.: 80.57%; 0-shot Acc.: 0.24%
Effective dimension
Separability score
Singular alignment
0.2
50
0.8
0.6
0
0.0
20
40
80
Layer Number
Layer Number
Layer Number
Composite alignment
Mean- based alignment
Var:based alignment
0.010
8-shot
0.02
2
O-shot
0.005
0.00
20
40
60
80
0
20
40
80
80
20
40
60
80
Layer Number
Layer Number
Layer Number
Figure 2: Comparison of trends in separability and alignment measures: ICL vs. zero-shot.
Under ICL., a clear phase transition emerges: separability increases first and then alignment surges. The effective dimension first rises and then declines. This pattern is missing in the zero-shot setting. Accuracy gains from ICL over zero-shot are reflected in alignment, not separability measures.
monotonously and the rise in alignment measures is not nearly pronounced. Hence, this phase transition pattern characterizes ICL as a two-stage process: an initial phase enhancing separability, followed by refinement phase that improves alignment.
ICL improves aceuracy primarily through alignment A key observation is that ICL and zero-shot settings show only a small gap in separability scores 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 main effect of ICL. (see Appendix E. for the practical implication of this), making alignment the primary bottleneck zero-shot performance. As shown in Figure (C), the high-separability direction in zero-shot hidden states is misaligned with the label- difference 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 sriaacce-bsse 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 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 (A), 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 remains-disting distinguishing it from zero-shot, where effective dimension rises monotonically-
the post-transition alignment surge dampened.
Phase transition holds in open-ended settings To test whether the phase transition pattern occurs in tasks beyond classification, i.c., open-ended generative tasks, where ICL is also frequently applied, we consider the following task: the queries assume the format "Praise/Critique (subject }", where (sub ject. ) is the name of food (e.g.. tapas). The labels are reviews for the subject satisfying the sentiment requirement, generated using GPT-4o [OpenAl et al., 2024], for instance, "Praise tapas - The tapas assortment at this restaurant absolutely delightful!". We track the layer-wise dynamics of the separability score and composite alignment" as representatives of the metrics evaluated on the 8-shot and O-shot hidden states of this task. The results in Figure 4 mirror those in Figure 2, where
*Since generation task has an indefinite and infinite label space. we let the label unembedding difference direction needed calculate composite alignment be that between "positive" and "negative".

---

Avemge
Omigotte alignment
60
Soparabiity scure
Seggulin algnmmmtt
Accurang
22
6--
66 Layer Nember
25
66
25
20
16
Demo
Lagar Namber
Layer Number
Layer Number
(a) Varying number of demonstrations
Liama2 70B: Random Selection Ace.: 80. 57,5 KNN Selection Acc.: 88.01% Effectiive dimension
Maama2-70: Normal Label, Acc.: 80.3 57%, Symbol Label Acc. 59.05%
Segarabiiity some
Singular alignment.
Effetiive dimeasiio
Seperablity scoze
Sinqular aligument
60
40
D.1
50
50 Layer Number
50 Layer Number
50 Lopee Number
50 Loper Namber
Layse Namber
Layer Number
Composie aligmment
Mean-basert alignment
Var. bareba alignment
Comp-ssite allgnment.
Mean based alignment
Var-based alignment
A
0.000
0.15
0.01
0.00
0.005
0.00
0.00
1,000
50 Layse Numbee
50 Number
50
Number
Number
Number
Number
Selection
Symbol Libei
(b) KNN demonstration selection
(c) Symbolized demonstration labels
Figure 3: Layer-wise trends of separability and alignment measures in different ICL settings. (A) Phase transition is evident under demonstration numbers from to 24. Accuracy improvements of increasing demonstrations reflected by consistently improving alignment measures.
(B) Changing
the demonstration selection method to kNN retrieval preserves phase transition and improves accuracy through enhancing alignment. (C) Using uninformative demonstration labels hurts accuracy due to decreased alignment of hidden states, yet a similar phase transition is evident.
the few-shot and zero-shot hidden states exhibit similar separability properties across layers, but the surge in composite alignment beginning from mid-layers specific to few-shot hidden states.
Alignment explains accu- racy differences across ICL settings We find that: accuracy differences be- tween ICL settings are con- sistently reflected in align- ment measures. Specif- ically, in Figure 3 (A), increasing the number of demonstrations from 4 to 24 produces minimal changes in separability, while alignment measures improve in consistent manner as the trend curves shifted versions of one another. Effective dimension also decreases steadily, suggesting a higher concentration of hidden states along uneebddien directions. Figure (B) similarly shows that the accuracy gains from kNN-selected demonstrations correspond to stronger alignment and lower effective dimension. In Figure (C), the accuracy drop from symbolic label replacement aligns with a flattened post-transition rise in alignment measures, resembling the case zero-shot.
1.0
Compo. Allgn.
Sep. Score
B-shot
0.002
D-shot
U-shot
0.000
(a) Layer-wise separability score:
(b) Layer-wise composite alignment:
shot vs. 8-shot
0-shot vs. 8-shot
Figure 4: (A) In the generation setting, the separability score exhibits almost identical trends in the 0-shot and 8-shot cases. (B) In the 8-shot
case, the phase-tree transition of composite alignment is pronounced.
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 case study on the SST-2 dataset to interpret hidden-star state alignment dynamics.
Surge in alignment concurs with the encoding 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 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 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 il-immediately post-transition. This confirms that post-transition layers elevate label-relevant
directions in both the semantic and variance structure of the hidden states.
(A) 1.0
Phase Transition
1.0
Phase Transition
Layer-wise pattern of semantic re- tention and filtering To better un- derstand how exactly layer updates after the phase transition encode rel- evant semantics into hidden states, we analyze the semantic informa- tion retained or removed by each layer. Given the SVDs of the cen- tered hidden states at two consecu- tive layers - and
Output Alignment
0.8-
0.0R
Composile Alienumntt
Outgut Aligmment
0.8
0.30
Singuu Allgnig momen
0.25
0.08
0.6-
0.6
0.20
0.4-
0.04
0.4
0.15
0.30
0.2-
0.02
0.2
0.05
0.0
0.00
0.0
0.00
20
40
60
40
60
80
Layer Number
Layer Number
(B)
EE,
Layer 21 opesdart. regastive, pea, zech, stetn, and spandast, regastve, ster, anth, sarth,
Layer
ipendatt,
X1
H, = we compute the pro- jection of the previous layer's singu- lar directions onto the next layer's i as V.VVVV E R The rows of this matrix with larger norms are the prin-
Eaper
riltrred (terovue))
amony Layers
Laperao Alken Balen, lor, erge. Tan, Boday, Ticheets, lez,
Layarde Bet, evel, is, enf, Tak, tatch, seal, tont, Stem.
V
Tagerre cot, ye, th, archell, cirl, schee, 365, recid, and
Layerse
iset,
Retained Semantion smong Lavers
cipal directions of H,-1 that are re-
27 ibser, orp, lif, seak, slaaporiss, anth, di dlaxgpoint, and
tained in H₁, 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, label- related semantic directions are consis- tently 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 cor- responds 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 presented in Appendix G.4.
16, pontetion,
anda, salenn,
criefe,
canesy oel, obes As, Stach, 358, fr. by. wuther, waeer, atherat Figure 5: Dynamics of alignment measures and semantics of post-transition hidden states. (A) Strong correlation between output and directional alignment; (B) Surge in alignment measures concur with encoding of label-related semanties; (C) Post-transition layers, except for the last ones, filter out unrelated semantics and retain the related ones. Refer to Appendix G for more semantics decoding cases.
TayerTe
Low-rank denoising enhances label-relevant semantics To better understand the retention- and-filtering process, we test whether filtering can be accelerated and label-relevant seman- tics enhanced by denoising the hidden states. Specifically, we apply a rank-10 SVD approxi- mation to the centered hidden states f1, retain- ing 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-transitio layers where irrelevant semantics are just beginning to be filtered. These results validate low-rank denoising and support semantic retention and filtering as core mechanism of post-transie post-transition layers. We explore the practical implication of this in Appendix E.2.
Acc, Improvement
Rank 10 Denoising
8.00
0.05
0.00
300
30
8o
layer
Figure 6: Accuracy gains of rank- 10 denoising. See Appendix G.5 for results in other settings.
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
8

---

Llama2-705h average across datasets
Normal Acc.: 80.57% w/o PTH Acc.: 33.11% w/o IH Acc.: 2.50% Random Heads Ablated Acc.: 79.13%
Effective dimension
Separability score
Singular alignment
100
0.8
0.2
0.1
80
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var-basse alignment
0.010
0.02
0.005
0.00
0.000
20
40
60
80
20
40
60
80
20
40
60
80
Layer Number
Layer Namber
Layer Number
Normal
w/o PTH
w/o IH
Random Heads Ablated
Figure 7: Effects of attention heads ablation the layer-wise separability and alignment measures. Ablating PTH significantly reduces separability, and ablating IH significantly reduces alignment.
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 -bbo-iinformetve,, 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][B][A]][]] [[A[[B]].-AAI[B,IIAA
a PTH at the final [A] attends to the immediately preceding [B,]. while an IH attends to tokens such as [B,]-[B] [-[,,, following earlier [A] positions. For example, in the ICL prompt "I don't like it. Answer: negative. I like it. Answer:", at the final "." 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 s hidden state. By analyzing their effects on separability and alignment across layers, gain more fine- ame-igained understanding of their role in ICL ICL-beyyond 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 Appendix H.
PTHs induce separability, IHS induce alignment The results in Figure confirm that PTHs and IHs are crucial not only for verbalizing ICL outputs [Cho al., 2025a] but also for structuring hidden states both separability- effect (as shown by overlapping black and gray curves). 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 aligmment) which concerns the separation margin of the label clusters, but has minimal impact other alignment measures. Conversely, ablating IHs leaves separability largely intact but severely disrupts all alignment measures, and has a far greater impact output alignment as measured by the accuracy.
and aligmment-we wise. Ablating random attention heads has a negligible
These results not only confirm two-stage ICL process-first driven by separability, then by align- ment-but also clarify its mechanism. PTHs, concen- trated in early layers [Cho et al., 2025a], enhance sep- arability by attending 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.
Table 1: Dataset-mean layer positions of top IH and PTH heads by percentage levels with statistically significant differences.
%
IH mean
PTH mean
p-value
1%
40.203
33.379
1.01 x 10-7
2%
40.884
36.459
1.30 x 10
5%
42.238
36,422
2,06 x 10-17
10%
41.535
35.850
9.72 10
Furthermore, in Appendix I.5, show that

---

PTHs and IHs contribute to ICL through the same mechanism even when the query's 's ground label- previouuy considered essential for IHs [Cho 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 pntaraaterzzatoonnnndeee
separability and alignment respectively-ww 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 Mann-Whitney Whitney test.
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 aligmment, does not show how they actively drive the geometric evolution of hid- den 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 PTHH-derieed vector at layer I = to track separability, and the IH-derived vector at L to track composite alignment. These layers correspond to where separability and alignment begin to shift markedly in Figure 2. Figure Sa shows that injecting PTH outputs as task vectors boosts separability (the green line surpassing the magenta line before convergence; see Appendix 1.2 for similar trends across models), while injecting IH task vectors restores the ICL-specific spike in alignment metrics in zero-shot hidden states.
Sep. Score
Compo. Align.
00.015
0.8
0.010
IIT TV
PTLI
0.005
0.6
0.000
8
16
21
32
40
68
56
68
72
60
16
24
32
40
68
56
64 72 80
Layer
Layer
(a) Layer-wise separability
score:
(b) Layer-wise composite alignment:
PTH task vector
IH task vector
Figure 8: (A) PTH task vector increases separability after injection:
(B) Injecting IH task vector recreates the phase- phase-transitiee transition in aligmment.
80
TH outputs task vectors Since have shown that IH outputs with ICI. demonstrations can compensate for alignment deficits in hidden states, have clarified how IH outputs as task vectors influence the forward computation of models: they steer hidden states align with label-token unembeddings so these tokens can be decoded. This also explains why IH outputs themselves decode label-related tokens [Todd 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-70f The substantial accuracy gains in Figure 9 support this conclusion.
80.57
61.11
60
Accuracy
40
20
0.24
0
0-shot TV
ICL Average effect
Figure 9:
across datasets of steering zero-shot hidden states using IH outputs as task vectors.
6 Conclusions and Limitations
This work unifies two major perspectives on in-context learning (ICL) the roles of special attention heads and task 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 gigemment-claiffyiin both the phase transition and the effectiveness of IH outputs as task vectors.
Conclusions
a geometric framework centered separability
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.
10

---

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
T. 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, 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 Interaational Conference on Neural Information Processing Systems, page 1877-1901, Online and Vancouver, Canada, 2020. URL https: /dl a.m. wwt.o///ooo///oo///ooooooooooo httggg/do//bb// /bs// 10. 00..55553395577
3495883.
S. Chan, A. Santoro. A. Lampinen, J. Wang. A. Singh, P. Richemond, J. McClelland, and F. Hill. Data distributional properties drive emergent in-context learning transformers. Proceedings of the 36th Advances in Neural Information Processing Systems, pages 18878-18891, New Orleans,
USA, 2022. URL https: / ht///raxivvv.rr/ee abs 22555
H. 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::tp/// /pw.eeveeveeeeeeeeeeeee
net/ / forum? d=xi ZpnVQQqQQ
H. Cho, Y. Sakai, M. Kato, K. Tanaka, A. Ishii, and N. Inoue. Token-based decision criteria are suboptimal in in-context learning. I.. Chiruzzo, Ritter, and L., Wang, editors, Proceedings of the 2025 Conference of the Nations the Americas Chapter of the Association for Computa- tional Linguistics: Human Language Technologies (Volume Long Pupers), pages 5378-5401, Albuquerque, New Mexico, Apr. 2025b. Association for Computational Linguistics. ISBN 979-8- 89176-189-6. URL https:ttppe..m ac aaaahoatttoy...... /20/ wwa.c-acnnnnngcoa .278/-
I. Dagan, O. Glickman, and B. Magnini. The pascal recognising textual entailment challenge. In Machine learning challenges workshop, pages 177-190. Springer, 2005. URL https: http:://t //1nkk..
springer. cwm/ chapter/10, 11tt////0000000117766116000
M.-C. De Marneffe, M. Simons, and Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pages 107-124, 2019. URL URLhttps. hwtpshtt/////. //js ub. uni-konstaaz de//sb////e htt.iiidex..pp.e hatbb/rttce//// view/ 601.
Q. 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:://aaaaee http://awxvw.////aaaa/./ org/absg////oo /ab////3b.. /2301. 00234.
N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, Askell, Bai A. Chen, Con- erly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1): 12, 2021. URL https: http///aanforroooo circuitst...... pub/ 2011/221//// framework/ connewwok/indxx.... 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 htthtt//:/// wttpvv..ivvvvv/////i//e.//.aaav. org/ abs 22112/11222 11805.
A. Grattafiori, A. Dubey, A. Jauhri, and A. P. et al. The llama 3 herd of models, 2024. URL https://arxiv. orgg als 2222244022211 21783.
W. Gurnee and M. Tegmark. Language models represent space and time, 2024. URL https: //arxiv.org/all 2310 02207.
D. Halawi, J.-S. Denain, and Steinhardt. Overthinking the truth: Understanding how language models process false demonstrations, 2024. URL https: /arxiv. htrg/abs/2300 09476.

---

S. Han, J. Song, J. Gore, and P. Agrawal. Emergence and effectiveness of task vectors in iontenttxt learning: An encoder decoder perspective, 2025. URL https://aaxaai wwt.aaiwwiiiaaaaaaaaoooooooorrot abs /bs//412. 2412. 12276.
R. Hendel, M. Geva, and A. Globerson. In-context learning creates task vectors. H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguistics; EMNLP 2023, pages 9318 93333, Singapore, Dec. 2023. Association for Computational Linguistics. doi: 10. 1188656511122222 h/oo///////222...nngnegeeem//// findings-ennll..244 URL https: :/ttps///an faclanthology.ooo .rrg/203. findings-emn-yy 624/.
G. 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: : aariivvv...../... wteeiiiorrr///e orr.or///aao... abs / 2502. 11019.
Jiang, Y. Dong, Zhou, and Z. Zhu. From compression to expansion: layerwise analysis of in-context learning, 2025b. URL
Kahardipraja. R. Achtibat, Wiegand, W. Samek, and S. Lapuschkin. The atlas of in-context Iearn- ing: How attention heads shape in-context retrieval augmentation. arXiv preprint arXiv:2505.1580) 2025.
A. 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 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-8- 89176- 195-7. URL htt.:::/aacaatttooorr.ooom wwgp://...../20200 findings-naacl.
A. Lee, X. Bai, I. Pres, M. Wattenberg, K. Kummerfeld, and R. Mihalcea. A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity, 2024. URL https:
//arxiv. org/ abs /ww./240111 221..21101......
D. 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: //taxaxivv.///////////a/////aavl www.iiiwog//..oeeeoo.ooooooooeee www.abs//////oooooeee/......aaae abs/240 11225.
X. Li and D. Roth. Learning question classifiers. In COLING 2002: The 19th International Confer- Computational Linguistics, 2002. URL httph:ht/////tt http:://wwwwwwlclw... org/ anthology/ C02-1155
J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context amples for GPT-3? E. Agirre, M. Apidianaki, and Vulić, editors, Proceedings of Deep Learning Inside Out (DeeLlO 2022): The 3rd Workshop
Knowledge Extraction and Inte- gration for Deep Learning Architectures, pages 100-114, Dublin, Ireland and Online, May 2022. Association for Computational Linguistics. doi: 10. 10.18655//120022deelleoo 1.10. URL https: alcl anthology orgg/022. deeeio-eeieee
S. 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. /ttt://///rrtt.oo///aaa//////// 2211 06668.
http//aaxvvvv/ org/abs
Z. Liu, J. Michaud, and M. Tegmark. Omnigrok: Grokking beyond algorithmic data, 2023, URL https://arxis org/ org/abs/////eoooooo/../o/// 2210.2212222 2210.01117. 01117.
B. 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: aaclanttooo log org /
Marks and M. Tegmark. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets, 2024. URL https: /ttp//axivv////////i///////aaaaa org/ hbs//bb//// /2bs 06824.
Merullo, C. Fickhoff, and E. Pavlick. Language models implement simple word2vec-style vector arithmetic, 2024. URL https: //trxivvvvv/////////////////aaav org/ abs/2305.16130.
12

---

T. Mikolov, K. Chen, G. Corrado, and J. Dean. Efficient estimation of word representations in vector space, 2013. URL http://aaaxvv...gabb.////s 1tt.aa////1311.11
nostalgebraist, Interpreting gpt: the logit lens, 2020. URL httpp::/ww.....eeeeee.
com/ /ostss//cRRBBBDDraGGGG6tr http:/ppwwwwwrrrrinttrerrreegnno -ttp:pprreeiinttttttt--eeeee-eee -ooppttttt---gttt---e
C. 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. Hattield-Dodds, D. Hernandez, 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: /ttpvv.xivivv./////////////..a org/orggg//s /bs/ 2209. 11895.
OpenAl, J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, Altenschmidt, S. Altman, S. Anadkat, R. Avila, Babuschkin, S. Balaji. Balcom, P. Bal- tescu, H. Bao, M. Bavarian, J. Belgum, I. Bello, et al. Gpt-4 technical report, 2024. URL https:/ /aresv//aavvv../ org/ abs 2303. 08774.
F. Park, Lee, E. Lubana, Y. Yang. M. Okawa, Nishi, M. Wattenberg, and H. Tanaka. Iclr: In- context learning of representations, 2024. URL https: http/aaxxvvv...... /ttrgg/ab///0000 00070.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pret- tenhofer, R. Weiss, V. Dubourg, Vanderplas, Passos, D. Cournapeau, M. Brucher, M. Per- rot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learn- ing Research, 12:2825-2830, 2011. URL hwtph:tts//// pedregosa 1 httppppdrrooaaaae pdf.
jmlr firg / paperss v2
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 2019. URL https: / d4mucffpsyywu wwwssyywwcooufroott.. net
hatpe//oooddd///ggggaeeeeeooodde modells pdff.
G. Reddy. The mechanistic basis 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.oo /abrg//b//////eeeeee//eeeeeeeeer 2312. 233558.
S. Sia, D. Mueller, and K. Duh. Where does in-context learning happen in large language models? In The eighth- Annual Conference Neural Information Processing Systems, 2024. URL https: cpenreview net./ feet/forume forum? LLLs jgbssgg55990 aan.
A. K. Singh, T. Moskovitz, F. Hill, 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 Forty- first International Conference on Machine Learning, 2024. URL https: /ttp/xixivvvoo////oo///a org./iioorraaeooo abs/ 2404. 07129.
R. 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, Baldwin, A. Korhonen, K. Liveseu, and S. Bethard, editors, Proceedings of the 2013 Conference Empirical Methods in Natural Language Processing, pages 1631-1642, Seattle, Washington, USA, Oct. 2013, Association for Computational Linguistics. URL https: haw.anhhoootccccccccc org/ D13-1170/.
Song, Z. Xu, and Y. Zhong. Out-of-distribution generalization via composition: lens through induction heads in transformers. Proceedings of the National Academy of Sciences, 122(6): e2417182122, 2025. URL /
abs 2408. 09503.
E. Todd, M. I.. Li, A. S. Sharma, Mueller, B. C. Wallace, and Bau. Function vectors in large language models, 2024. URL https: /ttp//iivvvv.////////aaa///aaavv
15213.
13

---

H. 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 (at.:3302.137711 2023. URL https: // warxiv.org/ablll... abs////////// 2302. 13971.
S. Wang and C. Manning. Baselines and bigrams: Simple, good sentiment and topic classification. In H. Li, C.-Y. Lin, M. Osbome, 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: http://c...ttooooooooo//////////
14

---

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 contri- butions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of 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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

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.
Justification:
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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide 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 properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
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 this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
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 a specific model and empirical evaluation, 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 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 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 submis- sions 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 reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We have submitted the data and code the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines Ihttps://nips cc/
publ c /gui/guuoug des guides/ Cuudeu s oonoolicoo cy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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 NeurlPS code and data submission guidelines (https: / /nipss cc sacc/publiie hwwp///guusseeeeeeee dediiddes//eSSeeeee Code Submissbieeee ssiiisiinn cassoooolli (oooicyy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
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 the core the paper to a level 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, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 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 emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the csGuidel NeurlPS Code of Ethics https: /tt////euur neurips iss. ceiiiii iies?
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 reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration 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.
18

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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.
19

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw s swi thoode 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 encouraged to reach out the asset's s creators.
13. New assets
Question: Are 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 generated in Appendix and Appendix D.
Guidelines:
The answer NA means that the paper does not release 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. Crowdsoureing 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 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 contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, 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 involve crowdsourcing nor research with human subjects.
Depending 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.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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
Does the paper describe the usage of LLMs if 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, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: LLM is used generate one of the datasets used in the experiments. We introduce this dataset in detail in Appendix D and Appendix Apart from that, LLM is not related to any aspects of the core methodology in this paper.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non ntn-standard components.
Please refer to our LLM policy (https: //n////nipppe ips LLM) for what should should not be described.
2025/
21

---

Appendices
A Proofs related to Theorem 1
We first show that S is attained at some point v² in the hypersphere gd-1 = {u u E Rdd | = though Sd-1 is an infinite set and S(u) is not continuous. Observe that each h; in partitions Sd-1 into two hyperhemispheres. One is {u u E R ||^|| = 1, u"h; 2 0} and the other is {u tu:ue R. ||||== 1, l, h, < 0}. Hence, as long as the number of queries in the dataset n is finite, hy.... hn will partition @d-1 into a finite number of "cells". Each cell corresponds the elements in Sd-1 that satisfy certain sign sequence [sgn(u' ...,, (u" hn)] {1,0, -13". Hence, S(u) is constant on each 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 E R2xd and the two rows are B and Eyg. Then we have
E,,h,
Acc
Σ 1 arg max
(3)
Σ
E.yh
Σ
1 ( yB h Eygh)
+
E
1( < h,
definition of argmax
(4)
n (EN)
IEN2
1 n IEN)
Σ
(( y/
Byg hi 2 0
+ Σ
([
Ey h, 0]
reorganizing terms
(5)
iEN
Ey E
Ey
Eys
1[(
}h: 0]
1[(
hh < 0]
linearity
(6)
EEy,
Eyg ||2
EEy,
Eyg ll
< sup
Σ
1(u h: 2 0) + Σ
lu h: < 0))
supremum
MERd |M||a=1
INN
=
(8)
Here, when E,, h h, without loss of generality, we break the tie by letting the predicted label to be yA for x;.
E.,h E. Elsanil B.ghi
When there are more tokens in V, it holds true that 1(y; = arg max
S 1 =
Vi
h. E ). Hence, we have Acc= = E" 16x = arg max.
E, h,
arg
B,ch S S*
1yE1.....91
Vj
By the definition of u", we have
1( T h: 2 0)+
Σ
1(u h; < 0)) =5"
(9)
iER
IEN2
If Ev, Eve
cu" for some positive constant then
[yya
Eyalz
E[
-Ey,) Eyn) 2 0]+ Σ 1( - -E,z) < 0)))
(10)
n
ieN:
1
1 1 T h 2 0)
Σ ieN2
1 T h < 0)).
(11)
n
iNN
We already showed that
E, h
1 ( arg max
(12)
" E, Ebr E,,h:
Σ 1[(E ya Eyg) h; 2 0
Σ 1[ - ) < 0)].
(13)
EN
IENg
22

---

Vi Nj UN and When V E, < max h,, hi, hyg hi,,,, Vi, we have
Er,h: h,
1 arg
(14)
11
Σ vyEtnoot
B,
h:
E. hi
1 arg max 1i
(15)
Σ
E
hi
vyEl.....
Acc.
(16)
Evi Hence, we have Acc = 5" when 7V E,h; <, max (E,, Ey, h:), and
CM'
Eyall
B Generalization of Theorem
In this section", we consider generalization of Theorem where "y are allowed to have < labels, i.e. y1 E (yi),....)m) Vi. Same as before, let Nj = (i: yz = yj}, = I,.., m. j = I,..,m m are pairwise disjoint and ,N = nI. Let 5,a(4) = |मा (Lie,, 1(" h; 2 0) NNN, 1(u" h; < <)) and let d = arg west- Sja(u) for j.k E 1,...,, m, and Sju (ja(4 S For each j and 4),x is guaranteed be in gd-1 because only [N [Ml hyperplanes involved in separating separating Sd-1, resulting in smaller number of cells (guaranteed to be finite) compared the case of Xi, having only y2 as labels. Then, we have the following result.
Theorem Acc S m|m (M The equality is achieved when for all unordered pairs (j,k),jke I,...,,, max byyr E.h, < Ey, Eyg = C for some positive constant Cjk-
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 yj and y& for each kj,k). the dataset classification accuracy is then apper-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 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.
Here, we focus on the multiway classification, thus we change our notation to label categories from ya and yg into Ym, where 2t.
23

---

Proof First, note the following relations
R
E hi
Acc
1 arg max
Σ
E hi
1 arg
Σ
E,y
Σ
Σ
1(E h 2 Ey Vi 1....m))
JR Σ
Σ
min jel.jel....m
(E h; 2 h,)
indicator function is either 1 or 0
1x
Σ
Σ
E
1 1E h, 2 hi))
h:))
m1
J.I..JEl.....
M
Σ
Σ
hE > E,,hi))
n(m-
J.
M
yg Eyj
L
h: 2 0)).
n(m
E By,l12
Note that the sum is over all ordered pairs (j,k) s.t. j.k m.... + Consider the case of
(k = 1.j = 2) and (k 2, = 1), then both 1 (E h; 2 Ey,hi h:) and ( 2 E h:)) appear in the
sum. Hence, the sum can be taken over all unordered pairs (j,k) and rewritten as:
E y, -By, h; 2 0) + 1
E y2 E h: 0))
E
Σ
(1(
(17)
n(m
E,, E
E E ll
< n1 1)
Σ
Σ
[N
(18)
j=1+
Hence, we have
Acc S n(m 1)
(N&l + (N,1S,
(19)
1=1 j=++1
As proof of concept, plug-in = 2, then the bound becomes (INi + N21)5 = since Ni +
[N2l = n, which matches the result in Theorem 1.
To establish the equality conditions, as in the proof of Theorem 1, if Ey CjuM), for some positive constant CjA,, then
Σ IEN
E[
E,,)
7 h
2
0]+
Σ [( E,,) h <0]
(20)
IEN;
Σ
0) Σ
1 1 < 0) + N
(21)
IE
iEhj
We also know that if max cV E,h, < max Ey h E,, hi), Vie N1 UN, then
Σ 1[ E /Th 2 0 +
Σ 1 [[ E,, ' 0)]
(22)
IEN
IEN,
Eryh: h
Σ
16y
arg max "
(23)
Σ E(gen)) E. h
iENjUN
E,, h;
Σ JENUUN
1 arg max
(24)
Σ
E,y h
UyE wyEl.... I......V!
[V]
24

---

Finally, note that
E D Eyfei B,
Acc
1 arg
(25)
m
E Ei
[y: arg max Σ
(26)
B,y
HINNUNN
yyl.....VV
E,y
because for each Nj. the term 1 arg max for each E Nj. Hence, combining everything, we have
appears (m - 1) times in the summation
E
ryet....
m
m
E
Acc
E iERYJuN
1 arg max
(27)
Σ VyEl....VV
E
m
Σ
Σ
(NN
MMDS
(28)
m
j=++1
The proof is thus complete.
C Calculation of geometric measures
1. Separability score We train 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 the logistic classifier, and set the number of iterations to be 100.
2. Output alignment We feed H into the output unembedding layer of the model get the model's predictions of the query labels based on H and calculate the percentage of correct predictions.
3. Singular alignment First obtain the centered version of H, denoted H. Then calculate its singular value decomposition, namely H UEV", and compute the absolute values of cosine similarities between the top-r singular vectors and the difference of label unembedding vectors Ey, Eyg The absolute value is taken since the sign of the singular vectors is arbitrary. In practice, calculate the values using the top-2 singular vectors and report the maximum the two.
E,, E.g 4. Variance-based alignment First normalize Eу₁- - Eyo as [Ey, Eygli Ey E,g T
Then the ratio of variance
E, Evp. is (w.r.t. its mean) of H explained along [E,, Eyg12 TE E,gle n
E.g
EEn which
EE E
take to be the variance-based alignment of H.
5. Mean-based alignment First compute the means of label clusters (subsets of H with each of the two different labels), i.e. [NA]
Σ h; and
Σ h; Then take the difference and compute
[Ngl
iENA E. Eng EsA
iENN E E
the value of its projection onto
i.e.
Σ FA
h
Σ h:). Then, iEN Ey,
Eop
TE
Ezg
iE
Erg
calculate the weighted average of the variances the two label clusters along
i.e.
|E E
[9x] (arlVar((h)
Bya -Byg JeN D+ Vai(()] BB Eyp
Esy EEg
( Finally, divide the projected difference
[E, Eig
between means by the square root of the weighted average of the projected variance.
6. Composite alignment Multiply the variance-based alignment by the mean-based alignment.
7. Effective dimension Given the SVD UEY and the singular values o oi,... Gn the effective
o
dimension of H is 1
25

---

Table Prompt templates and labels for different datasets.
Dataset
Template
Label
SST-2
[Sentence] Semiimet:: (Label)
positive / negative
SUBI
[Sentencc]] Type: [Tabe")
subjective objective
TREC
The question is: (Premise]? The answer abbreviation entity / description / human/ location/ number
Question: [Sentence] Type: (Label)
SNLI
truc maybe false
(Hypothesis) (Label)
RTE
The question is: (Prenisef? (Hypothesis) (Label)
The answer is:
true/ / false
CB
The question is: [Prenise]? (Rypotheats) (Tabel)
The answer is:
true maybe false
Famous People
(Ferson Name) is (Iabel)
actor politician/ singer scientist / writer athlete
Table 3: Mappings used to replace ground-trute -ruth labels of different datasets to symbols
Dataset
Label Mapping
SST-2
negative/positive - @/#
SUBJ
objective/subjeetive - @/#
TREC
abbreviation entity/ description / person/ / number location/ -> @///1/S/8/*
SNLJ
true/mayhe/false /#/#/!
RTE
true/false +
CB
true/maybe/false
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-44 [OpenAI et al., 2024]. The dataset has 180 datapoints, each is the name of a famous person with profession label from one of six categories: actor, politician, singer, scientist, writer, athlete. The dataset balanced as there are data points for each label.
ICL setting For each dataset except for the Famous People dataset, we select demonstrations from the train and queries from the test set, or the validation set 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 the train has more than 10000 entries, and the first 1000 data points in the test set/validation 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 of the query in l2 distance for each query. We use prompt formats detailed in Table 2.
Devices All experiments with Llama2-7B, Llama2-13B, Gemma-2B, and Gemma-7B conducted with an A800 GPU. All experiments with Llama3-888 Llama2-70B, and Llama3-70B are conducted with an H200 GPU.
For the label mapping experiment Subsection 5.1. We use the mappings listed in Table 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
26

---

prepended to the labels. This implies that rather than direct decoding we can instead use the model as feature extractor with which 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 100% 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 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) 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 classificr-be based method in that it assumes no knowledge of the ground- ground-truth label of each hidden state. The datasct-averae results demonstrate that the rank-5 denoising can indeed improve the classification performance of all models with the only exception of Gemma-7B.
Supplementary materials for Subsection 5.1
F1 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. can be drawn, Alignment measures explain the accuracy difference between ICL and cero-shot across models, and the layer-wise trends of the measures explain the phase transition pattern for all models.
F.2 Replication Figure 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 ICI. 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 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 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.
27

---

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 label- related tokens in the hidden states (as demonstrated by subplot (B) of Figure 5) is present in other datasets 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 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 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 Table 6, and aligns with the point where output alignment begins 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 compete with semantically relevant task tokens, marking the onset of aligmment 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-trut -auuh labels. Similar results for SST-2 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 irrelevant ones, and the final layers do the opposite. We show in Table 1 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 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.
the innitial layers after the phase transition. This confirms the
In Figure 43-48, we present results for other models. These are consistent with the findings from Llama2-700 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 natural language inference tasks tasks-likely because such tasks require reasoning over both 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 rank-r SVD approxi- mations of Llama2-70B hidden states across all datasets. Specifically, we investigate how increasingly
28

---

aggressive rank reductions- 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 some- times leads 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= I 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- digee--dmmensional 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 IHs and PTHs. Denote the queries as X5. each with token length s(x1).....s(ss))
Identification of PTHs For each X; the LLM will generate an attention tensor Attn; E [=xx where L is the total number of layers, N₂ is the number of attention heads per layer. Hence, for the nhh attention head layer I, there is corresponding attention matrix Ar)m(1,ng):= E (x)(())) where Atrn is the attention of this attention head at the 1Th token to ((x)
the jth token in xy. The PTH score of this attention head on x; is thus defined as Li= Attn (1, 1-1- which measures the total attention the attention head assigns at each token position to the previ-
ous token. Consequently, the PTH score of the ni attention head at layer on all the queries is
10 E Attn(L, (n)n((,nn)is-11 We calculate the PTH scores for all (1, pairs and choose the top 10% attention heads as the identified PTHS.
Identification of IHS For each x;, randomly select 8 demonstrations and prepare them to The resultant ICL prompt (denoted as X; with length (X;)), with slight abuse of notation, takes the format of (x,1) (yr.)).... (x/s) (vv.s), (xx) i, where (x,x) represents the sentence part of demonstration k ("I like this movie. Sentiment") and (Vix) the label ("positive") separated from the sentence part by colon (:), and (x²) 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 (vi.1).... (vis). Let Ti denote the indices of the label tokens in X,. The IH score of the ni, attention head at layer over the 50 queries is thus 50 Exel, E Attn(I, m, (x) i.e. we are measuring the total attention head assigns at the final Li=1
position to the positions of all the label tokens summed over all 50 queries. We calculate the IH
scores for all (l, pairs and choose the top 10% attention heads the identified IHs.
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 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.
29

---

I.3 Replication of Table 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 thhnne 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.
1.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 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.
1.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 al. [2022], where none of the in-context demonstrations share the query's label. This setting is useful for analyzing attention known for copy-like behavior-whis direct label copying is not possible.
especially IHs, which are
To this end, we construct a Famous People dataset following the procedures in Appendix D. Each data point is 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 query or a demonstration. For each query, we sample k demonstrations, ensuring that none share the query's 's profession label. We then ablate PTHs and IHs 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 '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 im- pairs 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 oppoottt to the standard case, where IH ablation has 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 further confirmed by the substantial accuracy gain observed when the query label is forced into the context.
We provide results analogous Figure 35 for other models in Figure 36-41, which demonstrate similar results.
1.6 Results of attention heads ablation in the zero-shot setting
In this section, we investigate the effeect 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 stantially 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 P'THH as key structural components of L.MMs.
demonstrations from which to copy. On the contrary, ablating PTHs sub-
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,
30

---

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 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 caleulating 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 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 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 surponnntsssucch Pre- vious Token Heads and Induction Heads in shaping model behavior, our study brings interpretability to 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 model is failing to align internal representations with task-relevant signals. Additionally, the -VD-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 privacy-sensitee settings where retraining is infeasible.
However, deeper understanding of model internals also introduces risks. The ability to precisely manipulate hidden states or attention mechanisms may enable adversarial behaviors, such as con- structing prompts that selectively suppress or amplify certain outputs for political or financial gain. Moreover, techniques for filtering and steering semantic content could 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 real-world systems.
Overall, we view our contributions as 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.
31

---

Llama3-8B: 8-shot Ace.: 76.20%; 0-shot Acc.: 0.30%
Effective dimension
Separability
Singular alignment
40
0.2
0.8
20
0.1
0.6
0.0
10
20
30
0
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Composite alignment
Mean- based alignment
Var.baased alignment
0.02
2
0.010
B-shot 0-shot
1
0.005
0.00
0.000
10
30
20
30
10
20
Layer Number
Layer Number
Layer Number
Figure 10: Comparison of trends in separability and alignment measures of Llama3-8B hidden states
between ICL and zero-shot
Llama3-70B: 8-shot Acc.: 81.09%; 0-shot Acc.: 0.00%
Effective dimension
Separability score
Singular alignment
0.8
0.10
20
0.6
0.05
Layer Number
Layer Number
Number
Composite alignment
Mlean-aaeed alignment
Var-based alignment
0.02
2
0.005
0-shot
0
0.000
80
20
Layer Number
Layer Number
Number
Figure 11: Comparison of trends in separability and alignment measures of Llama3-7008 hidden states between ICL and zero-shot
Llama2-7B: 8-shot Acc.: 75.13%; 0-shot Acc.: 0.31%
Effective dimension
Separability score
Singular alignment
40
0.8
0.2
20
0.6
0.1
10
20
30
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Composite alignment
Mean based alignment
Var-based aligmment
2
B-shot 0-shot
0.01
0.02
0.00
0.00
0
10
20
30
0
10
20
30
10
20
30
Layer Number
Layer Number
Layer Numher
Figure 12: Comparison of trends in separability and alignment measures of Llama2-7B hidden states between ICL and zero-shot
32

---

Llama2-132 8-shot Acc.: 74.79%; 0-shot Acc.: 0.36%
Effective dimension
Separability
Singular alignment
40
0.8
0.2
20
0.6
0.1
10
20
30
40
10
20
30
40
10
20
30
40
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-bases alignment
Var-baased aligmment
0.04
2
B-shot 0-shot
0.01
0.00
0.00
10
30
40
30
40
20
Layer Number
Layer Number
Layer Number
Figure 13: Comparison of trends separability and alignment measures of Llama2-13B hidden states between ICL and zero-shot
Gemma-2B: 8-shot Acc.: 66.16%; 0-shot Acc.: 0.00%
Effective dimension
Separability score
Singular alignment
0.8
20
0.6
0.1
13
10
10
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var.baase alignment
0.02
8-shot 0-shot
0.01
0.00
0.00
0
10
15
10
15
10
15
Layer Number
Layer Number
Layer Number
Figure 14: Comparison of trends in separability and alignment measures of Gemma-2B hidden states between ICL and zero-shot
Gemma-7B: 8-shot Acc.: 79.06%; 0-shot Acc.: 0.02%
Effective dimension
Separability score
Singular alignment
20
0.0
Layer Number
Layer Number
Layer Number
Composite alignment
Mean based alignment
Var. based alignment
0.050
8-shot 0-shot
0.025
0.01
0.000
0.00
10
20
10
20
10
20
Layer Number
Layer Number
Layer Number
Figure 15: Comparison of trends in separability and alignment measures of Gemma-7B hidden states between ICL and zero-shot
33

---

Limma3 80: Average scrous datasets
Soparabgg bo soore
Stnggi aliggemen entt
Composste ailpmmont
Acturasy
0.02
12
16
Number
Layer Number
Layer Number
Number
(a) Varying number of demonstrations
1lame 8.: Random Solection Acc.. 78.2 78.20%, KNN Selection Acc.: 84.03%
Llama3 00: Normal Label Acc.: 76.20%, Symbol Label Acc.: 57.05%
Effective cimension
Separability score
Siagular alignment
Effectiva dimenssoo
Separablity soore
Singulla aligamenn
0.8
0.3
0.6
0
20 Layer Number
20 Layer Number
20 Layer Number
20 Tayer Number
20 Laywer Number
20 Lager Number
Composite
alignment
Var-hased
Composlte
Meat-based alignment:
Var-baased aligrament
0.02
0.00
0.00
0.010
0.01
4
20
Layer Number
Layyr
Layur
Laysr Number
Luper
Handom Selection
K&N Seboction
Labsi
Symbol Label
(b) KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 16: Layer-wist wise trends of separability and alignment measures of Llama3-88 hidden states in different ICI. settings
LJama3 70B: Average across datasets
68-
Scarrbbity scure
6:1.0
6.225
Bigumest
0.033
0.025
Accurory
6.300
0.33
0.
0.00
Layer
Layer
Layer
Layer
(a) Varying number of demonstrations
1lama 3- 70B: Random Selection Act.: 81, 09%, KNN Selection Acc.: 89 49%- Exteetivee dimension
Liamaa: 70B: Nocmal Label Acc.; 81 090%%; Symbol Label Acc.: 64.67%
Separability souce
Singrular alignment
Effeccive dimensiion
Sepembility seare Singular allgmment
0.15
3.6
0.8
20
3.6
0.05
20
0.6
5.3
50
50
50
Lopee Nambee
Layer Number
Layer Number
Layer Number
Loper Nunther
Loyee Mumber
aligmment
Mean
Mar-bbaeed
Composite alignment
Meads-hssed alignment
Var.
alignment
005
50
Numberr
Number
Rondom Selection
KNN Seiection
Neemal Exbel
Symboi Label
KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 17: Layer-wise trends of separability and alignment measures of Llama3-700 70B hidden states in different ICL settings
34

---

Llmma 70: Average scrous datasets
Separabilis scure
Singnlie llggmmment
Compootto alsmmonnt
Accurssy
8.15
Estroctive
32
16
12
20
10
Number
Layer Number
Layer Number
Number
(a) Varying number of demonstrations
1.7m2278: Random Selection Arc:: KNN Salection Anc.: 84.07%
Llama2 7B: Normal Label Act.: 75. 13%; Symbel Label Acc.: 44.73%
Effective diinension
Separability scoce
Siagular alignment
Effecctiee dimension
Separablitty score
Singular allgnament
0.6
60
0.8
20
0.1
20
20
20 Layer Number
20
20 Layer Northe
20 Tayer Number
Layer Numthe
Layer Number
Layer Number
aligement
for-hased
Composito alignment
Mean-baased aligament
Var-hassen alignmment
Jumenn
Layer
Lager
Layer
Layer
Layer
Layer
Random Selection
XNV Sebection
Norrral Label
Symbol Label
(b) KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 18: Layer- wise trends of separability and alignment measures of Llama2-7B hidden states in
different ICI. settings
Llama2 13B: Average across datasets
0.8
Separabiity 10000
Sirguler allegrmmn Biggmmak
Compootti Elgnnmmnt
10.83
0.04
Infrectrve
12
8--
92
Number
Number
Number
(a) Varying number of demonstrations
1lama 2- 13B: Random Selection Anc.: 74, 79%, KNN Seiection Acc. 86, 18% Effective dimensioo
1lama2-13B: Normal Label Acc.: 74.79%, Symbol Label Ace.: 51.78%.
Sepambility scote
Singular alig Lignment.
Effecttive dimesssoon
Separability soore
Singular alignmmnt
3.8
0.
0.2
20
0.1
0-
20
40
20 Layer Nurther
20
40
20 Lague Number
20 Layse Number
Layer Sumber
Lopec Number
Layur Number
Composite allenment
algnment
Var-bassdd
Composite aligmmmnt
Var-based
5.02
0.04
n00
6.01
0.02
Random Sclection
KSN Sclection
Normal Label
Sysabol Label
(b) KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 19: Layer-wiss wise trends of separability and alignment measures of Llama2- 13B hidden states in different ICL settings
35

---

Gemma 2.: Average arross datasets
Saparaaae abdtt soore
Stnggi alrggem entt
Compostte alipmment
Accurasy
0:20
10.33
8:15
0.12
16
Number
Layer Number
Layer Number
Number
(a) Varying number of demonstrations
Gemma 2D: Random Selection Acc.: 66.16% KNN Selection. Acc.. 77.99%
Gemman Normal Label Acc.: 66.16%; Symbol Label Apc.: 40. 00.066%
E:focttive dimenssson
Separabiltty score
Singular alignment
Effective dimension
Separabilitt score
Singular aligameen
0.5
20
0.2
20
0.2
0.3
10
10 Layer Number
10
10 Laper Number
10 Layer Number
Luyet Number
Layer Number
Lapee Number
Composto aligament
algunent
alggument
Composite aligument
alignment
0.02
0.080
0.02
0.02
0.01
0.01
0.01
0.00
0.00
10
10
10
Layse Number
Layer Number
Layer Number
Lypee Mumber
Laper Number
Layer Number
Random Selection
INN Seiection
Normal Label
Synbbol Label
(b) KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 20: Layer-wise trends of separability and alignment measures of Gemma-20 hidden states in different ICL settings
Commo-77 7B: across datnsets
0.02
Couuaan score
Singulor aligemmnt
Aeturany
22-
0.22
Layer
Luyur
Layer Numbor
Layer
(a) Varying number of demonstrations
Gemma- 7B: Random Selection Acc.; 79.06%; KNN Selection Aoc.: 86.23% Fffective dimension
7B: Normal Label 79.06%, Symbol
Separiiility
Singuler alignment
Effective dimension
Separability
Singular alignment
0.8
0.6
0.6
20
Layer Number
Lapec Namber
Lapec Nanber
Number
Number
Number
Composlte
alignment
Var-based alignment
Composite alignment.
alignment
alignment
p.
0.02
3.05
Nunber
Namber
Random Selection
XNN Selection
Normal Label
Symbell Label
(b) KNN Demonstration selection
(c) Random symbols as demonstration labels
Figure 21: Layer- wise trends of separability and alignment measures of Gemma- 7B hidden states in different ICL settings
36

---

Cutgut Aigustont
Aliamn.er
5.20
Singuae Alignnmn
Cutgut Al gesttnn
10.055
Compaiii Alqnmma
62.
5.10
10.020
6.
5:10
5.00
0.000
40
(a) SUBJ
(b) TREC
0,020
6.8
Outgus Alqunoit
Singuler Rliggn Alignnent
Curput Alpument
2.13
8.2-
5.20
2.02
All
2.10
622
6.0
2.00
0.3
6:0
9.000
6:3
20
60
82
40
60
82
25
55
20
40
82
Lager Number
Layer Number
Layer Namber
Lager Number
(c) SNLI
(d) RTE
Klgnmantt
0.8
Compue all Alignmece
Output Nignmatt
0.0
Alggimett
0.c
0.c
0.2
caaca
cacoo
50
20
00
Layer Number
Lagar Number
(c) CB
Figure 22: Output and directional alignment of Llama2-70B ICL hidden states with label unembed- ding vectors on various datasets.
Table 4: Tokens decoded from top 10 singular directions from layers near the phase transition on SUBJ
Layer
Tokens
Layer 51 Layer
opinion, humor, genre, adm, facts, plot, description, explanation, pont, humor objective, humor.... adm, dram, cutbus, categor, pont, subject objective, humor, plot, fact, humor, enthus, explanation, character, hamor
Layer 53
Table 5: Tokens decoded from top 10 singular directions from layers near the phase transition TREC
Layer
Tokens
Layer
Ras, Flost, Physics. Heimat, Zob. rela, ijd. Ibl, alphabet, uther Geography, name, Physics, map. >, Elc, location, Vall, alphabet, Tests definitions, maps, Geography. amos, map. babel.
Layer 19
Layer
Table Tokens decoded from top 10 singular directions of layers near the phase transition on SNLI
Layer
Tokens
Layer 24 Layer
imagination, inen, exceptions, hner, igin, amerikanischer, Fich, float, upposite, orf dup, unknown, similarity, Dy. etc, nost, Fich, endorf, Correct erde, unknown, similarity. context, apper, architecture, unity, ambigu. py
Layer
Table 7: Tokens decoded from top 10 singular directions from layers near the phase transition RTE
Layer
Tokens
Layer 21 Layer
Eva, pul, ikel, orp. Mol, cak, Dal, adrat, hor, Ung correct, nor, aqu, Dal. subscribe, currency, sign correction, zero, tempo, Bob, VP. lak, FA, iva, currency, trag
Layer 23
37

---

Outbut Algrmment
+0 006
Oclput Alggument
1.08
Singulan Alignmet
Octput Altrmmnt
8.S
Outbut Algensnn
0.004.
St-
Complaite
0,002
62
6.-
25
22
Number
Number
(a) SUBJ
(b) TREC
0.025
Outgut Algemment
Singule Alggnmal
Octpuz Alganent
0.020
2.25
A20
2.20
64
0.025
0.2
+0 010
623
2.10
00.
6:0
0.000
633
0.0
0.000
00
25
15
23
10
25
15
23
30
(c) SNLI
(d) RTE
Outpet Migemmntt
-0 01
Klqument
0.2%
Output Alcimenn
2.00
0.8
3.25
3,04
19. 19.020
8.t.
5.15
Compostter
9.15
6.22
9.ots
6.2
5,00
5.20
0,000
0:3
20
25
10
25
15
20
50
15
20
30
Layer Number
Layer Number
Loyar Number
Layer Number
(e) CB
(f) SST-2
Figure 23: Output and directional alignment of Llama3-8B ICL hidden states with label unembedding vectors on various datasets.
Table 8: Tokens decoded from top singular directions from layers near the phase transition on CB
Layer
Tokens
Layer
TR. qual, Joan Icon, ebol, par,, Kontrola
Layer
deg, Cet, Boa, SV, ando, Simon, True. lav, ih81, opposition Layer negative, schop. ftrag, court, ando, inda, rib, Marí, counter,
Table 9: Tokens decoded from top 10 singular directions with labels replaced on SNLI
Layer
Tokens
Layer 23
Tier, aro, cum, Vector, FB, Eur, Fon, pending. Dre. Fant False, aten, beck, ardo, Sti, attan, Kaiser, Fant False, Cop, cpi. Dou, 15, Ra, hur. TF, Dre. nchm
Layer
Layer 25
Layer
false, combination, inois, multiply, scenario, group, action, group, prec, Solo False, H, xspace, @ neutral, @ large, sentiment, False False, conte,
Layer 45
Layer
neutral, cheer, group, ppo, monol
Layer
false, False, $., @ @, false, S., false, False, false, True, soo, twe, X, 1, :R, 4it, action, 1, ft Street
Layer
#f
Layer 80
38

---

10,000
Outbut Alennoen
.0.000
Alyneser
6.x
Singlar Algimmet
Outpet Algmimest
Composde Alignmes
B.c
Rdt.
5.06
6.2
0.002
19.000
30
25
40
82
Layor OF Number
Layer Number
Lay or Number
Layer Number
(a) SUBJ
(b) TREC
AZS
5.20
Qutput Algnnnest
Ouiput Abqnneent
Cutput Alenmeet
Algnnset
6.6-
too
Bagnn
ans
0.0075
0L0020
Comptaate
R.30
6:2-
822-
0.025
6.2-
0.0022
6.2
300
Lay or Number
Layer Number
Layer Numbor
Lager Number or
(c) SNLI
(d) RTE
Octtut Altemont
Gelput Algement
Singulan Alggmmaa
Output Alghonent
Allnmna
9.25
0.--
B.0-
5.20
0,0055
Compuaille
But-
0022
0025
6.2
500
5.00
42
Number
Number
CB
(f) SST-2
Figure 24: Output and directional alignment of Llama3-70B ICL hidden states with label unembed- ding vectors various datasets.
Table 10: Tokens decoded from top 10 singular directions with labels replaced on SST-2
Layer
Tokens
Layer
poste, Cant, Jug. Lars, >, zak, zig, geh, eg. alus roll, igi, endi, noreferrer, positive, zak, Aur. Mans, humor rott, Armen, Lau, noreferrer, negative, mee, cot, tu. fici
Layer 25
Layer 26
Layer 40
ipage, pra, Klc, estaven, Maj, CURL .OPT, extreme, hamor, G, anel ViewByld, pro, meno, Einzelnach, ziel, @, disappoint. Meyer, bumor negative, kinzelnach, #. humor, extreme, humor, embargo, #
Layer
Layer
Layer
negalive, noum, elt,
hyper. $., <0,, positive, inea, positive,, (1. *,, neutnal, Tru, ???, <so, 7, humor, Sent il positive, :, disappoint, positive, about, positive, wed,
Layer 79
Layer
Table 11: Tokens decoded from top retained and filtered directions in layer updates on SUBJ
Retained Tokens
Stage
Layer
Tokens
Stage
Layer Layer 59
humor, description, explanation. objective,
shock, subject, character, comparison,
burnor, plot, pra, facts, objective, hypoth, subject, met, explanation, subject objective, explan, prom, subject. character, objective, sci plot, observation, IM, Thom, verb, hyper, quar, Em, ps. Guillaume
Layer
Stage
Layer 80
Filtered Tokens
Stage
Layer
Tokens
Stage
Layer 54 Layer
oro, Kant, Harr, Werner, Bast, dosert, Kent, Uns, ola, Fall iter, Arg. hulten, enser, Kontrola, tring. atif, ias. Sof conj, dens, shal, amo, Lisa, lish, Soh, narrow, UMN, subject, Met, chron, intros, mix, emot, humor, Type, joctive, subject
Layer 66
Stage
Layer
39

---

Output Aigunoon
6a.
Singuler Algnniii
Cutput All guauuuee
0,020
Composta Algnima
4.18
8:16
5.20
+0 930
63
Output.
R
5.08
20
Layer
Layer
(a) SUBJ
(b) TREC
2.04
5.25
2.30
Cutyet Algmmontt
Curput Altumment
020
9.25
6:15
5.20
0.8
8.2
0:2
6.05
6.0
2.00
0.3
100
600
0.000
6:9
10
30
15
20
10
25
15
23
30
Layer Number
Layer Number
Lager Number
Layer Number
(c) SNLI
(d) RTE
Outout Alemtonn
2.13
Alqunsont
Cutput Nimmment
3,000
Xlemmin
0.3
5.20
8.15
6.8
Hi
2.04
Composd
pick
15
20
50
2.00
0.3
12
25
5.00
12
20
50
15
20
Layer Number
Lager Number
Lagar Number
Layer Number
(c) CB
(f) SST-2
Figure 25: Output and directional aligmment of Llama2-7B ICL hidden states with label unembedding vectors on various datasets.
Table 12: Tokens decoded from top retained and filtered directions in layer updates on TREC
Retained Tokens
Stage
Layer
Tokens
Stage
Layer 30
pione, taste, definition. temporal. procision, science, amen, lich, irectory, eata deseription, subst, person. numbers, interpret, location, location, date may, date, number, location, abbre, currency, location, number, description, Creative hol, Heidel, Desp, eth, olis, deseription, origin, Type, description, arto
Layer 50
Layer 70
Stage
Layer 80
Filtered Tokens
Stage
Layer
Tokens
Stage
Layer 30
ako, do, amen, ainer, overlay, aar, CHAP, storm. synchron, tober Gordon, Chinese, acci, Grace, penas, zak, gas, Kii, ps, prototype
Layer 50
Layer 70
lit, iomi, Sall, ätre, quet, thru, Sever number, list, discipline, organ, entertain, weather, hum, poll, coff
Stage
Layer 80
Table 13: Tokens decoded from top retained and filtered directions in layer updates on SNLI
Retained Tokens
Stuge
Layer
Stage
Layer 25
pac, (cio, charm, beck, hist, ma, iella, fact, contradiction, engol
Layer 50
probable, false, uncertain, solo, neutral, sender, unknown, gender, unlikely reverse, false, scenario. be, reverse, identity, gender, conce, Gram Dic,
Layer 70
Stage
Layer
Column, Commnn,, motor, .... ge, Sali, kern
Filtered Tokens
Stuge
Layer
Tokens
Stage
Layer 25
hrief, uclide, Ath, hul, références, FCC, Rot, nia, aho, obn, zu, umble, prepared, Mey, jer, Bras, irre, Maj
Taml
Layer 50
Layer 70
Yu, Mare, iner, artificial. overlap. ilder, aret, elev, constru, Sterm im, 
. Il, maybe. similarity, sem, habit, Tá, gender, action
Stage
Layer 80
40

---

0.012
Compuiile Alqmmmnt
Output Aigimeen
M
-0.00 010
Gutput Klgaasent
B.c.
5:15
Singuler Alignmni
Cutput Alpustunnt
10.20
Output Algumantt
1.126
0.008
00.00
10.022
6.1
0.000
63
Layer
Layer
Layer Number
(a) SUBJ
(b) TREC
Cutgus Algunont
Alqmzsnnt
1:25
5.20 Sangene
Outpuz Algunment
19,025
Allgunont
5.02
16
Compaster
6.2-
3,00
6.2
0.020
60
2.00
6.3
60-
0.000
00.
20
25
30
40
20
50
42
20
25
40
Luger Number
Lager Number
Lager Number
Lager Number
(c) SNLI
(d) RTE
Outyut Algmonn
Output Alguussnt
19.550
Cutput Almment
Domprasten Alinmmni
Alemmmnt
0,020
0.125
Ra
0.3
0.020
64
<0 ocs
0,000
0.022
0.000
03:
2.50
0:0
20
46
10
2:
30
20
25
30
40
Layer Number
Layar Number
Eagar Number
Layar Number
(c) CB
(f) SST-2
Figure 26: Output and directional alignment of Llama2-13B ICL hidden states with label unembed- ding vectors on various datasets.
Table 14: Tokens decoded from top retained and filtered directions in layer updates on RTE
Retained Tokens
Stage
Layer
Tokens
Stage
Layer 33
Wikipedia, disput, yes, Jakoh, alias, essel, indirect, False, false, alann, irrelevant, opposite, nost, naming, duplicates, misunder, Geography, location
Layer 50
Layer 68
X Geography, lament, sin, impossible, opinion, ken, numer, negative, true
Stage
Layer 80
Alc, Ad, Hob, Ke, lub, Terra, Har, Th, part, Andre
Filtered Tokens
Stage
Layer
Tokens
Stage
Layer 33
cust, Mün, Lucas, itel. Barb, inha, Hed, Fl, nica
Layer 50
noni, CA, Um, wob,, prot, jür, 10, ang, *6, Roman fert, CA, bach, Word, RL, gon, Charlic, MS, nen, Feld SR. tochni, gr. unclear. trag. opposite, Susan, estimate, rare, murder
Layer 68
Stage
Layer 80
Table 15: Tokens decoded from top retained and filtered directions in layer updates on CB
Retained Tokens
Stuge
Layer
Tokens
Stage
Layer 25
Bou, True, <os, Cet, ando, elin, nom, Pear, stabil, lav
Layer 49
true, unlikely. uncertain, absence, False, predictions, avoided, myth, prodictions, asym
Layer 70 Layer
true, false, agreement, may be, predictions, silent, negative, disapp, äv, predictions Giav, future, often, Frank, tend, imet, Sid, fran, Intent, maybe
Stage
Filtered Tokens
Stage
Layer
Tokens
Stage
Layer 25
slant, acci, I Laurent, cht, brie, chamber, olare, gart, vie Dra, cito, nick, Verm, moz, neglect, abel, WR, anger, uchar ottor, nation. Cooper, exhaust, reign, hem. Um. diag. wa false, esi,, False, Yes, often, aer, Mod. George
Layer 49
Layer 70
Stage
Layer 80
41

---

6.-
175
Octout Alignnett
Octgut Aligmnct
Outbut Allanoat
-0.150
6.20
10.125
0.013
6.3
6.2
6.2
62-
0.05
00
0.000
60
0.9
10
10
(a) SUBJ
(b) TREC
Octout Allnnoot
0.010
Allmnen
Sing lae Alginmen
Octpuz Alguneat
822
6.3
2.20
10 000
10
0.000
10
0.0
10 10
10
16
Layer Number
Layer Number
Lay oF Number
Layer Number
(c) SNLI
(d) RTE
Octout Abannent
Alqunent
0.84
5.20
Output Abpunscn
Alamment
-0 020
5.25
2,04
6:5x
0.005
0.2
Layer Number
Laver Numbor
Laver Number
Layer Number
(e) CB
(1) SST-2
Figure 27: Output and directional alignment of Gemma-21 2B ICL hidden states with label unembedding vectors on various datasets.
42

---

Outuut Alammont
Alqemeent
6.3-
Adgimme
Ousput Aliameat
4.34
5.20
2.02
Compualte
2,C3
Output
022
5.00
2.00
2.50
+0.000
10
20
10
2M
10
23
Layer Number
Layer Number
Lager ber
Layer Number
(a) SUBJ
(b) TREC
2.30
Output Algentont
Oupput Alimmeen
1.26
Output Almuneet
20.3
3.04
6:1
1..00
$.15
0.2
0.3
0.10
62
3.02
0.2
5.00
Layer Number
Layer Number
Layer Number
Lager Number
(c) SNLI
(d) RTE
Curput Alpmmtntt
Output Alpentonn
2.15
Algnnmez
2.03
2,20
Cat
3.02
6.
CA
82
20
20
Number
Number
Numbor
Number
CB
(i) SST-2
Figure 28: Output and directional alignment of Gemma-78 7B ICL hidden states with label unembedding vectors on various datasets.
Llama3-8B, average datasets
Normal Ace.: 76.20% w/o PTH Acc.: 31.49% w/o IH Acc.: 0.00% Random Heads Ablated Ace.: 47.86%
Effective dimension
Separability score
Singular alignment
0.2
100
0.8
0.1
0.6
0
Layer Number
Layer Number
Number
Composite alignment
Mean-based alignment
Var.based alignment
0.02
0.010
0.01
0.005
0.000
10
30
20
10
30
Layer Number
Layer Number
Layer Number
Normal
w/o PTH
Random Heads Ablated
Figure 29: Effects of attention heads ablation the layer-wise separability and alignment measures of Llama3-8B hidden states averaged across all datasets
43

---

Llama3-70B, average across datasets Normal Acc.: 81.09% w/o PTH Acc.: 26.89% w/o TH Acc.: 0.92% Random Heads Ablated Acc.: 77.89%
Effective dimension
Separability score
Singular alignment
200
0.1
100
0.6
0.0
20
60
80
60
80
20
40
60
80
Layor Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var-based alignment
0.02
0.005
20
40
60
80
20
40
60
80
20
40
60
80
Layer Number
Layer Number
Layer Number
Normal
Ablated
Figure 30: Effects of attention heads ablation the wise separability and alignment measures of Llama3-708 hidden states averaged across all datasets
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%
Effective dimension
Separability score
Singular alignment
0.2
0.1
0.0
10
20
30
20
10
30
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var.bbased alignment
0.00
10
20
30
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Normal
PTH
Random Heads Ablated
Figure 31: Effects of attention heads ablation the layer-wise separability and alignment measures of Llama2-7B hidden states averaged across all datasets
44

---

Llama2-13B, average across datasets Normal Acc.: 74.79% w/o PTH Acc.; 23.19% w/o TH Acc.: 1.16% Random Heads Ablated Acc.: 73.49%
Effective dimension
Separability score
Singular alignment
50
0.2
25
0.1
0.0
10
30
40
20
30
20
40
Layor Number
Layer Numbor
Layer Number
Composite alignment
Mean-based alignment
Var-based alignment
0.04
10
20
30
40
10
20
30
40
10
20
30
40
Layer Number
Layer Number
Layer Number
Normal
Ablated
Figure 32: Effects of attention heads ablation the wise separability and alignment measures of Llama2- 13B hidden states averaged across all datasets
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%
Effective dimension
Separability score
Singular alignment
0.8
15
15
10
15
Layer Number
Layer Number
Layer Number
Composite liggment
Mean-basee alignment
Var.baased alignment
0.02
0.01
0.01
0.00
5
10
15
0
10
15
0
10
15
Layer Number
Layer Number
Layer Number
Normal
PTH
Random Heads Ablated
Figure 33: Effects of attention heads ablation the layer-wise separability and alignment measures of Gemma-2B hidden states averaged across all datasets
45

---

Gemma-7B, average across datasets
Normal Ace.: 79.06% w/o PTH Acc.: 25.60% w/o IH Acc.: 5.66% Random Heads Ablated Ace.: 40.82%
Effective dimension
Separability score
Singular alignment
50
0.8
0.2
0.6
0.0
10
20
10
20
10
20
Layer Number
Layer
Lagor Number
Composite alignment
Mean-based alignment
Var. bas-e alignment
0.050
2
0.02
0.01
0.000
0.00
20
10
20
10
20
Layer Number
Layer Number
Layer Number
Normal
w/o PTH
w/o IH
Random Heads Ablated
Figure 34: Effects of attention heads ablation on the layer-wise separability and alignment measures of Gemma-7B hidden states averaged across all datasets
Llama2-70B Famous
Unseen label Acc.: 57.78% | w/o PTH Acc.: 1.11% w/o 1H Ace.: 37.78% Seen label Acc.: 78.89%
Effective dimension
Separability score
Singular alignment
20
20
80
60
80
20
40
60
Layer Number
Layer
Layer
Composite alignment
Mean-based alignment
Var.b aligmment
5
0.004
0.002
0.00
0
20
60
80
60
80
20
40
60
Layer Number
Layer Number
Layer Number
Unseen label
w/o PTH
w/o IH
Seen label
Figure 35: Effects of attention heads ablation on the wise separability and alignment measures of Llama2-70B hidden states on Famous People dataset.
46

---

Llama3-88 Pamous
Unseen label Acc.: 68.89% w/o PTH Acc.: 15.56% w/o IH Acc.: 6.67%
Seen label Acc.: 83.33%
Effective dimension
Separability score
Singular alignment
1.0
0.10
20
0.5
0.05
10
20
30
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var. based alignment
0.02
5.0
0.01
2.5
0.002
0.00
0.000
0
10
20
30
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Unseen label
w/o FTH
Seen label
Figure 36: Effects of attention heads ablation the layer-wise wise separability and alignment measures of Llama3-8B hidden states on Famous People dataset.
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%
Effective dimension
Separability score
Singular alignment
40
0.050
20
0.025
20
40
60
80
40
80
20
40
60
80
Layer Number
Layer Number
Layer Number
Composite alignment
Mean- basad alignment
Var.- har.-based alignment
0.002
0.005
0.001
0.000
0.000
20
40
60
80
20
40
60
80
20
40
60
80
Layer Number
Layer Number
Layer Number
Unseen label
w/o PTH
w/o
Seon label
Figure 37: Effects of attention heads ablation the layer-wiss wise separability and alignment measures of Llama3-78 70B hidden states on Famous People dataset.
47

---

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%
Effective dimension
Separability score
Singular alignment
0.15
40
0.10
20
0.05
10
30
20
Layer Number
Layer Number
Layer Number
Composite alignment
bosan-based alignment
Var. based alignment
0.0050
0.02
0.0025
0.00
0.0000
10
20
30
10
20
30
10
20
30
Layer Number
Layer Number
Layer Number
Unseen label
w/o PTH
w/o
Seen label
Figure 38: Effects of attention heads ablation on the wise separability and alignment measures of Llama2-7B hidden states Famous People dataset.
Llama2-13B Famous
Unseen label Acc.: 52.22% w/o PTH Acc.: 0.00% IH 0.00%
Seen label Acc.: 67.78%
Effective dimension
Separability
Singular alignment
1.0
0.10
0.05
0-
Layer Number
Layer Number
Number
Composite alignment
Mean based alignment
Var based aligmment
5.0
0.0050
2.5
0.01
0.0025
0.0
0.00
0.0000
0
10
20
30
40
10
20
30
40
10
20
30
40
Layer Number
Layer Number
Layer Number
Unseen label
PTH
TH
Seen label
Figure 39: Effects of attention heads ablation the layer-wise separability and alignment measures of Llama2-133 hidden states on Famous People dataset.
48

---

Gemma Famous
Unseen label Acc.: 60.00% w/o PTH Acc.: 0.00% w/o IH Acc.: 0.00%
Seen label Acc.: 73.33%
Effective dimension
Separability score
Singular alignment
1.0
0.10
0.5
0.05
15
10
10
15
Layer Number
Layer Number
Layer Number
Composite alignment
Mean-based alignment
Var-based alignment
0.02
5
0.004
0.002
0.00
0.000
15
10
10
Layer Number
Layer Number
Layer Number
Unseen lahel
w/o PTH
w/o TH
Seen label
Figure 40: Effects of attention heads ablation on the layer-wise separability and alignment measures of Gemma-2B hidden states on Famous People dataset.
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%
Effective dimension
Separability score
Singular alignment
20
0.10
0.5
0.05
10
20
10
20
10
20
Layer Number
Layer Number
Layer Number
Composite alignment
Mean basan alignment
Var. based alignment
5.0
0.0050
0.02
2.5
0.0025
0.00
0.0000
10
20
Layer Number
Layer Number
Layer Number
Unseen label
w/o PTH
w/o IH
Seen label
Figure 41: Effects of attention heads ablation on the layer-wist wise separability and alignment measures of Gemma-7B hidden states on Famous People dataset.
Taak 19 Deevisteg
Kaak
Decolsiog
(a) SST-2
(b) SUBJ
(c) TREC
Kask 10 Bossiding
(d) SNLI
(c) RTE
(f) CB
Figure 42; Accuracy gains of rank- denoising over layers on all datasets with Llama2-70B
49

---

Beak 10 Decolabyy
Beck 10 Decotabyy
Reak Depsliing
laps
lapsr
(a) SST-2
(b) SUBJ
(c) TREC
Blais Desoldog
(d) SNLI
(c) RTE
(f) CB
Figure 43: Accuracy gains of rank- 10 denoising over layers on all datasets with Llama3-8B
Rarch Deuedshuy
Raank 19 Deneldin
Iawer
layse
Sayee
(a) SST-2
(b) SUBJ
(c) TREC
Rack 18 Decelaiy
Rark Densbbigg
Decaining
(d) SNLI
(e)
RTE
(l) CB
Figure 44: Accuracy gains of rank-10 10 denoising over layers on all datasets with Llama3-70B
Beok 18 Deeedaae
Deeohiteg
Back Deroisiee
lapor
lapur
(a) SST-2
(b) SUBJ
(c) TREC
Kaak to Desedsiig
March tn benaaiing
(d) SNLI
(e) RTE
(l) CB
Figure 45: Accuracy gains of rank- 10 denoising over layers on all datasets with Llama2-7B
18
Bank 10 Denotdan
Bank 10 Desrblng
(a) SST-2
(b) SUBJ
(c) TREC
Kaed Deoaking
Arc. tongreeeeenn
mask Besabinn
Lagez
laser
Bayer
(d) SNLI
(c)
RTE
(f) CB
Figure 46: Accuracy gains of rank-10 denoising over layers on all datasets with Llama2-13B
50

---

Blac Deesiaing
Blank 10 Depohite
laged
110
layez
Sayes
(a) SST-2
(b) SUBJ
(c) TREC
Kach Brocking
Kookk Boaaldrn
Rack Besaaling
layer
Jays
byyer
(d) SNLI
(c) RTE
(f) CB
Figure 47: Accuracy gains of rank-10 denoising over layers on all datasets with Gemma-2D
Basi 18 Denshing
Barkh 18 Dwasblng
Runk 20 Desebing
(a) SST-2
(b) SUBJ
(c) TREC
Denudahog
Desebing
(d) SNLI
(c) RTE
(f) CB
Figure 48: Accuracy gains of rank- denoising over layers on datasets with Gemma-7B
Denolite
(a) SST-2
(b) SUBJ
(c) TREC
Disnobidgg
Aec Ioeg morree
Aawar
Lape
hove
(d) SNLI
(c) RTE
(f) CB
Figure 49: Accuracy gains of rank-5 denoising over layers on all datasets with Llama2-70B
Bank Dembhing
Bank Denstatig
Beck Dreatring
Lgne
laye
SST-2
(b) SUBJ
(c) TREC
Rask Denoking
Kook Disntiie
Kark Deaskiee
Lgve
Upper
(d) SNLI
(c) RTE
(f) CB
Figure 50: Accuracy gains of rank-2 denoising over layers on all datasets with Llama2-70B
51

---

(a) SST-2
(b) SUBJ
(c) TREC
Back
Brak Dreabing
Rack Deeeeidn
Joer
(d) SNLI
(e) RTE
(f)
Figure 51: Accuracy gains of rank-1 denoising layers on all datasets with Llama2-70B
52

---

Llama2-70B, average across datasets
Normal Acc.: 0.24% | w/o PTH Acc.: 0.00% w/o Acc.: 0.05% Random Heads Ablated Acc.: 0.08%
Effective dimension
Separability score
Singular alignment
100
0.8
0.075
50
0.050
0.025
80
40
80
20
Layer Number
Layer Number
Layer Number
Composite alignment
Mean based alignment
Var. War--based alignment
0.004
0.001
0.000
0.000
20
40
60
80
0
20
40
60
80
20
40
60
80
Layer Number
Layer Number
Layer Number
Ablated
Figure 52: Average effects across datasets of attention heads ablation on the layer-wise separability and alignment measures of Llama2-70B hidden states in the zero-shot setting.
Table 16: Mean layer values of IH and PTH along with p-values of Mann-Whitney U test on Llama3- 8B.
% Level
IH mean layer
PTH mean layer
p-value of Mann- Whinnyy U test
1%
15.883
9.5667
5.06 10
2%
16.392
12.658
3.99 x 10
5%
16.493
13.833
2.29 x 10
10%
16.435
14.364
4.46 x 10
Table 17: Mean layer values of IH and PTH along with p-values of Mann-W Shin-Whitney U test on Llama3- 70B.
Level
mean layer
PTH mean layer
p-value Mann- Mann-Whitney test
47.624
45.755
0,5573
2%
46.273
46.977
0.1347
5%
44.354
47.277
1.18 x 10-7
10%
43.278
44.257
8.38 x 10
53

---

1.0
Compo Align
0.0075
Sep. Score
O-shot
0.0050
8-shot
0.8
0-shot
0.0025
8-shot
0.0000
16
24
24
32
Layer
Layer
(b) Layer-wise composite alignment: O-shot vs. 8-shot Figure 53: Layer-wise trends of separability score and composite alignment of O-shot and 8-shot hidden states of Llama3-8B.
(a) Layer-wise separability score: O-shot vs. 8-shot
Compo. Align.
Score
O-shot 8-shot
0.002
0.001
0.66
0-shol
8-shol
0.000
24
32
40 Layer
48
56
64
72
16
24
32
48 Layer
64
72
80
(a) Layer-wise separability score: 0-shot vs. 8-shot
(b) Layer-wise composite alignment: O-shot vs. 8-shot
Figure 54: Layer-wise trends of separability score and composite alignment of O-shot and 8-shot hidden states of Llama3-70B.
Sep. Score
1.0
Compo. Align.
0.006
0-shot
8-shot
0.8
0.004
0.002
8-shot
0.000
Layer
Layer
(b) Laser-wise composite alignment: O-shot vs. 8-shot Figure 55: Layer-wise trends of separability score and composite alignment of O-shot and 8-shot hidden states of Llama2-7B.
(a) Layer-wise separability score: O-shot vs. 8-shot
Sep. Score
1.0
Compo. Align.
0.015
O-shol
0.010
8-shol
0.8
0-shot
0.005
8-shot
0.000
16
24
16
24
32
40
Layer
Layer
(b) Lase composite alignment: 0-shot vs. 8-shot Figure 56: Layer-wise trends of separability score and composite alignment of 0-shot and 8-shot hidden states of Llama2-13B.
(a) Layer-wise separability score: (-shoo vs. 8-shot
54

---

1.0
Compo. Align.
Sep. Score
0.03
0-shot
8-shot
0.8
0.02
0-shot
0.01
0.6
8-shot
0.00
8
16
Layer
Layer
(a) Layer- wise separability score: 0-shot 8-shot
(b) Layer-wise composite alignment: O-shot vs. 8-shot
Figure 57: Layer-wise trends of separability score and composite alignment of (-shot and 8-shot hidden states of Gemma-2B.
1.0
Compo. Align.
Sep. Score
0.015
O-shot
8-shot
0.8
0.010
0-shot
0.005
0.6
8-shot
0.000
16
24
16
24
Layer
Layer
(a) Lise separability score: 0-shot 8-shot
(b) Layer-wise composite alignment: O-shot vs. 8-shot
Figure 58: Layer-wise trends of separability score and composite alignment of (/-hot and 8-shot hidden states of Gemma-7B.
Compo. Align.
Sep. Score
0.010
PTH TV
as TV
0.005
PTH as TV
0.6
0.000
32
Layer
Layer
Layer-wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 59; Effects of injecting PTH and IH outputs as task vectors on the geometric properties of Llama3-8B shot hidden states.
Compo. Align.
Sep. Score
0.0010
PTH TV
0.8
0.0005
PTH as TV
0.6
IH
0.0000
24
32
40 Layer
48
56
64
8o
16
24
32
40 Layer
56
64
80
(a) Layer-wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 60: Effects of injecting PTH and IH outputs as task vectors on the geometric properties of Llama3-70B zero-shot hidden states.
55

---

Compo. Align.
0.04
Sep. Score
PIH as TV
0.8
IH as TV
0.02
PTH as TV
IH as TV
0.00
16
24
32
16
24
32
Layer
Layer
(a) Layer-wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 61: Effects of injecting PTH and IH outputs as task vectors on the geometric properties of Llama2-7B zero-shot hidden states.
Compo. Align.
0.075
Score
PTH as TV
TV
0.8
0.050
IH as TV
PTH as
0.025
IH as TV
0.000
16
24
32
40
16
24
32
40
Layer
Layer
(a) Layer-wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 62: Effects of injecting PTH and IH outputs as task vectors on the geometric properties of
Llama2- 13B zero-shot hidden states.
Compo. Align.
Sep. Score
T
0.02
PTH TV
PTH as TV III TV
0.8
1H
0.01
TV
IH as TV
0.00
16
Layer
Layer
(a) Layer-wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 63: Effects of injecting PTH and IH outputs as task vectors on the geometric properties of Gemma-2B zero-shot hidden states.
Compo. Align.
Score
FTH
PTH as IH as TV
0.8
0.01
IH as TV
0.00
16
24
16
Layer
Layer
(a) Layer-wis wise separability score: PTH task vector
(b) Layer-wise composite alignment: IH task vector
Figure 64; Effects of injecting PTH and IH outputs as task vectors the geometric properties of
Gemma-7B 7B zero-shot hidden states.
56

---

Table 18: Mean layer values IH and PTH along with p-values of Mann- Mann-Whinney U test on Llama2-
7B.
Level
mean layer
PTH mean layer
p-value Mann- Whitney test
%
15.883
8.5833
5.62 x 10-7
2%
16.133
10.108
9.78 x 10
5%
16.301
12.092
4.12 x 10
10%
16.405
12.559
7.27 10
Table 19: Mean layer values IH and PTH along with p-values of Mann- Hann-Whitney U test on Llama2- 13B.
Level
mean layer
PTH mean layer
p-value Mann- Whitney test
1%
21.312
17.094
1.75 x 10-2
2%
20.135
18.026
9.44 10
5%
19.960
16.927
6.08 x 10
10%
19.888
16.343
3.08 10 10-13
Table 20: Mean layer values of IH and PTH along with p-values of Mann-W Whin-Whitney U test on Gemma-
2B.
Level
IH mean layer
PTH mean layer
p-value of Mann- Whitney test
1 1
9.0000
1.6667
8.58 x 10-3
2%
10.917
5,0000
7.91 x 10-4
5%
11.571
4.4286
2.22 x 10 10-10
10%
12.214
6.1667
2.48 10
Table 21: Mean layer values of IH and PTH along with p-values of Mann-Wiin WhinneW U test on Gemma-
7B.
Level
IH mean layer
PTH mean layer
p-value Mann- Mann-Whitney test
1%
18.708
9.0833
2.01 x 10-5
2%
19.062
9.1250
1.96 x 10-6
5%
18.742
7.6212
6.31 x 10-21
10%
18.523
9.0530
4.44 10
57

---

0.8
0.7
0.7
Accurary
0.5
Aorunn
0.5
0.4
0.4
0.3
0.3
0.3
0.2
0.2
0.1
0.0
0.0
0.0
Zhro-shot vector
Zeto
Figure 65: Average effect across datasets of steering Llama3-8B zero-shot hidden states using task vectors created from IH out- puts.
Figure 66: Average effect across datasets of steering Llama3-70B zero-shot hidden states using task vectors created from IH out-
Figure 67: Average effect across datasets of steering Llama2-7B zero-shot hidden states using task vectors created from IH out- puts.
puts.
0.8
0.7
0.6
Acarsey
Acoracy
0.4
0.5
0.4
0.4
0.3
0.3
0.2
0.1
Zeroshot Task
Zero-shot Task
Zaro-shot Task
Figure 68: Average effect across datasets steering Llama2-133 zero-shot hidden states using task vectors created from IH out- puts.
Figure 69: Average effect across
Figure 70: Average effect across
datasets of steering Gemma-2B zero-shot hidden states using task vectors created from out-
datasets of steering Gemma-78 zero-shot hidden states using task vectors created from IH out- puts.
puts.
Table 22: Comparison classification accuracy: training classifier v.s. direct decoding with zero-shot prompts.
Model
0-shot (%)
0-shot-classifier (%)
Improvement
Llama3-8B
0.30%
73.71%
+73.41%
Llama3- 30B
0.00%
78.97%
+78.97%
Lllmm2-7B
0.31% %
75.63%
+75.32%
Ilama2- 13B
0.36%
79.71%
+79.35%
-lama2- 2-B0B
0.24%
79.02%
+78.78%
Gemma- 2B
0.00%
75.56%
+75.56%
Gemma-77 7B
0.02%
75.43%
+75.41%
58

---

Table 23: Comparison classification accuracy: training classifier v.s. direct decoding with zero-shot
prompts.
Model
8-shot (%)
8-shot-classiffer (%)
Improvement
Llama3-88 8B
76.20%
85.47%
+9.27%
Llama3-70B
81.01 %
84.46%
+3.45%
Ilama2-7B 7B
74.83%
85.59%
+10.76%
I.ama2- 13B
74.82%
85.44%
+10.62%
Llama2-70 -0B
80.52%
87.76%
+7.24%
Gemma-2B
67.03%
77.87%
+10.84%
Gemma-77 7B
79.19%
84.94%
+5.75%
Table 24: Comparison of classification accuracy: regular ICL v.s. low-rank denoising of hidden states
Model
8-shot (%)
8-shot-denoising (%)
Improvement
Llama3-88 BBB
76.20%
78.58%
+2.38%
Llama3-70B
81.01% %
82.40%
+1.39%
Llama2-7B
74.83%
77.68%
+2.85%
Llama2-
74.82%
79.52%
+4.70%
Llama2- 70B
80.52%
80,72%
+0.20%
Gemma-28 2B
67. 67.03%
67.25%
+0.22%
Gemma- BB
79.19%
66.11%
-13.08%
59


Reviews:
[{"review_id": "betcmNTOJu", "reviewer": "Reviewer_Ua9F", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors' rebuttal has addressed most of my concerns. The inclusion of the generation task strengthens the generality of the results. Moreover, the proposed denoising and classifier-based approaches offer practical methods for enhancing the classification performance of ICL. Overall, this work provides a comprehensive and in-depth analysis of how ICL works from the perspective of hidden representations. That said, I still find the contribution somewhat limited in scope, as it remains confined to classification tasks. As such, I will maintain my current rating.", "rebuttal": "**Dear Reviewer Ua9F,**\n\nThank you for your thoughtful review of our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.\n\n---\n\n## **W1: Limitations of solely focusing on classification tasks**\n\nWe appreciate your view that structuring our analysis solely around classification tasks seems limiting. Per your suggestion, and to address your concern that our geometric framework of ICL may not generalize to generation tasks, **we repeat our experiments on a text generation dataset** with the following setup. The demonstrations and queries take the format \"Write a favourable/unfavourable review of {subject}:\" where {subject} is the name of a kind of food (e.g., tapas). The labels are GPT-generated favourable or unfavourable reviews about the {subject}. Eight demonstrations and a query are concatenated to form each ICL-style prompt, and we compute our separability and alignment measures, as well as the effective dimension, on the dataset of prompts. The query labels needed for training the classifier to compute the separability score are determined based on whether a favourable or unfavourable review is requested in the query. We also compute the alignment measures with respect to the unembedding vectors of the \"positive\" and \"negative\" tokens to address the issue that labels for generation tasks comprise multiple tokens and take indefinite forms. Below in Table 1, we report the results concerning the separability score and singular alignment (as a representative alignment measure) across all layers of Llama3-8B in the zero-shot and ICL settings. (All experimental results presented in the rebuttal are conducted on Llama3-8B.)\n\n| Singular Alignment      | ...    | 14     | 15     | 16     | 17     | 18     | 19     | 20     | 21     | 22     | 23     | 24     | 25     | 26     | 27     | 28     | 29     | 30     | 31     | 32     |\n|-------------------------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|\n| zero-shot               | ...    | 0.0304 | 0.0026 | 0.0163 | 0.0132 | 0.0200 | 0.0158 | 0.0147 | 0.0076 | 0.0057 | 0.0084 | 0.0088 | 0.0054 | 0.0064 | 0.0135 | 0.0155 | 0.0278 | 0.0278 | 0.0310 | 0.0108 |\n| ICL                     | ...    | 0.0270 | 0.0360 | 0.0410 | 0.0361 | 0.0545 | 0.0899 | 0.0762 | 0.0836 | 0.0918 | 0.0802 | 0.0731 | 0.0717 | 0.0766 | 0.0781 | 0.0547 | 0.0498 | 0.0613 | 0.0626 | 0.0810 |\n\n| Separability Score          | ...    | 14     | 15     | 16     | 17     | 18     | 19     | 20     | 21     | 22     | 23     | 24     | 25     | 26     | 27     | 28     | 29     | 30     | 31     | 32     |\n|-------------------------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|\n| zero-shot               | ...    | 0.992  | 0.992  | 0.985  | 0.985  | 0.977  | 0.969  | 0.969  | 0.961  | 0.961  | 0.969  | 0.969  | 0.969  | 0.969  | 0.969  | 0.977  | 0.969  | 0.977  | 0.992  | 1      |\n| ICL                     | ...    | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      | 1      |\n*Table 1: Layer-wise trend of singular alignment and separability score for the generation task: ICL (8-shot) and zero-shot* \n\nThese results demonstrate the increase in alignment measures in the middle layers unique to the ICL setting and show that both ZSL and ICL hidden states exhibit high separability, thereby demonstrating the **generalizability of our framework to generation tasks**.\n\n---\n\n## **W2 & Q2: Explanation for why PTHs improve separability, IHs improve alignment**\n\nBelow is a more detailed exposition building on our explanation in Section 5.3. First, it is widely documented that the hidden states of texts with different semantic values (e.g., positive/negative) are linearly separable in LLMs ([1][2][3]). PTHs mainly attend to the query texts with positive or negative sentiment and thus can effectively encode such separability into the final token's hidden state through their outputs, formed by the query texts' hidden states via their embedding matrices. We have also conducted an experiment showing that ablating PTHs affects separability in the zero-shot case (where only the query texts exist), thus validating this explanation.\n\n| Separability Score (SST-2)          | ...    | 23     | 24     | 25     | 26     | 27     | 28     | 29     | 30     | 31     | 32     |\n|-------------------------------------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|\n| zero-shot                           | ...    | 0.927  | 0.917  | 0.927  | 0.929  | 0.927  | 0.927  | 0.922  | 0.922  | 0.911  | 0.917  |\n| zero-shot (PTH ablated)             | ...    | 0.819  | 0.817  | 0.819  | 0.814  | 0.821  | 0.823  | 0.817  | 0.810  | 0.826  | 0.833  |\n*Table 2: Layer-wise trend of separability score for the SST-2: zero-shot & PTH ablation* \n\nSecond, in the middle to late layers, IHs emerge and attend to the label tokens. Via the same mechanism discussed above, they encode the information of the semantically rich label tokens into the final token's hidden state, thus promoting alignment. The alignment between the query hidden states and their respective labels is achieved through IHs distributing attention strengths among the demonstration labels in context (discussed in detail in [4][5]). Our experiment in Appendix H.2 (showing that ablating IHs has less effect when the query label is not in context) also supports this explanation.\n\n---\n\n## **W3: Implications for ICL performance**\n\nWe agree with you that connecting our analysis to practical implications regarding ICL performance is important. We would like to first draw your attention to Section 5.2 and Fig. 4 in particular, where we present a considerable increase in logits lens accuracy achieved through low-rank denoising, which exploits the alignment structure of the hidden states with label unembedding vectors. We apply rank-5 denoising to Llama3-8B final layer hidden states. As a result, we can increase the final layer output accuracy averaged across six datasets as reported in Table 3. \n\nFurthermore, the high separability score of ICL hidden states presented in Section 5.2 and Fig. 2 implies that we can also improve ICL accuracy by training a classifier on several ICL hidden states and using its prediction on the hidden states of the remaining prompts to get the label, resulting in an ICL calibration label akin to the one proposed by [6]. Using this method, we also manage to significantly increase the average ICL accuracy as shown in Table 3.\n\n|Method|ICL|ICL+Denoising|ICL+Classifier|\n|------|-----|----------------|----------------|\n|Accuracy|76.20|78.59|85.47|\n*Table 3: Improving ICL accuracy using denoising or classifier as calibration methods* \n\nIn summary, **our findings lead to two methods to improve ICL performance that are lightweight, effective (the low-rank denoising method is also supervision-free), and can be seamlessly integrated into any ICL scenario involving text classification.**\n\n---\n\n## **Q1: Where PTHs attend**\n\nYes, you are correct—PTHs should attend to the query text at the final position, as per our response to W2 & Q2. The reason we wrote that PTHs attend to the answer here is that, in the identification of PTHs, due to complex tokenization issues, it can be difficult to accurately delineate the token range of the query text and calculate attention scores. That is why we identify PTHs as the heads that attend to the immediately preceding token at each position, which is why we wrote that PTHs attend to \"Answer\" at the final position. But this is still valid at the final position because the hidden state of \"Answer\" will encode information from the query text due to the autoregressive structure of Transformers. **We have made the change here according to your suggestion.**\n\n---\n\n# References\n[1] Marks, Samuel, and Max Tegmark. \"The geometry of truth: Emergent linear structure in large language model representations of true/false datasets.\"\n\n\n[2] Park, Kiho, Yo Joong Choe, and Victor Veitch. \"The linear representation hypothesis and the geometry of large language models.\"\n\n[3] Tigges, Curt, et al. \"Linear representations of sentiment in large language models.\"\n\n[4] Yu, Zeping, and Sophia Ananiadou. \"How do large language models learn in-context? query and key matrices of in-context heads are two towers for metric learning.\" \n\n[5] Halawi, Danny, Jean-Stanislas Denain, and Jacob Steinhardt. \"Overthinking the truth: Understanding how language models process false demonstrations.\" \n\n[6] Cho, Hakaze, et al. \"Token-based decision criteria are suboptimal in in-context learning.\""}, {"review_id": "R5t7yGcLX1", "reviewer": "Reviewer_4WZT", "summary": "The 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.", "questions": "- 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?", "limitations": "Yes, limitations are discussed.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "Thanks for the very thorough response. I'm now further convinced of the reported findings, as they seem to hold through across several different methods of ablation and analysis. I also thank the authors for not only verbally addressing my last question but even running further experiments to give better intuition. I will update my score accordingly.", "rebuttal": "**Dear Reviewer 4WZT,**\n\nThank you for your constructive feedback on our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.\n\n---\n\n## **W1 Part 1 and W2: Clarifications about the attention head ablation approach**\n\nWe appreciate your concerns regarding the details of our attention head ablation procedure. As described in **lines 314-319** in Section 5, we use a **zero-ablation** approach. Regarding your question about the number of heads ablated and how they are identified: we ablate the top 10% of heads ranked by IH and PTH scores, as well as a control group consisting of 10% randomly selected heads that exclude the top 10% IHs and PTHs. Further details on the computation of IH and PTH scores can be found in **Appendix G**.\n\nIn **Appendix G**, we explain that we use the first 50 queries from each dataset to compute these scores. For PTHs, each head's score is calculated by summing the attention weights it assigns to the immediately preceding token at all token positions across a query, and then summing over all 50 queries. For IHs, we first append 8 demonstrations to each query. Then, we compute each head's IH score by summing the attention weights it assigns to the 8 demonstration labels at the final \":\" position of the query, and then summing over all queries.\n\nIn response to your question about the **layerwise distribution of ablated top IHs**, we analyzed the top 10% IHs in Llama3-8B. Specifically, we divided the model's 32 layers into 8 equal intervals (1-4, 5-8, ...) and computed the proportion of IHs found in each interval:\n\n|Layer Interval|1-4|5-8|9-12|13-16|17-20|21-24|25-28|29-32|\n|--|---|---|----|-----|-----|-----|-----|-----|\n|Percentage|2.4510|6.0458|16.1765|20.2614|23.5294|18.4641|7.8431|5.2288|\n*Table 1: Distribution of ablated top 10 IHs across layer intervals*\n\nThe results show that IHs are most concentrated in the middle layers—precisely where we observe phase transitions in ICL—supporting our claim about a two-stage ICL mechanism.\n\n---\n\n## **W1 Part 2: Using mean ablation as an alternative approach**\n\nWe agree with your suggestion that supplementing our findings with **mean ablation** results strengthens their robustness. Accordingly, we repeated the experiments using mean ablation on Llama3-8B across all 6 datasets. We report the dataset-averaged layerwise dynamics of **separability** and **alignment** under different ablation conditions. We report **separability score** and **singular alignment** as a representative of all alignment measures.\n\n|Separability Score|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|--|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|ICL|...|0.8920|0.8992|0.8991|0.9073|0.9012|0.9106|0.9089|0.9081|0.9027|0.9010|0.8921|0.8856|0.8847|0.8848|0.8829|0.8808|0.8783|0.8764|0.8547|\n|PTH ablated|...|0.5584|0.5515|0.5434|0.5516|0.5562|0.5582|0.5574|0.5579|0.5478|0.5503|0.5433|0.5502|0.5556|0.5496|0.5516|0.5331|0.5368|0.5473|0.5225|\n|IH ablated|...|0.6473|0.6576|0.6838|0.6897|0.6990|0.6948|0.6956|0.6841|0.6712|0.6600|0.6584|0.6536|0.6504|0.6505|0.6352|0.6360|0.6340|0.6330|0.6176|\n|Random heads ablated|...|0.7105|0.7515|0.7659|0.7609|0.7738|0.7854|0.7852|0.7873|0.7885|0.7841|0.7824|0.7951|0.7902|0.7891|0.7830|0.7838|0.7829|0.7698|0.7248|\n\n|Singular Alignment|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|--|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|ICL|...|0.0291|0.0407|0.0634|0.0787|0.0948|0.1351|0.1375|0.1440|0.1649|0.1862|0.2016|0.1996|0.2090|0.2095|0.2041|0.2033|0.1917|0.1533|0.1015|\n|PTH ablated|...|0.0137|0.0150|0.0217|0.0221|0.0293|0.0373|0.0380|0.0428|0.0624|0.0816|0.0945|0.1262|0.1268|0.1278|0.1317|0.1456|0.1290|0.0883|0.0287|\n|IH ablated|...|0.0238|0.0234|0.0226|0.0201|0.0318|0.0375|0.0269|0.0332|0.0395|0.0314|0.0324|0.0313|0.0353|0.0342|0.0270|0.0285|0.0234|0.0216|0.0156|\n|Random heads ablated|...|0.0170|0.0260|0.0485|0.0604|0.0890|0.1369|0.1362|0.1417|0.1595|0.1722|0.1914|0.1923|0.2007|0.2065|0.1987|0.1945|0.1819|0.1345|0.0597|\n\n*Table 2: Impacts on separability score and singular alignment when IH, PTH, or random heads are ablated*\n\n\nThese results confirm our hypothesis: **PTHs primarily affect separability**, while **IHs primarily affect alignment**.\n\n---\n\n## **Q1: PTHs and mean-based alignment**\n\nAs explained in **lines 324–326**, ablating PTHs can substantially impact mean-based alignment because this metric depends on the norm of the difference between the hidden state cluster means projected onto the label unembedding direction. If the label-conditioned clusters are not separable (as happens when PTHs are ablated), this norm becomes small, reducing alignment. In contrast, ablating IHs affects the alignment between this difference and the label unembedding differnce, but does not collapse the cluster separation that much—hence could cause a smaller impact on separability.\n\n---\n\n## **Q2: Regarding the referenced paper**\n\nThank you for bringing up this relevant work. We believe the differences in findings can be attributed to two main factors:\n\n**First**, the two studies use different methods to identify IHs. The paper in question computes prefix matching scores on random symbol sequences (e.g., `abc...abc...`), whereas we evaluate attention to demonstration labels in ICL-formatted prompts, following the approach in [1] (the original paper proposing function vector heads). This method ensures that the identified IHs are more task-relevant. Indeed, [1] also observed that 3 of their top 10 function vector heads exhibited IH-like attention.\n\nTo validate this, we conducted an experiment comparing the impact of ablating IHs identified via the method in the referenced paper versus ours. We applied mean ablation on SST-2 and measured output alignment:\n\n|Output Alignment|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|---------------|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|IHs ablated (approach in paper)|...|0|0|0|0|0|0|0.0023|0.0023|0.0218|0.0310|0.0677|0.0482|0.0745|0.0975|0.0206|0.0057|0|0|0.0218|\n|IHs ablated (ours)|...|0|0|0|0|0|0|0|0|0|0|0|0|0|0|0|0|0|0|0|\n*Table 3: Ouput Alignment with ablation of IHs identified using two different criteria*\n\nThe results confirm that **the outputs of our IHs encode more task-specific semantics** and serve as stronger task vectors.\n\n**Second**, while top-2% IHs and function vector heads do not perfectly overlap, the paper also reports **strong correlation** between their induction and FV scores across all heads. For example, top function vector heads often fall in the 90–95% percentile of induction scores and vice versa. This supports our use of top IHs as **effective, albeit not optimal, FV heads**.\n\n---\n\n## References\n\n[1] Todd, Eric, et al. \"Function vectors in large language models.\""}, {"review_id": "h8wp1nj8j0", "reviewer": "Reviewer_kMpH", "summary": "This 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.", "questions": "1. 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.", "limitations": "See 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.", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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", "was_revised": true, "final_justification": "The manuscript currently exhibits an insufficient and sometimes incorrect understanding of task vectors and ICL, and the necessary citations and conceptual distinctions are largely missing.", "rebuttal": "**Dear Reviewer kMpH,**\n\nThank you for your careful review of our manuscript. Below we respond to the weaknesses and questions you raised.\n\n\n## **W1 & W2 part 1: Insufficient discussion of task vector literature**\n\nWe appreciate your suggestion to clarify our connection to the task vector literature. As you noted, the task vector literature is vast, with each work offering unique definitions and settings. **We reiterate that the main focus of our work is the geometry-inspired, process-oriented exposition of the ICL mechanism. That is why we focus only on task vectors in ICL settings—because they are relevant and serve as a motivation for our framework.** \n\nFollowing your suggestion, we will include a Discussion Section in the camera-ready version situating and fitting our work within the ICL-related task vector literature. The section will include the following parts:\n\n1. **In-weight Task Vector.** The term Task Vector originates from the parameter difference between a fine-tuned model on a specific task and its pre-trained counterpart, typically constructed via Linear Mode Connectivity [2–6, 9–13, 19, 20, 29]. However, ICL involves no explicit parameter updates. Thus, despite sharing the same term, task vectors in ICL and those based on in-weight changes are conceptually hard to align. **We will include a brief discussion of in-weight task vectors as necessary background.**\n2. **Linear Task Vector.** Linear Task Vectors [8, 14–16] typically shift hidden states from a query embedding (e.g., China, England) to a target expression (e.g., Beijing, London) in factual recall tasks. Our definition generalizes this idea: as shown in Figure 1, our task vectors also transform embeddings with mixed semantics (e.g., China, or a positive sentence) toward specific attributes (e.g., Beijing, or positive), and are not limited to simple translation in hidden space. This perspective both aligns with prior work and extends it. **We will include these references to highlight the connection.**\n3. **Black-box Task Vector.** Some works [1, 7, 17, 18] define task vectors loosely as any vector that enables patching-and-improvement, without specifying their form. Our additive task vector, after passing through Transformer layers, can similarly induce such effects, thus serving as a precursor to black-box task vectors. **We will clarify this point in the final version.**\n\n\n## **W3: ICA as an alternative method for interpreting hidden state dynamics**\n\nWe agree that ICA is a viable method for extracting semantic directions. However, in our context, our SVD-based approach is preferable for two reasons.\n\n**Theoretical reason.** Our analysis targets the geometric evolution of hidden states. The top right singular vectors used in our analysis identify directions of maximum variation, crucial for capturing alignment with unembedding vectors. In contrast, the ICA components (columns of $\\mathbf{A}$ in $\\mathbf{H} = \\mathbf{S}\\mathbf{A}$) **lack geometric significance.** The column order is arbitrary, rendering most experiments in Section 5.2 infeasible.\n\n**Practical reason.** **ICA requires $n\\geq d$**, but the 70B models we use have $d=8192$, exceeding most dataset sizes. The paper you mentioned uses $d=300$ and manually inspect activations of ICA components for semantic interpretation. In contrast, our semantic analysis depends on measuring component alignment with unembedding vectors which equals model dimension.\n\nNonetheless, per your suggestion, we repeated the semantic filtering and retention experiment on SST-2 with Llama3-8B. Below is the result at layer 25:\n\n|layer 25|retained tokens|filtered tokens|\n|---|---|---|\n|SVD|extra, hud, neutral, rated, herent|PointerException, ldr, anus, ipt, allon|\n|ICA|elo, raries, injury, .DropDown, offensive|dead, ets, seksi, @Spring, 견|\n\nIt shows that SVD and ICA yield similar outcomes; SVD even decodes the SST-2-relevant \"neutral\" token.\n\n## **W4 & Q3 Part 2: Limitations of focusing solely on classification tasks**\n\nPlease see Section **W1: Limitations of solely focusing on classification tasks** in our response to **Reviewer Ua9F**, where we show our geometric two-stage ICL hypothesis extends to text generation tasks.\n\n\n## **Q1: Comparison between [7]'s findings and ours**\n\nThank you for highlighting this relevant. Below, we outline two key distinctions.\n\n**Study setting.** Paper [7], despite having 'in-context learning' its title, studies rather a general RAG setting where the model can draw on both internal knowledge and external sources. Their \"context\" includes both prompt and retrieved information, which allows them to define head types (e.g., in-context heads) based on which part of the context influences outputs. **In contrast, our setting uses classical ICL: the model must learn an input-output mapping from a few-shot prompt and respond to a query.** This setting justifies our focus on classification, which exemplifies the mapping learning process. **The importance of IHs and PTHs in this setting has already been well-documented ([21]–[25]) and is considered common knowledge in the ICL interpretability community**, as stated in our Introduction and Related Work.\n\n**The approach of study** [7] identifies the attention heads with special roles by measuring the association between the activation of attention heads or their distribution of attention weights, to the output logits. For instance, their 'task heads' is defined to be the heads whose attention weights to the question in the prompt have a large effect on the model outputs. Then they inject such heads' outputs as task vector into the hidden state of a prompt with a different question and recovers the same answer. This approach of characterizing attention heads can be found in other studies([26],[27]), and is conceptually akin to the ablation-based approach in [23][24], in that they characterize heads only by **how much** they influence the model outputs, but not by **how** they influence the outputs through LLM internal workings\n\n**This is why our study is innovative because we propose a process-oriented perspective of analyzing the how attention heads inform LLM outputs**. Thus our approach to attention heads and function vectors is that we first find that IHs induce the increase in alignment in the two-stage ICL process, and then we succesfully construct task vectors out of IH outputs because we know that alignment is missing in the zero-shot setting. \n\n\n## **Q2, Q3 Part 1, W2 Part 2: Hendel et al.'s task vectors—linearity, location, and adaptability**\n\nWe respect your concern that the linearity of the task vector defined by Hendel et al. is evident in [15]'s setting (Country-Capital etc.) that mimics the simple word2vec scenario, but not as evident in classification tasks. We have addressed this issue in the footnote at **line 130**. The Hendel et al.-styled hidden state patching can be reformulated as an addition where the added vector is the difference between the ICL and zero-shot hidden states.  Since ICL hidden states are more aligned with label unembeddings than zero-shot ones, the difference vector also increases alignment when added to zero-shot states. We repeat Hendel et al.'s experiment with SST-2 and Llama3-8B at all layers to verify this. We have also decoded the top tokens from the patched ICL hidden state, the original hidden state and the difference vector extracted from one example. \n\n|Index|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|Accuracy| ...|0.0321|0.0883|0.2752|0.2282|0.2523|0.3165|0.3429|0.3372|0.3601|0.4587|0.4736|0.4817|0.5000|0.4759|0.4966|0.4977|0.4862|0.4885|0.5172|\n\n(Before layer 14, accuracy is ~0)\n\n|Layer 26|original $h$s|ICL $h$s|difference|\n|---|---|---|---|\n| Tokens| ..\\n, asiswa, _exempt, in, واء|negative, negatively, positive, positively, in|negative, positive, -negative, negative, Negative |\n\nThe results show that the benefit of intervention increases with layer-consistent with our finding that alignment improves in deeper layer. The decoded tokens show the difference vector can steer zero-shot states toward label-aligned directions. The ~50% accuracy is expected: since we patch using a dummy query's ICL state following Hendel et al., there's a 50% chance that its label matches the current one in the two-labeled SST-2. Note that Hendel et al.'s result **does not show that task vector intervention is not effective in late layers.** It is because that the patched ICL hidden state in deep layers aligns with the dummy query's label and on non-classification tasks the probability of two queries sharing a label is smaller. [15] also note the effectiveness of task vector in quasi-classification tasks (e.g. Name to Nationality) in late layers, and [28] find that task vectors can emerge in many layers from middle to late, depending on the nature of the task. These findings can reconcile the discrepancies between our findings and the results you mentioned concerning the where task vectors are most effective\n\n# References\nFollows the paper list you mentioned in the review...\n\n[21] Olsson, C. et al. \"In-context learning and induction heads.\"\n\n[22] Singh, A. K. et al. \"What needs to go right for an induction head?...\"\n\n[23] Cho, H. et al. \"Revisiting in-context learning inference circuit...\"\n\n[24] Crosbie, J., Shutova, E. \"Induction heads as an essential mechanism...\"\n\n[25] Song, J. et al. \"Out-of-distribution generalization via composition...\"\n\n[26] Jin, Z. et al. \"Cutting off the head ends the conflict...\"\n\n[27] Bansal, H. et al. \"Rethinking the role of scale for in-context learning...\"\n\n[28] Yang, L. et al. \"Yang, Liu, et al. \"Task vectors in in-context learning: Emergence, formation, and benefit.\"\n\n[29] Neyshabur, B. et al. \"What is being transferred in transfer learning?\""}, {"review_id": "xG4fBgsFht", "reviewer": "Reviewer_LiD7", "summary": "This 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.", "questions": "How 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?", "limitations": "adequately addressed", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "My concerns have been comprehensively addressed in the rebuttal, and I have upgraded my rating.", "rebuttal": "**Dear Reviewer LiD7,**\n\nThank you for your insightful review of our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.\n\n---\n\n## **W1 & Q1: Limitations of solely focusing on classification tasks**\n\nWe appreciate your view that structuring our analysis solely around classification tasks seems limiting. Per your instruction, and to address your concern that our geometric framework of ICL may not generalize to generation tasks, **we repeat our experiments on a text generation dataset** with the following setup. The demonstrations and queries take the format \"Write a favourable/unfavourable review of {subject}:\" where {subject} is the name of a kind of food (e.g., tapas). The labels are GPT-generated favourable or unfavourable reviews about the {subject}. Eight demonstrations and a query are concatenated to form each ICL-style prompt, and we compute our separability and alignment measures, as well as the effective dimension, on the dataset of prompts. The query labels needed for training the classifier to compute the separability score are determined based on whether a favourable or unfavourable review is requested in the query, and we also compute the alignment measures w.r.t. the unembedding vectors of the \"positive\" and \"negative\" tokens to address the issue that labels for generation tasks comprise multiple tokens and take indefinite forms. Below, we report the results concerning the separability score and singular alignment (as a representative of alignment measures) across all layers of Llama3-8B in the zero-shot and ICL settings. (All experimental results presented in the rebuttal are conducted on Llama3-8B.)\n\n|SingularAlignment|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|----------------|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|zero-shot|...|0.0304|0.0026|0.0163|0.0132|0.0200|0.0158|0.0147|0.0076|0.0057|0.0084|0.0088|0.0054|0.0064|0.0135|0.0155|0.0278|0.0278|0.0310|0.0108|\n|ICL|...|0.0270|0.0360|0.0410|0.0361|0.0545|0.0899|0.0762|0.0836|0.0918|0.0802|0.0731|0.0717|0.0766|0.0781|0.0547|0.0498|0.0613|0.0626|0.0810|\n\n|Separability|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|------------|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|zero-shot|...|0.992|0.992|0.985|0.985|0.977|0.969|0.969|0.961|0.961|0.969|0.969|0.969|0.969|0.969|0.977|0.969|0.977|0.992|1|\n|ICL|...|1|1|1|1|1|1|1|1|1|1|1|1|1|1|1|1|1|1|1|\n\n*Table 1: Layer-wise trend of singular alignment and separability score for the generation task: ICL (8-shot) and zero-shot* \nThe results demonstrate the rise in alignment measures in the middle layers that is unique to the ICL setting and show that both ZSL and ICL hidden states have high separability, thereby demonstrating the **generalizability of our framework to generation tasks**.\n\n---\n\n## **W2 & Q2: Limited Practical Implication & Improving Zero-Shot Performance**\n\nWe respect your concern about the practical significance of our work, particularly for zero-shot learning. We would like to point out that our findings about the highly separable zero-shot hidden states can directly be used to improve zero-shot performance, without having to explicitly change the alignment properties of the hidden states as you mentioned. Concretely, we can simply train a classifier on the hidden states of some zero-shot prompts as labels, and use its predictions on other prompts as the predicted labels. For the Llama3-8B model, this approach increases the **zero-shot accuracy averaged across datasets from 0.3 to 73.71** (i.e., the mean separability score for the Llama3-8B final layer hidden states presented in the paper). Apart from being lightweight, this approach also requires only a small amount of labeled data. Using just 30 labeled prompts for training achieves 72.94 accuracy on the SST-2 dataset (in the calculation of the separability score, half of the prompts in the dataset are used to train the classifier).\n\n---\n\n## **W3 Part 1: Direct Evidence for the Geometric Significance of IHs and PTHs**\n\nWe agree with your suggestion that evidence directly supporting the association between alignment and IHs, as well as separability and PTHs beyond the ablation experiments, will make our case more convincing. Hence, we conduct experiments where we construct task vectors using the outputs of top-scoring IHs and PTHs and add the task vectors to the final token's zero-shot hidden state. Due to the additive nature of attention heads' outputs, this experiment mimics the actual effect of IHs and PTHs on the hidden state residual stream. We inject the PTH task vector at layer 2 and the IH task vector at layer 16 of the 32-layered Llama3-8B. Then we directly record the dynamics of the separability and alignment measures near the layers of injection. We report the separability score for the PTH task vector and output alignment (logit lens accuracy) for the IH task vector (averaged across datasets). The results are reported below.\n\n|Separability Score|0|1|2|3|4|5|6|7|8|9|10|...|\n|----------------|--|--|--|--|--|--|--|--|--|--|--|---|\n|zero-shot (PTH as TV)|0.4491|0.4532|0.5189|0.6960|0.6924|0.7121|0.7396|0.7495|0.7689|0.7832|0.8059|...|\n|zero-shot|0.4491|0.4491|0.4651|0.5178|0.5548|0.5989|0.6340|0.6797|0.7199|0.7331|0.7873|...|\n\n|Output Alignment|...|14|15|16|17|18|19|20|21|22|23|24|25|26|27|28|29|30|31|32|\n|---------------|---|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|zero-shot (IH as TV)|...|0|0|0.4431|0.4524|0.4605|0.4526|0.4533|0.4533|0.4537|0.4803|0.5074|0.5158|0.4985|0.5261|0.4633|0.4736|0.4933|0.4553|0.5476|\n|zero-shot|...|0|0|0|0|0|0.0003|0|0.0093|0.0006|0.0009|0.0025|0.0015|0.0019|0.0032|0.0017|0.0023|0.0003|0.0010|0.0023|\n\n*Table 2: Dynamics in separability and alignment measure after task vector injection using PTH and IH outputs as task vectors*\n\nThe significant increase in separability and alignment after the injection of PTH and IH outputs respectively is strong and direct evidence of the geometric significance of PTHs and IHs.\n\n---\n\n## **W3 Part 2: Location and Identification of IHs and PTHs**\n\nPer your suggestion, we have calculated the mean layer index of the top 1%, top 2%, top 5%, and top 10% PTHs and IHs for Llama3.1-8B. We have also performed a Mann-Whitney U test to examine whether the difference in the layer distributions is significant, with PTHs significantly appearing earlier than IHs at all four levels. The results below show that PTHs indeed significantly precede IHs at all levels, thereby supporting our two-stage hypothesis regarding ICL. We will add these results, together with the explanation of how IHs and PTHs are identified which is currently in **Appendix G**, to Section 5.3 of the camera-ready version of this paper.\n\n|%Level|IH mean layer|PTH mean layer|p-value of Mann-Whitney U test|\n|------|-----------|------------|---------------------------|\n|1%|15.883|9.5667|5.0594e-07|\n|2%|16.392|12.658|3.9983e-05|\n|5%|16.493|13.833|2.2897e-05|\n|10%|16.435|14.364|4.4557e-06|\n*Table 3: Mean layer index of IHs and PTHs and the significance of difference in layer distributions at different percentage levels*\n\n---\n\n## **Q3: Implications for Downstream Performance**\n\nWe agree with you that connecting our analysis to downstream performance is important. Previously, in Section **W2 & Q2: Limited Practical Implication & Improving Zero-Shot Performance**, we already demonstrated how our findings can be leveraged to improve zero-shot performance. Moreover, our findings can also be leveraged to improve ICL performance. We would like to first draw your attention to Section 5.2 and Fig. 4 in particular, where we present a considerable increase in logits lens accuracy achieved through low-rank denoising, which exploits the alignment structure of the hidden states with label unembedding vectors. We apply rank-5 denoising to Llama3-8B final layer hidden states. As a result, we can increase the final layer output accuracy averaged across six datasets as reported in Table 4. \n\nFurthermore, the high separability score of ICL hidden states presented in Section 5.2 and Fig. 2 implies that we can also improve ICL accuracy by training a classifier on several ICL hidden states and using its prediction on the hidden states of the remaining prompts to get the label, resulting in an ICL calibration label akin to the one proposed by [1]. Using this method, we also manage to significantly increase the average ICL accuracy as shown in Table 4.\n\n|Method|ICL|ICL+Denoising|ICL+Classifier|\n|------|-----|----------------|----------------|\n|Accuracy|76.20|78.59|85.47|\n*Table 4: Improving ICL accuracy using denoising or classifier as calibration methods* \n\nIn summary, **our findings lead to two methods to improve ICL performance that are lightweight, effective (the low-rank denoising method is also supervision-free), and can be seamlessly integrated into any ICL scenario involving text classification.**\n\n\n# References\n\n[1] Cho, Hakaze, et al. \"Token-based decision criteria are suboptimal in in-context learning.\""}]

Author rebuttal:

Meta-review: This paper studies the mechanisms underlying the in-context learning capabilities of transformers.
Focusing on in-context classification tasks, the authors characterize the layer-wise dynamics by revealing the two evolution stages of the hidden states, corresponding to separability and alignment of query hidden states.
In particular, it is shown that separability happens in early layers and alignment develops in later layers, and certain attention heads with specific functioning contribute to such phenomena.
This provides a unified view of the role of attention heads and task vectors in the context of in-context learning.

The paper is well-written and the idea is interesting and clearly illustrated.
The authors have derived novel insights into the cross-layer evolution of the geometric properties of the hidden states, by providing both empirical evidence and theoretical justification.
The findings provide explanations of ICL from a geometric perspective, which are valuable for understanding the inner workings of transformers.

The authors are encouraged to incorporate the additional results obtained during rebuttal into the final version of the paper.
Also, as noted by the reviewers, some missing references on related works about understanding cross-layer dynamics of transformers for ICL should be included to provide a more comprehensive background.
Decision: Accept (poster)