Stochastic Transition-Map Distillation for Fast Probabilistic Inference

TL;DR

This paper introduces STMD, a teacher-free distillation method that accelerates diffusion model inference while maintaining probabilistic sampling, suitable for tasks like inverse problems and energy-based fine-tuning.

Contribution

STMD distills the full transition map of diffusion models into a simple, scalable framework without requiring pretrained teachers or trajectory caching.

Findings

STMD achieves fast stochastic sampling with one or few steps.

The method maintains high-quality image generation on MNIST, CIFAR-10, and CelebA.

Theoretical convergence bounds are established in Wasserstein distance.

Abstract

Diffusion models achieve strong generation quality, diversity, and distribution coverage, but their performance often comes with expensive inference. In this work, we propose Stochastic Transition-Map Distillation (STMD), a teacher-free framework for accelerating diffusion model inference while preserving probabilistic sample generation. In contrast to score-based diffusion models, whose denoising parametrization models the mean of the posterior distribution, STMD distills the full transition map associated with the sampling stochastic differential equation (SDE). We parameterize these SDE transitions with a conditional Mean Flow model, yielding a one- or few-step stochastic sampler that retains the transition structure of the underlying diffusion process. This perspective is especially useful for downstream tasks that require stochastic inference, such as diffusion posterior sampling, inverse problems, and energy-based fine-tuning. Compared to recent distillation methods, STMD requires no pretrained teacher, bi-level optimization, or trajectory simulation and caching, enabling efficient and scalable training. We derive convergence bounds for our method in the Wasserstein distance, providing a strong theoretical foundation for our approach, and validate STMD on various image generation examples on the MNIST, CIFAR-10, and CelebA datasets.