PackFlow: Generative Molecular Crystal Structure Prediction via Reinforcement Learning Alignment

Motivation

• Crystal packing determines material properties beyond molecular identity, including polymorph-dependent solubility and stability.CSP therefore seeks experimentally realizable packings from minimal chemical information.Source: §1 · p. 2

Computational bottleneck

• CSP remains difficult because candidate spaces combine lattice parameters, orientations, conformations, and symmetry, while accurate ranking requires costly quantum-level energetics.Hierarchical screening and repeated periodic geometry optimizations are used to manage this coupled generation-and-ranking bottleneck.Source: §1 · p. 2

Related work

• Generative approaches can learn packing distributions, but molecular crystals remain harder than inorganic crystals because weak non-covalent interactions and conformational flexibility enlarge the configuration space.OXtal demonstrates conditional molecular-packing generation but does not explicitly determine lattice parameters for direct relaxation and energy ranking.Source: §1 · p. 3

PackFlow framework

• PackFlow jointly predicts heavy-atom Cartesian coordinates and lattice parameters, enabling direct periodic relaxation and energy ranking within CSP pipelines.Its physics-alignment stage uses MLIP-derived energies and forces to steer proposals toward physically favorable regions.Source: §1 · p. 3

Evaluation and contribution

• PackFlow is evaluated against heuristic proposal engines on a broad unseen test set and two unseen CSP blind tests, with improved structural fidelity and lower-energy relaxed minima.The framework is positioned as a practical way to amortize the relax-and-rank bottleneck.Source: §1 · p. 3

2.1 PackFlow-enabled crystal structure prediction workflow

• PackFlow proposes heavy-atom coordinates and lattice parameters, followed by hydrogen-only relaxation, full crystal relaxation, and MLIP lattice-energy ranking.The generator is intentionally limited to candidate generation because OOD robustness is not guaranteed, making energetic refinement necessary.Source: §2.1 · p. 4

2.2.1 Problem setup and flow-matching objective

• The flow-matching model jointly transports noisy coordinates and lattice parameters back to crystal data by integrating learned ODEs from Gaussian noise.Coordinates and lattice use independently sampled flow times, and evaluation covers density, clashes, AMD, RDF Wasserstein distances, and generation time.Source: §2.2.1 · pp. 4, 6, §2.2.2 · p. 6, §2.2.3 · p. 8

2.2.3 Independent flow times for coordinates and lattice

• Independent coordinate and lattice flow times outperform shared times across most structural metrics, while removing bond information raises clash rate from 2.74% to 19.01%.Graphormer-style attention bias preserves global packing interactions while directly emphasizing covalent neighbors; unwrapped molecular representations avoid periodic bond discontinuities and were essential for effective training.Source: §2.2.3 · p. 8, §2.2.4 · p. 8, §2.2.5 · p. 9

2.3 Matching experimental statistics and symmetric attention

• PackFlow-Base reproduces experimental bond, angle, lattice, and packing distributions, while attention shifts from global context at high flow times to short-range refinement at low flow times.Attention also captures nearby atoms and symmetry-related intermolecular regions despite no explicit space-group enforcement.Source: §2.3 · p. 10

2.4.1 Alignment signal from heavy-atom energies and forces

• Physics alignment uses MLIP heavy-atom energies and force statistics with GRPO, within-context group comparisons, a single-time flow surrogate, and normalized advantage mixing.These choices avoid expensive trajectory-integrated likelihoods and manual reward-scale tuning while preserving geometric fidelity through KL regularization.Source: §2.4.1 · p. 10, §2.4.2 · p. 12, §2.4.3 · p. 13

2.4.4 Alignment improves physical validity and proximity to ground truths

• Physics alignment reduces clash rates from 2.53% to 1.53–1.71% for PackFlow-60M while improving RDF and/or AMD metrics at unchanged sampling cost.The hybrid λ = 0.5 setting provides a robust compromise, while smaller λ favors clash reduction and larger λ favors AMD proximity.Source: §2.4.4 · pp. 13–14, Table 2 · p. 14

