Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations

We thank Reviewer pehy for the thoughtful comments and for recognizing the significance and insights of our analysis. Below, we address each of the reviewer’s questions in detail:

Q1. Parameter Analysis in Practical LLMs

We agree that directly analyzing the optimal parameter structure of practical LLMs is a valuable and exciting direction. However, it poses substantial challenges and is beyond the current scope of our work. Unlike our theoretical setting, which explicitly minimizes the ICL risk, real-world LLMs are pretrained on diverse corpora with objectives not specifically designed for ICL. As a result, the parameter patterns in pretrained models may reflect a mixture of capabilities with ICL being just one among many, making principled isolation and analysis particularly difficult. Nevertheless, we believe our analysis lays foundational insights that will aid future efforts in understanding and designing LLMs specialized for in-context learning.

We acknowledge this limitation in Appendix D.3, noting that our theoretical analysis is conducted under a simplified linear-attention setting. Nonetheless, we believe our results offer clear mechanistic insights into the emergence, function, and limitations of task vectors, and can guide future improvements to ICL in practice. Concretely, our contributions have the following practical implications:

• Our analysis of embedding concatenation helps identify the optimal layer for task vector extraction.

• The bijection task counterexample clarifies when and why task vectors fail, despite ICL success.

• Our study of attention weights helps determine the preferred ordering of in-context demonstrations.

• The validation of multi-vector injection highlights how task vectors can be extended to few-shot settings.

• Our critical point analysis sheds light on the conditions that encourage task vector emergence during fine-tuning.

We hope these insights will inspire future work on analyzing or even designing architectures that more explicitly support ICL-style learning.

Q2. Justification for the Block Diagonal Formulation

The block-diagonal formulation of the $V_l$ and $Q_l$ matrices is a widely adopted assumption in theoretical studies of ICL for transformer models, as it facilitates tractable analysis [1, 14, 20, 25]. Prior work by Ahn et al. [1] demonstrates that the global minimizer of single-layer linear-attention transformers indeed exhibits such block-diagonal structure. Although finding exact solutions for multi-layer transformers is more involved, it is reasonable to conjecture that similar structural patterns hold. Empirically, we observe that when optimizing the full matrices, gradient-based training also tends to converge to block-diagonal solutions.

Intuitively, given the high dimensionality of hidden states in modern LLMs, it is plausible to assume that the $x$ and $y$ components can be projected into orthogonal or nearly orthogonal subspaces. This motivates a decomposition of the projection matrices $V_l$ and $Q_l$ into two separate parts that operate independently on $x_i$ and $y_i$, which can be equivalently formulated as the block-diagonal structures.

We will clarify and expand upon this rationale in the revised manuscript.

Q3. Performance of TaskV-M

We agree that the performance gains of TaskV-M over TaskV are not dramatic across all ICL tasks. However, the goal of TaskV-M is not to surpass state-of-the-art ICL techniques, but to demonstrate that the task vector framework can be systematically extended by injecting multiple vectors simultaneously. This is especially valuable for complex tasks that inherently require higher-rank representations.

Our results on bijection tasks clearly validate this motivation: TaskV-M yields notable improvements over the standard TaskV method. For other simpler tasks, the marginal gains from TaskV-M suggest that the expressiveness of $W$ may not be the primary performance bottleneck. We believe these insights facilitate the design of future ICL and task vector methods.

Q4. Further Experimental Results

Due to rebuttal constraints, we are unfortunately unable to include additional figures. However, we reproduce below extended results from Table 1 that illustrate the failure of task vectors on bijection tasks across a broader range of LLMs and varying numbers of input demonstrations. These results further support our claims that:

• Task vectors systematically fail on bijection tasks, even when increasing the number of demonstrations.

• The failure is consistent across multiple model architectures, reinforcing that the issue stems from a fundamental expressiveness limitation rather than model-specific artifacts.

Below are the summarized performances for various LLMs on bijection tasks with $n = 10$ demonstrations. Left is ICL accuracy, and right is task vector accuracy.

Below are the results for $n = 20$ demonstrations.

We will include these results and further saliency analysis results in the revision.

We thank Reviewer QZNA for the constructive feedback and for recognizing the breadth and depth of our contributions. Below, we address each of the reviewer’s questions in detail:

Q1. The Strength of our Theoretical Results

In addition to the breadth of our empirical and theoretical contributions, we would like to highlight the technical depth and novelty of the proof techniques underlying our main results. Extending the analysis to pairwise and triplet-formatted demonstrations necessitated a careful construction and exploration of symmetry in both the input and parameter spaces, going beyond standard linear regression setups. To the best of our knowledge, our work provides the first rigorous theoretical framework that explains the emergence and functionality of task vectors, offering a principled understanding of their behavior, grounded in a formal analysis rather than empirical observation alone. We will make this point clearer in the revision.

