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

• URL: http://arxiv.org/abs/2505.18752v1
• Date: Sat, 24 May 2025 15:42:20 GMT
• Title: Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning
• Abstract summary: In this paper, we analyze two geometric factors that govern performance: the separability and alignment of query hidden states.<n>Previous Token Heads drive separability, while Induction Heads and task vectors enhance alignment.<n>Our findings thus bridge the gap between attention heads and task vectors, offering a unified account of ICL's underlying mechanisms.
• Score: 2.4866936275046405
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Unleashing Multimodal Large Language Models for Training-free HOI Detection in the Wild [49.77540427384148]
  We present AgentHOI, a training-free, agentic framework that transfers the generalist multimodal reasoning capabilities of foundation models to HOI detection in the wild.<n>To address the challenges of incomplete interaction discovery and ambiguous localization in complex scenes, we introduce two key mechanisms.<n>AgentHOI achieves superior performance over state-of-the-art supervised and weakly supervised methods in real-world settings.
  arXiv  Detail & Related papers  (2026-07-15T14:30:20Z)
• KeepLoRA++: Continual Learning with Layer-Scaled Residual Gradient Adaptation [67.23448686598414]
  Continual learning for pre-trained vision-language models requires balancing three competing objectives.<n>This paper presents KeepLoRA++, balancing these objectives through a unified dual-dimensional knowledge retention mechanism.
  arXiv  Detail & Related papers  (2026-06-15T05:57:59Z)
• Mechanistic Insights into Functional Sparsity in Multimodal LLMs via CoRe Heads [48.341678148905466]
  We present an in-depth interpretability study that uncovers a profound structural property within Multimodal Large Language Models (MLLMs)<n>We identify and characterize a highly specialized subset of attention heads, referred to as Context-aware Retrieval (CoRe) heads.<n>Our findings reveal a structural principle of functional sparsity within MLLMs, refining the current understanding of mechanistic interpretability.
  arXiv  Detail & Related papers  (2026-06-04T08:18:31Z)
• Mechanisms of Object Localization in Vision-Language Models [10.367669666212473]
  We study the processes that support object localization in visually-grounded language models (VLMs)<n>We find that localization is driven by a containerization mechanism in which object-aligned tokens define the spatial extent of the object.<n>Only a very small set of attention heads mediates the causal effect for both classification and localization.
  arXiv  Detail & Related papers  (2026-05-19T12:56:32Z)
• Emergent Structured Representations Support Flexible In-Context Inference in Large Language Models [77.98801218316505]
  Large language models (LLMs) exhibit emergent behaviors suggestive of human-like reasoning.<n>We investigate the internal processing of LLMs during in-context concept inference.
  arXiv  Detail & Related papers  (2026-02-08T03:14:39Z)
• Cognitive Mirrors: Exploring the Diverse Functional Roles of Attention Heads in LLM Reasoning [54.12174882424842]
  Large language models (LLMs) have achieved state-of-the-art performance in a variety of tasks, but remain largely opaque in terms of their internal mechanisms.<n>We propose a novel interpretability framework to systematically analyze the roles and behaviors of attention heads.
  arXiv  Detail & Related papers  (2025-12-03T10:24:34Z)
• Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis [7.677300807530906]
  We show that TR heads promote task recognition by aligning hidden states with the task subspace, while TL heads rotate hidden states within the subspace toward the correct label to facilitate prediction.<n>Our framework thus provides a unified and interpretable account of how large language models execute ICL across diverse tasks and settings.
  arXiv  Detail & Related papers  (2025-09-29T01:28:37Z)
• Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations [19.539276425108987]
  This work proposes the Linear Combination Conjecture, positing that task vectors act as single in-context demonstrations formed through linear combinations of the original ones.<n>We show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis.<n>We predict the failure of task vectors on representing high-rank mappings and confirm this on practical LLMs.
  arXiv  Detail & Related papers  (2025-06-10T17:59:31Z)
• From Compression to Expansion: A Layerwise Analysis of In-Context Learning [20.64102133977965]
  In-context learning (ICL) enables large language models to adapt to new tasks without weight updates by learning from demonstration sequences.<n>We conduct a statistical geometric analysis of ICL representations to investigate how task-specific information is captured across layers.<n>Our findings reveal an intriguing layerwise dynamic in ICL, highlight how structured representations emerge within LLMs, and showcase that analyzing internal representations can facilitate a deeper understanding of model behavior.
  arXiv  Detail & Related papers  (2025-05-22T22:22:03Z)
• A Survey of Frontiers in LLM Reasoning: Inference Scaling, Learning to Reason, and Agentic Systems [93.8285345915925]
  Reasoning is a fundamental cognitive process that enables logical inference, problem-solving, and decision-making.<n>With the rapid advancement of large language models (LLMs), reasoning has emerged as a key capability that distinguishes advanced AI systems.<n>We categorize existing methods along two dimensions: (1) Regimes, which define the stage at which reasoning is achieved; and (2) Architectures, which determine the components involved in the reasoning process.
  arXiv  Detail & Related papers  (2025-04-12T01:27:49Z)
• How do Large Language Models Understand Relevance? A Mechanistic Interpretability Perspective [64.00022624183781]
  Large language models (LLMs) can assess relevance and support information retrieval (IR) tasks.<n>We investigate how different LLM modules contribute to relevance judgment through the lens of mechanistic interpretability.
  arXiv  Detail & Related papers  (2025-04-10T16:14:55Z)
• Attention Heads of Large Language Models: A Survey [10.136767972375639]
  We aim to demystify the internal reasoning processes of Large Language Models (LLMs) by systematically exploring the roles and mechanisms of attention heads.<n>We first introduce a novel four-stage framework inspired by the human thought process: Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation.<n>We analyze the experimental methodologies used to discover these special heads, dividing them into two categories: Modeling-Free and Modeling-Required methods.
  arXiv  Detail & Related papers  (2024-09-05T17:59:12Z)
• Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers [49.80959223722325]
  We study the distinction between feed-forward and attention layers in large language models.<n>We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning.
  arXiv  Detail & Related papers  (2024-06-05T08:51:08Z)
• On Understanding Attention-Based In-Context Learning for Categorical Data [49.40350941996942]
  We develop a network composed of attention blocks, with each block employing a self-attention layer followed by a cross-attention layer, with associated skip connections.<n>This model can exactly perform multi-step functional GD inference for in-context inference with categorical observations.
  arXiv  Detail & Related papers  (2024-05-27T15:03:21Z)
• Dual Contrastive Learning for General Face Forgery Detection [64.41970626226221]
  We propose a novel face forgery detection framework, named Dual Contrastive Learning (DCL), which constructs positive and negative paired data. To explore the essential discrepancies, Intra-Instance Contrastive Learning (Intra-ICL) is introduced to focus on the local content inconsistencies prevalent in the forged faces.
  arXiv  Detail & Related papers  (2021-12-27T05:44:40Z)
• Structure-Aware Feature Generation for Zero-Shot Learning [108.76968151682621]
  We introduce a novel structure-aware feature generation scheme, termed as SA-GAN, to account for the topological structure in learning both the latent space and the generative networks. Our method significantly enhances the generalization capability on unseen-classes and consequently improve the classification performance.
  arXiv  Detail & Related papers  (2021-08-16T11:52:08Z)
• DisenE: Disentangling Knowledge Graph Embeddings [33.169388832519]
  DisenE is an end-to-end framework to learn disentangled knowledge graph embeddings. We introduce an attention-based mechanism that enables the model to explicitly focus on relevant components of entity embeddings according to a given relation.
  arXiv  Detail & Related papers  (2020-10-28T03:45:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.