[論文レビュー] Crystal Structure Prediction by Joint Equivariant Diffusion

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

Machine Learning in Materials Science被引用数 21

ひとこと要約

DiffCSPは、結晶構造を予測するために格子ベクトルと分数座標を同時に生成する等変拡散モデルを導入し、周期的なE(3)不変性と分数座標を活用して、従来の方法よりCSPの効率と精度を向上させる。

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.

研究の動機と目的

• 結晶幾何学の対称性（平行移動、回転、周期性）を尊重して、Crystal Structure Prediction (CSP) の課題に対処する。
• 格子ベクトルと原子の分数座標を同時に出力する拡散ベースの生成器を開発する。
• 分数座標とフーリエ特徴を活用して周期性をモデル化し、拡散の安定性を向上させる。
• DiffCSPが、ab initio生成の拡張を含む、精度と効率の点で既存のCSP手法を上回ることを示す。

提案手法

• CSPを、Lが格子行列、Fが単位胞内の分数座標を含むp(L, F | A)を学習する形で定式化する。
• ジョイント拡散過程を用いる：L上での前向きDDPMはO(3)-等変デノイザー、Fには周期性を尊重するスコアベースのWrapped Normal拡散を適用。
• 結晶を分数座標で表現して周期性を内在的にエンコードし、 Fourier Transform-basedなメッセージパッシングを適用して周期的な関係を捉える。
• 予測におけるL同値性とF周期性を保証するため、EgNNを基盤とした同変デノイングモデルphi(L, F, A, t)を採用する。
• Lデノイザーに対するL2損失とFデノイザーに対するスコア整合損失で訓練する；LブランチではL^TLを介して不変性を、FブランチではWrapped Normal遷移を介して不変性を強制する。

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.

実験結果

リサーチクエスチョン

• 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?

主な発見

• DiffCSPは、Perov-5, Carbon-24, MP-20, MPTS-52データセットにおいて競合ベースラインより高い一致率と低いRMSEを達成する。
• 格子と分数座標の共同拡散は、分離した拡散戦略（L→FまたはF→L）より優れている。
• O(3)等変性と周期的平行移動不変性を維持することが性能にとって重要であり、これらを除くと結果が著しく劣化する。
• 周期性をエンコードするための傅立変換付き相対分数座標（FT）を使用すると、非FTバリアントより結果が大幅に改善される。
• DiffCSPは、DFTベースの方法と比較して計算時間を大幅に削減しつつ、構造予測の精度は競合するか優れている。
• DiffCSPは、拡散フレームワーク内で原子種Aを追加的にモデリングすることによりab initio結晶生成へ拡張可能である。

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