PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Emily Jin, Andrei Cristian Nica, Mikhail Galkin, Jarrid Rector-Brooks, Kin Long Kelvin Lee, Santiago Miret, Frances H. Arnold, Michael Bronstein, Avishek Joey Bose, Alexander Tong, Cheng-Hao Liu

arXiv:2512.06987v3cs.LGcond-mat.mtrl-sci

TL;DR

Crystal structure prediction seeks experimentally realizable 3D packings from 2D molecular graphs, a difficult problem because molecular crystals combine flexible conformations with weak, long-range periodic interactions. OXTAL uses an all-atom diffusion model with symmetry-aware augmentation and lattice-free S4 training, achieving strong packing and conformer recovery while remaining far cheaper than DFT-based methods.

  • Problem

    Crystal structure prediction must infer experimentally realizable molecular crystal packings from 2D graphs despite diverse conformations, weak long-range interactions, and complex periodic structures.

  • Method

    OXTAL is an all-atom diffusion model trained on over 600k experimental crystals, using Cartesian-coordinate augmentation and lattice-free S4 training to model conformations and periodic packing.

  • Results

    OXTAL outperforms existing ML-based ab initio CSP methods, achieving RMSD1 < 0.5 Å, packing similarity above 80%, and competitive performance at several orders of magnitude lower cost than DFT.

  • Takeaways & Limitations

    OXTAL captures diverse intra- and intermolecular interactions while sampling molecular crystal structures efficiently at all-atom resolution.

  • Takeaways & Limitations

    The S4 analysis assumes local interactions and bounds the sampled neighborhood statistic proportionally to its size.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Molecular CSP seeks experimentally realizable crystal packings from 2D molecular graphs, but must model coupled conformations, periodic intermolecular interactions, and kinetic accessibility. OXTAL addresses this with scalable all-atom diffusion modeling trained on diverse experimental crystals.

  • Crystal packing determines important pharmaceutical and materials properties, making prediction from 2D chemical graphs a major computational-chemistry challenge.
  • Classical CSP combines structure search with force-field or DFT ranking, yet typically generates 1,000–100,000 structures per molecule and often remains trapped in unfavorable local minima.
  • Small-molecule crystals require modeling diverse scaffolds, conformational flexibility, many molecular copies per unit cell, and weak periodic interactions at scale.
  • OXTAL learns crystal conformations and periodic packing from 2D molecular graphs using a 600k-structure dataset spanning rigid and flexible molecules, co-crystals, and solvates.
  • OXTAL discards explicit lattice symmetry representations in favor of Cartesian coordinates and SE(3) augmentation, and introduces lattice-free S4 training for long-range interactions.
  • OXTAL achieves RMSD1 < 0.5 Å and packing similarity above 80% within 30 samples, while costing several orders of magnitude less than DFT-based methods.

2 BACKGROUND AND PRELIMINARIES

The background formalizes molecular crystals as periodic atomistic structures with lattice and symmetry equivalences, then introduces continuous-time diffusion as a framework for sampling their conditional distributions. These foundations motivate learning models that respect crystal invariances while handling flexible conformations, periodic packing, and unknown stoichiometry.

  • 2.1 CRYSTAL REPRESENTATIONS: A periodic crystal is represented by lattice vectors L and an atomic basis B containing species and fractional coordinates, with molecular connectivity decomposing B into Z molecular graphs.
  • 2.1 CRYSTAL REPRESENTATIONS: Crystal descriptions are equivalent under rigid SE(3) transformations, lattice retilings, global translations, and integer-matrix supercell expansions.
  • 2.1 CRYSTAL REPRESENTATIONS: Ab initio CSP is posed as conditional probabilistic inference over equivalence classes of periodic structures, with experimentally realized distributions shaped by Gibbs free energy and kinetic accessibility.
  • 2.1 CRYSTAL REPRESENTATIONS: Learning must couple intramolecular conformation with intermolecular packing, navigate large rugged landscapes with weak long-range interactions, respect symmetry, and marginalize over unknown Z.
  • 2.2 CONTINUOUS-TIME DIFFUSION MODELS: Continuous-time diffusion models define a forward SDE that progressively corrupts data to a Gaussian prior and use the reverse-time SDE to generate samples from the target distribution.
  • 2.2 CONTINUOUS-TIME DIFFUSION MODELS: The forward and reverse processes are linked by the score ∇x log pt(xt), which diffusion models estimate through denoising objectives and use during reverse-time sampling.

