[Paper Review] 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 Applications4 citations

TL;DR

EquiFlow proposes a novel equivariant conditional flow matching model with optimal transport for 3D molecular conformation prediction, leveraging a modified Equiformer to encode atomic and bond features while preserving E(3) equivariance. It achieves state-of-the-art performance with 0.17Å RMSD on single-conformation prediction and superior diversity and accuracy on multi-conformation generation, outperforming prior diffusion and flow-based models on QM9 and GEOM-QM9 datasets.

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.

Motivation & Objective

• To address the limitations of slow training and poor utilization of high-degree features in existing deep learning models for 3D molecular conformation prediction.
• To enable simulation-free, efficient training and fast inference by applying conditional flow matching with optimal transport to molecular geometry generation.
• To improve model performance by directly modeling Cartesian coordinates while maintaining translational and rotational equivariance through a modified Equiformer architecture.
• To enhance both accuracy and diversity in multi-conformation prediction by leveraging optimal transport-based training objectives.
• To establish a new SOTA baseline for 3D molecular conformation generation using end-to-end, equivariant flow matching.

Proposed method

• EquiFlow employs conditional flow matching (CFM) with an optimal transport (OT) objective to train a vector field that maps noise distributions to molecular conformations in a differentiable, simulation-free manner.
• It uses a modified Equiformer model to encode atomic coordinates, types, and bond features into high-degree geometric representations while preserving E(3) equivariance.
• The model parameterizes the flow vector field using an equivariant neural network that respects translational and rotational symmetries of molecular systems.
• An ODE solver is used during inference to efficiently generate conformations from noise, enabling faster inference than SDE-based diffusion models.
• The training objective minimizes the OT distance between the predicted and target molecular conformation distributions, ensuring stable and accurate training.
• The framework directly models Cartesian coordinates without relying on intermediate invariant features, improving flexibility and representation richness.

Experimental results

Research questions

• RQ1 Can conditional flow matching with optimal transport be effectively applied to 3D molecular conformation prediction to achieve faster training and inference than diffusion models?
• RQ2 How does incorporating high-degree geometric features through a modified Equiformer improve the accuracy and diversity of molecular conformation generation?
• RQ3 To what extent does maintaining E(3) equivariance in the vector field parameterization enhance model generalization and performance on molecular geometry tasks?
• RQ4 Can a simulation-free, end-to-end flow matching framework outperform existing diffusion-based and geometry-based generative models on standard molecular benchmarks?
• RQ5 How does the proposed method balance accuracy and diversity in multi-conformation prediction, particularly under strict RMSD thresholds?

Key findings

• EquiFlow achieves a single-conformation prediction RMSD of 0.17Å on the QM9 test set, outperforming all prior state-of-the-art methods.
• On the multi-conformation GEOM-QM9 benchmark, EquiFlow achieves the highest COV-R (coverage recall) and lowest MAT-R (matching recall) scores, indicating superior diversity and accuracy.
• The model demonstrates state-of-the-art performance across all four RMSD-based metrics: COV-R, MAT-R, COV-P, and MAT-P, with COV-R reaching 92.3% at δ = 0.5Å.
• EquiFlow enables faster inference than SDE-based diffusion models by using an ODE solver instead of stochastic processes, significantly reducing inference time.
• The integration of high-degree features via the modified Equiformer leads to more expressive representations, improving generalization and predictive power.
• The optimal transport-based training objective ensures stable and efficient optimization, avoiding mode collapse and improving sample quality.