Stochastic Transition-Map Distillation for Fast Probabilistic Inference

• URL: http://arxiv.org/abs/2605.07661v1
• Date: Fri, 08 May 2026 12:31:44 GMT
• Title: Stochastic Transition-Map Distillation for Fast Probabilistic Inference
• Abstract summary: Diffusion models achieve strong generation quality, diversity, and distribution coverage, but their performance often comes with expensive inference.<n>In this work, we propose Transition-Map Distillation (STMD), a teacher-free framework for accelerating diffusion model inference.
• Score: 14.377745343756485
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Diffusion models for multivariate subsurface generation and efficient probabilistic inversion [0.0]
  Diffusion models offer stable training and state-of-the-art performance for deep generative modeling tasks.<n>We introduce a likelihood approximation accounting for the noise-contamination that is inherent in diffusion modeling.<n>Our tests show significantly improved statistical robustness, enhanced sampling of the posterior probability density function.
  arXiv  Detail & Related papers  (2025-07-21T17:10:16Z)
• Forward-only Diffusion Probabilistic Models [23.247850964456177]
  This work presents a forward-only diffusion (FoD) approach for generative modelling.<n>FoD directly learns data generation through a single forward diffusion process, yielding a simple yet efficient generative framework.<n>Our code is available at https://github.com/Algolzw/FoD.
  arXiv  Detail & Related papers  (2025-05-22T14:47:07Z)
• Generative Latent Neural PDE Solver using Flow Matching [8.397730500554047]
  We propose a latent diffusion model for PDE simulation that embeds the PDE state in a lower-dimensional latent space.<n>Our framework uses an autoencoder to map different types of meshes onto a unified structured latent grid, capturing complex geometries.<n> Numerical experiments show that the proposed model outperforms several deterministic baselines in both accuracy and long-term stability.
  arXiv  Detail & Related papers  (2025-03-28T16:44:28Z)
• Variational Schrödinger Diffusion Models [14.480273869571468]
  Schr"odinger bridge (SB) has emerged as the go-to method for optimizing transportation plans in diffusion models.<n>We leverage variational inference to linearize the forward score functions (variational scores) of SB.<n>We propose the variational Schr"odinger diffusion model (VSDM), where the forward process is a multivariate diffusion and the variational scores are adaptively optimized for efficient transport.
  arXiv  Detail & Related papers  (2024-05-08T04:01:40Z)
• Gaussian Mixture Solvers for Diffusion Models [84.83349474361204]
  We introduce a novel class of SDE-based solvers called GMS for diffusion models. Our solver outperforms numerous SDE-based solvers in terms of sample quality in image generation and stroke-based synthesis.
  arXiv  Detail & Related papers  (2023-11-02T02:05:38Z)
• A Geometric Perspective on Diffusion Models [57.27857591493788]
  We inspect the ODE-based sampling of a popular variance-exploding SDE. We establish a theoretical relationship between the optimal ODE-based sampling and the classic mean-shift (mode-seeking) algorithm.
  arXiv  Detail & Related papers  (2023-05-31T15:33:16Z)
• Reflected Diffusion Models [93.26107023470979]
  We present Reflected Diffusion Models, which reverse a reflected differential equation evolving on the support of the data. Our approach learns the score function through a generalized score matching loss and extends key components of standard diffusion models.
  arXiv  Detail & Related papers  (2023-04-10T17:54:38Z)
• ShiftDDPMs: Exploring Conditional Diffusion Models by Shifting Diffusion Trajectories [144.03939123870416]
  We propose a novel conditional diffusion model by introducing conditions into the forward process. We use extra latent space to allocate an exclusive diffusion trajectory for each condition based on some shifting rules. We formulate our method, which we call textbfShiftDDPMs, and provide a unified point of view on existing related methods.
  arXiv  Detail & Related papers  (2023-02-05T12:48:21Z)
• Score-based Continuous-time Discrete Diffusion Models [102.65769839899315]
  We extend diffusion models to discrete variables by introducing a Markov jump process where the reverse process denoises via a continuous-time Markov chain. We show that an unbiased estimator can be obtained via simple matching the conditional marginal distributions. We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.
  arXiv  Detail & Related papers  (2022-11-30T05:33:29Z)
• How Much is Enough? A Study on Diffusion Times in Score-based Generative Models [76.76860707897413]
  Current best practice advocates for a large T to ensure that the forward dynamics brings the diffusion sufficiently close to a known and simple noise distribution. We show how an auxiliary model can be used to bridge the gap between the ideal and the simulated forward dynamics, followed by a standard reverse diffusion process.
  arXiv  Detail & Related papers  (2022-06-10T15:09:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.