1 Introduction↩︎

Diffusion models have rapidly become the dominant paradigm for 3D molecular generation [1]–[4]. Yet even state-of-the-art models can produce samples that are chemically invalid or unstable. This becomes a practical bottleneck when scaling up the generation budgets: deciding which molecules merit expensive downstream evaluation (e.g., docking or wet-lab validation) currently relies on coarse rule-based filters or on the model’s own—often poorly calibrated [5], [6]—likelihood.

For predictive models, uncertainty estimation is the standard tool for flagging unreliable outputs [7], and a growing line of work has begun to adapt these ideas to diffusion models in the image domain [8]–[12], showing promising early results for flagging low-quality generated images. To the best of our knowledge, however, no analogous work exists for molecular diffusion, despite the aforementioned practical need for sample-quality signals in scientific domains where downstream evaluation is expensive.

In this work, we take a first step towards extending the existing diffusion uncertainty approaches to molecular domain. Our method is post-hoc: given a pretrained molecular diffusion model, we fit a Laplace approximation [13] to the denoising network and measure the variability of its noise predictions along the generation trajectory. We aggregate this variability across selected timesteps, atoms, and feature dimensions to obtain a single uncertainty score for each generated molecule. Empirically, we show that this score is informative of molecular sample quality: it is negatively correlated with established metrics such as molecular stability, atom stability, and validity, and it is more predictive than diffusion likelihood baseline in our QM9 [14] experiments. We further show that the score can be used for test-time scaling [15], [16] by oversampling molecules and filtering out high-uncertainty generations, improving sample quality without retraining the underlying generator.

2 Background↩︎

2.0.0.1 Diffusion Models

Diffusion models [17]–[19] are a class of generative models that learn to invert a gradual noising process. Given a data sample \(\boldsymbol{x}_0 \sim q(\mathbf{x})\), the forward process produces progressively noisier samples via \[\begin{align} \label{eq:forward} \boldsymbol{x}_t = \sqrt{\bar{\alpha}_t}\, \boldsymbol{x}_0 + \sqrt{1 - \bar{\alpha}_t}\, \boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon}\sim \mathcal{N}(0, I), \end{align}\tag{1}\] where \(\bar{\alpha}_t = \prod_{r=1}^t (1 - \beta_r)\) and \(\{\beta_r\}_{r=1}^T\) is a variance schedule with \(\beta_r \in (0,1)\), chosen such that \(\boldsymbol{x}_T\) is approximately standard Gaussian. For any \(s < t\), the reverse-time posterior is also Gaussian distributed \[\begin{align} q(\boldsymbol{x}_{s} \mid \boldsymbol{x}_t, \boldsymbol{x}_0) = \mathcal{N}\!\left(\boldsymbol{x}_{s}; \boldsymbol{\mu}_{s\rightarrow t}(\boldsymbol{x}_t, \boldsymbol{x}_0),\, \beta_{s \rightarrow t} I\right), \end{align}\] where \(\boldsymbol{\mu}_{s\rightarrow t}\) and \(\beta_{s \rightarrow t}\) can be obtained analytically [20]. Since \(\boldsymbol{x}_0\) is unknown at generation time, the learned reverse process \(p_\theta(\boldsymbol{x}_{s} \mid \boldsymbol{x}_t)\) is trained to match this posterior. In practice, the mean is often reparametrized in terms of the noise (\(\boldsymbol{\epsilon}\)) prediction which yields the simplified regression objective [20]¹ \[\begin{align} \label{eq:diff-loss} \mathcal{L}(\theta) = \mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{\epsilon},t} \left[ \|\boldsymbol{\epsilon}- \hat{\boldsymbol{\epsilon}}_t\|_2^2 \right], \end{align}\tag{2}\] with \(\hat{\boldsymbol{\epsilon}}_t := f_\theta(\boldsymbol{x}_t, t)\). After training, starting from \(\hat{\boldsymbol{x}}_T \sim \mathcal{N}(0, I)\) a new sample \(\hat{\boldsymbol{x}}\) is generated by iteratively predicting noise \(\hat{\boldsymbol{\epsilon}}_t\) and sampling a partially denoised \(\hat{\boldsymbol{x}}_{t-1} \sim p_{\hat{\theta}}(\cdot \mid \hat{\boldsymbol{x}}_t, \hat{\boldsymbol{\epsilon}}_t)\) along a decreasing sequence of timesteps \(t=T, T-1,\ldots, 1\), according to the transition rules of a chosen sampler, such as DDPM [20] or DDIM [21].

