Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers

2026-05-05

AI Summary

This work addresses the lack of theoretical understanding regarding how the geometry of task vectors in Transformers is shaped by training distributions and underpins both in-distribution and out-of-distribution reasoning. By training small-scale Transformers from scratch in a controlled synthetic setting, the study integrates task vector analysis with representational geometric modeling to uncover, for the first time, the intrinsic relationship among task vector geometry, training distribution, and generalization behavior. The findings reveal that in-distribution tasks are solved through Bayesian-like retrieval via convex combinations of task vectors, whereas out-of-distribution tasks rely on extrapolation within approximately orthogonal subspaces. Crucially, both mechanisms coexist within the same architecture, establishing a formal theoretical link between the geometric structure of task vectors and the model’s generalization capabilities.

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.

Problem

• task vector geometry
• out-of-distribution generalization
• task inference
• training distribution
• transformers

Innovation

• task vector geometry
• Bayesian task retrieval
• out-of-distribution generalization
• transformer interpretability
• latent-task inference