Reproduce “Solving Prior Distribution Mismatch in Diffusion Models via Optimal Transport” on OpenResearch

About this paper

Zhanpeng Wang, Shenghao Li, Jiameng Che, Chen Wang, Shangling Jui, Na Lei, Zhongxuan Luo

This research presents an Optimal Transport (OT)-based framework to address the prior distribution mismatch in Diffusion Models, a fundamental issue where the forward terminal distribution does not align with the reverse initial distribution. The method constructs an Optimal Transport map to align these distributions, leading to improved generation quality and accelerated sampling, with W2 distance reduced to 0.00 for prior alignment and competitive FID scores achieved with fewer inference steps.

Key insights

• The prior distribution mismatch in DMs is formally equivalent to the non-zero Signal-to-Noise Ratio (SNR) problem, leading to an irreducible potential gap that causes sampling trajectories to deviate from the true data distribution.
• The proposed OT-based prior error eliminator provides a quantifiable error upper bound for the generative distribution, which is dependent on the accuracy of score matching and the approximation quality of the Brenier potential.
• DM training and sampling processes can be framed as stages of time-dependent Optimal Transport, with the static OT map asymptotically converging to the dynamic probability flow, confirming its intrinsic compatibility with diffusion dynamics.

Summary by alphaXiv.