Beyond sampling, diffusion models also admit tractable likelihood evaluation. In the discrete-time setting we consider, an estimate of \(\log p_{\hat{\theta}}(\boldsymbol{x})\) for any sample \(\boldsymbol{x}\) is obtained by evaluating a variational lower bound (ELBO) [20], [22]; in the continuous-time formulation, the analogous quantity follows from solving the probability flow ODE [18].

3 Methods↩︎

We next describe our approach for estimating the uncertainty of molecules generated by diffusion models. The goal is to produce a single uncertainty score \(u \in \mathbb{R}_{\ge 0}\) that correlates with the quality of a generated sample \(\hat{\boldsymbol{x}} \in \mathbb{R}^{N_a \times D}\), where \(N_a\) is the number of atoms and \(D\) the number of per-atom features (including coordinates). At a high level, we measure the variability of the denoiser’s prediction across sampling steps, based on intuition that samples the model is internally uncertain about should exhibit more variable denoising trajectories, and aggregate this variability over time, atoms, and features into a single score (see 1).

3.0.0.1 Laplace Posterior

As a first step, we fit an approximate posterior over model parameters, \(q(\theta) \approx p(\theta \mid \mathcal{D})\). To remain post-hoc compatible with any pretrained molecular diffusion model, we follow prior work on uncertainty estimation in diffusion models [8], [11] and use the Laplace approximation [23], \[\begin{align} \label{eq:la} q(\theta) = \mathcal{N}(\theta \mid \hat{\theta}, \Sigma), \quad \Sigma = \big(\nabla^2_{\theta} \mathcal{L}(\theta; \mathcal{D}) \big|_{\hat{\theta}} + \lambda I\big)^{-1}, \end{align}\tag{3}\] where \(\hat{\theta}\) are the parameters of the pretrained diffusion model, \(\mathcal{L}\) is the diffusion training loss from Eq. 2 , and \(\lambda\) is the prior precision. To scale to models with many parameters, we follow [13] and apply probabilistic treatment only to the final layers of the network as well as use the empirical Fisher as a Hessian approximation.

3.0.0.2 Noise Prediction Variability

Next, we use the posterior \(q(\theta)\) to measure the variability of the denoiser \(f_{\theta}\) along a sampling trajectory \(\{\hat{\boldsymbol{x}}_t\}_{t=T,\ldots,0}\). Concretely, for a subset of sampling steps \(\mathcal{T} \subseteq \{0, \ldots, T\}\), we draw \(M\) weight samples from the posterior and compute the corresponding noise predictions, \[\begin{align} \boldsymbol{\epsilon}_t^m = f_{\theta_m}(\hat{\boldsymbol{x}}_t, t), \quad \theta_m \sim q(\theta), \end{align}\] and define the intermediate uncertainty \(\boldsymbol{u}_t \in \mathbb{R}_{\ge 0}^{N_a \times D}\) as their elementwise sample variance, \[\begin{align} \boldsymbol{u}_t = \frac{1}{M}\sum_{m=1}^M (\boldsymbol{\epsilon}_t^m - \bar{\boldsymbol{\epsilon}}_t)^2, \quad \bar{\boldsymbol{\epsilon}}_t = \frac{1}{M}\sum_{m=1}^M \boldsymbol{\epsilon}_t^m. \end{align}\] Empirically, we find that a small subset of sampling steps (\(|\mathcal{T}| \ll T\)) is sufficient for good performance, which limits the computational overhead of uncertainty estimation. Finally, we aggregate the per-step uncertainties \(\{\boldsymbol{u}_t\}_{t \in \mathcal{T}}\) into a single score \(u (\hat{\boldsymbol{x}})\) for the generated sample \(\hat{\boldsymbol{x}}\) by averaging across time steps, atoms, and feature dimensions: \[\begin{align} \label{eq:unc95score} u(\hat{\boldsymbol{x}}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [\boldsymbol{u}_t]_{n,d}. \end{align}\tag{4}\]

4 Experiments↩︎

In our experiments, we first establish that our proposed uncertainty score (Equation 4 ) is informative of sample quality by reporting its correlation with standard sample-level metrics used in molecular generation (4.1). Next, we explore how uncertainty can be used in a test-time compute setting [15], where we oversample generated molecules and use uncertainty as a filtering criterion (4.2). Finally, we report a series of ablations to support our proposed design choices (4.3).

4.0.0.1 Experimental details

We evaluate our approach on two open-source diffusion models for 3D molecular generation: i) EDM [1] is an E(3)-equivariant diffusion model that jointly diffuses atomic coordinates and features in data space, while ii) GeoLDM [2] performs diffusion in a learned equivariant latent space. For both models, we use the official pretrained checkpoints, as our uncertainty estimation is post-hoc and does not require retraining. We use two standard datasets: QM9 [14], containing small organic molecules with up to 9 heavy atoms, and GEOM-Drugs [24], a larger collection of drug-like molecules with substantially more atoms and conformational complexity. For the Laplace approximation, we use default hyperparameters (prior precision \(\lambda=1\), likelihood noise \(\sigma=1\)). We measure uncertainty only in the final \(10\%\) of generation process, i.e. \(t<100\) when using \(T=1000\), as we found that including earlier timesteps yields no or only marginal performance improvements (4.3).

