[论文解读] 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.

研究动机与目标

• 通过遵循晶体几何对称性（平移、旋转、周期性）来解决晶体结构预测（CSP）中的挑战。
• 开发一个扩散生成器，联合输出晶格向量和原子的分数坐标。
• 利用分数坐标和傅里叶特征来建模周期性并提高扩散的稳定性。
• 证明 DiffCSP 在准确性和效率方面优于现有的 CSP 方法，包括从头生成的扩展。

提出的方法

• 将 CSP 表述为学习 p(L, F | A) 的问题，其中 L 是晶格矩阵，F 包含单胞内的分数坐标。
• 采用联合扩散过程：对 L 进行正向 DDPM 的去噪器为 O(3) 等变性，对 F 进行基于分数的 Wrapped Normal 扩散以遵循周期性。
• 以分数坐标表示晶体以本质上编码周期性，并应用基于傅里叶变换的信息传递来捕捉周期关系。
• 采用建立在 EgNN 基础之上的等变去噪模型 phi(L, F, A, t)，以确保在预测中的 L 等价性和 F 的周期性。
• 用 L2 损失对 L 的去噪器进行训练，以及对 F 的去噪器进行分数匹配损失训练；通过 L^T L 在 L 分支中强制不变性，通过 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 一个联合扩散过程在晶格向量和分数坐标上是否能够在周期性 E(3) 不变性下准确建模晶体结构？
• RQ2 使用分数坐标结合基于傅里叶的特征是否能在基于扩散的 CSP 中优于基于笛卡尔坐标或多图表示的方法？
• RQ3 在精度和训练稳定性方面，联合 L-F 扩散与两阶段或分离方法相比表现如何？
• RQ4 该方法能否扩展为从头（ab initio）晶体生成，即在扩散框架中也对原子类型 A 进行扩散？
• RQ5 DiffCSP 相对于基于 DFT 的方法和其他学习型 CSP 模型在标准数据集上的经验性能如何？

主要发现

• DiffCSP 在 Perov-5、Carbon-24、MP-20 和 MPTS-52 数据集上达到更高的匹配率和更低的 RMSE，相较于竞争基线。
• 晶格和分数坐标的联合扩散优于分离的扩散策略（L→F 或 F→L）。
• 维持 O(3) 等变性和周期性平移不变性对性能至关重要；移除这些会显著降低结果。
• 使用傅里叶变换的相对分数坐标（FT）来编码周期性在结果上显著优于非 FT 的变体。
• 与基于 DFT 的方法相比，DiffCSP 在计算时间上显著降低，同时仍提供具有竞争力或更优的结构预测准确性。
• 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