2.5 Blind-test case studies show proximity to experimental polymorphs

• Before relaxation, PackFlow proposals were substantially closer in density to experimental structures than Genarris baselines, whose methods respectively under-compressed and over-compressed.Genarris methods could nevertheless have lower average pre-relaxation energies, but their density biases and structural distortions produced less favorable relaxed minima.Source: §2.5 · p. 15, Figure 5 · p. 17
• Within 100 samples, PackFlow’s best polymorphs were closer to experiment in density–relative-lattice-energy space, with relaxed lattice-energy differences of only a few kJ/mol.One case was slightly above the canonical ∼5 kJ/mol CSP target, while the other was well within it.Source: §2.5 · p. 15, Figure 5 · p. 17
• The case studies indicate that physically plausible packings with realistic densities and fewer short-range conflicts can improve relaxed outcomes despite occasionally higher initial energies.All evaluated structures were filtered to remove implausible molecular geometries, and relative energies used hydrogenated crystals referenced to the globally relaxed minimum-energy structure.Source: §2.5 · p. 15, Figure 5 · p. 17
• Physics alignment trains PackFlow with MLIP energy and force advantages, KL regularization, and advantage mixing, lowering test energies and forces while reducing atomic clashes.Varying the mixing parameter λ provides a smooth tradeoff between energy and force performance.Source: Figure 4 · p. 16
• After MLIP relaxation, PackFlow candidates reached lower-energy minima than Genarris candidates, while PackFlow-PA reached even lower minima than the base model.This reversed the pre-relaxation energy tendency, where Genarris methods could have lower average energies despite density distortions.Source: §2.5 · p. 15, Figure 5 · p. 17

PackFlow’s proposal framework

• PackFlow jointly samples heavy-atom coordinates and unit-cell lattice parameters, enabling direct integration with periodic relaxation and lattice-energy ranking.This proposal-stage design targets practical end-to-end CSP workflows rather than coordinate generation alone.Source: §3 · p. 18

Empirical performance

• PackFlow-Base reduces density error by up to 83% versus Genarris while sampling in approximately 0.1 s per structure, and its candidates relax closer to experimental polymorphs in blind tests.Relaxed lattice energies were within a few kJ/mol of the experimental polymorph in both case studies.Source: §3 · p. 18

Physics alignment

• Physics alignment lowers clash rate by up to an additional 40% versus PackFlow-Base at identical generation time and concentrates probability mass in physically favorable regions.The aligned models’ improved proposal distribution carries through relaxation in the blind-test analyses.Source: §3 · p. 18

Architectural choices

• Attention bias for covalent connectivity and independent coordinate–lattice time parameterization were important for reliable joint generation of heavy-atom crystal structures.These choices address unphysical structures and coupled denoising dynamics, respectively.Source: §3 · p. 18

Limitations

• The study is limited to homomolecular crystals, so extending PackFlow to co-crystals and solvates requires calibration across chemically heterogeneous components and solvent inclusion.The authors state that the modeling components are not inherently specific to single-component systems.Source: §3 · p. 18
• Physics alignment uses heavy-atom energies and force norms as proxies, which can reduce reward fidelity when hydrogens strongly affect contacts, electrostatics, or packing preferences.Full all-atom evaluation at every GRPO step was avoided because it would substantially increase training time and memory pressure.Source: §3 · pp. 18–19
• The relax-and-rank evaluation assumes low-energy MLIP basins approximate experimentally realized polymorphs, whereas kinetics, solvent pathways, temperature, and molecular flexibility can shift crystallization outcomes.The resulting rankings are thermodynamics-leaning within an experimentally biased proposal set rather than a complete kinetic model.Source: §3 · p. 19
• PackFlow was demonstrated at moderate scale—up to 60M parameters and approximately 80k crystals—leaving broader data, chemical coverage, and unit-cell diversity as scaling directions.The authors report that additional compute and data could improve geometric fidelity, clash rates, and low-energy relaxation outcomes.Source: §3 · p. 19

4.1 Architecture

