Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction

2026-09-30

AI Summary

This study addresses the generative modeling challenges in molecular crystal structure prediction arising from polymorphism, large unit cells, and complex packing arrangements by proposing CG-OMatG, an equivariant Riemannian flow model. The method introduces a coarse-grained hierarchical representation coupled with a rigid-body message-passing mechanism to jointly reconstruct molecular centroids, orientations, and lattice parameters. Furthermore, it pioneers the application of policy gradient reinforcement learning for fine-tuning flow models, effectively guiding the generation of low-energy stable structures. Experimental evaluations demonstrate that CG-OMatG achieves superior performance on crystal structure prediction blind-test benchmarks, yielding generated structures in strong agreement with experimental data. Consequently, this work significantly accelerates polymorph screening and the discovery of organic solid-state materials.

Abstract

Crystal structure governs material properties, making crystal structure prediction (CSP) a fundamental problem in materials science. Generative models are a promising approach for solving this problem, but the prevalence of polymorphism, coupled with large unit cells and complex packing geometry, makes the molecular CSP task challenging for existing models. To address this, we introduce Coarse-Grained Open Materials Generation (CG-OMatG), an equivariant Riemannian flow-based generative model. CG-OMatG predicts molecular crystal structures \textit{via} a coarse-grained, hierarchical representation. CG-OMatG treats molecules as rigid bodies---performing both inter- and intra-molecular message passing to construct a geometric representation for molecular packings---and learns to reconstruct molecule centroid positions, orientations, and lattice parameters, conditioned on chemical species and conformer geometry. We train the model on subsets of the Open Molecular Crystals (OMC25) and Cambridge Structural Database (CSD) datasets. Further, we fine-tune the model \textit{via} policy gradient reinforcement learning to steer the model towards generating low-energy candidate structures. We validate the generated structures on the CSP blind test benchmark, assessing agreement with experimentally determined crystals using COMPACK packing-similarity analysis. CG-OMatG exhibits strong performance for generative molecular crystal structure prediction, paving the way for accelerated polymorph screening and organic solid-state materials discovery.

Problem

• Crystal Structure Prediction
• Molecular Crystals
• Polymorphism
• Generative Models

Innovation

• Riemannian flow models
• Reinforcement learning
• Crystal structure prediction
• Coarse-grained representation
• Equivariant generative model