Wyckoff Transformer: Generation of Symmetric Crystals [ICML 2025]

Installation (PyPI)

WyFormer is published on PyPI. Be mindful of your PyTorch situation, and install:

pip install wyckoff-transformer

Inference

The pre-trained models are published on HuggingFace. To use a HuggingFace model run:

wyformer-generate <output-file.json.gz> --hf-model <model-name>

Using a local model directory (must contain best_model_params.pt, config.yaml, and wyckoff_processor.json):

wyformer-generate <output-file.json.gz> --model-path runs/<run-id>

See this repository for a demo of standalone inference using the PyPI package as a library.

WyFormer Generated Datasets

If you just need the generated datasets for benchmarking, they are available at Figshare, including both original and DFT: WyFormer; WyFormer, DiffCSP(++), SymmCD, MiAD, WyCryst, CrystalFormer.

Abstract

Crystal symmetry plays a fundamental role in determining its physical, chemical, and electronic properties such as electrical and thermal conductivity, optical and polarization behavior, and mechanical strength. Almost all known crystalline materials have internal symmetry. However, this is often inadequately addressed by existing generative models, making the consistent generation of stable and symmetrically valid crystal structures a significant challenge. We introduce WyFormer, a generative model that directly tackles this by formally conditioning on space group symmetry. It achieves this by using Wyckoff positions as the basis for an elegant, compressed, and discrete structure representation. To model the distribution, we develop a permutation-invariant autoregressive model based on the Transformer encoder and an absence of positional encoding. Extensive experimentation demonstrates WyFormer's compelling combination of attributes: it achieves best-in-class symmetry-conditioned generation, incorporates a physics-motivated inductive bias, produces structures with competitive stability, predicts material properties with competitive accuracy even without atomic coordinates, and exhibits unparalleled inference speed.

Local development & training

Installation

1. Clone the repository
2. Run uv venv --python 3.12
3. Install the dependencies, including torch. There are several options:

• Manually install torch with your local flavour, e.g., uv pip install torch --index-url https://download.pytorch.org/whl/cu130, then run uv pip install -e
• Configure uv.toml with your desired indices, see scripts/platforms/zeus/uv.toml and scripts/platforms/cpu/uv.toml

Either route installs stock PyXtal from PyPI (pinned to 1.1.5). This project used to carry a PyXtal fork for a check_wp bug -- one species' like-like distance tolerance applied to every pair -- but our fix was adopted upstream, and the fork is gone. 4. wandb library is used extensively and must be installed. Logging can be disabled via WANDB_MODE=disabled. Otherwise, log into Wandb. Internally, we use WANDB_ENTITY=symmetry-advantage.

Running a pilot model

To verify that the installation is working, run a pilot model. Next token prediction:

wyformer-cache-dataset mp_20
python scripts/tokenise_a_dataset.py mp_20 yamls/tokenisers/mp_20_sg_multiplicity.yaml --new-tokenizer
python scripts/train.py yamls/models/NextToken/v6/base_sg.yaml mp_20 cuda --pilot

This will train a model, and save the results in the runs folder. The files are:

• best_model_params.pt - the model weights chosen by the validation loss
• config.yaml - the configuration used for training
• wyckoff_processor.json - tokenizers and preprocessing metadata (token engineers)
• generated_wp_no_calibration.json.gz - Wyckoff representation of the generated structures (if configured to evaluate generation)
• generated_wp_temperature_calibration.json.gz - Wyckoff representation of the generated structures, with the temperature calibration applied (if configured to evaluate generation)

Training Data Preprocessing

The available datasets correspond to the folders in data and cdvae/data. Dataset idetifiers are the folder names, they are used throught the project. Note that some of the folders are symlinks. Available datasets (in GitHub): alex_mp_20, mp_20, mp_20_biternary (binary and ternary structures from MP-20), mpts_52, carbon_24, perov_5. It is also possible to download and use matbench_discovery_mp_2022 notebook and matbench_discovery_mp_trj_full notebook.

For any data to be used for training, we need to do two preprocessing steps.

Compute and cache symmetry information

wyformer-cache-dataset <dataset-name>

This runs pyxtal over every structure and writes the Wyckoff representation to cache/<dataset-name>/{train,val,test}.parquet, one file per split -- see the data store, which also says how to convert a cache built before 2026-09-22. Every numeric, boolean and string column of the split CSVs is carried through as a per-structure label; --scalar-columns narrows that to a chosen few.

