Scinetics-ai/all-atom-structure-tokenizer

All-Atom Structure Tokenizer

A PyTorch reference implementation of the tokenizer described in the Scinetics AtomLM-9b for Biomolecular Structure Prediction technical report. The encoder converts a heavy-atom assembly to one integer per chemical unit. The coordinate decoder reconstructs an assembly from those integers, its chemical inventory and Gaussian noise.

This release contains source code, an example and tests. Training starts from randomly initialized parameters. The repository does not distribute the authors' trained checkpoints or reproduce the report's benchmark scores. IDs produced by different checkpoints need not have the same structural meaning.

Run an example

From the repository root, install Python 3.10 or later and run:

python -m pip install -e ".[test]"
python -m pytest -q
python -m atomtoken train examples/peptide_ligand.json --profile tiny --steps 2 --checkpoint tiny.pt
python -m atomtoken encode examples/peptide_ligand.json --checkpoint tiny.pt --out tokens.json
python -m atomtoken decode examples/peptide_ligand.json --checkpoint tiny.pt --tokens tokens.json --out reconstructed.json

The example contains two peptide residues and three ligand heavy atoms. It yields five token IDs and twelve reconstructed coordinates in ångströms. This short run exercises the code path. Use the same checkpoint for both commands. To train the report-size architecture, select --profile report (the default), supply training records after train, and set --device cuda on a machine with CUDA PyTorch. The small trainer cycles through the input files with AdamW. It saves the model configuration and weights.

Architecture

Stage	Report configuration
Atom graph	All covalent bonds plus up to 32 intramolecular and 16 intermolecular spatial neighbors per receiver within 12 Å
Atom encoder	8 scalar/vector blocks, 256 scalar and 64 vector channels
Unit encoder	Attention pooling inside each unit, followed by 4 full-assembly attention layers at width 512
Invariant readout	Scalars, vector norms, inner products and ordered signed volumes, projected from 512 to 5
FSQ	Levels (8, 8, 8, 4, 4), straight-through training, mixed-radix IDs 0–8191
Coordinate decoder	12 unit transformer layers at width 768 and 4 equivariant atom blocks at scalar width 256
Sampling	Centered Gaussian noise, joint assembly coordinates, 50 Heun steps by default
Training	Proper-rotation augmentation, proper-rotation Kabsch coupling, masked flow matching, bond, angle, clash, chirality, distance and marginal usage terms

The first scalar has radix weight 1. The other weights are 8, 64, 512 and 2048. Decoding an ID first recovers five discrete scalar levels exactly. Conversion from these levels to atomic coordinates is learned and stochastic. In training, the decoder receives these same discrete levels with a straight-through gradient. Its public sampling interface accepts only chemistry and token IDs.

Molecular inputs

examples/peptide_ligand.json specifies the complete input schema. An assembly contains three ordered arrays.

• units. Each unit has a kind (protein, dna, rna, ligand, ion or modified), component name, chain, molecule and integer chain-relative position. Protein residues and nucleotides each form one unit. Each ligand heavy atom and ion forms its own unit. Use a shared molecule identifier for the atoms of a ligand.
• atoms. Each atom has a zero-based unit index, chemical element, and an atom-role name. Optional fields are charge (formal integer charge), aromatic (boolean), and xyz (three numbers in ångströms). An absent or null xyz marks an expected atom with missing coordinates. Preserve the expected atom inventory when decoding or training with missing data.
• bonds. List each covalent bond once as [atom_index_a, atom_index_b, bond_order]. Orders 1, 2 and 3 denote single, double and triple bonds. Use 4 for aromatic bonds. Include inter-chain covalent links.

Supply standardized heavy-atom records with a resolved conformer, the desired biological assembly and an explicit chemical graph. PDB/mmCIF biological assembly expansion, ligand bond perception and CCD standardization are upstream preprocessing steps. Preserve atom and unit ordering between encoding and decoding. The CLI stores a fingerprint of the chemical records alongside the IDs and verifies it before decoding. Removing all xyz fields preserves this fingerprint.

The encoder pools observed atoms. An entirely unobserved unit receives the sentinel -1 from encode. Supply a valid ID for this unit before sampling. During training, unobserved units receive random discrete conditioning and contribute no observed-coordinate targets. Decoder sampling uses the complete expected inventory.

Python API

import torch
from atomtoken import StructureTokenizer, load_assembly, tokenizer_loss

chem, xyz, observed = load_assembly("examples/peptide_ligand.json")
model = StructureTokenizer()  # Report configuration, new parameters.
loss, components = tokenizer_loss(model, chem, xyz, observed)
loss.backward()

model.eval()
ids = model.encode(chem, xyz, observed)
sample = model.decode(chem, ids, generator=torch.Generator().manual_seed(7))

Implementation details

The reference choices for attention head counts, stable hash buckets for chemical names, auxiliary loss coefficients and pair sampling are available in the code. Auxiliary distance and clash terms use bounded pair samples. The graph builder computes exact distances in query chunks. Atom search costs O(N²) time, and unit attention uses O(L²) memory for N atoms and L units. Use cropped assemblies when memory is limited. The CLI processes one assembly per optimization step.

The tests cover all 8,192 FSQ IDs, geometry under proper rotations and translations, the sign of local chiral volumes, edge types and caps, missing coordinates, gradient flow through the discrete bottleneck, and deterministic decoding with a fixed noise seed. The repository covers the tokenizer and its coordinate decoder. The AtomLM language-model backbone, ranker and benchmark runners are separate components.