Spatial Induction Heads: In-Context Learning of Multidimensional Cellular Automata

2026-10-06

AI Summary

This study addresses the challenge that serializing multi-dimensional data impedes Transformers from leveraging spatial neighborhoods for in-context learning. To overcome this, it proposes a "spatial induction head" that reconstructs local neighborhoods and matches historical configurations via a two-layer circuit. By integrating a Bayesian counting layer to transcend conventional contiguous block assumptions, the approach achieves multi-dimensional cellular automata prediction without coordinate biases. The findings reveal that the positional dimension depends exclusively on local neighborhoods and spatial dimensionality. Furthermore, the model demonstrates strong generalization to unseen rules, achieving near-perfect deterministic predictions and maintaining KL divergence below 0.005 nats for stochastic rules, thereby validating the effectiveness of the proposed spatial induction mechanism.

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.

Problem

• Induction Heads
• In-Context Learning
• Cellular Automata
• Spatial Routing
• Transformers

Innovation

• Spatial Induction Heads
• In-Context Learning
• Cellular Automata
• Gather-and-Match Circuit
• Mechanistic Interpretability