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

2025-05-24

AI Summary

This work addresses the lack of a unified geometric interpretation for latent state evolution in in-context learning (ICL). We propose the first cross-layer latent-space analysis framework that jointly models attention head functionality and task vector representations. Methodologically, we quantify inter-layer geometric separability and alignment of hidden states, perform attention head attribution, project task vectors, and conduct systematic ablation studies. Our results reveal that ICL performance is jointly determined by early-layer query-state separability and late-layer state alignment. Key findings include: (i) Previous Token Heads predominantly drive separability; (ii) Induction Heads synergize with task vectors to enhance alignment. Crucially, our framework establishes, for the first time, a geometric linkage between local attention mechanisms and global task representations—thereby substantially improving the interpretability and controllability of ICL dynamics.

Technology Category

• Machine Learning: Representation Learning
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Computer Vision: Interpretability, Explainability, and Transparency

Application Category

• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Linking attention heads and task vectors in ICL mechanisms
• Analyzing hidden state geometry for classification performance
• Unifying separability and alignment in layer-wise dynamics

Innovation

• Analyzes hidden state geometry in ICL
• Links attention heads and task vectors
• Reveals two-stage layer-wise mechanism