[论文解读] Score-Based Generative Models for Molecule Generation

Dwaraknath Gnaneshwar, Bharath Ramsundar+3 more|arXiv (Cornell University)|Mar 7, 2022

Machine Learning in Materials Science被引用 7

一句话总结

该论文提出了一种基于分数的生成模型，用于分子生成，其基于Transformer的分数函数在150万条ZINC分子的SELFIES表示上进行训练。该模型生成了多样、有效且新颖的化合物——实现了100%的有效性和新颖性——但生成的分子与训练分布显著偏离，表明在建模复杂化学空间方面仍存在改进架构的潜力。

ABSTRACT

Recent advances in generative models have made exploring design spaces easier for de novo molecule generation. However, popular generative models like GANs and normalizing flows face challenges such as training instabilities due to adversarial training and architectural constraints, respectively. Score-based generative models sidestep these challenges by modelling the gradient of the log probability density using a score function approximation, as opposed to modelling the density function directly, and sampling from it using annealed Langevin Dynamics. We believe that score-based generative models could open up new opportunities in molecule generation due to their architectural flexibility, such as replacing the score function with an SE(3) equivariant model. In this work, we lay the foundations by testing the efficacy of score-based models for molecule generation. We train a Transformer-based score function on Self-Referencing Embedded Strings (SELFIES) representations of 1.5 million samples from the ZINC dataset and use the Moses benchmarking framework to evaluate the generated samples on a suite of metrics.

研究动机与目标

• 探究基于分数的生成模型在从头分子生成中的可行性和有效性。
• 通过利用基于分数模型的架构灵活性，克服GAN和归一化流模型的局限性（如训练不稳定性和架构约束）。
• 建立使用SELFIES表示在分子数据上训练基于分数模型的基础最佳实践。
• 通过MOSES基准套件在有效性、唯一性、多样性和分布保真度等指标上评估模型性能。
• 探索将等变架构（如SE(3)-等变模型）集成到分数函数中的潜力，以提升分子生成效果。

提出的方法

• 训练一个基于Transformer的分数函数，以估计分子数据在SELFIES字符串潜在空间中对数概率密度的梯度（即分数）。
• 使用多噪声级别去噪分数匹配（σ₁ = 10，σₗ = 0.01，共350个级别）以稳定训练并改善分数估计。
• 采用退火Langevin动力学进行采样：从随机噪声开始，通过逐步降低噪声水平，利用分数函数迭代去噪。
• 将分子表示为展平为向量的一位编码SELFIES字符串，作为分数网络的输入。
• 使用如下目标函数优化分数函数：∑ᵢ λ(i)𝔼[‖sθ(x,i) − ∇ₓlog pσᵢ(x)‖²]，其中λ(i)为不同噪声水平的权重。
• 在单张V100 GPU上训练约12小时，并在约4小时内采样30,000个分子，使用T=10个时间步。

实验结果

研究问题

• RQ1 当在SELFIES编码的分子数据上进行训练时，基于分数的生成模型能否有效生成化学上有效且多样的分子？
• RQ2 基于Transformer的分数函数在标准分子生成基准测试中，与现有模型（如VAEs、GANs）相比表现如何？
• RQ3 生成的分子在多大程度上匹配训练数据（ZINC 1.6M）的统计分布？导致分布偏差的可能因素有哪些？
• RQ4 基于分数模型的架构灵活性是否能支持未来集成等变层（如SE(3)-等变GNN）以提升3D分子生成效果？
• RQ5 在基于分数的分子生成中，哪些超参数和训练策略（如噪声调度、模型深度）最为有效？

主要发现

• 该模型在生成的30,000个分子中实现了100%的有效性和100%的新颖性，表明其能够生成结构有效且完全新颖的化合物。
• 在1,000个分子中唯一性得分为88%，在10,000个分子中为82%，虽低于最先进VAE模型，但仍具竞争力，且通过超参数调优存在进一步提升空间。
• IntDiv₁和IntDiv₂得分分别为0.90和0.88，表明生成分子集具有强化学多样性且无模式崩溃现象。
• FCD/Test和FCD/TestSF得分分别为39.84和40.92，表明与训练数据存在显著分布偏差，提示生成分子远离真实分子分布。
• 该模型性能凸显了一项目前存在的权衡：在牺牲分布保真度的前提下实现了高多样性和高新颖性，提示需提升分数函数容量或引入更具归纳偏置的架构。
• 结果验证了基于分数模型用于分子生成的可行性，并为未来集成等变架构以更好地捕捉分子对称性并改进分布建模打开了新路径。