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

Kimia Kazemian; Menghan Xu; John Thickstun; Sarah Dean

cs.LGOct 6, 2026

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

Authors: Kimia Kazemian, Menghan Xu, John Thickstun, Sarah Dean

Organizations: Cornell University

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.

Figures & tables

Figure 1: Trained-model predictions on two held-out rules. Top row: Wolfram Rule 110, a Turing-complete deterministic ECA. Bottom row: Stavskaya, a stochastic CA with $p{=}0.29$ ; predictions are colored by the probability the model assigns to state 1. Teacher-forced predictions (middle column) condition on the ground-truth history at every step; autoregressive predictions (right column) are rolled out from the context alone, conditioning on the sample from the previous timestep.

Figure 2: Layer 1 gathers the routing cells; Layer 2 retrieves matching precedents. Attention weights from the shared-head model on 1D $V{=}2,k{=}3$ (elementary cellular automata), Wolfram Rule 110. The query cell at $z$ predicts the target cell at $z+\hat{u}_{d}$ (marked “?”). Layer 1 (left) concentrates on the four routing cells in the previous timestep, at offsets $\mathcal{U}=D_{K}\cup D_{Q}$ relative to $z$ , where $D_{Q}=D_{K}+\hat{u}_{d}$ gives the target cell’s parent offsets. Layer 2 (right) concentrates on matching precedents: earlier cells whose parent configurations (at offsets $D_{K}$ ) match the query configuration (at offsets $D_{Q}$ ). Their states are used to predict the target cell.

		Deterministic		Stochastic
Task	Arch.	TF cell acc	AR cell acc	AR cell acc ( $k$ -gram)	AR KL
1D $V{=}2,\;k{=}3$	Shared-head	1.0000	0.9999	0.5496 (0.5485)	0.0019
	Dedicated-heads	1.0000	1.0000	0.5488 (0.5485)	0.0044
1D $V{=}3,\;k{=}3$	Shared-head	1.0000	1.0000	0.3476 (0.3475)	0.0026
	Dedicated-heads	1.0000	1.0000	0.3475 (0.3475)	0.0004
2D VN $V{=}2,\;k{=}5$	Shared-head	1.0000	1.0000	0.5201 (0.5197)	0.0013

Table 1: Test-set prediction accuracy on held-out rules. Deterministic metrics use the ground-truth history (teacher forcing, TF) or greedy autoregressive (AR) rollout. Stochastic AR cell accuracy is reported for the model and the

k

-gram sampling baseline. Stochastic AR KL is

D_{\mathrm{KL}}(q_{\text{kgram}}\|q_{\text{model}})

between the

k

-gram and model predictive distributions on the model’s rollout.

Figure 3: Attention alignment and ablation. Left: Mean attention mass on the positions predicted by each construction: L1 routing cells at offsets

\mathcal{U}=D_{K}\cup D_{Q}

, and L2 matching precedents, earlier cells whose parent configurations match the query configuration. Right: Cell accuracy when attention to the L1 routing cells or L2 matching precedents is masked at evaluation time, compared with no ablation and controls that mask attention to the same number of randomly selected cells.

Figure 4: Linear probe results. Higher is better. The configuration probe measures classification accuracy for the

k

-cell neighborhood; the k-gram predictive probe measures

1-\overline{D_{\mathrm{KL}}(q_{\text{kgram}}\|q_{\text{probe}})}\,/\,D_{\mathrm{KL}}(q_{\text{kgram}}\|\text{Uniform})

, which the axis label abbreviates as

1-\mathrm{KL}_{\mathrm{probe}}/\mathrm{KL}_{\mathrm{uniform}}

. Solid lines: deterministic; dashed: stochastic. Vertical lines mark the stage at which the theoretical construction places each target.

Figure 5: t-SNE visualization of frozen representations from the Shared-head model on deterministic 1D $V{=}2,k{=}3$ . Top row (colored by neighborhood configuration): configurations cluster after the L1 MLP. Bottom row (colored by output value): outputs cluster after L2 attention.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Stage result	Downstream usage
Corollary E.6	Positional features for offset selection in previous time steps
Corollary F.5	Distinct color slots and storage costs (Construction II)
Corollary G.12	State blocks (approximate or exact) per routed offset
Theorem H.2	Equal-weight matching for empirical counts and Dirichlet smoothing
Lemmas H.3 , H.5 , H.6	Bounds on separation, dispersion, and predictive cost
Theorem H.9	Finite-parameter prediction guarantee

Appendix

Table 2: Dependencies between construction stages.

