EVA-Flow: Learning Shared Chemistry for Unified Environment-Aware Molecular Conformation Generation

Bing Yan; Anuroop Sriram; Zachary Ward Ulissi; Kyunghyun Cho; Ricky T. Q. Chen; Benjamin Kurt Miller

cs.LGOct 26, 2026

EVA-Flow: Learning Shared Chemistry for Unified Environment-Aware Molecular Conformation Generation

Authors: Bing Yan, Anuroop Sriram, Zachary Ward Ulissi, Kyunghyun Cho, Ricky T. Q. Chen, Benjamin Kurt Miller

Organizations: FAIR at Meta · New York University

Abstract

Molecular 3D structures underpin simulations of thermodynamics, kinetics, crystallization, and nucleation, making their accurate prediction central to drug discovery and materials design. However, molecular geometry can change dramatically across environments. Existing models either ignore environmental effects or train separate models for each environment. We hypothesize that transfer learning across contexts is possible for conformer generation and that a general model will perform better when training data are limited, as in ligand binding and crystallization. We introduce EVA-Flow, a unified framework for environment-aware molecular conformation generation that learns shared chemical structure and adapts it through environment-conditioned generative modeling. Across vacuum, protein-ligand docking, solvation, and crystal packing, pretraining and cross-environment finetuning provide their largest gains in the three data-scarce settings. Evaluation against multiple reference conformations shows that EVA-Flow covers a broad range of solvation conformations. Analyses of molecules shared across environments support its ability to generate different conformations for the same molecule in different environments.

Figures & tables

Figure 1: Architecture of EVA-Flow. (a) Training. The environment graph

G_{E}

is embedded by

f_{\mathrm{env}}

into a vector

e

. The encoder

q_{\phi}(z\mid G,x,e)

maps the molecular graph

G

, ground-truth positions

x

, and environment

e

to a latent

z

. The base network

f_{\mathrm{base}}(G,x,e,z)

produces a distribution

p_{0}(x_{0}\mid G,x,e,z)

from which

x_{0}

is sampled. The FM decoder learns

u_{t}^{\theta}(x_{t},G,z,e)

with stochastic interpolation

x_{t}=(1{-}t)\,x_{0}+t\,x_{1}+\sigma(t)\epsilon

, targeting

u_{t}=(x_{1}-x_{0})+\dot{\sigma}(t)\epsilon

. Total loss:

\mathcal{L}=\mathcal{L}_{\mathrm{FM}}+\mathcal{L}_{\mathrm{KL\text{-}z}}+\mathcal{L}_{\mathrm{KL\text{-}base}}+\mathcal{L}_{\mathrm{consistency}}

. (b) Inference. The encoder and base network are bypassed:

e=f_{\mathrm{env}}(G_{E})

z\sim\mathcal{N}(\mathbf{0},I)

, and

x_{0}\sim p_{\mathrm{harm}}(G)

are sampled independently. All three condition the FM decoder, which generates conformations from

t{=}0

t{=}1

using the stochastic Euler sampler described in Section 3.2 . The Env Network and FM Decoder share weights between training and inference.

Stage/ environment	Target $G$	$G_{\mathrm{E}}$ nodes	$G_{\mathrm{E}}$ edges	Context attributes / coordinates	Available at inference
Pretraining	Molecule atoms/bonds	Random 30–70% atom-induced subgraph	Retained covalent bonds	Atom features and masked-subgraph coordinates	N/A (pretraining only)
Vacuum	Molecule atoms/bonds	Bemis–Murcko scaffold atoms	Scaffold covalent bonds	Atom features; coordinates set to zero	Yes, constructed from input molecule
Docking	Ligand atoms/bonds	Protein-pocket heavy atoms	Atoms within the same residue	Atom features and pocket coordinates	Yes, extracted pocket
Solvation	Solute atoms/bonds	Atoms of a water proxy	Water covalent bonds	sCAV, vCAV, eSOLV, eNELEC; fixed proxy coordinates	Yes, solvent descriptors
Packing	Molecule atoms/bonds	Copy of molecular atoms	Molecular covalent bonds	RDKit geometry, space group, lattice $(a,b,c,\alpha,\beta,\gamma)$	Yes, crystal metadata

Table 1: Roles and construction of the target graph

G

and environment graph

G_{\mathrm{E}}

. The target graph has the same semantics in every setting; only the context representation changes. “Coordinates” indicates coordinates supplied as input, not the target coordinates to be generated.

