Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation
2026-06-01 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how uncertainty affects the generation of 3D molecular structures using diffusion models. They found that two types of uncertainty—one from the model’s learning process and one intentionally added during sampling—combine in a way that causes errors in the final molecules. These errors make it hard to produce chemically valid molecules, especially when high accuracy is needed. To solve this, the authors propose a new method called UCD that adjusts the sampling process to better handle uncertainty, resulting in improved molecular generation performance.
Bayesian inferenceepistemic uncertaintyaleatoric uncertaintydiffusion models3D molecular generationreverse diffusiondenoiseruncertainty calibrationsampling qualitychemical validity
Authors
Fang Wan, Jingxiang Qu, Yi Liu
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.