How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations

📄 arXiv: 2310.10616v1 📥 PDF

作者: Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, Yu Bai

分类: cs.LG

发布日期: 2023-10-16

💡 一句话要点

提出基于表示学习的复杂场景下变换器的上下文学习机制

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

关键词: 上下文学习 变换器 表示学习 复杂场景 机器学习 自然语言处理 深度学习

📋 核心要点

  1. 现有理论主要集中在简单函数类的学习，缺乏对复杂场景下上下文学习的深入理解。
  2. 本文通过构建合成的上下文学习问题，探讨变换器如何在复杂表示下进行学习。
  3. 实验结果表明，训练后的变换器在复杂场景中实现了接近最优的ICL性能，揭示了其内部机制。

📝 摘要（中文）

尽管基于变换器架构的大型语言模型在上下文学习（ICL）方面表现出色，但对其能力的理解仍处于早期阶段，现有理论主要集中在简单场景上。本文通过构建合成的上下文学习问题，探讨了更复杂场景下的学习机制。我们证明了变换器可以近似实现最优的ICL算法，并通过实验证实训练后的变换器在该设置中表现出接近最优的ICL性能，揭示了其内部机制与理论的一致性。

🔬 方法详解

问题定义：本文旨在解决变换器在复杂场景下的上下文学习能力不足的问题。现有方法主要关注简单函数，缺乏对复杂表示的理解。

核心思路：我们构建了具有组合结构的合成上下文学习问题，标签依赖于输入的复杂固定表示函数，结合每个实例不同的线性函数。通过这种设计，最优的ICL算法首先通过表示函数转换输入，然后在转换后的数据集上执行线性ICL。

技术框架：整体框架包括输入的表示转换和线性ICL两个主要阶段。首先，变换器的底层负责数据集的转换，随后上层进行线性ICL。

关键创新：本文的主要创新在于理论上证明了变换器可以近似实现复杂场景下的最优ICL算法，并通过实验证实了这一点。与现有方法相比，我们的研究揭示了变换器在复杂表示学习中的潜力。

关键设计：在实验中，我们设置了特定的参数和损失函数，确保变换器能够有效地进行表示学习和线性ICL，具体细节包括对网络层次的设计和训练策略的调整。

📊 实验亮点

实验结果显示，训练后的变换器在复杂场景下的ICL性能接近最优，且在不同层次上展现出不同的学习机制。具体而言，底层有效地转换数据集，而上层则独立实现线性ICL，验证了理论的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、计算机视觉等需要复杂表示学习的任务。通过深入理解变换器的上下文学习机制，未来可以提升模型在多样化场景下的表现，推动智能系统的实际应用。

📄 摘要（原文）

While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic understanding focus mostly on simple scenarios such as learning simple function classes. This paper takes initial steps on understanding ICL in more complex scenarios, by studying learning with representations. Concretely, we construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function, composed with a linear function that differs in each instance. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size. Empirically, we find trained transformers consistently achieve near-optimal ICL performance in this setting, and exhibit the desired dissection where lower layers transforms the dataset and upper layers perform linear ICL. Through extensive probing and a new pasting experiment, we further reveal several mechanisms within the trained transformers, such as concrete copying behaviors on both the inputs and the representations, linear ICL capability of the upper layers alone, and a post-ICL representation selection mechanism in a harder mixture setting. These observed mechanisms align well with our theory and may shed light on how transformers perform ICL in more realistic scenarios.