[논문 리뷰] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

Minkai Xu, Lantao Yu외 4명|arXiv (Cornell University)|2022. 03. 06.

Machine Learning in Materials Science인용 수 173

한 줄 요약

GeoDiff는 SE(3)-등가 확산 프레임워크로 원자 좌표를 직접 모델링하여 분자 구성(conformations)을 생성하며, 특히 큰 분자에서 구성 생성에 있어 최첨단 성능을 달성합니다.

ABSTRACT

Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GeoDiff for molecular conformation prediction. GeoDiff treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto-translational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-to-end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GeoDiff is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

연구 동기 및 목표

• 분자가 직접 3D 원자 좌표를 출력하는 확산 기반 생성 모델의 동기 부여와 개발.
• SE(3) 등가 설계를 통해 회전-병진 불변성과 CoM 비가변 초기화를 보장합니다.
• 분자 그래프를 조건으로 하는 조건부 생성을 위한 맞춤형 ELBO 목표를 갖춘 엔드 투 엔드 학습 가능한 프레임워크를 제공합니다.

제안 방법

• 고정 분산 스케줄을 가진 3D 좌표 C^t의 순방향 확산을 정의하여 C^T ~ N(0, I)가 되도록 합니다.
• 분자 그래프 G에 조건화하면서 점진적으로 노이즈를 제거하는 조건부 역 마르코프 체인 p_θ(C^{t-1}|G, C^t)을 정의합니다.
• SE(3) 등가 신경망(Graph Field Network)을 사용하여 반대 방향의 평균 μ_θ를 매개화하고 회전-병진 등가를 보장합니다.
• 양의 중심(CoM) 없이 불변 초기 밀도와 CoM-프리 마르코프 커널을 사용하여 평행 이동 불변성을 보장합니다; μ_θ와 네트워크 아키텍처를 통해 회전 등가성을 강제합니다.
• 가중 ELBO를 통해 ε와 ε_θ 간의 가중된 L2 손실로 축약되도록 학습하며, 등가성 강제를 위한 두 가지 전략(정렬 기반 및 체인 규칙 접근)을 적용합니다.
• C^T ~ N(0, I)에서 시작하여 p_θ(C^{t-1}|G, C^t)를 사용해 반복적으로 C^{t-1}를 샘플링하는 샘플링 알고리즘을 제공합니다.

실험 결과

연구 질문

• RQ1 3D 좌표에서 직접 작동하는 확산 모델이 불변 기하 중간체에 의존하는 모델과 비교해 경쟁력 있거나 우수한 구성 생성 성능을 달성할 수 있을까요?
• RQ2 기하 생성을 위한 학습과 추론 단계에서 회전-병진 불변성을 어떻게 보장할 수 있을까요?
• RQ3 분자 구성의 역 확산 단계를 모델링하기 위한 효과적인 매개변수화와 아키텍처(예: SE(3)-등가 네트워크)는 무엇인가요?
• RQ4 엔드-투-엔드 좌표 기반 확산 생성이 더 큰 약물과 같은 분자에 확장되고 다양하고 고품질의 구성들을 제공할 수 있나요?

주요 결과

• GeoDiff는 GEOM-Drugs에서 최첨단 ML 기준선보다 일관되게 우수하며 COV-R 및 MAT-R 지표에서 주목할 만한 개선을 보입니다.
• 큰 분자에 대해 baselines보다 더 높은 COV-R와 더 낮은 MAT-R를 달성하여 강한 커버리지와 정확성을 보여줍니다.
• GeoDiff-C는 일반적으로 GeoDiff-A보다 약간 더 나은 최적화를 제공하여 체인 규칙 기반 학습이 학습을 더 잘 이끈다는 것을 시사합니다.
• 이 방법은 생성된 구성의 다양성과 품질을 크게 향상시키며, 특히 큰 분자가 테스트되는 Drugs 데이터셋에서 그 효과가 두드러집니다.
• 정성적 샘플은 GeoDiff가 국소적 및 전역 구조 특징을 모두 효과적으로 포착함을 보여줍니다.