Research map / Architectures and efficiency / Attention and transformers

Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Trains small transformers from scratch on latent-task sequence mixtures (dice, linear regression, Markov chains) to show task vectors implement Bayesian retrieval in-distribution while OOD extrapolative learning uses a near-orthogonal subspace.

Problem

How task-vector geometry in hidden states is shaped by the training distribution, and what geometry supports in-distribution retrieval versus out-of-distribution generalization in transformers.

Approach

Formalizes task vectors as mappings of latents with four properties (P0-P3) linking hidden states to Bayesian posterior mixtures, and shows that increasing task diversity moves models into an extrapolative mode whose representations lie in a subspace nearly orthogonal to the task-vector subspace.

Key results

• Task vectors explain hidden states with residual variance ratio near 0 at the last position (e.g., 0.004-0.075 across layers) across dice, linear regression, and Markov experiments.
• Steering the task-vector simplex reduces KL from 0.11 to 0.03 in dice and from 0.23 to 0.06 in Markov chains; RMSE drops from 2.8 to 1.1 in linear regression.
• Intervening on a low-rank orthogonal subspace raises OOD loss by 95.8% (E1) and 99.6% (E3) of in-context gain while major-task loss changes only about 2-4%.

Surprising findings

• Hidden states encode the full Dyck prefix in non-Markovian planted Dyck, rather than a simple running count summary, so the task-vector heuristic assumption can fail.
• Projection R^2 of OOD hidden states on the task subspace drops toward zero only at large task diversity, indicating the orthogonal geometry is learned rather than present from the start.

Limitations

• Analysis relies on the first-order Markov property and small synthetic setups.
• The near-orthogonal subspace does not give a mechanistic account of OOD generalization.
• Theory is theory-inspired and not a formal theory of generalization or learning dynamics; the LLM experiment is only qualitative.

Experimental setup

Datasets: synthetic biased dice sequences synthetic in-context linear regression synthetic Markov chain sequences planted Dyck sequences word-function tasks for Qwen2.5-7B probing

Benchmarks: synthetic E1 dice synthetic E2 linear regression synthetic E3 Markov chains

Metrics: R^2 KL divergence RMSE residual variance ratio interaction eta^2 fractional loss degradation

Baselines: unsteered model random same-rank orthogonal direction

Builds on: linear representation hypothesis task vectors superposition principle in-context linear regression

Closest papers by meaning

• Platonic Task Arithmetic (NeurIPS 2026 · similarity 0.88)
• Variational Task Vector Composition (NeurIPS 2025 · similarity 0.88)
• Does learning the right latent variables necessarily improve in-context learning? (ICML 2025 · similarity 0.88)
• Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.87)
• Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning (NeurIPS 2026 · similarity 0.87)
• Disentangling Representations through Multi-task Learning (ICLR 2025 · similarity 0.87)