Title: Formulating Discrete Probability Flow Through Optimal Transport
Venue: NeurIPS 2023
Accepted: yes
Keywords: Discrete Probability Flow

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.

Formulating Discrete Probability Flow Through 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 .

1 Introduction

The emerging diffusion-based models [43, 20, 45, 46] have been proven to be an effective technique for modeling data distribution, and generating high-quality texts [31, 14], images [34, 11, 40, 37, 38, 21] and videos [22, 19, 39, 51, 17]. Considering their generative capabilities have surpassed the previous state-of-the-art results achieved by generative adversarial networks [11], there has been a growing interest in exploring the potential of diffusion models in various advanced applications [41, 33, 48, 55, 10, 32, 49, 52, 18, 53].

Diffusion models are widely recognized for generating samples in a stochastic manner [46], which complicates the task of defining an encoder that translates a sample to a certain latent space. For instance, by following the configuration proposed by [20], it has been observed that generated samples from any given initial point have the potential to span the entire support of the data distribution. To achieve a deterministic sampling process while preserving the generative capability, Song et al. [46] proposed the probability flow, which provides a deterministic map between the data space and the latent space for continuous diffusion models. Unfortunately, the situation differs when it comes to discrete models. For instance, considering two binary distributions ( P 0 = 1 2 , P 1 = 1 2 ) and ( P 0 = 1 3 , P 1 = 2 3 ) , there is no deterministic map that can transform the former distribution to the latter one, as it would simply be a permutation. Although some previous research has been

conducted on discrete diffusion models with discrete [24, 23, 4, 12, 9, 26, 16] and continuous [6, 47] time configurations, these works primarily focus on improving the sampling quality and efficiency, while sampling certainty has received less attention. More specifically, there is a conspicuous absence of existing literature addressing the probability flow in discrete diffusion models.

The aim of this study is to establish the fundamental theory of the probability flow for discrete diffusion models. Our paper contributes in the following ways. Firstly, we provide proof that under some conditions the probability flow of continuous diffusion coincides with the Monge optimal transport map during any finite time interval within the range of (0 , ∞ ) . Secondly, we propose a discrete analogue of the probability flow under the framework of optimal transport, which we have defined as the discrete probability flow . Additionally, we identify several properties that are shared by both the continuous and discrete probability flow. Lastly, we propose a novel sampling method based on the aforementioned observations, and we demonstrate its effectiveness in significantly improving the certainty of the sampling outcomes on both synthetic toy dataset and CIFAR-10 dataset.

Proofs for all Propositions are given in the Appendix. For consistency, the probability flow and infinitesimal transport of a process X t is signified by ˆ X t and ˜ X t respectively.

2 Background on Diffusion Models and Optimal Transport

First of all, we review some important concepts from the theory of diffusion models, optimal transport and gradient flow.

2.1 Continuous state diffusion models

Diffusion models are generative models that consist of a forward process and a backward process. The forward process transforms the data distribution p data ( x 0 ) into a tractable reference distribution p T ( x T ) . The backward process then generates samples from the initial points drawn from p T ( x T ) . According to [28], the forward process is modeled as the (time-dependent) Ornstein-Uhlenbeck (OU) process:

where θ t ≥ 0 , σ t > 0 , ∀ t ≥ 0 and B t is the Brownian Motion (BM). The backward process is the reverse-time process of the forward process [2]:

where ˜ B t is the reverse-time Brownian motion and p ( X t , t ) is the single-time marginal distribution of the forward process, which also serves as the solution to the Fokker-Planck equation [35]:

In order to train a diffusion model, the primary objective is to minimize the discrepancy between the model output s θ ( x t , t ) and the Stein score function s ( x t , t ) = ∇ x t log p ( x t , t ) [25]. Song et al. [45] demonstrate that, it is equivalent to match s θ ( x t , t ) with the conditional score function:

where λ t is a weighting function, t is uniformly sampled over [0 , T ] and p ( x t , t | x 0 , 0) is the forward conditional distribution.

It is noted that every Ornstein-Uhlenbeck process has an associated probability flow, which is a deterministic process that shares the same single-time marginal distribution [46]. The probability flow is governed by the following Ordinary Differential Equation (ODE):

In accordance with the global version of Picard-Lindelöf theorem [1] and the adjoint method[36, 7], the map

is a diffeomorphism ∀ t ≥ s > 0 . The diffeomorphism naturally gives a transport map.

2.2 Discrete state diffusion models

In the realm of discrete state diffusion models, there are two primary classifications: the Discrete Time Discrete State (DTDS) models and the Continuous Time Discrete State (CTDS) models, which are founded on Discrete Time Markov Chains (DTMC) and Continuous Time Markov Chains (CTMC), correspondingly. Campbell et al. [6] conducted a comparative analysis of these models and determined that CTDS outperforms DTDS. The DTDS models construct the forward process through the utilization of the conditional distribution q t +1 | t ( x t +1 | x t ) and employ a neural network to approximate the reverse conditional distribution

