Which Attention Heads Matter for In-Context Learning?

📄 arXiv: 2502.14010v1 📥 PDF

作者: Kayo Yin, Jacob Steinhardt

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-02-19

期刊: ICML 2025

💡 一句话要点

揭示上下文学习的关键：功能向量头而非归纳头驱动LLM的ICL能力

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

关键词: 上下文学习 大型语言模型 注意力机制 归纳头 功能向量头

📋 核心要点

  1. 现有研究对LLM上下文学习机制存在争议，归纳头和功能向量头的作用尚不明确。
  2. 通过消融实验，论文发现功能向量头是驱动上下文学习的关键，尤其是在大型模型中。
  3. 研究揭示了功能向量头和归纳头之间的联系，表明归纳头可能促进了功能向量头的学习。

📝 摘要（中文）

大型语言模型（LLMs）展现出卓越的上下文学习（ICL）能力，仅需prompt中的少量演示即可执行新任务。目前有两种机制被提出解释ICL：寻找并复制相关token的归纳头，以及通过激活计算ICL任务潜在编码的功能向量（FV）头。为了更好理解哪种机制驱动ICL，本文研究并比较了12个语言模型中的归纳头和FV头。通过详细的消融实验，我们发现few-shot ICL性能主要依赖于FV头，尤其是在更大的模型中。此外，我们发现FV头和归纳头之间存在关联：许多FV头在训练期间最初是归纳头，然后过渡到FV机制。这使我们推测，归纳促进了更复杂的FV机制的学习，而FV机制最终驱动ICL。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）上下文学习（ICL）机制中的一个核心问题：究竟是归纳头（induction heads）还是功能向量（FV）头在驱动ICL能力？现有研究对这两种机制的作用存在争议，缺乏明确的证据表明哪种机制更为重要。理解ICL的驱动机制对于提升LLM的性能和可解释性至关重要。

核心思路：论文的核心思路是通过详细的消融实验，系统性地评估归纳头和FV头对ICL性能的影响。通过逐步移除或修改这些attention head，观察模型在few-shot ICL任务上的表现变化，从而确定哪些head对ICL至关重要。此外，论文还研究了FV头和归纳头之间的关系，探索它们在训练过程中的演变。

技术框架：论文的技术框架主要包括以下几个步骤：1) 选择多个不同规模的LLM作为研究对象。2) 定义ICL任务，并构建相应的few-shot数据集。3) 识别模型中的归纳头和FV头。4) 设计消融实验，包括移除特定类型的head、修改head的权重等。5) 评估消融实验后模型在ICL任务上的性能变化。6) 分析FV头和归纳头在训练过程中的演变，探索它们之间的关系。

关键创新：论文最重要的技术创新点在于揭示了FV头在驱动ICL能力中的主导作用。与以往研究不同，该论文通过严谨的实验证明，few-shot ICL性能主要依赖于FV头，尤其是在大型模型中。此外，论文还发现了FV头和归纳头之间的联系，表明归纳头可能促进了FV头的学习，为理解ICL机制提供了新的视角。

关键设计：论文的关键设计包括：1) 精心设计的消融实验，能够准确评估不同类型head对ICL性能的影响。2) 采用多种不同规模的LLM，确保结论的普适性。3) 深入分析FV头和归纳头在训练过程中的演变，揭示它们之间的关系。4) 使用标准化的ICL评估指标，便于与其他研究进行比较。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，few-shot ICL性能主要依赖于FV头，尤其是在更大的模型中。消融FV头会导致ICL性能显著下降，而消融归纳头的影响相对较小。此外，研究发现许多FV头在训练初期表现为归纳头，随后逐渐演变为FV机制，这表明归纳头可能为FV头的学习提供了基础。

🎯 应用场景

该研究成果可应用于提升大型语言模型的上下文学习能力，例如优化模型结构、设计更有效的训练策略。理解ICL的驱动机制有助于开发更高效、更可控的LLM，并促进其在自然语言处理、智能对话、知识推理等领域的应用。

📄 摘要（原文）

Large language models (LLMs) exhibit impressive in-context learning (ICL) capability, enabling them to perform new tasks using only a few demonstrations in the prompt. Two different mechanisms have been proposed to explain ICL: induction heads that find and copy relevant tokens, and function vector (FV) heads whose activations compute a latent encoding of the ICL task. To better understand which of the two distinct mechanisms drives ICL, we study and compare induction heads and FV heads in 12 language models. Through detailed ablations, we discover that few-shot ICL performance depends primarily on FV heads, especially in larger models. In addition, we uncover that FV and induction heads are connected: many FV heads start as induction heads during training before transitioning to the FV mechanism. This leads us to speculate that induction facilitates learning the more complex FV mechanism that ultimately drives ICL.