# Stochastic Transition-Map Distillation for Fast Probabilistic Inference

arXiv: https://arxiv.org/abs/2605.07661 · HTML: https://mlsanity.com/paper/2605.07661/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 35 of its 54 references that are in the map.
- First posted on arXiv 8 May 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)
- Knowledge distillation / [Multimodal distillation](https://mlsanity.com/distill/multimodal-distillation/)

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model](https://mlsanity.com/paper/2212.00490/) (2022)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Tackling the Generative Learning Trilemma with Denoising Diffusion GANs](https://mlsanity.com/paper/2112.07804/) (2021)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed](https://mlsanity.com/paper/2101.02388/) (2021)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions](https://mlsanity.com/paper/2211.01916/) (2022)
- [Consistency Models Made Easy](https://mlsanity.com/paper/2406.14548/) (2024)
- [Flow map matching with stochastic interpolants: A mathematical framework for consistency models](https://mlsanity.com/paper/2406.07507/) (2024)
- [How to build a consistency model: Learning flow maps via self-distillation](https://mlsanity.com/paper/2505.18825/) (2025)
- [The Principles of Diffusion Models](https://mlsanity.com/paper/2510.21890/) (2025)
- [Transition Matching Distillation for Fast Video Generation](https://mlsanity.com/paper/2601.09881/) (2026)
- [Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps](https://mlsanity.com/paper/2602.05993/) (2026)

## Abstract (arXiv)

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.