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

AI summary

Overview

Research area: Mechanistic interpretability of transformers, in-context learning, and cellular automata (Cornell University; arXiv:2610.09124v1 [cs.LG], 06 Oct 2026).

Technical level: Advanced. The paper combines formal circuit constructions, attention-weight proofs, and controlled empirical verification; the appendices contain the full parameter-level constructions.

Scope: The paper formalizes, constructs, and empirically verifies how a two-layer transformer uses learned absolute positional embeddings to gather spatially scattered context from a serialized grid and then match that context against precedent to predict the next state.

What This Paper Is About

Induction heads are a well-established mechanistic explanation of in-context learning in sequential data, but existing theory assumes the relevant context sits in a contiguous block immediately before the query. When a spatial grid is flattened into a one-dimensional token stream, local neighbors get scattered to distant positions, so that assumption fails. This paper asks how a transformer addresses and assembles scattered tokens to reconstruct a spatial neighborhood, and how, once assembled, it optimally infers an unknown local dynamical rule without any coordinate-based spatial inductive bias.

Key Contributions

1. Task and theoretical framework. The authors formalize spatial in-context learning using multidimensional cellular automata, derive the closed-form Bayes-optimal predictor for the Dirichlet prior (Theorem 3), and introduce the spatial induction head (Definition 5): a two-layer circuit whose first layer gathers spatial dependencies at each token and whose second layer matches the resulting configuration against precedent.

2. Mechanistic constructions of the gather. Two explicit realizations are given (Theorem 7, parts I and II): one using a head per offset, and one using a single content-routed head with a separating MLP. Both route scattered spatial neighbors with a compact set of positional frequencies whose number is independent of grid volume and trajectory horizon (Lemma 6).

3. Bayesian matching. The matching layer is shown to perform Bayesian counting: once the gathered representation is available, a single Layer-2 head computes the posterior exactly under hard attention (Proposition 8) and arbitrarily closely otherwise (Theorem 9).

4. Empirical verification. Two-layer transformers trained on 1D and 2D cellular automata generalize to held-out rules, surpassing prior in-context learning results on cellular automata, while attention maps (Figures 2 and 3) and layerwise probes (Figure 4) support the predicted gather-and-match computation.

Main Findings

• Spatial routing is a distinct problem from sequential induction. Serializing a grid preserves adjacency only along the innermost axis and obscures adjacency along the others, so the k tokens of a parent configuration need not form a contiguous block. The authors define the routing set as U ≔ D_K ∪ (D_K + û_d), where D_K is the key-offset set and û_d is the successor offset. In examples on large enough grids: the 1D elementary neighborhood gives D_K = {−1,0,1}, k = 3, |U| = 4; the 2D von Neumann neighborhood gives k = 5 and |U| = 8; the 2D Moore neighborhood gives k = 9 and |U| = 12. In general k ≤ |U| ≤ min(2k, |G|).

• Routing width does not scale with grid volume or horizon. Lemma 6 shows a Layer-1 positional score that ranks previous-step cells at offsets N above all other visible cells can be realized with d_p = d_S + 2, where the spatial addressing width is d_S = O(|N| d log |N|) in general, or d_S = 2d when |N| = 1. This positional dimension is independent of |G| and T.

• Two concrete gather constructions. Construction I (Dedicated-heads) uses |U| heads, one per offset, and no Layer-1 MLP, with d_p = 2d + 3 and d_model = (|U| + 2)V + 2d + 6, recovering each routed state to any accuracy. Construction II (Shared-head) uses one head plus an MLP that separates routed states by offset, with d_MLP ≤ 2(V−1)|U|³ + 3 and d_p ≤ O(|U| d log |U|) + |U|³ + 3, recovering them exactly. Both widths are independent of |G| and T.

