CrystAF — Crystal AnyFlow

Few-step, all-atom molecular crystal structure generation, post-trained with a machine-learned interatomic potential as the only reward.

CrystAF distills a 50-step Clari crystal generator into a dual-time flow map U(z, r, t) that jumps z += (t − r) · U(z, r, t), then post-trains that flow map with an NFT objective whose reward is UMA-OMC energy, max force, stress, clash and cell volume — no PoseBusters term, no validity term. One 16-LoRA adapter serves NFE 8 / 16 / 32 / 50: you change only the evaluation time grid, never the weights.

Code, environment setup, and every eval script: https://github.com/HaCTang/Crystal-NFT

Results — no stereochemistry correction

Plain model quality. 200 CSD validation families × 20 samples, summary.paper_bootstrap, L1 EMD PDD, interval flow-map sampler, no inference-time correctors and no cell calibration. Report grid is t_i = (i/N)^rho.

	NFE	rho	PB % ↑	clash % ↓	Vol.Err ↓	EMD PDD ↓
Clari-M backbone (Heun)	50	—	85.30	9.05	1.83	10.19
Clari-L backbone (Heun)	50	—	85.00	7.75	1.60	9.77
CrystAF distilled (cont3), uniform grid	16	—	77.65	14.71	2.07	10.59
CrystAF distilled, report grid	8	0.30	83.67	30.73	2.46	11.54
CrystAF distilled	16	0.75	85.49	13.85	2.09	10.92
CrystAF distilled	32	1	89.32	11.33	1.95	10.44
CrystAF distilled	50	1	89.44	10.88	1.86	10.47
+ RL, UMA reward (seed 929)	8	0.30	85.86	35.07	3.66	13.26
+ RL, UMA reward	16	0.75	83.96	11.95	2.17	10.71
+ RL, UMA reward	32	1	93.01	7.98	1.88	10.25
+ RL, UMA reward	50	1	94.05	8.13	1.84	10.24
+ RL, UMA reward (seed 2029)	32	1	93.48	8.60	2.01	10.42
+ RL, PoseBusters-ranked reward	32	1	93.21	7.86	1.61	10.07
+ RL, PoseBusters-ranked reward	50	1	93.84	7.37	1.70	10.22

A general-purpose physical potential matches a reward designed for the metric, on every column, in half the epochs, with no PoseBusters term in it. 94.05 at NFE 50 is the highest pass rate anywhere in this project.

Two things to read carefully:

• The RL weights specialise toward NFE ≥ 32. At 8 and 16 steps they are worse than the distilled model they came from. Quote the cont3 rows for few-step numbers; do not extrapolate the 32-step weights downward.
• Reproducibility. PB ≈ 93 reproduces across seeds for both rewards. The PoseBusters reward's 7.86 / 1.61 does not — its second seed bottoms out at 8.60 / 1.88, so its defensible summary is PB ≈ 93, clash 8.6–8.9, Vol.Err 1.7–2.0. The UMA reward's two seeds are 93.01 / 7.98 / 1.88 and 93.48 / 8.60 / 2.01. Both seeds of both rewards are published here so you can check this yourself.

Scaling the UMA recipe to 240 families per epoch and a group of 150 is a negative result: it peaks at 93.29 (epoch 4), the same plateau the 80-family recipe reaches after the same number of optimiser steps, and then collapses at epoch 8. The checkpoint is published for completeness.

Results — with stereochemistry correction

The correctors are training-free and isometric, applied at sampling time: CRYSTAF_MIRROR_FIX=body CRYSTAF_STEREO_REFLECT=1 CRYSTAF_MMFF=1 CRYSTAF_RELAX_CLASH=1, plus CRYSTAF_VOL_SCALE, a lattice-only cell calibration that moves no atom.

	NFE	PB % ↑	clash % ↓	Vol.Err ↓	EMD PDD ↓	stereo % ↑
UMA-RL weights, no correction	50	94.05	8.13	1.84	10.24	50.47
UMA-RL + correctors	8	92.31	7.32	3.79	13.85	94.84
UMA-RL + correctors	16	91.66	3.13	2.21	11.27	95.53
UMA-RL + correctors	32	93.43	2.80	1.93	10.78	95.09
UMA-RL + correctors	50	93.78	2.20	1.80	10.69	95.34
UMA-RL + correctors + VOL_SCALE=0.9850	32	93.42	2.55	1.58	10.50	95.43
UMA-RL + correctors + VOL_SCALE=0.9850	50	93.57	2.80	1.64	10.61	95.34
same chain on the PoseBusters-reward weights	32	93.58	2.12	1.51	10.47	95.20
CRYSTAF_PCFM=rs (max chirality, no relaxation)	16	70.13	14.14	2.06	10.79	99.97

