Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B

📄 arXiv: 2504.00132v4 📥 PDF

作者: Aleksandra Bakalova, Yana Veitsman, Xinting Huang, Michael Hahn

分类: cs.CL, cs.LG

发布日期: 2025-03-31 (更新: 2025-09-17)

💡 一句话要点

通过因果干预揭示Gemma-2 2B模型上下文学习的“语境化-聚合”机制

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

关键词: 上下文学习 因果干预 大型语言模型 Gemma-2 2B 语境化 聚合 信息流 自然语言处理

📋 核心要点

  1. 大型语言模型的上下文学习能力强大，但其内部机制尚不明确，尤其是在如何整合少量样本信息方面。
  2. 论文提出“语境化-聚合”两步策略，模型先构建单个样本表示并语境化，再聚合表示以识别任务。
  3. 通过对Gemma-2 2B模型的因果干预分析，揭示了ICL中信息流动的具体过程和关键步骤。

📝 摘要（中文）

本文旨在揭示大型语言模型(LLMs)上下文学习(ICL)的内在机制。尽管已有大量研究关注ICL的行为特性以及其在小型模型中的涌现，但对于模型如何从少量样本提示中整合任务信息的机制仍不明确。本文采用因果干预方法，在Gemma-2 2B模型上针对五个自然语言ICL任务，识别信息流。研究发现，模型采用一种名为“语境化-聚合”的两步策略来推断任务信息：在较低层，模型构建单个少量样本的表示，并通过序列中少量样本输入和输出token之间的连接，利用前面的例子对这些表示进行语境化；在较高层，这些表示被聚合以识别任务并准备预测下一个输出。语境化步骤的重要性因任务而异，并且在存在模糊示例时可能变得更加重要。总而言之，通过提供严格的因果分析，本文的研究结果阐明了语言模型中ICL发生的机制。

🔬 方法详解

问题定义：现有研究对大型语言模型上下文学习(ICL)的机制理解不足，尤其是在模型如何从少量样本提示中提取和整合任务信息方面。现有的方法难以解释模型内部的信息流动和处理方式，缺乏对ICL过程的细粒度因果分析。

核心思路：论文的核心思路是，ICL过程可以分解为两个阶段：“语境化”和“聚合”。模型首先对每个少量样本进行编码，并利用序列中其他样本的信息进行语境化，然后将这些语境化的表示进行聚合，以识别任务并生成预测。这种设计基于假设：ICL并非简单地将所有样本视为独立个体，而是考虑了它们之间的相互关系。

技术框架：研究采用因果干预的方法，分析Gemma-2 2B模型在执行五个自然语言ICL任务时的信息流动。具体步骤包括：1) 选择合适的ICL任务；2) 构建包含少量样本的提示；3) 对模型中间层的神经元进行干预，例如激活或抑制特定神经元；4) 观察干预对模型输出的影响，从而推断信息流动的路径和关键节点。通过分析不同层的信息流动，区分“语境化”和“聚合”两个阶段。

关键创新：论文的关键创新在于提出了“语境化-聚合”的ICL机制，并使用因果干预方法对其进行了验证。与以往的研究相比，该方法更注重对模型内部信息流动的因果分析，而非仅仅关注输入输出之间的相关性。这种方法能够更深入地理解ICL的内在机制，并为改进ICL性能提供新的思路。

关键设计：论文的关键设计包括：1) 选择具有代表性的ICL任务，例如情感分类、主题识别等；2) 精心设计因果干预方案，选择合适的干预对象和干预方式，以最大程度地揭示信息流动；3) 使用Gemma-2 2B模型作为研究对象，因为它具有一定的规模和性能，同时又相对易于分析；4) 通过比较不同任务和不同样本的语境化程度，分析语境化步骤的重要性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

研究发现，Gemma-2 2B模型在执行ICL任务时，确实采用了“语境化-聚合”的策略。在较低层，模型构建单个少量样本的表示，并通过序列中少量样本输入和输出token之间的连接，利用前面的例子对这些表示进行语境化。在较高层，这些表示被聚合以识别任务并准备预测下一个输出。语境化步骤的重要性因任务而异，并且在存在模糊示例时可能变得更加重要。这些发现为理解ICL的内在机制提供了重要的证据。

🎯 应用场景

该研究成果可应用于提升大型语言模型的上下文学习能力，尤其是在样本数量有限或存在歧义的情况下。通过理解“语境化-聚合”机制，可以设计更有效的提示工程方法，或改进模型架构以更好地利用上下文信息。此外，该研究方法也可推广到其他类型的模型和任务中，为理解和改进人工智能系统的内部机制提供借鉴。

📄 摘要（原文）

In-Context Learning (ICL) is an intriguing ability of large language models (LLMs). Despite a substantial amount of work on its behavioral aspects and how it emerges in miniature setups, it remains unclear which mechanism assembles task information from the individual examples in a fewshot prompt. We use causal interventions to identify information flow in Gemma-2 2B for five naturalistic ICL tasks. We find that the model infers task information using a two-step strategy we call contextualize-then-aggregate: In the lower layers, the model builds up representations of individual fewshot examples, which are contextualized by preceding examples through connections between fewshot input and output tokens across the sequence. In the higher layers, these representations are aggregated to identify the task and prepare prediction of the next output. The importance of the contextualization step differs between tasks, and it may become more important in the presence of ambiguous examples. Overall, by providing rigorous causal analysis, our results shed light on the mechanisms through which ICL happens in language models.