Research map / Probabilistic and scientific modeling / Molecules and proteins

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

NeurIPS 2026 · Probabilistic and scientific modeling › Molecules and proteins · approach: Structure-based molecule design with diffusion · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Distills a 1000-step GeoLDM molecule generator into a 4-8 step student using distribution matching distillation, with EDM-style timestep respacing and a Jensen-Shannon regularizer for diversity.

Problem

Diffusion-based 3D molecule generators such as GeoLDM need ~1000 denoising steps, and existing few-step methods (12-50 steps) lose molecular stability; the goal is reliable high-quality 3D molecule generation in very few steps.

Approach

Adapt distribution matching distillation (reverse KL) to latent 3D molecule diffusion, fix its poor initialization with an EDM-inspired timestep respacing of the GeoLDM schedule, and add a discriminator-approximated Jensen-Shannon term to counter mode-seeking and restore diversity.

Key results

• On QM9 with 8 steps, FlashMol reaches 94.87% molecule stability versus 70.78% for the respaced GeoLDM and 89.4% for the 1000-step GeoLDM.
• With 4 steps on QM9, FlashMol achieves 83.89% molecule stability versus 12.43% for respaced GeoLDM, and 97.70% atom stability.
• On GEOM-DRUG with 5 steps, FlashMol reaches 89.98% atom stability at 8 steps and 84.72% at 4 steps, compared with 86.5% for the dataset and 84.4% for the 1000-step GeoLDM.

Surprising findings

• Respacing the timesteps of the teacher alone (no distillation) raises QM9 molecule stability at 16 steps to 94.61%, above trained few-step MOLTD at 16 NFEs.
• Adding a forward-KL term with lambda=0.1 to DMD sharply degrades stability (58.09%), while JS regularization improves diversity.

Limitations

• Diversity (Valid & Unique) declines as distillation training proceeds, and V&U is slightly below some baselines.
• Training is sensitive to learning rates; 4-step training needs a 10x larger learning rate and longer training causes overshooting.
• Evaluation relies on 3D metrics only; 2D-bond generators such as SemlaFlow were excluded from comparison.

Experimental setup

Datasets: QM9 GEOM-DRUG

Benchmarks: QM9 GEOM-DRUG

Metrics: atom stability molecule stability validity validity and uniqueness NFE conditional MAE

Baselines: EDM GeoLDM EDM-Bridge GDM-aug EquiFM GeoBFN GOAT SLDM GeoRCG MOLTD AccGeoLDM

Builds on: distribution matching distillation DMD2 GeoLDM EDM f-divergence GAN distillation

Closest papers by meaning

• GenMol: A Drug Discovery Generalist with Discrete Diffusion (ICML 2025 · similarity 0.89)
• DMol: A Highly Efficient and Chemical Motif-Preserving Molecule Generation Platform (NeurIPS 2025 · similarity 0.89)
• Efficient Molecular Conformer Generation with SO(3)-Averaged Flow Matching and Reflow (ICML 2025 · similarity 0.88)
• DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization (AAAI 2024 · similarity 0.88)
• Enhancing Generalizability in Molecular Conformation Generation with METRIZATION-Informed Geometric Diffusion Pretraining (AAAI 2025 · similarity 0.88)
• Swallowing the Bitter Pill: Simplified Scalable Conformer Generation (ICML 2024 · similarity 0.88)