• Matching implements Bayesian counting. Under hard attention, a single Layer-2 head that assigns the same score to every matching precedent and −∞ to all others outputs the empirical conditional transition law (n_{c,y}/n_c). Adding one prior token of value (1/V)1 and score a + log(Vα) makes the output exactly the posterior P̂_α(y|c) = (n_{c,y} + α)/(n_c + Vα) from Theorem 3.

• End-to-end guarantees. With Δ ≔ β(1 − (2k+2)ε̄) as the match–nonmatch separation, Theorem 9 shows: for stochastic rules, if N ≥ Vα and Δ ≥ log(4N/((n_c + Vα)κ)) and ε̄ ≤ κ/(16kΔ), the Layer-2 attention output read as a distribution satisfies D_KL(P̂_α(·|c) ∥ p) ≤ κ for any κ ∈ (0,1]; for deterministic rules (α → 0), if ε̄ < 1/(2k+2) and Δ > log(2√2 (N − n_c)/n_c), the predicted next state is exactly correct.

• Trained models recover unseen rules. Across three settings (1D V=2, k=3; 1D V=3, k=3; 2D von Neumann V=2, k=5) and both architectures, deterministic teacher-forced cell accuracy is 1.0000 in all six model/setting combinations, and autoregressive cell accuracy is 0.9999 or 1.0000. On stochastic rules, all autoregressive KL values fall below 0.005 nats, ranging from 0.0004 to 0.0044.

• Stochastic autoregressive scores track the k-gram ceiling. Stochastic autoregressive cell accuracy is near the posterior-predictive k-gram baseline in every case: Shared-head and Dedicated-heads score 0.5496 and 0.5488 against a baseline of 0.5485 (1D V=2, k=3); 0.3476 and 0.3475 against 0.3475 (1D V=3, k=3); 0.5201 and 0.5191 against 0.5197 (2D VN V=2, k=5). The authors note the ceiling is well below 1.

• Attention lands where the construction predicts. Figure 3 (left) reports that both architectures place substantial attention mass on the Layer-1 routing cells at offsets U and on the Layer-2 matching precedents; Figure 2 illustrates the pattern for the shared-head model on Wolfram Rule 110 (1D V=2, k=3).

• Ablations are causal. Masking attention to the theory-predicted cells at evaluation time reduces cell accuracy in all three settings, while masking attention to an equal number of randomly selected control cells has little effect (Figure 3, right).

• Probes match the predicted layerwise structure. Linear probes on frozen post-residual representations indicate that the neighborhood configuration becomes linearly readable after the Layer-1 MLP in Shared-head models and after Layer-1 attention in Dedicated-heads models, and that the predictive distribution becomes readable after Layer-2 attention in both (Figure 4). t-SNE on the Shared-head model for deterministic 1D V=2, k=3 shows representations clustering by neighborhood configuration after the L1 MLP and by output value after L2 attention (Figure 5).

• Qualitative examples. Figure 1 shows trained-model predictions on two held-out rules: Wolfram Rule 110 (a Turing-complete deterministic elementary cellular automaton) and the stochastic Stavskaya rule at p = 0.29, with teacher-forced and autoregressive columns.

• How rules are counted. A neighborhood of size k over V states admits V^k configurations and therefore V^(V^k) deterministic rules: 256 in the elementary case, 2^32 for von Neumann, and 2^512 for Moore. Stochastic rules form a continuum of dimension V^k(V−1).

Methodology in Plain English

The authors pick a setting where the "correct" answer is knowable in closed form. They define a stochastic cellular automaton on a torus: a grid of cells, each updated in parallel by a local kernel that maps a cell's k-cell neighborhood to a distribution over the V possible next states. Crucially, the kernel is drawn fresh for every trajectory from a Dirichlet prior and never shown to the model, so the model must infer the rule purely from earlier occurrences of the same neighborhood pattern in its context.

