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.

• The paper introduces CDS which leverages Conditional Interpolants and closed-form SDEs to decompose complex sampling into PT-based initialization and precise SDE transport.
• It demonstrates significant improvements in sample quality and computational efficiency over existing MCMC methods across varied multimodal distributions and high-dimensional benchmarks.
• Empirical results show CDS reliably recovers all target modes and outperforms baselines, especially in scenarios with expensive density evaluations.

Conditional Diffusion Sampling: A Training-Free Framework for Efficient Sampling from Unnormalized Multimodal Distributions

Introduction and Motivation

Sampling from unnormalized, multimodal target distributions using a minimal number of density evaluations is a central challenge in computational statistics, machine learning, and the molecular sciences. Traditional approaches such as Parallel Tempering (PT) provide robust global exploration by defining annealing paths bridging a tractable reference distribution and the complex target; however, their efficacy deteriorates when reference and target distributions have poor overlap, and their convergence is limited by schedule design and computational expense. Diffusion-based approaches, notably those utilizing stochastic interpolants and score-based models, yield continuous probabilistic bridges but at the cost of training expensive neural transport models, often rendering them impractical within the resource constraints of Markov Chain Monte Carlo (MCMC) settings.

The Conditional Diffusion Sampling (CDS) framework advances the state of the art by rigorously unifying and extending these paradigms. CDS introduces Conditional Interpolants—conditional stochastic processes with transport dynamics governed by exact, closed-form SDEs, eschewing neural or amortized approximation. CDS effectively decomposes the sampling problem into tractable subproblems: global initialization via PT targeting a nearly reference-like conditional distribution, and refinement via direct SDE-based conditional transport. This construction allows significant reduction of the exploration-exploitation trade-off inherent to previous methods and achieves substantially improved sample quality per density computation.

Figure 1: Overview of CDS—Stage 1 employs Parallel Tempering to sample from an initialization conditional distribution close to the reference; Stage 2 integrates the closed-form SDE to transport samples to the target.

Theoretical Foundations: Conditional Interpolants and Exact Dynamics

The crux of CDS is the derivation and utilization of Conditional Interpolants. Given a reference distribution (e.g., a standard Gaussian) and an unnormalized target $\nu$, a stochastic interpolant $F_t(z, x)$ defines a family of distributions interpolating between the reference at $t=0$ to the target at $t=1$ via a deterministic map parametrized by a reference anchor $z$. Conditioning on $z$ induces a conditional path $\nu_{t|z}$, whose explicit density follows from a change of variables.

In typical flow matching or diffusion models, the transport vector field is intractable and learned via neural surrogates; in CDS, for a large class of interpolants—most notably the linear interpolant $F_t(z, x) = (1-t)z + t x$—the SDE drift and associated score $\nabla \log \pi_{t|z}$ are computable in closed form without approximations or learning. The associated SDE,

$$dx_t = \left( u_{t|z}(x_t) + \frac{\sigma_t^2}{2}\nabla \log \pi_{t|z}(x_t) \right)dt + \sigma_t dW_t,$$

is well-posed for $F_t(z, x)$0, with initialization at a strictly positive $F_t(z, x)$1 required to avoid singularities arising from the degeneracy of interpolant inverses at $F_t(z, x)$2. Importantly, as the initialization time $F_t(z, x)$3, the conditional distribution $F_t(z, x)$4 becomes highly concentrated around $F_t(z, x)$5, thus sampling it efficiently becomes possible.

Figure 2: Density evolution for the linear interpolant—at early $F_t(z, x)$6, the conditional distribution concentrates around the anchor, facilitating efficient PT-based initialization.

Two-Stage Procedure: Initialization and SDE-Based Transport

CDS operates via a two-phase mechanism:

• Stage 1: Conditional Sampling. Given $F_t(z, x)$7 reference, draw $F_t(z, x)$8. As $F_t(z, x)$9 decreases, overlap between reference and conditional distribution increases, enhancing the efficiency and round-trip rates of PT. CDS exploits this by running PT not on the entire target but on $t=0$0 for small $t=0$1, reducing the global communication barrier and rendering the multimodal sampling problem more tractable.
• Stage 2: Conditional SDE Integration. Starting from $t=0$2, integrate the exact SDE (Euler–Maruyama discretization with optional correctors) to transport samples up to $t=0$3, targeting the original distribution. The closed-form drift and score enable unbiased, training-free refinement, which corrects imperfect initialization even in nonconvex, multimodal scenarios.

Notably, SDE-based continuous refinement empirically outperforms direct inverse mapping (i.e., instantaneously mapping $t=0$4), due to the ability of the diffusion to progressively correct for exploration errors accrued in Stage 1.

