Fast Organic Crystal Structure Prediction with Unit Cell Flow Matching

Motivation

• Several CPU-years per molecule make traditional CSP prohibitive for screening even modest virtual libraries.Traditional pipelines densely generate, filter, and relax candidate structures with increasingly costly quantum-chemistry methods.Source: §1 · p. 2

Clari

• Clari directly predicts a single unit cell with atom coordinates and lattice vectors, avoiding bulk expansion and triangle-update layers.Its flow-matching model uses a pair-bias architecture and removes redundant bulk representations.Source: §1 · p. 2, Figure 1 · p. 2

Results

• Clari surpasses OXtal in prediction quality while accelerating sampling by roughly 15–30× on CSP Blind Tests and 5–8× end-to-end with UMA energy ranking.The model jointly represents heavy atoms and hydrogens, enabling energy-based ranking without decoration or relaxation.Source: §1 · pp. 2–3

Implications

• The model's all-atom generation supports inference-time scaling through direct energy ranking and makes generative CSP suitable for large virtual-library screening.The stated use case is screening solid-state properties across large molecular collections.Source: §1 · p. 3

Evaluation

• The evaluation combines complementary metrics, systematic ablations, and a CSD Teaching Subset containing chemically complex systems underrepresented in prior benchmarks.The new split includes fullerenes, boranes, and organometallic complexes.Source: §1 · p. 3, Figure 2 · p. 3

Machine learning for accelerating CSP

• Traditional CSP generates and evaluates many candidates, so even accelerated ranking leaves the combinatorial burden of structure generation largely unchanged.ML interatomic potentials reduce energy and force-evaluation costs but do not alter the search paradigm.Source: §2 · p. 3

Generative CSP

• Generative CSP directly samples low-energy crystals, reducing reliance on exhaustive downstream ranking compared with conventional pipelines.OXtal introduced this paradigm for organic CSP, followed by alternative formulations such as PackFlow and MolCrystalFlow.Source: §2 · p. 3
• Current generative methods either use computationally intensive triangle updates over large crops or restrictive rigid-molecule representations.These choices limit efficiency or applicability to organic crystal generation.Source: §2 · pp. 3–4

Generative models for chemistry

• Generative models for inorganic crystals commonly parameterize unit cells and atomic positions explicitly, providing a related precedent for lattice-based modelling.The paper also draws on equivariant generative modelling of proteins, complexes, and molecules.Source: §2 · p. 4

3.1 Crystal structure prediction

• Clari represents each crystal as lattice, centered atom coordinates, atom features, and bond-type adjacency, predicting lattice and coordinates from a molecular graph.The model treats CSP as inference of (L, C) given (F, E), without using space-group input.Source: §3.1 · p. 4

3.2 Flow matching

• Flow matching transports a fitted source distribution into the data distribution by regressing target velocity along linear interpolants and integrating the learned vector field with an ODE solver.The loss separates lattice and coordinate terms, each computed with mean-squared error.Source: §3.2 · p. 4
• The training design explores timestep schedules and uses a data-informed lattice prior because a standard normal source produces unrealistic densities and degenerate lattices.The prior fits atom density, cell angles, and volume-normalized cell lengths from training data.Source: §3.2 · pp. 4–5, Figure 3 · p. 5
• Clari adds volume and periodic-pair-distance auxiliary losses, uses self-conditioning, and implements a scalable DiT with pair-bias attention instead of triangular operations.Self-conditioning improves performance but slows training by approximately 20%, while pair bias transmits geometric information through attention logits.Source: §3.2 · p. 5, §3.3 · p. 5

3.4 Optimal transport and augmentation

• Approximate optimal transport aligns lattices, atoms, and poses across symmetry groups using signed-permutation, Hungarian, and Kabsch steps.Periodic translations are handled by randomly translating cells, rewrapping component centroids, and recentering atomic coordinates.Source: §3.4 · p. 6

3.5 Inference-time scaling

• Best-of-N inference generates multiple all-atom crystals, ranks them with UMA, and selects low-energy candidates without hydrogen decoration or relaxation.This enables direct energy-based inference-time scaling.Source: §3.5 · p. 6

4.1 Experimental setup

• Sol@k counts a target as solved when at least one of k samples matches the ground truth in at least 8 of 15 molecules with RMSD15 < 2 Å and no inter-body clashes.The held-out test pool combines OXtal test crystal families with the CSD Teaching Subset.Source: §4.1 · pp. 6–7

4.2 Ablations