The grid is flattened into a single token stream in row-major order, with separators marking row and time boundaries. Because serialization scatters spatial neighbors, the authors work out exactly which offsets a given query token must reach: the query offsets D_Q = D_K + û_d for the target cell's parents, and the routing set U = D_K ∪ D_Q, which covers both what the query must read and what the query token must represent so it can serve as a precedent for later predictions.

They then do theory. They derive the Bayes-optimal predictor for the Dirichlet prior — a smoothed count of how often each state followed each configuration — and define which queries are "resolvable" (those whose configuration appeared earlier in the stream). They construct a two-layer circuit matching this, with Layer 1 gathering the routing cells using learned positional addressing, and Layer 2 matching the query configuration against earlier parent configurations and averaging the following states. They give two ways to build Layer 1 and bound the positional dimensions needed. They then prove the whole circuit approximates the Bayes posterior to any accuracy for stochastic rules and is exactly correct for deterministic ones.

Finally, they train two-layer transformers with learned absolute positional embeddings — no coordinate-based spatial bias — instantiating both gather constructions. They evaluate on rules held out from training, measure accuracy and KL, inspect attention maps, run masking ablations with matched random controls, and train linear probes on frozen representations at each layer to check whether the internal structure matches what the construction predicts.

Why This Matters

Impact on research. The paper extends induction-head theory beyond the contiguous-context assumption that has underpinned rigorous analyses of in-context learning for Markov chains. It shows that the classical picture is incomplete for multidimensional data, and it supplies a replacement circuit — a two-layer gather-and-match — with matching optimality guarantees and the routing machinery needed to handle scattered context. It also gives a concrete account of how learned absolute positional embeddings can encode grid geometry that they were never given.

Real-world applications. The paper does not build applications; the following are areas where its framing is directly relevant.

• Weather and climate modeling, where grid-structured data is often handled with coordinate-dependent attention or mesh graphs, in contrast to the geometry-free setting studied here.
• Vision transformers, which flatten image patches and use learned one-dimensional absolute position embeddings that have been shown to recover aspects of the underlying 2D grid structure.
• Scientific simulation and forecasting of dynamical systems, including discretized chaotic systems and hidden Markov processes, where models must infer local transition structure from context.
• Cellular-automaton-style systems more broadly, whose simple local rules can generate essentially unbounded complexity and which the authors describe as historically viewed as systems of locally interacting Markov chains on a multidimensional grid.

Industry relevance. Mechanistic understanding of how transformers route and match spatial context matters for models that process raster, gridded, or otherwise multidimensional data as token sequences. Knowing that a two-layer circuit suffices, and that the positional width needed for routing depends only on the neighborhood and spatial dimension rather than grid volume or horizon, gives concrete guidance for architecture and positional-embedding design in such systems.

Future Directions

• Beyond cellular automata. The authors frame their construction as a general template for understanding induction mechanisms in other kinds of spatially structured data, since their analysis deliberately isolates finite-state local dynamics.
• Alternative positional schemes. The model studied uses learned absolute positional embeddings; the paper's Appendix C.3 compares alternative positional embedding schemes, and the main text notes that positional embedding design is an open dimension of the question.
• Lower bounds and sharpness. Lemma 6 and Theorem 7 give achievable positional widths; whether these bounds are tight, and how the head-versus-nonlinearity tradeoff behaves in practice, is left open.
• Where sequence-relative routing fails. The paper proves that sequence-relative biases cannot route local neighborhoods on a torus in a single layer, motivating spatial addressing; how far that limitation propagates to deeper or differently biased models is not resolved.

Target Audience

Mechanistic interpretability and deep learning theory researchers who study in-context learning and transformer circuits; researchers working on structured, grid-valued, or spatiotemporal data who want to understand how models handle serialization; and machine learning graduate students comfortable with attention mechanics, Bayesian posteriors, and circuit-level proofs. The core constructions are mathematically demanding, but the framing of the problem — that flattening a grid hides the geometry a model must exploit — is accessible to a broader technical readership.

Authors’ 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.