[论文解读] 3D molecule generation by denoising voxel grids

Pedro O. Pinheiro, Joshua A. Rackers+7 more|arXiv (Cornell University)|Jun 13, 2023

Machine Learning in Materials Science被引用 6

一句话总结

VoxMol 提出了一种新颖的基于分数的生成模型，用于 3D 分子生成，通过去噪体素化的原子密度网格实现。该方法利用 Langevin 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 分子生成方法，以克服现有点云扩散模型的局限性。
• 将分子表示为规则体素网格上的连续原子密度网格，以实现更有效的生成建模。
• 通过避免预测原子数量以及分别处理混合的连续-离散特征，简化训练和推理过程。
• 在训练复杂度和推理时间方面，相较于最先进点云扩散模型，实现更具竞争力的生成质量。
• 通过行走-跳跃 Langevin MCMC 采样和基于峰值的坐标精炼，实现高效采样，以支持下游分子分析。

提出的方法

• 分子被表示为 3D 体素网格，其中每个体素通过以原子位置为中心的高斯函数编码原子密度。
• 训练一个去噪神经网络，将带有固定大方差高斯噪声的噪声体素网格映射到干净的真实分子网格。
• 该去噪网络近似平滑数据分布的分数函数，从而支持基于分数的生成采样。
• 分子通过两步过程生成：(i) 从平滑噪声分布出发进行 Langevin 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 在固定噪声水平下，采用行走-跳跃 Langevin 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