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

Peining Zhang; Jinbo Bi; Minghu Song

physics.chem-phNov 11, 2025

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

Authors: Peining Zhang, Jinbo Bi, Minghu Song

Organizations: University of Connecticut, Storrs, Connecticut 06269, USA · Institute of Health and Medicine, Hefei Comprehensive National Science Center, Hefei 230601, China

Abstract

Diffusion models show promise for 3D molecular generation, but face a fundamental trade-off between sampling efficiency and conformational accuracy. While flow-based models are fast, they often produce geometrically inaccurate structures, as they have difficulty capturing the multimodal distributions of molecular conformations. In contrast, denoising diffusion models are more accurate but suffer from slow sampling, a limitation attributed to sub-optimal integration between diffusion dynamics and SE(3)-equivariant architectures. To address this, we propose VEDA, a unified SE(3)-equivariant framework that combines variance-exploding diffusion with annealing to efficiently generate conformationally accurate 3D molecular structures. Specifically, our key technical contributions include: (1) a VE schedule that enables noise injection functionally analogous to simulated annealing, improving 3D accuracy and reducing relaxation energy; (2) a novel preconditioning scheme that reconciles the coordinate-predicting nature of SE(3)-equivariant networks with a residual-based diffusion objective, and (3) a new arcsin-based scheduler that concentrates sampling in critical intervals of the logarithmic signal-to-noise ratio. On the QM9 and GEOM-DRUGS datasets, VEDA matches the sampling efficiency of flow-based models, achieving state-of-the-art valency stability and validity with only 100 sampling steps. More importantly, VEDA's generated structures are remarkably stable, as measured by their relaxation energy during GFN2-xTB optimization. The median energy change is only 1.72 kcal/mol, significantly lower than the 32.3 kcal/mol from its architectural baseline, SemlaFlow. Our framework demonstrates that principled integration of VE diffusion with SE(3)-equivariant architectures can achieve both high chemical accuracy and computational efficiency.

Figures & tables

Figure 1: An overview of the VEDA framework, detailing its training and sampling processes. During Training (top), a clean molecule $\mathcal{M}_{0}$ is perturbed via Gaussian noise for coordinates ( $\mathbf{x}_{t}$ ) and categorical corruption for features ( $\mathbf{z}_{t}$ ), defined by Eq. 1 and Eq. 3 . The equivariant network $f_{\theta}$ is then trained to predict the original molecule by minimizing the combined Mean Squared Error (MSE) and Cross-Entropy (CE) loss in Eq. 13 . During Sampling (bottom), the process starts from a pure noise distribution (a Gaussian point cloud with uniform categorical features) and iteratively refines the sample over $K$ steps. Each integration step $i$ involves: (1) a noise injection from $(\mathbf{x}_{i},\mathbf{z}_{i})$ to $(\mathbf{\hat{x}}_{i},\mathbf{\hat{z}}_{i})$ ; (2) the network $f_{\theta}$ is applied to both coordinates and features, with preconditioning affecting only the coordinate predictions; it outputs $x_{\text{pred}}$ and the category probabilities $p^{(z)}_{\theta}$ (Eq. 7 ); and (3) an update combining a continuous Euler step and a Discrete Flow Matching sampler to obtain $\mathcal{M}_{i-1}$ , formally given in Eq. 14 and Eq. 15 . In the diagram, blue boxes represent the main network module, gray boxes are auxiliary operations, and arrows indicate the data flow.

Figure 2: The arcsin sampling scheduler is proposed to focus on the middle part of sampling where log(SNR) is close to 0

Methods	NFE ↓	Atom Sta(%)↑	Mol Sta(%)↑	Valid (%)↑	Valid & Unique(%)↑
EDM	1000	98.7	82	91.9	90.9
GeoLDM	1000	98.9	89.4	93.8	92.6
UniGEM	1000	99.0	89.8	95	93.2
EquiFM	210	98.9	88.3	94.7	93.5
GeoBFN	2000	99.3	93.3	96.9	92.4
GOAT	90	99.2	-	92.9	92.0

