Uncertainty Estimation for Molecular Diffusion Models

2026-06-11

AI Summary

Existing 3D molecular diffusion models struggle to provide reliable uncertainty estimates for identifying low-quality generated samples. This work introduces Laplace approximation into the molecular diffusion framework for the first time, proposing a post-hoc, training-free method that quantifies uncertainty by modeling the variability of noise predictions from the denoising network along the generation trajectory. By assigning an uncertainty score to each generated sample and integrating a test-time scaling strategy, the resulting uncertainty estimates exhibit a strong negative correlation with established sample quality metrics. This enables effective filtering of low-quality molecules and substantially enhances the model’s generative performance at inference time without requiring any retraining.

Abstract

Diffusion models have seen wide adoption for 3D molecular generation, yet they offer no principled signal of when a generated molecule is likely to be of low quality. We propose a post-hoc method for estimating per-sample uncertainty in pretrained molecular diffusion models. Building on a Laplace approximation of the denoising network, we measure the variability of the noise prediction across the generation trajectory. Empirically, we show that the resulting uncertainty score is informative of sample quality, exhibiting a negative correlation with established sample-level quality metrics. We further study how the proposed uncertainty score can be used to filter generated samples, improving model performance via test-time scaling.

Problem

• Uncertainty Estimation
• Molecular Diffusion Models
• Sample Quality
• 3D Molecular Generation

Innovation

• uncertainty estimation
• molecular diffusion models
• Laplace approximation
• noise prediction variability
• test-time scaling