Crystal Structure Prediction by Joint Equivariant Diffusion

TL;DR

This work tackles Crystal Structure Prediction by introducing DiffCSP, a joint diffusion model operating on lattice vectors and fractional coordinates to respect periodic E(3) invariances. By using fractional coordinates and Fourier-based features within an equivariant GNN framework, DiffCSP effectively models crystal geometry and symmetry, outperforming DFT-based methods in accuracy and efficiency. The approach is validated across multiple datasets and extended to ab initio crystal generation with energy-guided refinement, highlighting practical potential for rapid, reliable CSP. The work also provides thorough ablations, underscoring the importance of joint diffusion, invariance properties, and periodic representations.

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.

Figures (13)

• 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.
• 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 Gaussian noises, $\hat{{\bm{\epsilon}}}_{\bm{L}}$ and $\hat{{\bm{\epsilon}}}_{\bm{F}}$ are predicted by the denoising model $\phi$.
• Figure 3: Visualization of the predicted structures from different methods. We select the structure of the lowest RMSE over 20 candidates. We translate the same predicted atom by all methods to the origin for better comparison. Our DiffCSP accurately delivers high quality structure predictions.
• Figure 4: An example of periodic translation invariance. From the view of a unit cell, the atoms translated across the right boundary will be brought back to the left side.
• Figure 5: Overview of the original (a,b) and adapted (c,d) CDVAE. The key adaptations lie in two points. (1) We introduce an additional 1D prior encoder to fit the latent distribution of the given composition. (2) We initialize the generation procedure of the 3D decoder with the ground truth composition and keep the atom types unchanged to ensure the generated structure conforms to the given composition.

Theorems & Definitions (15)

• Definition 1: Permutation Invariance
• Definition 2: O(3) Invariance
• Definition 3: Periodic Translation Invariance
• Proposition 1
• Proposition 2
• Proposition 3
• Definition 4
• Lemma 1: xu2021geodiff
• proof
• Proposition 3