Fast Organic Crystal Structure Prediction with Unit Cell Flow Matching

Abstract: Organic crystal structure prediction (CSP) is a requirement for computational modelling of organic solids, but traditionally costs several CPU-years per molecule. Generative models such as OXtal dramatically reduce this cost by sampling stable organic crystal structures directly. However, OXtal forgoes explicit lattice parametrization in favour of modelling large crops of the bulk material with expensive triangle layers, which can incur a computational cost of minutes per molecule. In this paper, we reduce this to seconds with Clari, a large-scale flow matching model that generates redundancy-free unit cells and replaces triangle layers with pure pair-bias attention. Clari requires only atom types and bonds as input and does not need an RDKit-sanitizable input molecule, which expands its applicability to challenging chemistries such as fullerenes, metal complexes, and atom clusters. We further ablate key design choices such as auxiliary losses, timestep distributions, noise priors, and self-conditioning. On OXtal's test sets, we surpass OXtal's solve rate while obtaining a speedup of $15$-$30\times$. Because Clari also models explicit hydrogens, it supports inference-time scaling via direct energy ranking, without any decoration or relaxation step. When generating 150 crystals and selecting the top-30 by energy, we further improve solve rate while maintaining a speedup of $5$-$8\times$. We also introduce the CSD Teaching Subset as a new test split of diverse and complex molecules for future benchmarking. Our contributions enable CSP within seconds, making large-scale virtual screening of organic solids practical. Code is available at https://github.com/aspuru-guzik-group/clari.

• The paper introduces Clari, a flow-matching model that jointly generates lattice vectors and atomic coordinates for a single organic-crystal unit cell from a 2D molecular graph.
• Clari-L achieves 15–30× faster sampling than OXtal, generating 150 structures in 2.2 seconds on average and reaching higher solve rates across CSP Blind Tests.
• The paper shows that data-informed lattice priors, self-conditioning, auxiliary packing losses, and simplified pair-bias attention reduce clashes and improve crystal reconstruction without requiring relaxation or DFT polishing.

Overview

This paper introduces Clari, a flow-matching generative model for organic crystal structure prediction (CSP) that generates a single unit cell — atom coordinates and lattice vectors jointly — directly from a 2D molecular graph. The work targets the two principal inefficiencies of OXtal, the prior state-of-the-art generative CSP model: its bulk representation, which processes multiple symmetry-related molecular copies simultaneously, and its reliance on triangle-update layers inherited from AlphaFold3. By replacing bulk crops with redundancy-free unit cells and triangle operations with pure pair-bias attention, Clari achieves a $15$–$30\times$ sampling speedup over OXtal on the CSP Blind Tests while surpassing its solve rate, reducing per-molecule generation from minutes to seconds (2.2 s for 150 structures on average; 6.0 s including UMA energy ranking).

Method

Clari represents a crystal as a tuple of lattice matrix $\mathbf{L}$, Cartesian coordinates with zeroed centroid, atom features, and a bond adjacency matrix. The lattice's three primitive vectors are treated as three additional virtual tokens concatenated into a unified $(3+n)\times 3$ object, so lattice and atom tokens flow through the network identically without special pooling. The velocity field is trained with a linear-interpolant conditional flow matching objective, split into equal-weight lattice and coordinate mean-squared-error terms.

Several design choices distinguish the training recipe:

• Data-informed source prior: rather than a unit normal $p_0$, which concentrates mass at vanishing densities and produces degenerate lattices, the authors fit independent Gaussians to three decomposed lattice components — atom density, cell angles, and sorted volume-normalized lengths — closely matching CSD marginals.
• Auxiliary losses: a relative volume error loss and a pairwise periodic-distance loss (with an explicit steric-clash penalty) applied to the one-step endpoint estimate.
• Self-conditioning: recycling the previous timestep's endpoint estimate as input costs ~20% extra training time but yields consistent gains.
• Approximate optimal-transport coupling: samples are aligned across signed lattice permutations (brute force), atomic permutations (Hungarian algorithm over cached component isomorphisms), and rotations (weighted Kabsch), straightening trajectories and acting implicitly as augmentation.

