Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers
For interpretability and mechanistic understanding of transformers, this work provides a rigorous mathematical characterization of how task vectors enable dual inference modes, though in a controlled synthetic setting.
The paper studies how transformers infer latent tasks from context, showing that in-distribution behavior relies on convex combinations of learned task vectors (Bayesian retrieval), while out-of-distribution generalization arises from extrapolative learning in a nearly orthogonal subspace. These results connect task-vector geometry, training distributions, and generalization.
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.