• PackFlow tokenizes heavy atoms with coordinate, atom-descriptor, lattice-state, and time embeddings, then predicts coordinate and lattice vector fields with separate heads.Lattice context and flow-time embeddings are broadcast across atoms, while lattice prediction uses mean-pooled atom tokens.Source: §4.1.1 · p. 20, §4.1.3 · pp. 20–21

4.1.2 Encoding covalent connectivity

• Covalent connectivity is encoded either by additive attention biases along bond edges or by a Bond-GNN message-passing layer before transformer processing.These are mutually exclusive, ablation-compatible alternatives that preserve the remainder of the architecture.Source: §4.1.2 · p. 20

4.2 Physics alignment of flow matching via reinforcement learning

• Physics alignment post-trains PackFlow with GRPO using MLIP energy and force rewards, group-relative advantages, normalized multi-objective mixing, KL regularization, and a single-time score surrogate.Groups contain proposals from one crystal template so reward comparisons remain meaningful across polymorphs.Source: §4.2 · p. 21, §4.2.1 · p. 21, §4.2.2 · p. 21

4.2.2 Single-time surrogate score for importance-ratio and KL

• The single-time surrogate replaces expensive trajectory-wide likelihood calculations with one sampled time/noise evaluation for importance ratios and KL terms.It is an empirical proxy for the policy log likelihood rather than an exact endpoint likelihood, while retaining the within-group learning signal used by GRPO.Source: §4.2.2 · pp. 21–23

4.3 Evaluation metrics

• Evaluation measures density error, clash rate, AMD structural distance, short- and long-range RDF Wasserstein distances, and generation time for heavy-atom proposals.The RDF variants distinguish local from global packing statistics, and timing excludes hydrogen addition and downstream relaxation or ranking.Source: §4.3 · p. 24, §4.3.1 · p. 24, §4.3.2 · p. 24

4.4–4.6 Data, relaxation, energy, and baselines

• The CSD dataset uses family-disjoint splits and a 250-atom cap, while the CSP pipeline hydrogenates proposals, relaxes them with an MLIP, computes relative lattice energies, and compares against Genarris baselines.Genarris Plain uses geometric placement, whereas Rigid Press adds isotropic hard-sphere compression as a low-cost stability proxy.Source: §4.4.1 · p. 26, §4.4.2 · p. 26, §4.5.2 · p. 27

• A.S. led the research, designed PackFlow and its evaluation framework, and performed the evaluations and physics-based calculations.E.P. developed the post-training reinforcement-learning method, while A.S. and E.P. led result analysis.Source: §6 · p. 28
• R.G.-B. and A.S. conceived the project, with technical contributions from J.N., M.W., S.Q., and C.W.P.; senior authors supervised the work.Multiple authors wrote the initial manuscript, and all authors reviewed and edited the final version.Source: §6 · p. 28

• PackFlow is authored by Akshay Subramanian, Elton Pan, Juno Nam, Maurice Weiler, Shuhui Qu, Cheol Woo Park, Tommi S. Jaakkola, Elsa Olivetti, and Rafael Gomez-Bombarelli.Source: Supplementary Information for… · p. 33
• The authors are affiliated with MIT departments in materials science and engineering and electrical engineering and computer science, plus Samsung Display America Lab.Source: Supplementary Information for… · p. 33

Appendix A Validity of heavy-atom proxy rewards

• Physics alignment evaluates heavy-atom energies and forces before hydrogen addition and relaxation to reduce GRPO-update cost while retaining feedback on steric overlap, strain, and energetic favorability.The proxy is intended to approximate the within-group advantage signals produced by corresponding with-hydrogen evaluations.Source: Appendix A · p. 33, §A.1 · p. 33

A.1 Proxy versus with-H advantage signals

• The appendix compares group-standardized and clipped GRPO advantages from heavy-atom and hydrogenated structures across single- and multi-objective alignment settings.The comparison uses mixing weights λ ∈ {0, 0.5, 1}, matching the post-training modes considered in Fig. 4.Source: §A.1 · p. 33