Stereochemistry is not learned: the backbone's atom/bond features are identical for the two enantiomers, so the base model sits at chance on genuine R/S centres, and training-time conditioning is a measured negative result (chirality stayed at 49.21 while PB fell 11.3 points). The 95% comes entirely from the sampling-time correctors.

The chain overwrites rather than compounds — clash floors at 2.1–2.8 whichever reward is underneath — so a better policy buys improvement only in the columns the correctors do not target. Do not apply it at NFE 8: Vol.Err 3.79 / PDD 13.85 says the 8-step geometry is not good enough to survive the MMFF step.

Vol.Err 1.51–1.58 is a dispersion floor — two different calibration factors (0.9850, 0.9925) both land there, so the residual is spread, not bias. EMD PDD is the one column still short of Clari-L on the matched protocol (10.50–10.61 vs 9.77, i.e. 2.6–3.6 SE). The bare RL weights are much closer (10.24, 1.5 SE), so that cost belongs to the corrector chain rather than to the base model.

Files

File	Use
crystaf-uma-rl-seed929-epoch6.pt	The report checkpoint — UMA-reward RL. Reproduces the bold rows above (93.01 at NFE 32, 94.05 at NFE 50).
crystaf-uma-rl-seed2029-epoch6.pt	Second seed of the identical recipe — for the reproducibility claim.
crystaf-uma-rl-grid32-epoch12.pt	UMA reward rolled out on the NFE=32 report grid: most stable arm, best packing columns (Vol.Err 1.81, PDD 10.14).
crystaf-uma-rl-scaleup-epoch4.pt	Peak of the scaled-up run (240 families/epoch, group 150). A negative result; published for completeness.
crystaf-nft-mfpure-epoch12.pt	The PoseBusters-ranked reward, same architecture and twice the epochs.
crystaf-nft-seed2-epoch12.pt	Second seed of the PoseBusters-reward recipe — the reproducibility caveat above.
crystaf-cont3-step2000.pt	Distilled, pre-RL baseline.
clari-m-uma-nft-epoch1.pt, clari-l-uma-nft-epoch1.pt	The same UMA reward applied to the multi-step Clari generators (1000-family protocol: 88.43 → 91.61, 86.88 → 91.37).
molcrystalflow-uma-nft-phaseB-epoch12.pt	The same reward applied to the rigid-body MolCrystalFlow generator (lattice-volume RMAD 3.88 → 3.16).
teacher-rank800-pbnft-epoch1.pt	rank800 PB-NFT teacher (LoRA merged). Only needed to re-run distillation.

Each CrystAF file holds net_state_dict + ema_state_dict + meta; evaluate with the EMA weights. The Clari-M backbone is not redistributed here — fetch it from the-matter-lab/clari; you need it to build the DiT.

from huggingface_hub import hf_hub_download

student  = hf_hub_download("Haocheng1/CrystAF", "crystaf-uma-rl-seed929-epoch6.pt")
backbone = hf_hub_download("the-matter-lab/clari", "clari-med.ckpt")

Evaluation also needs CSD-derived tensors, which are CCDC-licensed and cannot be redistributed — build them yourself with scripts/build_clari_csd.sh (see doc/env.md).

Five things that silently produce wrong numbers

1. Sample with the interval flow map, not Clari's Heun sampler (MEANFLOW_SAMPLER_MODE=interval). Heun drops the second time argument.
2. Use the reported rho per NFE (t_i = (i/N)^rho): 0.30 at NFE 8, 0.75 at 16, 1 at 32/50. NFE 8 at rho=0.75 scores 48 instead of 83.67.
3. Report summary.paper_bootstrap.pb_score_pct, not raw mean_pb_score — invalid crystals naively score 1 and inflate it. Check mean_pb_valid ≈ 0.695.
4. Two stereo keys. mean_stereo_agreement_pct includes molecules with no R/S centre (79 here); mean_stereo_agreement_defined_pct counts only defined centres (95). Every stereo number above is the _defined one. Reading the other key produces an apparent 16-point regression that does not exist.
5. Do not compare across protocols. Every number here is 200 families × 20 samples. The same Clari-M checkpoint scores 85.30 on this subset and 88.43 on the published 1000-family set — about 3 PB points of protocol, which is larger than most effects being discussed. Measure your baseline on your own protocol first.

Judge changes on all four Table 1 columns, not PB alone. Repeat runs of an identical config vary by roughly ±0.5 PB / ±0.6 clash / ±0.3 PDD, and 60-family grids cannot rank configurations at all.

Caveat on the clash and volume columns

With the correctors on, clash and Vol.Err each have a corrector aimed directly at them, so they are no longer independent checks of packing quality — the clash relaxation optimises exactly what clash_rate measures. The bare-model table above is the one to read for packing quality. EMD PDD stays independent in both tables.

License

CC-BY-NC-4.0, inherited from the Clari model weights these derive from (Clari's code is MIT; its weights are CC-BY-NC-4.0). Non-commercial use only; please credit the-matter-lab/clari upstream.