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

📄 arXiv: 2509.24164v1 📥 PDF

作者: Haolin Yang, Hakaze Cho, Naoya Inoue

分类: cs.CL

发布日期: 2025-09-29

备注: 45 pages, 88 figures, 10 tables

💡 一句话要点

通过注意力头分析，在上下文学习中定位任务识别和任务学习

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文学习 注意力头分析 任务识别 任务学习 Logit归因 大型语言模型 可解释性

📋 核心要点

  1. 现有方法难以从机制上理解大型语言模型上下文学习(ICL)中任务识别(TR)和任务学习(TL)的具体作用。
  2. 提出基于任务子空间Logit归因(TSLA)的框架，识别专门执行TR和TL的注意力头，并分析其作用。
  3. 实验证明TR和TL头独立有效地捕获ICL的相应组件，并通过几何分析揭示其促进任务识别和预测的机制。

📝 摘要（中文）

本文通过调和两种主流视角，即注意力头的组件级分析和将ICL整体分解为任务识别（TR）和任务学习（TL），研究了大型语言模型中上下文学习（ICL）的机制基础。我们提出了一个基于任务子空间Logit归因（TSLA）的新框架，以识别专门用于TR和TL的注意力头，并展示了它们独特而互补的作用。通过相关性分析、消融研究和输入扰动，我们表明，所识别的TR和TL头独立且有效地捕获了ICL的TR和TL组件。通过对隐藏状态进行几何分析的引导实验，我们揭示了TR头通过将隐藏状态与任务子空间对齐来促进任务识别，而TL头在子空间内旋转隐藏状态以朝着正确的标签，从而促进预测。我们进一步展示了先前关于ICL机制的研究结果，包括归纳头和任务向量，如何与我们对TR-TL分解的注意力头级别分析相协调。因此，我们的框架为大型语言模型如何在各种任务和设置中执行ICL提供了一个统一且可解释的说明。

🔬 方法详解

问题定义：大型语言模型中的上下文学习(ICL)能力令人印象深刻，但其内部机制仍然是一个黑盒。现有方法要么侧重于注意力头的组件级分析，要么将ICL分解为任务识别(TR)和任务学习(TL)两个阶段进行整体分析，缺乏一个统一的框架来理解TR和TL在注意力头层面的具体作用。因此，如何定位并理解ICL过程中TR和TL的具体机制是本文要解决的核心问题。

核心思路：本文的核心思路是通过分析注意力头在执行TR和TL时的行为，来揭示ICL的内部机制。具体来说，通过提出任务子空间Logit归因(TSLA)方法，识别出专门负责TR和TL的注意力头，并分析它们如何协同工作以实现ICL。这种方法将组件级分析和整体分解相结合，提供了一个更全面和可解释的视角。

技术框架：该研究的技术框架主要包含以下几个阶段：1) 任务子空间Logit归因(TSLA)：提出TSLA方法，用于识别专门负责TR和TL的注意力头。2) 相关性分析、消融研究和输入扰动：通过这些实验验证所识别的TR和TL头是否独立且有效地捕获了ICL的相应组件。3) 引导实验和几何分析：通过对隐藏状态进行几何分析，揭示TR头如何通过将隐藏状态与任务子空间对齐来促进任务识别，以及TL头如何在子空间内旋转隐藏状态以朝着正确的标签，从而促进预测。4) 与现有ICL机制研究结果的协调：将本文的研究结果与先前关于ICL机制的研究结果（如归纳头和任务向量）进行协调，以提供一个更统一的解释。

关键创新：本文最重要的技术创新点在于提出了任务子空间Logit归因(TSLA)方法，该方法能够有效地识别出专门负责TR和TL的注意力头。与现有方法相比，TSLA方法能够更精细地分析ICL的内部机制，并揭示TR和TL在注意力头层面的具体作用。此外，本文还通过几何分析揭示了TR和TL头促进任务识别和预测的具体机制，为理解ICL提供了新的视角。

关键设计：TSLA方法的关键在于如何定义任务子空间以及如何计算Logit归因。具体来说，任务子空间是通过对不同任务的隐藏状态进行主成分分析(PCA)得到的。Logit归因则是通过计算每个注意力头对最终预测Logit的贡献来衡量的。此外，在引导实验中，通过调整隐藏状态的方向和大小，来模拟TR和TL头的行为，并观察对预测结果的影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，通过TSLA方法识别的TR和TL头能够独立且有效地捕获ICL的相应组件。消融实验显示，移除TR头会导致任务识别能力显著下降，而移除TL头则会影响预测的准确性。几何分析表明，TR头能够将隐藏状态与任务子空间对齐，而TL头则能够在子空间内旋转隐藏状态以朝着正确的标签。这些结果为理解ICL的内部机制提供了有力的证据。

🎯 应用场景

该研究成果可应用于提升大型语言模型的可解释性和可控性。通过理解ICL的内部机制，可以更好地设计提示工程，提高模型在特定任务上的性能。此外，该研究也有助于开发更高效的ICL方法，例如通过选择性地激活TR和TL头来减少计算成本。未来，该研究可以扩展到其他类型的模型和任务，进一步提升人工智能系统的可靠性和安全性。

📄 摘要（原文）

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.