symbol	meaning	fixed at
Grid and dual group
$d$	number of spatial dimensions	Def. 1
$W_{j}$	side length of axis $j$	Def. 1
$W_{\max}$	largest side length, $\max_{j}W_{j}$	Lem. E.3
$\mathcal{G}$ , $\lvert\mathcal{G}\rvert$	the torus $[W_{1}]\times\cdots\times[W_{d}]$ ; its volume $\prod_{j}W_{j}$	Def. 1
$z$ , $z_{j}$ , $\delta$	a cell of $\mathcal{G}$ ; its $j$ -th coordinate; an offset	Def. 1

Appendix

Table 3: Symbols, and where each is fixed.

	1D, $V{=}2$	1D, $V{=}3$	2D von Neumann
$\mathcal{G}$	$\mathbb{Z}_{16}$	$\mathbb{Z}_{32}$	$\mathbb{Z}_{6}^{2}$
$d$ , $W_{\max}$ , $\lvert\mathcal{G}\rvert$	$1$ , $16$ , $16$	$1$ , $32$ , $32$	$2$ , $6$ , $36$
$V$ , $T$	$2$ , $10$	$3$ , $12$	$2$ , $12$
$D_{K}$ , $k$	$\{-1,0,1\}$ , $3$	$\{-1,0,1\}$ , $3$	von Neumann, $5$
$\hat{u}_{d}$ , $\lvert\mathcal{U}\rvert$	$1$ , $4$	$1$ , $4$	$(0,1)$ , $8$
tokens per step, $N$	$17$ , $170$	$33$ , $396$	$42$ , $504$

Appendix

Table 4: The three experimental settings as tasks.

token	predicts	spatial target
cell $(t,z)$ , $z_{d}<W_{d}-1$	cell $(t,z+\hat{u}_{d})$	$z+\hat{u}_{d}$
cell $(t,z)$ , $z_{d}=W_{d}-1$	a separator	none
separator	first cell of the next row	$\kappa(n)+\hat{u}_{d}$

Appendix

Table 5: Prediction duties under Definition C.1 . For a non-row-final cell and for a separator, the predicted token is a cell and its spatial location is expressed through the constant successor offset

\hat{u}_{d}

. A row-final cell instead predicts a separator whose identity is fixed by position.

	1D, $V{=}2$	1D, $V{=}3$	2D von Neumann
$\lvert(\mathcal{U}{-}\mathcal{U})\setminus\{0\}\rvert$	$6$	$6$	$21$
degree bound on $M^{\star}$	$7$	$7$	$22$
subgroup $K$	$4\mathbb{Z}_{16}$	$4\mathbb{Z}_{32}$	$\langle(2,2)\rangle$
colors $M=M_{K}$	$4$	$4$	$12$

Appendix

Table 6: The coloring Construction II uses in each experimental setting. Every one is admissible by Lemma F.2 , and every one is far below the degree bound on the same line.

neighborhood	grid	$\lvert\mathcal{U}\rvert$	least periodic $M_{K}$	degree bound	$\lvert\mathcal{G}\rvert$
elementary	$\mathbb{Z}_{5}$	4	5	7	5
elementary	$\mathbb{Z}_{16}$	4	4	7	16
elementary	$\mathbb{Z}_{32}$	4	4	7	32
von Neumann	$\mathbb{Z}_{5}^{2}$	8	25	23	25
von Neumann	$\mathbb{Z}_{6}^{2}$	8	12	23	36
von Neumann	$\mathbb{Z}_{7}^{2}$	8	49	23	49

Appendix

Table 7: Colorings across grids, for reference.

M_{K}

is the least number of colors over periodic admissible colorings and is arithmetic in the side length. The degree bound

\lvert(\mathcal{U}-\mathcal{U})\setminus\{0\}\rvert+1

of Lemma F.3 is evaluated with the differences taken in

\mathbb{Z}^{d}

, so it is the same for every grid and bounds

M^{\star}

from above; the exact torus count can be smaller, as Remark F.4 shows on

\mathbb{Z}_{6}^{2}

. The bound never degenerates and

M_{K}

sometimes does: on

\mathbb{Z}_{5}^{2}

and

\mathbb{Z}_{7}^{2}

every nontrivial subgroup has a nonzero intersection with

\mathcal{U}-\mathcal{U}

, so

M_{K}=\lvert\mathcal{G}\rvert

and every cell takes a color of its own.

Construction I (Thm. G.9 )
	1D, $V{=}2$	1D, $V{=}3$	2D von Neumann
heads, $d_{\mathrm{MLP}}$	$4$ , $0$	$4$ , $0$	$8$ , $0$
$d_{S}=2d$ , $d_{p}=d_{S}{+}3$	$2$ , $5$	$2$ , $5$	$4$ , $7$
width of $B_{3}=\lvert\mathcal{U}\rvert V$	$8$	$12$	$16$
residual $d_{\mathrm{model}}$	$20$	$26$	$30$
block-error threshold $1/(2k{+}2)$	$1/8$	$1/8$	$1/12$

