Formulating Discrete Probability Flow Through Optimal Transport
This work addresses a foundational gap in discrete diffusion models for generative modeling, offering a theoretical framework that leads to improved sampling performance.
The paper tackles the lack of a deterministic probability flow theory for discrete diffusion models by proving that continuous probability flow is a Monge optimal transport map under certain conditions and extending this to discrete cases, enabling a new sampling method that outperforms previous discrete diffusion models on synthetic and CIFAR-10 datasets.
Continuous diffusion models are commonly acknowledged to display a deterministic probability flow, whereas discrete diffusion models do not. In this paper, we aim to establish the fundamental theory for the probability flow of discrete diffusion models. Specifically, we first prove that the continuous probability flow is the Monge optimal transport map under certain conditions, and also present an equivalent evidence for discrete cases. In view of these findings, we are then able to define the discrete probability flow in line with the principles of optimal transport. Finally, drawing upon our newly established definitions, we propose a novel sampling method that surpasses previous discrete diffusion models in its ability to generate more certain outcomes. Extensive experiments on the synthetic toy dataset and the CIFAR-10 dataset have validated the effectiveness of our proposed discrete probability flow. Code is released at: https://github.com/PangzeCheung/Discrete-Probability-Flow.