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
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 M0 is perturbed via Gaussian noise for coordinates ( xt ) and categorical corruption for features ( zt ), defined by Eq. 1 and Eq. 3 . The equivariant network fθ 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 (xi,zi) to (x^i,z^i) ; (2) the network fθ is applied to both coordinates and features, with preconditioning affecting only the coordinate predictions; it outputs xpred and the category probabilities pθ(z) (Eq. 7 ); and (3) an update combining a continuous Euler step and a Discrete Flow Matching sampler to obtain Mi−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 ↑
V&C ↑
Bond Length ( ×10−2 ) ↓
Bond Angles ↓
Tor. ↓
Median ΔErelax↓
Mean ΔErelax↓
Median RMSD ↓
Mean RMSD ↓
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 ( ΔErelax ), 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 ρ , (middle) noise injection level γ , and (right) sampling steps. Black circles indicate selected values ( ρ=2.5 , γ=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 ↑
Mol Stab ↑
Valid ↑
EnergyMMFF↓
StrainMMFF↓
RMSDMMFF↓
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.
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 ↑
V&C ↑
Bond Length ( ×10−2 ) ↓
Bond Angles ↓
Torsions ↓
Median ΔErelax↓
Mean ΔErelax↓
Median RMSD ↓
Mean RMSD ↓
SemlaFlow( γ =0)
100
0.969
0.920
3.10
2.06
6.05
32.3
91.0
0.273
0.358
SemlaFlow( γ =0.01)
100
0.969
0.928
2.90
1.92
5.47
31.7
72.6
0.252
0.368
SemlaFlow( γ =0.03)
100
0.973
0.934
2.88
1.93
5.66
31.9
67.5
0.250
0.382
SemlaFlow( γ =0.1)
100
0.978
0.933
2.87
1.92
5.56
31.6
67.0
0.269
0.383
SemlaFlow( γ =0.2)
100
0.975
0.932
2.93
1.95
5.68
33.8
77.6
0.254
0.384
SemlaFlow( γ =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+ϵ⋅γ , where ϵ∼N(0,1) . Although SemlaFlow ( Irwin et al. 2025 ) reports γ=0.2 in their paper, their open-source implementation uses γ=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 ηt (solid orange line) and a variable η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 η . The results show that the variable η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 ↑
V&C ↑
Bond Length ( ×10−2 ) ↓
Bond Angles ↓
Torsions ↓
Median ΔErelax↓
Mean ΔErelax↓
Median RMSD ↓
Mean RMSD ↓
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.
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.
Junyi An, Chao Qu, Yun-Fei Shi +3
Shanghai Academy of Artificial Intelligence for Science, SAIS · Artificial Intelligence Innovation and Incubation (AI3) Institute, Fudan University
The generation of accurate 3D molecular conformations is a pivotal challenge in computational chemistry and drug discovery. Recently, diffusion and flow matching models have achieved remarkable success. However, there is a critical misalignment between their mathematical formulation and the physical reality of molecules. Existing approaches predominantly treat molecules as unstructured point clouds in Cartesian space, overlooking the intrinsic hierarchical mechanics where bond lengths and bond angles are relatively stiff, whereas torsion angles constitute the dominant flexible degrees of freedom. This lack of manifold awareness forces models to relearn fundamental geometric constraints from scratch, often leading to physically implausible intermediate structures. To address this, we propose GO-Flow that aligns generative modeling with molecular geometry via manifold decomposition. Instead of forcing motion through Euclidean space, GO-Flow decomposes the generation process into three physically motivated subspaces: translation space with linear optimal transport, rotation space with geodesic flows on SO(3), and conformation space with entropic optimal transport. This decomposition injects geometric inductive biases and makes the generative paths better aligned with molecular degrees of freedom. When combined with equivariant neural architectures, it encourages rotation-consistent generation and improves geometric validity. Extensive experiments on GEOM-Drugs and GEOM-QM9 demonstrate that GO-Flow achieves state-of-the-art generation quality. Notably, by learning straighter probability paths on the correct manifolds naturally, our method enables high-fidelity sampling with as few as 50 steps, effectively bridging the gap between structural precision and computational efficiency.
Yunqing Liu, Yi Zhou, Wenqi Fan
The Hong Kong Polytechnic University Hong Kong, China
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× acceleration in sampling speed while maintaining high molecular quality.
Xinyuan Wei, Zian Li, Shaoheng Yan +2
Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University +2