Task Vectors in In-Context Learning: Emergence, Formation, and Benefits
Authors: Liu Yang
Ziqian Lin
Kangwook Lee
Dimitris Papailiopoulos
Robert D Nowak
Venue: COLM 2025 (2025)
Keywords: in-context learning; task vector
TL;DR: We study task vectors in transformers trained from scratch on synthetic tasks and find they emerge but remain indistinct. To enhance their formation, we introduce TVP-loss. Strong task vectors in deeper layers improve ICL on OOD prompts.

Abstract: In-context learning is a remarkable capability of transformers, referring to their ability to adapt to specific tasks based on a short history or context. Previous research has found that task-specific information is locally encoded within models, though their emergence and functionality remain unclear due to opaque pre-training processes. In this work, we investigate the formation of task vectors in a controlled setting, using models trained from scratch on synthetic datasets. Our findings confirm that task vectors naturally emerge under certain conditions, but the tasks may be relatively weakly and/or non-locally encoded within the model. To promote strong task vectors encoded at a prescribed location within the model, we propose an auxiliary training mechanism based on a task vector prompting loss (TVP-loss). This method eliminates the need to search for task-correlated encodings within the trained model and demonstrably improves robustness and generalization.

Published as a conference paper at COLM 2025
Task Vectors in In-Context Learning: Emergence, Formation, and Benefits
Liu Yang", Ziqian Lin"', Kangwook Lee",, Dimitris Papailiopoulos & Robert D. Nowak
"University of Wisconsin-Madison, "MMicrooftt Research yanu. zlin284, kangwook. leeeewisc.eeeee dimitris@papail. io, rdnowakewisc.edd edu
Abstract
In-context learning is a remarkable capability of transformers, referring to their ability to adapt to specific tasks based on short history or con- text. Previous research has found that task-specific information locally encoded within models, though their emergence and functionality remain unclear due to opaque pre-training processes. In this work, we investigate the formation of task vectors in controlled setting, using models trained from scratch on synthetic datasets. Our findings confirm that task vectors naturally emerge under certain conditions, but the tasks may be relatively weakly and/ or non-locally encoded within the model. To promote strong task vectors encoded at a prescribed location within the model, we pro- pose an auxiliary training mechanism based on a fask vector prompting loss (TVP-losss. This method eliminates the need to search for task-correlated en- codings within the trained model and demonstrably improves robustness and generalization."
Introduction
To understand the underlying mechanisms of in-context learning in transformers, re- searchers have probed pre-trained models from various perspectives, such as altering the labels in demonstrations (Min et al., 2022; Kim et al., 2022) and investigating circuit mechanisms (Elhage et al., 2021; Wang et al., 2023; Hanna et al., 2024; Singh et al., 2024). Additionally, controlled, small-scale studies have been conducted by training transformers from scratch to observe their in-context learning behavior on linear regression tasks (Garg et al., 2022; von Oswald et al., 2022; Lin & Lee, 2024), discrete functions (Bhattamishra et al., 2023), hidden Markov chains (Xie et al., 2021), and DFAs (Akyürek et al., 2024). Furthermore, theoretical approaches have also been applied to this problem (Xie et al., 2021; Lin & Lee, 2024; Giannou et al., 2023).
Among the various efforts probe pre-trained models, one significant line of research employs the concept of a "task vector", which is a vector in the model's 's weight or activation space" that encodes task-specific information. The concept of the task vector was first introduced by Ilharco et al. (2022), where it is defined as a direction in a model's weight space corresponding to particular task. Subsequently, Hendel et al. (2023) demonstrated that, given task demonstration as context, pre-trained large language model forms task vector in its activation space at certain layers. This task vector encodes only the task information and is independent of the specific demonstration of the task. By inserting the task vector directly into the model, it is able to perform the task without context or demonstration (i.e., zero-shot). We will refer to this as Task Vector Prompting (TVP) in this paper. Concurrently, Liu et al. (2023a); Todd et al. (2023); Merullo et al. (2023a); Li et al. (2024b); Saglam et al. have also identified a single vector that encodes the task information, albeit using different terminology. We omit the details here and refer the reader to the related work and the original papers for more information.
Motivated by prior observations of task vectors in pre-trained LLMs-where training conditions are inherently difficult to control-we investigate their emergence in controlled
1 the code is available at https:// //gi thub. www.:thh..ooooo..oeooooooootoet com/Leiay taskk vector prompting. 2The "activation space" refers the space where the output of each transformer layer resides.
1

---

Published as a conference paper at COLM 2025
Eo:Context Learning
Vanilla
1.0
Transforeer
Transforrer
Mean Squared Error
0.5
32
0.0%
In Context Learning Mode
Zero-Shot Mode
Task-Vectop Fr Proup
Vazilla
10:
Ours
Transtorner
Transtormer
0.5
5.0
(4g
Y1
Asest
40
Vector Pronpcing
of in-context samples
Figure 1: Overview of the transformer operating in int-oonext learning (ICL) and task vector prompting (TVP) modes. A transformer can be configured operate in ICL mode, using input-output pairs as prompts, or TVP mode, extracting task-specific embeddings for zero-shot predictions (architecture and training details in Section 3). In the bottom, the ICL and TVP performances of the vanilla- aanilla-trained model and our method are shown, with the dashed horizontal line indicating random prediction performance (i.e., task information is inferred). Compared to vanilla training, our approach enhances task- specific representations in the TVP mode while preserving comparable ICL performance.
setting by training transformers from scratch on synthetic catasetss As shown in Figure 1, when trained from scratch on the linear regression task defined as y; = w x,, the transformer (dashed line) demonstrates the ability to use in-context learning (ICL) to solve the task (bottom-left). We evaluate the trained model's performance in task vector prompting (TVP) mode (Figuree 1, middle), where the task vector is extracted from the in-context learning mode and injected back in a zero-shot manner, as defined by Hendel et al. (2023). The TVP performance is much better than chance (indicated by the horizontal dashed line), but a bit worse than ICL performance, which we attribute to the fact that the encoding of the task may not be strong and localized by normal training methods. To encourage the formation of strong and localized task vector, we propose an auxiliary training loss, called the lask veclor prompting loss (TVP-loss). In our new approach, the model is trained using the TVP-loss in addition to normal training losses. As illustrated in Figure 1, our approach (solid line) achieves comparable ICL performance to the vanilla model in the ICL mode (bottom-left) and its TVP performance (bottom-right right) is significantly improved and comparable to ICL performance with multi-shot demonstration.
Our main contributions and findings are outlined below:
Emergence of Task Vectors During Training. We investigate the effectiveness of task vector extraction methods-originalll proposed for pre-trained LLMs-when applied to small- scale models trained from scratch on synthetic datasets. Our findings show that task vectors can naturally emerge during training, consistent with the observations in Hendel et al. (2023) for pre-trained LLMs, provided appropriate input formats and sufficient model capacity. To better understand what promotes this emergence, we examine various hyperparameters, including model depth, sparse attention, and input embedding size.
Strong Task Vectors with Proposed Auxiliary Loss. Although task vectors naturally emerge, they are often weak and entangled with information from input queries. This prevents task vectors from representing purely task-specific knowledge. To overcome this limitation, we propose training algorithm that explicitly encourages the formation of task vectors independent of query information. This approach strengthens the task vector, ensuring that task-specific encoding is explicitly established within the model.
Task Vectors for In-Context Learning Robustness. Using our proposed training algorithm, we analyze the effects of TVP-loss on (a) synthetic, (b) formal language, and (c) natural language tasks, demonstrating enhanced robustness in in-context learning.
Related Works
The notion of "task vector" was first introduced by Ilharco al. (2022). Subsequently, Hendel et al. (2023), Merullo et al. (2023a), Yu et al. (2023), Liu et al. (2023a), Saglam et al. and Liet al. (2024b) demonstrated that a single vector in the model's activation space can encode learned functions in pretrained model. Specifically, Liu et al. (2023a) showed that, when presented with an in-context learning demonstration, the activation at each layer points to a subspace encoding the task information. Additionally, Todd et al. (2023) identified the "function vector," " another form of task vector, by averaging the causal attention heads, which can also guide the pre-trained language model's performance towards desired tasks.
2

---

Published as a conference paper at COLM 2025
Task vectors have also been identified under different modality (Luo et al., 2024; Hojel et al., 2024; Peng et al., 2024), and from different perspective, such as cognitive science (Piantadosi et al., 2024), and through theory perspective (Tao et al., 2024). Park et al. (2023) observe that concepts are represented linearly as directions in some representation space.
A recent study by Mittal et al. (2024) explores setting similar to ours, examining whether learning appropriate latent representations via a bottleneck architecture enhances robustness in in-context learning. Notably, our approach facilitates the natural emergence of latent representations within the model's architecture. Expanding on this, Kobayashi et al. (2024) demonstrate that bottleneck architectures improve compositional generalization, though their definition of compositionality pertains to meta-skill learning rather than the sequen- tial compositionality studied in our work. Additionally, Elmoznino et al. (2024) leverage bottleneck a architectures to analyze prequential ICL performance. Furthermore, Han et al. (2024) investigate the dynamics of task vector formation in both sparse linear regression and pre-trained LLMs, showing that as task vectors emerge, models simultaneously develop conditional decoding strategies, leading to improved in-context learning performance.
Due to space constraints, extended related work appears in Appendix A.
3 Task Vector Definition
In this section, we formally define the notion of task vector. Let F denote a class of functions or "tasks", and let X' and y be the input and output spaces, respectively. If f F is the task in a specific "context", then for any input x E X, the corresponding output is y = fx) E y: Consider Xiest X, and the transformer model M. We can measure the model's zero-shot performance by (Xiet)))), M,M,),,MP()( M M(Pquery)), )), where Pquery = [xiest] denotes the query input. The in-context learning performance with k-shot examples is measured by M(P))),, where Ps is the in-context prompt. The prompt contains the k-shot examples (xx), Xxr, (xx), Xtest]-
Consider a demonstrated prompt pe, where the superscript highlights its role in demonstrating task information. This prompt contains k in-context samples: pe = [x) fxx), Xk. fxx)] A task vector is an internal embedding т extracted when the model is presented with PP, which encodes the task at hand (f). Inserting the task vector into the model during zero-shot prompting is denoted by M(P M(Pquery' т). Note that the task vector T is extracted from the internal embeddings when inputting PP, not P&. therefore the task vector T does not have explicit knowledge of Xiest when encoding the task.
A task vector extractor g tries to locate this single embedding т in the model given the demonstrated prompt Pp i.e. == (MM(P))) Then the performance of this extractor can be measured by ((f(xtes), M (Pquery: T)). A task vector is considered successfully formed in the model M for the task f if the performance of the extracted task vector, i.e. task vector prompting performance, closely aligns with the in-context learning performance and outperforms zero-shot performance.
Hendel et al. (2023) confirmed that, in pre-trained large language models, g( (M(PP) corre- sponds to the output embedding at approximately the middle layers of the model during the forward pass, which is responsible for predicting f f( Similarly, Sia et al. (2024) found that the task in context is recognized by the model during the middle stage of the forward pass.
4 Emergence of Task Vectors in Fraisec-from-Seracch Models
4.1 Experimental Setup
We study task vector localization using GPT-2 decoder trained on randomly generated prompts for the following tasks (details in Appendix B.1):
1. Linear Regression. fxx)) w x; with x;, w IRd. We set d in our experiments.
2. Sinusoidal Regression. f(x;) (0.5 w Tx; b), where b (0,
3

---

Published as a conference paper at COLM 2025
3. Discrete Token Offset Prediction, f(x:) (a x; + mod C, where {1,2,3), b (0, sampled uniformly, and C as the vocabulary size (default: C = 1000). mod denotes the modulo (remainder) operation.
During training, prompts are generated on the fly. Specifically, for each prompt, a function f F randomly sampled according to the function distribution described for each task above. Subsequently, input tokens {x;),, are independently sampled from the correspond- ing input distribution. The function f then evaluated on these inputs to produce the target outputs, forming the in-context learning prompt Px fxx, (..)), Xk, f fxxx) Xrest] Let the
distribution of such prompts be denoted P. The transformer M parameterized by is then trained to minimize the following expected loss:
k min Σ (o(()))
where P; (x))). X,, examples. The loss function ((,,-) is defined as mean squared error for regression tasks and cross-entropy loss for discrete token tasks. In our experiments, we use GPTT-2 model with an embedding size of 64, 4 attention heads, and 3 layers, trained with a maximum context 63) and without positional embeddings (NoPE). We use the Adam
(x;)), fx;;), xi+1 represents the prompt prefix containing in-context
length of 63 (i.e.,& optimizer with a learning rate of 0.0001 and a batch size of 256, training for 300k iterations. With dynamically generated prompts, this corresponds to 76.8 million distinct prompts. All experiments are run on an NVIDIA GeForce RTX 3090.
We investigate various input formats and task vector extraction methods (details in Ap- pendix C.1 and Appendix C.2) to identify conditions under which task encoding emerges in froneedddfoo-mcccacchh transformers. Based on this investigation, we focus on the following setup, where task vectors emerge most distinctly:
imput format: Prompts with in-context examples are generated as Pr. (x))...................
fxxx X where z is special token placed at the beginning of the prompt to serve as a placeholder for injecting task encoding during zero-shot task vector prompting. During training, the embedding of is treated as a learnable parameter that is shared across all prompts (details in Appendix C.1).
task vector extractor method: inspired by Hendel et al. (2023), we locate the task vectors at the activations of the token z and {yj} in the input format mentioned above. Specifically, during task vector prompting, the embedding is copied into the corresponding position in a zero-shot model, and the zero-shot loss is measured. A lower zero-shot loss indicates that the embedding effectively encodes task-specific information (details in Appendix C.2).
As noted in Hendel et al. (2023), the extraction process involves identifying the layer that optimally encodes the task vector. The optimal layer index I* determined by evaluating task vectors extracted from each layer and selecting the one that minimizes the average loss across all N demonstrated prompts. This procedure ensures that the selected layer provides the most effective task vector for predicting the test outputs. A formal description of this procedure is provided in Appendix C.3.
Remark 4.1 (Task Vector Location). The work of Hendel et al. (2023) employs an in-context learning prompt format structured as P₂ x, z, (x1), X f (xx), I, where indicates the "maps-to" token. This setup differs from the prompt format used in our experiments. In Appendix C.1, we investigate various input formats to determine their effect on task vector emergence. We find that the emergence only when the input format for training is Pt = x (x)),. Xx, f(xx), Xxt1l, where no additional tokens placed between x and fx).
model exhibits task vector
4.2 Task Vector Emergence
Finding 1: Task vectors naturally emerge in small models trained from scratch.
In Figure 1, we demonstrate the presence of task encoding in the activations of token y; under the linear regression task, using the experimental setup described in Section 4.1
4