Panel of Figure 2 — Figure 2: Dataset composition. (a) Pretraining corpus by source ( $445{,}662$ molecule entries), counting SPICE and CREST-relaxed SPICE separately. (b) Finetuning corpus by environment ( $304{,}622$ molecule entries). (c) Reuse of all $272{,}641$ GEOM-DRUGS entries across pretraining and Vacuum finetuning; the remaining counts denote entries in the other datasets. (d) Data source and environment graph $G_{E}$ representation for each finetuning environment.

Panel of Figure 3 — Figure 3: Recall COV as a function of the RMSD threshold across the four environments. In every panel the four EVA-Flow training setups (Individual, Unified, Pretrain+Individual, Pretrain+Unified) are compared to retrained ETFlow. (a) Vacuum (cleaned GEOM-DRUGS ( Nikitin et al., 2025 ) ); reference points denote published Recall COV from Xu et al. (2022) ; Ganea et al. (2021) ; Jing et al. (2022) ; Wang et al. (2024) ; Hassan et al. (2024) . (b) Docking (PDBBind), (c) Solvation (Aquamarine), and (d) Packing (COD) compare models trained and evaluated under the same centered-conformation protocol. Dashed vertical lines mark the environment-specific RMSD thresholds ( $0.75/2.0/0.5/0.75\,\text{\AA}$ ).

Environment	Setup	Recall			Precision			sRMSD $\downarrow$
		COV $\uparrow$	AMR $\downarrow$		COV $\uparrow$	AMR $\downarrow$
		mean	mean	median	mean	mean	median
Vacuum	Pretrain+Unified	0.747 [-0.2ex] $\pm$ 0.013	0.468 [-0.2ex] $\pm$ 0.012	0.144 [-0.2ex] $\pm$ 0.024	0.463 [-0.2ex] $\pm$ 0.004	0.969 [-0.2ex] $\pm$ 0.006	0.900 [-0.2ex] $\pm$ 0.023	–
	Pretrain+Individual	0.733 [-0.2ex] $\pm$ 0.010	0.481 [-0.2ex] $\pm$ 0.012	0.146 [-0.2ex] $\pm$ 0.019	0.454 [-0.2ex] $\pm$ 0.004	0.980 [-0.2ex] $\pm$ 0.008	0.925 [-0.2ex] $\pm$ 0.016	–
	Unified	0.740 [-0.2ex] $\pm$ 0.007	0.456 [-0.2ex] $\pm$ 0.008	0.131 [-0.2ex] $\pm$ 0.011	0.463 [-0.2ex] $\pm$ 0.003	0.954 [-0.2ex] $\pm$ 0.005	0.879 [-0.2ex] $\pm$ 0.011	–
	Individual	0.728 [-0.2ex] $\pm$ 0.005	0.487 [-0.2ex] $\pm$ 0.004	0.134 [-0.2ex] $\pm$ 0.006	0.460 [-0.2ex] $\pm$ 0.003	0.978 [-0.2ex] $\pm$ 0.004	0.911 [-0.2ex] $\pm$ 0.017	–

Table 2: Conformation generation under the 4-conformer protocol (default) across four environments. Each model generates

4

conformers per molecule, compared with one reference. We report Recall and Precision COV (

\uparrow

, mean) and AMR (

\downarrow

, mean and median) at environment-specific RMSD thresholds:

0.5~\text{\AA}

for Solvation,

0.75~\text{\AA}

for Vacuum and Packing, and

2.0~\text{\AA}

for Docking. For Packing, we additionally report symmetry-aware RMSD (sRMSD,

\downarrow

, mean). Except for sRMSD, metric cells are mean

\pm

standard deviation over five independent SDE sampling seeds from one fixed trained checkpoint (Appendix C.1 ). Best means within each environment are bolded .

Setup	Recall				Precision
	COV $\uparrow$		AMR $\downarrow$		COV $\uparrow$		AMR $\downarrow$
	mean	median	mean	median	mean	median	mean	median
GeoDiff	0.421	0.378	0.835	0.809	0.249	0.145	1.136	1.090
GeoMol	0.446	0.414	0.875	0.834	0.430	0.364	0.928	0.841
Torsional Diff.	0.727	0.800	0.582	0.565	0.552	0.569	0.778	0.729
MCF-L	0.847	0.922	0.390	0.247	0.668	0.713	0.618	0.530

