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:

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.

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):

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

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:

@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},
}