Conditional Diffusion Sampling
Abstract
Sampling from unnormalized multimodal distributions with limited density evaluations remains a fundamental challenge in machine learning and natural sciences. Successful approaches construct a bridge between a tractable reference and the target distribution. Parallel Tempering (PT) serves as the gold standard, while recent diffusion-based approaches offer a continuous alternative at the cost of neural training. In this work, we introduce Conditional Diffusion Sampling (CDS), a framework that combines these two paradigms. To this end, we derive Conditional Interpolants, a class of stochastic processes whose transport dynamics are governed by an exact, closed-form stochastic differential equation (SDE), requiring no neural approximation. Although these dynamics require sampling from a non-trivial initialization distribution, we show both theoretically and empirically that the cost of this initialization diminishes for sufficiently short diffusion times. CDS leverages this by a two-stage procedure: (1) PT is used to efficiently sample the initial distribution, and then (2) samples are transported via the transport SDE. This combination couples the robust global exploration of PT with efficient local transport. Experiments suggest that CDS has the potential to achieve a superior trade-off between sample quality and density evaluation cost compared to state-of-the-art samplers.
Lay Summary
Many important problems in science require drawing samples from complex probability distributions: generating molecules, mapping all the ways a protein can fold, or navigating a complicated landscape with many peaks and valleys. This is especially difficult when the distribution has many separated regions of high probability, since existing methods either explore broadly but slowly, or move quickly but get trapped in local regions. Parallel Tempering is a widely used technique that tackles this by building a bridge from a simpler distribution to the target, but it can be computationally expensive. More recent diffusion-based methods offer a smoother way to build this bridge, but typically require approximations that add complexity and/or degrade performance. We introduce Conditional Diffusion Sampling (CDS), a framework that combines the strengths of both. We derive a mathematically exact diffusion process (requiring no approximations) that smoothly transports samples from one distribution to another. CDS first uses Parallel Tempering to efficiently identify promising regions of the distribution, then applies this diffusion process for precise, efficient transport to the target. Our results suggest that CDS achieves better sample quality with fewer expensive evaluations than current state-of-the-art methods, making it a promising tool for challenging problems in general sciences.