[논문 리뷰] MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation

Xingang Peng, Jiaqi Guan외 2명|arXiv (Cornell University)|2023. 05. 11.

Machine Learning in Materials Science인용 수 11

한 줄 요약

MolDiff는 Bond 가이던스가 있는 E(3)-동등 프레임워크를 사용하여 3D 분자에서 원자와 결합을 함께 확산시키고, 원자-결합 불일치를 해결하며 이전 모델보다 생성 품질이 높게 달성합니다.

ABSTRACT

Deep generative models have recently achieved superior performance in 3D molecule generation. Most of them first generate atoms and then add chemical bonds based on the generated atoms in a post-processing manner. However, there might be no corresponding bond solution for the temporally generated atoms as their locations are generated without considering potential bonds. We define this problem as the atom-bond inconsistency problem and claim it is the main reason for current approaches to generating unrealistic 3D molecules. To overcome this problem, we propose a new diffusion model called MolDiff which can generate atoms and bonds simultaneously while still maintaining their consistency by explicitly modeling the dependence between their relationships. We evaluated the generation ability of our proposed model and the quality of the generated molecules using criteria related to both geometry and chemical properties. The empirical studies showed that our model outperforms previous approaches, achieving a three-fold improvement in success rate and generating molecules with significantly better quality.

연구 동기 및 목표

• 3D 분자 확산 생성에서 원자-결합 불일치를 해결하는 것을 동기화합니다.
• 화학적 일관성을 유지하기 위해 원자와 결합을 동시에 샘플링하는 확산 모델을 제안합니다.
• 노드-링크 표현을 함께 업데이트하는 E(3)-등가 그래프 신경망을 설계합니다.
• 학습과 생성의 안정을 위한 Bond-먼저 노이즈 스케줄을 도입합니다.
• 생성된 분자의 기하학, 약물가능성, 구조적 합리성을 포착하는 평가 지표를 개발합니다.

제안 방법

• 3D 분자를 A,R,B로 구성된 M={A,R,B}로 모델링합니다(원자 타입 A, 좌표 R, 결합 B).
• 흡수형 노이즈를 갖는 원자 위치와 이산적 원자/결합 타입을 교란하는 순방향 확산 프로세스를 사용합니다.
• 원자와 결합에 대해 pθ(M^{t-1}|M^{t}) 를 예측하는 E(3)-등가 신경망으로 역확산을 수행합니다.
• 원자 및 결합 타입을 이산 변수로 취급하고 Bond-먼저 스케줄로 결합 및 원자 확산을 분리합니다.
• Bond 예측기가 로그 C의 기울기를 통해 원자 위치 잡음 제거를 안내하도록 하며, C는 Bond-예측기 신뢰도입니다.
• 노드(Vertex)와 에지(Edge) 표현을 업데이트하여 원자 타입과 결합 타입을 예측하는 E(3)-등가 메시지 전달 방식을 채택합니다(Eq. 4).
• 손실 L^{t-1}=L_pos^{t-1}+λ1 L_atom^{t-1}+λ2 L_bond^{t-1}으로 위치, 원자 타입, 결합 타입의 노이즈 제거 간의 균형을 맞추며 학습합니다.

실험 결과

연구 질문

• RQ1 원자와 결합의 공동 확산이 나중에 처리되는 결합 추가보다 3D 분자의 현실성 및 타당성을 향상시키나요?
• RQ2 Bond-먼저 확산 스케줄이 원자-결합 불일치를 줄이고 생성 중 화학적 토폴로지를 향상시키나요?
• RQ3 원자와 결합을 모두 업데이트하는 E(3)-등가 아키텍처가 생성 품질에 어떤 영향을 미치나요?
• RQ4 확산 기반 생성에서 기하학, 약물가능성, 합리성을 가장 잘 포착하는 지표는 무엇인가요?

주요 결과

타당성	연결성	성공률	참신성	다양성	고유성	유사성 값
0.997	0.996	0.993	0.972	0.769	0.986	0.634
EDM	0.447	0.830	0.371	1.000	0.729	1.000	0.441
Predict bond (Lookup table)	0.297	0.937	0.278	1.000	0.751	0.999	0.434
Predict bond (NN Predictor)	0.956	0.978	0.935	0.984	0.768	0.999	0.592
Predict bond (Openbabel)	0.992	0.969	0.961	0.986	0.762	0.999	0.576
Conti. diffusion	0.914	0.885	0.808	0.999	0.752	1.000	0.513
Conti. diffusion (scaling)	0.972	0.934	0.907	0.997	0.759	1.000	0.511
Add bond length loss	0.998	0.959	0.958	0.986	0.752	0.998	0.590

• MolDiff는 EDM에 비해 생성 성능을 크게 향상시키며 더 높은 타당성, 연결성, 거의 완벽한 성공률을 달성합니다.
• 원자 및 결합 타입의 이산 확산이 연속 확산 및 결합 할당의 후처리보다 다수의 기준에서 우수합니다.
• Bond-먼저 확산 스케줄은 결합 타입이 원자보다 우선적으로 확산되도록 하여 훈련을 안정시키고 이후 원자 정제를 통해 결합 일관성을 개선합니다.
• 원자 위치 잡음 제거 중 Bond 예측기 가이드 기울기를 도입하면 결합 길이 정확도와 전반적 분자 기하가 향상됩니다.
• MolDiff는 EDM에 비해 더 우수한 3D 기하 지표(더 낮은 RMSD, 더 나은 결합 각/이성질체 분포)와 더 현실적인 결합 타입 및 고리 통계를 제공합니다.
• MolDiff는 기준 EDM 대비 성공률을 세 배 향상시키고 약물가능성과 구조적 품질이 개선된 분자를 생성합니다.