nanocatalyst

A catalyst structure generation model you can train for under $20 (97 min on TPU v6e-8), or free via Google's TRC program. One flag — --depth 8 — configures everything: model size, learning rate, batch size, training horizon. Built on Karpathy's nanochat and JAX/Flax.

CatGPT (JACS 2024) proved that LLMs can generate novel catalyst structures, but reproducing it requires 2M training samples, 100M+ parameters, 30 epochs on an RTX 3090 ($200+, 26 days). nanocatalyst reaches 96% generation validity with 11x less data, 4x fewer parameters, and 10x lower cost ($18, 97 min), while achieving 100% uniqueness and novelty (vs. CatGPT's 20.8% and 9.1%).

Results

Model:    25.2M params (depth=8, n_embd=512, n_layer=8, n_head=8)
Data:     173,665 train / 9,644 val
Training: 20 epochs, 97 min on TPU v6e-8 (free via TRC)
Result:   val_loss=0.9518

Benchmark (CuPt3 + OH, T=0.8, top_k=40, 100 samples)

Check	Result	Note
Parseable	96/100
Element Match	96/100	Constrained decoding eliminates hallucination
Lattice Valid	96/100
Angles Valid	96/100
Volume Valid	96/100
Generation Validity	96/100 (96.0%)
Uniqueness	96/96 (100.0%)
Novelty	96/96 (100.0%)
Min Distance (≥ 0.5Å)	83/96 (86.5%)	Primary failure mode: overlapping atoms

Comparison

	CatGPT	CatGPT-BP	Ours
Training data	2M (OC20)	same	174K (OC20)
Parameters	~100M+	same	25.2M
Generation Validity	99.7%	99.7%	96.0%
Structural Validity	68.6%	100% (bypass)	96.0%
Uniqueness	—	20.8% (T=1.0)	100% (T=0.8)
Novelty	—	9.1% (T=1.0)	100% (T=0.8)
Bypass needed	No	Yes	No
Training cost	~$200 (RTX 3090, ~616 hrs)	same	~$18 (TPU v6e-8, 97 min)

Cost estimation: CatGPT training time estimated from paper (1.48 hr/1K steps × 416K steps ≈ 616 hrs ≈ 26 days). GPU cost based on RunPod RTX 3090 at ~$0.33/hr. TPU cost based on GCP on-demand v6e at $1.375/chip-hr × 8 chips. Our training was free via Google TRC.

Quick Start

Install

pip install -e .               # core (CPU/GPU)
pip install -e ".[tpu]"        # TPU backend
pip install -e ".[eval]"       # evaluation (pymatgen, matminer)

Data

# 1. Download OC20 S2EF (200K subset)
wget https://dl.fbaipublicfiles.com/opencatalystproject/data/s2ef_train_200K.tar
tar -xf s2ef_train_200K.tar -C data/raw/

# 2. Build dataset + tokenizer + pre-tokenize
python -m catalyst.data.make_dataset --src data/raw/s2ef_train_200K --dst data/processed
python -m catalyst.tokenizer.train_tokenizer --src data/processed/train.parquet --dst data/tokenizer_v3 --tokenizer-type wordlevel --digit-grouping pair
python -m catalyst.model.prepare_data --src data/processed/train.parquet --dst data/tokenized_v3/train.parquet --tokenizer data/tokenizer_v3 --seq-len 2048

Train

python -m catalyst.train --depth 8 \
    --data data/tokenized_v3 --tokenizer data/tokenizer_v3 \
    --epochs 20 --output-dir runs --run-name v3_depth8 \
    --mlp-type relu2 --logit-cap 15.0

Generate & Evaluate

# Generate
python -m catalyst.generate \
    --checkpoint runs/v3_depth8/ckpt_best --tokenizer data/tokenizer_v3 \
    --composition CuPt3 --adsorbate OH --n-samples 100 --temperature 0.8

# Evaluate
python -m catalyst.eval.benchmark \
    --checkpoint-dir runs/v3_depth8/ckpt_best --tokenizer data/tokenizer_v3 \
    --train-data data/tokenized_v3/train.parquet \
    --composition CuPt3 --adsorbate OH --n-samples 100 --temperature 0.8

Architecture

depth ──┬── n_embd  = depth x 64       (depth=8 → 512)
        ├── n_head  = n_embd / 64
        ├── n_layer = depth
        └── mlp_dim = n_embd x 2

Transformer block:
  x → RMSNorm → Attention (QK-norm, RoPE, causal) → +residual (λ)
    → RMSNorm → ReLU² MLP → +residual (λ)

Output: Embed → Blocks × n_layer → RMSNorm → lm_head (weight-tied) → softcap(15)

See ARCHITECTURE.md for full details (pipeline, tokenizer internals, training, constrained decoding, evaluation).

Project Structure

catalyst/
├── config.py              # Single --depth parameter
├── train.py               # pmap training, WSD schedule
├── generate.py            # Constrained decoding
├── loader.py              # Zero-copy parquet loading
├── sweep.py               # Depth sweep
├── hub.py                 # HuggingFace Hub download
├── model/
│   ├── transformer.py     # Flax transformer
│   └── prepare_data.py    # Pre-tokenize
├── data/
│   ├── serializer.py      # Structure ↔ text
│   └── make_dataset.py    # OC20 → parquet
├── tokenizer/
│   └── train_tokenizer.py # WordLevel (2-digit pair)
└── eval/
    ├── benchmark.py       # Validity, Uniqueness, Novelty
    ├── coverage.py        # Coverage Recall/Precision, EMD
    └── composition_coverage.py

Pretrained Checkpoints

pip install -e ".[hub]"

from catalyst.hub import download_checkpoint
ckpt_path = download_checkpoint("everythingchalna/nanocatalyst")

Documentation

• ARCHITECTURE.md — Full technical reference

Acknowledgments

Training compute was provided by the Google TPU Research Cloud (TRC) program.

Inspired by nanoGPT / nanochat (Karpathy) and CatGPT (Mok et al., JACS 2024).