3 OXTAL

OXTAL combines lattice-free S4 training with a non-equivariant all-atom diffusion architecture to model molecular conformations and periodic packing. The section covers shell-based cropping, model components, composite training losses, and a roughly 600k-crystal dataset.

  • 3 OXTAL: OXTAL combines lattice-free S4 training with a non-equivariant Transformer and data augmentation to decouple crystal generation from unit-cell and multiplicity representations.These choices target scalability for large molecular crystals with unknown multiplicity while exposing the model to long-range periodic cues.
  • 3.1 STOICHIOMETRIC STOCHASTIC SHELL SAMPLING (S4): S4 crops stoichiometry-preserving molecular shells around a sampled central molecule, capping atoms while retaining anisotropic contact networks and longer-range interactions.Shells are defined by minimum-image intermolecular distances and frontier subsampling when the token budget is exceeded.
  • 3.1 STOICHIOMETRIC STOCHASTIC SHELL SAMPLING (S4): S4 shell cropping outperforms centroid and kNN cropping, generalizes to long-range periodicity beyond training token sizes, and yields cropping-loss error that decreases with the cube root of token count.The method is motivated by local crystallization motifs propagating into larger periodic structures.
  • 3.2 MODEL ARCHITECTURE: OXTAL encodes atom-level physical and structural features, propagates single- and pairwise information with a Pairformer, and predicts denoised atomic positions using a 70M-parameter diffusion Transformer.The Pairformer is not explicitly equivariant, enabling training on larger sequences.
  • 3.3 TRAINING: Training combines mean squared error, smooth local distance, and distogram losses to capture global structure, local chemical environments, and binned pairwise distances across approximately 600k CSD crystals.The dataset supports training across molecular crystal structures curated from the Cambridge Structural Database.

4 EXPERIMENTS

OXTAL is evaluated against machine-learning and DFT-based CSP methods on benchmark tests and chemically diverse crystals. It combines strong structural recovery and packing similarity with substantially lower inference cost.

  • 4.1 CSP FOR RIGID AND FLEXIBLE MOLECULES: OXTAL outperforms existing ab initio ML methods across rigid and flexible molecular CSP metrics, with an order-of-magnitude improvement and the only approximate solves on the flexible dataset.The comparison uses 30 samples per target; Table 1 covers both rigid and flexible molecules.
  • 4.1 CSP FOR RIGID AND FLEXIBLE MOLECULES: OXTAL improves sample efficiency as the number of draws increases, with several rigid targets reaching packing and conformer success criteria in fewer than 10 samples.The reported RMSD15 improvement is log-linear among packing-similar predictions.
  • 4.2 CCDC CSP BLIND TESTS: Across the 5th–7th CCDC blind tests, OXTAL achieves the best or second-best performance with 30 samples and consistently the highest packing similarity rate.OXTAL generally recapitulates packing structures in 48–67% of samples, compared with 6–30% for DFT samples.
  • 4.2 CCDC CSP BLIND TESTS: OXTAL is over an order of magnitude cheaper at inference than DFT-based methods while maintaining strong packing similarity rates.Its upfront training cost is amortized across new molecules, supporting broad screening and dense posterior sampling before optional physics-based refinement.
  • 4.3 SURVEY OF CHEMICAL INTERPRETABILITY: The chemical survey finds accurate intra- and intermolecular interactions across flexible molecules, weak contacts, diverse packing motifs, polymorphs, and donor–acceptor co-crystals.Examples include a 6-mer peptide with 17 rotatable bonds, distinct experimental polymorphs, and alternating donor–acceptor π-stacked columns.
  • 4.3 SURVEY OF CHEMICAL INTERPRETABILITY: OXTAL samples occupy a relatively tight energy basin similar to physics-based submissions, suggesting avoidance of unphysical motifs without geometry relaxation.This analysis uses single-point GFN2-xTB calculations.

