Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers

📄 arXiv: 2406.03068v2 📥 PDF

作者: Lei Chen, Joan Bruna, Alberto Bietti

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

发布日期: 2024-06-05 (更新: 2025-03-06)

备注: ICLR 2025

💡 一句话要点

研究Transformer中前馈层与注意力层在知识存储和推理中的作用差异

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

关键词: Transformer 注意力机制 前馈网络 上下文推理 分布关联 梯度噪声 语言模型

📋 核心要点

  1. 大型语言模型依赖Transformer架构，但前馈层和注意力层在知识存储和推理中的作用机制尚不明确。
  2. 论文通过合成数据实验和理论分析，揭示了前馈层侧重于分布关联，而注意力层侧重于上下文推理。
  3. 实验结果表明，梯度噪声是导致前馈层和注意力层学习差异的关键因素，并在预训练模型中验证了这一结论。

📝 摘要（中文）

大型语言模型在上下文推理任务（如生成连贯语言）以及存储大量知识方面取得了成功。Transformer架构的核心是前馈层和注意力层，它们通常分别与知识和推理相关联。本文在一个受控的合成环境中，对这种区别进行了实证和理论研究，其中某些下一个token的预测涉及分布信息和上下文信息。我们发现前馈层倾向于学习简单的分布关联，例如bigram，而注意力层则侧重于上下文推理。我们的理论分析将梯度中的噪声确定为这种差异的关键因素。最后，我们通过对Pythia模型系列进行消融实验，展示了在预训练模型中如何出现类似的差异。

🔬 方法详解

问题定义：现有大型语言模型依赖Transformer架构，但对于Transformer中的前馈层和注意力层，它们分别在知识存储和上下文推理中扮演的角色和作用机制尚不明确。现有研究缺乏对二者差异的深入理解，阻碍了模型优化和能力提升。

核心思路：论文的核心思路是通过构建一个受控的合成环境，使得模型在预测下一个token时需要同时利用分布信息和上下文信息。通过分析前馈层和注意力层在学习过程中的行为差异，从而揭示它们各自擅长的功能。同时，通过理论分析，探究导致这种差异的根本原因。

技术框架：论文的技术框架主要包括以下几个部分：1) 构建合成数据集，该数据集包含分布信息（例如bigram）和上下文信息；2) 使用Transformer模型进行训练，并分别分析前馈层和注意力层的学习行为；3) 通过理论分析，研究梯度噪声对前馈层和注意力层学习的影响；4) 在预训练模型（Pythia模型系列）上进行消融实验，验证在真实场景下是否存在类似的差异。

关键创新：论文的关键创新在于：1) 通过构建合成数据集，实现了对前馈层和注意力层学习行为的精确控制和分析；2) 揭示了梯度噪声是导致前馈层和注意力层学习差异的关键因素，为理解Transformer的内部机制提供了新的视角；3) 通过在预训练模型上的消融实验，验证了在真实场景下也存在类似的差异，增强了研究的实用性。

关键设计：在合成数据集的设计上，论文精心构造了既包含简单的分布关联（如bigram），又需要上下文推理才能正确预测下一个token的数据。在理论分析中，论文关注梯度噪声的统计特性，并将其与前馈层和注意力层的学习行为联系起来。在消融实验中，论文选择Pythia模型系列，并设计了简单的推理任务，以便清晰地观察前馈层和注意力层的作用。

📊 实验亮点

实验结果表明，前馈层倾向于学习简单的分布关联（如bigram），而注意力层则侧重于上下文推理。理论分析表明，梯度噪声是导致这种差异的关键因素。在Pythia模型上的消融实验也验证了在预训练模型中存在类似的差异。这些发现为理解Transformer的内部机制提供了新的视角。

🎯 应用场景

该研究成果有助于更好地理解Transformer模型的内部机制，为改进模型设计、提升模型性能提供理论指导。例如，可以针对性地优化前馈层和注意力层的训练策略，或者设计新的网络结构，以更好地利用分布信息和上下文信息。此外，该研究还可以应用于知识图谱推理、对话系统等领域。

📄 摘要（原文）

Large language models have been successful at tasks involving basic forms of in-context reasoning, such as generating coherent language, as well as storing vast amounts of knowledge. At the core of the Transformer architecture behind such models are feed-forward and attention layers, which are often associated to knowledge and reasoning, respectively. In this paper, we study this distinction empirically and theoretically in a controlled synthetic setting where certain next-token predictions involve both distributional and in-context information. We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning. Our theoretical analysis identifies the noise in the gradients as a key factor behind this discrepancy. Finally, we illustrate how similar disparities emerge in pre-trained models through ablations on the Pythia model family on simple reasoning tasks.