Task Vectors in In-Context Learning: Emergence, Formation, and Benefits

We thank the reviewer for their thoughtful and constructive feedback. We are glad that the importance of studying task vectors, the clarity of our writing and visualizations, and the effectiveness of the proposed loss are recognized and appreciated. Below we address the reviewer’s concerns in detail.

Q1. 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?

We appreciate the question and the opportunity to clarify. The MetaICL setting refers to fine-tuning a pretrained language model on a diverse set of in-context learning (ICL) tasks, as introduced in [1], and then evaluating the model on held-out ICL tasks. In our experiments, both our method and the vanilla MetaICL baseline follow this process. The only difference is that our method incorporates the TVP-loss during fine-tuning. Evaluation is performed in ICL mode for both methods.

Since these held-out ICL tasks are not seen during training, and we do not explicitly supervise task vector formation for them, it is not expected that the model would perform strongly in TVP mode. However, we did evaluate TVP performance in the hr→lr setup. Specifically, we evaluated the TVP loss across 10 prompts to select the optimal layer index $l$, and then used the embedding from layer $l$ to perform task vector prompting for the remaining prompts within the task.

While the vanilla MetaICL model achieved a TVP performance of $31.81 \pm 0.13$, our method yielded $34.15 \pm 0.14$, suggesting that TVP-loss slightly improves task localization ability—even on tasks that were not explicitly supervised during training.

[1] Min, Sewon, et al. "Metaicl: Learning to learn in context." arXiv preprint arXiv:2110.15943 (2021).

Q2. In Fig 3, which base model does TVP evaluation use at the inference time?

In Figure 3, TVP evaluation is conducted using the models trained with the standard ICL loss (i.e., the $M_{\theta}$ model) for the vanilla baseline, and with both ICL loss and TVP-loss for our method.

It could be helpful to provide more insights on why adding the TVP loss can improve the ICL/MetaICL performance.

We appreciate this important question. Our insight is that TVP-loss introduces an architectural bias that guides the model to develop localized and consistent task-specific representations at particular layers. This design aligns with our empirical finding that such structured representations naturally arise in moderately sized models. In this sense, TVP-loss amplifies a tendency that the model already exhibits when its capacity is appropriately matched to the task.

We thank the reviewer for their thoughtful and constructive feedback. Below we address the reviewer’s concerns in detail.

Q1. Adding a loss that is identical to the metric being tested (TVP performance) will improve on that metric is obvious.

Thank you for the comment. While optimizing a loss aligned with the evaluation metric can lead to expected gains, our key contribution is showing that adding TVP-loss during ICL training (1) improves TVP performance, (2) does not degrade ICL performance, and (3) can improve OOD generalization. To our knowledge, this observed benefits are non-trivial and require empirical validation.

Q2. TVP performance of Vanilla model in Section 5 are nearly random.

Thank you for this insightful observation. As discussed in Figure 2, the emergence of task vectors is closely tied to model capacity. Specifically, task vectors tend to emerge only when the model is not overparameterized relative to the training dataset.

In Figure 1, the model has 3 layers, which is sufficient to learn the task but not deep enough to bypass the need for explicit task representation. In contrast, Figure 3 uses an 8-layer model for the linear regression task. This deeper model has enough capacity to solve the task directly in a single forward pass, without needing to encode task information in intermediate representations. As a result, the TVP performance of the vanilla 8-layer transformer is near random.

Q3. Is the performance on OOD prompts in 5.2 checked in ICL or TVP mode?

The performance on OOD prompts in Section 5.2 and throughout this paper is evaluated in ICL mode.

Q4. In Section 6, vanilla model in RegBench seems to be undertrained.

We appreciate this observation. We extended training for both the vanilla model and our method, and updated the ICL results in Figure 6 here. To confirm convergence, we provide training accuracy curves here, showing the mean DFA accuracy across context lengths converges.

In the updated figure, we still observe a gap in ICL and OOD performance, particularly when s=4. We believe the ICL and OOD improvement on RegBench reflects that TVP-loss helps structure the learning of these ICL tasks.

We also want to clarify that TVP-loss is not designed to explicitly optimize for OOD performance. Rather, the observed improvement is a byproduct of encouraging structured task representations through task vector formation.

Q5. In the natural language task, how is the model evaluated during inference?

During inference, we follow the ICL-style evaluation protocol from the original MetaICL paper [1]. For each test query, the model evaluates the conditional log-likelihood of each candidate answer given the prompt, and selects the answer with the lowest loss.

Q6. How is this layer chosen at inference time?

Since the model is evaluated purely in ICL mode on various held-out datasets, no task vector layer is selected or applied during inference. We also discuss the potential of reporting the TVP performance of the OOD prompt in response to Reviewer g1s8 Q1.

Q7. There is implicit training on validation performance in this method.

