README.md

---
license: apache-2.0
library_name: pytorch
tags:
- cellular-automata
- generative pretrained transformers
- forecasing
- rule inference
- interpretable modeling
---

# AutomataGPT

This Hugging Face repository hosts the paper-era datasets and checkpoint folders for **AutomataGPT**, a decoder-only transformer for **forward prediction** and **inverse rule inference** in two-dimensional binary deterministic cellular automata (CA).

AutomataGPT studies two tasks:

- **Forward problem**: predict the next CA state from a rules matrix and an initial condition.
- **Inverse problem**: infer a rules matrix from an initial condition and an observed future state.

This repository is intended as the **weights-and-data companion** to the main GitHub codebase.

## Links

- GitHub repository: https://github.com/lamm-mit/AutomataGPT
- arXiv paper: https://arxiv.org/abs/2506.17333

## Access

This is a **gated Hugging Face repository**. You must accept the repository access conditions before downloading files.

## What This Repository Contains

Although this is a Hugging Face **model** repository, it stores both:

- **CA datasets** in `CA_datasets/`
- **trained checkpoint folders** in `model_parameters/`

### Datasets

The repository currently includes the following dataset files:

- `CA_datasets/CA_dataset_2r_10kgpr.npy`
- `CA_datasets/CA_dataset_2r_100_val.npy`
- `CA_datasets/CA_dataset_5r_100gpr.npy`
- `CA_datasets/CA_dataset_10r_10kgpr.npy`
- `CA_datasets/CA_dataset_10r_100_val.npy`
- `CA_datasets/CA_dataset_100r_10kgpr.npy`
- `CA_datasets/CA_dataset_100r_10_val.npy`
- `CA_datasets/CA_dataset_100r_2gpr.npy`

### Checkpoints

The repository currently includes the following checkpoint folders:

- Forward, 2 rulesets:
  `model_parameters/10_23_2024_2_Ruleset_Broad_Ent_2024-10-23 17-07-12/`
  checkpoint: `AutomataGPT_2DBD_epoch_50.pt`

- Forward, 10 rulesets:
  `model_parameters/10_29_2024_10_Ruleset_Broad_Ent_2024-10-29 09-34-51/`
  checkpoint: `AutomataGPT_epoch_50.pt`

- Forward, 100 rulesets:
  `model_parameters/11_14_2024_100_Ruleset_Broad_Ent_2024-11-15 01-49-23/`
  checkpoint: `AutomataGPT_epoch_20.pt`

- Inverse, 2 rulesets:
  `model_parameters/10_30_2024_2_Ruleset_Inverse_Broad_Ent_2024-10-30 10-58-04/`
  checkpoint: `AutomataGPT_2DBD_10_rules_inverse_epoch_50.pt`

- Inverse, 10 rulesets:
  `model_parameters/10_30_2024_10_Ruleset_Inverse_Broad_Ent_2024-10-30 00-39-25/`
  checkpoint: `AutomataGPT_2DBD_10_rules_inverse_epoch_50.pt`

- Inverse, 100 rulesets:
  `model_parameters/11_11_2024_100_Ruleset_Inverse_Broad_Ent_2024-11-11 18-06-45/`
  checkpoint: `AutomataGPT_2DBD_100_rules_inverse_epoch_20.pt`

Each checkpoint folder also contains:

- `loss_data.csv`
- `*_info.txt`

## Recommended Checkpoints

For most users, the two most relevant paper-era checkpoints are:

- **Best forward checkpoint**:
  `model_parameters/11_14_2024_100_Ruleset_Broad_Ent_2024-11-15 01-49-23/AutomataGPT_epoch_20.pt`

- **Best inverse checkpoint**:
  `model_parameters/11_11_2024_100_Ruleset_Inverse_Broad_Ent_2024-11-11 18-06-45/AutomataGPT_2DBD_100_rules_inverse_epoch_20.pt`

## Important Note on Loading

These checkpoints are **not** packaged as `transformers.from_pretrained()` models.

They are raw **PyTorch `state_dict` checkpoints** trained with the code in the GitHub repository, using:

- a fixed tokenizer defined in `train_model.py`
- `x-transformers`
- custom forward/inverse prompt formatting

If you want to use these checkpoints, the recommended workflow is:

1. clone the GitHub repo
2. create the conda environment from the GitHub repo
3. download selected datasets/checkpoints from this Hugging Face repo into that clone
4. load the checkpoints with the GitHub utilities

