Where Should Physics Enter a Molecular Crystal Generator?

Haocheng Tang; Junmei Wang; Wengong Jin

q-bio.BMSep 28, 2026

Where Should Physics Enter a Molecular Crystal Generator?

Authors: Haocheng Tang, Junmei Wang, Wengong Jin

Organizations: Khoury College of Computer Science Northeastern University Boston, MA 02115, USA · School of Pharmacy University of Pittsburgh Pittsburgh, PA 15213, USA

Abstract

Generative models make molecular crystal structure prediction fast, but their samples still exhibit geometric and packing violations. Physics can be introduced during training, post-training, or inference, yet these choices are rarely compared with the generator and physical signal held fixed. We introduce CrystAF, an all-atom crystal flow-map generation model, and use it with the UMA interatomic potential to systematically study where physics should enter. Post-training learns physical preferences directly into CrystAF, improving molecular validity and crystal packing while leaving sampling unchanged: physics is paid for once during training rather than repeatedly at deployment. In contrast, UMA relaxation is effective at repairing local clashes but makes generation 6--26 $\times$ slower, while learning from relaxed targets provides little benefit. These routes are complementary rather than competing. Physics-informed post-training first shifts the generated distribution toward more physically reasonable structures, after which inexpensive inference-time corrections further remove clashes and restore stereochemistry that the generator cannot represent. Importantly, the same post-training strategy also improves the multi-step all-atom Clari-M and rigid-body MolCrystalFlow generators, demonstrating transfer across architectures and representations. Together, our results suggest a simple principle: learn reusable physical alignment into the generator, and reserve inference-time physics for residual constraints that are better corrected than learned.

Figures & tables

	Cost			Quality
Method	UMA	ms	GPU-h	PB $\uparrow$	Clash $\downarrow$	Vol.Err $\downarrow$	PDD $\downarrow$	Stereo $\uparrow$
No physics
Clari-M	0	289	—	85.30 $\pm 0.80$	9.05 $\pm 0.74$	1.83 $\pm 0.13$	10.19 $\pm 0.36$	–
CrystAF-base	0	396	—	89.32 $\pm 0.67$	11.33 $\pm 0.84$	1.95 $\pm 0.09$	10.44 $\pm 0.16$	–
Training-time learning
Relax-and-distill	0	396	6	88.20 $\pm 0.68$	11.08 $\pm 0.83$	1.74 $\pm 0.09$	10.21 $\pm 0.18$	–

Table 1: Controlled comparison of physics-injection routes. Means

\pm

SE over 200 validation families

\times

20 samples. CrystAF uses 32 NFE; Clari-M uses 99 NFE. “UMA” denotes potential calls per sample at inference; “ms” is wall-clock time per sample on one L40S; “GPU-h” is training cost beyond CrystAF-base. PB and clash are percentages; Stereo reports R/S agreement only when stereochemistry is explicitly enforced. Bold marks the best value per column.

Signal ablations — matched CrystAF-base initialization
		Quality
Variant	Epoch	PB % $\uparrow$	Clash % $\downarrow$	Vol.Err $\downarrow$	PDD $\downarrow$
CrystAF-base (no post-training)	—	89.32 $\pm$ 0.67	11.33 $\pm$ 0.84	1.95 $\pm$ 0.09	10.44 $\pm$ 0.16
+ $E$ , $F$ , $\sigma$ , eligibility	4	91.21 $\pm$ 0.46	9.42 $\pm$ 0.75	1.66 $\pm$ 0.09	9.99 $\pm$ 0.15
+ $E$ , eligibility	4	89.01 $\pm$ 0.52	9.39 $\pm$ 0.75	1.55 $\pm$ 0.08	9.94 $\pm$ 0.11
+ Eligibility, volume	4	91.52 $\pm$ 0.46	9.82 $\pm$ 0.77	1.62 $\pm$ 0.09	9.98 $\pm$ 0.12