q

t

|

t

+1

(

x

t

|

x

applications, it is preferable to parameterize this model using

t

+1

p

θ

|

0

t

) =

+1

q

t

+1

|

t

(

x

t

+1

|

x

t

)

q

t

(

x

q

t

+1

(

x

t

+1

)

[24, 4] and obtain

t

)

p

. In practical

θ

k

|

k

+1

through

In contrast to DTDS models, a CTDS model is characterized by the (infinitesimal) generator [3], or transition rate, Q t ( x, y ) . The Kolmogorov forward equation [13] is:

The generator of the reverse process can be written by [6, 47]:

There are various approaches to train the model, such as the Evidence Lower Bound (ELBO) technique [6], and the score-based approach [47]. It has been observed that the reverse generator can be factorized over dimensions, allowing parallel sampling for each dimension during the reverse process. However, it is important to note that this independence is only possible when the time interval for each step is small.

2.3 Optimal transport

The optimal transport problem can be formulated in two primary ways, namely the Monge formulation and the Kantorovich formulation [42]. Suppose there are two probability measures µ and ν on ( R n , B ) , and a cost function c : R n × R n → [0 , + ∞ ] . The Monge problem is

The measure T # µ is defined through T # µ ( A ) = µ ( T -1 ( A )) for every A ∈ B and is called the pushforward of µ through T.

It is evident that the Monge Problem (MP) transports the entire mass from a particular point, denoted as x , to a single point T ( x ) . In contrast, Kantorovich provided a more general formulation, referred to as the Kantorovich problem :

where Π ( µ, ν ) is the set of transport plans , i.e.,

where π x and π y are the two projections of R n × R n onto R n . For measures absolutely continuous with respect to the Lebesgue measure, these two problems are equivalent [50]. However, when the measures are discrete, they are entirely distinct as the constraint of the Monge Problem may never be fulfilled.

The reverse process is:

2.4 Fokker-Planck equation by gradient flow

According to [27], the Fokker-Planck equation represents the gradient flow of a functional in a metric space. In particular, for Brownian motion, its Fokker-Planck equation, which is also known as the heat diffusion equation, can be expressed as:

and it represents the gradient flow of the Gibbs-Boltzmann entropy multiplied by -1 2 :

It is worth noting that Eq. 15 is the gradient flow of Eq. 14 under the 2-wasserstein metric ( W 2 ).

Chow et al. [8] have developed an analogue in the discrete setting by introducing the discrete Gibbs-Boltzmann entropy:

and deriving the gradient flow using a newly defined metric (Definition 1 in [8]). Since the discrete model is defined on graph G ( V, E ) , where V = { a 1 , ..., a N } is the set of vertices, and E is the set of edges, the discrete Fokker-Planck equation with a constant potential can be written as:

where N ( i ) = { j ∈ { 1 , 2 , ..., N }|{ a i , a j } ∈ E } represents the one-ring neighborhood.

3 Continuous probability flow

3.1 The equivalence of Ornstein-Uhlenbeck processes and Brownian motion

The diffusion models that are commonly utilized in machine learning are founded on OrnsteinUhlenbeck processes. First of all, we demonstrate that it is feasible to deterministically convert a time-dependent Ornstein-Uhlenbeck process into a standard Brownian motion.

Proposition 1. Let X t and Y t be a time-dependent Ornstein-Uhlenbeck process and a Brownian motion respectively: dX t = -θ t X t dt + σ t dB (1) t , dY t = dB (2) t , where B (1) t and B (2) t are two independent Brownian motions and θ t ≥ 0 , σ t > 0 , ∀ t ≥ 0 . Let ϕ t = exp( ∫ t 0 θ τ d τ ) , β t = ∫ t 0 ( σ τ ϕ τ ) 2 d τ . Then X t coincides in law with ϕ -1 t Y β t .

Building upon the aforementioned proposition, the primary focus of this paper is centered around the standard Brownian motion dY t = dB t .

3.2 Probability flow is a Monge map

Khrulkov et al. [29] have proposed a conjecture that the probability flow of Ornstein-Uhlenbeck process is a Monge map. However, they only provided a proof for a simplified case. We demonstrate that under some conditions, the conjecture is correct.

It is important to highlight that the continuous optimal transports presented in this paper are defined exclusively with the cost function: c ( x, y ) = 1 2 | x -y | 2 .

Within the context of generative models, a collection of training samples denoted as { x i } N i =1 is typically provided, and these samples are intrinsically defined by a distribution:

where δ ( x ) represents the Dirac delta function. Given a Brownian motion with an initial distribution in the form of Equation (18), the single-time marginal distribution is [35]

The probability flow is defined as [46]:

According to [1, 36, 7], the solution exists for all t > 0 and the map ˆ Y t + s ( ˆ Y t ) is a diffeomorphism for all t > 0 , s ≥ 0 . We have discovered that ˆ Y t + s ( ˆ Y t ) is the Monge map under some conditions and the time does not reach 0 or + ∞ .