Table 1: Results on QM9. Top: bond‑implicit methods (bonds inferred from coordinates); Bottom: bond‑explicit methods (bonds generated directly). NFE = number of function evaluations (sampling steps). Best in each group in bold .

Model	NFE	Time (s)	MS $\uparrow$	V&C $\uparrow$	Bond Length ( $\times 10^{-2}$ ) $\downarrow$	Bond Angles $\downarrow$	Tor. $\downarrow$	Median $\Delta E_{relax}$ $\downarrow$	Mean $\Delta E_{relax}$ $\downarrow$	Median RMSD $\downarrow$	Mean RMSD $\downarrow$
EQGAT	500	3468	0.899	0.834	1.00	1.15	8.58	6.40	11.1	0.915	0.975
JODO	500	673	0.963	0.879	0.77	0.83	6.01	4.74	7.04	-	-
Megalodon-quick	500	-	0.944	0.900	0.66	0.71	5.58	3.19	5.76	-	-
FlowMol2	100	126.3	0.944	0.746	1.30	1.62	15.0	17.9	24.3	0.992	1.038
SemlaFlow	100	150.3	0.969	0.920	3.10	2.06	6.05	32.3	91.0	0.273	0.358
Megalodon-flow	100	-	0.987	0.948	2.30	1.62	5.58	20.9	46.9	-	-

Table 2: Performance comparison on the GEOM-DRUGS dataset. Our model, VEDA-S (bottom), is compared against denoising-based (top) and flow-based (middle) methods. Bold values indicate the best result. Metrics include Molecular Stability (MS), Validity & Connectivity (V&C), and Mean Absolute Error (MAE) for bond lengths, angles, and torsions. All metrics are lower-is-better except for MS and V&C. Time (s) reports the average wall-clock time for generating one molecule. NFE is the Number of Function Evaluations. RMSD values for some models are missing as they were not in the original benchmark. For each model, 5000 molecules were evaluated, the full results with confidence intervals and sampling at other steps are included in Appendix B .

Figure 3: Trade-off between Generation Quality and Computational Cost. The figure compares our model, VEDA-S (red), against flow-based (green) and denoising-based (blue) models. The horizontal axis represents the computational cost, measured by the Number of Function Evaluations (NFE). (Left) Quality measured by molecule Validity and Connectivity, where higher values are better. (Right) Quality measured by the median energy difference ( $\Delta E_{relax}$ ), where lower values are better.

Figure 4: Ablation study of key hyperparameters on GEOM-DRUGS: We report MMFF94 energy and validity when varying (left) arcsin noise factor $\rho$ , (middle) noise injection level $\gamma$ , and (right) sampling steps. Black circles indicate selected values ( $\rho{=}2.5$ , $\gamma{=}0.4$ , 100 steps).

Variant	OT	Noise	Pre	V&C ↑	M.S.↑	E. ↓	R. ↓
VEDA (Full)	✗	✓	✓	98.8	99.5	1.72	0.096
– Pre	✗	✓	✗	97.9	98.5	2.14	0.153
– Noise	✗	✗	✓	95.9	97.6	6.29	0.382
– Noise & Pre	✗	✗	✗	93.8	92.9	10.1	0.409
+ OT Align	✓	✗	✗	91.7	96.0	23.5	0.557
SemlaFlow	✓	✗	✗	92.0	96.9	32.3	0.273

Table 3: Ablation study on GEOM-DRUGS. Each variant enables (✓) or disables (✗) key components of VEDA: OT alignment (OT), noise injection (Noise), and preconditioning (Pre). V&C ↑ and M.S.↑ are reported as percentage-based metrics. The last two rows share the same component configuration and differ only in their time parameterization (VE vs. FM).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	QM9		GEOM-DRUGS
	VEDA-E	VEDA-S	VEDA-E	VEDA-S
