Out-of-distribution generalization via composition: a lens through induction heads in Transformers

📄 arXiv: 2408.09503v2 📥 PDF

作者: Jiajun Song, Zhuoyan Xu, Yiqiao Zhong

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

发布日期: 2024-08-18 (更新: 2024-12-28)

备注: 46 pages, 27 figures

💡 一句话要点

通过Transformer归纳头组合实现分布外泛化能力

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

关键词: 分布外泛化 Transformer 归纳头 自注意力机制 组合性 上下文学习 语言模型 表征学习

📋 核心要点

  1. 大型语言模型在分布外泛化方面面临挑战，现有方法难以解释其内在机制。
  2. 该研究提出通过Transformer中的归纳头组合来实现分布外泛化，关注模型如何学习和组合规则。
  3. 实验表明，模型可以通过组合自注意力层来学习规则，并发现共享潜在子空间作为组合的桥梁。

📝 摘要（中文）

大型语言模型（LLMs），如GPT-4，有时表现出创造力，能够解决新颖的任务，通常只需在提示中提供少量示例。这些任务要求模型在与训练数据不同的分布上进行泛化，这被称为分布外（OOD）泛化。尽管LLMs取得了巨大的成功，但它们如何实现OOD泛化仍然是一个未被充分探索的问题。本文研究了在实例根据隐藏规则生成的情况下OOD泛化，包括使用符号推理的上下文学习。模型需要在没有任何微调的情况下推断输入提示背后的隐藏规则。通过对合成示例的Transformer训练动态进行实证研究，并对各种预训练LLMs进行了广泛的实验，重点关注一种称为归纳头的组件。研究发现OOD泛化和组合是紧密相关的——模型可以通过组合两个自注意力层来学习规则，从而实现OOD泛化。此外，嵌入（或特征）空间中的共享潜在子空间通过对齐早期层和后期层，充当组合的桥梁，这被称为共同桥梁表示假设。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）在分布外（OOD）泛化方面的难题。现有方法缺乏对LLMs如何实现OOD泛化的深入理解，特别是当实例由隐藏规则生成时，模型如何在没有微调的情况下推断这些规则。现有方法难以解释LLM的内在机制，以及如何利用模型内部的组件来实现OOD泛化。

核心思路：论文的核心思路是，OOD泛化能力与Transformer模型中自注意力机制的组合能力密切相关。具体来说，模型可以通过组合两个自注意力层（即归纳头）来学习隐藏规则，从而实现OOD泛化。此外，论文提出了“共同桥梁表示假设”，认为嵌入空间中的共享潜在子空间充当了早期层和后期层之间组合的桥梁。

技术框架：论文的研究框架主要包括以下几个部分：1) 对合成数据集上的Transformer训练动态进行分析，以理解归纳头的行为；2) 在各种预训练LLMs上进行实验，验证归纳头在OOD泛化中的作用；3) 探索嵌入空间中的共享潜在子空间，验证“共同桥梁表示假设”。整体流程是先通过合成数据分析，然后在大规模预训练模型上验证，最后提出假设并进行验证。

关键创新：论文的关键创新在于：1) 揭示了Transformer模型中归纳头在OOD泛化中的作用，表明模型可以通过组合自注意力层来学习规则；2) 提出了“共同桥梁表示假设”，认为嵌入空间中的共享潜在子空间是实现组合的关键；3) 将OOD泛化与组合联系起来，为理解LLMs的泛化能力提供了一个新的视角。

关键设计：论文的关键设计包括：1) 使用合成数据集来控制隐藏规则的复杂性，便于分析模型的学习过程；2) 重点关注Transformer模型中的归纳头，通过分析其权重和激活来理解其行为；3) 使用线性探针等技术来探索嵌入空间中的共享潜在子空间；4) 实验中使用了多种预训练LLMs，以验证结果的泛化性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

该研究通过实验发现，Transformer模型可以通过组合两个自注意力层（归纳头）来学习规则，从而实现分布外泛化。此外，实验验证了“共同桥梁表示假设”，表明嵌入空间中的共享潜在子空间充当了早期层和后期层之间组合的桥梁。具体性能数据未知，但研究强调了归纳头和共享潜在子空间在OOD泛化中的重要作用。

🎯 应用场景

该研究成果可应用于提升大型语言模型在各种实际场景中的泛化能力，例如在零样本或少样本学习环境中解决新颖任务，或在面对数据分布偏移时保持性能稳定。此外，该研究有助于更好地理解LLMs的内部机制，为模型设计和优化提供指导。

📄 摘要（原文）

Large language models (LLMs) such as GPT-4 sometimes appear to be creative, solving novel tasks often with a few demonstrations in the prompt. These tasks require the models to generalize on distributions different from those from training data -- which is known as out-of-distribution (OOD) generalization. Despite the tremendous success of LLMs, how they approach OOD generalization remains an open and underexplored question. We examine OOD generalization in settings where instances are generated according to hidden rules, including in-context learning with symbolic reasoning. Models are required to infer the hidden rules behind input prompts without any fine-tuning. We empirically examined the training dynamics of Transformers on a synthetic example and conducted extensive experiments on a variety of pretrained LLMs, focusing on a type of components known as induction heads. We found that OOD generalization and composition are tied together -- models can learn rules by composing two self-attention layers, thereby achieving OOD generalization. Furthermore, a shared latent subspace in the embedding (or feature) space acts as a bridge for composition by aligning early layers and later layers, which we refer to as the common bridge representation hypothesis.