PackFlow: Generative Molecular Crystal Structure Prediction via Reinforcement Learning Alignment

Abstract: Organic molecular crystals underpin technologies ranging from pharmaceuticals to organic electronics, yet predicting solid-state packing of molecules remains challenging because candidate generation is combinatorial and stability is only resolved after costly energy evaluations. Here we introduce PackFlow, a flow matching framework for molecular crystal structure prediction (CSP) that generates heavy-atom crystal proposals by jointly sampling Cartesian coordinates and unit-cell lattice parameters given a molecular graph. This lattice-aware generation interfaces directly with downstream relaxation and lattice-energy ranking, positioning PackFlow as a scalable proposal engine within standard CSP pipelines. To explicitly steer generation toward physically favourable regions, we propose physics alignment, a reinforcement learning post-training stage that uses machine-learned interatomic potential energies and forces as stability proxies. Physics alignment improves physical validity without altering inference-time sampling. We validate PackFlow's performance against heuristic baselines through two distinct evaluations. First, on a broad unseen set of molecular systems, we demonstrate superior candidate generation capability, with proposals exhibiting greater structural similarity to experimental polymorphs. Second, we assess the full end-to-end workflow on two unseen CSP blind-test case studies, including relaxation and lattice-energy analysis. In both settings, PackFlow outperforms heuristics-based methods by concentrating probability mass in low-energy basins, yielding candidates that relax into lower-energy minima and offering a practical route to amortize the relax-and-rank bottleneck.

• The paper demonstrates a novel two-stage generative modeling framework that combines flow-matching pre-training with RL-based physics alignment to predict crystal architectures.
• It integrates joint generation of lattice parameters and heavy-atom coordinates using a bond-aware transformer encoder, significantly reducing clash rates and sample requirements.
• Empirical results indicate improved energy proximity to experimental structures, with lower density errors and enhanced convergence compared to heuristic baselines.

PackFlow: Reinforcement Learning-Enhanced Generative Modeling for Molecular Crystal Structure Prediction

Introduction and Motivation

PackFlow introduces a novel generative modeling framework to address the long-standing bottleneck in molecular crystal structure prediction (CSP): generating physically plausible candidate packings that can reliably relax into experimentally accessible polymorphs. The central challenges of molecular CSP stem from the combinatorial search over crystal packing configurations and the necessity of quantum-level accuracy for polymorph ranking. While existing proposal engines, such as Genarris, rely on heuristic or combinatorial enumeration, these methods often demand considerable computational resources for downstream energy evaluations and relaxation.

PackFlow distinguishes itself by combining a joint generative model for unit cell lattice parameters and heavy-atom coordinates with a reinforcement learning (RL)-based post-training stage (physics alignment) that leverages machine-learned interatomic potentials (MLIPs). The RL post-training utilizes energetic and force-based feedback, directly steering the proposal distribution towards regions of configuration space that are more likely to yield low-energy, physically valid crystals upon relaxation.

Model Architecture and Training Paradigm

PackFlow is formulated as a conditional flow-matching generative model. Given an input molecular graph, PackFlow jointly samples:

• Cartesian coordinates of all heavy atoms within the asymmetric unit,
• Unit cell lattice parameters (cell lengths and angles).

This two-channel generative task enables the direct output of periodic structures ready for energetic evaluation and relaxation—unlike prior models such as OXtal that focus solely on atomic coordinates, omitting explicit lattice generation.

The model employs a transformer encoder architecture. Atom types, Cartesian coordinates, lattice parameters, and flow time variables are embedded as atom-centric tokens, with independent flow times for the coordinate and lattice channels. Covalent bonding information is incorporated as an additive attention bias, enabling the network to explicitly prioritize local chemical environments without restricting global packing interactions—a crucial adaptation because downstream packing is governed by weak non-covalent interactions coupled with conformational flexibility.

Figure 1: Architectural components of PackFlow, including joint embedding of coordinates, atom types, lattice parameters, and independent flow times for per-atom tokens. Covalent bonding is encoded as a pairwise attention bias.

A critical preprocessing step is the canonical unwrapping of molecules across periodic boundaries, which ensures smooth, contiguous representations for intramolecular geometry and avoids spurious artifacts from coordinate wrapping.

Training proceeds in two stages:

1. Flow-matching pre-training: The model is optimized to match the joint data distribution of atomic coordinates and lattice parameters using a regression loss on time-dependent vector fields.
2. Physics alignment post-training: A policy optimization regime, group relative policy optimization (GRPO), is introduced. At each RL step, groups of candidate packings for the same contouring molecular specification are generated and evaluated using MLIP-derived heavy-atom energies and force norms. Post-training is directed through advantage-based feedback on relative performance within each group, mixing energy- and force-derived objectives to target both global thermodynamic stability and local physical validity.

   Figure 2: PackFlow workflow. (a) Generative proposal, hydrogenation, H-relaxation, and full relaxation with MLIP. (c,d) Pre-training via flow-matching and post-training using GRPO objective with MLIP-derived proxies.

Evaluation Metrics and Ablation Analyses

PackFlow's performance is assessed with a comprehensive suite of metrics targeting disparate CSP failure modes:

