Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation

2026-05-31

AI Summary

This work addresses the challenge in 3D molecular graph generation with diffusion models, where the coupling of epistemic and aleatoric uncertainties leads to variance inflation and mismatched sampling distributions, thereby compromising chemical validity. For the first time, this study elucidates the interaction mechanism between these two uncertainty types during diffusion-based inference and introduces a Bayesian neural network–based approach to model epistemic uncertainty and calibrate the reverse diffusion process. The proposed method features a simple architecture and computational efficiency, achieving state-of-the-art performance across multiple standard 3D molecular generation benchmarks. It significantly outperforms existing baselines, effectively enhancing both geometric accuracy and chemical validity of the generated molecules.

Abstract

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.

Problem

• epistemic uncertainty
• diffusion models
• 3D molecular generation
• uncertainty calibration
• chemical validity

Innovation

• Uncertainty Calibration
• Diffusion Models
• 3D Molecular Generation
• Epistemic Uncertainty
• Molecular Graphs