Q2. Constraint on $\Lambda_1$

We agree with the reviewer that setting $\Lambda_1 = I$ would impede the model from performing the embedding concatenation necessary for ICL, as our analysis shows. More precisely, let $\Lambda_1 = \begin{pmatrix} a & b \\\\ c & d \end{pmatrix}$; then, at least one of $b$ or $c$ must be non-zero to enable meaningful interaction between the covariates and responses.

However, in Theorem 1, we characterize all possible critical points of the expected ICL risk, regardless of whether they solve the problem. Since $\Lambda_1 = I$ is indeed a valid critical point, even if meaningless, we intentionally do not exclude this case from our formal analysis. We will make this nuance clearer in the final version.

Q3. Layers for Embedding Concatenation

Indeed, embedding concatenation can occur at multiple layers in linear transformers. However, our theoretical focus is on the first layer because it plays a pivotal role in enabling subsequent layers to perform gradient-based updates. Once the covariate-response pairs $(x_i, y_i)$ are concatenated in the first layer, any additional mixing in later layers does not contribute further to task inference under the ICL risk, as it merely reprocesses already concatenated embeddings. Thus, only the first-layer concatenation directly facilitates gradient descent optimization.

Q4. $\alpha_1$ and $\alpha_2$

We apologize for the lack of clarity. $\alpha_1$ and $\alpha_2$ are not necessarily equal in practice. In eq. (10), we omit them to simplify notation. Since the effective coefficient matrix is given by $W' = \alpha_2 Y\beta (\alpha_1 X\beta)^\top$, the overall scale is determined by the product $\alpha_1 \alpha_2$, and their individual values do not affect the optimization dynamics or our conclusions. We will clarify this simplification in the revised manuscript.

Q5. Task Vector Injection Layers

The arrow-token hidden states after the first linear-attention layer are theoretically optimal for task vector extraction. This is because subsequent layers apply preconditioned gradient descent updates (i.e., on the Gram matrix $X^\top X$), which distort the original demonstration representations. Therefore, extracting task vectors after the first layer preserves a cleaner representation aligned with the base mapping $W$. In practical LLMs, task vectors are often extracted from deeper layers because the early transformer layers are responsible for preprocessing token representations.

Q6. Emergence of Task Vectors

We agree with the reviewer that Proposition 3 allows for the stationary solution $\lambda = 0$. However, in practice, we find this trivial solution is attained with almost zero probability. Due to random initialization and the presence of gradient noise in early training, small but nonzero $\lambda$ values emerge. Since these weights do not contribute to minimizing the ICL risk, they are not penalized during optimization (even under $L_2$ regularization), and thus persist during the pretraining stage.

Once present, these nonzero $\lambda$ values tend to converge to nearby nontrivial stationary points described in Proposition 3. In our experiments with linear transformers, we have not observed convergence to $\lambda = 0$ or even $\lambda \approx 0$. Moreover, as shown in Proposition 6, the presence of external stochasticity (e.g., dropout) can further promote the persistence and stabilization of these weights, making the emergence of task vectors more robust in practical settings. We will incorporate the discussion in the revision.

Q7. Empirical Validation of High-Rank $W$

We appreciate the reviewer’s interest in the bijection task and the implications of our multi-vector injection strategy. To empirically verify that injecting multiple task vectors leads to a higher-rank coefficient matrix $W$, we conduct experiments using a 2-layer linear-attention model with hidden dimension $d = 4$. We vary the number of injected task vectors and report the top singular values of $W$ ($\lambda\_1$ to $\lambda\_4$), along with the corresponding ICL risk.

These results clearly demonstrate that increasing the number of injected task vectors enhances the rank of $W$, thereby reducing the ICL risk, verifying our analysis. We will add this experiment in the revision, either in the main paper or the appendix, depending on space.

Q8. Regarding the Linear Combination Conjecture

We appreciate the opportunity to clarify a potential misunderstanding regarding what is being linearly combined. In our main conjecture, the task vector is formed by a linear combination of the hidden states corresponding to the input demonstrations, not their raw token embeddings. We do not claim that task vector formation excludes non-linearity; in fact, we believe non-linear pre- and post-processing layers are essential to enhance the ICL ability of LLMs.