Table 2: Ablations of the physical signal and update rule. The first two blocks start from CrystAF-base and are read at epoch 4. The last block compares the UMA reward with a PoseBusters-ranked reward from a shared warm-started checkpoint (seed 929); second seeds are in Table S6 . Results use 200 families

\times

20 samples at 32 NFE. Epoch-6 results are in Table S11 .

	Rigid (50)	Flexible (50)	CSP5 (6)	CSP6 (5)	CSP7 (8)	Teaching (779)
Clari-M	$0.685{\pm}0.033$	$0.236{\pm}0.033$	$0.331{\pm}0.021$	$0.327{\pm}0.096$	$0.234{\pm}0.043$	$0.441{\pm}0.007$
CrystAF-base	$0.693{\pm}0.027$	$0.276{\pm}0.030$	$0.490{\pm}0.040$	$0.200{\pm}0.000$	$\mathbf{0.267}{\pm}0.094$	$0.456{\pm}0.007$
CrystAF-base + chain	$\mathbf{0.732}{\pm}0.026$	$0.308{\pm}0.032$	$0.431{\pm}0.087$	$0.455{\pm}0.136$	$0.240{\pm}0.035$	$\mathbf{0.486}{\pm}0.008$
CrystAF-UMA	$0.705{\pm}0.027$	$0.219{\pm}0.024$	$0.479{\pm}0.055$	$0.327{\pm}0.096$	$0.246{\pm}0.021$	$0.451{\pm}0.007$
CrystAF-UMA + chain	$0.708{\pm}0.028$	$\mathbf{0.310}{\pm}0.027$	$\mathbf{0.583}{\pm}0.100$	$\mathbf{0.519}{\pm}0.105$	$0.248{\pm}0.014$	$0.468{\pm}0.007$

Table 3: Structure recovery with

n_{s}=150

candidates and

k=30

lowest by UMA energy (

\mathrm{Sol}@k

; mean

\pm

SE). “+ chain” applies parity inversion, restrained MMFF, rigid declashing, and cell calibration. Clari-M uses 39 NFE; CrystAF uses 32 NFE. CSP subsets contain only 5–8 targets, so target-level uncertainty exceeds the reported SE. Bold : best per column.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

	protocol	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$
Clari-M	published, 1000 fam	87.34	9.56	1.59	9.56
	ours, 1000 fam	88.43 $\pm$ 0.33	8.57 $\pm$ 0.34	1.69 $\pm$ 0.04	9.82 $\pm$ 0.08
	ours, 200 fam	85.30 $\pm$ 0.80	9.05 $\pm$ 0.74	1.83 $\pm$ 0.13	10.19 $\pm$ 0.36
Clari-L	published, 1000 fam	85.89	7.69	1.50	9.28
	ours, 1000 fam	86.88 $\pm$ 0.39	6.92 $\pm$ 0.30	1.53 $\pm$ 0.04	9.43 $\pm$ 0.04
	ours, 200 fam	85.00 $\pm$ 0.86	7.75 $\pm$ 0.68	1.60 $\pm$ 0.09	9.77 $\pm$ 0.14

Appendix

Table S1: Clari-M and Clari-L results on the 1000-family and 200-family validation protocols. Both use 50 Heun steps, 20 samples per family, and identical metric code; “published” is the original paper’s Table 1.

signal term	PB pass	clash (cov. radii)	Vol.Err	EMD PDD	COMPACK
UMA energy, forces, stress	—	$\sim$	$\sim$	—	—
eligibility ( $d_{\min}\geq 0.8$ Å, density)	—	$\sim$	$\sim$	—	—
cell-volume term (optional)	—	—	✓	$\sim$	—

Appendix

Table S2: Overlap between the physical signal and the evaluation metrics. ✓: the metric is directly part of the signal;

\sim

: related but different quantity; —: not seen.

	NFE ( $\rho$ )	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$
CrystAF-base	8 (0.30)	83.67	30.73	2.46	11.54
	16 (0.75)	85.49	13.85	2.09	10.92
	32 (1)	89.32	11.33	1.95	10.44
	50 (1)	89.44	10.88	1.86	10.47
