[논문 리뷰] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation

Han Huang, Leilei Sun외 2명|arXiv (Cornell University)|2023. 05. 21.

Computational Drug Discovery Methods인용 수 5

한 줄 요약

이 논문은 분자의 원자 종류, 형식적 전하, 결합 유형 및 3D 좌표를 함께 생성하기 위해 분자의 그래프와 기하학적 구조 간의 상관관계를 모델링하는 새로운 Diffusion Graph Transformer를 활용한 통합 2D 및 3D 확산 모델인 JODO를 제안한다. 이 모델은 QM9와 GEOM-Drugs에서 무조건적 및 조건부 생성에서 최신 기술 수준(SOTA) 성능을 달성하여 타당성, FCD 및 스caffolding 유사도 측면에서 베이스라인을 초월하며, 빠른 소수단 샘플링과 역분자 설계를 가능하게 한다.

ABSTRACT

Designing new molecules is essential for drug discovery and material science. Recently, deep generative models that aim to model molecule distribution have made promising progress in narrowing down the chemical research space and generating high-fidelity molecules. However, current generative models only focus on modeling either 2D bonding graphs or 3D geometries, which are two complementary descriptors for molecules. The lack of ability to jointly model both limits the improvement of generation quality and further downstream applications. In this paper, we propose a new joint 2D and 3D diffusion model (JODO) that generates complete molecules with atom types, formal charges, bond information, and 3D coordinates. To capture the correlation between molecular graphs and geometries in the diffusion process, we develop a Diffusion Graph Transformer to parameterize the data prediction model that recovers the original data from noisy data. The Diffusion Graph Transformer interacts node and edge representations based on our relational attention mechanism, while simultaneously propagating and updating scalar features and geometric vectors. Our model can also be extended for inverse molecular design targeting single or multiple quantum properties. In our comprehensive evaluation pipeline for unconditional joint generation, the results of the experiment show that JODO remarkably outperforms the baselines on the QM9 and GEOM-Drugs datasets. Furthermore, our model excels in few-step fast sampling, as well as in inverse molecule design and molecular graph generation. Our code is provided in https://github.com/GRAPH-0/JODO.

연구 동기 및 목표

• 기존의 2D 분자 그래프나 3D 기하학을 별개로 모델링하는 생성 모델의 한계를 해결하기 위해, 이는 생성 품질과 후속 응용에 악영향을 미친다.
• 2D 결합 그래프와 3D 좌표를 동시에 모델링하는 통합 생성 프레임워크를 개발하여 상호 보완적인 분자 기술적 특징을 포착한다.
• 분자 구조와 기하학의 연합 분포를 학습하여 샘플의 정밀도를 향상시키고 효율적인 역분자 설계를 가능하게 한다.
• 기하학적 대칭성(SE(3)-불변성)과 그래프 순열 불변성을 유지하는 연속 시간 확산 프로세스를 설계한다.
• 후처리 또는 추가 시뮬레이션에 의존하지 않고도 빠르고 고품질의 샘플링과 양자 성질 기반 조건부 생성을 가능하게 한다.

제안 방법

• 2D 그래프와 3D 좌표 성분을 동시에 손상시키는 통합 노이즈 스케줄을 사용한 연속 시간 확산 모델을 제안한다.
• 관계 기반 어텐션 메커니즘을 활용하여 노드 및 엣지 표현을 동시에 업데이트하는 Diffusion Graph Transformer를 도입하며, 스칼라 특징과 기하 벡터를 통합한다.
• 노이즈 예측 대신 데이터 예측 모델을 사용하고, 자기 조건부 조건을 적용하여 이산적 구조 복구를 향상시킨다. 이는 마지막 샘플링 단계의 출력을 보조 조건으로 사용한다.
• 메시지 전파 아키텍처를 사용하여 데이터 예측 모델을 파arameter화하며, 분자 그래프 전반에 걸쳐 스칼라 및 3D 벡터 특징을 전파하고 업데이트한다.
• 기하학적 등변성(SE(3)-불변성)과 그래프 순열 불변성을 확보하기 위해 등변 및 불변 신경망 설계를 적용한다.
• 목표 양자 성질에 조건부로 생성되도록 모델을 확장하여 성질 목표 지향 분자 생성을 가능하게 한다.

실험 결과

연구 질문

• RQ1 2D와 3D를 동시에 모델링하는 통합 확산 모델은 단지 2D 그래프나 3D 좌표만 생성하는 모델보다 더 높은 정밀도로 완전한 분자를 생성할 수 있는가?
• RQ2 역소거 노이즈 과정 중에 분자의 위상 구조와 3D 기하학 간의 상관관계를 통합된 확산 프레임워크가 얼마나 효과적으로 모델링할 수 있는가?
• RQ3 관계 기반 어텐션과 기하 벡터 업데이트를 갖춘 제안된 Diffusion Graph Transformer는 생성 품질과 샘플링 효율성을 향상시키는가?
• RQ4 모델은 무조건적 및 조건부(성질 목표 지향) 분자 생성에서 최신 기술 수준 성능을 달성할 수 있는가?
• RQ5 기존의 베이스라인과 비교해 볼 때, 모델은 소수단 샘플링과 양자 성질을 위한 역설계에서 얼마나 우수한 성능을 보이는가?

주요 결과

• QM9 데이터셋에서 JODO-2D는 99.91%의 타당성(valency 검사 미적용)을 기록하여 CDGS(98.13%)와 DiGress(92.10%)를 포함한 모든 베이스라인을 앞서며 우수한 성능을 보였다.
• QM9에서 JODO-2D는 FCD 점수 0.472를 기록하여 CDGS(2.069)와 DiGress(3.597)보다 유의미하게 낮아, 분포 매칭 측면에서 열등한 성능을 보였다.
• MOSES 데이터셋에서 JODO-2D는 88.9%의 타당성, 100%의 유일성, 91.0%의 신선도를 기록하여 DiGress(85.7% 타당성, 95.0% 신선도)를 초월했으며, 다른 강력한 베이스라인과 유사하거나 이를 능가하는 성능을 보였다.
• 모델은 효과적인 양자 성질 기반 분자 설계를 위한 강력한 성능을 보이며, 역분자 설계에 있어 뛰어난 성능을 발휘했다.
• JODO는 고품질 출력을 제공하는 빠른 소수단 샘플링을 가능하게 하여 실생활 응용에서의 실용성을 입증했다.
• 제거 분석 결과, 자기 조건부 조건을 적용한 제안된 데이터 예측 모델과 Diffusion Graph Transformer가 생성 품질과 구조적 정밀도 향상에 크게 기여하는 것으로 확인되었다.