Search papers, labs, and topics across Lattice.
This paper investigates how transformers infer latent tasks from context, distinguishing between recognizing seen tasks (in-distribution) and adapting to novel ones (out-of-distribution). Using a controlled synthetic setting, the authors train small transformers and mathematically characterize the relationship between task-vector geometry and generalization. They find that in-distribution inference relies on Bayesian task retrieval via convex combinations of learned task vectors, while out-of-distribution inference leverages extrapolative task learning with representations orthogonal to the task-vector subspace.
Transformers generalize out-of-distribution not by clever interpolation, but by learning a separate, orthogonal representation subspace for unseen tasks.
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.
