PaperlayerView exact revision on arXiv

Source-linked AI summary

Crystal Structure Prediction by Joint Equivariant Diffusion

Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, Yang Liu

arXiv:2309.04475v2cond-mat.mtrl-scics.LG

TL;DR

CSP requires generating stable crystal structures while respecting translation, rotation, reflection, and periodicity symmetries. DiffCSP jointly diffuses lattices and fractional coordinates with a periodic-E(3)-equivariant denoiser. The paper reports strong performance across crystal datasets, comparisons with learning-based and DFT-based methods, and effectiveness after extension to ab initio generation.

  • Problem

    CSP generation must account for crystal-distribution symmetries and periodic atom arrangements, which make diffusion-model design challenging.

  • Method

    DiffCSP jointly diffuses lattice matrices and fractional coordinates using a denoising model designed to produce periodic-E(3)-invariant samples.

  • Results

    DiffCSP outperforms existing CSP methods and remains effective when extended to ab initio crystal generation.

  • Takeaways & Limitations

    Fractional coordinates provide an intrinsic representation of periodicity that supports the model’s symmetry-aware crystal generation process.

  • Takeaways & Limitations

    The model has slightly lower compositional validity, and further wet-lab experiments are needed to verify real-world effectiveness.

Abstract

from arXiv · show

Crystal Structure Prediction (CSP) is crucial in various scientific disciplines. While CSP can be addressed by employing currently-prevailing generative models (e.g. diffusion models), this task encounters unique challenges owing to the symmetric geometry of crystal structures -- the invariance of translation, rotation, and periodicity. To incorporate the above symmetries, this paper proposes DiffCSP, a novel diffusion model to learn the structure distribution from stable crystals. To be specific, DiffCSP jointly generates the lattice and atom coordinates for each crystal by employing a periodic-E(3)-equivariant denoising model, to better model the crystal geometry. Notably, different from related equivariant generative approaches, DiffCSP leverages fractional coordinates other than Cartesian coordinates to represent crystals, remarkably promoting the diffusion and the generation process of atom positions. Extensive experiments verify that our DiffCSP significantly outperforms existing CSP methods, with a much lower computation cost in contrast to DFT-based methods. Moreover, the superiority of DiffCSP is also observed when it is extended for ab initio crystal generation.

1 Introduction

Crystal structure prediction infers stable 3D structures from composition, but crystal symmetries and periodicity complicate diffusion-based generation. DiffCSP addresses these challenges through periodic-E(3)-equivariant joint generation of lattices and fractional coordinates.

  • Crystal structure prediction seeks stable 3D structures from chemical composition, supporting applications including drug, battery, and catalyst design.
  • DFT-based methods are computationally intensive, motivating generative models that learn stable-structure distributions directly from data.
  • Crystal generation must preserve E(3) invariance and periodicity because transformations of crystal coordinates do not change the physical structure distribution.
  • DiffCSP jointly diffuses lattice vectors and fractional coordinates with a denoising model theoretically proved to generate periodic-E(3)-invariant samples.
  • Fractional coordinates intrinsically encode periodicity and support Wrapped Normal modeling and Fourier-based denoising design.
  • DiffCSP is evaluated against learning-based and DFT-based methods and extended to ab initio generation, where the composition is also generated.

2 Related Works

Prior crystal-generation methods combine optimization, machine-learned energy prediction, graph representations, or diffusion, but do not uniformly address crystal symmetries and the CSP setting.

  • Traditional CSP combines DFT with optimization, while machine-learning methods can replace DFT as energy predictors within optimization pipelines.
  • Some direct generative methods learn stable structures but remain unaware of the full symmetries of crystal structures.
  • Periodic crystal models have used multi-graph edges or periodic-pattern encoding to represent interactions across adjacent lattice cells.
  • Diffusion models connect data and prior distributions through forward and backward Markov chains and have been applied to several scientific generation tasks.

3 Preliminaries