Table 3: Conformation generation for Vacuum under the 2K protocol : each model generates

2K

conformers per molecule, where

K

is the number of reference conformers. We report mean and median Recall and Precision COV (

\uparrow

) and AMR (

\downarrow

) at RMSD threshold

0.75~\text{\AA}

. Best results are bolded . The models are trained on the original GEOM-DRUGS dataset and evaluated on its full test set. Literature numbers for GeoDiff, GeoMol, Torsional Diffusion, MCF-L, and ETFlow (reported) are taken from Table 1 of Hassan et al. (2024) ; ETFlow (reported) refers to the ET-Flow base (

8.3

M) variant. ETFlow (retrained) and EVA-Flow are trained on all available conformers per molecule on the original GEOM-DRUGS and evaluated on the same test set.

Panel of Table 4 — Table 4: Reference and generated conformations (Pretrain+Unified) for molecules shared between Solvation and Packing or Solvation and Docking. The conformers are aligned to the Solvation reference.

Panel of Figure 4 — Figure 4: Shared-molecule analysis (Pretrain+Unified), using values from Table 5 . (a) Pairwise counts of shared molecules; darker cells indicate more molecules, and dashes mark the diagonal. (b, c) Mean heavy-atom RMSD (Å) after alignment for five comparisons. In both panels, $A$ denotes Solvation; $B$ denotes Packing in (b) and Docking in (c). “Own $A$ ” and “Own $B$ ” compare each conditioned generation with its own environment’s reference. $A\!\rightarrow\!B$ compares the $A$ -conditioned generation with the reference from $B$ , and conversely for $B\!\rightarrow\!A$ . “Generated A vs. generated B” compares the two environment-conditioned generations and corresponds to generation shift in Table 5 .

Pair ( $A/B$ )	$n$	own $A$	own $B$	cross $A$	cross $B$	generation shift
Solvation–Vacuum	142	0.12	0.64	3.41	3.31	3.41
Docking–Vacuum	558	0.75	0.76	1.22	1.21	1.20
Vacuum–Packing	487	0.33	0.35	0.64	0.65	0.64
Docking–Solvation	13	0.55	0.09	2.41	2.49	2.42
Packing–Solvation	23	0.06	0.09	1.49	1.48	1.48
Packing–Docking	64	0.28	0.36	0.45	0.38	0.42

Table 5: Shared-molecule analysis (Pretrain+Unified), corresponding to Figure 4 . All values are mean heavy-atom RMSD in Å.

Figure 5: Recall COV for the small model in the Unified setup across Vacuum, Docking, Solvation, and Packing. ‘EVA’ refers to the standard setup; ‘GCN’ uses non-equivariant networks for the encoder and base; ‘Gaussian’ uses a Gaussian base distribution instead of a Harmonic one; ‘NoConsist’ removes latent consistency loss; ‘NoBase’ replaces the learned base distribution network for initial positions

x_{0}

with a Harmonic distribution determined only by the graph; ‘NoFM’ removes the FM decoder and uses a regression decoder.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Histograms of atom and rotatable-bond counts for the pretraining datasets: GEOM DRUGS, GEOM QM9, SPICE, CREST-relaxed SPICE, and OE62.

Panel of Figure 7 — Figure 7: Histograms of atom and rotatable-bond counts for PDBBind, AQM, and COD. GEOM DRUGS is shared with the pretraining datasets (Figure 6 ).

Hyperparameter	Value
Batch size	32
FM decoder layers	20
Hidden channels	160
Attention heads	8
Latent dim ( $d_{z}$ )	128
Environment dim ( $d_{e}$ )	128

Appendix

Table 6: Hyperparameters for EVA-Flow.

Method	Sampler	Steps	R-COV $\uparrow$	Time (s/mol) $\downarrow$
ETFlow (retrained)	SDE	50	0.725	0.309
EVA-Flow (P+U)	SDE	50	0.747	0.266
EVA-Flow (P+U)	SDE	100	0.752	0.530

Appendix

Table 7: Per-molecule inference time on a single A100 GPU (GEOM-DRUGS,

1{,}000

molecules). EVA-Flow uses the Pretrain+Unified setup.

Panel of Figure 8 — Figure 8: Precision Coverage for all four environments: Vacuum, Docking, Solvation, and Packing.

Environment	Setup	R-COV median $\uparrow$	P-COV median $\uparrow$
Vacuum	Pretrain+Unified	1.00	0.50
	Pretrain+Individual	1.00	0.50
	Unified	1.00	0.50
	Individual	1.00	0.50
	ETFlow (retrained)	1.00	0.50