CrystAF-UMA	8 (0.30)	85.86	35.07	3.66	13.26
	16 (0.75)	83.96	11.95	2.17	10.71

Appendix

Table S3: Results for four sampling budgets with shared adapters. Weights are identical within each block; only

t_{i}=(i/N)^{\rho}

changes. Uniform 16-step sampling scores

77.65

and the tuned grid scores

85.49

; bold is best per column.

	CrystAF-UMA (released)	ablation arms	PB-rank reward
Init	base + 12 ep. PB-rank NFT	CrystAF-base	base + 12 ep. PB-rank NFT
LoRA rank / $\alpha$	$32$ / $64$	$32$ / $64$	$32$ / $64$
Trainable fraction	$3.77\%$	$3.77\%$	$3.77\%$
Rollout NFE / $\rho$	$16$ / $0.75$	$16$ / $0.75$	$16$ / $0.75$
Families per epoch	$20$	$20$	$20$
Ranks $\times$ candidates per family	$4\times 100$	$4\times 100$	$4\times 100$

Appendix

Table S4: Post-training configurations. The CrystAF-UMA column is the released model; the ablation column applies to every arm of Table 2 , which differ only as listed in Appendix A.24 .

240 families/epoch, group 150, 8 ranks
	families	epoch	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$
scaled up	480	2	92.92 $\pm$ 0.46	9.24 $\pm$ 0.75	2.07 $\pm$ 0.11	10.39 $\pm$ 0.18
scaled up	960	4	93.29 $\pm$ 0.36	8.73 $\pm$ 0.73	1.83 $\pm$ 0.10	10.17 $\pm$ 0.15
scaled up	1440	6	93.17 $\pm$ 0.40	8.33 $\pm$ 0.69	1.87 $\pm$ 0.10	10.26 $\pm$ 0.12
scaled up	1920	8	83.22 $\pm$ 0.71	22.83 $\pm$ 1.10	3.60 $\pm$ 0.14	12.21 $\pm$ 0.17
80 families/epoch, group 100, 4 ranks — the matched control

Appendix

Table S5: Scaling results for the physical-reward rollout at

\mathrm{NFE}{=}32

. Rows are aligned on families seen rather than epochs; the scaled-up run and its 80-family control are shown at matched family counts, with an additional scaled-up point at 1920 families. Bold is best per column.

	epoch	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$
CrystAF-UMA, seed 929	6	93.01 $\pm$ 0.43	7.98 $\pm$ 0.70	1.88 $\pm$ 0.10	10.25 $\pm$ 0.20
CrystAF-UMA, seed 2029	6	93.48 $\pm$ 0.36	8.60 $\pm$ 0.72	2.01 $\pm$ 0.08	10.42 $\pm$ 0.12
CrystAF-UMA, report grid	12	92.82 $\pm$ 0.42	8.72 $\pm$ 0.71	1.81 $\pm$ 0.10	10.14 $\pm$ 0.17
CrystAF-UMA, report grid	14	92.44 $\pm$ 0.44	9.02 $\pm$ 0.74	1.63 $\pm$ 0.09	10.03 $\pm$ 0.16
PB-rank reward, seed 929	12	93.21 $\pm$ 0.42	7.86 $\pm$ 0.67	1.61 $\pm$ 0.11	10.07 $\pm$ 0.23
PB-rank reward, seed 2029	12	92.56 $\pm$ 0.43	8.60 $\pm$ 0.71	1.88 $\pm$ 0.09	10.12 $\pm$ 0.14

Appendix

Table S6: Seed-to-seed variation at

\mathrm{NFE}{=}32

, 200 families

\times

20 samples. Both the released CrystAF-UMA model and the PoseBusters-ranked reward start from the same warm-started checkpoint (Appendix A.12 ); the report-grid rows roll out on the evaluation grid. PB-rank reward has

7.86

clash and

1.61

volume error for one seed, and

8.60

and

1.88

for the second.

CrystAF-UMA weights, no corrections
	NFE	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$	stereo % $\uparrow$
