Dear Reviewer LiD7,

Thank you for your insightful review of our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.

---

W1 & Q1: Limitations of solely focusing on classification tasks

We 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.)

Table 1: Layer-wise trend of singular alignment and separability score for the generation task: ICL (8-shot) and zero-shot The 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.

---

W2 & Q2: Limited Practical Implication & Improving Zero-Shot Performance

We 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).

---

W3 Part 1: Direct Evidence for the Geometric Significance of IHs and PTHs

We 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.

Table 2: Dynamics in separability and alignment measure after task vector injection using PTH and IH outputs as task vectors

The 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.

---

W3 Part 2: Location and Identification of IHs and PTHs

Per 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.

---

Q3: Implications for Downstream Performance

We 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.

Furthermore, 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.

In 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.

References

[1] Cho, Hakaze, et al. "Token-based decision criteria are suboptimal in in-context learning."

Dear Reviewer 4WZT,

Thank you for your constructive feedback on our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.

---

W1 Part 1 and W2: Clarifications about the attention head ablation approach

We 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.

In 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.

In 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:

The 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.

---

W1 Part 2: Using mean ablation as an alternative approach

We 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.

Table 2: Impacts on separability score and singular alignment when IH, PTH, or random heads are ablated

These results confirm our hypothesis: PTHs primarily affect separability, while IHs primarily affect alignment.

---

Q1: PTHs and mean-based alignment

As 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.

---

Q2: Regarding the referenced paper

Thank you for bringing up this relevant work. We believe the differences in findings can be attributed to two main factors:

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.

To 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:

The results confirm that the outputs of our IHs encode more task-specific semantics and serve as stronger task vectors.

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.

---

References

[1] Todd, Eric, et al. "Function vectors in large language models."

Dear Reviewer kMpH,

Thank you for your careful review of our manuscript. Below we respond to the weaknesses and questions you raised.

W1 & W2 part 1: Insufficient discussion of task vector literature

We 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.

Following 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:

1. 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.
2. 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.
3. 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.

W3: ICA as an alternative method for interpreting hidden state dynamics

We agree that ICA is a viable method for extracting semantic directions. However, in our context, our SVD-based approach is preferable for two reasons.

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.

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.

Nonetheless, per your suggestion, we repeated the semantic filtering and retention experiment on SST-2 with Llama3-8B. Below is the result at layer 25:

It shows that SVD and ICA yield similar outcomes; SVD even decodes the SST-2-relevant "neutral" token.

W4 & Q3 Part 2: Limitations of focusing solely on classification tasks

Please 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.

Q1: Comparison between [7]'s findings and ours

Thank you for highlighting this relevant. Below, we outline two key distinctions.

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.

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

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.

Q2, Q3 Part 1, W2 Part 2: Hendel et al.'s task vectors—linearity, location, and adaptability

We 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.

(Before layer 14, accuracy is ~0)

The 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

References

Follows the paper list you mentioned in the review...

[21] Olsson, C. et al. "In-context learning and induction heads."

[22] Singh, A. K. et al. "What needs to go right for an induction head?..."

[23] Cho, H. et al. "Revisiting in-context learning inference circuit..."

[24] Crosbie, J., Shutova, E. "Induction heads as an essential mechanism..."

[25] Song, J. et al. "Out-of-distribution generalization via composition..."

[26] Jin, Z. et al. "Cutting off the head ends the conflict..."

[27] Bansal, H. et al. "Rethinking the role of scale for in-context learning..."

[28] Yang, L. et al. "Yang, Liu, et al. "Task vectors in in-context learning: Emergence, formation, and benefit."

[29] Neyshabur, B. et al. "What is being transferred in transfer learning?"

Dear Authors,

Thank you for your response.

1. Task Naming First, I would like to clarify that the task your paper focuses on should more accurately be termed “single-token bi-classification” rather than simply “classification.” This distinction is crucial to inform the reader properly. Other tasks such as country-capital are also regarded as a classification over a set of classes, and calling one task “classification” while excluding the other introduces confusion.

2. Interpretation of Phases I respectfully disagree with your interpretation. My intention was to point out that your so-called “separability phase” (or clustering phase) corresponds to a high-level information extraction step—specifically, the extraction of low-dimensional task-relevant information. From an information-theoretic perspective, this aligns with the task vector formulation phase in [31, 15] or “Compression” phase described in [17].

