Technical reference for nanocatalyst, a minimal JAX/Flax transformer for catalyst structure generation. This document covers the full pipeline from raw OC20 data to evaluated generation results.
OC20 S2EF Tokenizer Pre-tokenized
(extxyz+txt) (vocab=186) (parquet)
| | |
v v v
make_dataset.py ---> train_tokenizer.py ---> prepare_data.py
| |
| parquet (text) | parquet (input_ids, loss_mask)
v v
train.py
|
| orbax checkpoint
v
generate.py
|
| generated texts
v
benchmark.py
|
v
Evaluation metrics
| Step | Module | Description |
|---|---|---|
| 1. Build dataset | catalyst.data.make_dataset |
Read OC20 S2EF extxyz+txt pairs, serialize each structure to canonical text with condition block, compute energy bin edges via quantiles, write train/val/test parquet splits. |
| 2. Train tokenizer | catalyst.tokenizer.train_tokenizer |
Train a WordLevel tokenizer with 2-digit pair encoding on the text parquet. Produces tokenizer.json with vocab=186 and special tokens at indices 0-3. |
| 3. Pre-tokenize | catalyst.model.prepare_data |
Tokenize all text samples to fixed-length input_ids, compute segment_ids (condition vs structure vs padding) and loss_mask (1 on structure tokens, 0 elsewhere). Output as tokenized parquet. |
| 4. Train | catalyst.train |
Distributed training via jax.pmap across TPU chips. WSD learning rate schedule, AdamW optimizer, orbax checkpointing. |
| 5. Generate | catalyst.generate |
Autoregressive generation with temperature sampling, top-k filtering, and optional constrained decoding. Uses fixed-length padded buffer to avoid XLA recompilation. |
| 6. Evaluate | catalyst.eval.benchmark |
Full benchmark: parse structures, validate lattice/angles/distances, compute generation validity, uniqueness (MD5 fingerprint), and novelty against training set. |
Each training sample is a condition block followed by structure text, wrapped in special tokens:
<|bos|><|cond|>
task=relax
ads=OH
composition=CuPt3
elements=Cu H O Pt
target_bin=bin2
<|sep|>
11.6878 9.6132 31.4619 94.0859 90.0001 114.1111
Pt 0.0000 0.0000 10.5985
Cu 1.9480 3.3735 10.5985
...
<|eos|>
The structure body begins with the lattice line (a, b, c, alpha, beta, gamma in Angstrom/degrees), followed by one line per atom (element symbol + Cartesian x, y, z coordinates in .4f format). Atoms are sorted by (z, element, x, y) for deterministic ordering.
The loss_mask array controls which tokens contribute to the training loss:
<|bos|><|cond|> task=relax ads=OH ... <|sep|> 11.6878 9.6132 ... Pt 0.0000 ... <|eos|> [PAD] ...
|_______________ loss_mask = 0 _______________| |_________ loss_mask = 1 __________| |___ 0 ___|
condition tokens structure tokens padding
- 0 on all condition tokens (from
<|bos|>through<|sep|>inclusive) -- the model reads the condition but is not penalized for predicting it. - 1 on all structure tokens (after
<|sep|>through<|eos|>) -- the model learns to generate coordinates and element placements. - 0 on all padding tokens.
This ensures the model focuses capacity on learning the structure distribution conditioned on the catalyst specification.
Coordinates are tokenized as 2-digit chunks rather than individual digits. This produces approximately 40% fewer tokens per coordinate while keeping the vocabulary small.
Single-digit: 10.5985 --> 1 0 . 5 9 8 5 (7 tokens)
2-digit pair: 10.5985 --> 10 . 59 85 (4 tokens)
The tokenizer uses HuggingFace's WordLevel model with a pre_tokenizer.Sequence that:
- Isolates newlines as their own tokens.
- Splits on spaces (removed as delimiters).
- Separates digit runs from non-digit characters.
- Splits digit runs into 2-character chunks via
Split(r'\d{2}|\d', behavior='isolated').
Total vocabulary: 186 tokens (vs CatGPT's ~1,125).
The small vocabulary comes from encoding coordinates as digit pairs (00-99) plus a decimal point, rather than CatGPT's approach of one token per coordinate value (0.000 through 1.000 = 1,001 tokens).
| Index | Token | Purpose |
|---|---|---|
| 0 | <|bos|> |
Beginning of sequence |
| 1 | <|cond|> |
Start of condition block |
| 2 | <|sep|> |
Separator between condition and structure |
| 3 | <|eos|> |
End of sequence |
Two normalizer rules preprocess text before tokenization:
- Equals splitting:
=is surrounded by spaces so it becomes a separate token.task=relaxbecomestask = relax(3 tokens). - Element symbol detection: A space is inserted before uppercase letters that follow lowercase/digit characters.
CuPt3becomesCu Pt3, preventing element symbols from merging with adjacent text.
A single depth parameter controls the entire model:
depth ---+--- n_embd = depth x 64 (e.g., depth=8 --> 512)
+--- n_head = n_embd / 64 (head_dim = 64, fixed)
+--- n_layer = depth
+--- mlp_dim = n_embd x 2
params ~ 12 x n_layer x n_embd^2
For depth=8: n_embd=512, n_head=8, n_layer=8, mlp_dim=1024, approximately 25.2M parameters.
+------------------+
| input x |
+--------+---------+
|
+--------v---------+
| RMSNorm |
+--------+---------+
|
+--------v---------+
| Multi-Head Attn |
| - QK Norm |
| - RoPE |
| - Causal Mask |
+--------+---------+
|
+--------------v--------------+
| + residual (learnable lambda)|
+--------------+--------------+
|
+--------v---------+
| RMSNorm |
+--------+---------+
|
+--------v---------+
| ReLU^2 MLP |
| up -> relu^2 |
| -> down |
+--------+---------+
|
+--------------v--------------+
| + residual (learnable lambda)|
+--------------+--------------+
|
+--------v---------+
| output |
+------------------+
input_ids [B, T]
|
v
Embed(vocab=186, n_embd)
|
v
TransformerBlock x n_layer
|
v
RMSNorm
|
v
lm_head (weight-tied with Embed)
|
v
logit softcapping: cap * tanh(logits / cap), cap=15.0
|
v
logits [B, T, vocab_size]
RMSNorm: Root Mean Square normalization with learned scale, no bias. Used in place of LayerNorm throughout.
QK Normalization: Query and key vectors are independently normalized by RMSNorm before the attention dot product. Stabilizes attention scores across layers.
RoPE (Rotary Position Embeddings): Precomputed frequency table applied to Q and K after QK-norm. Enables relative position awareness without learned position embeddings.
Causal mask: Lower-triangular boolean mask prevents attending to future positions.
ReLU^2 MLP: Two linear layers with squared ReLU activation (relu(x)^2). Hidden dimension is 3 * n_embd to match SwiGLU parameter count (6 * n_embd^2 per block).
Logit softcapping: cap * tanh(logits / cap) with cap=15.0 (Gemma-style). Bounds logits to [-15, 15], preventing extreme values during training and generation.
Residual scalars: Per-layer learnable parameters resid_lambda (initialized to 1.0) and x0_lambda (initialized to 0.1). The residual connection becomes:
x = resid_lambda * x + sublayer_output + x0_lambda * x0
where x0 is the RMSNorm of the initial embedding, providing a skip connection to the input representation.
Weight tying: The language model head reuses the embedding matrix via embed.attend(x), reducing parameter count and coupling input/output representations.
Training uses jax.pmap to distribute across all available TPU/GPU chips. Gradients and losses are averaged across devices via jax.lax.pmean. The batch is reshaped from [B, T] to [n_devices, B // n_devices, T] before each step.
Warmup-Stable-Decay (WSD) schedule with three phases:
LR
^
| ___________________________
| / \
| / stable (70%) \ cosine decay (30%)
| / \
| / warmup \
| / (2%) \
+---+-----+---------------------------+-+---> steps
0 2% 70% 100%
- Warmup (2% of total steps): Linear ramp from 0 to
lr_max. - Stable (70% of remaining steps): Constant at
lr_max. - Cosine decay (30% of remaining steps): Cosine annealing from
lr_maxto 0.
AdamW with:
- Learning rate: 1e-3
- Weight decay: 0.1
- Gradient clipping: global norm clipped to 1.0
Orbax (orbax.checkpoint.PyTreeCheckpointer) saves model parameters at each epoch. The best checkpoint (lowest validation loss) is saved separately as ckpt_best. Parameters are unreplicated (first device slice) before saving.
Zero-copy parquet loading via PyArrow's list_flatten + numpy reshape. This avoids the ~400M Python object creation overhead of to_pylist() for large datasets.
Model: 25.2M parameters
Data: 173,665 train / 9,644 val
Seq len: 2048
Batch: 32 (8 per device x 4 TPU chips)
Epochs: 20 (108,540 steps)
Time: 96 minutes on TPU v5p-8
Cost: ~$27 on-demand equivalent (free via Google TRC)
Final loss: train=0.8908, val=0.9288
After each newline token in the structure body, the model should produce an element symbol (for the next atom line) or EOS (to end generation). Constrained decoding enforces this rule at the logit level:
State machine:
[generating lattice / coordinates]
|
| newline token & in structure body
v
[element-constrained state]
|
| mask logits: only allowed_element_ids + eos_id
v
[element token sampled] --> back to unconstrained generation
The constraint logic in constrained_logits:
- Create a mask filled with
-inf. - Set mask to
0.0for each allowed element token ID and the EOS token. - Add mask to logits, zeroing out the probability of all disallowed tokens.
Allowed elements are derived from the condition block: for composition=CuPt3, adsorbate=OH, the allowed set is {Cu, Pt, O, H}.
This offloads element validity from model capacity to decoding logic. The model does not need to learn which elements are valid for each composition -- the logit mask enforces it. This frees model capacity to focus on learning the coordinate distribution, contributing to data efficiency (174K structures vs CatGPT's 2M).
Generation uses a fixed-size numpy buffer of length max_gen_len:
buf = np.zeros(max_gen_len, dtype=np.int32) # fixed shape
buf[:len(prefix_ids)] = prefix_ids # write prefix
# ...each step writes buf[cur_len] = next_token
token_buf = jnp.array(buf[None, :]) # always [1, max_gen_len]The JIT-compiled forward function always receives [1, max_gen_len] input, so JAX/XLA compiles the computation graph exactly once. Without this, each unique sequence length would trigger a separate compilation, adding tens of seconds per sample.
Evaluation uses a 4-level validation hierarchy. A structure must pass all levels to be counted as valid.
The generated text (after <|sep|>, before <|eos|>) is fed to parse_canonical_string. This checks:
- At least 2 lines (lattice + atoms).
- Lattice line has exactly 6 numeric values.
- Each atom line has exactly 4 fields (element + 3 coordinates).
- All element symbols are recognized by ASE.
The set of elements in the parsed structure must be a subset of the allowed elements from the condition block. Any element not in the allowed set is a hallucination.
Condition: composition=CuPt3, adsorbate=OH
Allowed: {Cu, Pt, O, H}
Generated: {Cu, Pt, O, H} --> PASS
Generated: {Cu, Pt, O, H, Ni} --> FAIL (Ni hallucinated)
Four sub-checks:
- Lattice valid: a, b, c > 0.
- Angles valid: 0 < alpha, beta, gamma < 180 degrees.
- Volume valid: unit cell volume >= 0.1 A^3.
- Min distance valid: minimum interatomic distance >= 0.5 A (with periodic boundary conditions via
atoms.get_all_distances(mic=True)).
Generation Validity = L1 (parseable) AND L2 (element match) AND L3 (structural validity)
Among all valid structures, uniqueness is the fraction with distinct MD5 fingerprints. The fingerprint is computed from:
- Cell parameters rounded to 2 decimal places.
- Atom records (element, x, y, z) sorted by (symbol, z, x, y), rounded to 2 decimals.
- MD5 hash of the concatenated string.
Uniqueness = |unique fingerprints| / |valid structures|
Among unique valid structures, novelty is the fraction whose fingerprint does not appear in the training set.
Novelty = |unique fingerprints not in training set| / |unique fingerprints|
| CatGPT | nanocatalyst (Ours) | |
|---|---|---|
| Training data | 2M (OC20) | 174K (OC20) |
| Parameters | ~100M+ | 25.2M |
| Gen. Validity | 99.7% | 95.0% |
| Struct. Validity | 68.6% | 95.0% |
| Uniqueness | 20.8% | 100% |
| Bypass needed | Yes (CatGPT-BP) | No |
| Training cost | -- | ~$27 (96 min, TPU v5p-8) |
Key differences:
- CatGPT achieves higher generation validity (99.7%) but only 68.6% structural validity due to overlapping atoms. Their bypass method (CatGPT-BP) skips atoms placed closer than 0.5A, reaching 100% structural validity at the cost of incomplete structures.
- nanocatalyst achieves 95.0% on both metrics without any post-processing. Constrained decoding eliminates element hallucination, and the modern architecture (RMSNorm, RoPE, softcapping) reduces overlapping atom generation.
- CatGPT's 20.8% uniqueness (at T=1.0) vs 100% (at T=0.8) suggests our model produces a more diverse distribution despite lower temperature.
git clone https://github.com/everythingchalna/nanocatalyst.git
cd nanocatalyst
pip install -e ".[tpu]" # or pip install -e . for CPU/GPU
pip install -e ".[eval]" # pymatgen, matminer for evaluationwget https://dl.fbaipublicfiles.com/opencatalystproject/data/s2ef_train_200K.tar
tar -xf s2ef_train_200K.tar -C data/raw/python -m catalyst.data.make_dataset \
--src data/raw/s2ef_train_200K \
--dst data/processedpython -m catalyst.tokenizer.train_tokenizer \
--src data/processed/train.parquet \
--dst data/tokenizer_v3 \
--tokenizer-type wordlevel \
--digit-grouping pairpython -m catalyst.model.prepare_data \
--src data/processed/train.parquet \
--dst data/tokenized_v3/train.parquet \
--tokenizer data/tokenizer_v3 \
--seq-len 2048
python -m catalyst.model.prepare_data \
--src data/processed/val.parquet \
--dst data/tokenized_v3/val.parquet \
--tokenizer data/tokenizer_v3 \
--seq-len 2048python -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.0python -m catalyst.generate \
--checkpoint runs/v3_depth8/ckpt_best \
--tokenizer data/tokenizer_v3 \
--composition CuPt3 --adsorbate OH --target-bin bin2 \
--n-samples 100 --constrained --temperature 0.8python -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.8See the README for additional details on TPU setup, depth sweeps, and advanced configuration.