sdeangroup/CA-Transformer-2D-Det-V2K5-Shared

CA-Transformer: 2D Deterministic, V=2, K=5, Shared Head

This repository contains the Shared Head checkpoint for the 2D deterministic V=2, K=5 cellular-automaton task introduced in Spatial Induction Heads: In-Context Learning of Multidimensional Cellular Automata.

The model performs in-context learning of cellular-automaton rules. It receives the first frames of a trajectory generated by an unknown rule and predicts the rest of the trajectory. The rule table is never given to the model, and test trajectories are generated by rules that were not seen during training.

Model details

Property	Value
Model key	2d-det-v2k5-shared
Task	2D deterministic cellular automata
Number of cell states V	2
Neighborhood size K	5 (von Neumann)
Grid	6×6 (torus, periodic boundaries)
Trajectory length T	12
Context length M	8
Architecture	Shared Head
Transformer layers	2
Attention heads per layer	[1, 1]
MLP per layer (1 = yes, 0 = no)	[1, 1]
Hidden dimension	1024
MLP dimension	4096
Positional encoding	learned absolute
Checkpoint file	checkpoint.pt

T is the total number of frames in a trajectory and M is the number of initial frames given as context. The model predicts frames M to T-1.

Architecture

The Shared-head architecture has two layers, each with one attention head and one MLP.

The paper's construction implements a two-stage spatial induction head:

1. Gather: the first layer uses relative positions to pull the cells of each spatial neighborhood from the previous frame into the current position, even though those cells can be far apart in the flattened sequence.
2. Match: the second layer compares the current neighborhood with earlier occurrences of the same configuration in the context and retrieves what came next.

See the paper for the full construction and mechanistic analysis.

Inputs and outputs

Cell states are token IDs 0, 1. The time-separator token has ID 2 and the row-separator token has ID 3.

A trajectory is serialized frame by frame. Each frame is flattened row by row with a row separator between rows, and a time separator is placed between consecutive frames. The model produces next-token logits over this sequence.

The inference utilities in the code repository convert between trajectories and token sequences, so you normally do not need to handle tokens yourself.

Training data

Each trajectory starts from an initial frame whose cells are sampled independently and uniformly over the V states. A single rule is applied at every step, with periodic boundary conditions.

There are 2^32 possible rules on the 5-cell von Neumann neighborhood (north, west, center, east, south). Each rule is grouped with its 0/1 complement, and rules are stratified by fold-k, the smaller of the number of 1s and the number of 0s in the rule table. The train, validation and test pools are disjoint at the class level, and validation and test rules are drawn balanced across fold-k.

Validation and test trajectories are kept only if all 32 neighborhood configurations appear in the context frames, so every entry of the rule table can be inferred from the context. Training trajectories are not filtered.

The original training set contained 50 million trajectories. The training script in the code repository generates fresh trajectories from the same distribution on the fly, and the released rule pools ship with the code.

Training procedure

Hyperparameter	Value
Optimizer	AdamW
Learning rate	0.001
Weight decay	0.2
Adam betas	(0.9, 0.95)
Schedule	cosine with linear warmup
Warmup ratio	0.004
Global batch size	1024
Dropout	0.0
Epochs (maximum)	5
Checkpoint selection	best held-out-rule validation teacher-forced sequence accuracy
Precision	bfloat16 autocast
Random seed	42

One epoch is 50 million training trajectories. Validation runs every 5000 optimizer steps.

Evaluation

The model is evaluated on test trajectories generated by held-out rules. Only cell tokens in frames M to T-1 are scored.

• Teacher-forced cell accuracy: each prediction is conditioned on the ground-truth history.
• Teacher-forced sequence accuracy: a trajectory counts as correct only if every scored cell is predicted correctly.
• Autoregressive cell accuracy: starting from the M context frames, predictions are decoded greedily and fed back into the model.

Paper-reported results

Metric	Value
Teacher-forced sequence accuracy	1.0000
Teacher-forced cell accuracy	1.0000
Autoregressive cell accuracy	1.0000

The paper-reported results for all released models are in configs/paper_results.json.

Usage

Clone and install the accompanying code. Install a PyTorch build that matches your hardware first.

git clone https://github.com/sdean-group/ICL-CA.git
cd ICL-CA
pip install -e .

Load the checkpoint and run inference:

from src import load_pretrained, generate_example, predict, plot_prediction

bundle = load_pretrained("2d-det-v2k5-shared")

# A rule drawn from the held-out test split.
example = generate_example(bundle.spec, seed=0)

result = predict(bundle, example)
print(result.metrics)
plot_prediction(example, result)

load_pretrained downloads checkpoint.pt from the Hugging Face Hub and caches it locally. A complete interactive example is in quickstart.ipynb.

Reproducing the evaluation

bash scripts/eval.sh 2d-det-v2k5-shared            # full evaluation
N=2000 bash scripts/eval.sh 2d-det-v2k5-shared     # quicker, fewer trajectories

Results are written to results/. To retrain this model from scratch:

bash scripts/train.sh 2d-det-v2k5-shared

Intended use

This checkpoint is intended for:

• reproducing the experiments in the paper;
• studying in-context learning of unseen cellular-automaton rules;
• investigating spatial induction heads and other learned transformer mechanisms;
• visualizing teacher-forced and autoregressive predictions;
• serving as a controlled testbed for mechanistic interpretability.

Limitations

• The model only works for the grid size, trajectory length, neighborhood, tokenization and boundary conditions listed above.
• It is trained only on synthetic cellular-automaton trajectories and is not a general-purpose sequence or language model.
• Reported performance follows the paper's evaluation protocol, including its rule splits and the context-coverage filter.

Model collection

This checkpoint is one of 12 models in the Spatial Induction Heads collection. The collection covers 1D V=2, K=3; 1D V=3, K=3; and 2D V=2, K=5 (von Neumann) tasks, each with deterministic and stochastic rules and with Shared-head and Dedicated-heads architectures.

Resources

• Paper: https://arxiv.org/abs/2610.09124
• Code: https://github.com/sdean-group/ICL-CA
• Project website: https://kimzemian.github.io/spatial-induction-heads/
• Model collection: https://huggingface.co/collections/sdeangroup/spatial-induction-heads

License

This model is released under the MIT License.

Citation

@misc{Kazemian2026:spatial,
  title         = {Spatial Induction Heads: In-Context Learning of Multidimensional Cellular Automata},
  author        = {Kimia Kazemian and Menghan Xu and John Thickstun and Sarah Dean},
  year          = {2026},
  eprint        = {2610.09124},
  archiveprefix = {arXiv},
  primaryclass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.09124}
}