---

Published as a conference paper at COLM 2025
with number of demonstrated prompts N 50. In the task vector prompting mode, the transformer achieves a mean squared error (MSE) of approximately 0.25, significantly lower than the random prediction baseline of around 1.2, indicating that the model successfully infers task information in its representation. Throughout the paper, we use N = 50 by default unless stated otherwise. Additionally, we examine the impact of different N values on task vector prompting performance, as detailed in Appendix C.6.
In the following sections, we analyze the impact of model depth and context length on task vector emergence across different problem dimensions (Section 4.2.1) and investigate the locations of task vectors (Appendix C.3.1). Additionally, we explore the effects of embedding size and sparse attention in Appendices C.4 and C.5, respectively.
4.2.1 Effects of Model Depth and Context Length on Task Vector Emergence
In the linear regression task, we examine how model depth and context length affect the emergence of task vectors by varying the problem dimension and model depth. We explore d E (,,,,,,7,8,9) ,), keeping the transformer's embedding size 64 and adjusting the model depth from L = 3 to L = 8. The transformer's performance in ICL mode (solid line) and TVP mode (line with triangular markers) is shown in Figure 2, with 50 for clarity.
L=4
L=5
L=6
L=7
L=8
1.0
Mean Squared Error
0.5
0.0
0
0-
20
40
20
40
60
0
20
40
60
d+4
d=6
d-8
1.0
0.5
40
20
d-5
d-7
d-9
# of in-context samples
Figure 2: Performance of the transformer in in-context learning (ICL) mode (solid line) and task vector prompting (TVP) mode (line with triangular markers) across problem dimensions d E (4,5, 6, 8,9) 9} and varying model depths L E (3,4,5,6,7 6,7,8}. The dashed line indicates the random-guest baseline. The results indicate that smaller problem dimensions (d 4 to 7) and shallower model depths (L = to I. 5) yield stronger and more stable task encoding in the TVP mode. However, task encoding remains noisy in most cases. Task vectors emerge more clearly when the ICT. loss plateaus but tend to disperse with increasing in-context length. In deeper models, task information appears to distribute across layers, reducing the distinctiveness of task vectors.
In the linear regression task, task vectors are most distinct with moderate model depth and just before the in-context learning loss plateaus.
Finding 2:
We observe the following characteristics of task vectors:
Along the Model Depth: Tasks with smaller problem dimensions (d = 4 to 7) exhibit clear pattern of specific model depths producing strong and stable task encoding, as illustrated in the plot. However, in most cases, task encoding remains noisy. Notably, as model depth increases, the clarity of task encoding diminishes. This could be due to the model's increased capacity to approximate the least square solution directly in a single forward pass, reducing its dependence on previously calculated task encoding.
Across the In-Context Length: Task vectors become more distinct as the in-context learn- ing loss plateaus, but they tend to disperse afterward. For example, when d = 5 with L == 4 or = 7 with L = 4,5,6, 6, task encoding is most evident around the 5-th in-context example but then disperses into random task encodings. This suggests that the model initially learns the task in a compact and focused manner; however, over time, the task information becomes distributed more broadly across the context.
5

---

