Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Pascal Jutras-Dube, Jiaru Zhang, Ziran Wang, Ruqi Zhang

TL;DR

The paper tackles the inefficiency of sampling from unnormalized densities by introducing Self-Distilled One-Step Diffusion Samplers (OSDS), which learn a step-conditioned probability-flow ODE to emulate many small steps in a single large step. To address the instability of ELBO estimates in few-step regimes, OSDS combines a forward-consistent state distillation with a volume-consistency regularizer and introduces deterministic-flow importance weights that bypass fragile backward kernels. Empirically, OSDS achieves competitive sample quality with orders-of-magnitude fewer network evaluations and provides robust log Z estimates across synthetic and Bayesian benchmarks, including strongly multimodal targets. The work thus offers a practical framework for fast, reliable sampling and evidentiary estimation from unnormalized densities, with potential impact on scalable Bayesian inference and statistical physics simulations.

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

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

TL;DR

The paper tackles the inefficiency of sampling from unnormalized densities by introducing Self-Distilled One-Step Diffusion Samplers (OSDS), which learn a step-conditioned probability-flow ODE to emulate many small steps in a single large step. To address the instability of ELBO estimates in few-step regimes, OSDS combines a forward-consistent state distillation with a volume-consistency regularizer and introduces deterministic-flow importance weights that bypass fragile backward kernels. Empirically, OSDS achieves competitive sample quality with orders-of-magnitude fewer network evaluations and provides robust log Z estimates across synthetic and Bayesian benchmarks, including strongly multimodal targets. The work thus offers a practical framework for fast, reliable sampling and evidentiary estimation from unnormalized densities, with potential impact on scalable Bayesian inference and statistical physics simulations.

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

Paper Structure

This paper contains 60 sections, 3 theorems, 65 equations, 3 figures, 4 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Markov chain Monte Carlo samplers.
  4. Normalizing flows.
  5. Generative flow networks.
  6. Bridge samplers.
  7. Diffusion‐based samplers.
  8. Acceleration techniques for diffusion models.
  9. Summary.
  10. Preliminaries
  11. Diffusion Samplers and SDEs
  12. The FB-RND Framework for Evidence Estimation
  13. One-step Transport and Evidence Estimation
  14. Learning One‑Step Transport via State Consistency
  15. RND base loss at small steps.
  16. ...and 45 more sections

Key Result

Proposition 1

Let ${\mathbb{P}}$ be the path measure of the forward SDE in eq:generative-sde with control $u$, and let $\reflectbox{$${\mathbb{P}}$$}$ be the path measure of the reverse-time SDE in eq:noising-sde with terminal density $p_{\rm{target}}$. Then for any path ${\mathbf{x}}$ we have

Figures (3)

  • Figure 1: Graphical illustration of state-space self-distillation.
  • Figure 2: Single‑step OSDS samples (blue) overlaid on target energy contours. Samples populate distinct high‑density regions, illustrating broad mode coverage in one step.
  • Figure 3: Absolute and relative NFE savings of OSDS over the baseline diffusion samplers PIS, DDS as a function of the number of generated samples (log scale).

Theorems & Definitions (7)

  • Proposition 1: Forward-backward Radon-Nikodym derivative (RND)
  • Proposition 2: Finite-step forward-backward bound
  • Proposition 3: Deterministic-flow IS weight, \ref{['sec:det-logz']}
  • proof
  • proof
  • Definition 1: Time-adjoint.
  • proof