The architecture is a DiT-style Transformer with gated attention, QKNorm, SwiGLU, AdaLN-Zero conditioning on timestep, lattice, and formulae, and additive pair-bias attention built from bond, topological, Cartesian, and periodic distance features. Notably, it is not equivariant by construction but learns invariances through augmentation at scale, following recent trends in biomolecular modelling. Two models are trained: Clari-M (88M parameters, 14 h on 4 H100s) and Clari-L (173M parameters, 17 h on 8 H100s), both for 150k steps on a filtered CSD corpus (~917k training crystals). Unlike prior work, hydrogens are retained and RDKit sanitization is not required, broadening applicability to fullerenes, metal complexes, boranes, and atom clusters.

Ablations

A staged ablation isolates each contribution. Starting from a base DiT (25.0% clash rate), adding lattice tokens is roughly quality-neutral but stabilizes optimization; auxiliary losses and self-conditioning reduce clash rate to 9.79% and improve all reconstruction metrics. Replacing the fitted lattice prior with a standard normal degrades volume error from 1.55% to 2.88%, confirming the importance of the data-informed source distribution. Scaling to Clari-L further improves all metrics (7.69% clash rate, 1.50% volume error, EMD PDD 9.28).

Timestep ablations yield a domain-specific finding that contradicts practice in protein and small-molecule generation: log discretization is uniformly harmful here, and mid-range-weighted distributions (logit-normal) improve packing metrics but degrade PoseBusters validity, sometimes catastrophically (31.5% when combined with log discretization). The authors interpret this as evidence that global packing geometry is established at intermediate flow times, so over-emphasizing late timesteps starves these stages. The final models use Beta(1.8, 1) with linear discretization as a compromise, though the authors concede the trade-off space remains underexplored.

Benchmark results

Evaluation uses COMPACK-based Sol@$k$: a target is solved if any of $k$ candidates matches the ground truth in ≥8/15 molecules with RMSD$_{15}$ < 2 Å and no clashes, evaluated as-generated with no relaxation or DFT polish. Key results:

Method	$n_s$	$k$	Rigid	Flexible	CSP5	CSP6	CSP7	Teaching
OXtal	30	30	0.300	0.220	0.167	0.200	0.125	–
Clari-L	30	30	0.731	0.287	0.681	0.355	0.245	0.461
Clari-L	150	30	0.772	0.346	0.789	0.480	0.263	0.484
Clari-L	400	200	0.919	0.596	0.975	0.729	0.566	0.669
Clari-L	1000	1000	0.940	0.760	1.000	0.800	0.875	0.763

Even without inference-time scaling ($30\times$0), Clari-M outperforms OXtal on every subset. With best-of-$30\times$1 sampling ranked by UMA energy — enabled by explicit hydrogen modelling, which permits direct energy scoring without decoration or relaxation — Clari-L at $30\times$2, $30\times$3 exceeds the average DFT participant in CSP5 and CSP7. On the newly introduced CSD Teaching Subset benchmark (773 diverse crystals including organometallics and VSEPR-diverse structures), Clari-L reaches Sol@1000 of 0.763. Under the stricter 15/15 match criterion, solve rates drop substantially (e.g., 0.524 on Teaching at $30\times$4), which the authors acknowledge limits the scope of headline numbers given the loose 8/15 threshold.

Limitations and open questions

The paper states several limitations plainly. Clari conditions on a stereochemistry-agnostic 2D graph, which is inadequate for pharmaceutical applications where chirality determines polymorphism; conformer or chiral-tag conditioning is left open. The model requires the number of molecules per unit cell ($30\times$5) as input, though sweeping common values mitigates this. Evaluation infrastructure is itself a bottleneck: graph-based similarity measures for multi-component, non-sanitizable crystals are immature, the current splits rely on refcode grouping plus heuristic filtering, and no fast robust crystal-similarity metric exists for large-scale evaluation. Clari also occasionally produces steric clashes and unphysical voids, motivating reward alignment or inference-time steering (e.g., toward powder diffraction patterns or space groups), none of which is implemented here. Finally, the authors note that unit-cell modelling may be ill-suited to amorphous systems, where lattice-free bulk approaches retain an advantage.

Conclusion

Clari demonstrates that two targeted changes — explicit unit-cell parametrization and elimination of triangle layers — suffice to make generative organic CSP both more accurate and roughly an order of magnitude faster than OXtal, bringing per-molecule prediction to seconds and making large-scale virtual screening of organic solids computationally practical. The systematic ablations, particularly the finding that crystal generation requires balanced attention to global packing rather than late-timestep detail alone, offer transferable guidance for flow-matching models of periodic structures.