Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective

Thank you for your detailed and thoughtful review. We greatly appreciate that you find “the authors make great work at demonstrating the emergence of the task encoding and task decoding mechanisms” and that you consider “the synthetic experiments … and the real LLM experiments to be well-designed and clearly executed.” We also appreciate your suggestion to add detailed clarifications and cite relevant works such as Olsson et al., which provided valuable feedback for improving our manuscript.

We now address the specific questions that the reviewer raised their review.

Provide precise definitions of the Task Decodability and patching procedure

Common Point #3 Raised by Reviewer RjY9, UVog

We define the patching procedure and task decodability more precisely below and will update the manuscript for the camera ready accordingly.

1. Task Decodabilty

Let $\mathcal{T}$ be the set of latent tasks. For each task $z \in \mathcal{T}$, sample N=100 datapoints $\{(x_i, y_i)\} _ {i=1}^N$ and collect the intermediate representations $\{(h_i)\} _ {i=1}^N$ from a chosen layer (e.g., the token embedding immediately after $x_i$). Label each representation $h_i$ with the corresponding task $z$. This yields a set $\mathcal{S} = \bigcup_{z \in \mathcal{T}} \{ (h_i, z) \mid i = 1, \dots, N \}$ , over all tasks $z$.

Given a query point $(h_i, z)$, we exclude $(h_i, z)$ from $\mathcal{S}$ (i.e., $\mathcal{S} \setminus(h_i,z)$) and find its $k$ nearest neighbors (with $k=10$) in the remaining set. We then use majority voting on those neighbors’ task labels to produce a predicted label $\hat{z}$. If $\hat{z} = z$, classification for the task label is correct. The TD score at this layer is the fraction of query points classified correctly:

$\text{TD} \coloneqq \frac{1}{|\mathcal{D}|}\sum_{(h_i, z)\in \mathcal{D}} \mathbf{1}\bigl[\hat{z} = z\bigr].$

2. Patching Procedure

We collect representations $\{h_i\} _ {i=1}^n$ at selected position and layer for set of in-context examples $\{(x_i,y_i)\} _ {i=1}^n$ by passing the prompt through the model and recording the activations at a specified layer after $x_i$. We then average these activations to obtain a single “task” representation $h_{\text{task}} = \frac{1}{n}\sum_i h_i$. For a new query $x_{\text{query}}$, we add $h_{\text{task}}$ into the same layer’s activations—effectively replacing the representation after $x_{\text{query}}$—and let the model continue forward.

Why does the model show maximal representation separation at layer 5 in the figure 2?

Thank you for noting that detail, as it has been a common point of interest for the reviewers (Refer to Common Point 1). Indeed, Figure 11 in the appendix shows that the separation of representations by concept emerges most distinctly at layer 5, which is why we focus on that layer in the main Figure 2.

Suggestions for citing several papers(Olsson, Catherine, et al. (2022) and Sivaramakrishnan, et al.(2023))

Thank you for these suggestions, as they seem very relevant. We will add these to our discussion of related works in the final camera ready.

Suggestions for adding a limitation about the discrepancy between the simple setup and LLMs for tasks with shared support.

We agree this is a notable limitation of our simpler mixture-of-sparse-linear-regression framework. Perhaps, the large-scale pretraining in LLMs allows them to handle overlapping task supports more effectively, whereas our synthetic setup sometimes struggles to separate representations for tasks with shared support. We will add this in our limitations section for our camera ready version.

How many in-context example do you give to the TestMSE loss in Figure 2?

The TestMSE in Figure 2[left] is mean squared error over sequence of 20 in-context examples. Specifically, TestMSE $\epsilon$ is given as $\epsilon = E[\frac{1}{20} \sum_{i=1}^{20} \left( y_i - \hat{y}_i \right)^2]$. We will clarify this point in the followup.

Why does self-pertubation improve performance at convergence (Line 206, right)?

Figure 7 that the reviewer is referring to shows test MSE (y-axis) vs. number of demonstrations (x-axis). Because the task dimension is 4, having fewer than 4 demonstrations leaves uncertainty about the underlying basis. In this case, self-perturbation (patching the task vector) provides extra information about the basis and thus improves performance in that regime. That’s precisely why once the number of demonstrations matches the task dimension (around 4), the model identifies the task basis, and patching offers little to no additional benefit.