Appendix

Table 8: What the two constructions require on the three experimental settings, at the colorings of Table 6 .

Dataset type	Train	Val-in (no filter)	Val-in (filter)	Val-out / Val	Test
Deterministic	No	No	Yes	Yes	Yes
Stochastic	No	—	—	Yes	Yes

Appendix

Table 9: Trajectory coverage filter by split. “Yes” means the split rejection-samples trajectories until all

V^{k}

neighborhood configurations appear in the context; “No” means no filter. Deterministic runs use five splits (train, val-in without filter, val-in with filter, val-out, test); stochastic runs use three (train, validation, test). The main-text deterministic validation refers to val-out; val-in is diagnostic only and not used for checkpoint selection.

Task	Rule type	Grid $\mathcal{G}$	$T$	$T_{0}$	$V^{k}$ configs
1D $V{=}2,\;k{=}3$	det / stoch	16	10 / 20	4 / 10	8
1D $V{=}3,\;k{=}3$	det / stoch	32 / 16	12 / 20	8 / 10	27
2D VN $V{=}2,\;k{=}5$	det / stoch	$6\times 6$	12 / 20	8 / 10	32

Appendix

Table 10: Grid and trajectory parameters for each experimental setting. Values shown as “det / stoch” when they differ between the two regimes.

Task	Arch.	heads	mlp	$d_{\text{model}}$	$d_{\text{ffn}}$	lr	wd	$\beta_{2}$	ep
1D $V{=}2,k{=}3$	Shared-head	$[1,1]$	$[1,1]$	128	512	$10^{-3}$	0.2	0.999	3
	Dedicated-heads	$[4,1]$	$[0,1]$	128	512	$10^{-3}$	0.2	0.999	3
1D $V{=}3,k{=}3$	Shared-head	$[1,1]$	$[1,1]$	256	1024	$5{\times}10^{-4}$	0.2	0.95	8
	Dedicated-heads	$[4,1]$	$[0,1]$	256	1024	$5{\times}10^{-4}$	0.2	0.95	8
2D VN $V{=}2,k{=}5$	Shared-head	$[1,1]$	$[1,1]$	1024	4096	$10^{-3}$	0.2	0.95	5
	Dedicated-heads	$[8,1]$	$[0,1]$	1024	4096	$10^{-3}$	0.2	0.95	5

Appendix

Table 11: Architecture and training hyperparameters: deterministic settings.

Task	Arch.	heads	mlp	$d_{\text{model}}$	$d_{\text{ffn}}$	lr	wd	$\beta_{2}$	ep
1D $V{=}2,k{=}3$	Shared-head	$[1,1]$	$[1,1]$	256	1024	$10^{-3}$	0.2	0.999	15
	Dedicated-heads	$[4,1]$	$[0,1]$	128	512	$10^{-3}$	0.2	0.999	8
1D $V{=}3,k{=}3$	Shared-head	$[1,1]$	$[1,1]$	256	1024	$5{\times}10^{-4}$	0.2	0.95	8
	Dedicated-heads	$[4,1]$	$[0,1]$	256	1024	$5{\times}10^{-4}$	0.2	0.95	8
2D VN $V{=}2,k{=}5$	Shared-head	$[1,1]$	$[1,1]$	1024	4096	$5{\times}10^{-4}$	0.2	0.95	10
	Dedicated-heads	$[8,1]$	$[0,1]$	512	2048	$10^{-3}$	0.2	0.95	5

Appendix

Table 12: Architecture and training hyperparameters: stochastic settings.

Task	Arch.	TF seq acc	TF cell acc	AR cell acc
1D $V{=}2,\;k{=}3$	Shared-head	0.9990	1.0000	0.9999
	Dedicated-heads	0.9994	1.0000	1.0000
1D $V{=}3,\;k{=}3$	Shared-head	0.9994	1.0000	1.0000
	Dedicated-heads	0.9989	1.0000	1.0000
2D VN $V{=}2,\;k{=}5$	Shared-head	1.0000	1.0000	1.0000
	Dedicated-heads	0.9989	1.0000	0.9999

Appendix

Table 13: Deterministic evaluation: full metrics on held-out test rules.

		TF KL		TF cell acc	AR
Task	Arch.	kg $\\|$ m	true $\\|$ m	(kgram)	cell acc (kgram)	KL kg $\\|$ m
1D $V{=}2,k{=}3$	Shared-head	0.0020	0.0185	0.7491 (0.7499)	0.5496 (0.5485)	0.0019
	Dedicated-heads	0.0046	0.0211	0.7469 (0.7499)	0.5488 (0.5485)	0.0044