A crystal is represented by a periodically repeated unit cell, and CSP learns the conditional distribution of lattice and fractional coordinates given composition.

  • A crystal consists of an infinite periodic atom arrangement whose smallest repeating unit is the unit cell.
  • The unit cell is represented by atom types A, Cartesian coordinates X, and lattice matrix L containing three basic periodicity vectors.
  • Integral translations of lattice vectors describe repeated cells throughout the periodic crystal.
  • Fractional coordinates use lattice vectors as bases, mapping f to Cartesian coordinates through x = Σ_i f_i l_i.
  • DiffCSP represents each crystal as M = (A, F, L), with F containing all atoms’ fractional coordinates in [0, 1)3×N.
  • CSP predicts L and F for each unit cell from composition A by learning p(L, F | A).

4 The Proposed Method: DiffCSP

DiffCSP constructs a symmetry-aware diffusion process that jointly generates lattice and fractional-coordinate variables while preserving permutation, orthogonal, and periodic-translation invariances.

  • 4.1 Symmetries of Crystal Structure Distribution: Crystal structure distributions require permutation, O(3), and periodic-translation invariance; the latter two are termed periodic E(3) invariance.
  • 4.1 Symmetries of Crystal Structure Distribution: Permutation invariance means reordering atoms does not change the conditional crystal distribution.
  • 4.1 Symmetries of Crystal Structure Distribution: O(3) invariance means rotating or reflecting lattice vectors leaves the distribution unchanged.
  • 4.1 Symmetries of Crystal Structure Distribution: Periodic translation invariance wraps translated fractional coordinates into [0, 1)3×N without changing the distribution.
  • 4.2 Joint Equivariant Diffusion: DiffCSP simultaneously diffuses L and F through forward noising and backward denoising conditioned on composition A.
  • 4.2 Joint Equivariant Diffusion: Invariant priors combined with equivariant Markov transitions ensure the recovered distribution satisfies periodic E(3) invariance.
  • 4.2 Joint Equivariant Diffusion: Lattice diffusion uses DDPM with a Normal prior, while fractional-coordinate diffusion uses score matching and Wrapped Normal distributions for its bounded cyclical domain.
  • 4.3 The Architecture of the Denoising Model: Fractional coordinates are used instead of Cartesian coordinates, eliminating the need for multi-graph periodic representations and naturally encoding periodicity.

5 Experiments

Experiments evaluate DiffCSP across crystal datasets, stable-structure prediction, DFT comparison, ablations, and ab initio generation. DiffCSP consistently outperforms competing methods, while ablations support joint lattice-coordinate diffusion and periodic-invariance design choices.

  • Experimental Setup: DiffCSP is evaluated on four datasets spanning distinct difficulty levels, including structures with up to 52 atoms per unit cell.The datasets are Perov-5, Carbon-24, MP-20, and the more challenging MPTS-52.
  • Experimental Setup: The evaluation matches generated candidates to ground-truth structures using Match rate and normalized RMSE, with k = 1 and k = 20 samples.Optimization methods instead provide 20 lowest-energy candidates selected from 5,000 structures generated during testing.
  • Stable Structure Prediction: DiffCSP consistently outperforms all other methods, including optimization and generative baselines, although every method degrades as atoms per cell increase.Optimization methods achieve Match rates below 10% on MPTS-52, underscoring the difficulty of the larger search space.
  • DFT-based Comparison: DiffCSP produces the best qualitative predictions in Figure 3 and is compared with USPEX on 15 selected binary and ternary MP-20 compounds.The comparison uses 20 candidates per compound; DiffCSP inference uses a model trained for 5.2 hours, while USPEX uses 20 generations and 20 populations.
  • Ablation Studies: Ablations show that breaking O(3) equivariance yields only 1.66% matched structures, while removing Fourier periodicity lowers Match rate from 51.49% to 29.15%.The ablations also report that sequential lattice-coordinate updates, non-periodic coordinate generation, and multi-graph representations reduce performance or sampling stability.
  • Ab Initio Crystal Generation: For ab initio generation, DiffCSP achieves comparable validity and coverage to prior methods while significantly outperforming them on property-statistics similarity.This extension applies discrete diffusion to atom types in addition to the model's structural generation process.

6 Discussions