Similarly, what you refer to as the “alignment phase”—which I continue to interpret as the task execution phase (involving the query and the task representation)—aligns with the task execution phase in [31, 15] or the “Expansion” phase in [17], where the system executes the task after compressing relevant information.

3. On Task Vectors and Experimental Observations The example I previously gave—“I don’t like it. Answer: positive. I like it. Answer:”—was intended to highlight a key distinction: in single-token bi-classification, the task vector defined in [31] (In-context learning creates task vectors) does not appear to exist. In [31], the task vector is such that, when attending to any query token in the latent space, it can produce the correct output with high probability in later layers.

However, based on my extensive experimental experience, I could not find such a vector in the bi-classification setting. That is precisely why bi-classification tasks are absent in [15] and [31]. (I wonder why you claimed [15] includes it.)

That said, the initial clustering phase does exist, even though it does not result in the formation of a task vector in the sense of [31]. This is not surprising—transformers are not inherently designed to form such task vectors; it occurs only in specific cases as studied in [31] etc.

4. Missing Literature: Wang et al. [34] and Beyond Importantly, the information aggregation process, which corresponding to your first phases during the bi-classification task—without forming a task vector—is studied in detail in Wang et al. [34] (EMNLP Best Paper 2023). This paper is currently missing from your manuscript, but it is highly relevant. Wang et al. show that during the clustering phase, label-word anchors encode the task information, and that the subsequent “execution phase” performs the classification based on these anchors. This supports the interpretation that your “alignment phase” corresponds to their “execution phase” (label prediction).

Notably, Wang et al. [34], as well as [31], [15], and [17], all observe that the first phase ends around similar layers—even though bi-classification in Wang et al. does not result in a task vector as defined in [31]. Therefore, your argument that these phases collapse into one in your setting is, in my view, not persuasive and lack of understanding.

5. Gaps in Literature Review and ICL Understanding Beyond the task vector and bi-classification literature, your manuscript also lacks adequate discussion of core In-Context Learning (ICL) works, especially [35] and [36]. This is particularly important given your attempt of developing a simple theory of ICL. In fact, the latent variable in the BMA perspective corresponds precisely to the first phase—recognition of the task identity—as emphasized in [15], [17], [31], and [34].

To summarize, my concerns regarding the paper remain. 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. I maintain my position that the task vector should reflect high-level task information extraction, which enables functionality—not low-level execution once that information is already extracted.

As for NeurIPS policy, I only regarded it as an arXiv paper, also with flaws, imperfection and missing discussions. The heavy issues in [17] do not influence my judgment of this manuscript.

Apologies for the delay in my response due to crazy schedule. Also wishing you a pleasant day!

Warm regards,

Reviewer kMpH

---

References

[34] Wang et al. Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning. EMNLP Best Paper 2023.

[35] Xie et al. An Explanation of In-Context Learning as Implicit Bayesian Inference

[36] Zhang et al. What and How Does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization

Dear Reviewer Ua9F,

Thank you for your thoughtful review of our manuscript, which we greatly appreciate. Below, we respond to the weaknesses and questions you raised.

---

W1: Limitations of solely focusing on classification tasks

We 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.)

These 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.

---

W2 & Q2: Explanation for why PTHs improve separability, IHs improve alignment

Below 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.

Second, 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.

---

W3: Implications for ICL performance

We 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.

Furthermore, 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.

In 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.

---

Q1: Where PTHs attend

Yes, 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.

---

References

[1] Marks, Samuel, and Max Tegmark. "The geometry of truth: Emergent linear structure in large language model representations of true/false datasets."

[2] Park, Kiho, Yo Joong Choe, and Victor Veitch. "The linear representation hypothesis and the geometry of large language models."

[3] Tigges, Curt, et al. "Linear representations of sentiment in large language models."

[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."

[5] Halawi, Danny, Jean-Stanislas Denain, and Jacob Steinhardt. "Overthinking the truth: Understanding how language models process false demonstrations."

[6] Cho, Hakaze, et al. "Token-based decision criteria are suboptimal in in-context learning."