1D $V{=}3,k{=}3$	Shared-head	0.0026	0.0881	0.5669 (0.5674)	0.3476 (0.3475)	0.0026
	Dedicated-heads	0.0004	0.0859	0.5675 (0.5674)	0.3475 (0.3475)	0.0004
2D VN $V{=}2,k{=}5$	Shared-head	0.0014	0.0300	0.7345 (0.7351)	0.5201 (0.5197)	0.0013

Appendix

Table 14: Stochastic evaluation: full metrics on held-out test trajectories. TF KL columns report teacher-forced divergence against the

k

-gram and true-rule baselines.

Figure 6: t-SNE of frozen residual-stream representations from the Dedicated-heads model on deterministic 1D $V{=}2,k{=}3$ . Top row: colored by neighborhood configuration. Bottom row: colored by the ground-truth output value. Under the Dedicated-heads construction, the configuration should become accessible after L1 attention rather than after the L1 MLP.

Figure 7: t-SNE of frozen representations from the Shared-head model on stochastic 1D $V{=}2,k{=}3$ . Top row: colored by neighborhood configuration. Bottom row: colored by the model’s predicted $P(y{=}1)$ at the corresponding query position. Because rules are drawn from a continuous Dirichlet prior, the output value at a fixed configuration is a sample from the true conditional distribution, so the model’s predictive distribution varies continuously across queries with the same configuration.

Figure 8: t-SNE of frozen representations from the Shared-head model on deterministic 2D von Neumann $V{=}2,k{=}5$ . Top row: colored by neighborhood configuration. Bottom row: colored by the ground-truth output value. The larger neighborhood ( $k=5$ , $V^{k}=32$ configurations) is harder to display in two dimensions than the ECA case; per-configuration clusters are correspondingly less separated.

	Deterministic (Rule 110)	Stochastic (Stavskaya)
Task	1D $V{=}2,k{=}3$	1D $V{=}2,k{=}3$
Training $\lvert\mathcal{G}\rvert,T,T_{0}$	$64,64,12$	$64,64,12$
Plotted $\lvert\mathcal{G}\rvert,T,T_{0}$	$64,64,12$	$64,64,12$
heads / mlp	$[4,1]/[0,1]$	$[1,1]/[1,1]$
$d_{\text{model}}$ / $d_{\text{ffn}}$	$512/2048$	$512/2048$
lr, wd, epochs	$5{\times}10^{-4}$ , 0.2, 10	$10^{-3}$ , 0.2, 8

Appendix

Table 15: Configuration of the visualization checkpoints used for Figure 1 . The deterministic panel uses greedy decoding; the stochastic panel uses temperature-1 sampling, and the displayed values are per-cell probabilities of state 1.

Explore similar work

Induction Heads Interpolate N-Grams
Jul 2, 2026Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman +1Transformer Architectures In-Context Learning
1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.
Invariant Learning Dynamics of Transformers in Inductive Reasoning Tasks
Jul 13, 2026Tiberiu Musat, Tiago Pimentel, Nicolas Zucchet +1Transformer Architectures Induction
ETH Zurich · Stanford University
Storing Is Not Remembering: LSTM-UT and Bounded Gated Memory for Looped Transformers
Sep 17, 2026Aras Kavuncu, Muhammad Burhan HafezTransformer Architectures Neural Cellular Automata
University of Southampton, Electronics and Computer Science

Jul 2, 2026, cs.LGJ/K move · Enter open · S save

Induction Heads Interpolate N-Grams

Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman, Nicolas Flammarion

1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.

Transformer Architectures In-Context Learning Unigram Tokenizer Induction

Open paper AskPDF arXiv

Induction heads are attention circuits believed to underlie in-context learning in transformers, yet a precise characterization of the estimators they implement remains elusive. We study transformers trained on order-

k

Markov chains and identify two complementary smoothing mechanisms. First, at finite attention-weight scale, the circuit implements a soft context-matching estimator: it aggregates contributions from exact and partial context matches, weighted exponentially by their overlap, and induces a data-dependent interpolation across context orders analogous to Jelinek-Mercer smoothing. Second, a beginning-of-sequence (BOS) token induces additive pseudo-counts, recovering Dirichlet-style smoothing. We construct a disentangled transformer implementing both mechanisms and show that trained transformers recover the predicted attention patterns. Across settings where pseudo-count smoothing is optimal or lower-order contexts provide structured evidence, trained transformers match or outperform classical count-based baselines. Our results bridge mechanistic interpretability of induction heads with classical statistical smoothing, revealing that transformers learn to regularize in-context estimation rather than simply count.

Abstract

Explore similar work

Induction Heads Interpolate N-Grams

Invariant Learning Dynamics of Transformers in Inductive Reasoning Tasks

Storing Is Not Remembering: LSTM-UT and Bounded Gated Memory for Looped Transformers