[Paper Review] Crystal Structure Prediction by Joint Equivariant Diffusion

Rui Jiao, Wenbing Huang+5 more|arXiv (Cornell University)|Jul 30, 2023

Machine Learning in Materials Science21 citations

TL;DR

DiffCSP introduces a joint, equivariant diffusion model that generates lattice vectors and fractional atom coordinates to predict crystal structures, leveraging periodic E(3) invariance and fractional coordinates to improve CSP efficiency and accuracy over prior methods.

ABSTRACT

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.

Motivation & Objective

• Address the Crystal Structure Prediction (CSP) challenge by respecting the symmetries of crystal geometry (translation, rotation, periodicity).
• Develop a diffusion-based generator that jointly outputs lattice vectors and fractional coordinates of atoms.
• Leverage fractional coordinates and Fourier features to model periodicity and improve diffusion stability.
• Demonstrate that DiffCSP outperforms existing CSP methods in accuracy and efficiency, including ab initio generation extensions.

Proposed method

• Formulate CSP as learning p(L, F | A) where L is the lattice matrix and F contains fractional coordinates within a unit cell.
• Use a joint diffusion process: forward DDPM on L with an O(3)-equivariant denoiser and a score-based, Wrapped Normal diffusion on F to respect periodicity.
• Represent crystals in fractional coordinates to encode periodicity intrinsically and apply Fourier Transform-based message passing to capture periodic relations.
• Employ an equivariant denoising model phi(L, F, A, t) built on EgNN foundations to ensure L-equivalence and F-periodicity in predictions.
• Train with L2 loss on the L denoiser and score-matching loss on the F denoiser; enforce invariance properties via L^TL in the L branch and wrapped normal transitions in the F branch.

Figure 1: (a) $\rightarrow$ (b): The orthogonal transformation of the lattice vectors. (c) $\rightarrow$ (d): The periodic translation of the fractional coordinates. Both cases do not change the structure.

Experimental results

Research questions

• RQ1 Can a joint diffusion process on lattice vectors and fractional coordinates accurately model crystal structures under periodic E(3) invariance?
• RQ2 Does using fractional coordinates plus Fourier-based features improve diffusion-based CSP over Cartesian-based or multi-graph representations?
• RQ3 How does joint L-F diffusion compare with two-stage or separated approaches in terms of accuracy and training stability?
• RQ4 Can the method be extended to ab initio crystal generation by also diffusing atom types A?
• RQ5 What is the empirical performance of DiffCSP against DFT-based methods and other learning-based CSP models on standard datasets?

Key findings

• DiffCSP achieves higher match rates and lower RMSE than competitive baselines across Perov-5, Carbon-24, MP-20, and MPTS-52 datasets.
• Joint diffusion on lattice and fractional coordinates outperforms separated diffusion strategies (L→F or F→L).
• Maintaining O(3) equivariance and periodic translation invariance is critical for performance; removing these degrades results dramatically.
• Using Fourier-transformed relative fractional coordinates (FT) to encode periodicity substantially improves results over non-FT variants.
• DiffCSP substantially reduces computation time compared to DFT-based methods while delivering competitive or superior structure-prediction accuracy.
• DiffCSP can be extended to ab initio crystal generation by additionally modeling atom types A within the diffusion framework.

Figure 2: Overview of DiffCSP. Given the composition ${\bm{A}}$ , we denote the crystal, its lattice and fractional coordinate matrix at time $t$ as ${\mathcal{M}}_{t}$ , ${\bm{L}}_{t}$ and ${\bm{F}}_{t}$ , respectively. The terms ${\bm{\epsilon}}_{\bm{L}}$ and ${\bm{\epsilon}}_{\bm{F}}$ are Gaussia