[논문 리뷰] Equivariant Diffusion for Molecule Generation in 3D

Emiel Hoogeboom, Victor García Satorras외 2명|arXiv (Cornell University)|2022. 03. 31.

Machine Learning in Materials Science인용 수 118

한 줄 요약

E(3) 등변 확산 모델(E(3) equivariant diffusion model, EDM)을 도입하여 3D 좌표와 원자 종류 특징을 함께 노이즈 제거하여 분자를 생성하고, 확률적 가능도 프레임워크를 제공하며 기존의 3D 분자 생성기들에 비해 실험적으로 큰 이점을 보임.

ABSTRACT

This work introduces a diffusion model for molecule generation in 3D that is equivariant to Euclidean transformations. Our E(3) Equivariant Diffusion Model (EDM) learns to denoise a diffusion process with an equivariant network that jointly operates on both continuous (atom coordinates) and categorical features (atom types). In addition, we provide a probabilistic analysis which admits likelihood computation of molecules using our model. Experimentally, the proposed method significantly outperforms previous 3D molecular generative methods regarding the quality of generated samples and efficiency at training time.

연구 동기 및 목표

• 유클리드 대칭(E(3))를 존중하는 3D 분자 생성의 필요성에 대한 동기 제시.
• 확산 프레임워크에서 연속 좌표와 이산 원자 특징을 함께 모델링하는 EDM 제안.
• 제안된 모델 하에서 생성된 분자의 가능도를 계산할 수 있는 tractable 가능도 확보.
• 품질과 효율성 면에서 기존의 3D 분자 생성 접근 방식 대비 실험적 개선 증명.

제안 방법

• 좌표 x와 특징 h에서 작동하는 등변적 노이즈 제거 확산 모델 정의.
• 등변적으로 노이즈 제거를 예측하는 EGNN 기반 동역학 φ를 사용.
• 가능도 목표를 L_t 항들과 ε를 예측하는 실용적 노이즈 매개화를 포함해 형식화.
• 원자 유형을 원-핫 인코딩으로 처리하고 x와 h에 대한 적절한 가능도 항을 적용.
• x에 대한 영(제로) 중심의 중력 제약을 도입하고 특징을 스케일링하여 학습 안정화.
• 조건부 생성을 가능하게 하기 위해 입력에 속성 c를 포함하여 denoising 네트워크를 조정.

실험 결과

연구 질문

• RQ1 E(3) 등변 확산 모델이 고정된 원자 배열 없이 3D 분자 입체구조와 원자 유형을 생성할 수 있는가?
• RQ2 EDM이 QM9에서 E-NF, G-Schnet과 같은 이전 등변 방법에 비해 로그 가능도와 분자 안정성에서 더 나은 성능을 보이는가?
• RQ3 가능도 계산이 가능한 상태로 더 큰 약물 유사 데이터셋으로도 효율적으로 학습 및 확장 가능한가?
• RQ4 연속 좌표와 이산 원자 특징이 공동 확산 과정에 어떻게 통합되는가?
• RQ5 분자 특성에 조건을 부여한 경우 목표 분자 생성을 위한 영향은 무엇인가?

주요 결과

지표	E-NF	G-Schnet	GDM	GDM-aug	EDM (ours)
NLL	-59.7	N.A	-94.7	-92.5	-110.7 ±1.5
원자 안정성 (%)	85.0	95.7	97.0	97.6	98.7 ±0.1
분자 안정성 (%)	4.9	68.1	63.2	71.6	82.0 ±0.4

• EDM은 바탕선 E-NF (-59.7), G-Schnet (N.A.), GDM (-94.7), GDM-aug (-92.5)보다 불활성 음의 로그 가능도(NLL)가 더 유리하게 나타남(-110.7 ± 1.5).
• 원자 안정성은 EDM이 98.7 ± 0.1%로 바탕선의 85.0%, 95.7%, 97.0%, 97.6%보다 높다.
• 분자 안정성은 EDM이 82.0 ± 0.4%로 바탕선의 4.9%, 68.1%, 63.2%, 71.6%보다 높다.
• EDM은 3D에서 더 많은 유효하고 고유한 분자를 생성하며(표 2 결과: 97.5 ± 0.2% 유효, H를 포함한 경우 94.3 ± 0.2% 유효 및 고유).
• EDM은 GEOM-Drugs와 같은 더 큰 데이터셋으로도 더 효율적으로 학습하고 확장 가능하며, 기존의 3D 확산 및 정규화 흐름 접근법을 능가한다.