校准扩散模型中的不确定性,提升3D分子生成精度与可靠性
Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation

- 引入不确定性校准机制,修正反向扩散过程中的认知误差
- 在多个标准数据集上显著提升生成质量,达新SOTA水平
- 适合需要高精度分子结构生成的药物研发与材料设计场景
贝叶斯推断通过将预测视为分布而非确定值,为神经网络中的认知不确定性建模提供了严谨框架。然而,基于扩散的3D分子图生成依赖于受严格化学约束的脆弱几何结构,推理对不确定性校准敏感。一个被忽视的问题是:学习得到的去噪器产生的认知不确定性,与反向扩散中故意注入的随机性不确定性相互作用,导致方差系统性膨胀,使真实分布与模拟分布不匹配。这一现象对高精度分子生成尤为有害,微小偏差即可能破坏化学有效性。本文从理论和实证角度分析了认知不确定性在扩散推理中的传播机制及其对采样质量的负面影响。基于此,提出UCD(Uncertainty-Calibrated Diffusion)方法,简单而有效,能校准反向扩散过程以考虑认知不确定性。在标准3D分子基准测试上的大量实验表明,UCD在多种基线方法上持续提升采样质量,建立3D分子扩散生成的新SOTA。代码已开源:https://github.com/jiuguaiwf/UCD。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。