Uncertainty Estimation for Molecular Diffusion Models

Technical Summary: Uncertainty Estimation for Molecular Diffusion Models

Problem Statement
Diffusion models have become the dominant paradigm for 3D molecular generation. However, a critical practical bottleneck exists: these models lack a principled signal to identify when a generated molecule is chemically invalid or unstable. While predictive models typically use uncertainty estimation to flag unreliable outputs, existing methods for diffusion models have primarily focused on the image domain. In the molecular domain, where downstream evaluation (e.g., docking or wet-lab validation) is expensive, relying on coarse rule-based filters or the model's own likelihood—which is often poorly calibrated—remains the standard. There is a lack of analogous work providing sample-quality signals for molecular diffusion models.

Methodology
The authors propose a post-hoc method to estimate per-sample uncertainty for pretrained molecular diffusion models without retraining the generator. The approach is built on the intuition that samples the model is internally uncertain about will exhibit more variable denoising trajectories.

1. Laplace Approximation: To remain compatible with any pretrained model, the method fits a Laplace approximation to the denoising network's parameters. Specifically, a posterior distribution $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ is constructed around the pretrained parameters $\hat{\theta}$. The covariance $\Sigma$ is approximated using the empirical Fisher information matrix (applied only to the final layers to ensure scalability) plus a prior precision term.
2. Noise Prediction Variability: During the generation trajectory, for a selected subset of timesteps $\mathcal{T}$, the method draws $M$ weight samples $\theta_m$ from the posterior $q(\theta)$. For each sample, the denoiser predicts noise $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$.
3. Uncertainty Score Aggregation: The intermediate uncertainty $u_t$ at each timestep is defined as the element-wise sample variance of these noise predictions across the $M$ samples. This per-step uncertainty is then aggregated across the selected timesteps, all atoms ($N_a$), and all feature dimensions ($D$) to produce a single scalar uncertainty score $u(\hat{x})$ for the generated molecule:
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   Empirically, the authors find that measuring uncertainty only in the final 10% of the generation process (where $t$ is small) is sufficient for performance, limiting computational overhead.

Key Contributions

• First Application to Molecular Diffusion: This work extends diffusion uncertainty approaches from the image domain to molecular generation, addressing the specific need for sample-quality signals in scientific domains with expensive downstream evaluation.
• Post-hoc Framework: The method requires no retraining of the underlying diffusion model, making it applicable to existing state-of-the-art architectures like EDM and GeoLDM.
• Test-Time Scaling: The paper demonstrates how the uncertainty score can be used as a filtering criterion in a test-time scaling setting (oversampling and filtering) to improve the quality of generated molecules.

Results
The authors evaluated their approach on the QM9 and GEOM-Drugs datasets using EDM and GeoLDM models.

• Correlation with Quality: On the QM9 dataset, the proposed uncertainty score exhibits a statistically significant negative correlation with established quality metrics (molecular stability, atom stability, and validity). This indicates that higher uncertainty scores correspond to lower-quality samples.
• Comparison to Likelihood: The uncertainty score consistently outperforms the diffusion model's negative log-likelihood (NLL) as a predictor of sample quality. For example, on EDM/QM9, the uncertainty score achieved a Spearman correlation of $-0.284$ with molecular stability, compared to $-0.150$ for NLL.
• Test-Time Scaling Performance: By generating 20K molecules and filtering to keep the 10K with the lowest uncertainty, the authors observed significant improvements in molecular stability ($\sim10\%$), atomic stability ($\sim1\%$), and validity ($\sim5\%$) for the EDM model. These gains were comparable to the improvement achieved by switching the underlying model architecture from EDM to GeoLDM.
• Ablation Studies:
  • Fisher Approximation: Removing the Fisher approximation (replacing the covariance with an isotropic Gaussian) yielded nearly identical performance, suggesting the score acts more as a sensitivity measure to small perturbations than a strict Bayesian epistemic uncertainty.
  • Timestep Selection: The signal is strongest at the "clean" end of the generation trajectory (small $t$), validating the choice to focus uncertainty measurement on the final steps.
• Limitations: The method did not yield meaningful improvements on the larger, more complex GEOM-Drugs dataset, indicating that the approach does not currently transfer to datasets with higher conformational complexity.

Significance and Claims
The paper claims to take a "first step" towards extending diffusion uncertainty approaches to the molecular domain. The primary significance lies in providing a practical, post-hoc tool for filtering generated molecules, thereby improving the efficiency of downstream evaluation pipelines. The authors position their uncertainty score as a natural replacement for NLL-based verifiers in adaptive test-time scaling methods.

The authors remain modest regarding their claims, explicitly noting that the approach currently fails on larger, more complex datasets (GEOM-Drugs) and that the underlying mechanism appears to be a sensitivity measure rather than true Bayesian epistemic uncertainty. They suggest future work should explore computing uncertainty in a "semantic" feature space using pretrained molecular encoders, analogous to recent successes in high-dimensional image generation.