[논문 리뷰] Crystal Structure Prediction by Joint Equivariant Diffusion

Rui Jiao, Wenbing Huang외 5명|arXiv (Cornell University)|2023. 07. 30.

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 방법보다 우수함을 입증하며, ab initio 생성 확장도 포함한다.

제안 방법

• CSP를 p(L, F | A)를 학습하는 문제로 공식화한다. 여기서 L은 격자 행렬이고 F는 단위 셀 내의 분수 좌표를 포함한다.
• 공동 확산 과정 사용: L에 대해 O(3)-등변 노이저를 갖는 순방향 DDPM과 F에 대해 주기성을 존중하는 스코어 기반 Wrapped Normal 확산을 사용한다.
• 주기성을 고유하게 인코딩하기 위해 결정 구조를 분수 좌표로 표현하고, 주기적 관계를 포착하기 위해 푸리에 변환 기반 메시지 전달을 적용한다.
• L-동등성 및 F-주기성을 예측에서 보장하기 위해 EgNN 기반 위에 구축된 등변적 노이즈 제거 모델 phi(L, F, A, t)을 사용한다.
• L denoiser에 대한 L2 손실과 F denoiser에 대한 스코어 매칭 손실로 학습한다; 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 격자 벡터와 분수 좌표에 관한 공동 확산 프로세스가 주기적 E(3) 불변성 하에서 결정 구조를 정확하게 모델링할 수 있는가?
• RQ2 분수 좌표와 푸리에 기반 특징을 사용하는 것이 직교 좌표 기반 또는 다중 그래프 표현에 비해 확산 기반 CSP를 향상시키는가?
• RQ3 공동 L-F 확산이 두 단계 또는 분리된 접근법과 정확도 및 학습 안정성 측면에서 어떻게 비교되는가?
• RQ4 원자 타입 A도 확산시켜 ab initio 결정 생성으로 방법을 확장할 수 있는가?
• RQ5 표준 데이터세트에서 DFT 기반 방법 및 다른 학습 기반 CSP 모델에 대한 DiffCSP의 실험적 성능은 어떤가?

주요 결과

• 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