[论文解读] EquiFlow: Equivariant Conditional Flow Matching with Optimal Transport for 3D Molecular Conformation Prediction

Qingwen Tian, Yuxin Xu+5 more|arXiv (Cornell University)|Dec 15, 2024

Mass Spectrometry Techniques and Applications被引用 4

一句话总结

EquiFlow 提出了一种新颖的等变条件流匹配模型，结合最优传输方法用于3D分子构象预测，利用改进的 Equiformer 编码原子和键特征，同时保持 E(3) 等变性。该方法在单构象预测中实现了 0.17Å 的 RMSD，且在多构象生成中表现出更优的多样性与准确性，在 QM9 和 GEOM-QM9 数据集上优于先前的扩散模型与流模型，达到新的 SOTA 水平。

ABSTRACT

Molecular 3D conformations play a key role in determining how molecules interact with other molecules or protein surfaces. Recent deep learning advancements have improved conformation prediction, but slow training speeds and difficulties in utilizing high-degree features limit performance. We propose EquiFlow, an equivariant conditional flow matching model with optimal transport. EquiFlow uniquely applies conditional flow matching in molecular 3D conformation prediction, leveraging simulation-free training to address slow training speeds. It uses a modified Equiformer model to encode Cartesian molecular conformations along with their atomic and bond properties into higher-degree embeddings. Additionally, EquiFlow employs an ODE solver, providing faster inference speeds compared to diffusion models with SDEs. Experiments on the QM9 dataset show that EquiFlow predicts small molecule conformations more accurately than current state-of-the-art models.

研究动机与目标

• 解决现有用于 3D 分子构象预测的深度学习模型在训练速度慢和高阶特征利用不足方面的局限性。
• 通过将条件流匹配与最优传输应用于分子几何生成，实现无需模拟的高效训练与快速推理。
• 通过改进的 Equiformer 架构，直接建模笛卡尔坐标，同时保持平移与旋转等变性，以提升模型性能。
• 通过基于最优传输的训练目标，提升多构象预测中的准确率与多样性。
• 通过端到端、等变的流匹配方法，建立 3D 分子构象生成的新 SOTA 基线。

提出的方法

• EquiFlow 采用条件流匹配（CFM）并结合最优传输（OT）目标，以可微、无需模拟的方式训练向量场，将噪声分布映射到分子构象。
• 使用改进的 Equiformer 模型将原子坐标、类型和键特征编码为高阶几何表示，同时保持 E(3) 等变性。
• 通过尊重分子系统平移与旋转对称性的等变神经网络参数化流向量场。
• 推理阶段使用 ODE 求解器，从噪声高效生成构象，相比基于 SDE 的扩散模型显著提升推理速度。
• 训练目标最小化预测与目标分子构象分布之间的 OT 距离，确保训练过程稳定且准确。
• 框架直接建模笛卡尔坐标，不依赖中间不变特征，提升了灵活性与表征丰富性。

实验结果

研究问题

• RQ1 能否将基于最优传输的条件流匹配有效应用于 3D 分子构象预测，实现比扩散模型更快的训练与推理速度？
• RQ2 通过改进的 Equiformer 引入高阶几何特征，如何提升分子构象生成的准确率与多样性？
• RQ3 在向量场参数化中保持 E(3) 等变性，在多分子几何任务中在多大程度上增强模型泛化能力与性能？
• RQ4 无模拟、端到端的流匹配框架能否在标准分子基准上超越现有的基于扩散与几何的生成模型？
• RQ5 所提出方法在多构象预测中如何在严格 RMSD 阈值下平衡准确率与多样性？

主要发现

• 在 QM9 测试集上，EquiFlow 实现了 0.17Å 的单构象预测 RMSD，优于所有先前的 SOTA 方法。
• 在多构象 GEOM-QM9 基准上，EquiFlow 达到最高的 COV-R（覆盖召回率）与最低的 MAT-R（匹配召回率）得分，表明其在多样性和准确性方面表现更优。
• 该模型在所有四个基于 RMSD 的指标（COV-R、MAT-R、COV-P 和 MAT-P）上均达到 SOTA 水平，其中在 δ = 0.5Å 时 COV-R 达到 92.3%。
• 通过使用 ODE 求解器而非随机过程，EquiFlow 的推理速度显著快于基于 SDE 的扩散模型，大幅缩短推理时间。
• 通过改进的 Equiformer 集成高阶特征，生成更具表现力的表征，提升了泛化能力与预测性能。
• 基于最优传输的训练目标确保了稳定高效的优化，避免了模式崩溃，提升了样本质量。