Learning rate	3e-4	3e-4	1e-3	3e-4
Warm up steps	2000	15000	2000	10000
Batch Size	32	N/A	4	N/A
Batch Cost	N/A	4096	N/A	4096
EMA decay rate	0.9999	0.999	0.9999	0.999

Appendix

Table 4: Hyperparameters used for training VEDA-E and VEDA-S models

Methods	NFE	Atom Stab $\uparrow$	Mol Stab $\uparrow$	Valid $\uparrow$	$\text{Energy}_{\text{MMFF}}$ $\downarrow$	$\text{Strain}_{\text{MMFF}}$ $\downarrow$	$\text{RMSD}_{\text{MMFF}}$ $\downarrow$
EDM	1000	81.3	N/A	92.6	-	-	-
GeoLDM	1000	84.4	N/A	99.6	-	-	-
UniGEM	1000	85.1	N/A	98.4
EquiFM	1000	84.1	N/A	98.9	-	-	-
GOAT	90	84.8	N/A	96.2	-	-	-
VEDA-E	50	88.7	3.6	97.3

Appendix

Table 5: Molecular generation results on the GEOM-DRUGS dataset. VEDA-S with Semla as the backbone. Molecule stability and validity are reported as percentages, while energy and strain energy are expressed in kcal·mol -1 . Results marked with ∗ were reproduced in our own experiments. All results are calculated based on 1000 randomly sampled molecules.

Panel of Figure 5 — Figure 5: Quality–efficiency trade‑off for VEDA‑E on QM9 under different sampling strategies, compared with EDM ( Hoogeboom et al. 2022 ) . Even at low NFE, our model sustains high validity and structural stability.

Model	NFE	MS $\uparrow$	V&C $\uparrow$	Bond Length ( $\times 10^{-2}$ ) $\downarrow$	Bond Angles $\downarrow$	Torsions $\downarrow$	Median $\Delta E_{relax}$ $\downarrow$	Mean $\Delta E_{relax}$ $\downarrow$	Median RMSD $\downarrow$	Mean RMSD $\downarrow$
SemlaFlow( $\gamma$ =0)	100	0.969	0.920	3.10	2.06	6.05	32.3	91.0	0.273	0.358
SemlaFlow( $\gamma$ =0.01)	100	0.969	0.928	2.90	1.92	5.47	31.7	72.6	0.252	0.368
SemlaFlow( $\gamma$ =0.03)	100	0.973	0.934	2.88	1.93	5.66	31.9	67.5	0.250	0.382
SemlaFlow( $\gamma$ =0.1)	100	0.978	0.933	2.87	1.92	5.56	31.6	67.0	0.269	0.383
SemlaFlow( $\gamma$ =0.2)	100	0.975	0.932	2.93	1.95	5.68	33.8	77.6	0.254	0.384
SemlaFlow( $\gamma$ =0.4)	100	0.970	0.926	2.88	1.93	5.61	31.6	71.0	0.256	0.387

Appendix

Table 6: Performance comparison on the GEOM-DRUGS dataset using the benchmark from Nikitin et al. (2025) . Noise is added according to

v=v+\epsilon\cdot\gamma

, where

\epsilon\sim\mathcal{N}(0,1)

. Although SemlaFlow ( Irwin et al. 2025 ) reports

\gamma=0.2

in their paper, their open-source implementation uses

\gamma=0

Figure 6: Comparison of re-masking strategies. The x-axis is the sampling steps. The performance of a Uniform Random strategy is compared against two confidence-based heuristics: selecting tokens with the lowest prediction confidence ( Confidence ) or the smallest probability margin ( Prob_margin ). The random strategy consistently achieves lower energy, lower optimization RMSD, and higher validity, indicating its superiority for this generative task.