no corrections	32	93.01 $\pm$ 0.43	7.98 $\pm$ 0.70	1.88 $\pm$ 0.10	10.25 $\pm$ 0.20	48.38
no corrections	50	94.05 $\pm$ 0.34	8.13 $\pm$ 0.71	1.84 $\pm$ 0.12	10.24 $\pm$ 0.29	50.47
+ parity inversion, constrained MMFF, rigid declash
+ corrections	8	92.31 $\pm$ 0.41	7.32 $\pm$ 0.73	3.79 $\pm$ 0.19	13.85 $\pm$ 0.20	94.84
+ corrections	16	91.66 $\pm$ 0.42	3.13 $\pm$ 0.49	2.21 $\pm$ 0.13	11.27 $\pm$ 0.29	95.53

Appendix

Table S7: Results for sampling-time correction-chain variants. These rows extend Table 1 ; the chain overwrites rather than compounds. The CrystAF-base block applies the identical chain and cell calibration to the weights before reinforcement learning. Bold is best per column.

recipe	stereo %	PB %	clash %	Vol.Err	EMD PDD
none	49.91	85.49	13.85	2.09	10.92
parity inversion	85.19	85.44	19.10	$\approx$ 2.09	—
+ MMFF	95.08	91.71	23.20	2.02	11.05
+ MMFF + declash	95.24	92.16	2.27	2.06	11.03
same, $\mathrm{NFE}{=}50$	95.59	92.63	2.40	1.78	10.57
constraint projection	99.97	70.13	14.14	2.06	10.79

Appendix

Table S8: Results for sampling-time recipes on the distilled student CrystAF-base. Rows use the same weights,

\mathrm{NFE}{=}16

unless noted, and involve no training.

	valid fraction, epoch 1–6						KL, epoch 1–6
full signal	0.70	0.41	0.56	0.62	0.39	0.14	0.001	0.007	0.018	0.043	0.235	1.127
$E,F,\sigma$	0.69	0.39	0.51	0.56	0.33	0.35	0.001	0.011	0.040	0.178	0.940	1.888
$E$ only	0.69	0.37	0.46	0.41	0.05	0.29	0.001	0.018	0.089	0.696	1.973	2.397
feasibility only	0.68	0.37	0.54	0.57	0.29	0.20	0.001	0.010	0.035	0.145	0.985	2.034
eligibility gate only	0.70	0.39	0.54	0.62	—	—	0.001	0.007	0.023	0.057	—	—
DiffusionNFT on $U$	0.71	0.55	0.73	0.76	0.61	0.68	0.000	0.002	0.004	0.012	0.038	0.128

Appendix

Table S9: Rollout valid fraction and epoch-mean KL per epoch (rank 0), matched arms from CrystAF-base.

method	eligible	$\Delta E$ / molecule (eV)	median $F_{\max}$ (eV/Å)	median $\lVert\sigma\rVert$	median $\|\Delta V/V\|$
CrystAF-base	0.853	$+0.00$	8.59	0.058	0.0386
post-training, full signal	0.819	$+0.91$	11.07	0.079	0.0331
post-training, $E,F,\sigma$	0.779	$+1.60$	11.65	0.095	0.0313
post-training, gate only	0.836	$+0.77$	10.65	0.075	0.0359
post-training, feasibility only	0.805	$+1.42$	11.60	0.089	0.0315
DiffusionNFT on $U$	0.886	$-0.27$	8.80	0.043	0.0480

Appendix

Table S10: UMA single-point diagnostics of generated structures (200 families

\times

20, NFE 32; post-trained arms at epoch 4). Eligible: passes the

0.8

Å distance and density check. Median

|\Delta V/V|

is against the experimental cell.

