Single-Step Consistent Diffusion Samplers

Pascal Jutras Dube, Patrick Pynadath, Ruqi Zhang

Keywords: Sampling, Diffusion Processes, Self-Consistency, Single-Step Inference

TL;DR: We propose consistent diffusion-based samplers that enables single-step sampling from an unnormalized target distribution.

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 instead 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.

Submission Number: 73