Title: Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations
Venue: NeurIPS 2025
Accepted: no
Keywords: transformer

in-context learning

task vector
Abstract: Task vectors offer a compelling mechanism for accelerating inference in in-context learning (ICL) by distilling task-specific information into a single, reusable representation. Despite their empirical success, the underlying principles governing their emergence and functionality remain unclear. This work proposes the *Linear Combination Conjecture*, positing that task vectors act as single in-context demonstrations formed through linear combinations of the original ones. We provide both theoretical and empirical support for this conjecture. First, we show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis. Next, we predict the failure of task vectors on representing high-rank mappings and confirm this on practical LLMs. Our findings are further validated through saliency analyses and parameter visualization, suggesting an enhancement of task vectors by injecting multiple ones into few-shot prompts. Together, our results advance the understanding of task vectors and shed light on the mechanisms underlying ICL in transformer-based models.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper attempts to provide a theoretical observation of how task vectors of ICL emerge within self-attention. Specifically, the authors simulate a linear regression task and augment it to an ICL-style linear regression form (triplet-formatted prompts). On these tasks, they identify low-loss solutions of self-attention in which the task vector can be described as a weighted average of the ICL demonstrations. As a limitation of such task vector, they investigate and successfully predict the failure of ICL (in both synthetic and real-world scenarios) under bijective mappings (or full-rank mappings). Furthermore, they present several empirical observations and applications in real LLMs, aiming to transfer their theoretical findings into practical settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The theoretical simulations presented by the authors are based on reasonable assumptions and yield credible results. Given that empirical findings on ICL have indicated the critical role of these structured tokens (such as the $\\rightarrow$ in this paper) in LLMs, theoretically investigating these tokens is valuable. Moreover, the theoretical results align well with existing empirical evidence (both of these based on feature correlation patterns (e.g., saliency maps or induction heads) or task vectors) thus offering potential theoretical support for how ICL emerges in LLMs.\n\n2. The discussion of bijective tasks (Proposition 4) is both reasonable and highly insightful. Intuitively, I believe that tasks where the input space is significantly larger than the output space (e.g., classification) operate under fundamentally different ICL mechanisms compared to bijective tasks (e.g., mapping countries to their capitals). This paper successfully brings attention to this issue and uses its theoretical framework to make predictions and offer explanations, which is highly commendable.\n\n3. The paper is well-written and easy to follow.\n\n**Weaknesses**\n\n1. The paper did not examine how the theoretically optimal parameter patterns are reflected in the parameters of real-world LLMs. Instead, the paper only investigated the alignment between the saliency patterns in actual LLMs and their theoretical results, which is insufficient to uniquely confirm the correctness of the theoretical framework. While I acknowledge that analyzing the parameters of real LLMs can be very challenging, the lack of such analysis inevitably weakens the credibility of the paper’s conclusions. Additionally, it remains unclear how the simulated simple linear regression tasks map onto LLMs trained on real-world data. Thus, although the authors propose an elegant theoretical framework, applying it to real-world LLMs still requires considerable effort.\n\n2. This paper seem to impose strong conditions in constructing the theoretical framework. For instance, in Equation (6), they assume that $V_l$ and $Q_l$ are block-diagonal matrices, which may lack intuitive or empirical justification. Specifically, if the goal is to simulate the training of real LLMs, where parameters are almost unconstrained, such explicit constraints may distance the current theoretical framework from practical scenarios. Did I miss any crucial justification for this assumption?\n\n3. The proposed TaskV-M application method appears to yield only marginal performance gains. While I appreciate the authors’ effort to translate their theoretical findings into practical applications, the method may not serve as a strong empirical validation of their conclusions.\n\n4. The empirical results (Fig. 3, Table 1) presented in the paper appear to be anecdotal. I would like to see additional results on more LLMs and input samples.\n\nQUESTIONS:\n1. The paper does not provide a detailed analysis of how the theoretical results are reflected in the parameters of real-world LLMs. Could the authors provide more insights into how the theoretical framework can be applied to real LLMs, or at least discuss the challenges, limitations, and potential future directions for such applications?\n2. Regarding Weakness 2, could you provide a detailed explanation (and incorporate it into the revised paper) of your justification for using such parameter constraints?\n3. Regarding Weakness 4, could you provide additional results on more LLMs and data samples to strengthen the empirical validation of the theoretical framework?\n\nI would be happy to raise the score to 5 if the author can properly address these concerns in the rebuttal.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the task vector method under the setting of ICL of linear regressions using linear transformers. The authors first show that if the $x_i$ and $y_i$ are separate tokens in the input, a critical point solution for transformers on the ICL loss is to mix up each $(x_i,y_i)$ pairs using the positional encoding (PE) at each layer and to perform GD++ after the first layer. Then they studied the setting where the input consists of triplets $(x_i,\\to,y_i)$. In this setting, there is a critical point solution that not only performs the GD++ and $x,y$ pairs mixing, but also uses the arrow token to attend to all $x_i$ and $y_i$ to make linear combinations as new examples of linear regressions. The authors then argue that the activations of these arrows (after the 1st layer) could be used as the task vectors (Task V), which are inserted after query ICL example to improve zero-shot performance. The traditional Task V method only uses the activation vector of the last arrow vector and the authors propose to use the activation vectors of all arrows, which is termed Task V-M. The authors also empirically show that Task V method does not work well in the bijection tasks and the explaination relies on the conjecture that the task vector is a linear combination of demonstrations.\n\nSTRENGTHS AND WEAKNESSES:\nStrength: \n1. This paper is of rich content: the analysis of loss landscapes for three settings of input, the analogy of the activation vector of arrows to the task vectors, the limitation of task vector method (bijection task) and empirical verifications.\n2. The writing is mostly clear.\n\nWeakness:\n1. Most explainations work for the 2-layer transformer while they are not clear to me for deeper models. I didn't find any discussion on that.\n2. It is not clear to me under what conditions the transformers will converge to the critical points that authors describe. There should be infinitely many critical points and some of them might be essentially different (please see the questions for details). This takes some more careful optimization dynamics analysis involving concrete training algorithms.\n\nI didn't read the proofs in the appendix so it is possible that there are some strengths/weaknesses that I'm not aware of.\n\nQUESTIONS:\n1. Is there any constraint on $\\Lambda_1$ in Theorem 1? It seems to me that $\\Lambda_1=\\begin{pmatrix} 0 & 1 \\\\\\ 1 & 0 \\end{pmatrix}$ would work while $\\Lambda_1=I$ doesn't.\n2. Line 135.  *\"As a result, an L-layer linear transformer allocates one layer for embedding concatenation and utilizes the remaining L − 1 layers to perform gradient descent.\"* Seems like the concatenation is happening at each layer instead of just the first layer from the statement of theorem 1 where $D\\in\\mathcal{S}_P^L$. Is there an explaination on that?\n3. Line 160. Is $\\alpha_1$ equal to $\\alpha_2$? In eq(10) they are equal but the Theorem 2 does not tell me anything about that.\n4. Line 161. *\"These vectors can then be injected into zero-shot prompts and function as single-token demonstrations\"* Are these task vectors the hidden activations of arrows after the first layer or every layer? The discussions are clear to me if they are extracted after the 1st layer. It is a bit hard to understand if they are extracted in deeper layers.\n5. Line 164. *\"task vectors naturally emerge from the optimization dynamic\"* Why would the transformers converge to such critical points? Specifically, in proposition 3 the $\\lambda$ could be $0$ and it corresponds to the case where there is no task vectors. In my understanding, as the authors also mentioned in line 238, the task vectors do not help with minimizing the ICL loss and it is possible to converge to a critical point where there is no task vector. \n6. The idea of bijection task is interesting. Is it possible to empirically show that Task V-M works for the bijection task as multiple task vectors (from multiple independent linear combinations) allow a $W$ of a higher rank.\n7. For ICL of linear regressions, it is somewhat reasonable to conjecture that task vectors are linear combinations of demonstrations. It is not clear to me if this conjecture makes sense for non-linear tasks. The linear combinations of demonstrations are easy to obtain only after the 1st layer if the 1st layer's $V$ matrix is identity. For tasks whose task vectors are extracted in deeper layers with non-linear transformers, the task vectors seem to be non-linear on the input. I'm wondering if the optimal depth to extract task vector is 1 for linear tasks, while it is deeper for non-linear tasks.\n8. The message from Fig 3 is unclear to me.\n\nI'm willing to raise the score if questions are well addressed.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper theoretically analyses the emergence of task vectors in bidirectional linear self-attention models trained on in-context linear regression tasks.\nIt finds that when data points are presented in-context as three tokens (x, ->, y) instead of two tokens (x, y), task vectors form as a weighted summation of all demonstrations.\nAs a result, task vectors are limited in expressivity since they can only realize rank-one coefficient matrices in the linear regression setting.\nBased on this insight, the paper conjectures that LLMs using task vectors should struggle on bijection tasks, e.g. $y = W x$ and $x = W y$, which can in general not be solved with rank-one matrices $W$.\nThis conjecture is validated empirically on pretrained large language models showing degraded performance on bijection tasks using task vectors while in-context learning performance is relatively high.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper makes an interesting case for understanding task vectors through its theoretical analysis that leads to the surprising finding that task vectors do not work well on bijection tasks.\nUnderstanding how and why task vectors work has the potential to significantly improve our understanding of transformer models.\nUnfortunately, the paper lacks clarity from the way it is structured and how it presents the main conclusions from its theoretical study.\nI am not very convinced by what is positioned as the main hypothesis, namely \"_The injected task vector functions as a single in-context demonstration, formed through a linear combination of the original demonstrations (hidden states)._\", given that the analyzed models are stacked, bidirectional linear self-attention models which -- as far as I understand, see questions -- almost by definition satisfy this hypothesis.\nImportantly, a dedicated related work section is completely missing and the result is not sufficiently contextualized within the research context.\n\nQUESTIONS:\n- What is the alternative hypothesis to the linear combination conjecture? What evidence would have disproved it? I am wondering this especially since the theoretical model is a stack of linear attention layers with skip connections without nonlinearities. Linear attention is basically a weighted combination of other tokens so I fail to see what type of result this hypothesis excludes. \n\n- As a follow-up to the previous questions, is there some meaning to the weights in this weighted combination? Maybe an even stronger empirical consequence of your result would be to show that task vectors can in practice be created without running the model simply by adding the examples as suggested by the finding? What prevents demonstrating this?\n\n- I am confused how to reconcile the results presented in Table 1 and Table 2. Table 1 using Llama-7B seems to suggest that task vectors do not work for bijection tasks whereas Table 2 using Llama-13B shows that bijection tasks can be solved with high accuracy using task vectors (TaskV rows).\n\n> Notably, the structure of $S_P$ closely aligns with our visualization of $D_l$ in Figure 2b, confirming our theoretical analysis\n\n- I am not sure I can follow this statement. What I can clearly see in the attention pattern is the concatenation operation. What other properties predicted by Theorem 2 are visible?\n\n- In the caption of Figure 3, you use a blue rectangle to refer to the ith demonstration but I am not sure what it visually corresponds to in the figure?\n\n> These vectors serve as redundancy against information loss induced by dropout, thereby improving robustness.\n> [...] While dropout is not always applied during LLM pretraining or fine-tuning, the injection of position encodings and use of normalization act as alternative sources of perturbation, thereby promoting the emergence of such redundancy.\n\n- I am not sure I follow this line of argument. Token-wise dropout is very uncommon as far as I know but I am not sure why position encodings or normalization would introduce a similar pressure to copy information. Could you elaborate on this point, i.e. how would they create a pressure to redundantly copy information? In general, I am surprised by the prominence of the conclusion that task vectors serve as a redundancy against dropout given that this is a relatively minor component of your theoretical investigation.\n\n- To what extent does the analysis rely on using the expected ICL risk instead of the more common autoregressive loss? I am asking this primarily out of curiosity whether you have considered it and less because I think it is necessarily a problem. However, I am wondering to what extent having a causal mask might change the main conclusions?\n\n\n## Suggestions\n\n- I find it misleading to refer to the models consisting of purely linear self-attention as (linear) transformers (e.g. L136) given that they lack the feedforward blocks. One suggestion would be to simply call them linear self-attention model instead.\n\n- I found it confusing that section 5 was called \"Further Discussions\". Based on this non-descriptive title and the position in the paper, I was expecting a related work section (which is missing). Maybe incorporating it into section 4 or finding a more expressive title would be helpful to guide the reader. \n\n> As a result, an L-layer linear transformer allocates one layer for embedding concatenation and utilizes the remaining L − 1\nlayers to perform gradient descent.\n\n- The split into a copy/concatenation operation and GD operation is similarly found in e.g. [1]. It think this should be mentioned.\n\n[1] Uncovering mesa-optimization algorithms in Transformers, von Oswald et al., 2023\n\nLIMITATIONS:\nA discussion of limitations is missing. I would have for instance expected some contextualization of the possible caveats of trying to extrapolate the results to full transformers given the simplified theoretical model (beyond the fact that this type of model has been considered by prior work).", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the Linear Combination Conjecture as a mechanism by which task vectors are formed from linear combinations of in-context example demonstrations. The authors use this assumption to derive the emergence of task vectors by analyzing critical points of linear transformers trained on in-context learning (ICL) triplet examples for random linear regression. Next, the paper correctly predicts that modern LLMs with task vectors perform poorly on bijection tasks. Finally, this conjecture is utilized to add multiple task vectors to a few-shot prompt, resulting in improved performance.\n\nSTRENGTHS AND WEAKNESSES:\n### Quality\n- Proofs and experimental methodology appear to be correct\n- Each of these points provide evidence supporting the central claim of Linear Combination Conjecture\n\n### Significance\n- Important contribution that increases understanding of in-context learning, a fundamental component of many LLM applications.\n- Potentially high impact if it leads to follow up works that further enhance the task vector method\n\n### Originality\n- This paper seems to build on existing work (task vectors, loss landscape analysis) in a novel manner to produce novel insights.\n\n### Clarity\n- Generally well-organized and a pleasure to read. Ideas are presented clearly with sufficient background information\n- Notation for dimensions is confusing at times, in particular the relationship between $d_p$ (number of tokens) and $n$ (number of examples) could be more explicit.\n\n- Typo: In the equation below line 402, I believe $\\tilde{R}_i$ should be $R_i$\n\nQUESTIONS:\n1. Are Tables 1 and 2 directly comparable? For bijection tasks, TaskV performs better than the ICL Baseline in Table 2 but worse than ICL in Table 1.\n2. In addition to Sections 6.3 and C.2, another multi-vector variant could extract multiple vectors from each arrow in a few-shot prompt and then inject them into the same corresponding positions. Would this provide additional insight to support or refute the Linear Combination Conjecture?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the underlying mechanisms of task vectors, a technique used to accelerate inference in in-context learning (ICL) by condensing task-specific information into a single reusable representation. The authors propose the Linear Combination Conjecture, which suggests that task vectors function as single in-context demonstrations formed via linear combinations of original demonstrations. The paper provides both theoretical justification and empirical evidence, showing that task vectors naturally arise in linear transformers trained on triplet-formatted prompts, as revealed through loss landscape analysis. The authors also predict and confirm that task vectors struggle to represent high-rank mappings in practice. Additional validation through saliency analysis and parameter visualization supports their conclusions. To address limitations, the study proposes injecting multiple task vectors into few-shot prompts to improve performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper presents a strong theoretical foundation to support the proposed Linear Combination Conjecture, offering valuable insights into the formation and function of task vectors. The authors conduct comprehensive zero-shot and few-shot experiments, which effectively validate the core ideas of the work across practical settings.\n\nWeaknesses:\nThe paper does not discuss its limitations, leaving questions about the generalizability and constraints of the proposed approach.\nThe heavy use of mathematical notation and formulas may hinder readability and make the paper difficult to follow for some readers.\nIt remains unclear whether the proposed method is applicable to multimodal task.\n\nQUESTIONS:\nSee the weaknesses.\n\nLIMITATIONS:\nA discussion on the potential limitations of the method would strengthen the paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}