5 RELATED WORKS

Prior CSP approaches rely on expensive search and energy evaluations, while related generative models target inorganic crystals or proteins with more constrained interaction spaces. OXTAL instead addresses flexible molecular crystals through direct all-atom generation.

  • Physical approaches to crystal structure prediction: Classical CSP relies on search over predefined spaces and many expensive energy evaluations, whereas OXTAL avoids explicit energy-function calls.Recent alternatives replace DFT with machine-learning interatomic potentials, but OXTAL targets direct generation and reports lower speed and inference costs.
  • Generative models for inorganic crystal structure generation: Inorganic crystal generators typically model smaller unit cells dominated by strong interatomic bonding, unlike molecular crystals with flexible packing.This distinguishes their structural regime from molecular CSP.
  • Protein structure prediction: Protein structure models operate with a smaller residue vocabulary and evolutionary information, whereas general molecular CSP spans a larger chemical space without that information.The paper places OXTAL in a broader and less constrained generative setting.

6 DISCUSSION

OXTAL learns molecular conformations and periodic packing jointly from 2D graphs using symmetry-aware augmentation and lattice-free S4 training. It delivers strong accuracy and low cost, while leaving ranking, relaxation, crystallization context, solve rate, and sample efficiency for future improvement.

  • 6 DISCUSSION: OXTAL learns the joint distribution of molecular conformations and periodic packing from 2D molecular graphs using symmetry-aware augmentation and S4 instead of explicit equivariance and unit-cell parametrization.The approach learns periodic motifs from locally consistent neighborhoods at all-atom scale.
  • 6 DISCUSSION: OXTAL achieves state-of-the-art results among ab initio ML methods and competitive packing similarity at several orders of magnitude lower cost than DFT-based methods.The chemical survey also supports diverse interaction modeling while avoiding unphysical motifs.
  • 6 DISCUSSION: Future improvements include reliable ranking, local relaxation, crystallization-context conditioning, and higher solve rate and sample efficiency.The cited discussion identifies solvent and temperature as examples of crystallization context.

CONTRIBUTIONS STATEMENT

The paper combines shell-based theoretical analysis with implementation details for scalable crystal-structure modeling. Its technical design uses locality, uniform-density assumptions, and lattice-inspired cropping to control boundary effects.

  • CONTRIBUTIONS STATEMENT: The project combines model development, theoretical work, data processing, baseline experiments, and inference analysis across the listed contributors.The current model was developed by EJ, CL, and JRB, while theoretical work, data processing, baseline experiments, and inference analysis had separate leads.
  • A THEORY: S4 constructs crops by sampling complete molecular shells around a central molecule until the token budget is reached.Shells are defined using intermolecular distances and a predefined contact radius; the resulting crop is approximated by a ball centered on the central molecule.
  • A THEORY: Under locality and uniform-density assumptions, the theory bounds cropping-boundary effects and shows that boundary surface becomes less problematic as token count grows.The assumptions require interactions to vanish beyond a finite radius and token counts to scale with volume.
  • A THEORY: The analysis motivates shell cropping because kNN and centroid crops can miss distant or anisotropic interactions important for crystallization.The examples contrast hydrogen bonding with π-π stacking and show centroid crops missing peripheral interactions in elongated molecules.
  • B TECHNICAL DETAILS: The implementation uses AlphaFold3-style atom-level processing and diffusion denoising, trained across multinode GPU hardware.The model code is reported as publicly available, and experiments use NVIDIA L40S and H100 GPUs.

B.1 TRAINING DATA PROCESSING

