Progressive Tempering Sampler with Diffusion

TL;DR

This work introduces Progressive Tempering Sampler with Diffusion (PTSD), a framework that unifies parallel tempering with diffusion-based neural samplers by learning diffusion models across temperatures and guiding high-temperature samples to approximate lower-temperature targets. PTSD leverages temperature guidance, buffers, truncated importance sampling, and local parallel-tempering refinements to transfer information across temperatures and produce uncorrelated, well-mixed samples at the lowest temperature with far fewer target evaluations. Key contributions include a concrete extrapolation mechanism for scores between temperatures, a sequential training pipeline across a temperature ladder, and refinement techniques to mitigate approximation bias. Empirically, PTSD achieves state-of-the-art target-evaluation efficiency across multiple benchmarks and scales to molecular settings like alanine dipeptide, demonstrating practical impact for efficient unnormalized-density sampling in complex, multimodal landscapes.

Abstract

Recent research has focused on designing neural samplers that amortize the process of sampling from unnormalized densities. However, despite significant advancements, they still fall short of the state-of-the-art MCMC approach, Parallel Tempering (PT), when it comes to the efficiency of target evaluations. On the other hand, unlike a well-trained neural sampler, PT yields only dependent samples and needs to be rerun -- at considerable computational cost -- whenever new samples are required. To address these weaknesses, we propose the Progressive Tempering Sampler with Diffusion (PTSD), which trains diffusion models sequentially across temperatures, leveraging the advantages of PT to improve the training of neural samplers. We also introduce a novel method to combine high-temperature diffusion models to generate approximate lower-temperature samples, which are minimally refined using MCMC and used to train the next diffusion model. PTSD enables efficient reuse of sample information across temperature levels while generating well-mixed, uncorrelated samples. Our method significantly improves target evaluation efficiency, outperforming diffusion-based neural samplers.

Figures (13)

• Figure 1: Sample error ($\mathcal{W}_2$ distance) and target evaluation times for several diffusion (and control)-based neural samplers on Many-Well-32 target, including DDS, iDEM, BNEM, and CMCD, with our proposed approach. We include the results obtained by first running PT and fit a diffusion model post hoc for comparison.
• Figure 2: Illustration of parallel tempering with three temperatures.
• Figure 3: The training process of PTSD. We unroll the training process into a sequence of target temperatures. We first initialize buffers and models at the highest two temperatures and generate samples from lower temperatures sequentially using the temperature guidance to extrapolate to lower temperatures.
• Figure 4: "Self-bootstrapping" training of neural samplers.
• Figure 5: First two panels: Marginal density of the interatomic distance on buffers at two temperatures, along with the marginal density of a diffusion model fit to the buffer. Right: The temperature-guided extrapolation based on the higher-temperature models. AD: auto-diff extrapolation; ME: model extrapolation; RS: score rescaling heuristic; GT: ground truth; PTSD: our temperature guidance.