Temperature-Regulated Stochastic Sampling for Diffusion-Based High-Quality Molecular Generation
Abstract
Despite the strong performance of diffusion models in molecular generation, the sampling process itself remains largely understudied, with most existing methods adopt off-the-shelf sampling strategies developed for natural images. Yet unlike natural images, molecular distributions are governed by physical laws and are sharply concentrated, making them particularly challenging for standard samplers, which frequently produce invalid or unphysical structures. In this work, we revisit diffusion sampling from a unified stochastic differential equation (SDE) perspective and introduce a general framework parameterized by two interpretable controls: \emph{stochasticity} and \emph{temperature}. Our theoretical analysis reveals that stochasticity accelerates the decay of sampling error, while temperature directly controls distribution sharpness, enabling concentration on physically plausible configurations. Building on these insights, we develop \textsc{TReaSSure} (Temperature-Regulated Stochastic Sampling), a lightweight, training-free strategy that combines stochasticity scheduling with temperature annealing to better capture the sharply concentrated distributions of molecular data. Extensive experiments on small molecule generation, protein structure prediction, and protein design demonstrate that \textsc{TReaSSure} consistently improves generation quality and produces more physically valid structures. Empirical analyses further corroborate our theoretical findings. The method harnesses pretrained models without retraining, underscoring its generality and practical value across diverse molecular domains.