Proposition 2. Given that Y 0 follows the initial condition (18), and all x i s lie on the same line, the diffeomorphism ˆ Y t + s ( ˆ Y t ) is the Monge optimal transport map between p B ( x, t ) and p B ( x, t + s ) , ∀ t > 0 , s ≥ 0 .

There is a counterexample [30] to demonstrate that the probability flow map does not necessarily provide optimal transport. It is important to note that their case differs from our assumptions in two ways. Firstly, they consider the limit case of ˆ Y + ∞ ( ˆ Y 0 ) . Secondly, the initial distribution of the counterexample does not conform to the form specified in Equation (18). Therefore, their counterexample is not applicable to our situation.

It has been shown that the heat diffusion equation can be regarded as the gradient flow of the GibbsBoltzmann entropy concerning the W 2 metric [27]. As W 2 is associated with optimal transport, it is reasonable to anticipate that the "infinitesimal transport" ˆ Y t + dt ( ˆ Y t ) is optimal [29].

In order to interpret the concept of "infinitesimal transport", we utilize the generator of the process Y t . Let C 2 c ( R n ) denote the set of twice continuously differentiable functions on R n with compact support. The generator A t is defined as follows [35]:

It is straightforward to verify that

We define the "infinitesimal transport" to be the diffeomorphism ˜ Y t + s ( ˜ Y t ) where ˜ Y t + s evolves according to the following equation

with the initial condition ˜ Y t = ˆ Y t . The generator of ˜ Y t + s is

Proposition 3. Given any t > 0 , there exists a δ t > 0 s.t. ∀ 0 < s < δ t , the diffeomorphism ˜ Y t + s ( ˜ Y t ) with the initial condition ˜ Y t = ˆ Y t is the Monge optimal transport map.

Let us return to the original Ornstein-Uhlenbeck process X t . As it is merely a deterministic transformation of the Brownian motion Y t , we can anticipate that the probability flow of X t , denoted by ˆ X t , will be a Monge map. In fact, this expectation holds true:

Proposition 4. Given that X 0 follows the initial condition (18), and all x i s lie on the same line, the diffeomorphism ˆ X t + s ( ˆ X t ) is the Monge optimal transport map for all t > 0 , s ≥ 0 .

4 Discrete probability flow

The continuous probability flow is deterministic, which means the "mass" at ˆ Y t is entirely transported to ˆ Y t + s during the time interval [ t, t + s ] . However, it is widely acknowledged that for discrete distributions µ and ν , there may not exist a T such that T # µ = ν . As a result, discrete diffusions cannot possess a deterministic probability flow. To establish the concept of the discrete probability flow , we employ the methodology of optimal transport. First of all, a discrete diffusion model is proposed as an analogue of Brownian motion. Secondly, we modified the forward process to create an optimal transport map, which is used to define the discrete probability flow. Finally, a novel sampling technique is introduced, which significantly improves the certainty of the sampling outcomes.

4.1 Constructing discrete probability flow

It is demonstrated that the process described by Equation (17) is a discrete equivalent of the heat diffusion process (14) [8]. We adopt this process as our discrete diffusion model and represent it in a more comprehensive notation.

The discrete diffusion model has K dimensions and S states. The states are denoted by i = ( i 1 , i 2 , . . . , i K ) , where i j ∈ { 1 , 2 , . . . , S } . The Kolmogorov forward equation for this process is

where P i j ( t | s ) means P ( x t = j | x s = i ) and Q D is defined as:

where d D ( i, j ) = ∑ K l =1 | i l -j l | . If we let the solution of the Equation (25) be denoted by P D ( t | s ) and assume an initial condition P 0 , the single-time marginal distribution can be computed as follows:

It is noteworthy that the process defined by Q D is not an optimal transport map, as there exist mutual flows between the states (i.e., there exists two states i , j with Q i j > 0 and Q j i > 0 ). Therefore, we propose a modified version that will be proved to be a solution to the Kantorovich problem, namely, an optimal transport plan. The modified version is defined by the following Q :

where

In order to avoid singular cases, We define Q i j ( t ) to be 0 when P Di ( t ) = 0 . In fact, it is easy to verify that P Di ( t ) > 0 for all t > 0 , i ∈ { 1 , 2 , . . . , K } . We will show that the process defined by Q is equivalent in distribution to the one generated by Q D .

Proposition 5. The processes generated by Q D and Q have the same single-time marginal distribution ∀ t > 0 .

Proposition 6. Given any t > 0 , there exists a δ t > 0 s.t. ∀ 0 < s < δ t , the process generated by Q provides an optimal transport map from P D ( t ) to P D ( t + s ) under the cost d D .

Proposition 6 demonstrates that Q D generates a Kantorovich plan between P D ( t ) and P D ( t + s ) under a certain cost function. On the other hand, the continuous probability flow is the Monge map between p B ( x, t ) and p B ( x, t + s ) . Therefore, it is reasonable to define the process defined by Q D as the discrete probability flow of the original process defined by Q .

