Where Should Physics Enter a Molecular Crystal Generator?
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 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 | Clash | Vol.Err | PDD | Stereo |
| No physics | ||||||||
| Clari-M | 0 | 289 | — | 85.30 | 9.05 | 1.83 | 10.19 | – |
| CrystAF-base | 0 | 396 | — | 89.32 | 11.33 | 1.95 | 10.44 | – |
| Training-time learning | ||||||||
| Relax-and-distill | 0 | 396 | 6 | 88.20 | 11.08 | 1.74 | 10.21 | – |
| Quality | |||||
| Variant | Epoch | PB % | Clash % | Vol.Err | PDD |
| Signal ablations — matched CrystAF-base initialization | |||||
| CrystAF-base (no post-training) | — | 89.32 0.67 | 11.33 0.84 | 1.95 0.09 | 10.44 0.16 |
| + , , , eligibility | 4 | 91.21 0.46 | 9.42 0.75 | 1.66 0.09 | 9.99 0.15 |
| + , eligibility | 4 | 89.01 0.52 | 9.39 0.75 | 1.55 0.08 | 9.94 0.11 |
| + Eligibility, volume | 4 | 91.52 0.46 | 9.82 0.77 | 1.62 0.09 | 9.98 0.12 |
| Rigid (50) | Flexible (50) | CSP5 (6) | CSP6 (5) | CSP7 (8) | Teaching (779) | |
|---|---|---|---|---|---|---|
| Clari-M | ||||||
| CrystAF-base | ||||||
| CrystAF-base + chain | ||||||
| CrystAF-UMA | ||||||
| CrystAF-UMA + chain |
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| protocol | PB % | clash % | Vol.Err | EMD PDD | |
|---|---|---|---|---|---|
| Clari-M | published, 1000 fam | 87.34 | 9.56 | 1.59 | 9.56 |
| ours, 1000 fam | 88.43 0.33 | 8.57 0.34 | 1.69 0.04 | 9.82 0.08 | |
| ours, 200 fam | 85.30 0.80 | 9.05 0.74 | 1.83 0.13 | 10.19 0.36 | |
| Clari-L | published, 1000 fam | 85.89 | 7.69 | 1.50 | 9.28 |
| ours, 1000 fam | 86.88 0.39 | 6.92 0.30 | 1.53 0.04 | 9.43 0.04 | |
| ours, 200 fam | 85.00 0.86 | 7.75 0.68 | 1.60 0.09 | 9.77 0.14 |
| signal term | PB pass | clash (cov. radii) | Vol.Err | EMD PDD | COMPACK |
| UMA energy, forces, stress | — | — | — | ||
| eligibility ( Å, density) | — | — | — | ||
| cell-volume term (optional) | — | — | ✓ | — |
| NFE ( ) | PB % | clash % | Vol.Err | EMD PDD | |
|---|---|---|---|---|---|
| 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 |
| 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 / | / | / | / |
| Trainable fraction | |||
| Rollout NFE / | / | / | / |
| Families per epoch | |||
| Ranks candidates per family |
| families | epoch | PB % | clash % | Vol.Err | EMD PDD | |
|---|---|---|---|---|---|---|
| 240 families/epoch, group 150, 8 ranks | ||||||
| scaled up | 480 | 2 | 92.92 0.46 | 9.24 0.75 | 2.07 0.11 | 10.39 0.18 |
| scaled up | 960 | 4 | 93.29 0.36 | 8.73 0.73 | 1.83 0.10 | 10.17 0.15 |
| scaled up | 1440 | 6 | 93.17 0.40 | 8.33 0.69 | 1.87 0.10 | 10.26 0.12 |
| scaled up | 1920 | 8 | 83.22 0.71 | 22.83 1.10 | 3.60 0.14 | 12.21 0.17 |
| 80 families/epoch, group 100, 4 ranks — the matched control | ||||||
| epoch | PB % | clash % | Vol.Err | EMD PDD | |
| CrystAF-UMA, seed 929 | 6 | 93.01 0.43 | 7.98 0.70 | 1.88 0.10 | 10.25 0.20 |
| CrystAF-UMA, seed 2029 | 6 | 93.48 0.36 | 8.60 0.72 | 2.01 0.08 | 10.42 0.12 |
| CrystAF-UMA, report grid | 12 | 92.82 0.42 | 8.72 0.71 | 1.81 0.10 | 10.14 0.17 |
| CrystAF-UMA, report grid | 14 | 92.44 0.44 | 9.02 0.74 | 1.63 0.09 | 10.03 0.16 |
| PB-rank reward, seed 929 | 12 | 93.21 0.42 | 7.86 0.67 | 1.61 0.11 | 10.07 0.23 |
| PB-rank reward, seed 2029 | 12 | 92.56 0.43 | 8.60 0.71 | 1.88 0.09 | 10.12 0.14 |
| NFE | PB % | clash % | Vol.Err | EMD PDD | stereo % | |
| CrystAF-UMA weights, no corrections | ||||||
| no corrections | 32 | 93.01 0.43 | 7.98 0.70 | 1.88 0.10 | 10.25 0.20 | 48.38 |
| no corrections | 50 | 94.05 0.34 | 8.13 0.71 | 1.84 0.12 | 10.24 0.29 | 50.47 |
| + parity inversion, constrained MMFF, rigid declash | ||||||
| + corrections | 8 | 92.31 0.41 | 7.32 0.73 | 3.79 0.19 | 13.85 0.20 | 94.84 |
| + corrections | 16 | 91.66 0.42 | 3.13 0.49 | 2.21 0.13 | 11.27 0.29 | 95.53 |
| 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 | 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, | 95.59 | 92.63 | 2.40 | 1.78 | 10.57 |
| constraint projection | 99.97 | 70.13 | 14.14 | 2.06 | 10.79 |
| 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 |
| 0.69 | 0.39 | 0.51 | 0.56 | 0.33 | 0.35 | 0.001 | 0.011 | 0.040 | 0.178 | 0.940 | 1.888 | |
| 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 | 0.71 | 0.55 | 0.73 | 0.76 | 0.61 | 0.68 | 0.000 | 0.002 | 0.004 | 0.012 | 0.038 | 0.128 |
| method | eligible | / molecule (eV) | median (eV/Å) | median | median |
|---|---|---|---|---|---|
| CrystAF-base | 0.853 | 8.59 | 0.058 | 0.0386 | |
| post-training, full signal | 0.819 | 11.07 | 0.079 | 0.0331 | |
| post-training, | 0.779 | 11.65 | 0.095 | 0.0313 | |
| post-training, gate only | 0.836 | 10.65 | 0.075 | 0.0359 | |
| post-training, feasibility only | 0.805 | 11.60 | 0.089 | 0.0315 | |
| DiffusionNFT on | 0.886 | 8.80 | 0.043 | 0.0480 |
| Arm (epoch 6) | PB % | clash % | Vol.Err | EMD PDD |
|---|---|---|---|---|
| full signal: , , , eligibility, cell volume | 89.41 0.50 | 10.60 0.78 | 1.57 0.09 | 9.94 0.13 |
| , , , eligibility | 87.91 0.55 | 10.64 0.81 | 1.65 0.09 | 10.01 0.11 |
| , eligibility | 87.47 0.58 | 10.06 0.79 | 1.63 0.08 | 10.00 0.11 |
| feasibility only: eligibility, cell volume (no UMA) | 88.17 0.57 | 10.71 0.81 | 1.68 0.08 | 10.00 0.10 |
| DiffusionNFT on , full signal | 93.07 0.38 | 10.41 0.82 | 2.90 0.12 | 11.09 0.16 |
| relax-and-distill (training-time) | 86.25 0.76 | 12.29 0.85 | 1.69 0.10 | 10.16 0.16 |
| (a) Other generators | ||||
| Generator (state) | Reward | Metric | base | post-training |
| MolCrystalFlow (rigid bodies) | UMA , , , cell | Vol. RMAD % | 3.88 0.08 | 3.16 0.06 |
| Clari-M (all atom, 99 NFE) | UMA , , , eligibility | PB % | 88.43 0.33 | 91.61 0.28 |
| (b) NFT vs. Flow-GRPO (MolCrystalFlow, Vol. RMAD % ) | ||||
| Stage | wall-clock, NFT / GRPO | NFT | Flow-GRPO | |
| 2 | 0.97 h / 0.71 h | 3.27 0.06 | 3.53 0.09 | |
| Rigid (50) | Flexible (50) | CSP5 (6) | CSP6 (5) | CSP7 (8) | Teaching (779) | |
|---|---|---|---|---|---|---|
| Clari-M | ||||||
| CrystAF-base | ||||||
| CrystAF-base + chain | ||||||
| CrystAF-UMA | ||||||
| CrystAF-UMA + chain |
| Rigid (50) | Flexible (50) | CSP5 (6) | CSP6 (5) | CSP7 (8) | Teaching (779) | |
|---|---|---|---|---|---|---|
| Clari-M | ||||||
| CrystAF-base | ||||||
| CrystAF-base + chain | ||||||
| CrystAF-UMA | ||||||
| CrystAF-UMA + chain |