FlashMol: High-Quality Molecule Generation in as Few as Four Steps

2026-05-07

AI Summary

Traditional diffusion models for 3D molecular conformation generation typically require hundreds to thousands of sampling steps, rendering them impractical for large-scale virtual screening due to computational inefficiency. Existing acceleration strategies often compromise sample stability or diversity. To address this, this work proposes FlashMol, the first method to adapt distribution matching distillation (DMD) to molecular conformation generation. By integrating timestep reordering, reverse KL divergence minimization, and Jensen–Shannon divergence regularization, FlashMol achieves high-quality and diverse 3D conformations in as few as four sampling steps. Evaluated on the QM9 and GEOM-DRUG datasets, FlashMol matches or surpasses the performance of a 1000-step teacher model while accelerating sampling by up to 250×, effectively balancing efficiency, stability, and generation fidelity.

Abstract

Generating chemically valid 3D molecular conformations is critical for computational drug discovery. Classical diffusion-based models like GeoLDM perform well but require hundreds of steps, making large-scale in silico screening impractical. Recent efforts on few-step molecular generation have accelerated this process to 12-50 steps, but they often largely sacrifice sample stability. In this work, we present FlashMol, an ultra-fast molecule generative model producing high-quality molecular conformations in as few as 4 steps. To achieve this, we adapt distribution matching distillation (DMD) - a reverse KL-divergence minimization objective - to the molecular domain for effective distillation. Considering the local minimization behavior of DMD, we respace the molecule generation timesteps, providing the generator with much better initialization and enables effective distillation. Additionally, to mitigate the mode-seeking behavior of DMD and improve diversity, we further regularize it with a Jensen-Shannon divergence term, which incorporates the mean-seeking behavior of the forward KL divergence. Extensive experiments on QM9 and GEOM-DRUG datasets demonstrate that FlashMol matches and even surpasses the original 1000-step teacher, achieving up to 250$\times$ acceleration in sampling speed while maintaining high molecular quality.

Problem

• molecule generation
• 3D molecular conformations
• few-step sampling
• sample stability
• computational drug discovery

Innovation

• FlashMol
• distribution matching distillation
• few-step molecular generation
• 3D molecular conformation
• Jensen-Shannon divergence regularization