• Auxiliary volume and pairwise-distance losses plus self-conditioning produce notable improvements across all validation metrics.Self-conditioning adds an extra half-batch forward pass during training and slows training by approximately 20%, but does not add inference cost.Source: §4.2 · p. 8
• A fitted lattice prior outperforms a standard normal prior, whose reconstruction performance falls close to the simpler baseline and whose degraded EMD PDD suggests unrealistic sparse unit cells.The standard normal prior marginally improves clash rate but worsens reconstruction metrics.Source: §4.2 · p. 8
• Timestep choices trade off packing and validity: mid-range-focused training improves packing metrics, whereas Beta(1.8, 1) yields better PoseBusters validity and is used in the final models.Log discretization is broadly harmful, particularly for clash rate and EMD PDD.Source: Table 2 · p. 7, §4.2 · p. 8
• Scaling from Clari-M at 88M to Clari-L at 173M parameters improves reconstruction while making training only approximately 20% slower.The larger model nearly doubles parameter count without doubling training time.Source: §4.2 · p. 8

4.3 Comparison to baselines

• Inference-time energy ranking improves selection performance, but increasing the sampling budget does not always improve Sol@30 because ranking noise can intervene.The authors tune ns = 150 on a held-out set and use UMA energy to retain the lowest-energy candidates.Source: §4.3 · pp. 8–9, Figure 5 · p. 9
• Clari-L with top-30 energy selection achieves the strongest performance among the authors’ methods on every OXtal benchmark subset, while Clari-M without scaling outperforms OXtal across all subsets.Generating 150 structures takes 2.2 seconds per molecule, or 6.0 seconds when downselecting to 30 with UMA energy ranking.Source: Table 3 · p. 8, §4.3 · p. 9
• Clari-L reaches Sol@1000 = 0.763 on the CSD Teaching Subset, covering diverse chemistry including metal complexes, fullerenes, and atom clusters.The filtered subset contains 773 crystals, and these examples are not RDKit-sanitizable for OXtal-style input preparation.Source: §4.3 · pp. 9–10

• 15–30× faster than OXtal on CSP Blind Tests, or 5–8× faster end-to-end with UMA energy ranking, while improving prediction quality.Clari operates on a single unit cell and avoids bulk expansion and triangle-update layers.Source: §5 · p. 10
• Joint modelling of heavy atoms and hydrogens enables direct energy-based ranking and best-of-N inference-time scaling without decoration or relaxation.Source: §5 · p. 10
• Explicit lattice modelling reduces redundant computation, cutting crystal generation from minutes to seconds and supporting screening of large virtual libraries.Few-step generation methods are described as complementary to Clari’s architectural speedups.Source: §5 · p. 10
• Clari’s explicit-lattice representation is complementary to bulk approaches, which may better accommodate amorphous systems with large or ill-defined unit cells.Source: §5 · p. 10
• The method conditions on a stereochemistry-agnostic 2D graph, requires the number of molecular copies Z, and can generate steric clashes or unphysical voids.The paper also notes limitations from approximate similarity metrics and non-principled dataset splitting.Source: §5 · p. 10
• Accelerated CSP may support screening in pharmaceuticals and organic materials, while energetic-material applications require responsible use within safety and regulatory frameworks.Source: §5 · p. 10

A.1 Crystal processing pipeline

• The CSD is accessed through its Python API for metadata and ConQuest for raw CIF and MOL2 files.Source: §A.1 · p. 17
• The pipeline filters metadata, parses asymmetric units, caches component isomorphisms, expands space-group operations, removes polymers and clashes, and featurizes unit cells.Source: §A.1 · p. 17
• Isomorphisms are cached on the asymmetric unit and extended after replication, avoiding repeated expensive checks while retaining one matching isomorphism when available.Source: §A.1 · p. 17
• Distance filtering uses periodic distances, covalent radii, and thresholds of 0.6 for metals or bonded atoms and 1.0 otherwise.The thresholds are selected from CSD periodic-distance distributions, with a special case for exactly superimposable isomorphic components.Source: §A.1 · p. 17

A.2 Dataset splitting

• 917,014 training, 1,048 validation, and 2,996 test examples are split across 859,866, 1,000, and 919 families, respectively.The test count includes held-out crystals with more than 512 atoms in the unit cell.Source: §A.2 · pp. 17–18
• The held-out test pool combines OXtal test families and the CSD Teaching Subset, excluding six-letter refcode families and shared RDKit-sanitizable components with more than seven heavy atoms.Source: §A.2 · p. 17
• CSD distance distributions distinguish bonded neighbors from nearest unbonded neighbors, including metal-containing and nonmetal pairs, with histograms truncated at the lower 0.5% tail.Source: Figure 6 · p. 18