DiffCSP’s diffusion process is designed to preserve crystal symmetries while jointly modeling lattice and fractional-coordinate variables. Experiments support its applicability, though compositional validity and real-world experimental verification remain limitations.

  • Limitations: DiffCSP achieves slightly lower compositional validity in the ab initio generation evaluation.The authors identify composition generation as a limitation and suggest more powerful atom-type generation methods.
  • Limitations: Further wet-lab experiments are needed to better verify the model’s effectiveness in real applications.This is identified as a separate limitation of the experimental evaluation.
  • Method: DiffCSP jointly optimizes lattice and fractional coordinates within a diffusion-based framework tailored to crystal symmetries.The method is presented as a learning framework whose diffusion process jointly handles lattice and fractional-coordinate variables.
  • Symmetry guarantees: The generation process guarantees invariant intermediate distributions under the necessary transformations.This follows from the symmetry-preserving construction of the diffusion process.
  • Symmetry guarantees: The symmetry proof relies on a G-invariant prior and G-equivariant Markov transitions to establish invariance of the generated marginal distribution.The lemma provides the general diffusion-process argument used in the proposition.
  • Symmetry guarantees: O(3)-equivariance of the lattice denoising term implies that the marginal lattice distribution p(L0) is O(3)-invariant.The proposition connects the denoiser’s equivariance condition with invariance of the generated lattice distribution.

A.2 Proof of Proposition 2

The section establishes that DiffCSP’s wrapped-normal coordinate transitions preserve periodic translation symmetry, yielding a periodic E(3)-invariant generated distribution. It also describes the sampling procedures and the adaptations needed to condition CDVAE on composition.

  • Diffusion on fractional coordinates: Wrapped normal distributions model fractional coordinates on the periodic domain [0,1), whose period is 1.This avoids applying an ordinary normal diffusion directly to the cyclical, bounded fractional-coordinate space.
  • Proof of periodic translation symmetry: If the fractional-coordinate denoiser is periodic translation invariant, the wrapped-normal transition is periodic translation equivariant.The proof uses invariance of the wrapped normal under simultaneous periodic translation of its arguments.
  • Proof of periodic translation symmetry: The predictor and corrector transitions are periodic translation equivariant, so the marginal distribution p(F0) is periodic translation invariant.The prior U(0, 1) is itself periodic translation invariant, completing the argument for the sampled fractional coordinates.
  • Generated-distribution invariance: DiffCSP’s lattice score is O(3)-equivariant and its fractional-coordinate score is periodic translation invariant, making the generated distribution periodic E(3) invariant.The construction therefore incorporates orthogonal symmetry and periodic translation symmetry into generation.
  • Periodic boundary shifting: Changing the unit-cell origin is equivalent to applying a periodic translation to fractional coordinates, so different periodic boundaries represent the same infinite crystal.The bijection between unit cells with origins p1 and p2 establishes this equivalence.
  • CDVAE adaptation and sampling: The adapted CDVAE conditions generation on composition by adding a 1D prior encoder and keeping atom types unchanged during Langevin dynamics.Its generation procedure initializes the 3D decoder with the ground-truth composition, while DiffCSP uses predictor-corrector sampling for F0.

B.4 Hyper-parameters and Training Details

The experiments vary sampling, confidence ranking, and diffusion schedulers to assess prediction quality, diversity, and candidate selection. DiffCSP consistently benefits from more samples and outperforms baselines across the reported sampling settings.

  • Sampling and evaluation: Match rate increases for all methods as more candidates are sampled, while DiffCSP outperforms baselines at every tested sample count.The comparison uses match rate to evaluate whether at least one generated candidate matches the ground-truth structure.
  • Diversity: DiffCSP and CDVAE are compared on diversity using mean and maximum pairwise CrystalNN fingerprint L2 distances across 20 samples.The diversity comparison is reported for Perov-5 and MP-20.
  • Candidate ranking: Three confidence models rank generated candidates from energy, ground-truth matching probability, or composition–structure contrastive similarity.The models are Energy Predictor, Match Discriminator, and Contrastive Scorer.
  • Candidate ranking: Match Discriminator and Contrastive Scorer perform relatively better, but both remain below the oracle ranker.The oracle uses negative RMSD against the ground truth as the ranking score.
  • Noise schedulers: For lattices, linear scheduling is comparable to cosine scheduling, whereas sigmoid scheduling hinders performance.The original lattice scheduler is cosine with s = 0.008; the fractional-coordinate scheduler uses σmax = 0.5.
  • Noise schedulers: Ab initio generation results compare different atom-type schedulers, with DiffCSP using a cosine scheduler on A.The comparison is summarized in the scheduler ablation table.