Furthermore, the "infinitesimal transport" of the discrete process, which is defined by d ds ˆ P ( t + s ) = ˆ P ( t + s ) Q ( t ) , also provides an optimal transport map.

Proposition 7. Given any t > 0 , there exists a δ t > 0 s.t. ∀ 0 < s < δ t , the process above provides an optimal transport map from ˆ P ( t ) to ˆ P ( t + s ) under the cost d D .

4.2 Sampling by discrete probability flow

In order to train the modified model, we employ a score-based method described in the Score-based Continuous-time Discrete Diffusion Model (SDDM) [47]. Specifically, we directly learn the conditional probability P θ ( i l ( t ) | { i 1 , . . . , i l -1 , i l +1 , . . . , i K } ( t )) . According to proposition 5, it follows that P θ = P θ D , and consequently, the training process is identical to that of [47]. For the sake of brevity, we will employ the notation P θ i l | i \ i l ( t ) to replace P θ ( i l ( t ) | { i 1 , . . . , i l -1 , i l +1 , . . . , i K } ( t )) .

The generator of the reverse process is

̸

We use the Euler's method to generate samples. Given the time step length ϵ , the transition probabilities for dimension l is:

̸

When ϵ is small, the reverse conditional distribution has the factorized probability:

In this way, it becomes possible to generate samples by sequentially sampling from the reverse conditional distribution 32.

Transition to higher probability states The reverse process of the continuous probability flow, as described in Equation (20), causes particles to move towards areas with higher logarithmic probability densities. As the logarithm function is monotonically increasing, this reverse flow pushes particles to higher probability density states. This phenomenon is also observed in the discrete probability flow. By examining the reverse generator, as shown in Equation (30), it can be determined that the transition rate R i j ( t ) > 0 only when the destination state j has a higher probability than the source state i . This implies that transitions only occur in higher probability states. In contrast, the original continuous reverse process (2) and the discrete reverse process from (10) allow any transitions.

Reduction of Standard Deviation We measure the certainty of the sampling method by the expectation of the Conditional Standard Deviation (CSD):

where Std ( X s | X t ) = Var 1 2 ( X s | X t ) = E 1 2 X s [ X s -E X s [ X s | X t ] | X t ] . CSD s,t ( X ) is 0 when the process is deterministic, such as the continuous probability flow. In the discrete situation, there does not exist any deterministic map. However, our discrete probability flow significantly reduces CSD s,t ( X ) . Table 2 presents numerical evidence of this phenomenon. Therefore, we posit that the discrete probability flow enhances the certainty of the sampling outcomes.

5 Related Work

The concept of probability flow was initially introduced in [46] as a deterministic alternative to the Itô diffusion. In the work [44], they presented the Denoising Diffusion Implicit Model (DDIM) and demonstrated its equivalence to the probability flow. Subsequently, [29] investigated the relationship between the probability flow and optimal transport. They hypothesized that the probability flow could be considered a Monge optimal transition map and provided a proof for a specific case. Additionally, they conducted numerical experiments that supported their conjecture, showing negligible errors. However, [30] has discovered an initial distribution that renders probability flow not optimal.

The discrete diffusion models were first introduced by [43], who considered a binary model. Following the success of continuous diffusion models, discrete models have garnered more attention. The bulk of research on discrete models has focused primarily on the design of the forward process [24, 23, 4, 5, 26, 16, 9]. Continuous time discrete state models were introduced by [6] and subsequently developed by [47].

6 Experiments

We conduct numerical experiments using our novel sampling method by Discrete Probability Flow (DPF) on synthetic data. The primary goal is to demonstrate that our method can generate samples of comparable quality with higher certainty.

Experiments are conducted on synthetic data using the same setup as SDDM [47], with the exception that we replaced the generator Q with Equation (26). In addition to the binary situation ( S = 2 ) studied in [47], we also perform experiments on synthetic data with the state size S set to 5 and 10. To evaluate the quality of the generated samples, we generated 40,000 / 4,000 samples for binary data / other type of data using SDDM and DPF, and measured the Maximum Mean Discrepancy (MMD) with the Laplace kernel [15]. The results are shown in Table 1. It can be seen that the MMD value obtained using DPF is slightly higher than that of SDDM, which may be attributed to the structure of the reverse generator 10. Specifically, DPF approximates an additional term, Q t ( y, x ) , with the neural network, which potentially introduces additional errors to the sampling process, leading to a higher MMD value compared to SDDM. However, such difference is minimal and does not significantly impact the quality of the generated samples. As evident from the visualization of the

SDDM

DDM_OT

2spirals

8gaussians

checkerboard

circles

moons

pinwheel

swissroll

SDDM

DPF

SDDM

DPF

2spirals

8gaussians checkerboard

circles moons

pinwheel swissroll

2spirals

8gaussians

checkerboard

circles

moons

pinwheel

swissroll

DPF