In the definition of the TD score, how do you define the set of representations over which you select the neighbors? Also what is N=100?

We define the set of representations by first sampling 100 datapoints (N=100) from each task and collecting their representations. We then compute the TD score by performing kNN classification on these representations.

We are grateful for your careful review. You highlight a key point that “the different encoding progress for different tasks is really interesting,” which resonates strongly with our primary objective of illustrating how task encodings develop and shape downstream ICL performance.

Below, we address your comments and questions.

Why does the model learn to encode B1 first in the synthetic experiment?

Although the four bases are theoretically equally difficult (same rank), random factors like data order and seed can cause the model to decode one basis first. As shown in Fig. 8, different seeds lead to different bases being prioritized.

How does our observation explain or contradict Yang et al. 2025?

We agree that the formation of distinct task vectors depends heavily on task formalism (how distinct task are defined). When the tasks share a common structure (e.g., the variant of regression tasks in Yang et al. [1] and Von Oswald et al. [2]), the model has little incentive to learn distinct representations, and simple attention-based algorithm[2] may suffice. This aligns with the Kim et al. [3], who report the underlying common structure between the aforementioned regression tasks.

Our synthetic sparse regression experiments are designed with algorithms that differ across tasks(basis), effectively enforcing a algorithmic diversifications and formation of distinct task vectors(albeit with varying quality). From these experiments, we report that the encoding and decoding algorithms emerge simultaneously, and that the quality of the learned encoding (measured by TD score) correlates with ICL performance.

We acknowledge that real-world tasks lie on a spectrum: some share overlapping structures, while others require truly orthogonal algorithms. This naturally places them between the two extremes of having fully separate versus completely shared task representations. Nonetheless, our key takeaway is that when tasks are sufficiently distinct, the quality of the learned task encodings can serve as useful indicator of ICL performance.

[1] Yang, Liu, et al. (2025).

[2] Von Oswald, Johannes, et al. (2023).

[3] Kim, Jaeyeon, et al. (2024).

What is the criterion for layer selection for Fig. 2?

