Technical Summary: Stochastic Transition-Map Distillation for Fast Probabilistic Inference

Problem Statement

Diffusion models have established state-of-the-art performance in generation quality, diversity, and distribution coverage compared to predecessors like GANs and VAEs. However, this performance is often accompanied by expensive inference costs due to the necessity of integrating continuous-time ordinary differential equations (ODEs) or stochastic differential equations (SDEs) over many steps.

While existing acceleration techniques such as distillation (e.g., Consistency Models) and Optimal Transport-based approaches exist, they often face a trade-off between generation quality/diversity and inference complexity. A critical limitation of current distillation methods is their reliance on deterministic inference (e.g., via the probability-flow ODE). Given a fixed noise sample, the output is unique. This restricts their utility in downstream tasks that inherently require stochastic inference, such as diffusion posterior sampling for inverse problems, energy-based fine-tuning, or improving alignment with user queries. Furthermore, many existing distillation frameworks require a pretrained teacher model, bi-level optimization, or the simulation and caching of trajectories, which complicates training and scalability.

Methodology: Stochastic Transition-Map Distillation (STMD)

The authors propose Stochastic Transition-Map Distillation (STMD), a teacher-free framework designed to accelerate diffusion model inference while preserving the stochastic transition structure of the underlying process.

Core Concept

Unlike score-based diffusion models that typically model the mean of the posterior distribution, STMD distills the full transition map associated with the sampling SDE. The method learns a conditional generator that predicts the transition probability $p(x_{t'} | x_t)$ for $t' < t$ in a single or few steps, avoiding the costly numerical integration of the reverse SDE.

Technical Formulation

1. Conditional Mean Flow (MF): The framework parameterizes the SDE transitions using a conditional Mean Flow model. Instead of learning a deterministic flow map, it learns the average velocity along the trajectory of the diffusion process.
2. x0-Parametrization: To model the transition $p(x_{t'} | x_t)$, the authors utilize an $x_0$-parametrization. They express the transition as an integral over the posterior $p(x_0 | x_t)$ and the known bridge measure $p(x_{t'} | x_0, x_t)$. Since the bridge measure for variance-preserving (VP) SDEs is available in closed form, the problem reduces to learning a conditional model $p_\theta(x_0 | x_t)$.
3. Training Objective: The model $u_\theta$ is trained to approximate the conditional mean velocity. The authors define a Conditional Mean Flow (CMF) loss function:
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   where $z_s$ is a straight-line interpolant between the data $x_0$ and a prior sample $z_1$, conditioned on the diffusion state $x_t$ and time $t$. The $\text{sg}$ operator denotes the stop-gradient.
4. Teacher-Free Training: The training process does not require a pretrained teacher model. It relies on samples from the joint distribution $p(x_0, x_t)$, which can be generated by sampling $x_0$ from the data distribution and applying the forward SDE kernel (which is known in closed form).

Inference

Inference is performed by integrating the learned conditional Mean Flow backward in time. The process involves:

1. Sampling an initial noise $x_1$ (or $z_1$).
2. Iteratively applying the learned velocity field $u_\theta$ to step from time $s$ to $s-\Delta s$.
3. At each diffusion time step $t$, sampling from the bridge distribution $p(x_{t-\Delta t} | x_0, x_t)$ to update the state, effectively simulating the stochastic transition in large steps.

Key Contributions

1. Novel Framework: The proposal of STMD, a distillation framework based on conditional Mean Flows that enables learning transition maps of high-dimensional diffusion SDEs using a single-step or few-step generator via a regression objective, without any pretrained teacher.
2. Theoretical Convergence: The derivation of convergence bounds for the Mean Flow algorithm in the 2-Wasserstein distance. To the authors' knowledge, this is the first convergence analysis of the Mean Flow algorithm in this metric. They extend this to prove that their learned conditional transition map converges to the transition kernel of the underlying reverse SDE.
3. Empirical Validation: Demonstration of the method on image generation and inpainting tasks across MNIST, CIFAR-10, and CelebA datasets, showing competitive performance with baselines while retaining stochasticity.

Experimental Results

The authors evaluated STMD on three datasets:

• MNIST: Using a U-Net backbone, the model was compared against the original Mean Flow algorithm and a Denoising Diffusion Probabilistic Model (DDPM). The authors utilized the Fréchet Distance (FD) in the latent space of a classifier (as FID is unreliable for MNIST). Results showed that STMD performs comparably to the baselines while supporting stochastic sampling.
• CIFAR-10: Similar experiments were conducted using a U-Net. The model achieved competitive Fréchet Inception Distance (FID) scores compared to DDPM and Mean Flow baselines for a fixed number of training iterations, while maintaining the full stochasticity of the diffusion sampler.
• CelebA: Experiments utilized a latent DiT architecture. The model generated unconditional samples with an FID of 8.28 (using 4 inference steps and 2 Mean Flow steps) and successfully performed image inpainting tasks.

In all cases, the models were trained from scratch with the same network backbones and iteration counts as the baselines to ensure fair comparison.

Significance and Claims

The paper claims that STMD addresses the gap between fast deterministic distillation and the need for stochastic inference in downstream tasks. By distilling the transition map rather than just the flow map, STMD preserves the stochastic structure of the diffusion process.

The authors emphasize that their approach is teacher-free, eliminating the need for bi-level optimization or trajectory caching, which makes the training process more efficient and scalable. The theoretical contribution of establishing 2-Wasserstein convergence bounds provides a rigorous foundation for the Mean Flow approach, bridging a gap between theoretical guarantees and practical application.

The work positions itself as a solution for tasks requiring probabilistic inference, such as inverse problems and posterior sampling, where deterministic distillation methods fall short. The authors remain modest regarding the current performance, noting that the limited number of training iterations prevents state-of-the-art results, but the comparative study validates the method's efficacy relative to existing baselines.