数据多样性决定Transformer是学规律还是记位置,影响模型泛化能力。
From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers
- 通过分析触发词间距的多样性,判断模型学归纳头还是位置捷径。
- 当触发词间距差异小时,模型能泛化到未见场景;反之则失效。
- 实验验证:丰富上下文分布可稳定诱导出泛化能力强的归纳头。
Transformer 可实现可泛化的算法(如归纳头)或简单的位置捷径(如记忆固定输出位置)。本文研究预训练数据分布如何引导浅层 Transformer 选择其中一种行为。聚焦于一个最小任务——在特殊触发词第二次出现时复制其后一个词——我们对单层 Transformer 的梯度训练进行了严格分析。在无限与有限样本情形下,均证明了学习机制的转变:若输入序列具有足够多样性,以触发词间距离的“最大和比”较低为指标,模型会采用归纳头并泛化至未见上下文;反之,当该比值较大时,模型采用位置捷径,无法进行分布外(OOD)泛化。我们还揭示了预训练上下文长度与 OOD 泛化之间的权衡关系,并推导出每样本计算成本最低的最优预训练分布。最后,通过受控合成实验验证理论预测:拓宽上下文分布能稳健诱导出归纳头,实现 OOD 泛化。结果揭示了预训练 Transformer 的算法偏见,并为数据驱动控制其学习行为提供概念性指导。
原文摘要 · Abstract (English)
Transformers can implement both generalizable algorithms (e.g., induction heads) and simple positional shortcuts (e.g., memorizing fixed output positions). In this work, we study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other. Focusing on a minimal trigger-output prediction task -- copying the token immediately following a special trigger upon its second occurrence -- we present a rigorous analysis of gradient-based training of a single-layer transformer. In both the infinite and finite sample regimes, we prove a transition in the learned mechanism: if input sequences exhibit sufficient diversity, measured by a low ``max-sum'' ratio of trigger-to-trigger distances, the trained model implements an induction head and generalizes to unseen contexts; by contrast, when this ratio is large, the model resorts to a positional shortcut and fails to generalize out-of-distribution (OOD). We also reveal a trade-off between the pretraining context length and OOD generalization, and derive the optimal pretraining distribution that minimizes computational cost per sample. Finally, we validate our theoretical predictions with controlled synthetic experiments, demonstrating that broadening context distributions robustly induces induction heads and enables OOD generalization. Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。