2spirals

8gaussians

checkerboard

circles

moons

pinwheel

swissroll

SDDM

DPF

SDDM

DPF

2spirals

8gaussians

checkerboard

circles

moons

pinwheel

swissroll

𝑥

!

distributions obtained from SDDM and DPF in Figure 1, it is clear that DPF can generate samples that are comparable to those generated by SDDM.

In addition, we also compare the sampling certainty of DPF and SDDM by computing CSD s,t using a Monte-Carlo based method. Specifically, we set s = 0 and t = T , and sample 4,000 x t s with 10 x s s for each x t . We then estimate E ( x s | x t ) and Std ( x s | x t ) using the sample mean and sample standard deviation, respectively. The results of certainty are presented in Table 2. Our findings indicate that DPF significantly reduces the CSD , which suggests a higher certainty. Additionally, we visualize the results of 4,000 generated samples (in blue) from a single initial point (in red) in the binary case in Figure 2. It is apparent that the sampling of SDDM exhibits high uncertainty, as it can sample the entire pattern from a single initial point. In contrast, our method reduces such uncertainty and is only able to sample a limited number of states.

To provide a more intuitive representation of the generated samples originating from various initial points, we select 20 × 20 initial points arranged in the grid, and distinguish them using different colors. Subsequently, we visualize the results by sampling 10 outcomes from each initial point, as shown in Figure 3. We observe that the visualization of SDDM samples appears disorganized, indicating significant uncertainty. In contrast, the visualization of DPF samples exhibits clear regularity, manifesting in two aspects: (1) the generated samples from the same initial point using DPF are clustered by color, demonstrating the better sampling certainty of our DPF. (2) Both of the generated samples and initial points are colored similarly at each position. For example, in the lower

𝑥

!

𝜏

DPF (ours)

right area, a majority of the generated samples are colored purple, which corresponds to the color assigned to the initial points x T in that area. This observation demonstrates that most of the sampling results obtained through DPF are closer to their respective initial points, aligning with our design intention of optimal transport. It is worth noting that similar phenomena are observed across different state sizes, and we have provided these results in the Appendix.

Finally, we extended our DPF to the CIFAR-10 dataset, and compare it with the τ LDR-0 method proposed in [6]. The visualization results are shown in Figure 4. It can be seen that our method greatly reduces the uncertainty of generating images by sampling from the same initial x T . Detailed experimental settings and more experimental results are presented in the Appendix.

7 Discussion

In this study, we introduce a discrete counterpart of the probability flow and established its connections with the continuous formulations. We began by demonstrating that the continuous probability flow corresponds to a Monge optimal transport map. Subsequently, we proposed a method to modify a discrete diffusion model to achieve a Kantorovich plan, which naturally defines the discrete probability flow. We also discovered shared properties between continuous and discrete probability flows. Finally, we propose a novel sampling method that significantly reduces sampling uncertainty. However, there are still remaining aspects to be explored in the context of the discrete probability flow. For instance, to obtain more general conclusions under a general initial condition, the semi-group method [54] could be employed. Additionally, while we have proven the existence of a Kantorovich plan in a small time interval, it is possible to extend this to a global solution. Moreover, the definition of the probability flow has been limited to a specific type of discrete diffusion model, which also could be extended to a broader range of models. These topics remain open for future studies.

8 Acknowledgments and Disclosure of Funding

We would like to thank all the reviewers for their constructive comments. Our work was supported in National Natural Science Foundation of China (NSFC) under Grant No.U22A2095 and No.62072482.

References

Herbert Amann. Ordinary differential equations: an introduction to nonlinear analysis , volume 13. Walter de gruyter, 2011.

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

William J Anderson. Continuous-time Markov chains: An applications-oriented approach . Springer Science & Business Media, 2012.

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Sam Bond-Taylor, Peter Hessey, Hiroshi Sasaki, Toby P Breckon, and Chris G Willcocks. Unleashing transformers: parallel token prediction with discrete absorbing diffusion for fast high-resolution image generation from vector-quantized codes. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXIII , pages 170-188. Springer, 2022.

Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems , 35:28266-28279, 2022.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Shui-Nee Chow, Wen Huang, Yao Li, and Haomin Zhou. Fokker-planck equations for a free energy functional or markov process on a graph. Archive for Rational Mechanics and Analysis , 203:969-1008, 2012.

Max Cohen, Guillaume Quispe, Sylvain Le Corff, Charles Ollion, and Eric Moulines. Diffusion bridges vector quantized variational autoencoders. arXiv preprint arXiv:2202.04895 , 2022.

Guillaume Couairon, Jakob Verbeek, Holger Schwenk, and Matthieu Cord. Diffedit: Diffusionbased semantic image editing with mask guidance. In International Conference on Learning Representations , 2023.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems , volume 34, pages 8780-8794, 2021.

Patrick Esser, Robin Rombach, Andreas Blattmann, and Bjorn Ommer. Imagebart: Bidirectional context with multinomial diffusion for autoregressive image synthesis. Advances in Neural Information Processing Systems , 34:3518-3532, 2021.

