Physics improves molecular crystal generation training and fixing
Where Should Physics Enter a Molecular Crystal Generator?
Artificial IntelligenceComputational Engineering, Finance, and Science
Summary
Predicting how molecules form crystals can be done by computer models, but their guesses often have problems with shapes and how molecules fit together. The authors made a new model called CrystAF and studied when to add physics knowledge: during training, after training, or when making predictions. They found teaching physics after training helps the model create more realistic molecules without slowing down later, and fixing problems during prediction can catch small errors. This combined approach works for different types of models and improves crystal predictions.
What this means in practice
- •For materials engineering teams: Improve molecular crystal design tools by integrating physics after training to generate more accurate crystal structures rapidly.
- •For computational chemistry developers: Enhance crystal generation software with post-training physics alignment and fast inference corrections to reduce molecular clashes and errors.
Authors
Haocheng Tang, Junmei Wang, Wengong Jin
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.