OXtal: An All-Atom Diffusion Model for Organic Crystal Structure Prediction (2512.06987v1)

Abstract: Accurately predicting experimentally-realizable 3D molecular crystal structures from their 2D chemical graphs is a long-standing open challenge in computational chemistry called crystal structure prediction (CSP). Efficiently solving this problem has implications ranging from pharmaceuticals to organic semiconductors, as crystal packing directly governs the physical and chemical properties of organic solids. In this paper, we introduce OXtal, a large-scale 100M parameter all-atom diffusion model that directly learns the conditional joint distribution over intramolecular conformations and periodic packing. To efficiently scale OXtal, we abandon explicit equivariant architectures imposing inductive bias arising from crystal symmetries in favor of data augmentation strategies. We further propose a novel crystallization-inspired lattice-free training scheme, Stoichiometric Stochastic Shell Sampling ($S^4$), that efficiently captures long-range interactions while sidestepping explicit lattice parametrization -- thus enabling more scalable architectural choices at all-atom resolution. By leveraging a large dataset of 600K experimentally validated crystal structures (including rigid and flexible molecules, co-crystals, and solvates), OXtal achieves orders-of-magnitude improvements over prior ab initio machine learning CSP methods, while remaining orders of magnitude cheaper than traditional quantum-chemical approaches. Specifically, OXtal recovers experimental structures with conformer $\text{RMSD}_1<0.5$ Å and attains over 80\% packing similarity rate, demonstrating its ability to model both thermodynamic and kinetic regularities of molecular crystallization.

• The paper introduces an all-atom diffusion model (OXtal) that achieves state-of-the-art accuracy in predicting organic crystal packings using the novel S⁴ local-to-global sampling strategy.
• OXtal leverages a transformer-derived architecture with rich atomic features and composite loss functions to outperform ML baselines and reduce computational cost by over 10×.
• Experimental evaluations show that OXtal delivers high packing similarity, conformer recovery, and chemical realism across both rigid and flexible molecular systems.

OXtal: An All-Atom Diffusion Model for Organic Crystal Structure Prediction

Problem Statement and Motivation

Crystal structure prediction (CSP)—inferring the three-dimensional packing of molecules within a periodic lattice solely from two-dimensional chemical graphs—remains a formidable challenge in computational chemistry. Accurate CSP is crucial for pharmaceuticals and materials science, as crystal packing directly dictates macroscopic properties (e.g., solubility, charge transport, stability). Traditional CSP approaches rely on expensive quantum chemical calculations—such as density functional theory (DFT)—and often enumerate upwards of $10^4$–$10^5$ candidate packings per target. While recent generative models have advanced protein and inorganic structure prediction, the intrinsic diversity, conformational flexibility, and multiplicity ($Z$) in molecular crystals demand substantially higher model expressiveness, scalability, and sample efficiency.

OXtal: Methodological Advances

The paper introduces OXtal, a $\sim$100M-parameter, all-atom diffusion model for ab initio molecular CSP. Rather than imposing crystal symmetry via explicit equivariant architectures, OXtal leverages data augmentation to achieve $\textrm{SE}(3)$ invariance, which, combined with large-scale training ($\sim$600k experimental structures), enables scalability across diverse organic compounds—including flexible molecules and multicomponent co-crystals.

Central to OXtal’s training framework is Stoichiometric Stochastic Shell Sampling ($S^4$), a lattice-free, crystallization-inspired cropping strategy. $S^4$ exposes the model to local-to-global packing cues without explicit lattice parameterization:

1. Local-to-Global Training via Shell Crops: Training samples consist of concentric “shells” of molecules around a random central entity, constructed by intermolecular distance rather than $k$NN or centroid heuristics. This captures weak and long-range packing cues essential for correct crystallization.
2. Scalability and Generalization: $S^4$ ensures the cropping-induced loss is dominated by bulk interior terms, with boundary errors vanishing as $O({T}^{-1/3})$ (see Appendix, Proposition 1), allowing crops up to hundreds of atoms to generalize to much larger periodic contexts.

   Figure 1: Molecular crystal structures generated by OXtal (color) compared to ground truth (grey).

Model Architecture Overview:

• Atom Encoder: Embeds 2D molecular graphs with rich physical and chemical features (atomic numbers, partial charges, bond types, 3D geometry from GFN2-xTB reference conformers).
• Pairformer Trunk: A transformer-derived architecture computes single and pairwise atom representations, adapted from AlphaFold3 but simplified to atom tokens without explicit residue or MSA channels.
• Diffusion Head: Incorporates per-atom attention encoders/decoders with a 70M-parameter diffusion transformer, parameterizing the denoising process.

Training employs composite losses, including MSE, smooth local distance difference test (sLDDT), and a distogram loss, to ensure accuracy in both global crystal packing and local chemical environment.

Figure 2: Example crystal packing generated by A-Transformer, AssembleFlow, and OXtal.

Experimental Evaluation: Baseline and Benchmark Performance

