Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis

• URL: http://arxiv.org/abs/2509.24164v1
• Date: Mon, 29 Sep 2025 01:28:37 GMT
• Title: Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis
• Abstract summary: 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.
• Score: 7.677300807530906
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We investigate the mechanistic underpinnings of in-context learning (ICL) in large language models by reconciling two dominant perspectives: the component-level analysis of attention heads and the holistic decomposition of ICL into Task Recognition (TR) and Task Learning (TL). We propose a novel framework based on Task Subspace Logit Attribution (TSLA) to identify attention heads specialized in TR and TL, and demonstrate their distinct yet complementary roles. Through correlation analysis, ablation studies, and input perturbations, we show that the identified TR and TL heads independently and effectively capture the TR and TL components of ICL. Using steering experiments with geometric analysis of hidden states, we reveal 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. We further show how previous findings on ICL mechanisms, including induction heads and task vectors, can be reconciled with our attention-head-level analysis of the TR-TL decomposition. Our framework thus provides a unified and interpretable account of how large language models execute ICL across diverse tasks and settings.

Related papers

• Disentangling meaning from language in LLM-based machine translation [24.40574806667368]
  We study sentence-level Machine Translation from a mechanistic perspective.<n>We decompose MT into two subtasks: producing text in the target language and preserving the input sentence's meaning.<n>We show that modifying just 1% of the relevant heads enables instruction-free MT performance comparable to instruction-based prompting.
  arXiv  Detail & Related papers  (2026-02-04T14:40:53Z)
• Connecting the Dots: Training-Free Visual Grounding via Agentic Reasoning [63.109585527799005]
  GroundingAgent is a visual grounding framework that operates without task-specific fine-tuning.<n>It achieves an average zero-shot grounding accuracy of 65.1 % on widely-used benchmarks.<n>It also offers strong interpretability, transparently illustrating each reasoning step.
  arXiv  Detail & Related papers  (2025-11-24T03:11:08Z)
• Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning [2.4866936275046405]
  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.
  arXiv  Detail & Related papers  (2025-05-24T15:42:20Z)
• Illusion or Algorithm? Investigating Memorization, Emergence, and Symbolic Processing in In-Context Learning [50.53703102032562]
  Large-scale Transformer language models (LMs) trained solely on next-token prediction with web-scale data can solve a wide range of tasks.<n>The mechanism behind this capability, known as in-context learning (ICL), remains both controversial and poorly understood.
  arXiv  Detail & Related papers  (2025-05-16T08:50:42Z)
• 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)
• Unveiling In-Context Learning: A Coordinate System to Understand Its Working Mechanism [28.751003584429615]
  Large language models (LLMs) exhibit remarkable in-context learning capabilities. Recent research presents two conflicting views on ICL. We provide a Two-Dimensional Coordinate System that unifies both views into a systematic framework.
  arXiv  Detail & Related papers  (2024-07-24T05:26:52Z)
• Latent Causal Probing: A Formal Perspective on Probing with Causal Models of Data [3.376269351435396]
  We develop a formal perspective on probing using structural causal models (SCM) We extend a recent study of LMs in the context of a synthetic grid-world navigation task. Our techniques provide robust empirical evidence for the ability of LMs to induce the latent concepts underlying text.
  arXiv  Detail & Related papers  (2024-07-18T17:59:27Z)
• A Theoretical Analysis of Self-Supervised Learning for Vision Transformers [66.08606211686339]
  Masked autoencoders (MAE) and contrastive learning (CL) capture different types of representations.<n>We study the training dynamics of one-layer softmax-based vision transformers (ViTs) on both MAE and CL objectives.
  arXiv  Detail & Related papers  (2024-03-04T17:24:03Z)
• Identifying Semantic Induction Heads to Understand In-Context Learning [103.00463655766066]
  We investigate whether attention heads encode two types of relationships between tokens present in natural languages. We find that certain attention heads exhibit a pattern where, when attending to head tokens, they recall tail tokens and increase the output logits of those tail tokens.
  arXiv  Detail & Related papers  (2024-02-20T14:43:39Z)
• What In-Context Learning "Learns" In-Context: Disentangling Task Recognition and Task Learning [24.395288160951118]
  Large language models (LLMs) exploit in-context learning (ICL) to solve tasks with only a few demonstrations. We characterize two ways through which ICL leverages demonstrations. We show that models can achieve non-trivial performance with only TR, and TR does not further improve with larger models or more demonstrations.
  arXiv  Detail & Related papers  (2023-05-16T18:05:19Z)
• Contrastive Instruction-Trajectory Learning for Vision-Language Navigation [66.16980504844233]
  A vision-language navigation (VLN) task requires an agent to reach a target with the guidance of natural language instruction. Previous works fail to discriminate the similarities and discrepancies across instruction-trajectory pairs and ignore the temporal continuity of sub-instructions. We propose a Contrastive Instruction-Trajectory Learning framework that explores invariance across similar data samples and variance across different ones to learn distinctive representations for robust navigation.
  arXiv  Detail & Related papers  (2021-12-08T06:32:52Z)
• Multilingual Multi-Aspect Explainability Analyses on Machine Reading Comprehension Models [76.48370548802464]
  This paper focuses on conducting a series of analytical experiments to examine the relations between the multi-head self-attention and the final MRC system performance. We discover that passage-to-question and passage understanding attentions are the most important ones in the question answering process. Through comprehensive visualizations and case studies, we also observe several general findings on the attention maps, which can be helpful to understand how these models solve the questions.
  arXiv  Detail & Related papers  (2021-08-26T04:23:57Z)

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.