[논문 리뷰] A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining

Shengchao Liu, Weitao Du외 3명|arXiv (Cornell University)|2023. 05. 28.

Computational Drug Discovery Methods인용 수 6

한 줄 요약

이 논문은 2D 구조와 3D 기하학적 구조에서의 공동 표현을 학습함으로써 다중 모odal 전훈을 가능하게 하는 군 대칭 확률 미분 방정식 모델인 MoleculeSDE를 제안한다. 입력 공간에서 직접 2D 그래프에서 3D 구형 구조로, 그리고 그 반대로 3D 구조에서 2D 그래프로 생성하기 위해 SE(3)-동변성 및 반사 반대칭성 SDE를 활용함으로써 MoleculeSDE는 더 낮은 상호정보량 추정과 32개의 후행 작업 중 26개에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Molecule pretraining has quickly become the go-to schema to boost the performance of AI-based drug discovery. Naturally, molecules can be represented as 2D topological graphs or 3D geometric point clouds. Although most existing pertaining methods focus on merely the single modality, recent research has shown that maximizing the mutual information (MI) between such two modalities enhances the molecule representation ability. Meanwhile, existing molecule multi-modal pretraining approaches approximate MI based on the representation space encoded from the topology and geometry, thus resulting in the loss of critical structural information of molecules. To address this issue, we propose MoleculeSDE. MoleculeSDE leverages group symmetric (e.g., SE(3)-equivariant and reflection-antisymmetric) stochastic differential equation models to generate the 3D geometries from 2D topologies, and vice versa, directly in the input space. It not only obtains tighter MI bound but also enables prosperous downstream tasks than the previous work. By comparing with 17 pretraining baselines, we empirically verify that MoleculeSDE can learn an expressive representation with state-of-the-art performance on 26 out of 32 downstream tasks.

연구 동기 및 목표

• 기존의 다중 모달 분자 전훈 방법이 표현 공간에서 상호정보량(MI)을 근사함으로써 구조적 정보 손실을 겪는 한계를 해결하기 위해.
• 표현 공간이 아닌 압축된 표현 공간이 아닌 입력 공간(2D 구조 및 3D 기하학)에서 직접 생성 모델링을 수행함으로써 MI 추정 정확도를 향상시키기 위해.
• 물리적 대칭성을 갖춘 SDE를 사용하여 2D 구조에서 고품질의 3D 구형 구조를 생성함으로써 후행 기하학적 모델링을 지원하기 위해.
• 다양한 후행 분자 성질 예측 및 구형 구조 생성 작업에서 최신 기술 수준 성능을 달성하기 위해.
• 분자 표현 학습을 위한 통합 프레임워크를 제공하기 위해, 대비 학습과 생성 학습을 결합하고 군 대칭 제약 조건을 통합하기 위해.

제안 방법

• MoleculeSDE는 3D 구형 구조에서 2D 구조로, 그리고 2D 구조에서 3D 구조로의 역방향 노이즈 제거 과정을 모델링하기 위해 확률 미분 방정식(SDE) 프레임워크를 사용한다.
• 2D에서 3D로의 생성에는 3D 기하학에서 회전 및 이동 대칭성을 유지하기 위해 SE(3)-동변성 SDE를 사용한다.
• 3D에서 2D로의 생성에는 오직 불변 특징(예: 노드 및 엣지 유형)에 집중하기 위해 SE(3)-불변 SDE를 사용한다.
• 표현 정렬과 생성 정밀도를 동시에 최적화하기 위해, EBM-NCE를 사용한 대비 학습과 생성 학습 목표를 통합한다.
• 신경망 아키텍처에 군 대칭성(SE(3)-동변성 및 반사 반대칭성)를 명시적으로 구현하여 기하학적 불변성을 유지한다.
• 프록시 기반 MI 추정을 피하기 위해 데이터 공간에서 직접 생성 모델링을 수행함으로써 더 낮은 상호정보량 경계를 확보하고 정보 손실를 감소시킨다.

실험 결과

연구 질문

• RQ1 대칭성 SDE를 사용하여 입력 공간에서 직접 생성 모델링을 수행할 경우, 프록시 기반 방법보다 2D 분자 구조와 3D 구형 구조 간의 상호정보량 추정이 더 정밀한가?
• RQ2 SDE 기반 생성 과정에서 SE(3)-동변성 및 반사 반대칭성을 강제 적용할 경우, 생성된 3D 구형 구조의 품질과 물리적 타당성이 향상되는가?
• RQ3 공동 2D-3D 전훈을 활용하여 MoleculeSDE가 다양한 후행 분자 성질 예측 작업에서 최신 기술 수준 성능을 달성할 수 있는가?
• RQ4 기존의 2D, 3D 및 다중 모달 전훈 기반 모델들과 비교했을 때, MoleculeSDE는 표준 벤치마크에서 성능과 효율성 측면에서 어떻게 다른가?
• RQ5 MoleculeSDE는 2D 구조만으로도 고품질의 3D 구형 구조를 얼마나 잘 생성할 수 있는가? 이는 후행 기하학적 모델링을 가능하게 하는가?

주요 결과

• MoleculeSDE는 32개의 후행 작업 중 26개에서 최신 기술 수준 성능을 달성하였으며, 다양한 분자 성질 예측 벤치마크에서 17개의 전훈 기반 모델보다 뚜렷하게 뛰어난 성능을 보였다.
• MD17 힘 예측 벤치마크에서 MoleculeSDE(VE)는 아스피린의 평균 절대 오차(MAE)가 0.421로, 두 번째로 좋은 성능을 보인 GeoSSL-DDM(0.427)을 능가했다.
• GEOM QM9에서의 구형 구조 생성 성능 측정에서 MoleculeSDE는 커버리지(COV) 92.37%와 평균 절대 오차(MAT) 0.2423 Å를 기록하여 이전 방법들인 DMGC 및 ConfGF를 뛰어넘었다.
• GEOM Drugs에서 MoleculeSDE는 85.42%의 커버리지와 0.9485 Å의 MAT 성능을 기록했으며, RMCF-C(87.12% COV, 0.749 Å MAT)보다 커버리지에서는 떨어지지만 MAT 성능에 거의 도달했다.
• MoleculeSDE의 계산 비용은 단일 V100 GPU에서 에포크당 약 30분으로, 더 단순한 2D 또는 3D SSL 방법보다는 높지만, 뛰어난 성능에 비해 정당화된다.
• 제거 분석 결과, 대칭성 SDE와 직접적인 입력 공간 생성을 사용할 경우, GraphMVP와 같은 프록시 기반 방법보다 더 정확한 MI 추정과 더 나은 후행 성능을 달성함을 확인했다.