SemlaFlow -- Efficient 3D Molecular Generation with Latent Attention and Equivariant Flow Matching

TL;DR

The paper tackles the bottleneck of slow sampling and limited chemical validity in unconditional 3D molecular generation. It introduces Semla, an $E(3)$-equivariant latent-attention architecture, and SemlaFlow, a flow-matching model that jointly generates topology and 3D conformations. SemlaFlow achieves state-of-the-art results with as few as 20 sampling steps, offering substantial speedups and proposing energy-based benchmarks to better assess conformer quality. The work highlights evaluation gaps in 3D generation and demonstrates that combining latent attention with equivariant flow matching yields practical, scalable molecules suitable for downstream drug design tasks.

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.

Figures (5)

• Figure 1: An overview of the inference procedure for SemlaFlow. Noise is firstly sampled from the prior distribution over atom types, bond types and coordinates. We then integrate the vector field by applying our current prediction to the SemlaFlow model and taking a step in the direction of the model's prediction. To improve readability noisy samples are displayed with bonds inferred based on the 3D coordinates.
• Figure 2: Architectural overview of one Semla layer.
• Figure 3: An overview of a full Semla model. A Semla model is created from a stack of Semla layers where the first layer encodes bond information into the node features and the final layer generates bond information from the pairwise features.
• Figure 4: Examples of molecules which would be considered valid by RDKit.
• Figure 5: Random samples from a SemlaFlow model trained on GEOM Drugs. These samples were generated using 100 ODE integration steps.