We clarify that the validation set is strictly held out from the training set. The decision to use adaptive layer selection was made based on its performance on the hr→lr task as opposed to select a single layer (as shown in Q8). For the other two tasks, we applied the same hyperparameter setup without further tuning.

We want to note that our goal is to study if there is any benefit using TVP-loss in the ICL performance of the large-scale natural language models.

Q8. What is the performance of the Natural Language Task if you choose and set a single layer as the other experiments?

Thank you for this insightful question. In the hr->lr task, we conducted additional experiments where the task vector location was pre-assigned to specific layers: 12, 14, 16, 18, and 20. The results are as follows:

In the MetaICL setup, the fine-tuning dataset contains a variety of tasks. Fixing the task vector at the 12th layer gives similar performance to vanilla, but performance drops consistently at deeper layers. We hypothesize that this is because different tasks encode best at different layers, and fixing the location limits the model’s ability to represent all tasks effectively. Adaptive selection during training provides the flexibility needed to capture task-specific information more robustly.

We thank the reviewer for their thoughtful and constructive feedback. We are glad that our controlled framework, the simple and effective TVP-loss, and the extensive experiments across diverse tasks are recognized and appreciated. Below we address the reviewer’s concerns in detail.

Q1. Lack of theoretical explanation for why TVP-loss works. Why does TVP-loss work?

Thank you for this important question. TVP-loss can be viewed as “working” in two ways: improving task vector identifiability (TVP performance) and improving OOD ICL performance.

For the first, TVP-loss is explicitly designed to align intermediate representations with task-specific vectors. Minimizing this loss directly improves task vector quality, as supported by standard empirical risk minimization theory.

For the second, we offer the following intuition: TVP-loss introduces a structural bias that encourages the model to encode task information consistently at specific layers. This regularization effect aligns with our empirical finding that task vectors naturally emerge in moderately sized models, suggesting that TVP-loss reinforces a behavior the model is already inclined toward under the proper capacity. A full theoretical explanation is left to future work.

Q2. Synthetic nature of tasks limits conclusions about real-world applicability

Prior work [1] has demonstrated the emergence of task vectors in real-world LLMs under simple ICL tasks. Additionally, in Section 6.2, we show that applying TVP-loss during fine-tuning improves ICL performance on the MetaICL benchmark using a GPT2-Large (774M parameter) model. These results support our claim that task vectors can emerge in real-world LLM and that explicitly encouraging their formation enhances generalization in more realistic, natural language settings.

[1] Hendel, Roee, Mor Geva, and Amir Globerson. "In-context learning creates task vectors." arXiv preprint arXiv:2310.15916 (2023).

Q3. TVP-loss adds computational overhead; scalability to larger models is untested

Thank you for raising this point. We acknowledge that TVP-loss introduces additional training-time overhead due to the extra forward pass and potential layer selection. In our experiments, this results in a modest 1.5 to 3x increase in training time for both small models and GPT2-Large (774M). Inference under ICL mode incurs no extra cost compared to the baseline.

We note that no effort was made to optimize efficiency. We leave such improvements to future work, as our current focus is to characterize the emergence of task vectors and demonstrate their benefits.

Q4. Hyperparameter choices require more explanation

Thank you for pointing this out. We will expand the description of our hyperparameter selection process in the Appendix. Specifically, for the synthetic tasks, we follow the previous setup [1,2]; for the formal language tasks, we follow the original paper [3, 4]'s setup, only adding a TVP-loss to the training procedure. For the MetaICL tasks, we also follow the experiment setup outlined in the original paper [5].

[1] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.
[2] Lin, Ziqian, and Kangwook Lee. "Dual operating modes of in-context learning." Forty-first International Conference on Machine Learning. 2024.
[3] Xie, Sang Michael, et al. "An explanation of in-context learning as implicit bayesian inference." arXiv preprint arXiv:2111.02080 (2021).
[4] Akyürek, Ekin, et al. "In-context language learning: Architectures and algorithms." arXiv preprint arXiv:2401.12973 (2024).
[5] Min, Sewon, et al. "Metaicl: Learning to learn in context." arXiv preprint arXiv:2110.15943 (2021).

We thank the reviewer for recognizing our controlled synthetic setup and thorough ablations. We address your concern below.

"Task vectors vanish when the model depth is large enough"—disagrees with behavior in large language models.

While standard LLMs are indeed deep, they are typically trained on datasets that are far larger and more diverse. In our synthetic experiments (Figure 3 and Figure 10), we observe that for linear regression task, the task vectors emerge most clearly in moderately capacity models--those with enough capacity to support structured task encoding, but not so much that task information can be handled without forming explicit representations. We conjecture that this regime--where model capacity is more closely aligned with task complexity--might share some properties with the conditions under which task vector emergences in large pretrained models [1]. We leave a deeper investigation of this connection to future work.

[1] Hendel, Roee, Mor Geva, and Amir Globerson. "In-context learning creates task vectors." arXiv preprint arXiv:2310.15916 (2023).