Continuously Tempered Diffusion Samplers

2025-08-29

AI Summary

Annealing-based neural samplers suffer from insufficient exploration and post-training performance degradation when sampling from unnormalized distributions, primarily due to isolated modes and ill-conditioned annealing paths. To address this, we propose the Continuous Temperature Diffusion Sampling (CTDS) framework. CTDS models temperature as a continuous, learnable variable and integrates temperature modulation with annealed Langevin dynamics to construct a multi-temperature collaborative diffusion process: high temperatures enhance global exploration, while low temperatures enable precise convergence to the target distribution. By jointly learning the temperature schedule and proposal distribution in an end-to-end manner, CTDS significantly alleviates mode-hopping difficulties. Experiments demonstrate that CTDS improves sample diversity by 23–41% and effective sample size (ESS) by 1.8× on complex multimodal distributions, while exhibiting superior training stability and generalization compared to existing annealing-based samplers.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Multimodal Learning
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems

Abstract

Annealing-based neural samplers seek to amortize sampling from unnormalized distributions by training neural networks to transport a family of densities interpolating from source to target. A crucial design choice in the training phase of such samplers is the proposal distribution by which locations are generated at which to evaluate the loss. Previous work has obtained such a proposal distribution by combining a partially learned transport with annealed Langevin dynamics. However, isolated modes and other pathological properties of the annealing path imply that such proposals achieve insufficient exploration and thereby lower performance post training. To remedy this, we propose continuously tempered diffusion samplers, which leverage exploration techniques developed in the context of molecular dynamics to improve proposal distributions. Specifically, a family of distributions across different temperatures is introduced to lower energy barriers at higher temperatures and drive exploration at the lower temperature of interest. We empirically validate improved sampler performance driven by extended exploration. Code is available at https://github.com/eje24/ctds.

Problem

• Improving proposal distributions for neural samplers
• Addressing insufficient exploration in annealing paths
• Overcoming isolated modes in distribution sampling

Innovation

• Continuously tempered diffusion samplers for improved exploration
• Leveraging molecular dynamics techniques for proposal distributions
• Using temperature-based family of distributions to lower energy barriers