Transformers learn spatial patterns from flattened cellular automata data
Spatial Induction Heads: In-Context Learning of Multidimensional Cellular Automata
Machine LearningArtificial Intelligence
Summary
It’s hard for AI models to learn from data that’s spread out in space when it’s been flattened into a sequence. The authors studied how transformer models learn local rules that govern grid-based patterns called cellular automata, even when they only see a long, flat list of data points. They uncovered a special two-step process inside transformers that first gathers neighboring spatial information, then matches it with past patterns to make accurate predictions. This allows the model to guess future states of the system almost perfectly, even without being told about spatial positions explicitly.
What this means in practice
- •For machine learning engineers: Build transformer models that predict complex spatial patterns from flattened multidimensional data without explicit spatial inputs.
- •For computational physics teams: Simulate and forecast multidimensional cellular automata dynamics accurately using sequence-based transformers trained with spatial induction mechanisms.
Authors
Kimia Kazemian, Menghan Xu, John Thickstun, Sarah Dean
Abstract
Induction heads provide a mechanistic account of in-context learning in sequential data, but existing theory largely assumes that the context relevant to a prediction forms a contiguous block. In multidimensional data, serialization breaks this assumption by scattering spatial neighbors across distant positions in the token sequence. We study how transformers overcome this routing problem in multidimensional stochastic and deterministic cellular automata, where each trajectory is generated by an unknown local rule and presented as a flattened sequence without an explicit coordinate-based spatial inductive bias. We introduce spatial induction heads, two-layer gather-and-match circuits in which the first layer reconstructs the relevant spatial neighborhood and the second matches the resulting configuration against earlier occurrences. We give two explicit realizations of the gather and show that the positional dimension required for spatial routing depends only on the local neighborhood and spatial dimension, not on grid volume or trajectory horizon. We further construct a matching layer which implements Bayesian counting. The end-to-end circuit can approximate the Bayesian posterior arbitrarily closely for stochastic rules and can predict exactly for deterministic rules. Empirically, trained two-layer transformers generalize to unseen rules in one and two dimensional settings, achieving near-perfect deterministic rollouts and less than 0.005 nats KL from the Bayes-optimal predictor on stochastic rules. Attention patterns and layerwise probes align with the predicted gather-and-match computation, providing mechanistic evidence for spatial induction in trained transformers.