• Advantage vector alignment measures cosine similarity between proxy- and with-H-based advantage vectors, while sign agreement measures the fraction sharing the same sign.Source: §A.2 · pp. 34–35
• High values indicate that the two reward schemes would reweight samples similarly and assign consistent good/bad labels within proposal groups.Source: §A.2 · p. 35

• Across 90 crystals with average group size K = 20, heavy-atom proxy rewards yield strongly aligned advantage signals with corresponding with-H advantages for all three post-training modes.Source: §A.3 · p. 35
• After identical GRPO standardization and clipping, heavy-atom energy and force rewards closely track with-H learning signals, supporting them as efficient physics-alignment proxies.Source: §A.3 · p. 36

• A covalent-connectivity graph and BFS traversal reconstruct each molecule by applying minimum-image lattice shifts to bonded atoms.Source: Appendix B · p. 36
• The procedure recenters molecules around the atom nearest the geometric centroid, producing deterministic coordinates invariant to input atom ordering.Source: Appendix B · pp. 36–37
• PackFlow-Base uses graph conditioning with RDKit features and attention bias, while Cartesian coordinates and lattice parameters are trained with flow matching.Source: Appendix B · p. 37

Appendix C Symmetries and gauge-fixing in the data representation

• Molecules are unwrapped across periodic boundaries and represented canonically, avoiding discontinuous bond vectors caused by wrapped coordinates.Source: Appendix C · pp. 37–38

C.1 Global translation

• Global translation is removed by mean-centering, so training and sampling occur in a centered coordinate frame.Source: §C.1 · p. 38

C.2 Rotation

• Because PackFlow is not rotationally invariant, lattices use a deterministic Pymatgen embedding with c parallel to z and a right-handed cell convention.Source: §C.2 · p. 39

Appendix D Genarris sampling and shortlisting strategy

• Genarris generates at least the target number of valid structures per space group, then shortlists exactly K proposals using space-group-stratified quotas and random deficit filling.Source: Appendix D · p. 39

• PackFlow models reach significantly lower AMD L∞ than Genarris baselines at faster rates across the entire test set.The comparison uses mean AMD L∞ scores averaged over the test set.Source: Appendix E · p. 40, Figure E.1 · p. 40

Appendix F Stability of post-training

• β = 0 causes significant instability in heavy-atom energies and forces during physics alignment.The instability appears as large training spikes, motivating a nonzero KL contribution.Source: Appendix F · p. 40, Figure F.1 · p. 41
• Larger group sizes produce more well-behaved physics-alignment training, whereas small groups lead to instability.The effect is shown through heavy-atom energy and force behavior during post-training.Source: Appendix F · p. 41, Figure F.1 · p. 41

Appendix G Visualization filters for Fig. 5

• The visualization filter removes relaxed structures with clearly unphysical intramolecular geometries without changing generation or relaxation.It is used only to produce clean density–energy summaries for Fig. 5.Source: Appendix G · p. 41, §G.1 · p. 42

G.1 Bond-length validity (covalent radii) filter

• Bond-length validity requires every template connectivity bond to fall within a covalent-radius-based tolerance under periodic boundary conditions.Distances use minimum-image geometry, and expected lengths derive from tabulated covalent radii.Source: §G.1 · pp. 41–42

I.1 RDKit atom and bond features

• PackFlow encodes molecular identity and connectivity as RDKit atom and bond graphs with discrete feature vectors.Atom and bond attributes are mapped to fixed-vocabulary integer indices, including a final miscellaneous bucket for unseen values.Source: §I.1 · p. 42, §I.1.1 · p. 42, §I.1.2 · p. 42

I.1.2 Atom features (featurize atoms)

• Atom features incorporate ring metadata and updated valence-related properties before extraction, while bond features include bond type and conjugation.These attributes are summarized in Tables I3 and I4.Source: §I.1.2 · p. 42, §I.1.3 · p. 42, Table I3 · p. 44

I.1.3 Bond features (featurize bond)

• Physics alignment shifts generated structures toward better clash and short-range RDF values, indicating improved handling of short-range unphysical contacts.The comparison is reported for base and physics-aligned models in the corresponding scatter plots.Source: Figure H.1 · p. 43