Published as a conference paper at COLM 2025
5 A New Training Algorithm to Encourage the Formation of Task Vector
As evidenced in the previous section, it is difficult to locate single vector in the trained- from-scratch small-scale transformer model that cleanly encodes the task at hand. In this section, we propose a training algorithm that explicitly encourages the formation of task vectors (as defined in Hendel et al. (2023)) in the in-context learning process.
Through this algorithm, which is a straightforward extension of the task vector definition, we obtain a model with the task vector explicitly formed. As demonstrated in Section 5.2, comparing this model to the vanilla model (where task vectors are not explicitly formed) reveals that task vector formation enhances the model's robustness in in-context learning tasks, particularly with ditto--ditribbution prompts. This highlights the value of task vector formation in improving generalization and interpretability in in-context learning.
5.1 A New Training Algorithm to Encourage Task Vector Formation
To encourage the formation of task vectors, we include the performance metric of the task vector in the training loss. Specifically, following the notation in Section 4, let each random prompt be P2- = [z, f (x)),.... X2,, fxxx), (xxx),x22 xx-1] to be sampled following the prompt dis- tribution P, and a random test prompt Iquery- = (tes)] We train the transformer Ma by optimizing the following loss:
min E Ep-p
Σ
e( ((((x))) + (x(est)) Ma Ma(Qquery ((Pueryy(()))
(1)
ICL-less
TVP-less
where I! represents the hidden state at the I-th layer of the transformer, extracted from the i-th in-context example. Instead of adaptively identifying the layer where the task vector resides, we simplify the process by directly designating specific layer / as the location for forming the task vector. This ensures consistent representation of task-specific information and facilitates training. We illustrate the training algorithm in Figure 13. In Section 5.2, we further explore the impact of this hyper-paramett choice on the model's performance.
The loss comprises two complementary components: (1) In-Context Learning Loss (ICL- loss): The term (x;+1)), M4 (Pt)) trains the model to predict the output f(x++1) for the (i+1)-th example, based on the preceding context P,. Summing over k examples ensures the model learns from the entire context effectively. (2) Task Vector Prompting Loss (TVP- loss): The term (x(est), My ()) with т = h! evaluates the model's ability to use the injected task vector 11,, derived from the hidden state at the i-th in-context example, to predict the test output f (Ntes). This term encourages the model to encode task-specific information in the task vector т. -A evidence of this, in Figure 1, although the model trained with ICL-loss already exhibits emergence of task vector, the model with TVP-loss further improves the encoded task vector, achieving a lower TVP error.
5.2 Experimental Results on Synthetic Tasks
We apply this training algorithm to the aforementioned three synthetic tasks: (1) the linear regression task, (2) the sinusoidal regression task, and (3) the discrete token offset task described in Section 4.1. For all experiments, the transformer model is configured with total of 8 layers. For consistency to the previous results, we set N 50.
Finding 3: The proposed TVP-loss enhances task vector prompting performance, aligning it closely with in-context learning performance.
For clarity, we evaluate the in-context learning and task vector prompting performance at the 63rd context (i.e., the final query) for models trained with and without the proposed TVP-loss. These results remain consistent across various context lengths, with full details provided in Appendix D.2. In this experiment, we designate task vectors to form at the 1st, 3rd, 5th, and 7th layers. As illustrated in Figure 3, for the 8-layer transformer, the task

---

Published as a conference paper at COLM 2025
[CL
L Linear Regress ion
Sinusoidal Regressine ion
offsel
10a
10
10
Vanilaa
0.0 Vanilla
Vanilla
TV Layer TO
Figure 3: Model performance with vanilla and TVP-loss training. This figure compares models trained with vanilla training and TVP-loss across layens where task vectors are formed. The upper row shows in-context learning (ICL) and task vector prompting (TVP) performance on logarithmic scale for clarity, while the bottom row depicts diatt-ddditributioon (OOD) prompt performance. Dashed lines represent the performance of the vanilla-trained model in each respective setting (ICL, TVP, and OOD), providing baselines for comparison. Each column corresponds to task: linear regression (left), sinusoidal regression (middle), and discrete token offset (right). The horizontal axis indicates the layer which task vectors are formed. ICL performance: Models with TVP-loss achieve similar ICL performance to vanilla-trained models. TVP performance: TVP-loss improves TVP performance, especially when task vectors form in intermediate or later layers. OOD robustness: TVP-loss enhances OOD generalization, notably in complex tasks like sinusoidal regression and discrete token offset, with task vectors formed at specifie layers.
vector prompting performance of the vanilla-trained model is nearly random. In contrast, models trained with the TVP--loss exhibit task vector prompting performance that closely matches their in-context learning performance, as long as the task vector layer is set to an intermediate or later layer rather than the initial layers of the model.
Performance on OOD Prompts. Forcing the task vector to form specific layer can be interpreted as introducing an implicit bottleneck architecture within the model's forward pass. In this section, we examine the model's generalization ability when the task vector is constrained to form at a particular layer. Specifically, for the three synthetic tasks, we evaluate performance on the OOD prompts described in Appendix B.1, with additional OOD tasks and their results are presented in Appendix D.3.
Note that during training, none of these OOD settings were encountered. As shown in Figure 3, models with task vectors formed at specific layers exhibit equal or improved performance on OOD prompts compared to vanilla-trained models. For the linear regression task, forming the task vector at the 5th and 7th layer yields slightly better OOD performance.
The sinusoidal regression task can be consider as a compositional task, where fi (xx) = w" x, and f2(y) syn(y). Then the OOD task modifies the f2(y) sin(y) to fz(y) = vy. In this case, models with task vectors show enhanced generalization to the modified compositional task, indicating that capturing task representations at intermediate layers improves OOD generalization. We provide an extended analysis of the compositional functions in Appendix D.4.
In the discrete token offset task, models with task vectors handle context changes signifi- cantly better than those trained without the auxiliary loss. This suggests that the auxiliary loss enables the model to adapt more effectively by filtering out unrelated context at earlier positions, compared to vanilla-trained models.
6 In-Context Learning Beyond Synthetic Tasks
Finding 4: TVP-loss improves in context learning performance on out-of out-of- distributiio prompts.
In this section, we extend our study beyond the simple synthetic tasks on regression and token offsets, and investigate the impact of the TVP-loss in synthetic formal language in Section 6.1 as well as the natural language tasks in Section 6.2.
6.1 Synthetic Formal Language Task
We conduct experiments on two benchmarks: the Generative In-Context Learning (GINC) dataset introduced in Xie et al. (2021), and the RegBench dataset from Akyürek et al. (2024). In both datasets, each context consists of tokens per example. For instance, an in-context prompt with context length k = 3 and example length 4 looks like: b d / e f g j k 1, while a corresponding zero-shot prompt is m n o p. Further experimental details are provided in Appendix B.2.
GINC Dataset. Though the model is not Meta-ICL trained-meaninee meaning the pre-training distribution differs from the prompt distribution-traan training with the TVP-loss still enables
7

---

Published as conference paper at COLM 2025
In-Context Learning
(b) Task-vector Pro
(C) Ho1sy
(d) Change
5010
Regtench
5
sasples
Figure 4: GINC (top) and RegBench (boftom) dataset in-context learning performance. Comparison of the in-context learning performance for the vanilla-traing model (dashed line) and the model trained with TVP-loss at the 5th layer for GINC and the 4th layer for RegBench (solid line). In all cases, the model with task vector formation outperforms the vanilla-trained model.
the formation of task vectors during in-context learning. We present the result in Figure (top). We measure both model's performance on (a) in-context learning, (b) task vector prompting, (c) noisy in-context learning, where the in-context token is replaced with random token with probability 0.1, and (d) when the underlying HMM is changed after the second context.
RegBench Dataset. In Figure 4 (bottom), we present results on: (a) in-context learning, (b) task vector prompting, (c) noisy in-context learning, where 4 out of 32 in-context labels are replaced with random tokens, and (d) performance when the underlying DFA of the first 4 out of samples is changed another DFA. When evaluating the next predicted token in the in-context learning prompt, we follow the same setup as in Akyürek et al. (2024): as long as the next predicted token belongs to the outgoing edge set of the current DFA state, is considered correct prediction; otherwise, it is incorrect. Intuitively, as k increases, more transitions are presented to the transformer. This requires the model to approximate larger 1--ram distribution, leading to a degradation in performance.
Performance Analysis. Across both datasets, models trained with TVP-loss consistently outperform vanilla-trained models. Specifically, in the GINC dataset, TVP-loss improves the model's ability to learn in-context, even when trained only on next-token prediction tasks, and enables task vector formation without directly optimizing for in-context learning during pre-training. As shown in Figure 4, models trained with TVP-loss achieve slight improvements in ICL performance and significant gains in TVP performance, Additionally, these models exhibit increased robustness to label noise and -uttof-disttibution (OOD) in-context samples, consistently outperforming vanilla-trained models in OOD scenarios.
6.2 Natural Language Task
Motivated by the OOD improvements observed with TVPP-loss on GINC and RegBench, we further assess its effectiveness in natural language tasks using the CrossFit benchmark (Ye et al., 2021) following the MetaICL settings in Min et al. (2021). Specifically, we fine- tune pretrained GPT-2 Large model (774M parameters), modifying the original MetaICL procedure by training across all context lengths rather than using fixed context length (k = 16), which allows the model to generalize to varying context lengths at inference time.
Since the meta-training set contains diverse range of tasks, the optimal layer for forming task vectors may vary across tasks. To accommodate this, we do not fix a single task vector layer during fine-tuning with TVP-loss. Instead, for each training prompt, we evaluate the T'P--los at every layer from 8 to 15 (out of 36 total layers) and select the layer that minimizes the loss. Formally, recall the TVP-loss equation (1) for layer I defined
TVP
e(((sus)) Ma (PPPurryy
For each training step and each prompt, instead of minimizing e. e.V with predefined we minimize (TVP minng (TVP. This allows the model to determine the most effective task vector location per prompt.
During inference, following Min et al. (2021), each test query is paired with set of can- didates (e.g., classification labels or answer choices for QA), and the model selects the candidate with the highest conditional probability.
8

---

Published as a conference paper at COLM 2025
Table presents results under three task transfer settings in the CrossFit benchmark follow- ing MetalCL settings. In each case, we train on a set of mnta-trainnnng tasks and evaluate on distinct set of target tasks to assess generalization. The three settings are:
1. High-resource >> Low-resource (HR- Datasets with 10,000 or more training axamples-refereee to as high-resouree tasks are used for meta-training, while the remaining tasks are used for evaluation.
2. Non-Class + Class: The model is trained tasks that do not involve classification and then evaluated on classification tasks.
3. Non-Paraphrase + Paraphrase: The model is trained tasks not with paraphrase detection and evaluated on tasks with paraphrase detection.
We report in-context learning performance (with context length k = 16) on the target tasks, showing the mean and standard deviation over inference seeds. In addition to overall target task performance, we also separately evaluate unseen domain target tasks-datased whose topics (e.g., finance, poetry, climate, or medical) do not appear during meta-training- to assess cross-domain generalization. As shown in Table 1, models meta-trained with TVP-loss consistently outperform those without TVP-loss. We hypothesize that TVP-loss serves a regularizer, promoting better OOD generalization by encouraging the model to form task-specific encoding.
Lastly, we evaluate the performance of task vector prompting with N 10 demonstrations. For each test query, we follow the same inference strategy as in standard in-context learning: pairing the query with set of candidate answers. Instead of passing the full demonstration context, we insert the task vector extracted from the context into specific layer determined by the demonstrations. The task vector prompting results are reported in the lower half of Table 1. Notably, although the evaluation tasks were never seen during training, TVP-llos still provides a modest performance gain HR-LR and Nan-Paraphrase-- Parraphraaa settings, suggesting that it also improves task localization ability.
Method
HR LLR
non-Class -Caass-s-laass
non-Paraphruse- -aaraaareee--aappruree
MetalCL
44.88 = 0.29 / 39.14 3.20
37.97
1 39.55
41.07 2.57 / 34,05 0.00 41.62 1.73 34.05 0.00
ICL
TVP--Ios (Ours)
42.10 3.67
33.78 244
MetalCL
31.93 0.20
0.34
30.97 2.00
TVP
TVP- (Ours)
34.41 0.25
22.71 : 1.56
33.45 + 0.82
Table ICL and TVP Accuracy of Fine-tuned GPT-2 Large on the CrossFit Benchmark. We evaluate performance under three task transfer settings: HR-LR (High-Resource to Low-Resource), non- Fach cell reports the mean ICI accuracy (+ std) over three inference seeds using & 16 in-context examples, shown as: overall performance performance on target tasks from unseen domains. Bold values indicate that the difference between MetaICL and ours exceeds the sum of their standard deviations.
Class-+Class, and
6.3 Effects of Various Hyperparameter Choices
To further understand the robustness and sensitivity of our approach, we analyze the impact of key hyperparameters on the HR-+L task:
(A) the range of candidate layers used for selecting the task vector location, and (B) the weight WTVP in the combined objective LICL. WTVP LTvP (see Equation 1).
Table reports ICL and TVP performance under different configurations.
Layer range: Expanding the candidate range from the default (8-15) to a broader span (8-18) generally improves ICL performance but degrades TVP accuracy, potentially due to noisier task vector selection when the range is too wide. Conversely, narrowing shifting the range (8-12, 6-13, or 10-17) slightly degrades both ICL and TVP performance. These results suggest that the 8-15 range strikes balance between flexibility and task vector quality.
Loss weight wrvp: Varying the weight of the TVP-loss the joint objective reveals a trade- off: smaller weights (e.g., WTVP 0.3) slightly improve ICL performance but reduce TVP

---

Published as a conference paper at COLM 2025
accuracy, whereas larger weights (e.g., WTVP 3) hurt both metrics metrics-likklyy likely due to over- regularization. We find that WTVP = yields the best overall balance between in-context learning and task vector generalization.
Range wryve
8-15
8-12
8-18
6-13
10-17
8-13
8-15
0.3
ICI
45.60 1: 0.61
45.33 1 0.30
46.00 1: 0.09
45.45
0.52
45.60 1: 0.47
45.93
1. 0.45
44.53
0.40
TVP
34.41 + 0.25
34.62 + 0.33
32.91 0.55
33.99
+ 0.22
33.28 + 0.68
33.83 0.10
33.65 0.11
Table 2: ICL and TVP Accuracy of HR Task under Various Hyperparameter Setup.
Discussion and Conclusions
Factors Affecting Task Vector Emergence. The emergence of the task vector in pre-trained large language models may attributed to multiple factors, such as the model's capacity, scaling laws, and the diverse tasks encountered during pre-training. While Hendel et al. (2023) and related studies identify and analyze the existence of task vectors, they do not study the factors that may effect their emergence and performance. This gap in under- standing motivates our investigation: we reproduce this phenomenon in the small-scale setting when trained from scratch, gaining insight into the factors influencing task vector emergence, such as input format, model depth, and context length. For instance, in the linear regression task, when the model is trained with prompts where and y alternate closely, model with moderate depth encourages the emergence of task vectors, which is most evident before the in-context learning loss plateaus.
Training with TVP-Loss. While task vectors can emerge through normal training pro- cesses, adding the TVP-loss to the training process encourages the formation of task vectors at prescribed locations and leads to improved accuracy and robustness. Across various benchmark datasets, we have demonstrated that prescribing the task vector to an intermedi- ate or later layer of the model results in comparable in-context learning performance while forming strong task vectors. Depending on the task, task vectors formed at specific layers can enhance the model's 'ss robustness to out-of-distributtin (OOD) prompts, most likely at the intermediate or later layer.
Potential Applications of Task Vectors Task vectors can serve as soft prompt, effectively compressing the entire context into a single vector representation. In practical scenarios, where the number and nature of tasks are unknown during pre-training and only demon- strations are provided to the model, task vectors enable both context summarization and task identification. Specifically, model with task vectors formed gains the ability to identify the underlying task from demonstrations, cluster tasks using the extracted task vectors, and subsequently perform zero-shot inference with these extracted vectors.
Another practical benefit is during inference, attention computations over earlier tokens can be masked out after the prescribed I-th layer, restricting attention to only the task vector and the query. This approach mirrors the findings of Sia et al. (2024), who demonstrated 45% reduction in computation for pre-trained LLMs. Looking ahead, incorporating the TVP-loss as an auxiliary objective during training offers promising approach to enhance task-specific representation and overall performance.
References
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. ArXiv, abs/2221.15661, 2022. URL https://api. httt::::/app..emmnniccchhohooo wtw.:icccccoooaoooorrrrrrssoDo org/CorpusII: 254043880
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learn- ing: Architectures and algorithms. ArXiv, abs/2401.12973, 2024. URL https://api.
semanticscholar. wwweer.org//oee (ttp/CoouuID::70550000
Yu Bai, Fan Chen, Haiquan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. ArXiv, abs/ abs62230.04467, 2023. URL https://apt htt.tttp::/ppp..smaticclllol
semanticscholar wttp://ssssooooooooo//oooo/////a
10

---

Published as a conference paper at COLM 2025
Yu Bai, Heyan Huang, Cesare Spinoso-Di Piano, Marc-Antoine Rondeau, Sanxing Chen, Yang Gao, and Jackie Chi Kit Cheung. Identifying and analyzing task-encoding tokens in large language models. 2024. URL https: www.ppsssaaacccllaa.ooomm wwt.ccccoooaa..oorrruuuesoeeesee org/CorpusID:: 267068913.
Róbert Belanec, Simon Ostermann, Ivan Srba, and Mária Bieliková, Task prompt vectors: Effective initialization through multi-task soft-prompt transfer. ArXiv, abs/ 0bss/200.000 2408.01119, 2024. URL https: /api. semanticscholar. Swwpihollrrrrg///ooooooo org/CorpusID : 271693681.
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor V. Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. ArXiv, abs/2303.00 2/S//230.0811,, 2023. URL attpps//aap) seaatticccat org/
ItrpusID:25550492
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. aiv preprint ar orXim:2004.051: 2020.
S. Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and Ilms by learning to learn discrete functions. ArXiv, abs/ 2bs///300330 2023. URL https: http:://app..smantiicchoo...
263320005....
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss: Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXio, abs/ 220/5.005.1155 14165, 2020.
Stephanie C.
Chan, Ishita Dasgupta, Junkyung Kim, Dharshan Kumaran, Andrew Kyle Lampinen, and Felix Hill. Transformers generalize differently from information stored in context vs in weights. ArXiv, abs abs12220.05675, 2022a. URL https: //api. semanticscholar. arg/CorrpusD::22246775.
Stephanie C. Y. Chan, Adam Santoro, Andrew Kyle Lampinen, Jane X. Wang, Aaditya K Singh, Pierre H. Richemond, Jay Meclelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. ArXiv, abs/2205.05055 2022b. URL httpss//aaps eeemanticsccolllee :otparrooo///r 248665718..
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. 2022. URL https: /ppi. semanticscholar org/CorpusID: 258686544.
Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. Analyzing transformers in embed- ding space. Xiv preprint ar arXiv:2209.0283 022555,, 2022.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds. Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical frame- work for transformer circuits. Transformer Circuits Thread, 2021. https://trant ttaes://tranformer- dircuits.pub/211 /001 http:/pp/p///////eoort....
Eric Elmoznino, Tom Marty, Tejas Kasetty, Leo Gagnon, Sarthak Mittal, Mahan Fathi, Dhanya Sridhar, and Guillaume Lajoie. In-context learning and occam's razor. arx preprint
(rXio:2410.14086, 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. ArXiv, abs/ /200..00001 01066, 2022.
11

---

Published as a conference paper at COLM 2025
Marta Garnelo, Dan Rosenbaum, Chris J. Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Jimenez Rezende, and S. M. Ali Eslami. Conditional neural processes, ArXiv, abs/ fbss/180.001113,, 2018. URL https://app semanticscholar.cm org/
CorpusID:49574999.
Mor Geva, Avi Caciularu, Ke Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. ArXiv, abs/22033 2203.2220,4680 2022a. URL https://apii semnnticcccaa hwtpaaar.rrr//rrrsi::2277777 /tt 247762385.
Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg, Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. ar arv preprint arXiv:2203.14630, 2022b.
Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris P'apailioooulos. Looped transformers programmable computers. ArXiv, abs/2301.13196, 2023.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. ArXiv, abs/ 2310..22,, 2023. URL https://api. wwwanticscholeeeee..e www.ttticccccoooooo/ooooooo org/CorpusID: 263608983.
Wes Gurnee and Max Tegmark. Language models represent space and time. ArXio, abs/2310.02077, 2023. URL httppss//ppp..emmttiiiiicoo wwmanticschollaat httpat/tiicc.ooooooooo////ooo//a orgg/orppuID::236887755
Seungwook Han, Jinyeop Song, Jeff Gore, and Pulkit Agrawal. Emergence of abstractions: Concept encoding and decoding mechanism for in-context learning in transformers. 2024.
of
URL https://api. wwt.deesmaaaacoo.l.....o wwwpoollarrrrrgroeeeeeeeeeee
274799679
Michael Hanna, Sandro Pezzelle, and Yonatan Belinkov. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms. ArXio, Abbs/2403.177 2024.
URL https://api. semant semanticscholar www.oo///////.o Corpus 268 268691 935.
Yun He, Huaixiu Steven Zheng, Yi Tay, Jai Gupta, Yu Du, V. Aribandi, Zhe Zhao, Yaguang Li, Zhaoji Chen, Donald Metzler, Heng- Tze Cheng, and Ed H. Chi. Hyperprompt: Prompt- based task-conditioning of transformers. In International Conference Machine Learning, 2022. URL httpps:/appi. semanticscholar. Crgp/orp 247218862.
Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vec- tors. ArXiv, abs/2310.15916, 2023. URL https://api..smmnttccchollll. org/CorpusID: 264439386.
semanticscoolarr
Alberto Hojel, Yutong Bai, Trevor Darrell, Amir Globerson, and Amir Bar. Finding visual task vectors. ArXiv, abs/ 2404.05729, 2024. URL https://api. htttp::/app..smmmniicccee wwwanticschollettt org/CCrpusID:: 269005382.
Edward S. Hu, Kwangjun Ahn, Qinghua Liu, Haoran Xu, Manan Tomar, Ada Langford, Dinesh Jayaraman, Alex Lamb, and John Langford. The belief state transformer. 2024.
URL https: http:////app.ssmnticscchhohooooo semanticscholar httpttiiicccoolaaooooooorrrrreie
273707334
Shih-Cheng Pin-Zu Li, Yu-Chi Hsu, Kuang-Ming Chen, Yu Tung Lin, Shih-Kai Hsiao, Richard Tzong-Han Tsai, and Hung yi Lee. Chat vector: A simple approach to equip lims with instruction following and model alignment in new languages. In Amnual Meeting of the Association for Computational Linguistics, 2023. URL hittps://apl: semanticschol.a. org/Corpus DB/CCorpsSD::222533322 268253329.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arith- metic ArXio, Abs//222.04488, 2022. URL httpss///apl https:://api..mmmaiccchhtte semantccscoolarr .aar..ogg/Cee www./oorgsssi.... 254408495.
Pavan Kalyan, Preeti Rao, Preethi Jyothi, and Pushpak Bhattacharyya. Emotion arithmetic: Emotional speech synthesis via weight space interpolation. Interspeech 2024, 2024. URL https:/ //api. semanticscholas ott.choo.....ooo/.oooppo/2211110
12

---

Published as a conference paper at COLM 2025
Ji-Hun Kang, Jaehwan Lee, Mun-Hak Lee, and Joon-Hyuk Chang. Whisper multilingual downstream task tuning using task vectors. Interspeech 2024, 2024. URL https://api. semanticscholar wwtg/CorrusI:::23355
Junyeob Kim, Hyuhng Joon Kim, Hyunsoo Cho, Hwiyeol Jo, Sang-Woo Lee, Sang goo Lee, Kang Min Yoo, and Taeuk Kim. Ground-truth labels matter: A deeper look into input- label demonstrations. ArXia, abs/2205. 12685, 2022. URL https: //app. semanticscholas
org/CorpusID: 249062788.
Seijin Kobayashi, Simon Schug, Yassir Akram, Florian Redhardt, Johannes von Oswald, Razvan Pascanu, Guillaume Lajoie, and João Sacramento. When can transformers compositionally generalize in-context? ArXiv, abs/2407.12275, 2024. URL https:
sespiiismmtticschhee http:/oo....o.r//roussi:::111111 271244618..
Jannik Kossen, Yarin Gal, and Tom Rainforth. In-context learning learns label relation- ships but is not conventional learning. 2023. URL https://ap..smmaniiccchoo semanticscholar... org/ CorpusID: 2601 260123277. 25327.
Vedang Lad, Wes Gurnee, and Max Tegmark. The remarkable robustness of Ilms: Stages of inference? ArXiv, abs/2406.1900 2024. URL https://api. wewaiieeeaaiiiiiiii wwwmmannccccolaa.orr..
CorpusID:270764655.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learn- ing. ArXiv, abs/ Abss/206.14892, 2023. URL https://api. wwwpaantisccolla.oo/oooooooro 259262142.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Conference on Empirical Methods in Natural Language Processing, 2021.
URL https://api. http:ss////pp.semaniishhhhoooo.
www.oorg//oooo CorpusID 233296808.
Dongfang Li, Zhenyu Liu, Xinshuo Hu, Zetian Sun, Baotian Hu, and Min Zhang. In- context learning state vector with inner and momentum optimization. 2024a. URL
https://api. semanti wwtp:/oorpssD:2299900 269187789
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. 2023.
Zhuowei Li, Zihao Xu, Ligong Han, Yunhe Gao, Song Wen, Di Liu, Hao Wang, and Dim- itris N. Metaxas. Implicit in-context learning. ArXiv, 2405/.40.,6660 2024b. URL
https://api. semanti .rgg/CorpusID 269983389.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. ArXiv, abs/2230.08866, 2023. URL //ap:://api. semanticscholar org/CorpusID:
htt//CorpuIDD:2239997788 263909278.
Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. arXiv preprint ar arXiv:2402.18819, 2024.
Daria Lioubashevski, Tomer M. Schlank, Gabriel Stanovsky, and Ariel Goldstein. Looking beyond the top-1: Transformers determine top tokens in order. ArXiv, abs/2410.22210,, 2024. URL https: /api..pp../// wwwanticscholaieee. org/CorpusID) 2 273655045
Michael Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, 14, 2001.
Sheng Liu, Haotian Ye, Lei Xing, and James Y. Zou. In-context vectors: Making in con- text learning more effective and controllable through latent space steering. ArXio, abs/231 2ts//3111.6666 2023a. URL htttps://app semanticscholar org/CorpusID:
sat.aiiicccooooooo////rrrssss 265149781.
Xiao Liu, Yanan Zheng, Zhengxiao Du, Ming Ding, Yujie Qian, Zhilin Yang, and Jie Tang. Gpt understands, too. ArXiv, abs/2103.100855 2021. URL https: //pii. semanticscholar. org/CoorusID::22220 2322669696.
13

---

Published as a conference paper at COLM 2025
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: Contextual sparsity for efficient lims at inference time. In International Conference on Machine Learning, 2023b. URL https://ppi.semmnticchheee wwwanticscholar...... .or.aa.rrrr//ooeeeeee wwtp/CoppsSI::::...ooomm 26081 15690.
Quanyu Long, Yin Wu, Wenya Wang, and Sinno Jialin Pan. Decomposing label space, format and discrimination: Rethinking how Ilms respond and solve tasks via in-context learn- ing. ArXiz, abs/2404.07546,, 2024. URL https: /ttps://apil semant i cscholar. wwweerorrggooooeeeeeee C 269042687.
Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. Sparser faster and less is more: Efficient sparse attention for long-range transformers. arXiv preprint ar XrXoo2240..1674) 2024.
Grace Luo, Trevor Darrell, and Amir Bar. Task vectors are cross-modal. 2024. URL
https://app semanticscholar. hrt.aaarrrr//orrssi22235555
Ang Lv, Kaiyi Zhang, Yuhan Chen, Yulong Wang, Lifeng Liu, Ji-Rong Wen, Jian Xie, and Rui Yan. Interpreting key mechanisms of factual recall in transformer-based language models. ArXiv, 22//224431.111,11 2024. URL https://app http::://app.smmmticccoo..o.... semanticscholar org/CorpusID: 268732668.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. In Neural Information Processing Systems, 2022. URL https:
//api. seappi.sematiicccoooeooooo wwwp///r////oD///
ott.://mmn..ccco..oooooooooooooo 255825985.
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple
word2vec-style vector arithmetic. ArXio, abs/2305.16130, 2023a. URL https://api.
http://cchhhcaacoororrrrrrspseee 258387799.
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. A mechanism for solving relational tasks in transformer language models. 2023b.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. ArXiv, abs/2110.15943,, 2021. URL https: semanticschon org/ 20028810:240288833
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Conference on Empirical Methods Natural Language Processing, 2022.
Sarthak Mittal, Eric Elmoznino, L'eo Gagnon, Sangnie Bhardwaj, Dhanya Sridhar, and Guil- laume Lajoie. Does learning the right latent variables necessarily improve in-context learn- ing? ArXia, hass/2405.1110 2024. URL https://app semanti aichoolarr org/CorpusID: 270095225.
Jesse Mu, Xiang Lisa Li, and Noah D. Goodman. Learning to compress prompts with gist tokens. ArXiv, abs/2304.0844 08467, 2023. URL https://api. httpss///ppi.smmttiiscchiiie wwwantccshollar...mm cem.anntcccoooaa.ror//ooooaaaa www.aaar.rr///rrerrr org/CorpusID: 258179012.
Samuel Muller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. ArXiv, abs/ 2112.10510, 2021.
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalehi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre-Luc Bacon. Bridging state and history representations: Understand- ing self-predictive rl. ArXiv, abs/2401.08889, 2024. URL htttps://ppi. semanticscoola.. Ttg/CorpusDD:22702774.3
nostalgebraist. interpreting gpt: the logit lens., 2020. URL http:://wwBssooaaaaavvvvrrteeeee sepppreeinggg-tt--eee looelloogtteees
com/
Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen. What in-context learning "learns" in-context: Disentangling task recognition and task learning. In Anmual Meefing of the Association for Computational Linguistics, 2023, URL https:// /appi.//pi... semanticscholar.org .rg//
CorpusID:258740972.
14

---

Published as a conference paper at COLM 2025
Core Francisco Park, Andrew Lee, Ekdeep Singh Lubana, Yongyi Yang, Maya Okawa, Kento Nishi, Martin Wattenberg, and Hidenori Tanaka. Iclr In-context learning of
representations.
URL o0
Kiho Park, Yo Joong Choe, and Victor Veitch. The linear representation hypothesis and the geometry of large language models. ArXie, abs/2311.03658,, 2023. URL https:
sew.eticsmniicccii wwwpcchlar.rrg/oorrppppeee wtgg/CorpssI::222
Yingzhe Peng, Chenduo Hao, Xu Yang, Jiawei Peng, Xinting Hu, and Xin Geng. Learnable in-context vector for visual question answering. ArXiv, abs/2406.13185, 2024. URL
https://apl.. semantisschoooarr
Minh Pham, Kelly O. Marshall, Chinmay Hegde, and Niv Cohen. Robust concept erasure using task vectors. ArXiv, abs/2404.03631,, 2024. URL https://api. semanticscholar.
org/CorpusID: 268889344.
Jason Phang. Investigating the effectiveness of hypertuning via gisting. ArXiv, abs/ 2Dss/2400.16817. 2024. URL htt.htttss/ppp..ssaatiiccollllll
semantisccollrr.rrg//roaa www./oorrsii:266000000 2680332888.
Steven T Piantadosi, Dyana C.Y. Muller, Joshua S. Rule, Karthikeya Kaushik, Mark I. Goren- stein, Elena R Leib, and Emily Sanford. Why concepts are (probably) vectors. Trends in Cognitive Sciences, 28:844-856, 2024. URL https://api. www.ssiimaaattnnieeeeee org/CorpusID: 271721102.
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint ar axXf::1911.02972, 2019.
Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli.
The effects of pretraining task diversity on in-context learning of ridge regression. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
Siyu Ren, Qi Jia, and Kenny Q. Zhu. Context compression for auto-regressive transformers with sentinel tokens. In Conference on Empirical Methods in Natural Language Processing, 2023. URL https: /api. wwwanticscholeeeeeoooe.oooococee
org/CorpusIDD /t: 263909553.
Baturay Saglam, Zhuoran Yang, Dionysis Kalogerias, and Amin Karbasi. Learning task representations from in-context learning. In ICMI. 2024 Workshop on In-Context Learning.
Pratyusha Sharma, Jordan T. Ash, and Dipendra Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. ArXiv, ab3/2212.13558 2023. URL https: /api.. semanticscholl org/ Cogp/CrrpssI.... (0: 266435969.
Suzanna Sia, David Mueller, and Kevin Duh. Where does in-context translation happen in large language models. ArXiv, abs/2403.04510, 2024. URL https://p// semanticscholar. org/CorpusID:2282260 268264275.
Aaditya K. Singh, Ted Moskovitz, Felix Hill, Stephanie C. Y. Chan, and Andrew M. Saxe. What needs to go right for an induction head? a mechanistic study of in- context learning circuits and their formation. ArXiv, abs Abs/2404.77129, 2024. URL https://api. wew.aannnniiiss.....oooooooooroo .rrp/Corpus10::29933311
Siva K. Swaminathan, Antoine Dedieu, Rajkumar Vasudeva Raju, Murray Shanahan, Miguel Lázaro-Gredilla, and Dileep George. Schema-leaming and rebinding as mechanisms of in-context learning and emergence. Arviv, abs/2307.00 01201, 2023. URL https://api.
semanticscholar. org : 259342490.
Zhixu Tao, Ian Mason, Sanjeev Kulkarni, and Xavier Boix. Task arithmetic through the lens of one-shot federated learning. 2024. URL http:::/ap..semmaticsceee semanticscholar. CIr 274306177.
Eric Todd, Millicent Li, Amab Sen Sharma, Aaron Mueller, Byron C. Wallace, and David Bau. Function vectors in large language models. ArXiv, abs/2310.15213, 2023. URL https://api. semanticscholar httpaar.rr//goooso/024442333
15

---

Published as a conference paper at COLM 2025
Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ArXio, abs/ /t///2222.07677/ 2022.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymy- rov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimization algorithms in transformers. ArXiv, abss/330.008855 2023, URL attpss//appi. semanticccamaa org/ CorpusID:261 696852.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in gpt-2 small. ArXio, abs/2211.00593, 2022. URL https://api. semanti orgg/CorpuSID:2532222222
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. Conference on Empirical Methods in Natural Language Processing, 2023. URL https://app wemmanticscoolaat www.orgg///ooooo wttpg//cCID::::ii/ 258841111
Chris Wendler, Veniamin Veselovsky, Giovanni Monea, and Robert West. Do llamas work in english? on the latent language of multilingual transformers. ArXiv, abs/ 2402.10588 2024. URL https: / /api.......... semanticscholar
Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Ad- vances in Neural Information Processing Systems, volume 36, pp. 36637-36651. Curran Associates, Inc., 2023. URL httptttss//prrrceeeing..... neurips cwtpcper.ppppeeeeeeeceeeeeeee wwt.feellleeppe...../ pdf..
Shiguang Wu, Yaqing Wang, and Quanming Yao. Why in-context learning models are good few-shot learners? In International Conference Learning Representations (ICLR), 2025. URL / wtt.:e//ooooummiiiiilsseeeeeeee
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. ArXiv, abs/2111.02080, 2021.
Zhaozhuo Xu, Zirui Liu, Beidi Chen, Yuxin Tang, Jue Wang, Kaixiong Zhou, Xia Hu, and Anshumali Shrivastava. Compress, then prompt: Improving accuracy-effcciency trade-off 11m inference with transferable prompt. ArXiv, abss23305.1186. 2023. URL https://api. semanticccholoar wtt.aant...ccc.....oooooooooooo.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable
narrow model selection capabilities in transformer models. ArXia, abs/2211..00 2023a.
URL h otrp://eeeettccoo...oo.ooooooorr 2649333329
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Can transformer models generalize via in-context learning beyond pretraining data? In NeurlPS 2023 Workshop on Distribution Shifts: New Frontiers with Foundation Models, 2023b.
Qinyuan Ye, Bill Yuchen Lin, and Xiang Ren. CrossFit: A few-shot learning challenge for cross-task generalization in NLP. In Marie-Francinn Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7163-7189, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.1 (0.18653/v11/21 org//0211 emn.p-main.
emnlp-main.577 URL
Qinan Yu, Jack Merullo, and Ellie Pavlick. Characterizing mechanisms for factual recall in language models. ArXie, absi,aass/2a0 221001210,,9 2023. URL hntps:://pp. semanticscholar
org/CorpusID: 264439114.
Jooyeol Yun and Jaegul Choo. Scaling up personalized aesthetic assessment via task vector customization. ArXiv, abs/ abs//240.07717, 2024. URL https: /taps://api... semanticscholar.oo CorpusID: 221pslDD:7718899..
16

---

Published as a conference paper at COLM 2025
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generaliza- tion. ArXiv, abs/2305.194D) 2023. URL https://apt sew..iissmttccchollrrr wwt.//orps//o///// 258987402.
Haiyan Zhao, Heng Zhao, Bo Shen, Ali Payani, Fan Yang, and Mengnan Du. Beyond single concept vector: Modeling concept subspace in lims with gaussian distribution. 2024. URL
https://api. commanticscholal wttpecchoola...oooooossssiss../e orgg/CorpusD::23922271
Yufan Zhuang, Chandan Singh, Liyuan Liu, Jingbo Shang, and Jianfeng Gao. Vector-icl: In-context learning with continuous vector representations. 2024. URL https://api.
semanticscholar.cm :ra.aarrioor//oo 273227918
A Related Work (cont.)
Understanding Pre-trained Large Language Model's Depth-Dependent Behavior. Re- searchers have extensively studied the layer- and depth-dependent behaviors of large language models (LLMs) along the forward pass. nostalgebraist (2020), Belrose al. (2023), and Wendler al. (2024) analyzed how GPT models gradually determine top tokens through the logit lens and embedding space (Dar et al., 2022; Geva et al., 2022b). Lioubashevski et al. (2024) further highlighted that transformers identify top-k tokens sequentially across layers. Beyond token predictions, several works investigate layer-specific encoding of information. For instance, Lv et al. (2024), Yu et al. (2023), Meng et al. (2022), and Gurnee & Tegmark (2023) explore how factual associations are stored in a layer-dependent manner. Additionally, Liu et al. (2023b) studied contextual sparsity changes across layers, while Sia et al. (2024) observed that "task recognition" typically occurs during the middle stages of the forward pass. More broadly, phase transitions in transformer behavior across layers have been discussed by Lad et al. (2024).
Leveraging these insights into layer-dependent behaviors, several works propose strategies to enhance LLM performance. Sharma et al. (2023) demonstrated that selectively apply- ing low-rank decomposition in specific layers can improve reasoning ability in question- answering tasks. Meanwhile, Sia et al. (2024) and Liu et al. (2023b) introduced methods to reduce inference computation by identifying and leveraging specific layer behaviors. The concept of task vectors introduced in Hendel et al. (2023) also reflects a layer-dependent behavior in the context of in-context learning.
Understanding Pre-trained Large Language Model's In-Context Learning Behavior Since the launch of GPT-3 (Brown et al., 2020) and the observation that, during inference, the model can leverage few-shot examples to improve performance, researchers have been keenly interested in understanding the mechanisms of this in-context learning behavior. Several studies have explored the significance of labeling in context demonstrations (Min et al., 2022; Kim et al., 2022; Kossen et al., 2023; Long et al., 2024), while others have examined the phenomenon through circuit mechanisms (Elhage et al., 2021; Wang et al., 2022; Singh et al., 2024; Hanna et al., 2024; Wang et al., 2023). Research rch by Dai et al. (2022), Geva al. (2022a), and Merullo et al. (2023b) links in-context learning behavior to implicit internal weight and activation updates within the model. To better understand "learning" in-context, several works (Pan et al., 2023; Sia et al., 2024) have studied its two distinct phases-tast recognition and task learning-usinn
controlled experiments designed disentangle these phases. Park et al. (2024) studied the phase transition along the context length and depth. These studies focus on pre-trained large language models without access to the pre-training procedure. The literature on embedding editing and task vector is discussed in the main text.
Understanding In-Context Learning Behavior in the Controlled Small-Scale Setting on Data-Fitting Problems. Researchers have also examined in-context learning behavior in small-scale models, where models are trained from scratch to perform in-context learning tasks. Garg et al. (2022); Akyürek et al. (2022); von Oswald et al. (2022; 2023) investigated the
17

---

Published as a conference paper at COLM 2025
ability of transformers to learn regression problems, interpreting the model as performing a single step of gradient descent. This line of research was extended further in different pre-training mixtures of tasks (Raventos et al., 2023; Yadlowsky et al., 2023a), the ability to generalize to unseen tasks (Yadlowsky et al., 2023b), explanations via the Bayesian optimal estimator (Zhang et al., 2023; Bai et al., 2023; Muller et al., 2021), the viewpoint of generalization error (Li et al., 2023), and the perspectives of task retrieval and task learning (Lin & Lee, 2024; Wies et al., 2023).
Beyond regression, several studies have explored in-context learning abilities in other domains, such as reinforcement learning (Lee et al., 2023; Lin et al., 2023), discrete function learning (Bhattamishra et al., 2023), factorial hidden Markov chains (Xie et al., 2021), and deterministic finite automata (DFAs) (Akyürek et al., 2024). Swaminathan et al. (2023) studied the ICL mechanism with the clone-structured causal graphs to understand the schema learning, retrieval and rebinding. Additionally, Chan al. (2022a;b) investigated the behavior of in-context learning and in-weight learning on the Omniglot datasets.
Task Vector Applications Li et al. (2024a) utilized this task vector to enhance test-time adaptation, while Pham al. (2024) employed it to erase malicious concepts during pre- training. Beyond activations, task information can also be encoded in a task token (Bai et al., 2024), or a pause token can be used to gain extra computation time (Goyal et al., 2023). Zhao et al. (2024) propose to present each task concept with Gaussian distribution. Following the notion of "task vector" the weight space, several works utilize the weight space difference for chatbot on new language (Huang et al., 2023; Kang et al., 2024), emotion transfer (Kalyan et al., 2024), aesthetic assessment (Yun & Choo, 2024), and soft prompt initialization (Belanec et al., 2024).
Soft Prompt and Efficient Adaptation of Transformers. In this work, we primarily follow the definition of a task vector as described by Hendel al. (2023). This task vector effectively compresses the information in the context into single vector, which can be injected into the model to perform zero-shot inference. The notion of a task vector also related to the concept of soft prompt (Lester et al., 2021; He et al., 2022; Liu et al., 2021; Xu et al., 2023; Kang et al., 2024), where the context information is similarly compressed. Beyond learning soft prompts that encode instructions, researchers have also explored context compression by modifying the attention mask to ensure the summarization of context at specific token (Ren et al., 2023; Mu et al., 2023; Phang, 2024). In the format of ICI., the researcher also study the continuous representation in context (Zhuang et al., 2024), as reminiscent of the soft encoding.
Although we propose a training algorithm to encourage the formation of task vectors, our primary goal is not compare our method of encoding task vectors with these approaches in soft prompt learning. Instead, we aim to understand the benefits of incorporating task vectors into in-context learning.
Latent Representations and Sufficient Statistics in In-Context Learning. The TVP-loss in our approach can be interpreted as encouraging the model to form a latent representation that serves sufficient statistic of the task. Prior works have explored similar ideas in different contexts. As discussed in the main text, Hendel et al. (2023) observe that pre- trained LLMs first summarize the task into a latent vector representation, which is then used to guide task execution. Mittal et al. (2024) investigate whether learning an appropriate latent representation-via a bottleneck architecture-cal improve robustness in in-context learning. Additionally, Kobayashi et al. (2024) show that such bottleneck architectures can enhance compositional generalization. These approaches are closely related to amortization- based meta-learning frameworks (Garnelo et al., 2018; Wu et al., 2025), and to methods that learn belief-state representations (Hu et al., 2024) or predictive state representations in reinforcement learning (Littman & Sutton, 2001; Ni et al., 2024). Our work connects to this line by explicitly supervising the formation of task-representative latent vectors through the proposed TVP-loss.
18

---

Published as a conference paper at COLM 2025
B Experimental Setup
B.1 Synthetic Experiments
Following the setup in (Lin & Lee, 2024; Garg et al., 2022), the input vectors x; E Rd are sampled from a normal distribution N (0, 1). The linear function weights w E R which define the task, are sampled from a mixture of d Gaussians, each with means corresponding to the standard basis vectors: 14i E R with Mij 1 if =, and 0 otherwise. All Gaussians share a covariance matrix Σ = I, and each component is selected with equal probability. The target outputs are computed as f(x)) w x;. We set d = 6 in our experiments.
Linear Regression.
Sinusoidal Regression. Following the same sampling strategy as the linear regression task, the target outputs are defined f(x;) s0.5 w x; + b), where b (0,1). In this task, w and determines the underlying function.
Discrete Token Offset Prediction. In this task, we set the vocabulary size to be C, and each input token x; is discrete value drawn from the set {0, C C- 1}. The target output is defined as f(x:) = (a x; + b) mod C, where a E [1,2,3), E [0.1,2) are uniformly sampled, and mod denotes the modulo (remainder) operation. The task is uniquely determined by the choice of a and b, resulting in total possible tasks. By default,
set C 1000.
For the three synthetic tasks, we evaluate the following out-of-distribution (OOD) tasks:
1. Linear Regression: We introduce an OOD prompt using the quadratic regression task, defined as f(x;) = TT (xi xx), where denotes element-wise multiplication.
2. Sinusoidal Regression:
consider the OOD prompt where the task is fxx)
/ //max(05. w x; h, + where h, ~ (O,0))) 1).
3. Discrete Token Offset: the OOD prompt switches the context at the 6th position.
B.2 In-Context Learning for Formal Language
We provide the full experimental setup details for the two formal language datasets below.
B.2.1 The GINC Dataset: the Factorial Hidden Markow Chain.
The GINC dataset is introduced in Xie et al. (2021) to study the behavior of language pretraining and the emergence of in-context learning. We follow the setup as described in Xie et al. (2021): we define a uniform mixture of HMMSs over a family of five concepts. Each entity and property is assigned unique tokens, resulting in total vocabulary size of 100. To generate samples, we first select an HMM from the mixture and use it to produce documents containing 576 tokens. Next, we randomly insert 192 dummy tokens, denoted as "-," into these documents. These dummy tokens are labeled as -000 to be ignored during processing, resulting in documents with a total length of 768 tokens.
We use the GPT-2 model with 8 layers, 8 heads, and an embedding dimension of 256. The context length of this transformer is set to 768. To train the model to form task vector at the 1-th layer, we sample 10 extra tokens from the same HMM and randomly insert dummy token "_" This forms the "zero-shot" document. The task vector loss is then calculated by injecting the /-th layer's hidden states at the token "_" into the "_" token in the zero-shot document. This training pipeline (shown in Figure 17) differs from the Meta-ICL. pipeline described earlier (Figure 13) because GINC specifies pre-trained distribution that is different from the prompt distribution. In GINC, there are no explicit input-output pairs; instead, the entire trajectory performs one task. Therefore, we encode the task vector into the inserted token "_" in the trajectory.
To evaluate in-context learning performance, we follow the setup described in Xie et al. (2021) and sample prompts from the five concepts used during training. An example of an in-context prompt with context length of k = =3 and example length 4 is: b d
19

---

Published as a conference paper at COLM 2025
f - k Similarly, an example zero-shot prompt with example length would be: m n o p. Here, the input x represents a short document, such as a b c, and the task f corresponds to an underlying hidden Markov model. The output f(x) is the next sampled token in the document. Unlike the synthetic tasks described earlier, this function f is stochastic.
To evaluate task vector prompting performance, we extract the hidden states at the i-th "-" token and inject them into the I-th layer activations of the zero-shot prompt at its "-" token. This setup aligns with the input format described in Hendel et al. (2023). Notably, because the first "-" token in the in-context prompt appears after an initial document of the HMM (e.g,, the b document in the earlier example) and is then injected into the zero-shot prompt, additional context information from this HMM document may be included, which could lead to better task vector prompting performance compared to in-context learning performance. Through out the experiment, we set N 1, ie. the task vector is extracted from a single demonstrated prompt.
B.2.2 The RegBench Dataset: the Probabilistic Finite Automata.
The RegBench dataset is introduced in Akyürek et al. (2024) to study the behavior of language pretraining and the emergence of n-gram head. We follow and revise the setup of Akyürek et al. (2024): we sample 100 deterministic finite automata (DFAs) for training. At each training iteration, we randomly generate 5,000 training sequences from these DFAs. Each DFA has fixed initial state, and each state transits to another state given an input token. For each DFA, the total vocabulary size (i.e., the number of unique tokens) is set to 40, the maximum number of states is set to 20, and each state is allowed maximum of 10 outgoing edges. To generate samples (sequences) from a DFA, we convert it to probabilistic finite automaton (PFA) by assigning uniform transition probabilities across all outgoing edges, following Akyürek et al. (2024).
Each prompt contains 32 samples (sequences) generated by the sampled DEA and those example are with length s separated by the special token "I". To help the formation of task vectors, we insert dummy token ">" before the last token of the example. After inserting >, prompt with context length k = 3 and example length 4 would be like " efg>>d/efghh| jijki""
We use the GPT-2 model with 8 layers, 2 heads, and an embedding dimension of 128 following Akyürek et al. (2024). For each training sequence, we sample an additional example from the same DFA and insert dummy token > before the last token of the example. This serves as a zero-shot prompt. Similar to the GINC dataset, the input x corresponds to the generated sequence of tokens from the PFA. The task f is determined by the PÉA's transition function, and the output fxx) is the next token sampled by the DFÁ. Importantly, this transition function is stochastic as well. During training, use the next token prediction loss on all tokens except the special tokens ">" and "I". During inference, we evaluate task vector prompting performance with N = 1.
C Supplements on the Trained-from-Srracch Model
C.1 Effects of Input Formats
In this section, we examine different input formats for in-context learning prompts and their effects on the emergence of task encoding. The various prompt formats are detailed in Table 3. For clarity, we refer to the token positions where task vector presence is examined as the task token locations.
We focus on the case where the problem dimension is d = 6 and the model depth is L = 3, as exhibits the most pronounced task encoding, as shown in Figure 2. In Figure 5, we evaluate various task vector extraction methods (detailed in Appendix C.2), with their respective task vector prompting performance shown as line with triangular markers. Notably, only the model with the input format (*xy) demonstrates task encoding that surpasses random
20

---

Published as a conference paper at COLM 2025
baselines. By default, use the format (*xy) throughout the main paper unless specified otherwise.
C.1.1 Attention Cluster Visualization
To understand the non-random performance observed in the task vector prompting mode, we analyze how the model encodes task information using attention maps (left) and a PCA of activations at each layer while varying the number of in-context examples (right), when input format is *xy in Figure 6 and when input format is x->y in Figure 7. In transformer models, attention mechanisms enable each input token to dynamically focus on specific parts of the input data by assigning importance to them. This process can be visualized using attention maps, which quantify the relationships between input tokens. These maps are computed using query and key vectors: the query represents the current focus of interest, while the key interact with the query to calculate importance scores (via their dot product), determining how relevant each input element is to the query.
Recall that we follow the experimental setup in Section 4.1 and conduct experiments on the linear regression task using 3-layer transformer with 63-shot context. On the left, the attention map of a specific head in 2nd layer shows that the activations at x; position primarily attend to the activations of the preceding V;-1, where task information is stored. Furthermore, the activations at the y/i location attend to both themselves and the preceding y; 1, enabling the model to update task information online. On the right, we present a PCA of the activations for token y; (i (0, 4, 6,8)) across layers I,, evaluated on three different linear functions: w = [1, [1,0,0,0,0,0,", w² (0,0,,,0,0000 10,0,1,0,0,0] and W3 = [0,0,0,0,1,0]?
When the input format is *xy, at the input to the third layer, the activations for different tasks begin to form distinct clusters, though there is still considerable overlap and dispersion among the clusters. However, for x->y, across different layers, the attention maps do not exhibit meaningful task information encoded in the patterns. Consequently, we showcase only one example attention map to illustrate the absence of distinct task encoding in this setup. Additionally, the PCA of the activations reveals that the encoded information is nearly identical across tasks.
C.1.2 Discrepancy with Pre-Trained LLM
In the pre-trained LLM (Hendel et al., 2023), the input format uses x->y. There is a clear difference between the pre-trained LLM and the task information encoded. In the pre-trained LLM, task information is primarily stored in the "maps-to" token (->). Conversely, in the trained-from-scraat transformer, task information is stored in the y token (i.e., the "label" token), but only when x and alternate closely in the input format.
transformer in how
We hypothesize that this discrepancy arises because pre-trained LLMs learn the semantic meaning of the "maps-to" symbol (->) from their extensive pre-training corpus, enabling them to use -> as an anchor for task summarization. Additionally, this delimiter -> is needed to separate the and entries. In contrast, for the hained-from-ssrated transformer, the -> token functions more like a <pause> token (Goyal et al., 2023), providing additional computational resources rather than semantic significance. Moreover, the training procedure ensures that x and y tokens occupy fixed positions, removing the need for delimiter to separate them.
Furthermore, Wang et al. (2023) show that label words themselves can serve as anchors for aggregating task information in context. This observation aligns with our finding that task information is encoded in the y tokens. While this result is not explicitly framed within the task vector framework, it highlights the critical role of label tokens in task encoding.
C.2 Various Task Vector Locating Methods
In this section, we analyze different methods for extracting task vectors. The main question we address is: If task encoding is detectable by one task vector extraction method, will it also be detectable by another method? Our focus is not on the quality of the extracted task vector but
21

---

Published as a conference paper at COLM 2025
Table Different prompt formats examined. We analyze the input format based on the settings described in Garg et al. (2022) and Hendel et al. (2023). Here, z represents a special token, Pquery denotes the test query format, and Pt refers to the k-shot context. We examine the activations at A,, referred to as the task token location.
P 'query
K-shot examples Pk
Task Token Location Ay
(*xy) (Gang et al., 2022)
Niest
f Nkr (xx Xx+1]
[o::::]: (f(x,)) and
(x-Py) Hendel et al. (2023)
Ntest
X, (x) ), (xt 1
A[1::33:
(Mx-2y) Hendel et al. (2023)
Niett
2,81, 2, (x1), N2 (20) x3-1
:10::3]: (f(x;) and
ICL
Nean Squared
1.5
Error
1.0-
X->y
1.0-
+x->y
0.5
0.5
0.0
20
40 60
0
20
40
60
20
40 60
of in-context samples
Figure 5: In-context learning (ICL. and task vector prompting (TVP) performancc across input formats and fask vector extraction methuds. This figure compares task vector prompting performance across four task vector extraction methods (detailed in Appendix C.2) and three input formats: (*xy), (x->y), and (xx->y). Method (A) is the default task vector extraction method, and (*xy) is the default input format used throughout the main paper. As shown, for the (*xy) input format, noticeable task encoding is observed, with performance exceeding random predictions, regardless of the extraction method. In contrast, for the (x->y) and (*x-2y) input formats, no task encoding can be reliably extracted by any method, resulting in task vector prompting performance that is nearly random.
rather on whether the method can detect task encoding that is non-random. Below, we present task vector extraction methods inspired by Hendel et al. (2023); Liu et al. (2023a); Li et al. (2024b), which are representative approaches in the task vector literature.
Following the definition in Section 3, we study the following task vector extractor methods:
(A) Task Token's Output Embedding Hendel et al. (2023) have observed that in pre- trained language models, the I-th layer output embedding of the k-th task token, denoted as bi encodes the task vector. We follow this setup and define the task vector extractor as follows: Let T It then
M(Pqeery; I() == replace task token's s 1-th layer embedding with hn-
This notion of the task vector indicates the independence of the encoded task information from the demonstrated prompt.
(B) Task Token's Output Embedding Difference Liu et al. (2023a) pointed out that when inputting two sequences x and y, the difference in the hidden space corresponds to the mapping from to y in the pre-trained model. Though their setting is slightly different from the in-context learning format, we leverage this idea and propose the following way to extract the task embedding: Following the definition of hi from above, and let Ri be the hidden state of the model when inputting the test prompt Pquery, then T= Rt j Arl indicating the difference between the k-th context embedding and uninformative context, then
((qquery; Ahk) add Ark to task token's 1-th layer embedding.
This notion of the task vector loosens the constraint of absolute independence of the task from the demonstrated prompt.
(C) Principal Direction of Task Token's Output Embedding Difference Liu et al. (2023a) also observed that taking the principal direction along the difference of hidden states
22

---

Published as a conference paper at COLM 2025
Query
input
91
32
Y3
X44
Y5
Y6
37
in-context samples
Figure 6: Attention map and PCA visualization of activaatioss across three different linear functions. (Left) For the linear regression task described in Section 4.1, the attention map illustrates how each query (rows) attends all available keys (columns), with each row summing 1. The heatmap reveals that the activations at the x; positions predominantly attend the activations of the preceding Vi-1- where task information tion is stored (highlighted by the black boxes). Additionally, the activations at the yi positions attend to both themselves and the preceding V/-1, enabling the online updating of task information (highlighted by the white boxes). (Right) PCA visualizations of token yi activations C [0, 4, 6, 8)) across layers reveal that task-specific clusters (three colors correspond three different tasks) begin to emerge at the output of the 2nd layer, indicating that the model progressively encodes task information as depth increases.
08
Query
L=1
044
imput
2
8
Key
$ of in-context samples
Figure 7: Attention map and PCA visualization of activations for three linear functions with imput format (x->y). Following the same setup as in Figure 6, we plot the attention map and PCA visualization of the -> token's activations. (Left) The attention map at the 2nd layer highlights interactions between task tokens (white boxes) and the attention from tokens to task tokens (black boxes). Unlike the format (*xy), this format does not exhibit meaningful attention pattern where x; tokens attend to previous task tokens, or task tokens attend prior task tokens. (Right) PCA visualizations show that the activations the -> token encode nearly identical information across different context lengths and different tasks, suggesting the absence of distinct task encoding dynamics in this input format.
aligns more closely with the task. Motivated by their observation and setting, we propose the following task vector extractor: Following the definition of Alr,, from above, we set PCA(A))) then
M(Pquery; (CAAUA)):PAAAAMM))) PCA(AN)))
== add PCA(A))) to task token's 1-th layer embedding.
We also re-scale the updated hidden states to preserve the model's original capabilities. This approach differs slightly from the original setup in Liu et (2023a), where PCA is applied to the concatenated hidden states across all layers. Their method is specifically designed for alignment tasks in natural language settings, where the context x; and y,, are input separately. In our case, the model is trained in an in-context learning (ICL) format, we adhere to the input prompt structure where x; and y/ are provided alternately in one prompt. We also attempted to follow their setup by applying PCA to the concatenated
23

---

Published as a conference paper at COLM 2025
hidden states across all layers; however, this approach resulted in near-random task vector prompting performance.
(D) Linear Combination Task Token's Output Embedding to the Query Output Embed- ding Li et al. (2024b) propose to learn the coefficient of the linear combination between the task token's output embedding to the query output embedding. Motivated by this, we let r = (" then
M(Pquery; (1g)e = replace task token's 's 1-th layer embedding wi th
at 1 + B/
AI E [L]-
where leee is the output embedding I-th layer for the test query, and at, and & are learnable parameters. To learn these coefficients, we use constant learning rate of 0.01, the A optimizer, and train for total of 100 epochs. Following the initialization setup in et al. (2024b), the coefficients are initialized as af = 0.1 and B =1
C.2.1 Task Vector Extractor Performance
We evaluate the performance of various task vector extraction methods across different input formats, as described in C.1,, and present the results in Figure 5. When using the *xy input format, the task vector extraction methods demonstrate noticeable task encoding. However, for the other input formats, informative task encoding is detected, regardless of the extraction method used.
It is worth noting that the learnable task vector extraction method (D) performs slightly worse than the other three methods, possibly due to its initialization. We did not focus on further improving this method, as the primary goal of this section is demonstrate the detectability of task encoding rather than to optimize its quality.
C.3 Formal Description of Layer Localization for Task Vector
For each task f, we generate N demonstrated prompts (PP where the same function f is consistently used across all prompts, and the input to each prompt p is uniformly sampled. Let Pquery, with test input X represents the corresponding test queries for the j-th demonstrated prompt. For a task vector extracted from a k-shot prompt, the optimal layer index |* is computed as:
1" = arg min Σ l (f(stes), (PQuery:((M((PP)))
where ((M(PP"),, I) denotes the task vector extracted from the /-th layer when processing the demonstrated prompt P/ and M(Pquery: T) represents the model's prediction for the query prompt Piueryy where the task vector I extracted from the I-th layer is copied to the corresponding position in the model's representation when processing Pauery. A lower loss l indicates that the corresponding layer provides a more effective encoding of task-specific information.
C.3.1 Task Vector Layer Localization in Linear Regression Tasks
When extracting the task vector, the extractor must determine which layer the task vector resides in, based on the N demonstrated prompts. For linear regression task with the exper- imental setup described in Section 4.1 and Appendix B.1, we measure model's performance in the TVP mode when the extracted task vector is placed at each layer and average the loss across varying context lengths. Figure (left) demonstrate that among the 3 layers, the 2nd layer is predominantly selected as the location where the task vector resides.
24

---

Published as a conference paper at COLM 2025
L-4
L-5
L-6
1.0
1.0
1.0
Layer usage
0.8-
Iope
0.8-
0.4
layer 2
lyyer
byee
Agss
Jager
txpst
0.2
0.2
0.2
0.0
20
40
0.0
20
40
0.0
26
40
60
in-context samples
Error
Layer
Figure 8: Layer selection distribution (upper row) and averaged task vector prompting (TVP) performance (lower row) for lask vector emergence in Linear Regression with problem dimension d 6 and model depths == (left), L (middle), and L (rightt) Upper row: The frequency of each layer being selected as the task vector location across varying numbers of in-context examples is shown. As model depth increases, task information becomes less concentrated in any single layer, resulting in a more distributed selection of task vector locations. Lower row: The averaged TVP performance for each layer, measured across various context lengths, is presented. The dashed line indicates random-guu guess performance, representing a scenario where task information is inferred. These results demonstrate that for deeper models (especially information, and task vector performance approaches random, highlighting the distributed nature of task encoding in deeper architectures.
5, 6), no single layer consistently encodes significant task
This observation contrasts with the phenomenon reported in pre-trained LLMs by Hendel et al. (2023), where task vectors for in-context learning tasks primarily emerge in the early layers out of a total of around 20 layers. In our experiments, however, we find that task vectors consistently emerge in the penultimate layer for shallow models (e.g., layers). As the model depth increases (e.g., 4 or more layers), task vectors no longer emerge distinctly at any single layer, shown in Figure Instead, task information becomes distributed across multiple layers, reducing the distinctiveness of any specific layer as the task vector's 's locatio focation. As shown in Figure for deeper models, particularly with L and 6, no single layer consistently encodes task information, and the task vector prompting performance appears random. From this analysis, we conclude that in the linear regression task, and under the experimental setup we examined, task vectors either emerge in the penultimate layer when the model is shallow or fail to emerge entirely when the model deep.
One factor to consider is the nature of the task at hand. For the linear function class, the "task" appears to be identifying the parameter w, after which the "task execution" phase simply requires the model to compute w" x to generate predictions. To explore this further, examine the sinusoidal regression task using the same GPT-2 model configured with embedding size of 64, 4 attention heads, and layers. Additionally, we analyze the discrete token offset task with larger embedding size of 256.
As shown in Figure 9 (middle), for this 6-layer transformer model, the 3rd layer is primarily selected the task vector location. This finding aligns with observations in pre-trained LLMs and provides the insight that the task vector's 's emergence layer in trained-from-scratel models depends on the characteristics of the task being performed. Additionally, for the discrete token offset task, the 5th layer is most frequently selected as I*, with other layers also being selected occasionally lower probabilities.
C.4 Effects of Model Capacity
In this section, we examine the model with same depth, but different embedding size, to investigate the effect of model capacity. Specifically, for the three tasks mentioned in Section 4.1, we study the feaned-from-scrattee model's performance model depth 4, 6 or 8, with embedding size 64 (with notation "S") or 256 (with notation "L"), and present the result in Figure 10.
25

---

Published as a conference paper at COLM 2025
Linear Regression
Sinusoidal Regression
Offset
Error
10
lrgHt
2
11340 2
3
10
Figure 9: Task vector prompting performance for task vector emergence in: Linear Regression (left), Sinusoidal Regression (middle), and Discrete Token Offset (right). We present the averaged task vector prompting (TVP) performance for each layer, measured across various context lengths. The dashed line represents random-guess performance, indicating no task information is inferred.
As illustrated in the figure, increasing the embedding size allows the model to discover shortcut solutions for regression tasks. Specifically, the model approximates the least squares solution in a single forward pass, eliminating the need to store intermediate task encodings. Conversely, for the discrete token offset task, the model operates on tokens represented in learned embedding space. To solve f(x) to another using the task-specific parameters a and b. We hypothesize that the nature of token embeddings makes it less straightforward for the model to learn shortcut solution for this task. As result, the model is compelled to store the task information explicitly. Consequently, increasing the model capacity provides more space for storing and organizing task encodings, leading to clearer task representations as capacity grows.
+ b, the model must map one token
Linear Regression
Sinuseidal Regression
Offset
LACS)
1463 16R3
L415) LA(S)
14L1
L5(3)
1863
1033
Error
55
0-
0
0 20 40
60
20
40 60
of in-context samples
Figure 10: Effects of model embedding dimension on task vector prompting performance. This figure shows the task vector prompting performance of models with different embedding sizes (64 for notation "S" and 256 for notation "L"). For the regression tasks (Linear and Sinusoidal), larger embedding sizes make task encoding less prominent, likely because the increased capacity allows the model to find "shortcut" solutions, such as approximating the least-squares solution directly, without the need to store intermediate task information. In contrast, for the discrete token offset task, larger embedding size enables the model to store task information in more separable and structured manner.
C.5 Effects of Sparse Attention
Sparse attention (Lou et al., 2024) has been proposed to improve transformer efficiency with minimal performance degradation. In in-context learning tasks (detailed in Section 4.1), model could potentially solve these tasks using sparse attention: For each query, the model only needs to attend to itself and the previously stored task information. Specifically, each x; token needs to attend two tokens (x; and y;-1), while each y; token needs to summarize the task information by attending to three tokens (yi 1/-1, and x;).
In this section, we investigate the impact of explicitly enforcing sparse attention constraints on in-context learning and task vector prompting performance. Using the setup described in Section 4.1, we apply a sliding window strategy (Qiu al., 2019; Beltagy al., 2020) combined with causal attention, with window sizes of 3, 5, and 7. We evaluate the transformer with various depth on the linear regression task with = (where task vectors emerge) and d = 9 (where task vectors nearly absent), and present the results in Figure 11.
As shown in the figure, models trained with sparse attention demonstrate better task encoding compared to those with full-window causal attention. For instance, when and L 5, task encoding in the full-window model fades after the first few in-context
26

---

Published as a conference paper at COLM 2025
vanilla
s-3
5-5
S-7
dm6 L
Lm4
Lm6,
Mean Squared
0.5
Error
0-9.L-5
d-9. L-6
d-9 9-7
20
40 60
20
20 40 60
in-context samples
Figure 11: Task vector prompting performance for models trained with sliding window sparse attention. We compare models trained with full-window causal attention (vanilla) and sparse attention using sliding window sizes 5, and Solid lines represent the in-context learning performance, while lines with extra triangular marker indicate the task vector prompting performance. As shown, sparse attention enhances task vector emergence but leads to a slight degradation in in-context learning performance.
In-Context Learning
Task-Vector Pronpting
Ours
1.0
N-2
0.5
N=10
20
40
60
0.0
20
40
60
N-50
Figure 12: Impact the Number of Demonsirated Prompts (N) Task Vector Prompting Performancc. Increasing N improves task vector prompting performance, larger values yield more concentrated task vectors. Across N settings, model trained with TVP-loss ("Ours") consistently outperforms the vanilla model, demonstrating better task vector encoding. The dashed line represents random-guess performance.
examples. In contrast, models trained with sliding window attention (window sizes 3, 5, and 7) maintain stable task encoding across varying context lengths.
However, while the sliding window improves task vector emergence, it slightly degrades in-context learning performance because each query has access to fewer tokens during training. Moreover, task encoding stability is not always guaranteed: larger window sizes improve in-context learning performance but degrade task vector prompting performance.
C.6 Effects of Number of Demonstrated Prompts N
Throughout the paper, we N 50 by default when evaluating task vector prompting performance. Intuitively, larger N values yield more concentrated and purified task vectors. To illustrate the effect of N, we evaluate task vector prompting performance for N 1,2,5, 10, following the setup in Section 4.1, and present the results in Figure 12. In this experiment, the transformer trained on linear regression task with problem dimension of d = 6 and transformer depth L = 3. For the model trained with TVP-loss, the task vector explicitly formed at the 2nd layer.
As shown in the figure, when N = 1, the vanilla model yields random task vector prompting performance. With increasing N, task vector performance improves steadily. Notably, the model trained with TVP-loss (denoted as "Ours" in the figure) demonstrates better task vector encoding compared to the vanilla model across all values. Because the linear functions in regression tasks have weights that exist in a continuous space and can be very close to each other, separating tasks with only a single demonstration challenging. As result, it requires N = 5 to achieve task vector prompting performance comparable to its in-context learning performance. For the vanilla model, however, even with N 50, task vector prompting performance remains slightly worse than in-context learning performance.
27

---

Published as a conference paper at COLM 2025
D Supplements on Training with TVP-Loss
D.1 Demonstrating Our Method in Meta-ICL Training Fine-Tuning
We illustrate the training algorithm in Figure 13, where the red arrow indicates the flow of the gradient.
Yeer
Transtorner Block
Transformer Black
prodiction
Transformer Block
Tr frrsrrrr Block
entor back- propogation
D1
X/2
Ye
(z >iest)
Few-shot context
zero-shot query P qurry
Figure 13: Demonstration of our training algorithm. vanilla Meta-ICL training, the model is updated using the ICL-loss signal from the few--hot context. To encourage the formation of task vectors, we also explicitly include the TVP-loss from the zero-shot query. This means the model is asked to predict Vest when only Xtest and the injected hidden states are given. In the given illustrated example, there are in total layer in the transformer model, and set = 1 encourage the formation of the task vector the first transformer block's output.
D.2 Full Results for Section 5.2
In Section 5.2, we present the in-context learning and task vector prompting performance at the 63rd context position for models trained with and without the TVP-loss, along with the in-context learning performance on out-of-distribution prompts. Figure 14 extends these results by showing the full performance across varying context lengths. This analysis complements the main results, demonstrating that the observations at the 63rd position are not unique but instead hold consistently across different context lengths.
D.3 Performance on Additional OOD Tasks
For the experiment result presented Section 5.2, we evaluate the model's performance additional out-of-distribution (OOD) prompts for the regression tasks, specifically: (a) logistic regression, (b) in-context learning with an outlier context presented with probability of 0.1, where the outlier is defined as with noisy labels (noise level 0.3), (d) scaled weights (x3).
and y 1. (c) in-context learning
For the linear regression training task, models with task vectors formed at the 7th layer (i.e., the last layer) perform equal to or slightly better than the vanilla-trained model, except when the weights and inputs are scaled. This suggests that models trained with TVP-loss exhibit reduced robustness when the input range is significantly altered. For the sinusoidal regression task, models with task vectors formed at the 3rd layer exhibit equal or improved OOD performance.
D.4 In-Context Multi-Function Learning
To further explore compositional tasks, investigate the model's ability to generalize to unseen functions by training on multiple compositional functions. Specifically, we follow the linear regression task's sampling strategy to sample and train on the following function classes:
1. Square-Root Regression: f(x;) = Vmax(0.5. w x; + h,0) + b, where h, ~ (0,11).
2. Sinusoidal Regression: f(x;) (0.5 T b), where (0,1))
28

---

Published as a conference paper at COLM 2025
OOD
1.0
Linear Regresston
40
20
40
60
1.5
vanilla
Sinusofdal Regression
ours (1 13
1.0
0.25
0.25
ours (f-
ourS (1
10-
of In-context samples
Figure 14: Model Performance with Vamilla and TVP-Loss Training. We compare the in-context learning (ICL), task vector prompting (TVP), and out-of-distribution (OOD) performance of models trained with vanilla training (black) and those trained with our TVP-loss, where task vectors are formed at different layers
1,3,5,7) Full results are shown across varying numbers of in-context examples for three tasks: linear regression (top row), sinusoidal regression (middle row), and discrete token offset task (bottom row).
L Linear Regression
S Sinusoidal Regress fon
2+30P
10a
Logistic Regression
6x1e11
1x10
4x2811
20
40
60
20
60
100
6x30-1
4x3c-1
Outlier
10-1
2
10-1
60
40
6x28
vanilla
ours ()
Input
ours ()
0
20
40
60
20
40
60
101
Sx 101
Scaled
10
20
60
Figure 15: bn- Context Learning Performancc on OOD Prompis for Models Trained with Regression Tasks. We compare the in-context learning performance of the vanilla-trainet model (labeled oanilla) and models trained with auxiliary loss to form task vectors different layers (labeled ours (D)). We evaluate the models on OOD prompts, including: (a) logistic regression, (b) outlier in context, (c) noisy input labels with noise level of 0.3, (d) scaled weights (x3).
29

---

Published as a conference paper at COLM 2025
000:
000: sine
000: signoid
000: quad
0.02
05
2
00
0.00
Vani
Vanilla
Vanilla
TV Layer ID
Figure 16: Out-of- Distribution (OOD) Performance with Vanilla and TVP-Loss Training. This figure compares the OOD prompt performance of models trained with vanilla training and TVP-loss across layers where task vectors are formed. Dashed lines represent the performance of vanilla-trained models for reference. The horizontal axis indicates the layer in which task vectors are formed. TVP-loss improves OOD generalization, particularly for the sine and quadratic regression tasks, while achieving comparable performance to vanilla training for the sigmoid task, and slightly worse performance on square root regression task..
1 3. Sigmoid: f(xi) = 1+exp(-0.55 -055- xx+b)' where ~ (0,
4. Quadratic Regression: f(x;) 0.5( T -h)2
+ where h, (0,1).
During training, we designate one function i E {1,2,3,4} as the cut-of-distribution (OOD) function and train the transformer using in-context prompts sampled from the remaining three function classes, each sampled uniformly. We then evaluate the model's compositional generalization on the OOD function.
We train an 8-layer transformer with 4 attention heads and an embedding dimension of 256, keeping other training hyperparameters consistent with Section 4.1. To examine the effect of task vector location, we test placing the task vector at the 3rd, 5th, or 7th layer. The results, shown in Figure 16, represent the averaged in-context learning performance across context lengths. These results demonstrate that TVP-loss enhances OOD generalization, particularly for the sine and quadratic regression tasks, achieves comparable performance to vanilla training for the sigmoid task, and shows slightly reduced performance for the square root regression task.
D.5 Indentification of task vector in GINC dataset
In GINC, there explicit input-output pairs; instead, the entire trajectory performs one task (hidden markov model). Therefore, we encode the task vector into the inserted token in the trajectory. We provide illustration of our training algorithm in Figure. 17.
ab
ah
ba
Transformer Block
Transformer Block
prediction
Transtormer Block
Transformer Black
back- propogation
ae context Pr
ab
zero-shot query P
Piverry
Figure 17: Demonstration of our training algorithm in GINC dataset. In contrast to the Meta-ICI. format described in Figure 13, in GINC dataset, there trajectory performs one single task defined by hidden markov model. Therefore, we randomly insert several 4" tokens into the trajectory in the fhw-shot mode, while inserting one +" token into the trajectory in the zero-shot mode. The task vector at the /-th layer is then injected into the embedding in the zero-shot mode.
explicit input-output pairs. Instead, the entire
30

---

Published as a conference paper at COLM 2025
D.6 Fully Converged Performance the RegBench Dataset
As noted by Reviewer GxZE, our results on the RegBench dataset (Figure 4) were based on undertrained models. To address this concern, double the training duration and report updated performance in Figure 18 and Figure 19.
Figure 18 shows the next-token prediction accuracy over the course of training. Both the baseline and our TVP-loss model converge by epoch 40, indicating comparable performance at converging. However, as shown in Figure 19, our method continues to demonstrate better ICL and TVP performance, especially at prompt length s = 4. These results suggest that TVP-loss provides structural benefit by encouraging more robust and transferable task representations, leading to stronger ICL generalization.
s=2
s=d
6
all costat Tengin
Reghorh
0.0
c.e
0.6
0.6
D-2
50
103
100
200
100
200
200
Training
Stegs
Training Steps
Training Steps
Training Steps
Figure 18: Training accuracy on the RegBench dataset over 40 epocks. Both the baseline and TVP-loss models converge in next-token prediction.
Dn-Dntest Learning
Vester Prarpting
(c) Wolos Irput
(4) Duange
Cargee
Figure 19: In-context learning (ICL) accuracy on RegBench af convergence, The TVP-loss model consis- tently outperforms the baseline across different prompt lengths, particularly at 4.
31


Reviews:
[{"review_id": "LSVh9lxOXw", "reviewer": "Reviewer_9mvp", "summary": "This study examines how task‑specific representations (“task vectors”) arise inside transformers during in‑context learning. Training models from scratch on synthetic data, the authors show that such vectors do emerge, but are often diffuse or weakly expressed. To concentrate and strengthen them at chosen layers, they introduce a task‑vector prompting loss (TVP‑loss) applied during training. TVP‑loss removes the need for post‑hoc searches for task‑correlated activations and leads to models with more robust, localized task vectors that generalize better across contexts.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The authors designed a good synthesized experiment that is small-scaled and easier to control. Thus they did a full set of ablation study to investigate the key factors for the natural emergence of task vectors.\n2. This work proposed a simple yet effective way to strengthen the task encoding and improve the performance of the later task vector prompting.\n3. The experiments are intensive and the results are informative.", "reasons_to_reject": "The emergence of task vectors during training, as one of the major contributions, is not well supported by the results in Section 4.2.1. While the authors' conclusion is that task vectors are most distinct with **moderate** model depth, my interpretation can be that the \"task vectors vanish when the model depth is large enough\", which is the usual case for large language models. The condition for the emergence shown in Section 4.2.1 is so rigid that there is no specific implication on the LLM scenario.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DLITp7QC6T", "reviewer": "Reviewer_BR3X", "summary": "This paper discusses task vector formation in transformers​ for ​improving robustness and generalization in in-context learning (ICL)​. The paper (1) claims that task vectors emerge weakly and non-locally in pre-trained models, limiting their utility, and (2) proposes a ​task vector prompting loss (TVP-loss) to explicitly guide their formation during training. Experimental results on ​synthetic regression, formal language, and natural language tasks​ show that TVP-loss significantly enhances zero-shot performance compared to baseline methods.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "S1. A ​controlled experimental framework​ to study task vector emergence from scratch, addressing gaps in understanding how task information is encoded in pre-trained LLMs.\n\nS2. A simple yet effective auxiliary training mechanism with a task vector prompting loss (TVP-loss),  that encourages task vectors to form at prescribed model layers, improving robustness to out-of-distribution (OOD) prompts.\n\nS3. Expensive experiments across multiple tasks show its effectiveness.", "reasons_to_reject": "W1. Although empirical results are strong, the paper lacks a ​theoretical explanation​ for why TVP-loss works.\n\nW2. The synthetic nature of tasks limits conclusions about real-world applicability.\n\nW3. TVP-loss adds computational overhead. Its scalability to larger models/architectures is also untested.\n\nW4. The selection of some hyperparameters may require more explanation.", "questions_to_authors": "Why does TVP-loss work?", "extra_scores": {}}, {"review_id": "XlYtYNxnll", "reviewer": "Reviewer_GxZE", "summary": "The paper looks into the formation of task vectors in small transformer models trained on synthetic datasets, and studies this over a set of hyperparameters and input structures. It also proposes a loss to aid in the formation of task vectors in such models.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Current work focuses on looking at task vectors in pre-trained LLMs, so there is novelty in studying them in smaller transformers. \n- There are clean experiments conducted to show the emergence of task vectors over the course of training transformers. There are sufficient ablations to show the effect of multiple model and training hyperparameters.", "reasons_to_reject": "- It is somewhat obvious already that adding a loss that is identical to the metric being tested (TVP performance) will improve on that metric. In this regard the new loss seems to not be an insightful discovery. It would improve the paper to show other benefits of adding this loss and improving tvp performance. The OOD performance is an attempt at this but as highlighted later, it is not as convincing.\n- It is quite unclear to me why the vanilla models are significantly underperforming in the experimental sections 5 and 6. As noted on line 248, in figure 3 the tvp performance of the vanilla model is near random. This is a surprising result because regularly trained transformers do exhibit evidence of task vectors, even in this paper itself in Figure 1. This extreme result is cause for concern. Similarly in section 6 it is not clear why the vanilla model should perform worse on ICL. \n- The benefits of the new loss on OOD are unconvincing. In section 5.2 and figure 3, all (but a few) of the results are near random performance. Even the ones that do perform better than vanilla are not significantly better. In light of the close-to-random performance this result seems not indicative of an improved model.\n\n  In Section 6 and figure 4, again the claimed result only seems to be true for a few settings, s=6 and s=9. As noted before, for these specific results the vanilla model seems to be undertrained even for ICL performance, and the OOD results simply mirror this lack of performance. \n  \n  In summary, both experiments are not convincing enough to claim an improvement in OOD generalization as a result of the TVP loss. \n- What is the performance of the Natural Language Task without implicitly selecting the best layer to minimize TVP loss? How is this layer chosen at inference time? \nThe current setting serves as an unfair comparison to the baseline, as there is implicit training on validation performance in this method. The extra overhead of calling the model 8 times to check performance on the test query is significant.", "questions_to_authors": "- Since adding the tvp loss supposedly increases robustness, do you expect the inverse to be true? Would techniques to improve robustness affect the tvp performance?\n- What is the performance of the Natural Language Task if you choose and set a single layer as the other experiments?\n- Why is vanilla transformer performance low in sections 5 and 6? \n- Is the performance on OOD prompts in 5.2 checked in ICL or TVP mode?", "extra_scores": {}}, {"review_id": "5vq6aLYyiw", "reviewer": "Reviewer_g1s8", "summary": "This paper studies the emergence of task vectors during LM training, and then proposes a loss to encourage the formation of strong task vectors. Experiments show that the proposed loss further improves models ICL performance and the generalization ability on OOD data.\n\nThe paper is overall well-written, and the visualizations are clear. The proposed loss is well-motivated with validation on experiments.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The problem of studying task vectors and using them to enhance prompting is important and can benefit LM tuning efficiency. \n\n2. The paper is well-written, with experimental settings and results clearly described. The experiments are thorough, conducted on both synthetic and real datasets. \n\n3. The proposed loss is well-motivated, and the experiments further validate its efficacy.", "reasons_to_reject": "1. The experiments on natural language tasks are only conducted under the MetaICL setting. What are the reasons behind that and how does the proposed loss perform in the general ICL/TVP setting?", "questions_to_authors": "1. In Fig 3, which base model does TVP evaluation use at the inference time? The model $M _\\theta$ trained together with in-context loss or the untrained base model $M$?\n\n2. It could be helpful to provide more insights on why adding the TVP loss can improve the ICL/MetaICL performance.", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper addresses the emergence of task vectors during training and suggests a loss that encourage such emergence. The reviewers appreciated the empirical setting, and the effectiveness of the new training approach.
Decision: Accept