# Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning

arXiv: https://arxiv.org/abs/2505.18752 · HTML: https://mlsanity.com/paper/2505.18752/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 28 references that are in the map.
- First posted on arXiv 24 May 2025.

## Builds on (most cited first, 17 of 17)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets](https://mlsanity.com/paper/2310.06824/) (2023)
- [In-Context Learning Creates Task Vectors](https://mlsanity.com/paper/2310.15916/) (2023)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Function Vectors in Large Language Models](https://mlsanity.com/paper/2310.15213/) (2023)
- [A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity](https://mlsanity.com/paper/2401.01967/) (2024)
- [The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction](https://mlsanity.com/paper/2312.13558/) (2023)
- [The mechanistic basis of data dependence and abrupt learning in an in-context classification task](https://mlsanity.com/paper/2312.03002/) (2023)
- [What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation](https://mlsanity.com/paper/2404.07129/) (2024)
- [Overthinking the Truth: Understanding how Language Models Process False Demonstrations](https://mlsanity.com/paper/2307.09476/) (2023)
- [Out-of-distribution generalization via composition: a lens through induction heads in Transformers](https://mlsanity.com/paper/2408.09503/) (2024)
- [The Geometry of Prompting: Unveiling Distinct Mechanisms of Task Adaptation in Language Models](https://mlsanity.com/paper/2502.08009/) (2025)
- [Revisiting In-context Learning Inference Circuit in Large Language Models](https://mlsanity.com/paper/2410.04468/) (2024)
- [Token-based Decision Criteria Are Suboptimal in In-context Learning](https://mlsanity.com/paper/2406.16535/) (2024)

## Built on by (most cited first, 10 of 10)

- [CoT Vectors: Transferring and Probing the Reasoning Mechanisms of LLMs](https://mlsanity.com/paper/2510.00579/) (2025)
- [Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insight](https://mlsanity.com/paper/2509.24169/) (2025)
- [Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis](https://mlsanity.com/paper/2509.24164/) (2025)
- [HIFICL: High-Fidelity In-Context Learning for Multimodal Tasks](https://mlsanity.com/paper/2603.12760/) (2026)
- [The Geometric Price of Discrete Logic: Context-driven Manifold Dynamics of Number Representations](https://mlsanity.com/paper/2603.23577/) (2026)
- [Context-Gated Associative Retrieval: From Theory to Transformers](https://mlsanity.com/paper/2605.10970/) (2026)
- [Few-Shot Demonstrations Elicit the Use of In-Context World Representations in LLMs](https://mlsanity.com/paper/2609.24352/) (2026)
- [Can Linguistic Reasoning Vectors Enhance Multimodal Reasoning Ability?](https://mlsanity.com/paper/2609.31140/) (2026)
- [Are In-Context Images Worth 10 Dimensions?](https://mlsanity.com/paper/2609.37659/) (2026)
- [Rules Amortize, Pairings Don't: Linguistic Structure Determines What Latent Task Representations Can Replace In-Context Learning](https://mlsanity.com/paper/2610.00526/) (2026)

## Abstract (arXiv)

The unusual properties of in-context learning (ICL) have prompted investigations into the internal mechanisms of large language models. Prior work typically focuses on either special attention heads or task vectors at specific layers, but lacks a unified framework linking these components to the evolution of hidden states across layers that ultimately produce the model's output. In this paper, we propose such a framework for ICL in classification tasks by analyzing two geometric factors that govern performance: the separability and alignment of query hidden states. A fine-grained analysis of layer-wise dynamics reveals a striking two-stage mechanism: separability emerges in early layers, while alignment develops in later layers. Ablation studies further show that Previous Token Heads drive separability, while Induction Heads and task vectors enhance alignment. Our findings thus bridge the gap between attention heads and task vectors, offering a unified account of ICL's underlying mechanisms.