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

TL;DR

<3-5 sentence high-level summary> The paper investigates how pretraining data diversity determines whether a transformer learns a generalizable induction-head mechanism or a brittle positional shortcut for a simple trigger-output copying task. By analyzing gradient-based training of a shallow transformer in both population and finite-sample regimes, it introduces the max-sum ratio $R(\mathcal{D}_\ell)$ as the key quantity governing the transition between mechanisms. The results show that sufficiently diverse pretraining data leads to robust OOD generalization via induction heads, while low diversity biases the model toward position-based memorization that fails out-of-distribution. The work also derives an optimal linear-in-context-length pretraining distribution to minimize computational cost and validates predictions with synthetic experiments and broader gradient-based settings.

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.

Figures (5)

• Figure 1: Two mechanisms for the associative copying task [..., t, o, ..., t] $\mapsto$ o. In the pretraining data, the size of irrelevant tokens before the occurrence of the first and second trigger $\ell$ remains fixed per sequence, hence allowing two solutions: $(i)$positional shortcut that outputs the token at position $(T+1)/2$ for input length $T$; and $(ii)$induction head using token embedding $\boldsymbol{e}$, which finds the queried token and returns the ensuing token. Whereas on OOD sequences with varying $\ell_1\neq\ell_2$, only $(ii)$ remains a valid solution.
• Figure 2: Attention heatmaps of $\boldsymbol{W}^{KQ}$ when the pretraining sequence diversity is small (left) and large (right). In the left figure, there is a strong positional shortcut that links position 9 to position 5 (the correct position in pretraining data), whereas in the right figure, the trigger positions are more dispersed, weakening this shortcut. Instead, a signal corresponding to induction head -- detecting tokens after trigger -- becomes dominant.
• Figure 3: Out-of-distribution accuracy map over varying $\ell_\mathrm{min}$ (vertical) and $\ell_\mathrm{max}$ (horizontal); moving right indicates greater diversity of $\ell$ in the pretraining distribution.
• Figure 4: Map of two types of errors due to the positional shortcut. Note that both errors can probabilistically coincide with the correct answer, and such cases are not excluded.
• Figure 5: Accuracy map over $\ell_\mathrm{min}$ (vertical), $\ell_\mathrm{max}$ (horizontal) for a 3-layer transformer trained with AdamW.

Theorems & Definitions (26)

• Definition 1: Data Distribution
• Definition 2
• Definition 3
• Remark 1
• Definition 4
• Example 1
• Theorem 5: Infinite Sample Setting
• Remark 2
• Theorem 6: Finite Sample Setting
• Lemma 7: Informal