sdeangroup/CA-Transformer-1D-Stoch-V2K3-Shared

CA-Transformer: 1D Stochastic, V=2, K=3, Shared Head

This repository contains the Shared Head checkpoint for the 1D stochastic V=2, K=3 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	1d-stoch-v2k3-shared
Task	1D stochastic cellular automata
Number of cell states V	2
Neighborhood size K	3 (left, center, right)
Grid	16 (ring, periodic boundaries)
Trajectory length T	20
Context length M	10
Architecture	Shared Head
Transformer layers	2
Attention heads per layer	[1, 1]
MLP per layer (1 = yes, 0 = no)	[1, 1]
Hidden dimension	256
MLP dimension	1024
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. For stochastic rules this amounts to Bayesian counting.

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.

A trajectory is serialized frame by frame, with a time separator 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.

Each trajectory independently samples a new stochastic rule: for every neighborhood configuration, the distribution over next-cell states is drawn from a Dirichlet(1, ..., 1) distribution. Because the rule distribution is continuous, independently generated trajectories almost surely use different rule tables, including across the training, validation and test data.

Validation and test trajectories are kept only if all 8 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.02
Global batch size	1024
Dropout	0.0
Epochs (maximum)	15
Checkpoint selection	best validation teacher-forced cell 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.

The model is compared with the true rule and with a causal add-one k-gram baseline, which counts how often each neighborhood configuration was followed by each state earlier in the same trajectory. This baseline is the Bayes-optimal predictor under the Dirichlet(1, ..., 1) prior. KL divergences are in nats and written as KL(reference ‖ predictor).

• Teacher-forced: each prediction is conditioned on the ground-truth history.
• Autoregressive: starting from the M context frames, the model samples each cell at temperature 1 and feeds its samples back in; the k-gram baseline is rolled out the same way on its own samples. Accuracy compares the sampled cells with the ground-truth trajectory.

Paper-reported results

Metric	Value
TF KL(q_kgram ‖ q_model)	0.0020
TF KL(p_true ‖ q_model)	0.0185
TF model cell accuracy	0.7491
TF k-gram cell accuracy	0.7499
AR model cell accuracy	0.5496
AR k-gram cell accuracy	0.5485
AR KL(q_kgram ‖ q_model)	0.0019

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("1d-stoch-v2k3-shared")

# A fresh Dirichlet(1, ..., 1) rule, never seen in training.
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 1d-stoch-v2k3-shared            # full evaluation
N=2000 bash scripts/eval.sh 1d-stoch-v2k3-shared     # quicker, fewer trajectories

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

bash scripts/train.sh 1d-stoch-v2k3-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}
}