Figure 7: Hyperparameter tuning for DFM sampling strategies. We compare the performance of two distinct sampling approaches derived from the Discrete Flow Matching (DFM) framework. The first approach, based on the transition rate formulation of Campbell et al. (2024) , is evaluated with both a fixed $\eta_{t}$ (solid orange line) and a variable $\eta_{t}$ schedule (dashed blue line). The second approach is the simplified velocity-field method from Gat et al. (2024) , denoted as Gat’DFM (dotted green line). Performance is measured across three metrics—Energy per Atom, Optimization RMSD, and Validity—as a function of the categorical noise level hyperparameter $\eta$ . The results show that the variable $\eta_{t}$ schedule provides the most robust trade-off, significantly minimizing optimization RMSD while maintaining high structural validity across a wide range of noise levels.

Model	NFE	MS $\uparrow$	V&C $\uparrow$	Bond Length ( $\times 10^{-2}$ ) $\downarrow$	Bond Angles $\downarrow$	Torsions $\downarrow$	Median $\Delta E_{relax}$ $\downarrow$	Mean $\Delta E_{relax}$ $\downarrow$	Median RMSD $\downarrow$	Mean RMSD $\downarrow$
EQGAT	250	0.318	0.605	2.63	3.01	17.63	34.95	76.49	0.974	1.044
EQGAT	500	0.899	0.834	1.00	1.15	8.58	6.40	11.1	0.915	0.975
JODO	500	0.963	0.879	0.77	0.83	6.01	4.74	7.04	-	-
Megalodon-quick	500	0.944	0.900	0.66	0.71	5.58	3.19	5.76	-	-
FlowMol2	50	0.750	0.694	1.87	1.99	16.22	23.9	31.34	0.947	1.024
FlowMol2	100	0.830	0.768	1.54	1.78	15.45	21.0	27.1	0.992	1.038

Appendix

Table 7: Performance comparison on the GEOM-DRUGS dataset following the benchmark of Nikitin et al. (2025) . Results for FlowMol2, EQGAT, SemlaFlow, and VEDA-S are reported under different sampling steps. VEDA-S is evaluated using 5000 generated samples at 50 and 100 steps, with standard deviations computed over 5 subsets. All other methods are evaluated on 1,000 generated samples.

Explore similar work

Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation
Mar 10, 2026Junyi An, Chao Qu, Yun-Fei Shi +33D Molecular Generation Molecular Generation
Shanghai Academy of Artificial Intelligence for Science, SAIS · Artificial Intelligence Innovation and Incubation (AI3) Institute, Fudan University
Geometric Flow Matching for Molecular Conformation Generation via Manifold Decomposition
May 25, 2026Yunqing Liu, Yi Zhou, Wenqi FanFlow Matching Generative Modeling
The Hong Kong Polytechnic University Hong Kong, China
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

Mar 10, 2026J/K move · Enter open · S save

Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation

Junyi An, Chao Qu, Yun-Fei Shi, Zhijian Zhou+2

Shanghai Academy of Artificial Intelligence for Science, SAIS · Artificial Intelligence Innovation and Incubation (AI3) Institute, Fudan University

3D Molecular Generation Molecular Generation Diffusion Models 3D Diffusion Models Diffusion Model Acceleration Diffusion Timestep Scheduling

Open paper AskPDF arXiv

Recent 3D molecular generation methods primarily use asynchronous auto-regressive or synchronous diffusion models. While auto-regressive models build molecules sequentially, they're limited by a short horizon and a discrepancy between training and inference. Conversely, synchronous diffusion models denoise all atoms at once, offering a molecule-level horizon but failing to capture the causal relationships inherent in hierarchical molecular structures. We introduce Equivariant Asynchronous Diffusion (EAD) to overcome these limitations. EAD is a novel diffusion model that combines the strengths of both approaches: it uses an asynchronous denoising schedule to better capture molecular hierarchy while maintaining a molecule-level horizon. Since these relationships are often complex, we propose a dynamic scheduling mechanism to adaptively determine the denoising timestep. Experimental results show that EAD achieves state-of-the-art performance in 3D molecular generation.

Abstract

Explore similar work

Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation

Geometric Flow Matching for Molecular Conformation Generation via Manifold Decomposition

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