Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers

• URL: http://arxiv.org/abs/2605.03780v1
• Date: Tue, 05 May 2026 14:07:55 GMT
• Title: Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers
• Abstract summary: Transformers are effective at inferring the latent task from context via two inference modes.<n>Recent interpretability studies have identified from middle-layer representations task-specific directions.<n>We show that two inference modes can coexist within a single model.
• Score: 6.89278796238822
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformers are effective at inferring the latent task from context via two inference modes: recognizing a task seen during training, and adapting to a novel one. Recent interpretability studies have identified from middle-layer representations task-specific directions, or task vectors, that steer model behavior. However, a lack of rigorous foundations hinders connecting internal representations to external model behavior: existing work fails to explain how task-vector geometry is shaped by the training distribution, and what geometry enables out-of-distribution (OOD) generalization. In this paper, we study these questions in a controlled synthetic setting by training small transformers from scratch on latent-task sequence distributions, which allows a principled mathematical characterization. We show that two inference modes can coexist within a single model. In-distribution behavior is governed by Bayesian task retrieval, implemented internally through convex combinations of learned task vectors. OOD behavior, by contrast, arises through extrapolative task learning, whose representations occupy a subspace nearly orthogonal to the task-vector subspace. Taken together, our results suggest that task-vector geometry, training distributions, and generalization behaviors are closely related.

Related papers

• PACT: Preserving Anchored Cores in Task-vectors for Model Merging [68.52455853496585]
  Model merging aims to combine multiple task-specific fine-tuned models into a single multi-task model.<n>Most existing model merging approaches follow the Task Arithmetic paradigm.<n>We propose PACT, which preserves the anchored task-specific cores (i.e., LBW dimensions) within task vectors by aligning their complements with the subspace of the pre-trained weights.
  arXiv  Detail & Related papers  (2026-06-17T02:48:35Z)
• Purifying Task Vectors in Knowledge-Aware Subspace for Model Merging [83.5273168208788]
  Model merging aims to integrate task-specific abilities from individually fine-tuned models into a single model without extra training.<n>The merged model often suffers from notable performance degradation due to the conflicts caused by task-irrelevant redundancy in task vectors.<n>We propose Purifying TAsk Vectors (PAVE) in knowledge-aware subspace to overcome these challenges.
  arXiv  Detail & Related papers  (2025-10-16T14:02:57Z)
• Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations [19.539276425108987]
  This work proposes the Linear Combination Conjecture, positing that task vectors act as single in-context demonstrations formed through linear combinations of the original ones.<n>We show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis.<n>We predict the failure of task vectors on representing high-rank mappings and confirm this on practical LLMs.
  arXiv  Detail & Related papers  (2025-06-10T17:59:31Z)
• When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers [64.1656365676171]
  Task arithmetic refers to editing the pre-trained model by adding a weighted sum of task vectors.<n>This paper theoretically prove the effectiveness of task addition in simultaneously learning a set of irrelevant or irrelevant tasks.<n>We prove the proper selection for task arithmetic to achieve negation to out-of-domain tasks.
  arXiv  Detail & Related papers  (2025-04-15T08:04:39Z)
• Provable Benefits of Task-Specific Prompts for In-context Learning [44.768199865867494]
  In this work, we consider a novel setting where the global task distribution can be partitioned into a union of conditional task distributions.<n>We then examine the use of task-specific prompts and prediction heads for learning the prior information associated with the conditional task distribution using a one-layer attention model.
  arXiv  Detail & Related papers  (2025-03-03T22:37:03Z)
• Task Vectors in In-Context Learning: Emergence, Formation, and Benefit [17.72043522825441]
  We investigate the formation of task vectors in a controlled setting using models trained from scratch on synthetic datasets.<n>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.<n>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.
  arXiv  Detail & Related papers  (2025-01-16T01:54:23Z)
• Editing Models with Task Arithmetic [69.97273155842966]
  Changing how pre-trained models behave is a common practice when developing machine learning systems. We build task vectors by subtracting the weights of a pre-trained model from the weights of the same model after fine-tuning on a task. We show that these task vectors can be modified and combined together through arithmetic operations such as negation and addition.
  arXiv  Detail & Related papers  (2022-12-08T05:50:53Z)
• Fast Inference and Transfer of Compositional Task Structures for Few-shot Task Generalization [101.72755769194677]
  We formulate it as a few-shot reinforcement learning problem where a task is characterized by a subtask graph. Our multi-task subtask graph inferencer (MTSGI) first infers the common high-level task structure in terms of the subtask graph from the training tasks. Our experiment results on 2D grid-world and complex web navigation domains show that the proposed method can learn and leverage the common underlying structure of the tasks for faster adaptation to the unseen tasks.
  arXiv  Detail & Related papers  (2022-05-25T10:44:25Z)
• OCEAN: Online Task Inference for Compositional Tasks with Context Adaptation [150.1979017130774]
  We propose a variational inference framework to perform online task inference for compositional tasks. Our framework supports flexible latent distributions based on prior knowledge of the task structure and can be trained in an unsupervised manner.
  arXiv  Detail & Related papers  (2020-08-17T04:50:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.