# Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction

arXiv: https://arxiv.org/abs/2609.39773 · HTML: https://mlsanity.com/paper/2609.39773/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 50 of its 89 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 30 of 50)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Crystal Diffusion Variational Autoencoder for Periodic Material Generation](https://mlsanity.com/paper/2110.06197/) (2021)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [SE(3) diffusion model with application to protein backbone generation](https://mlsanity.com/paper/2302.02277/) (2023)
- [Crystal Structure Prediction by Joint Equivariant Diffusion](https://mlsanity.com/paper/2309.04475/) (2023)
- [Riemannian Continuous Normalizing Flows](https://mlsanity.com/paper/2006.10605/) (2020)
- [A micro Lie theory for state estimation in robotics](https://mlsanity.com/paper/1812.01537/) (2018)
- [FlowMM: Generating Materials with Riemannian Flow Matching](https://mlsanity.com/paper/2406.04713/) (2024)
- [UMA: A Family of Universal Models for Atoms](https://mlsanity.com/paper/2506.23971/) (2025)
- [Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts](https://mlsanity.com/paper/2503.02819/) (2025)
- [Neural Manifold Ordinary Differential Equations](https://mlsanity.com/paper/2006.10254/) (2020)
- [Space Group Constrained Crystal Generation](https://mlsanity.com/paper/2402.03992/) (2024)
- [Normalizing Flows on Riemannian Manifolds](https://mlsanity.com/paper/1611.02304/) (2016)
- [Orb: A Fast, Scalable Neural Network Potential](https://mlsanity.com/paper/2410.22570/) (2024)
- [Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions](https://mlsanity.com/paper/2110.05064/) (2021)
- [Neural Ordinary Differential Equations on Manifolds](https://mlsanity.com/paper/2006.06663/) (2020)
- [Open Materials Generation with Stochastic Interpolants](https://mlsanity.com/paper/2502.02582/) (2025)
- [Guiding Generative Models to Uncover Diverse and Novel Crystals via Reinforcement Learning](https://mlsanity.com/paper/2511.07158/) (2025)
- [Equivariant Manifold Flows](https://mlsanity.com/paper/2107.08596/) (2021)
- [Deep Reinforcement Learning for Inverse Inorganic Materials Design](https://mlsanity.com/paper/2210.11931/) (2022)

## Abstract (arXiv)

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 textitvia 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 textitvia 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.