From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers

TL;DR

This paper analyzes how the diversity of pretraining data influences whether transformers learn generalizable induction heads or rely on positional shortcuts, affecting their out-of-distribution generalization.

Contribution

It provides a rigorous theoretical analysis of how data diversity steers transformers toward induction heads or shortcuts, supported by synthetic experiments.

Findings

Diverse data promotes induction head learning and OOD generalization.

Large trigger-to-trigger distance ratios lead to positional shortcuts.

Optimal pretraining distribution balances context length and computational cost.

Abstract

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.