E Learning Curves of Different Variants

Learning curves show that multi-graph variants are harder to optimize, while Fourier transformation improves early convergence by preserving periodic translation invariance.

  • Learning curves: Multi-graph variants have higher training and validation losses because disturbed lattices produce significantly varying edge constructions.The varying edges complicate the training procedure.
  • Learning curves: Fourier transformation helps the model converge faster initially by expanding relative coordinates while maintaining periodic translation invariance.The reported benefit concerns early training convergence.

F.1 Implementation Details

The implementation comparison evaluates DiffCSP against USPEX on selected MP-20 compounds using matched candidate counts and reports structure accuracy and inference cost.

  • Experimental setup: The evaluation selects 10 binary and 5 ternary MP-20 compounds and compares 20 candidates from DiffCSP with 20 candidates from USPEX.USPEX uses 20 generations and 20 populations per compound, while DiffCSP samples 20 candidates directly.
  • Results: Overall results summarize match rate, averaged RMSD over matched structures, and averaged inference time across the 15 compounds.The aggregate comparison is provided in Table 10.
  • Results: DiffCSP correctly predicts more structures with higher match rate and significantly lower time cost than USPEX.This is the reported outcome of the selected-compound comparison.
  • Experimental setup: Table 9 reports the minimum RMSD among 20 candidates for USPEX and DiffCSP, with N/A indicating no matching candidate.The table compares per-compound structural matching quality.

G Extension to More General Tasks

DiffCSP is extended beyond fixed-composition CSP to ab initio generation and property optimization by generating atom types and adding energy guidance. These extensions retain the model’s symmetry design and show lower-energy structure search.

  • Task extensions: CSP generates stable structures from fixed composition, while ab initio generation additionally requires generating atom types.DiffCSP first enables atom-type generation to extend beyond the fixed-composition setting.
  • Ab initio generation: DiffCSP supports discrete multinomial and continuous one-hot diffusion approaches for generating atom-type compositions.The two approaches are denoted DiffCSP-D and DiffCSP-C.
  • Ab initio generation: The joint training objective weights lattice and fractional-coordinate losses equally for CSP, while adding an atom-type loss for ab initio generation.The reported settings are λL = λF = 1, λA = 0 for CSP and λL = λF = 1, λA = 20 for ab initio generation.
  • Ab initio generation: The number of atoms is sampled from the training distribution because it remains unchanged during generation, while the number of elements is modeled conditionally.The distribution is expressed as p(M, N) = p(N)p(M|N).
  • Ab initio generation: DiffCSP-C yields better validity and coverage than DiffCSP-D on MP-20, while DiffCSP-D generates structures with more element types than the data distribution.The continuous variant is selected for subsequent experiments.
  • Property optimization: Energy guidance adds gradients from a time-dependent property predictor while preserving periodic E(3) invariance.The predicted energy is periodic E(3) invariant, and the corresponding gradients guide L, F, and A during sampling.
  • Property optimization: Energy-guided sampling uses 100 sampled test structures and evaluates optimization across T′ values from 100 through 1,000.The guidance magnitude is set to s = 50, with T = 1,000 maximum sampling steps.
  • Property optimization: With fixed composition, energy-guided sampling searches novel crystal structures with lower energies than existing structures in the reported VASP evaluation.The comparison visualizes eight cases before and after optimization.

I Error Bars

Additional experiments report error bars for Perov-5 and MP-20, while visualizations compare prediction accuracy, structural diversity, and periodic translation behavior across methods.

  • I Error Bars: Table 15 reports Perov-5 and MP-20 results with error bars, showing similar results to the single-run evaluation in Table 1.The authors obtain these results by performing two additional inferences for each generative method.
  • I Error Bars: DiffCSP provides more accurate predictions than baseline methods in the additional structure visualizations.Figure 11 uses translated structures for clearer visual comparison.
  • I Error Bars: Sixteen generated structures across Perov-5, Carbon-24, and MP-20 illustrate DiffCSP's ability to generate diverse structures.These structures are visualized in Figure 12.
  • I Error Bars: The generation process shows M0 as a periodic translation of the MP-20 ground-truth structure, consistent with the desired periodic translation invariance.Figure 13 includes translated structures with the same atom moved to the origin for comparison.
Loading 2309.04475v2…