Docking	Pretrain+Unified	1.00	0.50

Appendix

Table 8: Median Recall and Precision COV under the 4-conformer protocol.

Environment	Mean $\uparrow$	Median	Std
Vacuum	0.915	0.903	0.620
Docking	2.227	1.828	1.508
Solvation	0.456	0.359	0.315
Packing	0.967	0.912	0.634

Appendix

Table 9: Pairwise-RMSD diversity (Å) among the four generated conformers per molecule (Pretrain+Unified).

Training	Protocol	Refs/mol	Samples/mol	R-COV $\uparrow$	R-AMR $\downarrow$	P-COV $\uparrow$	P-AMR $\downarrow$
Single conformer	4-conf.	1	4	0.936	0.157	0.693	0.431
Single conformer	2K	mean 37	$2K$	0.981	0.102	0.541	0.591
Multiple conformers	2K	mean 37	$2K$	0.962	0.111	0.518	0.612

Appendix

Table 10: Solvation evaluation with single- and multi-reference protocols.

Environment	Method	R-COV $\uparrow$	R-AMR $\downarrow$
Vacuum	ETKDG+MMFF	0.409	1.18
	EVA-Flow (P+U)	0.747	0.468
Docking	ETKDG+MMFF	0.581	2.09
	EVA-Flow (P+U)	0.684	1.899
Solvation	ETKDG+MMFF	0.006	3.42
	GFN2-xTB+ALPB-water	0.006	3.38

Appendix

Table 11: Non-learned reference methods under the 4-conformer protocol.

Model	Training set	Test set	R-COV $\uparrow$	R-AMR $\downarrow$	P-COV $\uparrow$	P-AMR $\downarrow$
Published ETFlow checkpoint	Original	Original	0.549	0.766	0.339	1.057
ETFlow (retrained)	Original	Original	0.712	0.562	0.704	0.594
EVA-Flow	Original	Original	0.721	0.550	0.707	0.589
Published ETFlow checkpoint	Original	Cleaned	0.517	0.799	0.294	1.193
ETFlow (retrained)	Original	Cleaned	0.759	0.505	0.629	0.688
EVA-Flow	Original	Cleaned	0.763	0.503	0.632	0.686

Appendix

Table 12: Matched

2K

evaluation (

0.75

Å) on original and cleaned GEOM-DRUGS test sets. All rows within a block use identical evaluation code and test references.

Metric	Value
UMA relaxation converged	99.9% (3,995/4,000)
Mean $\Delta E$ after relaxation	0.302 kcal/mol
Median $\Delta E$ after relaxation	0.115 kcal/mol
Fraction with $\Delta E<1$ kcal/mol	97.1%
Mean RMSD shift during relaxation	0.056 Å
Fraction with RMSD shift $<0.1$ Å	84.2%

Appendix

Table 13: UMA relaxation of EVA-Flow Vacuum conformers (

4{,}000

conformers from

1{,}000

molecules).

Pair ( $A/B$ )	$n$	own $A$	own $B$	cross $A$	cross $B$	generation shift
Solvation–Vacuum	142	0.29	0.90	3.35	3.25	3.32
Docking–Vacuum	558	0.95	0.81	1.17	1.19	1.16
Vacuum–Packing	487	0.38	0.43	0.64	0.59	0.62
Docking–Solvation	13	0.63	0.35	2.34	2.27	2.23
Packing–Solvation	23	0.14	0.14	1.48	1.47	1.50
Packing–Docking	64	0.36	0.46	0.51	0.37	0.40

Appendix

Table 14: Shared-molecule analysis after excluding all shared molecules from finetuning.

Environment	Ablation	R-COV $\uparrow$	R-AMR $\downarrow$	P-COV $\uparrow$	P-AMR $\downarrow$
Solvation	Full	0.949	0.153	0.780	0.367
Solvation	NoSolvDesc	0.905	0.191	0.734	0.399
Packing	Full	0.644	0.664	0.410	1.073
Packing	NoLattice	0.575	0.762	0.360	1.154
Docking	Full	0.684	1.899	0.512	2.471
Docking	NoPocket	0.400	2.590	0.257	3.197

Appendix

Table 15: Environment-feature sensitivity (inference-time zeroing, no retraining). Thresholds: Solvation

0.5

Å, Packing

0.75