OXtal is rigorously benchmarked against both modern ML baselines (AssembleFlow, A-Transformer, AlphaFold3 zero-shot) and traditional DFT-based approaches—especially those featured in the CCDC CSP5, CSP6, and CSP7 blind tests.

Quantitative Metrics:

• Packing Similarity Rate (Pac$_S$, Pac$_C$): Fraction of generated samples/crystals matching experimental packing over a 15-molecule cluster via CSD COMPACK.
• Conformer Recovery (Rec$_S$, Rec$_C$): RMSD$_1 < 0.5$Å on internal coordinates.
• Collision Rate (Col$_S$): Fraction of unphysical samples with severe steric clashes.

Key Results:

• On both rigid and flexible molecule datasets, OXtal exceeds all ML baselines by an order of magnitude in packing similarity and conformer recovery. For rigid systems, Pac$_C = 1.0$, Rec$_C = 0.96$, Col$_S = 0.011$. For flexible systems, OXtal is the only model to achieve any approximate solves (Pac$_C = 0.9$, Rec$_C = 0.4$).
• Few-sample efficiency is demonstrated: most targets achieve correct packing within 10–30 samples, compared to hundreds-thousands for DFT techniques.

  Figure 3: OXtal sample efficiency for 10 rigid and flexible molecules.

• In all three recent CSP blind tests (CSP5-CSP7), OXtal matches or surpasses DFT-based approaches in packing similarity rate with orders-of-magnitude lower computational cost. For instance, in CSP7, OXtal's Pac$_C$ is 0.875 with 30 samples per target compared to DFT's 0.511 with potentially several thousand samples.
• OXtal is >10× more cost-efficient per successful packing prediction, as illustrated by normalized cloud compute estimates.

  Figure 4: Packing similarity rate per crystal relative to average inference cost for CSP competition methods. OXtal shown in red.

Analysis of Chemical Realism and Diversity

Comprehensive chemical analysis reveals that OXtal’s generative samples recover realistic intramolecular geometries for highly flexible molecules (e.g., 6-mer peptides, drugs), and capture strong as well as weak intermolecular interactions— hydrogen bonding, halogen bonds, $\pi$-stacking, and more—in alignment with experimental structures. The model generalizes to chemically and structurally diverse packings, including large, herringbone, layered, and brickwork motifs.

Figure 5: Truncated distribution of intermolecular distances in the processed training dataset.

Importantly, OXtal samples encompass multiple experimental polymorphs, demonstrating support for modeling both thermodynamic and kinetic crystallization basins. The model also handles complex, multicomponent co-crystals, accurately reproducing donor–acceptor interactions and extended periodicity in electronic materials.

Figure 6: Examples of OXtal generated co-crystal structures (color) compared against experimental structures (gray).

Energetic plausibility is further validated via single-point GFN2-xTB calculations: OXtal-generated structures fall within the narrow, experimentally relevant energy basin occupied by stable DFT samples, despite no explicit energy minimization.

Figure 7: GFN2-xTB single point energy analysis of ground truth, DFT submissions, and OXtal samples.

Limitations and Future Directions

Although OXtal achieves state-of-the-art ab initio CSP performance, several areas for future improvement are identified:

• Ranking and Local Relaxation: Integration of a robust ranking stage and/or post hoc local refinement could further improve solve rates and RMSD distributions.
• Conditional Generation: Incorporating synthesis context (e.g., solvent, temperature) for environment- or polymorph-specific prediction.
• Expanded Scope: Extensions to metals, metal–organic frameworks, or explicit treatment of disorder and high-Z′ systems.
• Enhanced Sampling/Score-based Inference: Further sample efficiency improvements leveraging score-based inference or advanced denoising samplers.

Theoretical and Practical Implications

OXtal demonstrates that large-scale, symmetry-agnostic, all-atom diffusion models—when paired with effective cropping, data augmentation, and massive experimental data—can outperform both symmetry-constrained ML and physics-based DFT methods on molecular CSP. The $S^4$ protocol allows for efficient, boundary-controlled local-to-global learning in periodic systems, addressing a major scalability bottleneck in generative modelling for materials.

Practically, OXtal enables rapid, high-throughput screening of organic solids and drug candidates, facilitating materials discovery cycles which are infeasible for classical DFT. The ability to generalize periodic motifs from finite blocks built only from chemical graphs is a powerful paradigm for AI-driven design in chemistry and materials science.

Conclusion

OXtal establishes a new state-of-the-art for ab initio crystal structure prediction, providing accurate, sample-efficient, and computationally tractable models for the all-atom generation of realistic molecular crystals. The methodology bridges gaps between classic search, ML-based potentials, and full generative ab initio structure prediction. As generative modeling and structural datasets continue to grow, this approach is poised to significantly accelerate and democratize the design of organic materials and pharmaceuticals.

Figure 8: Example set of structures generated by OXtal for various crystals (labeled by CSD ID).