Arm (epoch 6)	PB % $\uparrow$	clash % $\downarrow$	Vol.Err $\downarrow$	EMD PDD $\downarrow$
full signal: $E$ , $F$ , $\sigma$ , eligibility, cell volume	89.41 $\pm$ 0.50	10.60 $\pm$ 0.78	1.57 $\pm$ 0.09	9.94 $\pm$ 0.13
$E$ , $F$ , $\sigma$ , eligibility	87.91 $\pm$ 0.55	10.64 $\pm$ 0.81	1.65 $\pm$ 0.09	10.01 $\pm$ 0.11
$E$ , eligibility	87.47 $\pm$ 0.58	10.06 $\pm$ 0.79	1.63 $\pm$ 0.08	10.00 $\pm$ 0.11
feasibility only: eligibility, cell volume (no UMA)	88.17 $\pm$ 0.57	10.71 $\pm$ 0.81	1.68 $\pm$ 0.08	10.00 $\pm$ 0.10
DiffusionNFT on $U$ , full signal	93.07 $\pm$ 0.38	10.41 $\pm$ 0.82	2.90 $\pm$ 0.12	11.09 $\pm$ 0.16
relax-and-distill (training-time)	86.25 $\pm$ 0.76	12.29 $\pm$ 0.85	1.69 $\pm$ 0.10	10.16 $\pm$ 0.16

Appendix

Table S11: The matched arms of Table 2 read out at epoch 6. All of them except DiffusionNFT on

U

have passed the drift point flagged by the training-side signals (Table S9 ).

(a) Other generators
Generator (state)	Reward	Metric	base	$+$ post-training
MolCrystalFlow (rigid bodies)	UMA $E$ , $F$ , $\sigma$ , cell	Vol. RMAD % $\downarrow$	3.88 $\pm$ 0.08	3.16 $\pm$ 0.06
Clari-M (all atom, 99 NFE)	UMA $E$ , $F$ , $\sigma$ , eligibility	PB % $\uparrow$	88.43 $\pm$ 0.33	91.61 $\pm$ 0.28
(b) NFT vs. Flow-GRPO (MolCrystalFlow, Vol. RMAD % $\downarrow$ )
Stage	wall-clock, NFT / GRPO	NFT	Flow-GRPO	$\Delta$
2	0.97 h / 0.71 h	3.27 $\pm$ 0.06	3.53 $\pm$ 0.09	$+0.27$

Appendix

Table S12: The same post-training route on other generators, each on its native protocol. (a) UMA energy, force, and stress with a representation-specific feasibility term. Clari-M starts from the public checkpoint and is evaluated on 1000 validation families. (b) NFT and Flow-GRPO on MolCrystalFlow with matched reward, schedule, and stage budget.

	Rigid (50)	Flexible (50)	CSP5 (6)	CSP6 (5)	CSP7 (8)	Teaching (779)
Clari-M	$0.825{\pm}0.015$	$0.422{\pm}0.023$	$0.644{\pm}0.058$	$0.564{\pm}0.078$	$0.408{\pm}0.086$	$0.641{\pm}0.006$
CrystAF-base	$0.850{\pm}0.020$	$0.469{\pm}0.021$	$0.713{\pm}0.112$	$0.625{\pm}0.152$	$0.528{\pm}0.090$	$0.652{\pm}0.006$
CrystAF-base + chain	$0.859{\pm}0.017$	$0.499{\pm}0.028$	$0.872{\pm}0.117$	$0.697{\pm}0.121$	$0.595{\pm}0.112$	$0.667{\pm}0.006$
CrystAF-UMA	$0.844{\pm}0.022$	$0.490{\pm}0.026$	$0.812{\pm}0.055$	$0.625{\pm}0.152$	$0.437{\pm}0.074$	$0.635{\pm}0.006$
CrystAF-UMA + chain	$0.829{\pm}0.020$	$0.522{\pm}0.029$	$0.773{\pm}0.080$	$0.501{\pm}0.116$	$0.447{\pm}0.066$	$0.664{\pm}0.006$

Appendix

Table S13: Structure recovery with

n_{s}=400

k=200

(bootstrap mean

\pm

SE over samples). Same models, samplers, and COMPACK settings as Table 3 .

	Rigid (50)	Flexible (50)	CSP5 (6)	CSP6 (5)	CSP7 (8)	Teaching (779)