Å, Docking

2.0

Å.

Metric	Setup	Vacuum	Docking	Solvation	Packing
R-COV $\uparrow$	EVA (Full)	0.740	0.674	0.930	0.640
	GCN	0.670	0.665	0.803	0.605
	Gaussian	0.405	0.237	0.822	0.312
	NoConsist	0.338	0.577	0.134	0.373
	NoBase	0.321	0.577	0.255	0.398
	NoFM	0.000	0.126	0.006	0.005

Appendix

Table 16: Full ablation metrics (Unified setup, small model). Thresholds: Vacuum/Packing

0.75

Å, Solvation

0.5

Å, Docking

2.0

Å. Best per metric/environment in bold .

Environment	Individual	Unified	Pretrain+Ind.	Pretrain+Unified
Vacuum	0.728	0.740	0.733	0.747
Docking	0.553	0.674	0.656	0.684
Solvation	0.414	0.930	0.936	0.949
Packing	0.512	0.640	0.633	0.644

Appendix

Table 17: Full architecture (reported model): Recall COV at the environment threshold across all four training regimes.

Environment	Individual	Unified	Pretrain+Ind.	Pretrain+Unified
Vacuum	0.727	0.708	0.688	0.672
Docking	0.470	0.660	0.614	0.628
Solvation	0.503	0.675	0.771	0.771
Packing	0.493	0.622	0.615	0.613

Appendix

Table 18: Latent-free model (learned base and latent removed): Recall COV at the environment threshold across all four training regimes. Compared with the full model (Table 17 ), Individual performance is nearly unchanged in Vacuum, higher in Solvation, and lower in Docking and Packing. All twelve transfer comparisons favor the full model, with the largest gaps in Solvation.

Environment	Dataset	Unified/ep	Unified tot.	Indiv./ep	Indiv. tot.
Vacuum (DRUGS)	272,641	60,000	48.0M	272,641	218.1M
Packing (COD)	20,765	20,000	16.0M	20,765	16.6M
Solvation (AQM)	1,276	1,276	1.02M	1,276	1.02M
Docking (PDB)	9,940	9,940	7.95M	9,940	7.95M

Appendix

Table 19: Top: total molecule exposure across training. Bottom: balanced vs. default sampling (Recall COV at the environment threshold).

$\epsilon$	R-COV $\uparrow$	R-AMR $\downarrow$	P-COV $\uparrow$	P-AMR $\downarrow$
0.00 (default)	0.747	0.468	0.464	0.979
0.01	0.738	0.470	0.455	0.983
0.05	0.735	0.493	0.459	0.981
0.10	0.745	0.471	0.464	0.973
0.20	0.730	0.486	0.452	0.985
0.50	0.734	0.480	0.453	0.982

Appendix

Table 20: Inference noise-injection sweep on Vacuum (Pretrain+Unified checkpoint, SDE

50

steps): perturbing the harmonic prior by

\epsilon\cdot\mathcal{N}(\mathbf{0},I)

Model	Active dims	KL/node (nats)	Pooled $\\|\mu\\|$	Environment-dependent
Full	0/128	$\approx 10^{-6}$	$\approx 10^{-3}$	No (cosine 1.000)
NoBase	3/128	up to $\sim 3$	$\approx 0.21$	Yes

Appendix

Table 21: Latent activation in the full and NoBase models.

Environment	KL/node mean	KL/node median	Active dims	Posterior $\|\mu\|$	R-COV random	R-COV zero	P-COV random	P-COV zero
Vacuum	$1.5{\times}10^{-6}$	$1.4{\times}10^{-6}$	0/128	$8.2{\times}10^{-5}$	0.707	0.706	0.433	0.433
Docking	0.0476	$1.4{\times}10^{-6}$	0/128	$1.9{\times}10^{-4}$	0.661	0.661	0.497	0.497
Solvation	$1.6{\times}10^{-6}$	$1.5{\times}10^{-6}$	0/128	$8.2{\times}10^{-5}$	0.803	0.809	0.580	0.591
Packing	$1.6{\times}10^{-6}$	$1.5{\times}10^{-6}$	0/128	$8.2{\times}10^{-5}$	0.649	0.649	0.419	0.418

Appendix

Table 22: Latent-collapse diagnostics and intervention. “Random” samples

z\sim\mathcal{N}(\mathbf{0},I)

; “zero” replaces the matched latent draw by zero. COV values use the environment-specific thresholds.

