Single-Step Consistent Diffusion Samplers

TL;DR

This work tackles the challenge of sampling from unnormalized densities $p_{\rm target}=\rho/Z$ with intractable $Z$ by introducing two consistency-based diffusion approaches. CDDS distills a pretrained diffusion model into a single-step sampler using incomplete trajectories, while SCDS trains a time-and-step-size conditioned control to permit single-step sampling without any pretrained teacher, jointly learning large-step shortcuts via a self-consistency objective. Across benchmarks including GMM, Image, Funnel, MW54, MW52, and LGCP, CDDS and SCDS achieve high-fidelity samples with far fewer function evaluations (often under $1\%$ of traditional diffusion costs) and, in the case of SCDS, can estimate the partition function $Z$ through the Radon–Nikodym framework. These results demonstrate that consistency-based diffusion substantially reduces sampling complexity for unnormalized targets while maintaining competitive accuracy, with SCDS offering data-free learning and adjustable inference budgets.

Abstract

Sampling from unnormalized target distributions is a fundamental yet challenging task in machine learning and statistics. Existing sampling algorithms typically require many iterative steps to produce high-quality samples, leading to high computational costs that limit their practicality in time-sensitive or resource-constrained settings. In this work, we introduce consistent diffusion samplers, a new class of samplers designed to generate high-fidelity samples in a single step. We first develop a distillation algorithm to train a consistent diffusion sampler from a pretrained diffusion model without pre-collecting large datasets of samples. Our algorithm leverages incomplete sampling trajectories and noisy intermediate states directly from the diffusion process. We further propose a method to train a consistent diffusion sampler from scratch, fully amortizing exploration by training a single model that both performs diffusion sampling and skips intermediate steps using a self-consistency loss. Through extensive experiments on a variety of unnormalized distributions, we show that our approach yields high-fidelity samples using less than 1% of the network evaluations required by traditional diffusion samplers.

Figures (5)

• Figure 1: Consistency distilled diffusion samplers learn to map consecutive intermediate states (black and gray dots) along partial ODE trajectories (green curve) directly to the terminal state.
• Figure 2: Graphical illustration of the training procedure for SCDS over the path space. First, the SDE trajectory (white) is simulated to compute the sampling loss $\mathcal{L}_{S}$. Next, a timestep $t$ and a step size $d$ are randomly sampled. From ${\mathbf{x}}_t$ on the simulated SDE trajectory, we execute two consecutive steps of size $d$ (red) along the PF-ODE trajectory (pink), obtaining the target ${\mathbf{x}}_{t+2d}^{\prime}$. Finally, the shortcut step of size $d$ (orange) predicts ${\mathbf{x}}_{t+2d}$ directly from ${\mathbf{x}}_t$, and the self-consistency loss $\mathcal{L}_{SC}$ minimizes the squared difference between ${\mathbf{x}}_{t+2d}$ and the two-step target ${\mathbf{x}}_{t+2d}^{\prime}$, ensuring multi-scale consistency.
• Figure 3: Visualization of the GMM and MW54 tasks. CDDS and SCDS recover all modes in just a single sampling step.
• Figure 4: Comparison of Sinkhorn distance for a range of NFEs between the proposed consistency samplers (CDDS, SCDS) and diffusion-based samplers (PIS, DDS, DIS). For most targets, CDDS and SCDS show competitive Sinkhorn values with baselines with much lower NFEs.
• Figure 5: Loss curves for the samplers studied in this paper. SCDS and CDDS exhibit stable learning across most settings, except for the image target distribution, where all samplers—except CDDS—show instability. Notably, the self-consistency loss and the sampling loss remain relatively independent.

Theorems & Definitions (3)

• Theorem 4.1
• Theorem 4.1
• proof