[论文解读] SemlaFlow -- Efficient 3D Molecular Generation with Latent Attention and Equivariant Flow Matching

Ross W. Irwin, Alessandro Tibo+2 more|arXiv (Cornell University)|Jun 11, 2024

Innovative Microfluidic and Catalytic Techniques Innovation被引用 4

一句话总结

SemlaFlow 提出了一种可扩展的 E(3)-等变消息传递架构 Semla，结合尺度最优传输与流匹配，实现了高效的 3D 分子生成。其仅需 20 步采样即可达到最先进性能——相比先前方法提速 100 倍——同时生成的分子具有低应变能和更高的构象质量，且化学上有效。

ABSTRACT

Methods for jointly generating molecular graphs along with their 3D conformations have gained prominence recently due to their potential impact on structure-based drug design. Current approaches, however, often suffer from very slow sampling times or generate molecules with poor chemical validity. Addressing these limitations, we propose Semla, a scalable E(3)-equivariant message passing architecture. We further introduce an unconditional 3D molecular generation model, SemlaFlow, which is trained using equivariant flow matching to generate a joint distribution over atom types, coordinates, bond types and formal charges. Our model produces state-of-the-art results on benchmark datasets with as few as 20 sampling steps, corresponding to a two order-of-magnitude speedup compared to state-of-the-art. Furthermore, we highlight limitations of current evaluation methods for 3D generation and propose new benchmark metrics for unconditional molecular generators. Finally, using these new metrics, we compare our model's ability to generate high quality samples against current approaches and further demonstrate SemlaFlow's strong performance.

研究动机与目标

• 为解决 3D 分子生成模型中采样速度慢这一关键瓶颈，该问题限制了其在强化学习和蛋白口袋引导设计中的应用。
• 通过引入一种新型 E(3)-等变架构 Semla，提升 3D 分子生成器的可扩展性与效率，使其优于现有消息传递模型。
• 通过将等变最优传输扩展至考虑分子尺寸，提出尺度最优传输，以降低最优传输训练中的传输成本。
• 提出新的评估指标——特别是能量与应变能——以超越标准的有效性与相似性指标，更全面评估生成分子构象的质量。
• 证明仅通过 20 步即可生成高质量、化学上有效的 3D 分子，使其适用于药物发现流程中的实际部署。

提出的方法

• SemlaFlow 使用流匹配训练生成模型，将噪声分布传输至分子坐标的数据分布。
• 该模型采用一种新型 E(3)-等变消息传递架构 Semla，相比先前的全连接模型，显著提升了参数效率与推理速度。
• 引入尺度最优传输作为新的最优传输公式，考虑分子尺寸，降低传输成本并改善噪声与数据分布之间的对齐。
• 模型通过条件流匹配进行训练，其中神经网络回归向量场，引导基于 ODE 的从噪声到数据的生成过程。
• 训练目标通过时间依赖的条件流最小化期望向量场误差，实现无需模拟的端到端训练。
• 推理过程使用 ODE 求解器集成学习到的向量场，采样步数可配置以平衡速度与质量。

Figure 1 : An overview of our approach. Firstly, prior coordinate noise is sampled according to the respective scale of the molecules (atom and bond type noise is not shown for brevity). Our trained model, SemlaFlow , transports this noise to the data distribution to generate a molecular sample.

实验结果

研究问题

• RQ1 与全连接基线相比，可扩展的 E(3)-等变消息传递架构是否能显著提升 3D 分子生成模型的效率与可扩展性？
• RQ2 通过尺度最优传输将等变最优传输扩展至考虑分子尺寸，是否能降低传输成本并提升生成质量？
• RQ3 结合尺度最优传输的流匹配能否训练出在仅 20 步采样下即达到最先进性能的 3D 分子生成器？
• RQ4 生成分子的能量与应变能与真实数据及 SOTA 模型相比如何？这些指标是否能更有效地评估构象质量？
• RQ5 当前评估基准在多大程度上未能捕捉生成分子构象的物理合理性？如何改进？

主要发现

• SemlaFlow 在基准数据集上仅用 100 步采样即达到最先进性能，在有效性、稳定性与构象质量方面优于现有模型。
• 仅用 20 步采样，SemlaFlow 即可匹配 SOTA 模型性能，同时实现采样时间 100 倍加速，对应推理效率提升两个数量级。
• SemlaFlow 生成的分子应变能更低（127.5 ± 2.1），低于 EQGAT-diff（150.2 ± 1.8），表明构象更稳定。
• 尺度最优传输至关重要：消融研究显示其显著降低传输成本，并改善能量与应变能指标。
• 在 100 步时，SemlaFlow 生成分子的最小化能量（88.9 ± 1.3）低于 EQGAT-diff（95.2 ± 0.22），表明构象优化更优。
• 所提出的评估指标——能量与应变能——揭示当前基准低估了构象质量，SemlaFlow 在这些更具物理意义的指标上优于 SOTA。