DPM-OT: A New Diffusion Probabilistic Model Based on Optimal Transport

TL;DR

DPM-OT introduces a novel diffusion probabilistic model leveraging optimal transport to generate high-quality images rapidly, reducing sampling steps to around 10 while maintaining stability and minimizing mode mixture.

Contribution

It proposes a unified framework using optimal transport for fast diffusion models, significantly improving sampling speed and sample quality with theoretical guarantees.

Findings

Achieves high-quality samples within ~10 steps

Reduces mode mixture in short-step sampling

Outperforms existing methods in speed and quality metrics

Abstract

Sampling from diffusion probabilistic models (DPMs) can be viewed as a piecewise distribution transformation, which generally requires hundreds or thousands of steps of the inverse diffusion trajectory to get a high-quality image. Recent progress in designing fast samplers for DPMs achieves a trade-off between sampling speed and sample quality by knowledge distillation or adjusting the variance schedule or the denoising equation. However, it can't be optimal in both aspects and often suffer from mode mixture in short steps. To tackle this problem, we innovatively regard inverse diffusion as an optimal transport (OT) problem between latents at different stages and propose the DPM-OT, a unified learning framework for fast DPMs with a direct expressway represented by OT map, which can generate high-quality samples within around 10 function evaluations. By calculating the semi-discrete optimal transport map between the data latents and the white noise, we obtain an expressway from the prior distribution to the data distribution, while significantly alleviating the problem of mode mixture. In addition, we give the error bound of the proposed method, which theoretically guarantees the stability of the algorithm. Extensive experiments validate the effectiveness and advantages of DPM-OT in terms of speed and quality (FID and mode mixture), thus representing an efficient solution for generative modeling. Source codes are available at https://github.com/cognaclee/DPM-OT