[논문 리뷰] Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu, Alexander K. Powers외 3명|arXiv (Cornell University)|2023. 05. 02.

Model Reduction and Neural Networks인용 수 32

한 줄 요약

GeoLDM은 3D 분자 생성을 위해 불변 스칼라와 등가 텐서로 구성된 SE(3)-등가 잠재 공간을 갖춘 잠재 확산 프레임워크를 도입하며, 이전 방법들보다 더 높은 타당성(validity)과 제어 가능성을 달성합니다.

ABSTRACT

Generative models, especially diffusion models (DMs), have achieved promising results for generating feature-rich geometries and advancing foundational science problems such as molecule design. Inspired by the recent huge success of Stable (latent) Diffusion models, we propose a novel and principled method for 3D molecule generation named Geometric Latent Diffusion Models (GeoLDM). GeoLDM is the first latent DM model for the molecular geometry domain, composed of autoencoders encoding structures into continuous latent codes and DMs operating in the latent space. Our key innovation is that for modeling the 3D molecular geometries, we capture its critical roto-translational equivariance constraints by building a point-structured latent space with both invariant scalars and equivariant tensors. Extensive experiments demonstrate that GeoLDM can consistently achieve better performance on multiple molecule generation benchmarks, with up to 7\% improvement for the valid percentage of large biomolecules. Results also demonstrate GeoLDM's higher capacity for controllable generation thanks to the latent modeling. Code is provided at \url{https://github.com/MinkaiXu/GeoLDM}.

연구 동기 및 목표

• roto-translation 등가성을 존중하는 기하학적 인식 잠재 공간에서 3D 분자 생성을 발전시킨다.
• Invariant와 equivariant 구성요소를 모두 갖춘 잠재 코드를 생성하는 기하학적 자동인코더를 도입한다.
• 잠재 공간에서 확산 모델을 학습시켜 3D 분자의 타당성, 다양성 및 제어 가능한 생성을 개선한다.
• 잠재 GeoLDM이 QM9 및 DRUG 벤치마크에서 생성 지표를 향상시키고 제어 가능한 생성을 가능하게 함을 보여준다.
• 복제 및 추가 연구를 위한 공개 코드를 제공한다.

제안 방법

• EGNN(Eequivariant Graph Neural Networks)을 사용하여 각 노드의 불변(z_h) 및 등가(z_x) 구성요소를 갖는 잠재 공간으로 매핑하는 기하학적 자동인코더를 구성한다.
• 잠재 전이를 등가적으로 설계하고 잠재 사전치를 불변으로 만들어 결과적 주변 분포가 SE(3)-불변이 되도록 SE(3) 불변성을 보장한다.
• time-conditional EGNN으로 구현된 denoising 네트워크 epsilon_theta를 사용한 구조화된 잠재 공간에서 잠재 확산 모델(LDM)을 학습시킨다.
• 두 단계 학습을 채택한다: 먼저 재구성과 정규화를 통해 기하학적 AE를 학습하고, 그 다음 인코더를 고정하고 잠재 DM을 학습한다; 생성을 위해 p_theta(z_x,z_h)를 사용하고 디코딩은 p_xi(x,h|z_x,z_h)로 수행한다.
• 노이즈 제거기와 인코더/디코더를 속성 벡터 s(연결(concatenation) 방식, 교차 주의(attention) 대신)으로 조건화하여 제어 가능한 생성을 지원한다.
• Optionally 학습 분포 p(N)에서 분자 크기 N을 샘플링하고 그에 따라 생성한다.

Figure 1: Illustration of GeoLDM . The encoder ${\mathcal{E}}_{\phi}$ encodes molecular features ${\mathbf{x}},{\mathbf{h}}$ into equivariant latent variables ${\mathbf{z}}_{\textnormal{x}},{\mathbf{z}}_{\textnormal{h}}$ , and the latent diffusion transitions $q({\mathbf{z}}_{{\textnormal{x}},t},{\m

실험 결과

연구 질문

• RQ1 기하학 인식 잠재 공간(불변 및 등가 구성 요소)을 다루는 잠재 확산 모델이 3D 분자 생성 품질을 향상시킬 수 있는가?
• RQ2 잠재 공간에서 SE(3) 등가성을 강제하는 것이 원자 피처 공간에서 작동하는 잠재 모델에 비해 화학적 타당성과 구조적 다양성을 높이는가?
• RQ3 GeoLDM이 목표 화학적 속성을 위한 제어 가능한 생성을 얼마나 가능하게 하며 표준 3D 분자 벤치마크에서의 성능은 어떠한가?
• RQ4 QM9 및 DRUG 데이터셋에서 무조건적 생성과 조건부 생성 간의 GeoLDM의 성능 차이는 어떠한가?

주요 결과

데이터	원자 상태 (%)	분자 상태 (%)	유효성 (%)	유효성 및 고유성 (%)	원자 상태 (%)	유효성 (%)
QM9	ENF	85.0	4.9	40.2	39.4	G-Schnet	95.7	68.1	85.5	80.3	GDM	97.0	63.2	-	-
QM9	GDM-aug	97.6	71.6	90.4	89.5	EDM	98.7	82.0	91.9	90.7
QM9	EDM-Bridge	98.8	84.6	92.0*	90.7	GraphLDM	97.2	70.5	83.6	82.7
QM9	GraphLDM-aug	97.9	78.7	90.5	89.5	GeoLDM	98.9	89.4	93.8	92.7
DRUG	GeoLDM	84.4	99.3

• GeoLDM은 벤치마크에서 더 높은 생성 품질을 달성하며, 대형 생물분자에서 baselines 대비 최대 7% 높은 타당성(valid rate)을 달성한다.
• 불변 및 등가 잠재 변수를 모두 사용하는 잠재 모델링은 스칼라만 사용하는 모델에 비해 타당성과 다양성이 우수하다.
• 조건부 생성 실험은 잠재 공간 설계와 연결(concatenation)에 의한 속성 조건화의 간단함으로 인해 제어 가능성이 향상되었음을 보여준다.
• GeoLDM은 학습된 주변 분포에서 SE(3)-불변 특성이 강하게 나타나 회전/병진에 대한 일반화에 도움을 준다.
• 실증 결과는 GeoLDM이 다수의 지표(타당성, 안정성 및 타당성×고유성)에서 기존의 몇몇 3D 분자 생성 방법을 능가함을 보여준다.
• 저자들은 재현과 추가 연구를 위한 프로젝트 사이트의 코드를 제공한다.

Figure 2: Molecules generated by GeoLDM trained on QM9 (left three) and DRUG (right four).