[논문 리뷰] Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation

Tuan Le, Julian Cremer외 3명|arXiv (Cornell University)|2023. 09. 29.

Machine Learning in Materials Science인용 수 6

한 줄 요약

이 논문은 연속적인 원자 좌표, 범주형 원자/결합 특징, 시간에 따라 변하는 손실 가중치를 활용하여 QM9와 GEOM-Drugs에서 최고 성능을 기록하고, 더 빠른 훈련과 추론을 달성하는 3차원 신규 분자 생성을 위한 새로운 E(3)-등변성 확산 모델인 EQGAT-diff를 소개한다. 또한, 암시적 수소를 포함한 PubChem3D에서의 사전 훈련을 통해 이식 가능성을 입증하고, 하이브리드화 상태와 같은 화학적으로 정보적인 특징을 통합하여 유효성을 향상시킨다.

ABSTRACT

Deep generative diffusion models are a promising avenue for 3D de novo molecular design in materials science and drug discovery. However, their utility is still limited by suboptimal performance on large molecular structures and limited training data. To address this gap, we explore the design space of E(3)-equivariant diffusion models, focusing on previously unexplored areas. Our extensive comparative analysis evaluates the interplay between continuous and discrete state spaces. From this investigation, we present the EQGAT-diff model, which consistently outperforms established models for the QM9 and GEOM-Drugs datasets. Significantly, EQGAT-diff takes continuous atom positions, while chemical elements and bond types are categorical and uses time-dependent loss weighting, substantially increasing training convergence, the quality of generated samples, and inference time. We also showcase that including chemically motivated additional features like hybridization states in the diffusion process enhances the validity of generated molecules. To further strengthen the applicability of diffusion models to limited training data, we investigate the transferability of EQGAT-diff trained on the large PubChem3D dataset with implicit hydrogen atoms to target different data distributions. Fine-tuning EQGAT-diff for just a few iterations shows an efficient distribution shift, further improving performance throughout data sets. Finally, we test our model on the Crossdocked data set for structure-based de novo ligand generation, underlining the importance of our findings showing state-of-the-art performance on Vina docking scores.

연구 동기 및 목표

• 기존의 확산 모델이 3차원 분자 생성에서 약한 성능을 보이는 문제점, 특히 큰 분자에서의 성능 부족과 데이터 부족 문제를 해결하기 위해.
• 연속적 상태 공간과 이산적 상태 공간, 손실 가중치 전략을 중심으로 E(3)-등변성 확산 모델의 설계 공간을 체계적으로 탐색하기 위해.
• 시간에 따라 변하는 손실 가중치와 향상된 특징 공학을 통해 훈련 수렴, 추론 속도, 샘플 품질을 향상시키기 위해.
• 암시적 수소를 포함한 대규모 PubChem3D에서의 사전 훈련을 통해 소규모 데이터셋에서의 효율적인 파라미터 조정을 가능하게 하기 위해.
• 하이브리드화 상태와 같은 화학적으로 유의미한 특징을 확산 과정에 통합하여 분자의 유효성을 향상시키기 위해.

제안 방법

• 연속적인 3차원 원자 좌표와 범주형 원자/결합 유형을 포함한 디노이징 확산 프레임워크 내에서 E(3)-등변성 그래프 어텐션 네트워크인 EQGAT-diff를 제안한다.
• 신호 대 잡음비(SNR) 기반의 시간에 따라 변하는 손실 가중치를 활용하여 훈련 수렴 속도를 가속화하고 샘플 품질을 향상시킨다.
• 하이브리드 데이터 모odal 접근법을 도입: 연속 좌표, 이산 원자 유형, 이산 결합 유형을 별도의 디노이징 헤드로 처리한다.
• DDIM 추론을 수정된 이산 업데이트 방식으로 활용하여 범주형 특징의 정밀도를 유지하면서도 추론 속도를 가속화한다.
• 암시적 수소를 포함한 PubChem3D 데이터셋에서의 무조건적 사전 훈련을 수행하여, 하류 작업에서의 효율적인 파라미터 조정을 가능하게 한다.
• 하이브리드화 상태와 아로마티시티와 같은 화학적으로 의미 있는 특징을 입력 노드 및 엣지 특징에 통합하여 화학적 유효성을 향상시킨다.

실험 결과

연구 질문

• RQ1 3차원 분자 확산 모델에서 시간에 따라 변하는 손실 가중치(예: SNR 기반)가 균일하거나 이산적 손실 가중치보다 훈련 수렴과 샘플 품질 측면에서 어떻게 비교되는가?
• RQ2 연속적인 원자 좌표를 사용하는 것과 이산 표현을 사용하는 것의 성능 및 추론 속도에 미치는 영향은 무엇인가?
• RQ3 대규모 3차원 분자 데이터셋(예: PubChem3D)에서 사전 훈련한 모델이 소규모이자 더 복잡한 데이터셋에서 최소한의 추가 훈련으로 효과적으로 파라미터 조정 가능한가?
• RQ4 하이브리드화와 아로마티시티와 같은 화학적으로 정보적인 특징이 확산 기반 생성에서 생성된 분자의 유효성과 다양성에 어떤 영향을 미치는가?
• RQ5 매개변수화 방법의 선택(예: $x_0$ 대비 $x_t$)이 3차원 분자 확산 모델의 성능과 안정성에 얼마나 큰 영향을 미치는가?

주요 결과

• EQGAT-diff는 QM9와 GEOM-Drugs에서 최고 성능을 기록했으며, QM9에서 유효성 96.8%와 성공률 96.6%를 달성하여 이전 모델을 능가한다.
• 절단된 SNR(t) 기반의 시간에 따라 변하는 손실 가중치는 훈련 수렴과 샘플 품질을 크게 향상시키며, 오직 100개의 타임스텝만 사용할 경우 추론 시간을 최대 5배까지 단축시킨다.
• 100개 타임스텝과 SNR 기반 가중치로 훈련한 모델이 500개 타임스텝과 균일 가중치로 훈련한 모델보다 우수한 성능을 보이며, 적응형 손실 스케줄링의 효과를 입증한다.
• PubChem3D에서 사전 훈련한 후 Crossdocked와 같은 타겟 데이터셋에서의 파라미터 조정은 몇 차례의 파라미터 조정 반복만으로도 뛰어난 성능을 달성하며, 효율적인 분포 이동을 가능하게 한다.
• 하이브리드화와 아로마티시티 특징을 통합하면 분자의 유효성과 다양성이 향상되며, GEOM-Drugs에서 42.2%의 다양성 점수를 기록하여 기준 모델보다 뚜렷이 높다.
• DDIM 추론은 타임스텝을 500 미만으로 줄일 경우 샘플 품질 향상에 기여하지 않았지만, 100개 타임스텝과 SNR 기반 가중치로 훈련한 모델은 높은 성능을 유지하면서도 5배 빠른 추론 속도를 확보했다.