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

TL;DR

SemlaFlow proposes a scalable E(3)-equivariant message passing architecture, Semla, combined with scale optimal transport and flow matching to enable efficient 3D molecular generation. It achieves state-of-the-art performance with just 20 sampling steps—delivering a 100x speedup over prior methods—while generating chemically valid molecules with low strain energy and improved conformational quality.

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.

Motivation & Objective

• To address the critical bottleneck of slow sampling in 3D molecular generation models, which limits their use in reinforcement learning and protein-pocket guided design.
• To improve scalability and efficiency of 3D molecular generators by introducing a novel E(3)-equivariant architecture, Semla, that outperforms existing message-passing models.
• To reduce transport cost in optimal transport-based training by extending equivariant optimal transport to account for molecular size, introducing scale optimal transport.
• To propose new evaluation metrics—specifically energy and strain energy—to better assess the quality of generated molecular conformations beyond standard validity and similarity metrics.
• To demonstrate that high-quality, chemically valid 3D molecular samples can be generated with only 20 steps, enabling practical deployment in drug discovery pipelines.

Proposed method

• SemlaFlow uses flow matching to train a generative model that transports noise distributions to the data distribution of molecular coordinates.
• The model employs a novel E(3)-equivariant message passing architecture, Semla, which improves parameter efficiency and inference speed compared to prior fully connected models.
• Scale optimal transport is introduced as a new OT formulation that accounts for molecular size, reducing transport cost and improving alignment between noise and data distributions.
• The model is trained using conditional flow matching, where a neural network regresses the vector field to guide the ODE-based generation process from noise to data.
• The training objective minimizes the expected vector field error using a time-dependent conditional flow, enabling end-to-end training without simulation.
• The inference process integrates the learned vector field using an ODE solver, with sampling steps configurable to balance speed and quality.

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.

Experimental results

Research questions

• RQ1 Can a scalable E(3)-equivariant message passing architecture significantly improve the efficiency and scalability of 3D molecular generation models compared to fully connected baselines?
• RQ2 Does extending equivariant optimal transport to include molecular size via scale optimal transport reduce transport cost and improve generation quality?
• RQ3 Can flow matching with scale optimal transport train a 3D molecular generator that achieves state-of-the-art performance with only 20 sampling steps?
• RQ4 How do energy and strain energy of generated molecules compare to real data and SOTA models, and can these metrics better assess conformational quality?
• RQ5 To what extent do current evaluation benchmarks fail to capture the physical plausibility of generated molecular conformations, and how can they be improved?

Key findings

• SemlaFlow achieves state-of-the-art performance on benchmark datasets with only 100 sampling steps, outperforming existing models in validity, stability, and conformational quality.
• With just 20 sampling steps, SemlaFlow matches the performance of SOTA models while achieving a 100x speedup in sampling time, corresponding to a 2-orders-of-magnitude improvement in inference efficiency.
• Molecules generated by SemlaFlow exhibit lower strain energy (127.5 ± 2.1) compared to EQGAT-diff (150.2 ± 1.8), indicating better conformational stability.
• The use of scale optimal transport is critical: ablation studies show it significantly reduces transport cost and improves energy and strain energy metrics.
• SemlaFlow generates molecules with lower minimised energy (88.9 ± 1.3) than EQGAT-diff (95.2 ± 0.22) at 100 steps, indicating better conformational optimization.
• The proposed evaluation metrics—energy and strain energy—reveal that current benchmarks underestimate conformational quality, and SemlaFlow outperforms SOTA in these more physically meaningful metrics.