William Feller. On the theory of stochastic processes, with particular reference to applications. In Selected Papers I , pages 769-798. Springer, 2015.

Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. arXiv preprint arXiv:2210.08933 , 2022.

Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research , 13(1):723-773, 2012.

Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10696-10706, 2022.

William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Advances in Neural Information Processing Systems , volume 35, pages 27953-27965, 2022.

Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-or. Prompt-to-prompt image editing with cross-attention control. In International Conference on Learning Representations , 2023.

Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Jonathan Ho, Chitwan Saharia, William Chan, David J. Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research , 23(47):1-33, 2022.

Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. In Advances in Neural Information Processing Systems , 2022.

Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037 , 2021.

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

Daniel D Johnson, Jacob Austin, Rianne van den Berg, and Daniel Tarlow. Beyond inplace corruption: Insertion and deletion in denoising probabilistic models. arXiv preprint arXiv:2107.07675 , 2021.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokkerplanck equation. SIAM journal on mathematical analysis , 29(1):1-17, 1998.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. arXiv preprint arXiv:2206.00364 , 2022.

Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, and Ivan Oseledets. Understanding DDPM latent codes through optimal transport. In International Conference on Learning Representations , 2023.

Hugo Lavenant and Filippo Santambrogio. The flow map of the fokker-planck equation does not provide optimal transport. Applied Mathematics Letters , 133:108225, 2022.

Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. In Advances in Neural Information Processing Systems , volume 35, pages 4328-4343. Curran Associates, Inc., 2022.

Zhen Liu, Yao Feng, Michael J. Black, Derek Nowrouzezahrai, Liam Paull, and Weiyang Liu. Meshdiffusion: Score-based generative 3d mesh modeling. In International Conference on Learning Representations , 2023.

Chenlin Meng, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2108.01073 , 2021.

Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Mcgrew, Ilya Sutskever, and Mark Chen. GLIDE: Towards photorealistic image generation and editing with text-guided diffusion models. In Proceedings of the 39th International Conference on Machine Learning , volume 162, pages 16784-16804, 2022.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

LS PONTRJAGIN. The mathematical theory of optimal processes. Interscience , 1962.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Ludan Ruan, Yiyang Ma, Huan Yang, Huiguo He, Bei Liu, Jianlong Fu, Nicholas Jing Yuan, Qin Jin, and Baining Guo. Mm-diffusion: Learning multi-modal diffusion models for joint audio and video generation. arXiv preprint arXiv:2212.09478 , 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems , volume 35, pages 36479-36494, 2022.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(4):4713-4726, 2023.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:2211.16750 , 2022.

Haoru Tan, Sitong Wu, and Jimin Pi. Semantic diffusion network for semantic segmentation. In Advances in Neural Information Processing Systems , volume 35, pages 8702-8716, 2022.

Guy Tevet, Sigal Raab, Brian Gordon, Yoni Shafir, Daniel Cohen-or, and Amit Haim Bermano. Human motion diffusion model. In International Conference on Learning Representations , 2023.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Vikram Voleti, Alexia Jolicoeur-Martineau, and Chris Pal. Mcvd - masked conditional video diffusion for prediction, generation, and interpolation. In Advances in Neural Information Processing Systems , volume 35, pages 23371-23385, 2022.

Qiang Wang, Haoge Deng, Yonggang Qi, Da Li, and Yi-Zhe Song. Sketchknitter: Vectorized sketch generation with diffusion models. In International Conference on Learning Representations , 2023.

Sirui Xu, Yu-Xiong Wang, and Liangyan Gui. Stochastic multi-person 3d motion forecasting. In International Conference on Learning Representations , 2023.