4.1 Uncertainty correlates with sample quality↩︎

To assess whether our proposed uncertainty is predictive of sample quality, we report its correlation with a set of well-established metrics for molecular generation in 1 on QM9 dataset. Across both models, our uncertainty score exhibits a statistically-significant negative correlation, indicating that the measured variability is picking up signal informative of sample quality. Moreover, uncertainty exhibits a stronger negative correlation compared to using diffusion model log-likelihood \(\log p_{\hat{\theta}}(\hat{\boldsymbol{x}})\) (NLL). For example, using the EDM model uncertainty achieves \(-0.284\) correlation with molecular stability, while likelihood gets \(-0.15\).

4.2 Test-time scaling via uncertainty↩︎

We next examine whether the observed correlation between our uncertainty score and molecular sample quality can be exploited in practice. Following a test-time scaling setting [15], we improve the quality of a fixed-size set of generated molecules via oversampling and filtering: we generate \(N\) samples, with \(N\) varying from \(10\)K to \(20\)K, and retain the \(10\)K with the lowest estimated uncertainty (or, as a baseline, the lowest diffusion NLL). The \(N{=}10\)K case corresponds to no filtering.

Focusing first on QM9 (2), we find that uncertainty-based filtering meaningfully improves molecular stability, atomic stability, and validity across both EDM and GeoLDM. For EDM, oversampling to \(20\)K and filtering down to \(10\)K yields gains of \({\sim}10\%\) in molecular stability, \({\sim}1\%\) in atomic stability, and \({\sim}5\%\) in validity. Notably, these test-time gains are comparable in magnitude to the improvement from switching the underlying model from EDM to GeoLDM at the \(10\)K generation budget, highlighting test-time compute as a complementary axis for improving molecular generation. The gains come with a modest tradeoff in diversity: uniqueness drops by roughly \(1\%\). Finally, the stronger negative correlation of uncertainty relative to NLL translates into larger downstream gains: uncertainty-based filtering outperforms NLL-based filtering across both models and all three quality metrics (molecular stability, atomic stability, validity) on QM9.

On GEOM-Drugs (3), by contrast, test-time filtering fails to meaningfully improve any of the three sample quality metrics, both for our uncertainty score and for NLL. We leave an investigation of why uncertainty-based filtering does not transfer to this larger, more complex dataset to future work.

4.3 Ablations↩︎

We begin our ablations by isolating the contribution of the Laplace approximation (Equation 3 ). We compare the quality signal from our uncertainty score against a simpler alternative that omits the Fisher approximation of the Hessian and instead samples from an isotropic Gaussian centered at the MAP [25]: \[\begin{align} \label{eq:laplace95ablate} q_{\lambda}(\theta) = \mathcal{N}(\theta \mid \hat{\theta}, \lambda^{-1} I) \: . \end{align}\tag{5}\] As shown in 2, the simpler, Fisher-free variant performs almost on par with the original Laplace approximation in terms of yielding uncertainties that correlate (negatively) with the quality metrics like molecular stability. This suggests that the signal driving our score is closer to a sensitivity measure — the variability of model predictions under small isotropic perturbations of the MAP — than to ‘true’ epistemic uncertainty in a strict Bayesian sense.

Finally, we ablate the choice of timesteps \(\mathcal{T}\) at which we measure uncertainty along the generation trajectory. As shown in 4, the timesteps closest to the clean sample (small \(t\)) correlate most negatively with sample quality (atomic stability), supporting our choice to concentrate the uncertainty measurement at the end of the generation process.

5 Conclusion↩︎

We have introduced a post-hoc uncertainty estimation approach for molecular diffusion models based on noise-prediction variability along the sampling trajectory, and empirically validated its effectiveness on QM9. A central limitation is that the approach does not transfer to the larger, more complex GEOM-Drugs dataset, which we leave to future work. One promising direction is to compute uncertainty in a ‘semantic’ feature space—i.e., based on a pretrained molecular encoder—analogous to the approach shown to be effective for uncertainty estimation in high-dimensional image generation [11]. Lastly, our uncertainty score is a natural drop-in replacement for the NLL-based verifier in recently proposed adaptive test-time scaling methods [16], which represents another promising future direction.

References↩︎