## Download Example

The snippet below downloads the best inverse checkpoint and one evaluation dataset into the current working tree:

```python
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="lamm-mit/AutomataGPT",
    repo_type="model",
    local_dir=".",
    allow_patterns=[
        "model_parameters/11_11_2024_100_Ruleset_Inverse_Broad_Ent_2024-11-11 18-06-45/*",
        "CA_datasets/CA_dataset_100r_2gpr.npy",
    ],
)
```

## Load the Best Inverse Checkpoint

The example below follows the original GitHub loading path for the best inverse model.

```python
import os
import torch
from train_model import build_tokenizer
from x_transformers import Decoder, TransformerWrapper
from x_transformers.autoregressive_wrapper import AutoregressiveWrapper

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

fast_tokenizer = build_tokenizer()
vocab_size = len(fast_tokenizer.get_vocab())
max_length = 555

model = TransformerWrapper(
    num_tokens=vocab_size,
    max_seq_len=max_length,
    attn_layers=Decoder(
        dim=256,
        depth=6,
        heads=4,
        rotary_pos_emb=True,
        attn_flash=True,
    ),
)
model = AutoregressiveWrapper(model, mask_prob=0.15)

ckpt_path = os.path.join(
    "model_parameters",
    "11_11_2024_100_Ruleset_Inverse_Broad_Ent_2024-11-11 18-06-45",
    "AutomataGPT_2DBD_100_rules_inverse_epoch_20.pt",
)

state = torch.load(ckpt_path, map_location=device, weights_only=False)
model.load_state_dict(state)
model.to(device)
model.eval()
```

## One Inverse Inference Example

This example uses one sample from `CA_dataset_100r_2gpr.npy` and prompts the inverse model to predict a rules matrix from `(IC, GS2)`:

```python
import numpy as np
import torch
from train_model import tokenize_to_tensor, decode_from_tensor

sample = np.load("CA_datasets/CA_dataset_100r_2gpr.npy", allow_pickle=True)[0][0]

ic_str = " ".join(map(str, sample["tensor"][0].flatten().astype(int)))
gs2_str = " ".join(map(str, sample["tensor"][1].flatten().astype(int)))
prompt = f"[BOS] [BIC] {ic_str} [EIC] [BGS2] {gs2_str} [EGS2] [R]"

prompt_ids = tokenize_to_tensor(fast_tokenizer, [prompt]).to(device)
eos_tok = tokenize_to_tensor(fast_tokenizer, [fast_tokenizer.eos_token]).to(device)

with torch.no_grad():
    output_ids = model.generate(
        prompts=prompt_ids,
        seq_len=max_length,
        cache_kv=True,
        temperature=0,
        eos_token=eos_tok,
    )

decoded = decode_from_tensor(fast_tokenizer, output_ids[:1])[0]
predicted_rm = decoded.split("[R]", 1)[1].split("[EOS]", 1)[0].strip()

print(predicted_rm)
```

## Repository Layout Summary

```text
CA_datasets/
  CA_dataset_2r_10kgpr.npy
  CA_dataset_2r_100_val.npy
  CA_dataset_5r_100gpr.npy
  CA_dataset_10r_10kgpr.npy
  CA_dataset_10r_100_val.npy
  CA_dataset_100r_10kgpr.npy
  CA_dataset_100r_10_val.npy
  CA_dataset_100r_2gpr.npy

model_parameters/
  10_23_2024_2_Ruleset_Broad_Ent_2024-10-23 17-07-12/
  10_29_2024_10_Ruleset_Broad_Ent_2024-10-29 09-34-51/
  10_30_2024_2_Ruleset_Inverse_Broad_Ent_2024-10-30 10-58-04/
  10_30_2024_10_Ruleset_Inverse_Broad_Ent_2024-10-30 00-39-25/
  11_11_2024_100_Ruleset_Inverse_Broad_Ent_2024-11-11 18-06-45/
  11_14_2024_100_Ruleset_Broad_Ent_2024-11-15 01-49-23/
```

## Citation

If you use this repository, please cite:

```bibtex
@misc{berkovich2025automatagptforecastingrulesetinference,
      title={AutomataGPT: Forecasting and Ruleset Inference for Two-Dimensional Cellular Automata},
      author={Jaime A. Berkovich and Noah S. David and Markus J. Buehler},
      year={2025},
      eprint={2506.17333},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2506.17333},
}
```