From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers
Ryotaro Kawata, Yujin Song, Alberto Bietti, Naoki Nishikawa, Taiji Suzuki, Samuel Vaiter, Denny Wu
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 develop a rigorous analysis of gradient-based training of a single-layer transformer.In the infinite-sample regime, 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 that generalizes to unseen contexts; by contrast, when this ratio is large, the model resorts to positional memorization and fails to generalize out-of-distribution. Extending to finite samples, we quantify the number of pretraining sequences required to cross this transition and derive the optimal context-length distribution that minimizes sample complexity while ensuring induction-head learning. Finally, we validate our theoretical predictions with controlled synthetic experiments, demonstrating that broadening context distributions robustly induces induction heads and enables length extrapolation. Our results shed light on the algorithmic biases of pretrained transformers and offer practical guidelines for data-driven control of their learned behaviors.