The pipeline builds a filtered CSD training set, samples stoichiometric molecular shells, and predicts atomistic structures with an AlphaFold3-style diffusion architecture. Evaluation jointly measures collisions, packing, conformer recovery, and near-correct crystal recovery.

  • B.1 TRAINING DATA PROCESSING: 594,202 crystals remain after filtering CSD entries by structural quality, ambient single-crystal diffraction, molecular validity, known space group, and unit-cell size.Near-duplicate polymorphs are collapsed to reduce oversampling, and test-family or test-component leakage is excluded.
  • B.2.1 STOICHIOMETRIC STOCHASTIC SHELL SAMPLING (S4): S4 samples a central molecule, assigns other molecules to distance shells, and adds complete shells while preserving stoichiometry under the token budget.If no complete shell fits, adaptive sampling selects molecules by type proportions; implementation uses rcut = 4.5 Å, pmax = 0.8, and Tmax = 640.
  • B.2.2 ARCHITECTURE: The architecture processes atom-level tokens with an AF3-style Pairformer and diffusion head to output 3D coordinates.MSA pathways are disabled, while recycling and atom-attention encoding and decoding are retained.
  • B.2.3 LOSSES: The loss suite combines global and local structural objectives, including SmoothLDDT over atom pairs within 15 Å and multiple distance thresholds.SmoothLDDT replaces binary threshold tests with sigmoid-based comparisons after optimal rigid alignment.
  • B.2.4 MODEL HYPERPARAMETERS: Training uses Adam with specified scheduling and reports the model trained for 110,000 steps, while inference generates 30 seeded samples per target.The diffusion transformer uses crop and token size 640, with training split evenly between S4 and kNN cropping.
  • C.1 SELECTION OF PERFORMANCE METRICS: The evaluation metrics jointly probe collision freedom, packing similarity, conformer recovery, and holistic near-correct recovery.Crystal-level rates use OR aggregation across samples, whereas sample-level rates characterize the full proposal distribution.

C.2 MOLECULAR CSP ON RIGID AND FLEXIBLE DATASETS

The experiments evaluate OXTAL on balanced rigid and flexible molecular-crystal sets against ab initio machine-learning baselines and historical blind-test submissions. Results emphasize packing similarity, sample efficiency, and the advantages of S4 cropping.

  • C.2 MOLECULAR CSP ON RIGID AND FLEXIBLE DATASETS: The rigid and flexible benchmarks each contain 50 CSD molecular crystals, with flexibility defined primarily by rotatable-bond count and molecular context.The flexible set includes molecules with more than three rotatable bonds under the stated general definition.
  • C.3 Ab initio MACHINE LEARNING BASELINES: The baseline comparison includes A-Transformer, AssembleFlow-Atom, and AlphaFold3, evaluated with 30 samples per crystal target.AssembleFlow cannot handle flexible molecules, while AlphaFold3 is included as a general-purpose biomolecular generative baseline rather than a CSP-specific model.
  • C.4 CCDC CSP BLIND TEST DETAILS: OXTAL’s sample-efficiency plots show rapid recovery of both molecular conformers and periodicity on flexible and blind-test targets.The reported plots track RMSD1 and RMSD15 as the number of samples increases, while a beeswarm example uses 30 OXTAL samples.
  • C.4 CCDC CSP BLIND TEST DETAILS: Historical blind tests show that DFT ranking is reliable but remains computationally expensive and challenged by kinetic effects, disorder, flexibility, and high Z′.The paper positions OXTAL as directly generating high-quality structures rather than following the traditional generate–optimize–rank pipeline.
  • C.4 CCDC CSP BLIND TEST DETAILS: Blind-test metrics average per-group sample and crystal rates, counting groups with no submission for a crystal as misses.Blind Test 7 uses two stoichiometries for one target and counts eight total targets for crystal-level metrics.
  • C.4.1 INFERENCE TIME COST: Inference costs are estimated by pricing reported CPU and L40S GPU wall-clock times using September 2025 AWS on-demand rates.The method uses $0.0425 per vCPU-hour for normalized CPU time and $1.861 per L40S instance-hour.
  • D.1 CROPPING METHOD ABLATION: S4 performs best across all reported cropping metrics and datasets, with centroid-radius cropping second best in the ablation.The comparison is performed at 50k training steps, using 30 generated samples per target.

D.2 S4 RADIUS SIZE ABLATION