Thank you for highlighting this point—it aligns with a common reviewer concern (Refer to Common Point #1 in RjY9’s rebuttal). As shown in Fig. 11, concept separation is most distinct at layer 5, which is why we select it in Fig. 2.

In Table 1, prompting yields a high TD score but poor ICL performance on the Bitwise Arithmetic task. How can this observation be better explained?

While prompting (i.e., modifying only the input text) leads to strong TD scores, it does not always result in correspondingly high accuracy—particularly compared to finetuning cases. As discussed in Appendix G, we interpret this results with caution: high decodability observed in the early layers during prompting may not reflect meaningful task encoding-decoding. Instead, it could stem from spurious correlations introduced by differences in the input token distribution. The bitwise arithmetic experiments seem sensitive to these correlations, which can inflate decodability as shown in Figure 21 on the right. Nevertheless, prompting still consistently produces higher TD scores and accuracy than the pretrained model.

Why does learning the task encoding by finetuning only the early layers improve performance? Why is learning encoding more important than decoding?

That’s an insightful question. Our hypothesis is that, since the LLM is pretrained on vast amounts of data, it often already has the “decoding” circuits (e.g., it can do XOR/AND) and mainly needs the right trigger(encoding signal) to activate them. Thus, directly changing the earlier layers (task encoding) can often be critical. However, for truly novel tasks that the model hasn’t seen (e.g., coding in an esoteric programming language [1]), it may need to learn both the task representation and the decoding algorithm from scratch.

[1] https://en.wikipedia.org/wiki/Esoteric_programming_language

Two messages on emergence of task vectors and relationship between TD and performance.

Our paper indeed focuses on two key findings: (1) how task vectors emerge during pretraining and (2) the quantification of task vector quality. While these findings may initially appear distinct, they form a cohesive narrative as follows.

We first study the direct phenomenology of task vector emergence with the synthetic experiments and establish the foundation of our task encoding-decoding framework. Then, for the second point, we demonstrate how this framework is practically useful by using task decodability to predict downstream ICL performance in natural setting—bridging mechanistic insights with applied relevance.

Therefore, we believe the two claims offer a comprehensive encoding-decoding framework for understanding and evaluating ICL capabilities.

Thank you for your positive and thoughtful review. We appreciate that you see our “theoretical logic as comprehensive, extending from experimental observations in regression tasks to natural language experiments,” and that you find our encoder-decoder framework provides “new theoretical foundations and empirical support for understanding the performance mechanisms and optimization methods of large language models in ICL.”

Below, we address your comments and questions.

The task complexity adopted in the paper is relatively low.

Common Point #2 Raised by Reviewer RjY9, EJ7t

We agree that our selected tasks are relatively simple compared to MMLU for example, but we intentionally focus on standard tasks (e.g., Oswald et al., Garg et al., Chen et al.) to rigorously study the emergence and effectiveness of task vectors using clearly distinguishable concepts (e.g., AND/OR in Bitwise Arithmetic, Pronoun/Verb in POS tagging).

For more complex settings in our synthetic setup, we experimented with higher numbers of bases and overlapping/shared bases between tasks (see Appendix D.1 and D.2). Interestingly, we observed that models tend to develop shared algorithms for overlapping concepts while distinctly separating non-overlapping ones—driving significant improvements in ICL performance.

[1] Von Oswald, Johannes, et al. "Transformers learn in-context by gradient descent." ICML. PMLR, 2023.

[2] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." NeurIPS 35 (2022): 30583-30598.

[3] Chen, Xingwu, Lei Zhao, and Difan Zou. "How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression." arXiv preprint arXiv:2408.04532 (2024).

What is the criterion for layer selection for task vector analysis and how is it related to Todd et al., 2023?

We refer the reviewer to our rebuttal to Reviewer RjY9 (under Common Points #1) for how we select the layer based on the highest TD score.

This approach is aligned with Todd et al. 2024. Todd et al. observe that task vectors at different layers result in varying ICL performances (Figure 2c in their paper). Building on this, we propose TD score to quantify the “quality” of task vectors. We demonstrate in Section 4 that (1) TD score is predictive of downstream ICL performance directly and (2) both finetuning and prompting for ICL can be interpreted under this framework of an encoder-decoder.

What data and procedure is used for UMAP in Figure 2?

To generate the UMAP in Figure 2, we randomly draw 100 samples for each basis. Next, we collected $y$-token representations of 5th layer at 20th demonstrations and plotted the UMAP with the parameters of n_neighbors=15 and min_dist=0.1. This specific layer was selected based on the maximum TD score (See Common Point #1 under Reviewer RjY9 and Figures 10 and 11 in the manuscript).

Minor improvements in supplementary presentation

We appreciate the feedback and will refine the clarity of our supplementary results for the camera-ready version.

In Figure 7 (a), why do the solid and dashed lines coincide when the number of samples in the first column increases, but not in the first row?

We hypothesize that the phenomenon occurs because the decoding algorithm for B2,B3,B4 may have developed after the 5th layer at epoch 20. This allows perturbations to B1 introduced at the 5th layer to be corrected as there are additional encoding layers before decoding occurs, when sufficient demonstrations are provided.

In contrast, the decoding algorithm for B1 appears to be implemented directly at the 5th layer, as evidenced by the separated representations visible in the UMAP visualization in Figure 10. Consequently, the model cannot be recover from perturbations from B1 to B2, B3, or B4 at the 5th layer.

In Figure 7, why is the effect eliminated as the number of samples increases?

This is an example of when LLMs can recover the latent task against the adversarial task vector perturbation (patching wrong task vector). Previous works by Usman et al. [1]. and Chen et al. [2] show that higher number of in context examples can give stronger ICL signal and override the effect of adversarial demonstration attacks. In line with these findings, we hypothesize that a similar mechanism operates in our experiment (Figure 7), mitigating the impact of adversarial task vector perturbations. When there is uncertainty about the underlying task basis (# of demonstrations << task dimension), both self-perturbation and adversarial perturbation are more effective.

[1] Anwar, Usman, et al. "Adversarial robustness of in-context learning in transformers for linear regression." arXiv preprint arXiv:2411.05189 (2024).

[2] Cheng, Chen, et al. "Exploring the robustness of in-context learning with noisy labels." ICASSP 2025-2025 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2025.

Thank you for your positive and thoughtful review. We appreciate your recognition of Task Decodability as “a compelling new measure” and that our activation patching analysis “provides convincing evidence [of transformers’] encoder-decoder structure.” We also value your suggestion to generalize to Mamba, and address your comments and questions below.

What is the criterion for layer selection for task vector analysis?

Common Point #1 Raised by Reviewer bnoD, RjY9, UVog

We select the layer with the highest TD score as the best encoding of a task. For the synthetic task (Section 3.3), we computed TD scores across all layers (Figure 10) and chose the fifth layer (layer index 4 in zero-based Figure 11). UMAP visualizations in Figure 11 confirm that representations become separable precisely when the TD score peaks.

For clarity, figures illustrating TD scores and UMAP visualizations are:

• Synthetic tasks : Figure 10, Figure 11
• Natural language tasks : Figure 16

Does the work generalize to various other architectures like Mamba?

We explored Mamba-2 8B and found our insights to generalize. We conducted the same studies on natural tasks on the pretrained Mamba-2 8B model (results in https://sites.google.com/view/icl-encoder-decoder/home). The results show that the TD score correlates with ICL accuracy (Figure A), and separable representations emerges around layer 33 (Figure B). These results suggest that our encoder–decoder framework and the predictive power of TD score on ICL performance do indeed generalize beyond standard transformer-based models to recurrent models.

Selected natural tasks are simple, so can you try on MMLU per subject?

We refer the reviewer to Common Point #2 in rebuttal for reviewer EJ7T. In summary, in benchmark like MMLU the boundaries between tasks are inherently blurred, making it challenging to rigorously study the encoder-decoder framework compared to our controlled experiments.

What differentiates this work from previous studies on mechanistic interpretability of ICL, such as Core et al. (2024)?

We concisely compare our contributions from prior work:

1. Different Task Setup
   • Core et al. (2024): Adopts a relatively homogeneous Markov chain task and examines how the system transitions between different algorithmic stages (unigram/bigram retrieval and inference).
   • Our Work: Focuses on the heterogeneous tasks—a setting more reflective of how LLMs are trained and used in practice (where tasks vary substantially). We analyze emergence of distinct task vectors during pretraining.
2. Measured quality of task-specific vectors and its relation to downstream ICL peformance
   • Hendel et al. & Todd et al.: Demonstrate that large language models have task vectors but leave open questions about how these vectors emerge and why their effectiveness varies across tasks.
   • Our Work: Addresses these open questions by analyzing when and how “task-specific” representations emerge, and how their “quality”(TD score) can be measured and used to predict downstream ICL performance.

Try Kim et al. experimental setup with different regression tasks

As per the reviewer’s suggestion, we replicated the regression task mixture experiment from Kim et al., training a small transformer on linear, quadratic, sparse linear, and leaky ReLU regression tasks (results available at https://sites.google.com/view/icl-encoder-decoder/home). The model successfully learns the different tasks but fails quadratic regression (Figures A, C). TD scores and UMAP visualizations (Figures B, D, E) indicate no clear separation of representations, suggesting the absence of distinct task vectors. We interpret these findings as evidence that the transformer employs similar underlying algorithms across these tasks (e.g., implementing leaky ReLU via linear regression with sign adjustments). This is consistent with Kim et al.’s observation that a common structure between their regression tasks facilitates learning (Figure 6 and Section 4.2 in their paper).

Can you clarify the notation and definition in supplementary material for Section B?

We will improve the decodability definition like the following in Appendix B.

B.3 (Decodability). For a given decoder $G: \mathbb{R}^{d_{emb}} \rightarrow \mathcal{Z}$ and a specific latent variable $z$, the decodability measures how accurately the correct latent variable is inferred from representations. Representation is distributed as $E(z,D)$, and the inferred latent variable is $\hat{z} \equiv G(E(z,D))$ .

1. One-hot Accuracy : $A_{\text{1-hot}}(z) = E[\mathbf{1}\bigl[\hat{z} = z\bigr]]$
2. f-divergence : $A_{f}(z) = D_{f}(\hat{z} \parallel z)$, where $f$ is some f-divergence metric

Our Task Decodability is the one-hot accuracy with decoder of k-nearest neighbor majority-voting algorithm.

For a definition of Task Decodability, we refer reviewer to Common Point #3 in rebuttal to Reviewer EJ7T.