--max-sites N drops structures with more than N Wyckoff sites. Nothing truncates a structure: keeping its first N sites would silently change its composition while leaving its energy labels attached. Symmetry tolerance is --symmetry-precision / --symmetry-a-tol and is not set automatically; the datasets whose name includes a tolerance were built by passing it.

Tokenization

The tokenization script serves two purposes: it produces the mapping from the real data to token ids, and saves the resulting tensors. To produce a new tokenizer:

python scripts/tokenise_a_dataset.py <dataset-name> <path-to-tokenizer-yaml> --new-tokenizer

Tokenizer configs are stored in yamls/tokenisers. The processor is saved to cache/<dataset-names>/tokenisers/**.json, preserving the folder structure of the config.

Alternatively, you can use a cached tokeniser. This is important when a model that was trained on one dataset is applied to a different dataset.

python scripts/tokenise_a_dataset.py <dataset-name> <path-to-tokenizer-yaml> --tokenizer-path cache/<dataset-names>/tokenisers/<tokenizer-name>.json

Training

python scripts/train.py <path-to-model-yaml> <dataset-name> <device>

The model weights are saved to runs/<run-id>, and to WanDB, along with the processor metadata. See here for the list of configs. Adding --pilot will run the model for a small number of epochs.

Several GPUs of one node

torchrun --standalone --nproc-per-node <n-gpus> scripts/train.py <path-to-model-yaml> <dataset-name> cuda

train_batch_size stays the global batch, so a config trains the same run on any number of GPUs. See distributed training.

Resuming an interrupted run

Training writes runs/<run-id>/last_checkpoint.pt every optimisation.checkpoint_period epochs (defaulting to validation_period) and once more when the loop ends. It holds the weights, the optimiser and schedule state, the RNG and the loaders' shuffle position, so a run continued from it takes the same steps it would have taken had it never stopped -- unlike best_model_params.pt, which holds weights alone and restarts the optimiser cold.

python scripts/train.py <path-to-model-yaml> <dataset-name> <device> --resume <run-id>

This logs back into the same W&B run rather than opening a second one, and refuses to start if the config given differs from the one that run recorded, since the schedule and the optimiser are carried over from the checkpoint. A run that died before writing its first checkpoint has to be started over.

Preparing Representative Checkpoints

To train and prepare representative checkpoints for datasets like alex_mp_20 or mp_20, you can follow this end-to-end pipeline. Please replace <dataset-name> with your target dataset (e.g., alex_mp_20 or mp_20).

First, cache and tokenize the dataset:

wyformer-cache-dataset <dataset-name>
python scripts/tokenise_a_dataset.py <dataset-name> yamls/tokenisers/<dataset-name>_sg_multiplicity.yaml --new-tokenizer

Before training, ensure that your model configuration file points to the correct tokenizer. For example, in yamls/models/NextToken/v6/base_sg_schedule_free.yaml, update the tokenizer name to match your dataset:

tokeniser:
  name: <dataset-name>_sg_multiplicity

Then, initiate the training run:

python scripts/train.py yamls/models/NextToken/v6/base_sg_schedule_free.yaml <dataset-name> <device>

Generating structures

Wyckoff representations

Wyckoff representations are produced and stored in WanDB during model training. The wyformer-generate CLI (installed with the package) generates them from a trained model. Using a HuggingFace model:

wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name>

Using a W&B run:

wyformer-generate <output-file> --wandb-run <wandb-id> --use-cached-tensors

Using a local model directory (must contain best_model_params.pt, config.yaml, and wyckoff_processor.json):

wyformer-generate <output-file> --model-path runs/<run-id>

Note that the code does not automatically download Wandb artifacts, you need to do it manually, and place them in the runs folder when restoring via run ID.

To constrain generation to specific elements:

wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name> --required-elements Li-S --allowed-elements Li-S-P-O

To override the space group distribution from a cached dataset:

wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name> --sg-dist mp_20

3D Structures

There are two ways to generate 3D structures from Wyckoff representations: DiffCSP++ and CHGNet. They later can be relaxed with CHGNet and/or DFT.

DiffCSP++

Wyckoffs can be relaxed with modified DiffCSP++ code

CrySPR + MACE

CrySPR scheme using pyxtal and a MACE ML force field is integrated directly into the package. Install the optional extra first:

pip install "wyckoff-transformer[relax]"

Then run from the command line:

wyformer-cryspr WyckoffTransformer_mp_20.json \
    --model https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model \
    --output-dir results/ --start 0 --end 1000
# model_name defaults to the model file stem
head results/2023-12-10-mace-128-L0_energy_epoch-249_results.csv
model,id,formula,energy,energy_per_atom
...
2023-12-10-mace-128-L0_energy_epoch-249,35,H6O8Si2,-97.98,...

To process multiple Wyckoff genes in parallel, increase the worker count:

wyformer-cryspr WyckoffTransformer_mp_20.json \
  --model https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model \
  --output-dir results_parallel/ --start 0 --end 1000 --workers 8

URL-based models are downloaded once and cached in ~/.cache/wyckoff_transformer/mace_models/. A local path is accepted too: --model /path/to/model.model.

With --workers > 1, the CLI uses spawned worker processes. Each worker constructs its own MACE calculator, and the process environment is forced to single-threaded BLAS/OpenMP settings (OMP_NUM_THREADS=1, MKL_NUM_THREADS=1, OPENBLAS_NUM_THREADS=1) to avoid CPU oversubscription. This is usually the right mode for CPU-bound batch relaxation. On GPU, start conservatively and scale --workers only if the device has enough memory for multiple concurrent calculators.

Output layout is identical to the CHGNet variant below. Key options:

• --n-trials N — number of random PyXtal trials per structure (default 6)
• --fmax F — force convergence criterion in eV/Å (default 0.01)
• --model-name NAME — label for the results CSV (default: model file stem)
• --device auto|cpu|cuda — PyTorch device selection (default: auto)
• --workers W — number of parallel worker processes; each worker builds its own calculator (default 1)

Each PyXtal trial is relaxed in stages: first with a FixSymmetry constraint (a fix-cell warm-up, then cell + positions), then without it, and finally a rattle — positions and cell perturbed by a finite amount and relaxed again unconstrained, kept only if it wins 1 meV/atom. The rattle is what lets a structure leave a symmetric stationary point at all: the forces along symmetry-breaking modes vanish identically there, so gradient descent alone cannot, and the unconstrained stage took zero steps in 78% of trials. Per-trial CIFs are numbered by stage; ${full_formula}_kept.cif is the structure the trial kept, rattle.json records what the rattle did, and min_e_strc.cif links to the kept CIF of the lowest-energy trial.

CHGNet relaxation

The structures from all models can be optionally relaxed with CHGNet.

$ cp scripts/cryspr_chgnet.py mp_20/WyckoffLLM-naive/DiffCSP++/parsed_materials_10000_pyxtal.json_structures.json.gz /your/working/dir/
$ cd /your/working/dir/
$ gzip -d parsed_materials_10000_pyxtal.json_structures.json.gz
$ python ./cryspr_chgnet.py 0 -1 ./parsed_materials_10000_pyxtal.json_structures.json wylm-dcpp
$ head wylm-dcpp_id_formula_energy.csv
model,id,formula,energy
wylm-dcpp,0,Y4Ho4Ir8,-128.26178
wylm-dcpp,1,La4Cu4Si8,-89.13443
wylm-dcpp,2,Rb3Br3,-20.12497
wylm-dcpp,3,Gd1Ni2,-26.83080
wylm-dcpp,4,Pb3O18,-96.37995
wylm-dcpp,5,Tl2Au2S4Br4,-41.40260
wylm-dcpp,6,Ru1As2In3Ce3,-49.81073
wylm-dcpp,7,Pd4As4Ni4Se4,-73.34074
wylm-dcpp,8,Na4Lu4F16,-149.18192

• Output files are in the same format as above (CrySPR + CHGNet).
• ${reduced_formula}_${full_formula}_cell+pos.cif is the CHGNet relaxed structure.

DFT relaxation

We followed the Materials Project protocol, atomate2.vasp.flows.mp.MPGGADoubleRelaxStaticMaker. There isn't much to add, as the rest of the details of running DFT, unfortunately, depend on the HPC setup, and VASP is not open source. Here is the code to run at ASPIRE2.

DFT fixed-hull screening

wyformer-dft-screen ranks generated genes before an expensive DFT campaign. It combines the provenance-aware composition-floor ensemble with the gene attainable- energy critic, comparing both independently to the same immutable LeMat-Bulk PBE hull:

uv run wyformer-dft-screen generated/<run>/wyckoff_genes.json.gz \
    --formula-ensemble runs/formula_energy/ensemble.pt \
    --regressor-path runs/<gene-energy-run> \
    --top 1000 --out generated/<run>/dft_screen.csv

This command uses no MLIP energies and performs no relaxation or active learning; the shortlist feeds the external DFT workflow. See the DFT fixed-hull adversarial screening design for the objective, estimator boundaries, ablations, and reporting protocol.

Generative novelty screening

The energy screen trades novelty for stability, so it pairs with a second lever read off the generator itself. wyformer-gene-novelty scores each gene by -log p(gene) under the model that produced it -- a continuous novelty estimator that needs no reference set (AUC 0.93 against the fingerprint lookup, 0.81 against post-relaxation structure novelty):

uv run wyformer-gene-novelty generated/<run>/wyckoff_genes.json.gz \
    --model-path runs/<run> --condition energy_above_hull=0 \
    --permutation-samples 64 --device cuda \
    --out generated/<run>/gene_novelty.csv

Score the pool at the condition it was generated at. See generative novelty screening for what the density is, how the two levers combine, and what the combination is worth.

Ranking model variants (wyformer-protocol)

For comparing WyFormer variants during development, wyformer-protocol runs the whole funnel from sampled Wyckoff genes to MetaSUN in one command:

uv run wyformer-protocol generated/<run>/wyckoff_genes.json.gz \
    --output-dir generated/<run>/protocol \
    --cores 20            # or --devices cuda:0,cuda:1

Filters that need no potential — validity, uniqueness, novelty against LeMat-Bulk — run first, so only gene-novel structures are relaxed, and the relaxation itself uses one PyXtal trial and two stages. That is about 3.5× cheaper than the full CrySPR protocol at equal statistical power. The headline number is metasun_per_sampled_gene in <output-dir>/funnel.json.

--mlip is restricted to potentials LeMat-Bulk publishes a convex hull for (ORB by default), because e_above_hull is only meaningful when the energy and the hull come from the same model. Structure validity, BAWL fingerprinting and the hull energy are implemented in wyckoff_transformer.evaluation and pinned to LeMat-GenBench's own implementations by tests, so nothing imports that package at runtime.

See the de novo ranking protocol for the measurement rationale, the sample sizes it can resolve, required reference data, and known limitations, and every e_hull in this repository before comparing a hull energy from one part of the codebase with one from another — six things carry that name, and they differ in reference set, energy source and sign convention.

Generated Data Analysis

Storage

Public Figshare

Most analyzed datasets in a uniform format are available at Figshare.

Private Dropbox

The raw files are stored in a private Dropbox. To pull:

rclone copy "NUS_Dropbox:/Nikita Kazeev/Wyckoff Transformer data/generated.tar.gz" . --progress
tar -xvf generated.tar.gz

To push:

tar --use-compress-program=pigz -cvf generated.tar.gz generated
rclone copy generated.tar.gz "NUS_Dropbox:/Nikita Kazeev/Wyckoff Transformer data/" --progress

Tar is used to handle the large number of small files, and pigz is used to speed up the compression. The Dropbox folder is private, if you are a collaborator, please contact for access.

Preprocessing

In order to be analyzed the data must be preprocessed and cached. To preprocess all generated datasets in generated/datasets.yaml:

uv run python scripts/cache_generated_datasets.py

It supports filtering by dataset and transformations, e. g.:

uv run python scripts/cache_generated_datasets.py --dataset mp_20 --transformations DiffCSP++ DFT

Completing this step will enable loading the data with evaluation.generated_dataset.GeneratedDataset.from_cache

Metric computation

The ICML 2025 results were computed by the notebooks in ICML_eval. They include, but not limited to the following metrics:

1. S.U.N. - the fraction of stable, unique, and novel structures.
2. S.S.U.N. - the fraction of symmetric, stable, unique, and novel structures.
3. Space Group $\chi^2$ - the $\chi^2$ statistic of the space group distribution between the generated and the test set.
4. P1 - the fraction of generated structures that lack internal symmetries, i. e. belong to space group P1.
5. Property similarity and naive validity from Xie et al.