• Density error: Quantifies deviation from experimental crystal densities, probing global unit cell plausibility.
• Clash rate: Evaluates the proportion of atoms in severe short-range overlap under periodic boundary conditions, acting as a fast local validity proxy.
• AMD $L_\infty$ distance: Measures the periodic average-minimum-distance descriptor between generated and ground truth structures, assessing proximity to experimental polymorphs.
• RDF Wasserstein distances (short- and long-range): Capture distributional similarities in packing statistics for both local and global intermolecular structure.

Ablation studies demonstrate the necessity of attention bias for covalent bonding incorporation (removal increases clash rates by nearly an order of magnitude) and the benefit of independent flow times for lattice and coordinate generation.

Empirical Results

Geometric and Physical Fidelity

PackFlow-Base matches experimental distributions of bond lengths, bond angles, lattice lengths, and lattice angles across the test set, confirming successful learning of both intramolecular and packing statistics.

Figure 3: PackFlow-generated distributional statistics match experimental ground truths for bond lengths, angles, and unit cell parameters. Attention maps demonstrate transition from global to local interaction focus during generation.

The physics alignment (PA) stage systematically reduces both the heavy-atom clash rate (from 2.53% to as low as 1.53%) and improves convergence towards low-energy configurations. The trade-off between energy and force objectives is smoothly controlled via a mixing parameter $\lambda$, permitting robust compromise solutions for practical CSP workflows.

Figure 4: Physics alignment post-training lowers heavy-atom energies/forces and reduces clashes via within-group policy improvement with KL-regularization.

Comparative Performance: Blind Test Case Studies

On two CSP blind test examples (one rigid and one flexible molecule), PackFlow proposals are notably closer to experimental densities pre-relaxation and, after MLIP relaxation, reach minima with lower energies and better agreement in density-energy space with experimentally determined polymorphs than Genarris baselines. In both cases, PackFlow’s predictions are competitive with the known target structures, with relaxed lattice energy differences within a few kJ/mol of experiment—a meaningful threshold for CSP success.

Figure 5: Comparative analysis of PackFlow and Genarris proposals and relaxed outcomes on CSP blind test examples. PackFlow achieves superior density and energy proximity to experiment, both before and after relaxation.

Sample efficiency analysis reveals that PackFlow achieves lower AMD distances to ground truth with fewer generated candidates than heuristic baselines, amplifying its value in resource-constrained CSP pipelines.

Figure 6: PackFlow exhibits superior sample efficiency compared to Genarris for approaching experimental structures.

A further analysis reveals that physics alignment specifically resolves short-range unphysical contacts, as evidenced by shifts in clash rate vs. RDF (short) scatter plots.

Figure 7: Physics alignment post-training preferentially resolves short-range packing pathologies, improving local structure.

Theoretical and Practical Implications

PackFlow’s architecture and training regime signal multiple advancements to generative modeling under physical constraints:

• Joint lattice-coordination generation closes a key gap for generative models in periodic materials CSP, ensuring proposals are directly compatible with periodic energetic evaluation and relaxation.
• Bond-aware attention establishes an effective, low-overhead mechanism for encoding chemical constraints without sacrificing global modeling capacity—critical for organic molecular crystals where both local chemistry and long-range packing are important.
• Physics alignment via RL enables explicit distribution steering towards low-energy, physically plausible configurations, amortizing a portion of the computational burden traditionally associated with high-cost relax-and-rank CSP.
• Advantage mixing in RL post-training obviates the need for manual reward scaling, allowing flexible targeting of energetic or force objectives.
• Proxy-based rewards based on heavy-atom-only energies/forces are demonstrated to be tightly aligned (cosine similarity 0.79–0.91) with full hydrogenated structure metrics, enabling efficient RL without the need for all-atom relaxation during training.

Limitations and Future Directions

Current PackFlow models are limited to homomolecular crystals and moderate unit cell sizes, in part due to $O(N^2)$ attention bias memory scaling. Extending to binary/multicomponent crystals and longer sequences will require architectural modifications to maintain computational tractability. Additionally, physics alignment uses heavy-atom proxy rewards, and would benefit from multi-fidelity training incorporating periodic hydrogenated energies in at least a subset of updates.

In terms of CSP workflow, the generative model—while vastly improving proposal quality—remains reliant on downstream relaxation for out-of-distribution generalization and full energetic ranking. Integrating kinetic accessibility (rather than purely thermodynamic probability mass) with generative models could be a key avenue for progress.

Model scaling and incorporation of FlashAttention or block-sparse attention may further improve generalization, geometric fidelity, and reduce the need for large batch relaxation campaigns.

Conclusion

PackFlow represents a substantial step forward in generative modeling for molecular CSP by integrating physically grounded architecture, explicit chemical priors, and reinforcement learning-based physics alignment. It provides a scalable, physically aligned proposal engine that outperforms heuristic baselines in candidate plausibility, relaxation outcomes, and sample efficiency while preserving computational tractability.

The framework establishes a generalizable blueprint for coupling generative models with post-hoc preference alignment using proxy evaluations, suggesting broad applicability to other structured generative tasks in materials science and chemistry. The demonstrated efficacy on CSP benchmarks and the modularity of the RL alignment approach position PackFlow as a candidate for integration into large-scale, autonomous polymorph screening pipelines and for extension to more complex, multi-molecular crystal environments.