FlashMol: High-Quality Molecule Generation in as Few as Four Steps
TL;DR
This work presents FlashMol, an ultra-fast molecule generative model producing high-quality molecular conformations in as few as 4 steps, and adapt distribution matching distillation (DMD) -- a reverse KL-divergence minimization objective -- to the molecular domain for effective distillation.
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: by applying respacing to the GeoLDM teacher (GeoLDM-respaced), we achieve comparable results to the 1000-NFE teacher with as few as 16 NFEs. After further distillation, FlashMol is able to sample with only 4 NFEs while maintaining high molecular quality, achieving up to 250$\times$ acceleration compared to the GeoLDM teacher.