S4 performance changes little across reasonable shell radii, while a smaller 50M-parameter OXTAL still outperforms existing ab-initio ML methods.

  • D.2 S4 RADIUS SIZE ABLATION: The current S4 implementation may slightly favor smaller shell radii because the token budget limits the number of full shells it can include.
  • D.2 S4 RADIUS SIZE ABLATION: All three S4 radius settings outperform existing cropping methods, indicating robustness to reasonable radius choices.The main model uses a 4.5 Å radius, with a second shell at 9 Å chosen to reflect layered crystallization.
  • D.3 MODEL SIZE ABLATION: The smaller 50M-parameter OXTAL still outperforms existing ab-initio ML methods, supporting the benefits of its lattice-free, non-equivariant architecture at reduced scale.

E CONFORMER ANALYSIS

OXTAL remains robust to substantially different input conformers and can extend its predictions to larger periodic crystal blocks. Its conformer recovery stays stable as block size grows, while packing accuracy degrades only modestly.

  • E CONFORMER ANALYSIS: The model’s final conformers can differ substantially from the conditioning conformers, showing that diffusion refines rather than merely reproduces the input geometry.
  • E CONFORMER ANALYSIS: OXTAL’s predicted packings are not significantly affected by alternative or perturbed input conformers, even when they differ from the crystal conformation.The tested inputs include RDKit-ETKDG, universal-forcefield, MMFF94, GFN2-xTB, and ground-truth conformers.
  • F.1 INFERENCE ON LARGER CRYSTAL BLOCKS: As more molecule copies are generated, RMSD1 remains roughly constant while RMSD15 increases slightly, with larger blocks mostly staying below the 2.0 Å approximate-solve threshold.
  • F.1 INFERENCE ON LARGER CRYSTAL BLOCKS: OXTAL generalizes from local S4 neighborhoods to larger periodic packings, including an approximately solved ANTCEN structure containing over 2,400 tokens.

F.2 ADDITIONAL SAMPLE EVALUATION FOR CSP BLIND TEST 5

Using more generated candidates improves OXTAL’s per-crystal performance and brings its approximate solve rate to that of DFTavg. The surrounding analyses situate this result among physics-based, ML-potential, and generative CSP approaches.

  • F.2 ADDITIONAL SAMPLE EVALUATION FOR CSP BLIND TEST 5: Increasing OXTAL’s number of generated samples raises per-crystal performance and brings its approximate solve rate equal to DFTavg.
  • G ADDITIONAL RELATED WORK: Traditional CSP methods use search or sampling with force fields, quantum chemistry, or machine-learned potentials to evaluate candidate structures.
  • G ADDITIONAL RELATED WORK: Generative CSP research spans diffusion, symmetry-aware diffusion, flow matching, and language-model approaches, especially for inorganic crystals.
  • H ADDITIONAL CHEMICAL ANALYSIS: OXTAL examples include co-crystals, flexible conformers, intermolecular interactions, and polymorphs compared with experimental structures.

H.1 ENERGY ANALYSIS

GFN2-xTB analysis finds that raw OXTAL generations avoid catastrophic energetics and have a narrow distribution comparable to stable DFT basins. Qualitative samples also show varied packing, although planar arrangements may be preferred.

  • H.1 ENERGY ANALYSIS: The energy comparison uses GFN2-xTB single-point calculations on experimental structures, physics-based submissions, and OXTAL samples.
  • H.1 ENERGY ANALYSIS: Raw OXTAL samples avoid energetically catastrophic motifs and show a narrow energy distribution comparable to the stable basin of DFT samples despite lacking relaxation.The analysis is intended as a consistency check on gross energetic behavior, not a quantitative ranking of polymorph stability.
  • H.2 DIVERSITY OF GENERATED SAMPLES: Generated samples include both repeated orientations and more complex herringbone packings, including distinct orientations for the two components of the XATJOT co-crystal.
  • H.2 DIVERSITY OF GENERATED SAMPLES: OXTAL may prefer planar packings because they are more prevalent in its training data, though additional noise-scheduler tuning may mitigate this tendency.
Loading 2512.06987v3…