On the uncertainty in atomistic diffusion models

Diffusion models are state-of-the-art for generating atomistic structures, yet they still produce physically invalid samples. While uncertainty quantification (UQ) could mitigate this, traditional methods remain expensive due to the sequential nature of diffusion models. In this work, we probe the efficacy of a lightweight UQ-head, akin to heteroskedastic regression. We systematically leverage this UQ-head across three granularities: post-hoc filtering and two UQ-conditioned sampling interventions (molecule correction and atom forcing). We also show that chemical composition strongly confounds uncertainty, and therefore propose a normalised UQ-score to mitigate this bias. After empirically establishing that our scores reliably predict chemical validity, we evaluate how each mechanism, using both standard and normalised uncertainty, impacts sample quality, providing critical insights into the practical trade-offs of UQ-guided atomistic structure generation.