[논문 리뷰] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation

Clément Vignac, Nagham Osman외 2명|arXiv (Cornell University)|2023. 02. 17.

Machine Learning in Materials Science인용 수 10

한 줄 요약

MiDi는 엔드-투-엔드 확산을 통해 분자 그래프와 3D 컨포머를 함께 생성하며, 이전 방법에 비해 GEOM-DRUGS에서 안정적이고 유효한 분자 생성 비율이 높다.

ABSTRACT

This work introduces MiDi, a novel diffusion model for jointly generating molecular graphs and their corresponding 3D arrangement of atoms. Unlike existing methods that rely on predefined rules to determine molecular bonds based on the 3D conformation, MiDi offers an end-to-end differentiable approach that streamlines the molecule generation process. Our experimental results demonstrate the effectiveness of this approach. On the challenging GEOM-DRUGS dataset, MiDi generates 92% of stable molecules, against 6% for the previous EDM model that uses interatomic distances for bond prediction, and 40% using EDM followed by an algorithm that directly optimize bond orders for validity. Our code is available at github.com/cvignac/MiDi.

연구 동기 및 목표

• 그래프 구조와 3D 구성을 함께 포착하는 분자의 엔드-투-엔드 생성을 촉진한다.
• 그래프와 좌표를 동시에 디노이즈하는 확산 기반 모델을 제안한다.
• 새로운 rEGNN-강화 트랜스포머 아키텍처로 3D-등방성 처리(3D- Equivariant 처리)를 향상시킨다.
• GEOM-DRUGS와 QM9 데이터셋에서 무조건 생성 성능이 우수함을 입증한다.

제안 방법

• 원자 유형, 전하를 노드 특징으로 하는 3D에 임베딩된 그래프로 분자를 표현하고, 결합 유형을 엣지 특징으로 사용한다.
• 좌표(연속)와 이산(원자 유형, 전하, 결합)을 혼합한 확산 프로세스를 구성요소별 잡음 일정으로 사용한다.
• 좌표와 결합 유형을 원자 유형과 전하보다 우선하도록 적응형 잡음 스케줄을 도입한다.
• 평행이동 불변이 아닌 특징을 확장한 Relaxed Equivariant Graph Neural Network (rEGNN) 레이어를 개발한다.
• SE(3) 등방성을 유지하는 Transformer 기반 디노이징 네트워크에 rEGNN을 도입한다.
• 좌표 회귀와 이산 특징에 대한 교차 엔트로피를 결합한 손실로 학습한다.

Figure 1 : Samples from our model. MiDi generates simultaneously a 2D graph structure and a 3D conformation that is consistent with this structure.

실험 결과

연구 질문

• RQ1 확산 모델이 엔드-투-엔드 분화 가능하게 분자 그래프와 3D 좌표를 함께 학습할 수 있는가?
• RQ2 적응형 구성요소별 잡음 스케줄이 생성 분자의 품질과 안정성을 향상시키는가?
• RQ3 완화된 SE(3)-등방성 아키텍처(rEGNN)가 표준 EGNN보다 그래프+컨포머 생성을 향상시키는가?
• RQ4 QM9와 GEOM-DRUGS에서 OpenBabel 기반의 포스트-프로세싱 및 기존 확산 기반 방법과 비교했을 때 MiDi의 무조건 분자 생성 성능은 어떠한가?

주요 결과

• MiDi는 GEOM-DRUGS에서 안정 분자 비율이 높은 성과를 보이며 EDM 기반 방법 및 Open Babel 포스트-프로세싱을 능가한다. (Adaptive MiDi: Mol stable 91.6%, At stable 99.8%, Validity 77.8%, Connected 100.0%).
• QM9에서 MiDi는 그래프 기반 지표에서 기본 EDM 모델을 능가하지만, 간단한 데이터셋에서는 Open Babel 최적화가 여전히 강력하다.
• MiDi는 2D 그래프 구조와 3D 컨포머를 함께 엔드투엔드 미분 가능하게 생성하여 별도의 결합 예측 단계 없이 구현한다.
• 적응형 잡음 스케줄은 샘플 품질을 개선하고 좌표와 결합 유형을 원자 유형과 전하보다 먼저 우선하도록 모델을 유도한다.
• MiDi는 대규모 약물 유사 데이터셋(GEOM-DRUGS)으로 확장 가능하며 사후 결합 예측에 의존하는 방법보다 더 강한 엔드-투-엔드 성능을 보여준다.

Figure 2 : The noise schedule is tuned separately for each component. Atom coordinates and bond types are denoised earlier during sampling, while atom types and formal charges are updated later in the process. Experimentally, the adaptive schedule allows to obtain better 3D conformers and more stabl