# Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation

arXiv: https://arxiv.org/abs/2606.01595 · HTML: https://mlsanity.com/paper/2606.01595/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 50 references that are in the map.
- First posted on arXiv 1 Jun 2026.

## Builds on (most cited first, 29 of 29)

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [GraphVAE: Towards Generation of Small Graphs Using Variational Autoencoders](https://mlsanity.com/paper/1802.03480/) (2018)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [Geometric Latent Diffusion Models for 3D Molecule Generation](https://mlsanity.com/paper/2305.01140/) (2023)
- [Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules](https://mlsanity.com/paper/1906.00957/) (2019)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems](https://mlsanity.com/paper/2307.08423/) (2023)
- [Less is more: sampling chemical space with active learning](https://mlsanity.com/paper/1801.09319/) (2018)
- [ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs](https://mlsanity.com/paper/2206.08515/) (2022)
- [Variational Bayesian Last Layers](https://mlsanity.com/paper/2404.11599/) (2024)
- [Learning Hierarchical Protein Representations via Complete 3D Graph Networks](https://mlsanity.com/paper/2207.12600/) (2022)
- [BayesDiff: Estimating Pixel-wise Uncertainty in Diffusion via Bayesian Inference](https://mlsanity.com/paper/2310.11142/) (2023)
- [DiffMS: Diffusion Generation of Molecules Conditioned on Mass Spectra](https://mlsanity.com/paper/2502.09571/) (2025)
- [Single-model uncertainty quantification in neural network potentials does not consistently outperform model ensembles](https://mlsanity.com/paper/2305.01754/) (2023)
- [UniGEM: A Unified Approach to Generation and Property Prediction for Molecules](https://mlsanity.com/paper/2410.10516/) (2024)
- [Diffusion Model Guided Sampling with Pixel-Wise Aleatoric Uncertainty Estimation](https://mlsanity.com/paper/2412.00205/) (2024)
- [From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control](https://mlsanity.com/paper/2502.02205/) (2025)
- [RISE: Radius of Influence based Subgraph Extraction for 3D Molecular Graph Explanation](https://mlsanity.com/paper/2505.02247/) (2025)
- [Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment](https://mlsanity.com/paper/2506.10186/) (2025)
- [UnPose: Uncertainty-Guided Diffusion Priors for Zero-Shot Pose Estimation](https://mlsanity.com/paper/2508.15972/) (2025)
- [Controllable 3D Molecular Generation for Structure-Based Drug Design Through Bayesian Flow Networks and Gradient Integration](https://mlsanity.com/paper/2508.21468/) (2025)
- [GAGA: Gaussianity-Aware Gaussian Approximation for Efficient 3D Molecular Generation](https://mlsanity.com/paper/2507.09043/) (2025)

## Abstract (arXiv)

Bayesian inference provides a principled framework for modeling epistemic uncertainty in neural networks by treating predictions as distributions rather than deterministic values. Meanwhile, diffusion-based models for 3D molecular graph generation operate on fragile geometric structures governed by strict chemical constraints, making inference highly sensitive to uncertainty miscalibration. A largely overlooked issue is that epistemic uncertainty arising from the learned denoiser interacts with the aleatoric uncertainty intentionally injected during reverse diffusion, leading to systematic variance inflation and a mismatch between the true distribution and the simulated distribution. This effect is particularly detrimental for high-precision molecular generation, where even small deviations can violate chemical validity. In this work, we provide a theoretical and empirical analysis of how epistemic uncertainty propagates through diffusion inference and degrades sampling quality. Building on this investigation, we propose UCD (Uncertainty-Calibrated Diffusion), a simple yet effective method that calibrates the reverse diffusion process to account for epistemic uncertainty. Extensive experiments on standard 3D molecular benchmarks demonstrate that UCD consistently improves sampling quality across diverse baseline methods, establishing new state-of-the-art performance for 3D molecular diffusion. The code is available at https://github.com/jiuguaiwf/UCD.