Solving Prior Distribution Mismatch in Diffusion Models via Optimal Transport

TL;DR

This paper introduces a novel optimal transport-based framework to precisely match distributions in diffusion models, effectively eliminating prior errors and significantly improving sampling accuracy and efficiency.

Contribution

It proposes a theoretically grounded OT map approach to eliminate prior distribution mismatch in diffusion models, enhancing their performance.

Findings

Prior error is fully eliminated both theoretically and practically.

The method achieves precise distribution matching in diffusion processes.

Experimental results show improved generation quality and sampling efficiency.

Abstract

Diffusion Models (DMs) have achieved remarkable progress in generative modeling. However, the mismatch between the forward terminal distribution and reverse initial distribution introduces prior error, leading to deviations of sampling trajectories from the true distribution and severely limiting model performance. This issue further triggers cascading problems, including non-zero Signal-to-Noise Ratio, accumulated denoising errors, degraded generation quality, and constrained sampling efficiency. To address this issue, this paper proposes a prior error elimination framework based on Optimal Transport (OT). Specifically, an OT map from the reverse initial distribution to the forward terminal distribution is constructed to achieve precise matching of the two distributions. Meanwhile, the upper bound of the prior error is quantified using the Wasserstein distance, proving that the prior error can be effectively eliminated via the OT map. Additionally, by deriving the asymptotic consistency between dynamic OT and probability flow, this method is revealed to be highly compatible with the intrinsic mechanism of the diffusion process. Experimental results demonstrate that the proposed method completely eliminates the prior error both theoretically and practically, providing a universal and rigorous solution for optimizing the performance of DMs.