Our conjecture posits that practical LLMs perform task vector formation and prediction in middle transformer layers using mechanisms akin to those observed in linear-attention models, specifically embedding concatenation followed by gradient-like updates. The input $(x_i, y_i)$ pairs in our linear model analysis correspond to the hidden states of those tokens after early-layer processing in real-world LLMs. Therefore, while the optimal extraction depth is 1 in our linear setting, it lies deeper in real-world LLMs. This is consistent with our cross-validation results, where the empirically optimal extraction layer immediately follows the layer displaying the strongest weighted summation pattern (see Figure 3b).

What we aim to emphasize with this conjecture is not the “linear combination” itself — an inherent property of attention mechanisms where hidden states are aggregated via attention-weighted summations over value projections, but rather the implication it carries: task vectors cannot encode the full information of the task. Instead, they represent at most a single, prototypical demonstration, limiting their intrinsic expressiveness and capacity to generalize to complex tasks.

Q9. Interpreting Figure 3

Figure 3 provides empirical support for our conjecture that practical LLMs adopt similar inner mechanisms to those observed in our linear-attention model analysis. Specifically, the inference stage of ICL consists of:

• Early layers perform representation preprocessing;

• Intermediate layers execute embedding concatenation across $(x_i, y_i)$ pairs;

• Task vector then emerges via attention-weighted summation over these embeddings;

• Final layers handle gradient-based updates and generate the prediction.

For a conceptual overview, please refer to Figure 1. Figure 3a and 3b visualize saliency maps that clearly show embedding concatenation and task vector formation patterns. In particular, we observe that each $y_i$ token attends to its corresponding $(x_i, y_i)$ pair (concatenation), and the final arrow token aggregates across all $y_i$ tokens (weighted summation), closely mirroring our main conjecture. Together, these findings suggest that real-world LLMs implement structurally similar learning procedures for ICL. We will add the above description in the revision.

We thank Reviewer tBHD for the insightful comments and for acknowledging the significance of our work in advancing the understanding of transformer models. Below, we address the reviewer’s questions in detail:

Q1. Regarding the Linear Combination Conjecture

We appreciate the opportunity to clarify a potential misunderstanding regarding what is being linearly combined. In our main conjecture, the task vector is formed by a linear combination of the hidden states corresponding to the input demonstrations, not their raw token embeddings. We do not claim that task vector formation excludes non-linearity; in fact, we believe non-linear pre- and post-processing layers are essential to enhance the ICL ability of LLMs.

Our conjecture posits that practical LLMs perform task vector formation and prediction in middle transformer layers using mechanisms akin to those observed in linear-attention models, specifically embedding concatenation followed by gradient-like updates. The input $(x_i, y_i)$ pairs in our linear model analysis correspond to the hidden states of those tokens after early-layer processing in real-world LLMs. This interpretation is further supported by our saliency analysis, which reveals embedding concatenation and aggregation patterns in the middle layers of pretrained LLMs.

What we aim to emphasize with this conjecture is not the “linear combination” itself — an inherent property of attention mechanisms where hidden states are aggregated via attention-weighted summations over value projections, but rather the implication it carries: task vectors cannot encode the full information of the task. Instead, they represent at most a single, prototypical demonstration, limiting their intrinsic expressiveness and capacity to generalize to complex tasks. We will clarify this in the revision.

Q2. Alternative Hypothesis for Task Vectors

To our knowledge, there is currently no established theoretical framework explaining how task vectors emerge and function. A commonly held belief is that task vectors directly encode the task mapping $f: \mathcal{X} \to \mathcal{Y}$ into a single vector [7]. However, this hypothesis fails to explain why task vectors underperform on simple bijection tasks, while standard ICL succeeds.

Our work provides a refinement of this view: task vectors encode limited information, essentially equivalent to a single demonstration, thus explaining their failure in tasks requiring higher-rank representations, such as bijections. We believe this insight is a substantial step toward understanding the internal mechanisms of transformer models. If we have misunderstood the reviewer’s comments, we would appreciate further clarification and would be happy to address them.

Q3. Availability of Running-Free Task Vector Extraction

Thank you for suggesting this creative approach. As noted above, since task vectors are formed via linear combinations of the model's internal hidden states, they inherently depend on processing the prompt through the model. Consequently, extracting task vectors without executing the model may not be feasible in general.

Regarding the weights used in forming these linear combinations, our analysis (see Figure 3c) shows that causal attention induces an increasing weighting pattern (i.e., later demonstrations contribute more to the task vector). This finding aligns with our practical observations.

Q4. Table 1 and Table 2

Thank you for pointing out this possible source of confusion. The settings for Table 1 and Table 2 are indeed different:

• Table 1: Task vectors are injected into zero-shot prompts (TV).

• Table 2: Task vectors are injected into few-shot prompts (TaskV).

Given this, it is not surprising that TaskV achieves higher performance on bijection tasks in Table 2 due to the extra information provided by few-shot demonstrations. In contrast, TaskV yields ~50% accuracy on bijection tasks in the zero-shot setting (Table 2), consistent with random guessing. In the revision, we will make this distinction more explicit.

Q5. The structure of $\mathcal{S}_P$

We appreciate the reviewer’s careful examination of Figure 2b. As observed, each $3 \times 3$ sub-matrix in $D_l$ contains only two non-zero entries: the top and bottom cells in the second column, besides the concatenation patterns represented by $\Lambda_1$. This structural pattern aligns precisely with the $\Lambda_4 \otimes \Lambda_5$ component described in Theorem 2, which plays a central role in task vector formation.

Q6. Blue bars in Figure 3c

Thank you for raising this point. In Figure 3c, the blue bars correspond to the left y-axis $\\|z_{\text{tv}}^{-i} - z_{\text{tv}}\\|$, where we examine the contribution of each individual demonstration to the final task vector. For each demonstration, we introduce a perturbation and compute the deviation of the resulting task vector $z_{\text{tv}}^{-i}$ from the original $z_{\text{tv}}$. The observed pattern closely matches the theoretically predicted weight distribution shown by the orange line, thereby validating our analysis. We will clarify this in the revision.

Q7. Redundancy Against Dropout

We would like to clarify that we do not claim dropout as the main cause of task vector emergence, but rather that it amplifies the tendency towards such redundancy. Importantly, the $\Lambda_4$ weights for task vector formation emerge naturally because of the random initialization and noisy gradients in early updates, even without dropout. As these weights do not contribute to the ICL risk, they are not penalized during optimization and persist as part of the model’s learned structure.

We adopt token-wise dropout mainly due to its theoretical simplicity, and to intuitively demonstrate how random information loss promotes redundancy. While we acknowledge that token-wise dropout is rare in practical LLMs, we believe the broader insight generalizes: any mechanism that perturbs hidden states during training can push toward such redundancy.

In this spirit, we mention positional encodings and normalization layers as alternative sources of perturbations in the hidden state space that introduce extra scaling and shifting. These facts may subtly distort hidden states and lead to the same conclusion.

Note that these analyses are not the core contribution of our paper, but are offered as a potential supporting mechanism that may promote the emergence of task vectors. We will revise the manuscript to make this distinction clearer.

Q8. Expected ICL Risk vs. Auto-regressive Loss

We appreciate this insightful distinction. Our current analysis focuses on the ICL setting commonly used in recent works on task vectors [7, 12, 16], where prediction is modeled as a single-token generation task. It would be interesting to investigate whether our analysis could be extended to the auto-regressive loss.

Firstly, we consider combining the causal mask with the current expected ICL risk. We believe the current techniques used in our proof can be extended to causal linear attention after minor modifications. In particular, we expect that the structures of $A_l$, $B_l$, and $C_l$ remain unchanged, while $D_l$ becomes block-diagonal with upper-triangular sub-blocks. Therefore, our main conclusions are expected to hold under causal attention.

Further extension to auto-regressive loss is equivalent to optimizing the ICL risk over varying numbers of demonstrations in expectation. Since our current results characterize the critical points for arbitrary $n$, we believe they also provide insight into the stationary behavior of such auto-regressive training. One expected difference lies in $\mathcal{S}_P$, such that the diagonal sub-blocks may no longer be identical.

We have conducted preliminary experiments using causal linear attention, and the learned parameters match our analysis above. Due to rebuttal constraints, we cannot include additional plots, but we will add these results and visualizations in the revised version.

Q9. Related Works and Other Suggestions

We thank the reviewer for these helpful suggestions. In response, we have: (1) renamed the model to improve clarity; (2) restructured Section 5 for better organization and narrative flow; and (3) added the missing citation. Due to space constraints, we initially incorporated discussions of related work into the introduction. Following the reviewer’s suggestion, we will reorganize these discussions into a dedicated Related Work section in the revised version.

Q10. Limitations of our Work

We want to clarify that the limitation discussion is provided in Appendix D.3 due to limited space. Our analysis is conducted on simplified settings, particularly linear-attention models, which may not capture the full complexity of real-world LLMs. However, the primary goal of our work is not to exhaustively model practical LLMs, but to provide clear, mechanistic insights into the emergence, functionality, and limitations of task vectors in ICL. Our theoretical results offer several actionable contributions:

