[논문 리뷰] 3D molecule generation by denoising voxel grids

Pedro O. Pinheiro, Joshua A. Rackers외 7명|arXiv (Cornell University)|2023. 06. 13.

Machine Learning in Materials Science인용 수 6

한 줄 요약

VoxMol은 볼륨화된 원자 밀도 격자에서 노이즈 제거를 통해 3D 분자의 새로운 스코어 기반 생성 모델을 제안한다. 이는 랭그비안 MCMC 샘플링과 단일 노이즈 제거 단계를 활용하여 점군 확산 모델보다 더 빠르고 간단하게 분자를 생성하며, 무조건적 3D 분자 생성 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

We propose a new score-based approach to generate 3D molecules represented as atomic densities on regular grids. First, we train a denoising neural network that learns to map from a smooth distribution of noisy molecules to the distribution of real molecules. Then, we follow the neural empirical Bayes framework (Saremi and Hyvarinen, 19) and generate molecules in two steps: (i) sample noisy density grids from a smooth distribution via underdamped Langevin Markov chain Monte Carlo, and (ii) recover the "clean" molecule by denoising the noisy grid with a single step. Our method, VoxMol, generates molecules in a fundamentally different way than the current state of the art (ie, diffusion models applied to atom point clouds). It differs in terms of the data representation, the noise model, the network architecture and the generative modeling algorithm. Our experiments show that VoxMol captures the distribution of drug-like molecules better than state of the art, while being faster to generate samples.

연구 동기 및 목표

• 기존의 점군 확산 모델의 한계를 극복하는 새로운 방법을 개발하여 무조건적 3D 분자 생성을 수행한다.
• 정규 볼륨 격자 위에서 연속적인 원자 밀도 격자로 분자를 표현하여 더 효과적인 생성 모델링을 가능하게 한다.
• 원자 수 예측이 필요 없고 혼합된 연속-이산 특성 유형을 별도로 처리할 필요 없이 학습 및 추론을 단순화한다.
• 최신 기술 수준의 점군 기반 확산 모델과 비교해 학습 복잡도를 낮추고 추론 시간을 단축하면서도 경쟁 가능한 생성 품질을 달성한다.
• 보조 분자 분석을 위한 효율적인 샘플링을 위해 워크점프 랭그비안 MCMC 및 피크 기반 좌표 정밀 조정을 가능하게 한다.

제안 방법

• 각 볼륨 요소는 원자 위치에 중심을 둔 가우시안 유사 함수를 사용해 원자 밀도를 표현하는 3D 볼륨 격자로 분자를 표현한다.
• 노이즈가 첨가된 볼륨 격자(고정된 큰 분산을 가진 가우시안 노이즈)를 입력으로 받아 실제 분자 격자로 복원하는 노이즈 제거 신경망을 학습시킨다.
• 노이즈 제거 네트워크는 부드러운 데이터 분포의 스코어 함수를 근사하여 스코어 기반 생성 샘플링을 가능하게 한다.
• 분자는 두 단계 과정을 통해 생성된다: (i) 부드러운 노이즈 분포에서 랭그비안 MCMC 샘플링을 수행하고, (ii) 단일 단계 노이즈 제거를 통해 깔끔한 분자를 복원한다.
• 원자 좌표는 생성 후 볼륨 피크 탐지 및 최적화 기반 정밀 조정을 통해 회복되며, 기하학적 정확도를 향상시킨다.
• 볼륨화 과정에서 원자 반지름을 고정된 0.5 Å로 설정했으며, 이는 다양한 데이터셋에서 가장 우수한 성능을 보였다.

Figure 1: Voxelized molecules generated by our model and their corresponding molecular graphs. Left, samples from a model trained on QM9 dataset ( $32^{3}$ voxels). Right, samples from a model trained on GEOM-drugs ( $64^{3}$ voxels). In both cases, each voxel is a cubic grid with side length of $.2

실험 결과

연구 질문

• RQ1 볼륨화된 원자 밀도 격자 기반 스코어 기반 생성 모델이 무조건적 3D 분자 생성에서 점군 확산 모델을 능가할 수 있는가?
• RQ2 연속적이고 격자 기반 표현을 사용할 경우, 이산 원자 점군 대비 학습 및 추론이 간소화되는가?
• RQ3 볼륨화 과정에서 원자 반지름의 선택이 모델 성능과 샘플 품질에 어떤 영향을 미치는가?
• RQ4 고정된 노이즈 수준을 가진 워크점프 랭그비안 MCMC 샘플링이 분자의 다양성과 구조적 정확도를 어느 정도 유지하는가?
• RQ5 볼륨 피크에서 유도된 좌표 정밀 조정이 생성된 분자의 유효성과 기하학적 정확도를 크게 향상시킬 수 있는가?

주요 결과

• VoxMol은 QM9 및 GEOM-Drugs 데이터셋에서 무조건적 3D 분자 생성에서 최신 기술 수준의 성능을 달성하며, 경쟁 가능한 생성 품질을 확보한다.
• 노이즈 수준 σ = 1.0일 때 좌표 정밀 조정을 사용하면 QM9에서 68.6%의 유효성과 96.3%의 유일성을 달성했으며, 정밀 조정 없이 사용할 경우 55.0%의 유효성과 100%의 유일성으로 성능이 떨어진다.
• GEOM-Drugs 데이터셋에서는 σ = 1.0일 때 정밀 조정을 통해 45.1%의 유효성과 99.7%의 유일성을 확보했으며, 정밀 조정 없이 사용할 경우 32.5%의 유효성으로 성능이 저하된다.
• 모델는 노이즈 수준에 관계없이 실제 분자와 유사한 형태 다양성을 보이며, NPMIR 분포에서 막대 모양, 디스크 모양, 구형 모양의 구조를 모두 포함한다.
• 볼륨화 과정에서 원자 반지름을 0.5 Å로 고정할 경우 가장 우수한 성능을 보였으며, 0.25 Å일 경우 학습 수렴이 어려워지고, 더 큰 반지름은 샘플 품질을 떨어뜨린다.
• 점군 기반 확산 모델과 비교해 학습이 더 단순하고 분자 생성 속도가 더 빠르며, 원자 수 예측이나 혼합 특성 유형 처리가 필요 없어진다.

Figure 2: (a) A representation of our denoising training procedure. Each training sample is corrupted with isotropic Gaussian noise with a fixed noise level $\sigma$ . The model is trained to recover clean voxels from the noisy version. (b) Graphical model representation of the walk-jump sampling sc