Guardado en:
Detalles Bibliográficos
Autores principales: Kawata, Ryotaro, Song, Yujin, Bietti, Alberto, Nishikawa, Naoki, Suzuki, Taiji, Vaiter, Samuel, Wu, Denny
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:https://arxiv.org/abs/2512.18634
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866915688537391104
author Kawata, Ryotaro
Song, Yujin
Bietti, Alberto
Nishikawa, Naoki
Suzuki, Taiji
Vaiter, Samuel
Wu, Denny
author_facet Kawata, Ryotaro
Song, Yujin
Bietti, Alberto
Nishikawa, Naoki
Suzuki, Taiji
Vaiter, Samuel
Wu, Denny
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2512_18634
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers
Kawata, Ryotaro
Song, Yujin
Bietti, Alberto
Nishikawa, Naoki
Suzuki, Taiji
Vaiter, Samuel
Wu, Denny
Machine Learning
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.
title From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers
topic Machine Learning
url https://arxiv.org/abs/2512.18634