• Our analysis of embedding concatenation helps identify the optimal layer for task vector extraction.

• The bijection task counterexample clarifies when and why task vectors fail, despite ICL success.

• Our study of attention weights helps determine the preferred ordering of in-context demonstrations.

• The validation of multi-vector injection highlights how task vectors can be extended to few-shot settings.

• Our critical point analysis sheds light on the conditions that encourage task vector emergence during fine-tuning.

We believe that, while these insights are grounded in a tractable, simplified theoretical model, the conclusions drawn from it provide valuable guidance for better understanding and improving the ICL capabilities of modern transformer-based models. We will add these discussions and move the limitation section to the main paper in the revision.

We thank Reviewer eQNr for the thoughtful feedback and for recognizing the importance of our contribution to the understanding of ICL. Below, we address the reviewer’s comments and questions in detail:

Q1. Relation Between $d_p$ and $n$

We apologize for the confusion here. The value of $d_p$ depends on the structure of the input prompt, as summarized below:

We will improve the presentation by clarifying the explicit relations between $n$ and $d_p$ to avoid confusion at their first appearance.

Q2. Typo: $\tilde{R}_i$ Should Be $R_i$

Thank you for catching this typo. We have corrected it in the revised version.

Q3. Clarification on Table 1 and Table 2

The performance of ICL (in Table 1) and Baseline (in Table 2) should not be directly compared, as they are based on different prompt settings. Specifically:

• ICL in Table 1 uses many-shot prompts ($n=10$ demonstrations), while task vectors are injected into zero-shot prompts. As expected, task vectors underperform standard ICL, since they serve primarily as an acceleration mechanism.

• Baseline in Table 2 uses few-shot prompts (ranging from 0 to 4 demonstrations), the same setting for task vector injection. These results show that injecting task vectors improves ICL performance under these constrained settings.

We will make this difference more explicit in the revision.

Q4. The Multi-Vector Variant

Thank you for highlighting the alternative method of extracting multiple task vectors from a single prompt. This approach is indeed consistent with our theoretical findings. As discussed in Proposition 3, the task vector weights that emerge at each arrow ($\to$) token are approximately orthonormal, suggesting they encode distinct information subsets and can be simultaneously injected to enhance model performance (e.g., by increasing the rank of the induced coefficient matrix $W$).

We have also tried this method before. Below is a comparison between the current multi-vector method (TaskV-M) and this single-prompt variant (TaskV-MS) under a 4-shot setting:

While TaskV-MS also delivers strong performance, it slightly underperforms TaskV-M. We believe this is due to the causal attention mechanism in real LLMs, where earlier arrow tokens can only aggregate information from a subset of demonstrations. Nonetheless, TaskV-MS is a promising alternative for accelerating inference. We will add these results as part of the revised paper.

We thank Reviewer jCtd for the constructive comments and for acknowledging the significance of our theoretical analysis and its insights into understanding task vectors. Below, we respond to the reviewer’s specific concerns:

Q1. Discussion of Potential Limitations

We want to clarify that the limitation discussion is provided in Appendix D.3 due to limited space. We acknowledge that our theoretical analysis focuses on a simplified linear-attention transformer architecture, which may not fully encapsulate the complexity of real-world LLMs. However, we provide several strong empirical observations that support the practical relevance of our theory: (1) consistent failure of task vectors on bijection tasks, (2) matching saliency patterns in real LLMs that align with our theoretical predictions, and (3) consistent improvements from multi-vector injection. Together, these results offer meaningful insight into the emergence, functionality, and limitations of task vectors in practical transformer models.

Q2. Heavy Use of Mathematical Notation

To improve accessibility, we will add a dedicated notation table in the appendix to clearly define all symbols used throughout the paper. Additionally, we emphasize key takeaways and theoretical insights in bold text to help readers grasp the main conclusions. We will also add summaries and interpretations in plain language to help readers understand the key ideas easily.

Q3. Extension to Multimodal Settings

While our current analysis primarily focuses on text-based LLMs, it naturally extends to any ICL scenario where the input follows a triplet-formatted structure (e.g., "a $\to$ b, c $\to$ d, e $\to$"). This includes certain multimodal models that adopt similar prompting schemes, such as the one proposed in [A]. In particular, we believe our multi-vector injection strategy applies to multimodal task vectors. Formally extending our theoretical framework to the multimodal setting is an exciting and valuable direction. We have noted this in our future work.

[A] Huang Brandon, et al. Multimodal task vectors enable many-shot multimodal in-context learning, 2024.