[论文解读] Diffusion-based Molecule Generation with Informative Prior Bridges

Lemeng Wu, Chengyue Gong+3 more|arXiv (Cornell University)|Sep 2, 2022

Cell Image Analysis Techniques被引用 19

一句话总结

本文提出了一种新颖的框架，通过在信息性先验桥中注入物理和统计先验，实现基于扩散模型的分子与3D点云生成。这些先验桥是随机过程，可保证终端输出匹配目标数据分布。采用基于李雅普诺夫函数的方法构建这些桥梁，该方法提升了生成质量、稳定性和点云均匀性，在更少采样步数下实现最先进性能，且无需架构约束即可高效整合先验信息。

ABSTRACT

AI-based molecule generation provides a promising approach to a large area of biomedical sciences and engineering, such as antibody design, hydrolase engineering, or vaccine development. Because the molecules are governed by physical laws, a key challenge is to incorporate prior information into the training procedure to generate high-quality and realistic molecules. We propose a simple and novel approach to steer the training of diffusion-based generative models with physical and statistics prior information. This is achieved by constructing physically informed diffusion bridges, stochastic processes that guarantee to yield a given observation at the fixed terminal time. We develop a Lyapunov function based method to construct and determine bridges, and propose a number of proposals of informative prior bridges for both high-quality molecule generation and uniformity-promoted 3D point cloud generation. With comprehensive experiments, we show that our method provides a powerful approach to the 3D generation task, yielding molecule structures with better quality and stability scores and more uniformly distributed point clouds of high qualities.

研究动机与目标

• 解决将强物理与统计先验整合到基于扩散的生成模型中以实现分子与3D点云生成的挑战。
• 开发一种灵活、与架构无关的先验注入方法，可在不修改模型结构的前提下提升生成质量与稳定性。
• 通过在扩散过程中嵌入促进均匀性的力，实现均匀且高质量的3D点云生成。
• 提供一种系统性的、基于李雅普诺夫函数的桥梁构建方法，确保终端输出的确定性。
• 在减少采样步数的前提下，展示在分子生成与点云均匀性方面达到最先进性能。

提出的方法

• 将扩散桥梁构建为随机过程，其在固定终端时间确定性地生成给定数据点，从而确保先验信息引导的生成。
• 采用基于李雅普诺夫函数的方法系统设计并确定桥梁的漂移项，实现稳定且可解释的先验整合。
• 引入能量函数——瑞利势能与统计差距能量——以促进3D生成任务中的点分布均匀性。
• 应用物理信息能量函数（如键长、键角与二面角势能）引导分子生成，同时保持生成灵活性。
• 训练神经扩散模型以模仿先验桥梁，有效引导反向过程生成真实且高质量的输出。
• 采用联合学习方案，使桥梁参数（如 α）与扩散模型端到端联合优化。

实验结果

研究问题

• RQ1 信息性先验桥梁是否能在不修改架构的前提下显著提升基于扩散模型的分子生成质量与稳定性？
• RQ2 如何系统性地将物理与统计先验嵌入扩散过程，以引导生成向真实且均匀的输出？
• RQ3 先验信息引导的桥梁是否能减少采样步数，同时在3D点云生成中保持或提升生成质量？
• RQ4 不同先验力（如瑞利势能与统计差距能量）对点云均匀性与形状真实性有何影响？
• RQ5 与标准扩散模型相比，通过桥梁实现的先验整合对训练效率与收敛速度有何影响？

主要发现

• 所提方法在分子生成任务中达到最先进性能，优于现有物理信息基线，在质量与稳定性指标上均表现更优。
• 仅使用10次扩散步数，该方法的性能即匹配或超越使用100次步数的标准扩散模型，表明收敛速度更快。
• 在桥梁过程中引入瑞利能量与统计差距能量，显著提升了点云的均匀分布程度，其中后者展现出更强鲁棒性。
• 在100步设置下，采用先验桥梁的方法在所有评估类别中，于MMD与COV指标上均优于基线扩散模型。
• 通过直接将先验知识注入训练过程，该方法实现了高质量且稳定的分子生成，无需架构约束。
• 可视化结果表明，先验信息引导的桥梁能生成更真实且分布更均匀的点云，但瑞利能量因强排斥作用偶发引入异常点。