Kôsaku Yosida and JG Taylor. Functional analysis. Physics Today , 20(1):127-129, 1967.

Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojcic, Or Litany, Sanja Fidler, and Karsten Kreis. Lion: Latent point diffusion models for 3d shape generation. In Advances in Neural Information Processing Systems , volume 35, 2022.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper answers two important questions for the diffusion modelling community: \n\n(1) How to interpret the encoding given by the ODE transport map of a diffusion model - it is proven to be an optimal transport map.\n\n(2) How can this ODE type integration be translated to discrete diffusion models - a low stochasticity sampler is derived by again drawing on optimal transport ideas.\n\nThe low stochasticity sampler is interpreted on 2d datasets.\n\nSTRENGTHS:\nI like this paper as it very satisfyingly answers the above two important research questions in a very neat way using optimal transport.\nRegarding the continuous state space results, the paper has made a significant contribution to the understanding of the 'latent space' induced by a diffusion model integrated using the probability flow ODE. This result has been conjectured before and this paper has now proven the optimal transport interpretation of the flow for a realistic set of assumptions.\n\nRegarding the discrete state space results, the paper has made another significant contribution by formulating what an equivalent 'ODE' integration scheme could be for discrete state spaces, overcoming the challenge imposed by the discrete nature preventing a deterministic map from always being found. Formulating the problem as an optimal transport one seems a very natural way to overcome this. The intuition behind the derived optimal transport reverse rate is also very natural being a bias to only ever move to higher probability states and never flowing backward in such a way as to maintain the same marginals. This is really neat. I went through the proofs for propositions 5-7 and they seem correct to me.\n\nThe significance of this paper for the community will be large. For the discrete state spaces I can envision multiple works investigating this new reverse process formulation, for example, given that less transitions occur during sampling, it may be possible to achieve better sampling speed just as ODEs are easier to integrate in continuous state spaces.\n\nWEAKNESSES:\nThe paper is weak on the experimental side but I believe the strength of the theoretical insights outweigh this.\n\nThe paper has not provided any experiments to demonstrate the results on continuous state spaces however I believe these would mostly be similar to Khrulkov et al. The strength of the theoretical results stands alone. Note that I have not the expertise in this area to check the proofs.\n\nThe discrete state space experiments are toy but do very clearly present the idea. It is a shame that the authors did not include a practically relevant data type e.g. text as it would have been very interesting to investigate 'discrete interpolations' using this reduced stochasticity sampler. The paper would greatly benefit from an experiment of this type.\n\nThe presentation of the paper is poor in places although the overall flow of the paper is good.\n\nIn eq(16) you should define what a neighborhood is properly.\nOn line 178 you should define what you mean by mutual flow.\n\nTypos:\n\neq (2) dt is in the brackets as well for drift\n\nline 72 x_{t+t} should be x_{t+1}\n\nagain line 74 k+t should be k+1\n\neq 17 extra bracket\n\neq (26) s=0\n\nline 174 I assume you mean the sum is to K not S?\n\nEdit after rebuttal: I have read the author response and the new experiments alleviated my main concern and I have raised my score to 7.\n\nQUESTIONS:\nThe formulation for discrete optimal transport had to include a metric on the discrete state space, which may not be natural for some settings e.g. when the data has no intrinsic ordering. How would you interpret your results in this case, and could the metric perhaps be changed here.\n\nIn figures 9 and 10 in the appendix it seems that with bigger state sizes the difference between your low stochasticity sampler and the normal sampler are reduced. Do you expect your results to hold on practical datasets that are quite high dimensional and have a large number of states? Did you test this at all?\n\nLIMITATIONS:\nIt is unclear how the method scales to high dimensional / large number of state size datasets.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBuilding upon the concept of probability flows, the paper establishes a clear link between diffusion processes and optimal transport. It first shows that the continuous probability flow corresponds to a Monge map for every finite-length time interval. It then states and proves analogous statements in the context of continuous-time discrete space (CTDS) models. These results follow the intuition that diffusion should be asymmetric to avoid _mutual flows_ between states. Finally, the authors design a practical sampling method that substantially reduces the variance of end positions of particles, given their initial status.\n\nSTRENGTHS:\nThe paper presents a well-executed analysis of diffusion processes from the perspective of probability flows. It modifies a previous result by Khrulkov et al. [29] to make it applicable to the case of Brownian motion with initial distribution given by Eq. 17, i.e., that of a delta train induced by samples.\nTo analyze the discrete case, the authors revise the (reversed) heat diffusion process by preventing moves to states with lower probability. This procedure, reminiscent of well-known sampling techniques, leads to easy-to-analyze forward (Eq. 27) and backward transition rates. The authors show that the resulting process displaces mass optimally with respect to the L1 cost.\n\nWEAKNESSES:\nThe notation used throughout the paper is at times hard to parse: Several similar symbols are used concurrently but there is no concise overview of their meaning (and differences). This makes the presentation less intuitive and requires more effort from the reader. In particular, the interpretation of “infinitesimal transport” would benefit from a revision aimed at better disambiguating between processes $\\hat{Y}$ and $\\tilde{Y}$, and between generators $A$ and $\\hat{A}$. Also, the expanded state notation (i.e., $i_1, …, i_K$) is never needed in the main text and could be dropped altogether.\nIt would also be good to include additional derivations (in the appendix), such as the one of the reverse-time generator $R$ (eq. 29), which are not entirely trivial.\n\n**Experiments**\n\nThe authors include several experiments on toy datasets that showcase the effectiveness of the proposed sampling strategy in reducing the variance. They nicely illustrate the theoretical statement which links the process generated by $Q$ to an optimal transport map. However, they do not describe any practical scenarios in which this feature could be necessary or desirable. In particular, reduced variance in sampling seems to imply that slight biases in the initial distribution of particles could lead to unfaithful reconstructions of the other marginal, i.e., by creating “holes” that the optimal transport map fails to fill.\nI am also skeptical of the choice of experiments. Why is the current selection restricted to synthetic and low-dimensional settings? Did you run out of time or is the method's applicability limited?\n\n**Miscellaneous**\n\nTo ease the description of sampling processes it would be helpful to use the differential (or discrete increment) of time flowing backward, e.g., in Eq. 2.\n\nFurther, some **typos**:\n\n- In Table 1, given that the MMD value for the SDDM simulation on the checkerboard dataset has the wrong sign.\n- In line 72: $(x_{t+t})$.\n- In Eq. 19: missing a ½ factor?\n- In the Discussion section.\n\nQUESTIONS:\nWhat are possible practical applications of the proposed sampling method?\n\nWhat are possible alternative approximations of the quantity $Q_t$, which is likely the cause of the inferior quality of reconstructed marginals?\n\nWhy do points in Figure 2 appear to form rectangles, e.g., in the 2spirals, pinwheel, and swissroll datasets?\n\nEven though the infinite horizon case is outside the scope of the paper, it would be interesting to briefly discuss why the arguments presented in the paper break down in that case. Is this somehow related to mixing times of Markov chains? Additional insight, e.g., in the discussion, could provide a valuable starting point for further research.\n\nLIMITATIONS:\nThe practical relevance of the proposed sampling framework is not discussed. It is unclear whether the reduced variance achieved by it could be of use in real-world scenarios and justify the (admittedly) inferior quality of marginal reconstruction quality of DPF.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors present a theoretical way to define the probability flow for continuous-time diffusion models with discrete state-space, a notion which is largely used for continuous state-space diffusion models, but does not generalize to the discrete setting. They bring several contributions:\n\n- In the continuous setting, they show that the probability flow (under finite time horizon) of the Ornstein-Uhlenbeck process, starting from a finite collection of samples, is the optimal Monge map with respect to the quadratic cost, thus proving the conjecture of [1].\n\n- Then, they extend the notion of probability flow in the discrete setting, when the noising process $P_D$ is chosen as the discrete analogue of the Brownian motion. To do so, they consider a modified version of the transition rate of $P_D$ (with its reverse generator still tractable) and show that it generates a Kantorovitch plan between $P_D(t)$ and $P_D(t+s)$ under a certain cost function, for $t>0$, $0<s<\\delta_t$. Relying on the correspondence between optimal transport and probability flow that they highlight in the continuous case, they define the process hence generated as the discrete probability flow of the original process. \n \n- Relying on the score-based framework in continuous time developed by [2], they train two discrete diffusion models, one with the true noising process $P_D$, the other one with the modified noising process generated by $Q$, called Discrete Probability Flow (DPF). They show experimentally that DPF decreases the uncertainty occurring in the sampling phase, without hurting the quality of the generated samples, and give illustration of how DPF empirically satisfies optimal transport in this setting.\n\n[1] Understanding ddpm latent codes through optimal transport, Khrulkov et al., 2022.\n\n[2] Score-based Continuous-time Discrete Diffusion Models, Sun et al., 2022.\n\nSTRENGTHS:\n- Strong theoretical result in the continuous setting. Really good contribution in the field of discrete state space diffusion models, to bridge the gap between diffusion model and optimal transport.\n\n- Paper very well-written, easy to follow.\n\nWEAKNESSES:\nThe scalability of the experiments.\n\nQUESTIONS:\nIn Proposition 6, the existence of a Kantorovitch plan in the discrete setting is not provided for the whole time interval. Do you observe empirically any change in the transport of the particles on the whole time interval, that could  be aligned with this theoretical gap ?\n\nLIMITATIONS:\nThe notion of probability flow is only defined for the noising process corresponding to the Brownian motion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a formulation of discrete probability flow using tools of Optimal Transport. First establishing conditions under which continuous probability flows are OT maps, the authors establish a version for discrete case. The authors then define a discrete probability flow using the connection to OT. Experiments compare the proposed method with SDDM showing that the proposed method yields comparable quality with reduced uncertainty.\n\nSTRENGTHS:\nThe paper develops discrete probability flows using optimal transport with applications to diffusion modeling. Probability flows, OT, and diffusion models are all topics of great interest in the literature, which is a strength of the paper. I enjoyed the clear exposition of detailed math. The results were also clearly presented an accurately described. Overall, I found the paper to be a quite enjoyable and interesting read.\n\nWEAKNESSES:\nI see no major weaknesses of the paper. Always, one can ask for more simulations and comparisons, and that would strengthen the paper. However, the work as presented is already quite strong and interesting, and I do not see further simulations or experiments to be necessary for the author's intended contribution. \n\nDetailed comments: \n- Line 102 \"This holds under the 2...\" It is a bit unclear what is meant. The equation 14 holds? The interpretation from 13->14? Also, under what assumptions? \n- Line 105: \"using a newly defined metric\" what is this metric? \n- Line 107, a bit more detail about the neighborhood and eq 16 seems helpful. \n- Line 121: \"This paper\" --> \"We\" would be clearer. (Also \"This paper\" in other places.)\n- I don't believe SDDM is defined anywhere. It is also a bit surprising to have not discussed in detail SDDM since that is the main comparison.\n\nQUESTIONS:\nNone.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}