B.1 Source distribution

• The lattice source distribution independently fits atom density and cell angles with univariate Gaussians, while normalized cell lengths use a correlated multivariate Gaussian.The resulting lattice is reconstructed and randomized by signed permutation and rotation.Source: §B.1 · p. 18

B.2 Architecture

• AdaLN modulates LayerNorm outputs using conditioning features, and the architecture embeds sequence and pair features before trunk processing.Source: §B.2 · p. 18, §B.3 · p. 19
• PairAttention injects pairwise information through a linear transformation of pair features, while the main trunk alternates conditioned attention and transition blocks.Source: §B.2 · p. 18, §B.3 · p. 19

B.3 Inference batching

• Sampling replicates each crystal B(n) times for parallel inference, using smaller batch sizes as atom count increases and halving B after out-of-memory errors.Source: §B.3 · p. 18

B.4 Auxiliary losses

• The auxiliary volume and pair losses are applied between the one-step estimate and ground-truth structure.Source: §B.4 · p. 19
• The pair loss matches periodic distances and adds a clash term over atom pairs selected by distance criteria, using an approximation that is usually exact for crystals.Source: §B.4 · p. 19

B.5 Optimal-transport coupling

• The procedure alternates permutation and pose alignment after an initial lattice-only alignment, then constructs the interpolant from the approximately aligned pair.Pose alignment uses a weighted Kabsch procedure that gives the three lattice rows equal total weight to the N atom rows.Source: §B.5 · pp. 19–20
• Lattice alignment exactly minimizes the Frobenius distance over signed permutations and rotations, using enumeration and the Kabsch algorithm.Source: §B.5 · p. 19
• Atomic permutation alignment uses cached graph isomorphisms, RMSD-based component matching, and the Hungarian algorithm rather than searching all possible isomorphisms.This is an approximate alignment because only one cached isomorphism per component is considered.Source: §B.5 · p. 20

B.6 Hyperparameters

• The medium and large models use a 2.66× feedforward expansion, Muon for Transformer matrices, Adam for other parameters, and EMA parameters for sampling and evaluation.Source: §B.6 · p. 20

C Metrics

• Quality metrics assess inherent plausibility, whereas reconstruction metrics assess agreement with a reference crystal.Clash rate and PoseBusters are quality metrics; the remaining listed metrics are reconstruction metrics.Source: C Metrics · p. 20
• Clash rate reports the fraction of crystals containing at least one pair of components whose periodic atomic distance is below the sum of covalent radii.Source: C Metrics · pp. 20–21
• PoseBusters reports average component validity only over applicable RDKit-processable components and crystals containing at least one such component.Components outside the supported chemistry or processing conditions are excluded before averaging.Source: C Metrics · p. 21
• EMD PDD measures the L1 earth mover’s distance between pointwise distance distributions computed from the nearest 100 neighbours of each point.Source: C Metrics · p. 21
• COMPACK RMSD15 measures heavy-atom packing through a 15-molecule cluster match, and a target is solved when one candidate meets the match, RMSD15, and collision criteria.The comparison inherits OXtal’s tolerances and requires RMSD15 < 2.0 Å with at least 8 of 15 molecules matched and no detected collisions.Source: C Metrics · p. 21
• The protocol evaluates 119 OXtal-test crystals without relaxation, energy minimization, or DFT polishing, using only UMA-based ranking as post-processing.The set contains rigid, flexible, and CSP blind-test crystals, with reference matching based on specified CSD entries.Source: C Metrics · p. 21

D Additional results

• The section provides additional supporting figures and tables for the reported results.Source: D Additional results · p. 22
• The reconstruction curves compare models A and B on a 512-crystal validation subset using three samples per crystal and sample-wise mean aggregation.These results were taken from preliminary training logs.Source: Figure 7 · p. 22
• ns = 150 achieves the best Sol@30 solve rate on 50 random CSD Teaching Subset examples.The estimate uses 1000 samples per example and the mean over 5000 bootstrap resamples.Source: Figure 8 · p. 22
• Table 6 reports Sol@k metrics from Table 3 with bootstrap means and standard errors over 5000 resamples where appropriate.Source: Table 6 · p. 23
• Table 7 reports the same Sol@k metrics under a stricter 15-molecule match requirement, again with bootstrap means and standard errors where appropriate.Source: Table 7 · p. 24
• Example CSD Teaching Subset structures achieve full 15/15 molecule matches in COMPACK with reported RMSD15 values.Source: Figure 9 · p. 25