One-Step Diffusion Samplers via Self-Distillation and Deterministic Flow

TL;DR

This paper introduces a one-step diffusion sampler that uses self-distillation and deterministic flow to generate high-quality samples efficiently in a single step, reducing computational costs.

Contribution

It proposes a novel one-step diffusion sampling method with a step-conditioned ODE and a deterministic flow importance weight, improving efficiency and stability in sampling and evidence estimation.

Findings

Achieves high-quality samples with only one or few steps.

Maintains robust ELBO estimates with fewer network evaluations.

Performs competitively on synthetic and Bayesian benchmarks.

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. We introduce one-step diffusion samplers which learn a step-conditioned ODE so that one large step reproduces the trajectory of many small ones via a state-space consistency loss. We further show that standard ELBO estimates in diffusion samplers degrade in the few-step regime because common discrete integrators yield mismatched forward/backward transition kernels. Motivated by this analysis, we derive a deterministic-flow (DF) importance weight for ELBO estimation without a backward kernel. To calibrate DF, we introduce a volume-consistency regularization that aligns the accumulated volume change along the flow across step resolutions. Our proposed sampler therefore achieves both fast sampling and stable evidence estimate in only one or few steps. Across challenging synthetic and Bayesian benchmarks, it achieves competitive sample quality with orders-of-magnitude fewer network evaluations while maintaining robust ELBO estimates.