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

TL;DR

FlashMol is an ultra-fast molecule generation model that produces high-quality 3D conformations in as few as four steps, significantly accelerating computational drug discovery processes.

Contribution

It introduces a novel adaptation of distribution matching distillation for molecular generation, enabling high-quality results in only four steps.

Findings

Achieves up to 250× faster sampling than traditional models.

Maintains high molecular quality comparable to 1000-step models.

Demonstrates effectiveness on QM9 and GEOM-DRUG datasets.

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.