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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Wei, Xinyuan, Li, Zian, Yan, Shaoheng, Zhou, Cai, Zhang, Muhan
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911661066027008
author Wei, Xinyuan
Li, Zian
Yan, Shaoheng
Zhou, Cai
Zhang, Muhan
author_facet Wei, Xinyuan
Li, Zian
Yan, Shaoheng
Zhou, Cai
Zhang, Muhan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_07020
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle FlashMol: High-Quality Molecule Generation in as Few as Four Steps
Wei, Xinyuan
Li, Zian
Yan, Shaoheng
Zhou, Cai
Zhang, Muhan
Machine Learning
Artificial Intelligence
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.
title FlashMol: High-Quality Molecule Generation in as Few as Four Steps
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2605.07020