Weight change	Vacuum	Docking	Solvation	Packing
latent-KL $0.1\times$	$-0.018$	$-0.010$	$-0.020$	$-0.007$
latent-KL $10\times$	$-0.045$	$+0.018$	$-0.141$	$-0.030$
base-KL $0.1\times$	$-0.027$	$+0.004$	$+0.019$	$-0.008$
base-KL $10\times$	$-0.014$	$+0.032$	$+0.019$	$-0.015$
consistency $0.1\times$	$-0.018$	$+0.009$	$-0.064$	$-0.010$
consistency $10\times$	$-0.004$	$+0.028$	$+0.025$	$-0.001$

Appendix

Table 23: One-factor-at-a-time loss sensitivity: change in Recall COV from unit weights.

Environment	Coordinate content of $G_{E}$	$\Delta$ R-COV	$\Delta$ R-AMR
Vacuum	zero-coordinate scaffold	$-0.003$	$-0.011$
Docking	physical pocket coordinates	$+0.028$	$-0.082$
Solvation	fixed proxy; scalar descriptors	$-0.153$	$+0.139$
Packing	RDKit proxy and lattice metadata	$-0.006$	$+0.015$

Appendix

Table 24: Equivariant environment encoder minus matched GCN control. Positive

\Delta

R-COV and negative

\Delta

R-AMR favor the equivariant encoder.

Reference selection	R-COV $\uparrow$	R-AMR $\downarrow$	P-COV $\uparrow$	P-AMR $\downarrow$
c1 (default)	0.803	0.291	0.580	0.538
Minimum eSOLV	0.777	0.319	0.588	0.538
Deterministic random	0.803	0.292	0.584	0.538

Appendix

Table 25: Aquamarine test-reference sensitivity for a fixed model and sampling seed (

0.5

Å threshold).

Training-target selection	R-COV $\uparrow$
c1 (default rule)	0.580
Minimum eSOLV	0.561
Deterministic random	0.618

Appendix

Table 26: Aquamarine training-target sensitivity in matched 200-epoch runs.

Panel of Figure 9 — Figure 9: Pairwise heavy-atom RMSD ( Å ) between reference and generated conformers in the original overlap sets: (a) Solvation-Docking ( $n=12$ ), and (b) Packing-Docking ( $n=59$ ). In each heatmap we report the mean RMSD after alignment for: ref_A $\leftrightarrow$ ref_B, gen_A $\leftrightarrow$ ref_A, gen_A $\leftrightarrow$ ref_B, gen_B $\leftrightarrow$ ref_A, gen_B $\leftrightarrow$ ref_B, and gen_A $\leftrightarrow$ gen_B (A/B = the two tasks in the panel). Darker cells indicate lower RMSD (better agreement).

Panel of Table 27 — Table 27: Visualization of reference and generated conformations for the shared molecules in Packing-Docking. The molecules are aligned to the Docking reference conformers.

3D Molecular Generation

3D molecular generation creates molecules by modeling their chemical structures and three-dimensional geometries.

See all 55 papers

FlashMol: High-Quality Molecule Generation in as Few as Four Steps
May 7, 2026Xinyuan Wei, Zian Li, Shaoheng Yan +2Diffusion Distillation Few-Step Diffusion Sampling
Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University +2
Generating Physically Consistent Molecules with Energy-Based Models
May 18, 2026Christoph Griesbacher, Lea Bogensperger, Andreas Habring +1Energy-Based Models 3D Molecular Generation
Graz University of Technology · University of Zurich
VEDA: Generation of 3D Molecules via Variance-Exploding Diffusion with Annealing
Nov 11, 2025Peining Zhang, Jinbo Bi, Minghu Song3D Molecular Generation Diffusion Models
University of Connecticut, Storrs, Connecticut 06269, USA · Institute of Health and Medicine, Hefei Comprehensive National Science Center, Hefei 230601, China

May 7, 2026J/K move · Enter open · S save

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

Xinyuan Wei, Zian Li, Shaoheng Yan, Cai Zhou+1

Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University+2

Diffusion Distillation Few-Step Diffusion Sampling 3D Molecular Generation Molecular Generation Distribution Matching Distillation Conformer

Open paper AskPDF arXiv

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.

Abstract

3D Molecular Generation

Explore similar work

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

Generating Physically Consistent Molecules with Energy-Based Models

VEDA: Generation of 3D Molecules via Variance-Exploding Diffusion with Annealing