Clari-M	$0.670{\pm}0.026$	$0.231{\pm}0.030$	$0.333{\pm}0.003$	$0.326{\pm}0.097$	$0.199{\pm}0.065$	$0.397{\pm}0.007$
CrystAF-base	$0.643{\pm}0.032$	$0.232{\pm}0.031$	$0.578{\pm}0.103$	$0.200{\pm}0.000$	$0.189{\pm}0.072$	$0.423{\pm}0.008$
CrystAF-base + chain	$0.640{\pm}0.029$	$0.239{\pm}0.028$	$0.459{\pm}0.078$	$0.200{\pm}0.000$	$0.202{\pm}0.062$	$0.411{\pm}0.007$
CrystAF-UMA	$0.640{\pm}0.018$	$0.180{\pm}0.024$	$0.649{\pm}0.143$	$0.327{\pm}0.096$	$0.204{\pm}0.061$	$0.427{\pm}0.007$
CrystAF-UMA + chain	$0.624{\pm}0.028$	$0.262{\pm}0.033$	$0.477{\pm}0.059$	$0.200{\pm}0.000$	$0.229{\pm}0.048$	$0.423{\pm}0.007$

Appendix

Table S14: Structure recovery at

n_{s}=k=30

(bootstrap mean

\pm

SE over samples).

Explore similar work

CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models
May 9, 2026Chengqian Zhang, Yucheng Jin, Duo Zhang +2Crystal Interatomic Potentials
AI for Science Institute, Beijing, China · Center for Data Science, Peking University, Beijing, China · LMAM and School of Mathematical Sciences, Peking University, Beijing, China +3
uFlowCSP: Crystal Structure Prediction using Mean flow generative models
Sep 9, 2026Sourin Dey, Dipannoy Das Gupta, Lai Wei +2Crystal Structure Prediction Crystal Structure
Department of Computer Science and Engineering University of South Carolina Columbia, SC 29201 · Department of Computer Science and Cybersecurity University of North Georgia Columbia, SC 29201
Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction
Sep 30, 2026Thomas Egg, Harry Winston Sullivan, Maya M. Martirossyan +8Crystal Structure Prediction Molecular Generation
Center for Soft Matter Research, Department of Physics, New York University, New York 10003, USA · Simons Center for Computational Physical Chemistry, Department of Chemistry, New York University, New York 10003, USA · Department of Chemical Engineering and Materials Science, University of Minnesota, Minneapolis, MN 55455, USA +7

May 9, 2026, cond-mat.mtrl-sciJ/K move · Enter open · S save

CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models

Chengqian Zhang, Yucheng Jin, Duo Zhang, Tiejun Li+1

AI for Science Institute, Beijing, China · Center for Data Science, Peking University, Beijing, China · LMAM and School of Mathematical Sciences, Peking University, Beijing, China+3

Crystal Interatomic Potentials Molecular Generation Generative Models

Open paper AskPDF arXiv

Crystal generative models mainly learn what stable crystals look like, with little explicit supervision for what makes them stable. We reveal a substantial representation gap between state-of-the-art crystal generative models and pretrained universal machine learning interatomic potentials (MLIPs) via energy probing, and show this gap can be closed by a simple training-time alignment. We propose Crystal REPresentation Alignment (CrystalREPA), a plug-and-play framework that aligns the atom-wise hidden states of generative encoders with frozen MLIP representations through an element-aware contrastive objective, transferring stability-aware atomistic priors with marginal training overhead and no additional inference cost. Across three generative frameworks, ten MLIP teachers, and two benchmark datasets, CrystalREPA consistently improves the thermodynamic stability, structural validity, and structural fidelity of generated crystals. Equally important, we find that an MLIP's transfer effectiveness is poorly predicted by its accuracy on standard leaderboards (e.g., Matbench Discovery) but strongly predicted by the distinguishability of its atom-wise representation space, yielding a practical, accuracy-independent criterion for selecting MLIP teachers for generative transfer.

Abstract

Explore similar work

CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models

uFlowCSP: Crystal Structure Prediction using Mean flow generative models

Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction