Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation

技术摘要：用于可靠 3D 分子图生成的确定性校准扩散模型

问题陈述
扩散模型（DMs）已成为 3D 分子图生成领域的主导范式，通过迭代去噪来重建原子类型和 3D 坐标。然而，这些模型运行在受严格化学约束控制的脆弱几何结构之上，使得推理过程对不确定性失校准（uncertainty miscalibration）高度敏感。

本文识别了一个此前被忽视的失效模式：认知不确定性（epistemic uncertainty，源于由于数据有限或随机训练导致的学习器建模不完善）与偶然不确定性（aleatoric uncertainty，即在反向扩散过程中注入的有意随机噪声）之间的相互作用。在标准的扩散推理中，去噪器被视为确定性的。作者认为，去噪器预测中固有的认知不确定性实际上充当了额外的噪声源。这种相互作用导致反向转移核（reverse transition kernel）出现系统性的方差膨胀。随着推理过程从 $t=T$ 进行到 $t=0$，这种膨胀的方差不断累积，导致模拟分布偏离真实的训练数据分布。对于 3D 分子生成而言，即使是微小的偏差也会违反化学有效性（例如，错误的键长或化合价），从而导致不稳定或无效的分子结构。

方法论：不确定性校准扩散（UCD）
作者提出了 UCD（Uncertainty-Calibrated Diffusion，不确定性校准扩散），这是一种推理时校准策略，旨在通过校准反向扩散过程来应对认知不确定性，而无需修改底层的模型架构或训练目标。

1. 理论基础：

   • 作者对反向更新步骤进行了形式化处理。他们证明了如果去噪器预测 $\epsilon_\theta$ 包含被建模为加性高斯噪声 $\delta_t \sim \mathcal{N}(0, \tau_t^2 I)$ 的认知不确定性，则反向转移的有效方差将从标称的 $\sigma_t^2$ 增加到 $\tilde{\sigma}_t^2 = \sigma_t^2 + \eta_t^2$，其中 $\eta_t$ 是不确定性量级与扩散调度（diffusion schedule）的函数。
   • 他们提供了理论分析，表明这种方差膨胀会导致不确定性膨胀采样器轨迹与理想轨迹之间的 Kullback-Leibler (KL) 散度发生累积增加。此外，他们推导出的界限表明，这种局部方差膨胀会传播到终点，导致最终生成的样本产生可测量的偏差。

2. 算法实现：

   • 不确定性估计： UCD 采用了一种轻量级的贝叶斯近似，即使用 蒙特卡洛 Dropout（Monte Carlo dropout）。在推理期间，通过启用 Dropout 运行 $M$ 次去噪器，以估计平均预测值和认知方差 ($\tau_t^2$)。
   • 方差校正： UCD 不再使用用于偶然成分的固定噪声调度，而是动态调整注入噪声的规模。具体而言，它通过减少偶然噪声的方差来补偿估计出的认知方差，从而确保总体的有效方差与预期的理论调度相匹配。
   • 通用采样器： 该方法被实现为一个与各种扩散骨干网络（包括使用随机 SDE 或确定性 ODE 的网络）兼容的插件模块。它通过从噪声注入项中减去估计出的不确定性诱导方差，来修改反向更新方程。

核心贡献

1. 识别方差膨胀： 本文识别并形式化了认知不确定性与偶然不确定性之间的相互作用是导致分子扩散中方差膨胀的主要原因，解释了为什么标准推理会在长轨迹中降低样本质量。
2. 理论表征： 作者提供了严谨的理论分析，表征了认知不确定性如何通过反向链进行传播，通过路径空间 KL 散度量化了轨迹级偏差，并建立了终点偏差的界限。
3. UCD 框架： 他们提出了 UCD，一种简单、与模型无关的推理时校准方法。它不需要重新训练，不需要架构更改，也不需要额外的可学习参数，因此适用于多种现有的扩散骨干网络。

实验结果
作者在标准的 3D 分子基准测试集 QM9（小有机分子）和 GEOM-Drugs（较大的类药分子）上评估了 UCD。他们在多个最先进的基线模型上测试了该方法，包括 EDM、GeoLDM 和 RADM。

• 无条件生成： UCD 在所有指标（原子稳定性、分子稳定性、有效性和有效性 $\times$ 唯一性）上均一致地提升了性能。例如，在 QM9 上，UCD 将 EDM 的分子稳定性从 82.0% 提升至 87.2%，将有效性从 91.9% 提升至 94.3%。在 GeoLDM 和 RADM 上也观察到了类似的增益。
• 有条件生成： 在 QM9 上的属性控制生成任务（预测如 $\alpha$、$\Delta\epsilon$ 等属性）中，结合 UCD 的 RADM 相比基线实现了更低的平均绝对误差 (MAE)，表明其更好地遵循了目标属性约束。
• 效率： 该方法引入的开销极小。由于 Dropout 是标准操作，训练成本几乎可以忽略不计。推理成本随 Dropout 次数 ($M$) 线性扩展，但对于较小的 $M$（例如 5–10），虽然壁钟时间（wall-clock time）有所增加，但幅度适中，且能带来显著的质量提升。

意义与主张
本文声称，通过弥合理论扩散模型与其实际部署之间的差距，UCD 建立了 3D 分子扩散领域新的最先进性能。其重要性在于解决了一个根本性的可靠性问题：本文指出，在分子生成等高精度任务中，标准推理中对去噪器的“确定性”假设是存在缺陷的。通过显式地校准噪声注入以抵消模型不确定性，UCD 稳定了反向轨迹，确保生成的分子保持化学有效性和几何准确性。作者强调，这种方法是对现有架构创新的补充而非替代，为构建更可靠的生成式模型提供了一条原则性的路径。