Figure 3: Comparison between SDE-based transport and direct inverse mapping—continuous SDE-based correction yields higher-fidelity samples and improved accuracy, especially in complex targets.

Empirical Evaluation: Benchmarking and Sample Efficiency

CDS is evaluated across eight targets spanning Gaussian mixtures (GM), Lennard-Jones (LJ) clusters, Alanine Dipeptide (ALDP), and Bayesian Neural Network (BNN) posteriors, encompassing low to high-dimensional, highly multimodal, and physically relevant scenarios.

Performance is assessed by Pareto fronts over sample quality metrics (Wasserstein distance, KL divergence, test NLL, etc.) versus the number of density evaluations. Results demonstrate that CDS attains or exceeds state-of-the-art results, achieving higher mean hypervolume ratios (HVR) than NRPT, OASMC, DiGS, HMC, and MALA across all domains. This advantage is especially pronounced in large-scale and computationally expensive regimes (e.g., ALDP and BNN), where density evaluations dominate cost.

Figure 4: Pareto fronts for all tasks—CDS demonstrates a superior balance between convergence/coverage and computational budget, evidenced by higher Mean HVR with fewer density evaluations.

Qualitative analysis confirms that CDS robustly recovers all modes (e.g., all clusters in GM, both conformational states in ALDP), and matches ground-truth target statistics (e.g., energy distributions in LJ) better than baselines. Local samplers (MALA, HMC, NUTS, MAMS) display rapid collapse in the presence of multimodality; DiGS and OASMC are competitive for low dimensions but do not scale efficiently.

Figure 5: CDS recovers target distribution energy histograms with high fidelity in molecular benchmarks, surpassing alternatives, especially in large or structured state spaces.

Ablations and Design Implications

Extensive ablation studies clarify the sensitivity and optimality of key CDS hyperparameters:

• Stage 1 Sampler: PT invariably outperforms local samplers for initialization from $t=0$5. Alternative annealing-based methods (e.g., OASMC) are competitive but less robust.
• Initialization Time $t=0$6: There is a non-trivial trade-off; excessively small $t=0$7 yields overly peaked conditionals and reduces replica overlap in PT, while large $t=0$8 devolves the method into conventional PT targeting the full $t=0$9. Optimal $t=1$0 is task-dependent, with communication efficiency maximized at intermediate values (supported by GCB/RT analysis).
• Step Allocation: Performance improves if more computational budget is expended in Stage 1 for exploratory targets; however, in high-dimensional or stiff tasks, sufficient SDE integration steps in Stage 2 become critical.
• Corrector and Noise Scheduling: Corrector steps benefit low noise regimes but are less influential as diffusion variance increases.

  Figure 6: Global Communication Barrier (GCB) as a function of $t=1$1—reduced $t=1$2 enhances PT communication until degeneracy, indicating a nonmonotonic optimum.

Implications, Limitations, and Outlook

CDS demonstrates that training-free, conditional SDE-based sampling offers a compelling regime for high-fidelity sampling when target density evaluations are expensive but sample quality is paramount. The absence of amortization and neural training distinguishes CDS from recent diffusion-based samplers, placing it firmly in the high-cost-per-evaluation, high-accuracy regime characteristic of scientific and Bayesian inference applications.

From a theoretical perspective, the closed-form nature of CDS's interpolant and drift construction sidesteps the optimization and convergence limitations of learned generative models. The approach also suggests a flexible interface: the conditional interpolant may, in principle, be replaced by non-linear or geometry-aware maps, potentially further mitigating pathologies (e.g., in high-energy regimes of molecular systems).

On the other hand, CDS inherits limitations from the choice of interpolant; the use of a linear interpolant can induce numerical instabilities for distributions with strong singularities (notably in the LJ and ALDP tasks), and improper selection of $t=1$3 can result in suboptimal communication or degeneracy. Addressing these is a promising direction: geometric or Riemannian interpolants or adaptive $t=1$4 selection could offer further gains.

Future improvements may include:

• Designing flexible, task-adapted interpolants exploiting target geometry
• Automated $t=1$5 selection strategies, e.g., via online minimization of the Symmetric KL divergence or learned schedule
• Extensions to constrained state spaces or manifold-valued data
• Integration with amortized proposals for hybrid regimes

Conclusion

CDS provides an exact, training-free mechanism for sampling from unnormalized, multimodal targets by synthesizing global PT-based initialization of conditional interpolants with local, closed-form diffusion transport. It robustly outperforms current MCMC paradigms in terms of sample quality and evaluation efficiency, especially in regimes where density evaluations are expensive and high-dimensional multimodality is present. The framework’s flexibility, analytical tractability, and strong empirical performance across domains mark it as an important advance in the theory and practice of MCMC and generative modeling for scientific inference.