Formulating Discrete Probability Flow Through Optimal Transport

TL;DR

This paper establishes a theoretical foundation for discrete probability flow using optimal transport principles, connecting it to continuous models, and introduces a new sampling method that improves outcome certainty.

Contribution

It formulates the discrete probability flow based on optimal transport theory and proposes a novel sampling method outperforming previous discrete diffusion models.

Findings

Validated on synthetic and CIFAR-10 datasets

Achieved more certain outcomes in generated samples

Connected discrete probability flow to Monge optimal transport

Abstract

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.