Formulating Discrete Probability Flow Through Optimal Transport
Authors: Pengze Zhang
Hubery Yin
Chen Li
Xiaohua Xie
Venue: NeurIPS 2023 poster (2023)
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
Pengze Zhang Sun Yat-sen University zhangpz3@mai12 edu. en
Hubery Yin" WeChat, Tencent Inc. hubery@tencent.comm
Chen Li WeChat, Tencent Inc. chaselli@tencente com
Xiaohua Xie f
Sun Yat-sen University xiexiaoh6@mail. edu. cn
Abstract
Continuous diffusion models are commonly acknowledged to display a deter- ministic 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 prob- ability flow the Monge optimal transport map under certain conditions, and also present equivalent evidence for discrete cases. In view of these find- ings, we are then able to define the discrete probability flow in line with the principles of optimal transport. Finally, drawing upon our newly established def- initions, we propose a novel sampling method that surpasses previous discrete diffusion models in its ability to generate more certain outcomes. Extensive periments on the synthetic toy dataset and the CIFAR- dataset have validated the effectiveness of our proposed discrete probability flow. Code is released https: http//ggtttttttt/ee / hatpaaaactttuniiisssrte-rrrababa
Introduction
The emerging diffusion-bassee models 43
20
45
46
have been proven to be an effective technique
for modeling data distribution, and generating high-quality texts [3] 14], images [34] 40 37
38 [1] and videos (22) 19 39 [7]. Considering their generative capabilities have surpassed the previous state-of-the-ard results achieved by generative adversarial networks (T]) there has been growing interest in exploring the potential of diffusion models in various advanced applications
40
T8
Diffusion models are widely recognized for generating samples in a stochastic manner [46], which complicates the task of defining an encoder that translates sample 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 deterministic sampling process while preserving the generative capability, Song er 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 (Po 2 Pi = 1) and (Po 1
3, Pi 2), there is no deterministic map that can transform the former distribution to the latter one, as it would simply be permutation. Although some previous research has been
'Equal contribution. This work was done when Pengze Zhang was an intern at WeChat. i Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

conducted on discrete diffusion models with discrete [24 23 4 12, 9 26 16 and continuous 66 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 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, oo). Secondly, we propose 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 novel sampling method based on the aforementioned observations, and we demonstrate effectiveness 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 infinitesimall transport of a process X, signified by X, and X, respectively.
2
Background 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 forward process and backward process. The forward process transforms the data distribution Pdtta (zo) into a tractable reference distribution pr (IT). The backward process then generates samples from the initial points drawn from (xr)). According to (28), the forward process is modeled as the (time-dependent) process:
-aattebeUhhleebeck (OU)
dx,- = - dt + add,,
(1)
where 0, 2 0, 0,V1 Vt 0 and B, is the Brownian Motion (BM). The backward process is the reverse-time process of the forward process (2):
dX, 1- Vixx x, logp(X,, t)]dt + o,dB.,
where B, the reverse-time Brownian motion and p(X,,, t) is the single-time timeee-timeeeee marginal distribution of the forward process, which also serves as the solution to the Fokker- Planck equation [5]
x 0t p(z.t) = 6 v (xp(z, t)) + o p(xx,).
In order to train a diffusion model, the primary objective is minimize the discrepancy between the model output sp (xx, t) and the Stein score function s(z,, = Vz, log (og,( 1) (25). Song et al. [45 demonstrate that, is equivalent to match s0 (Ire () with the conditional score function:
0* arg min E. {AeErro. [11s0(r),)) Va, log p(z,, Mo. 0)]
where X₁ is weighting function, t is uniformly sampled over [0.2] and p(z,, t f| 0) is the forward conditional distribution.
noted that every Ornstein-Uhlenber 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):
1 70 (X,,)))) /
dx; = -,
In accordance with the global version of Picard-Lindelis theorem (1] and the adjoint method 36 7 the map
_ - R Ta .
(6)
x, x,+ X.
is a diffcomorphism Vt 2 0. The diffcomorphism naturally gives 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 el al. [6] conducted a comparative analysis of these models and determined that CTDS outperforms DTDSS The DTDS models construct the forward process through the utilization of the conditional distribution (c+1p((22+11 (ma) and employ neural network to
approximate the reverse conditional distribution Qr|++1 (Ze (xa+1) = 9 ( In practical
g+ 11:1(2()1)
applications, it preferable to parameterize this model using po x+1
24
and obtain P through
Pkjk+1 (2x
Σ gk|k+ 1, 2 (20
(50((z6/0)) Σ (1+1 06+110 (2)) Polk-41 (2o (---1)--
In contrast to DTDS models, a CTDS model is characterized by the (infinitesimal) generator 3 or transition rate, Qx y). The Kolmogorov forward equation [3 is:
a ot 2r (1)la)))))))
Σ
dejs (
(8)
The reverse process is:
as late U (z,124) = Σ
(9)
The generator the reverse process can be written by 16 477:
Re x) = qc(x)) Qe v) (y))
Σ No
( (wol()))))ee (2, y)-
(10)
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, interval for each step is small.
important to note that this independence only possible when the time
2.3 Optimal transport
The optimal transport problem be formulated in two primary ways, namely the Monge formulation and the Kantorovich formulation [42). Suppose there are two probability measures p and on (IR",B) and cost function I IR" [0, +oo]. The Monge problem is
(MP) inf { c(x, T(x)) dp(x)) Tap =v
(11)
The measure Tafe is defined through Teu(A) = 4(T '(A)) for every A e B and is called the pushforward of μ through T.
is evident that the Monge Problem (MP) transports the entire mass from a particular point, denoted x, single point T(x). In contrast, Kantorovich provided more general formulation, referred to as the Kantorovich problem:
(K)) inf U cdy: : E 1 ?
(12)
where [1 v) is the set of transport plans, i.c.,
-0 l1,,v) = fr E P(R" x R") (Tz) = p4 (Ty)) where म, and ny are the two projections of R x R onto R", 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.
(13)
3

---

2.4 Fokker-Planck equation by gradient flow
According [27]. the Fokker-Plann Planckk equation represents the gradient flow of functional in a metric space. In particular, for Brownian motion, its Fokker-Planck equation, which is also known as the heat diffusion equation, can expressed as:
① हe plz, 1) A t),
1
(14)
and it represents the gradient flow of the Bobbs-Boltzmann entropy multiplied by -
1 1 / p(x) log dz. 2 5m
S())
(15)
2
It is worth noting that Eq. 15 is the gradient flow of Eq. 14 under the 2-wasserstein metric (W2).
Chow al. 8 have developed analogue in the discrete setting by introducing the discrete
Gibbs-Bolee Boltzmnnnnnnnn entropy:
S Σ Pi log pt,
(16)
and deriving the gradient flow using newly defined metric (Definition 1 in (). Since the discrete model is defined on graph G(V, E), where V {ai...., aN J is the set of vertices, and E is the set of edges, the discrete Fokker-Planck equation with a constant potential can written as:
d
di
Σ Pj Phr JEN(
(17)
where ((i) = E e (1,2, (HH2,,,,,NH,,,, } E} represents the one-ring neighborhood.
Continuous probability flow
3.1
The equivalence of Ornstein-Chlenbeck processes and Brownian motion
The diffusion models that are commonly utilized in machine learning are founded on Orsttei- Uhlenbeck processes. First of all, we demonstrate that it is feasible to deterministically convert time-dependent
process into a standard Brownian motion.
Proposition Let X, and Y, be time-dependent Ornstein- Whhstein-UUllenbeck process and a Brownian motion respectively: dX, = - + o,dB(") dY, = dB((2,, where B( and B,?) are two independent Brownian motions and 0, 2 0, a. > 0, vt 2 0. Let de = =exp( ( 0, dr), Be. = (=)) dr. Then X, coincides in law with o, 8.. Building upon the aforementioned proposition, the primary focus of this paper is centered around the standard Brownian motion dY, d.B..
3.2 Probability flow a Monge map
Khrulkov el al. 29 have proposed a conjecture that the probability flow of Crnstein-Uhlenbed process is Monge map. However, they only provided proof for simplified case. We demonstrate that under some conditions, the conjecture is correct.
is important highlight that the continuous optimal transports presented in this paper are defined exclusively with the cost function: c(I, 3) = 1/ - y/
Within the context of generative models, a collection of training samples denoted as {r:) is typically provided, and these samples intrinsically defined by distribution:
N
p(x, 0)
8(x xz);
(18)

---

where (()) represents the Dirac delta function. Given Brownian motion with an initial distribution
the form of Equation (8). the single-time marginal distribution 335
pn(a, () (2ml)= & exp(
(19)
21
The probability flow defined as 46]
1 Vg log pg PY t)dt.
dY,
(20)
2
According to [136]] 71. the solution exists for all > 0 and the map (Y,) is a diffeomorphism for all > 0, s 2 O. We have discovered that (Y,) is the Monge map under some conditions and the time does reach 0 or +Oc.
Proposition 2. Given that Yo follows the initial condition 18 and all x;s lie on the same line, the diffeomorphism (Yi) the Monge optimal transport map between pB t) and pz(x, s), v 0, 2 O.
There is counterexample 30 demonstrate that the probability flow map does not necessarily provide optimal transport. is important to note that their case differs from our assumptions in two ways. Firstly, they consider the limit case of (Y)). Secondly, the initial distribution of the counterexample does not conform to the form specified in Equation (8). Therefore, their counterexample is not applicable to our situation.
It has been shown that the heat diffusion equation can be regarded the gradient flow of the Gibbs- Boltzmann entropy concerning the W₂ metric [27]. As W2 is associated with optimal transport, is reasonable to anticipate that the "infinitesimal transport" dz (Y,) optimal 229
In order to interpret the concept of "iniinitesimal transport". we utilize the generator of the process Y, Let (IR") denote set of twice continuously differentiable functions on Rn with compact support. The generator A, is defined follows [35]
f(f(+s:) f(Y)) Wf E (RR).
A.f lim Ar-0!
(21)
At
It is straightforward to verify that
A, = Va 1
1).
(22)
2
We define the "infinitesimal transport" to be the diffeomorphism (Y.) where Yes. evolves according to the following equation
1 - (Y 2
(23)
with the initial condition Y, = Yr. The generator of Ye4 is
A - Vy,,, log (Yg
(24)
Proposition 3.
Given any t 0. there exists 6, S.t. 4O < 8,, the diffeomorphism
(Y)) with the initial condition Y, Y, is the Monge optimal transport map.
Let us return to the original Ornstein- Cereeib--hheeeckk process X,. it is merely a deterministic transfor- mation the Brownian motion Y, we can anticipate that the probability flow of X,, denoted by X,, will be Monge map. In fact, this expectation holds true:
Proposition 4. Given that Xa follows the initial condition 78 and all x; lie on the same line, the diffeomorphism X (Xt) is the Monge optimal transport map for all t > 0, 0.
5

---

Discrete probability flow
The continuous probability flow is deterministic, which means the "mass" at Yú is entirely transported to YI+s during the time interval [t, ( s). However, it is widely acknowledged that for discrete distributions lt and V, there may not exist a such that Taje = v. As 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, discrete diffusion model is proposed as 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, 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 discrete equivalent of the heat diffusion process (4) [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 (i1, i2,................ .., ix), where ij E (1, The Kolmogorov forward equation for this process is
d di ()(())
Σ (1))) (1),
(25)
where (t(ss) means P(x, j|æ,, i) and QD is defined as:
dp(i,)) 1,
QD =
Σ (((ddb((()))))) QD dp(i,j) = 0, 0,
(26)
otherwise,
where (i.j) = - |ù jil- If we let the solution of the Equation 25 be denoted by Po (t
and assume initial condition Po. the single-time marginal distribution can be computed follows:
PDi(t) Po:(t)=E PoyQubi (t))
(27)
noteworthy that the process defined by QD not an optimal transport map, as there exist mutual flows between the states (i.e., there exists two states i, j with Ø3 > 0 and Q' >> 0). Therefore, we propose 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:
(t)) Pp.())
dp(i,) 1,
11(1)
Σ [((kkdd(i (1),
dp(i,j) = 0, otherwise.
(28)
0.
where
Ir x > 0,
(x)
(29)
(0, x s O.
In order to avoid singular cases, We define (t( to be 0 when PD:(t) 0. In fact, it is easy to verify that (t.(t) 0 for all é > 0 E 21, K}. We will show that the process defined by Q equivalent distribution to the one generated by QD-
Proposition 5. tion Vt > 0.
The processes generated by QD and Q have the same single-time marginal distribu-
Proposition Given any >> O, there exists a 6, > S.I. 4A < 8G, the process generated by Q provides an optimal transport map from Pp(t) to Pp (t + 8) under the cost dD-

---

Proposition 6 demonstrates that QD generates a Kantorovich plan between Pb(t) and Pp + s) under certain cost function. On the other hand, the continuous probability flow is the Monge map between pn (n 1) and ps(2, + s). Therefore, it is reasonable to define the process defined by Qn as the discrete probability flow of the original process defined by Q.
Furthermore, the "infinitesimal transport" of the discrete process, which is defined by di P(1 s) P(t + s)Q(t) also provides an optimal transport map.
Given any t > O, there exists a; > s.t. A < s < òt, the process above provides an optimal transport map from P(t) to P(t + s) under the cost dp-
Proposition 7.
4.2 Sampling by discrete probability flow
In order to train the modified model, we employ score-based method described in the Score-based Continuous-tine Discrete Diffusion Model (SDDM) 47]. Specifically, we directly learn the condi-
tional probability pe (t) (i....., 42-1, M++1, ..... (t)). According to proposition 5, it follows that pd pe D. and consequently, the training process is identical to that of [47]. For the sake of brevity, we will employ the notation po igla,; (t) to replace pa pa (i((t) (fi,..., i- (411............... , (t)).
The generator of the reverse process is
PR RELU PE Benny
Doing (8)
1), [[((kdd((()))))) (t),
dp(i, 1 and is + jts
(1)
(30)
dp(i,)) O, otherwise.
0,
We use the Euler's method to generate samples. Given the time step length e, the transition probabili- ties for dimension / is:
c (it 1 (R)
it - i(t),
()((tt)
(31)
ut(t - c) i(t)
When is small, the reverse conditional distribution has the factorized probability:
po (ii(-e))): - e))((() = -T (ir(1 - (i((t=))i(())
(32)
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 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 (t) 0 only when the destination state has higher probability than the source state i. This implies that transitions only occur in higher probability states. contrast, the original continuous reverse process 2 and the discrete reverse process from (1O) allow any transitions.
Reduction of Standard Deviation We measure the certainty of the sampling method by the expectation of the Conditional Standard Deviation (CSD):
CSD,,((X) Ex. [Sdd (t,)(,,)X)),,
(33)
where Std(X,|X,) Var (X, = B Xx Ex, [x, CSD,,, is 0 when the process deterministic, such as continuous probability flow. In the discrete situation, there does not exist any deterministic map. However, our discrete probability flow significantly reduces CSD,,, (X). Table presents numerical evidence of this phenomenon. Therefore, we posit that the discrete probability flow enhances the certainty of the sampling outcomes.
7

---

Table 1: Comparison of generation quality for SDDM and DPF, in terms of MMD with Laplace kernel using bandwith=0.11 Lower values indicate superior quality.
2spirals
8gaussians checkerboard
circles
moons
pinwheel
swissroll
discrete dimension = 32, state size =2
SDDM DPF (ours)
2.18e-06 1.89e-05
4.28e-06
1.33e-06
6.22e-06 5.62e-06
2.10e-06 1.60e-05
4.27e-06
1.09e-05
2.22e-05
3.27e-05 2.42e-05
2.18e-05
discrete dimension = 16, state size =5
SDDM DPF (ours)
2.06e-
1.1e-4
2.43e-4
1.74e-4
2.20le-4
3.37e-4
1.43e-4
3.87e-4
5.87c-4
4.93e-4
3.83e-4
3,43e-4
6.64e-4
3.20c-4
discrete dimension = 12, state size
10
SDDM DPF (ours)
5.52e-4 7.19c-4
3.0le-4
4.39c-4
4.22e-4
2.7le-4
2.90c-4
3.39c-4
3.49e-4
5.99e
6.65e-4
4.34e-4
4.14
5.17e-4
Table 2: Comparison of certainty for SDDM and DPF, in terms of CSD on 4,000 initial points, each of which has generated samples. Lower values indicate superior certainty.
2spirals
8gaussians
checkerboard
circles
moons
pinwheel
swissroll
discrete dimension = 32, state size =2
SDDM DPF (ours)
14,3053 2.1719
14.1882
14.7433
14.4327
14.1739
14.0450
14.0548
1.7945
2.0693
1.7210
2.0573
2.1834
1,8892
discrete dimension 16, state size =
SDDM DPF (ours)
14,4645
14.6143
14.6963
14.4807
14.2397
14.2466
14.2659
1.9711
1.9367
1.4172
1.7185
1.7668
1.9633
1.6665
discrete dimension 12, state size 10
SDDM DPF (ours)
12.8463
12.7933
13.0158
12.9232
12.6663
12.7634
12.7880
1,8123
1.3178
1.1348
1.4625
1.4859
1.8435
1.5227
Related Work
The concept probability flow was initially introduced in 46 as 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 Monge optimal transition map and provided 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 26 16 9] Continuous time discrete state models were introduced by (6] and subsequently developed by [47]
Experiments
We conduct numerical experiments using our novel sampling method by Discrete Probability Flow (DPF) 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 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, Qe(y, x), with the neural network, which potentially introduces additional errors to the sampling process, leading to a higher MMD value compared SDDM. However, such difference is minimal and does not significantly impact the quality of the generated samples. As evident from the visualization of the

---

S0DM
DPF
2spirals
Sgaussions
checkerboard
circles
moons
pinwhoel
swissroll
Figure 1: Visualization of the generation quality generated binary samples for SDDM and DPF.
5000
DPF
2spirals
8gaussians
checkerboard
drcles
moons
pinwheel
swissroll
Visualization of the generating certainty generated binary samples for SDDM and DPF. All the samples (in blue) are randomly generated from the single initial point (in red).
Figure 2:
S90X
DPF
Zspirals
Bgaussians
checkerboard
circles
moons
pinwheel
swissroll
Xy
Figure 3: Visualization of the generated binary samples from the given initial points ær. Different colors distinguish the generated samples from different initial points TT.
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 using a Monte-Cart based method. Specifically, we set s 0 and T, and sample 4,000 2;5 with 10 x,s for each We then estimate [xx,) and Std 22 (x)] using the sample mean and sample standard deviation, respectively. The results of certainty are presented Table Our findings indicate that DPF significantly reduces the CSD, which suggests higher certainty. Additionally, we visualize results of 4,000 generated samples (in blue) from a single initial point (in red) in the binary case 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 limited number of states.
in Figure 2
To provide more intuitive representation of the generated samples originating from various initial points, we select 20 x 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 colored similarly at each position. For example, in the lower

---

Xy
rL.DR-O
DPF (ours)
Figure 4: Image modeling CIFAR-10 dataset. The figure is divided into three groups: initial points Ir, sampling results of rLDR-0, and sampling results of our DPF. For each row, the sampled images obtained from the same initial point.
right area, majority of the generated samples are colored purple, which corresponds to the color assigned to the initial points 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-I0 dataset, and compare it with the rLDR-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 zy. Detailed experimental settings and more experimental results are presented in the Appendix.
7 Discussion
In this study, we introduce 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 discrete diffusion model to achieve Kantorovich plan, which naturally defines the discrete probability flow. We also discovered shared properties between continuous and discrete probability flows. Finally, propose novel sampling method that significantly reduces sampling uncertainty. However, there are still remaining aspects to be explored the context of the discrete probability flow. For instance, obtain more general conclusions under general initial condition, the semi-group method [54) could be employed. Additionally, while we have proven the existence of a Kantorovich plan in small time interval, it is possible extend this to global solution. Moreover, the definition of the probability flow has been limited to 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. U22A9929 and No.62072482.
10

---

References
[1] Herbert Amann. Ordinary differential equations: an introduction to nonlinear analysis, vol- ume 13. Walter de gruyter, 2011.
[2] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
[3] William J Anderson. Continuous- time Markov chains: An applications-oriennted approach. Springer Science Business Media, 2012.
[4] 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: 14:17981-17993 2021.
[5] 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- rector-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.
[6] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 55:28266-282777 2022.
[7] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[8] Shui-Nee Chow, Wen Huang, Yao Li, and Haomin Zhou. Fokker-planck equations for free energy functional or markov process on graph. Archive for Rational Mechanics and Analysis, 003:969-1008, 2012.
[9] Max Cohen, Guillaume Quispe, Sylvain Le Corff, Charles Ollion, and Eric Moulines. Diffusion bridges vector quantized variational autoencoders. arXiv preprint arXiv:2202. 048000.0.48..... 2022.
[10] Guillaume Couairon, Jakob Verbeek, Holger Schwenk, and Matthieu Cord. Diffedit: Diffusion- based semantic image editing with mask guidance. In International Conference on Learning Representations, 2023.
[11] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. In Advances Neural Information Processing Systems, volume 34, pages 8780-8794, 2021.
[12] 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-355 2021.
[13] William Feller. On the theory of stochastic processes, with particular reference to applications. In Selected Papers I, pages 769-798. Springer, 2015.
[14] Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. arXiv preprint arXiv:22/0.08933, 2022.
[15] 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.
[16] 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.
[17] 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.

---

[18] Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfrr Aberman, Yael Pritch, and Daniel Cohen-ur Prompt-to-pronee image editing with cross-atrention control. In International Conference on Learning Representations, 2023.
[19] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02.00 2022.
[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[21] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47):1-33, 2022.
[22] 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.
[23] Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.0003: 2021.
[24] Emiel Hoogcboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34: 14454-1122465 2021.
[25] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[26] Daniel D Johnson, Jacob Austin, Rianne van den Berg, and Daniel Tarlow. Beyond in- place corruption: Insertion and deletion in denoising probabilistic models. arXiv preprint arXiv:2107.07675 2021.
[27] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
[28] Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. arXiv preprint 2rXiv:2206.00366 2022.
[29] Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, and Ivan Oseledets. Understanding DDPM latent codes through optimal transport. In International Conference on Learning Representations, 2023.
[30] Hugo Lavenant and Filippo Santambrogio. The flow map of the fokker-planck equation does not provide optimal transport. Applied Mathematics Letters, 133:108225, 2022.
[31] Xiang Li, John Thickstun, Ishaan Guirajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. In Advances in Neural Information Pro- cessing Systems, volume 35, pages
Curran Associates, Inc., 2022.
[32] 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.
[33] 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.
[34] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, 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.
[35] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
12

---

[36] LS PONTRIAGIN. The mathematical theory of optimal processes. Interscience, 1962.
[37] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25, 2022.
[38] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[39] 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:22/2.094788 2022.
[40]
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-36499, 2022.
[41] 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-472) 2023.
[42] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[43] Jaseha
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[44] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[45] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[46] 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 2rXiv:20111134555 2020.
[47] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:2211.16750 2022.
[48] 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.
[49] Guy Tevet, Sigal Raab, Brian Gordon, Yoni Shafir, Daniel Cohen-or, and Amit Haim Bermano. Human motion diffusion model. In International Conference Learning Representations, 2023.
[50] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
[51] Vikram Voleti, Alexia Dolicoeur-Martinees and Chris Pal. Mevd masked conditional video diffusion for prediction, generation, and interpolation. In Advances in Neural Information Processing Systems, volume 35, pages 23371-23388 2022.
[52] Qiang Wang, Haoge Deng, Yonggang Qi, Da Li, and Yi-Zhe Song. Sketchknitter: Vectorized sketch generation with diffusion models. In International Conference on Learning Representa- tions, 2023.
[53] Sirui Xu, Yu-Xiong Wang, and Liangyan Gui. Stochastic multi-person 3d motion forecasting. In International Conference on Learning Representations, 2023.
[54] Kōsaku Yosida and JG Taylor. Functional analysis. Physics Today, 20(1): 1967.
[55] Xiaohui Zeng, Arash Vahdar, Francis Williams, Zan Gojeic, 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.
13


Reviews:
[{"review_id": "m29qGzxypG", "reviewer": "Reviewer_Rnpe", "summary": "The 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", "questions": "None. ", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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. ", "weaknesses": "I 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Unclear expression \"This holds under the 2...\" in Line 102.**\n\nAs the gradient flow of the same functional is different under different metrics, we would like to emphasize here that Eq. (13) is the gradient flow of Eq. (14) under the $W_2$ metric. Thanks for your advice. We will refine our expression in our final version.\n\n> **\"metric\" in Line 105.** \n\nThis metric is defined in [1] as Definition 1. Due to its reliance on several definitions and theorems in [1], it is not easy to rephrase it in a concise manner. Thank you for the reminder, we will provide explanations in the footnotes to help readers understand it better.\n\n> **Definition of “neighborhood” in Line 107.**\n\nThe 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 $N(i)$ represents the one-ring neighborhood of $i$ in the graph, which can be written as: $N(i) = \\\\{j \\\\in \\\\{1, 2,..., N\\\\}|\\\\{a_i, a_j\\\\} \\in E\\\\}$. We will further supplement this definition in line 107.\n\n> **This paper\" --> \"We\" would be clearer in Line 121.**\n\nWe acknowledge the suggestion and will modify the wording from \"this paper\" to \"we\".\n\n> **Details of the method SDDM.**\n\nThe SDDM (Score-based continuous-time Discrete Diffusion Model) is the method proposed in [2]. We will further supplement this definition in line 199, and provide more details on the SDDM method.\n\n---\n**Reference:**\n\n[1] Chow S. et al.. Fokker–planck equations for a free energy functional or markov process on a graph. Archive for Rational Mechanics and Analysis, 2012.\n\n[2] Sun, H. et al.. Score-based continuous-time discrete diffusion models. ICLR, 2023."}, {"review_id": "D55gEE8q0W", "reviewer": "Reviewer_wEPe", "summary": "In 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.", "questions": "In 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 ?", "limitations": "The notion of probability flow is only defined for the noising process corresponding to the Brownian motion.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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.", "weaknesses": "The scalability of the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Scalability of the experiments.**\n\nPlease refer to our general response.\n\n> **In 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\nWe limit the time frame to a narrow interval, as the validity of the proof hinges on the constancy of the sign of $P_i(t) - P_j(t)$.\nAlternatively, if both equations in Eq. (70) are established concurrently, a contradiction arises whereby 2 equals 0.\nConsequently, the KKT condition cannot be satisfied by any suitable Lagrange multipliers, thereby rendering the plan sub-optimal.\n\nFrom an intuitive standpoint, if we assume $P_i(t) > P_j(t)$ in the interval $[t,t + \\epsilon / 2)$ and $P_i(t) < P_j(t)$ in $(t+\\epsilon/2, t+\\epsilon]$, it follows that $\\Pi^i_j > 0$ and $\\Pi^j_i > 0$. Assuming $\\Pi^i_j > \\Pi^j_i$, we can demonstrate that the given plan is sub-optimal. If we define a new plan as ${\\Pi^{\\*}}^i_j = \\Pi^i_j - \\Pi^j_i$ and ${\\Pi^{\\*}}^j_i = 0$, we can verify that the resultant plan ${\\Pi^{*}}$ incurs a lower transportation cost than $\\Pi$. \nThe preceding derivation establishes the tightness of our announcement, indicating that the optimal transport plan cannot be extended across the entire time interval.\n\nIn order to confirm the existence of such a scenario, we explicitly construct it in the case where $K=1$.\nSince $P(t) = P(0) e^{Q_{D} t}$, we can obtain the analytical solution through eigen decomposition. The eigen decomposition of $Q_D$ can be obtained through difference equations, yielding the following outcomes: $\\lambda_i = 2 cos(i \\pi / S) - 2$ and $v_i = (1, cos(\\theta_i/2), ... , cos((2S-1)\\theta_i/2))$, where $\\theta_i = \\arccos((\\lambda_i + 2) / 2)$. \nSubsequently, we can assess a basic scenario wherein $S = 3$ and $P(t = 0) = (0.1, 0, 0.9)$. It can be observed that $P_{0}(t = 0) > P_{1}(t = 0)$ and $P_{0}(t = 0.1) < P_{1}(t = 0.1)$.\n\n> **The notion of probability flow is only defined for the noising process corresponding to the Brownian motion.**\n\nThanks for your comment. In contrast to continuous processes, which necessitate the stochastic term to be a Brownian motion, there are few assumptions regarding the discrete stochastic term. As a result, the consideration of the drift term becomes unnecessary as it can be assimilated into the stochastic term. However, it is worth exploring the potential distinctive properties that may arise from treating these two terms separately. We intend to investigate this particular scenario in future research endeavors."}, {"review_id": "e9VUVrFcLL", "reviewer": "Reviewer_v2T5", "summary": "Building 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", "questions": "What 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", "limitations": "The 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Notation is hard to parse ...**\n\nWe apologize for not describing the notations. We use $X_t$ and $Y_t$ to represent the OU process and Brownian motion. For consistency, we adopt the notation '^' and '~' to signify the corresponding probability flow and infinitesimal transport. The same convention applies to the generator $A$; however, we apologize for the omission of '^' in Eq. (20) and Eq. (21).\n\nThanks for your advice. In the revised version, we will use $i$ instead of $(i_1, ..., i_k)$, and include certain non-trivial derivations.\n> **Practical scenarios.**\n\nPlease refer to our general response.\n> **Slight biases in the initial distribution ... unfaithful reconstructions of the marginal ...**\n\nWhen the initial distribution is biased, the marginal distribution will also change, and almost all diffusion models and their sampling methods, such as continuous probability flow, suffer from this problem. Nonetheless, if the initial distribution is free of holes, the marginal distribution will also be free of holes.\n> **Choice of experiments, scalability.**\n\nPlease refer to our general response.\n> **Differential of time flowing backward.**\n\nThank you for your suggestion. We will include the backward differential equation in our paper as follows: $\\frac{\\partial}{\\partial s}q_{s|t}(x_s|x_t)=-\\sum_y q_{s|t}(y|x_t)R_t(y,x_s), s<t$.\n> **Typos ... negative MMD ... missing a ½ factor ...**\n\n(1) We use the unbiased estimation $\\text{MMD}_u^2$ from Eq.(3) of [1]. As mentioned in [1], it is possible for $\\text{MMD}_u^2$ to produce negative values. We will provide additional explanations in our paper.\n\n(2) We appreciate your kind reminder and will correct these equations for the missing factor of ½.\n> **Practical applications.**\n\nPlease refer to our general response.\n> **Alternative approximations of $Q\\_t$ & inferior quality of reconstructed marginals.**\n\nIn the sampling process, two terms are present on the right side of Eq. (9): $\\frac{q_t(y)}{q_t(x)}$ and $Q_t(y, x)$. In SDDM, only one term, i.e., $\\frac{q_t(y)}{q_t(x)}$ is estimated using a neural network, as ${Q_D}_t(y, x)$ is known. In contrast, our method evaluates both terms using quantities approximated by the neural network, as $Q_t(y, x)$ depends on $\\frac{q_t(y)}{q_t(x)}$ (Eq. 27). This approximation may lead to slightly inferior quality than the SDDM using precise ${Q_D}_t(y, x)$. Since this discrepancy arises from the neural network fitting error, we currently have no feasible alternative approximations to achieve a superior outcome.\n\nIn order to improve the outcome, it is possible to formulate a discrete counterpart of the technique presented in reference [2], in which the diffusion models are expressed within a variational framework, and the forward process is subject to optimization rather than being predetermined. Through this approach, it may be feasible to approximate the entire right-hand side of Eq.(9), rather than solely $q_t(x)$.\n> **Figure 2 appears to form rectangles.**\n\nThis arises from the construction of our synthetic dataset. We construct the synthetic dataset $S$ states and $K$ dimensions by encoding the x-axis and y-axis coordinates of the toy dataset (normalized to [0, 1]) into K/2-bit S-ary. This is equivalent to dividing the data space into $S^{K/2} \\times S^{K/2}$ rectangle regions, where the first few dimensions determine the approximate location of the data. Since our proposed method significantly reduces the uncertainty, each dimension (including the first few dimensions) has only a limited number of possible values. As a result, the points in Figure 2 appear to form rectangles.\n> **Infinite horizon case.**\n\nThe infinite horizon case is not addressed in our study due to the presence of singularities that pose significant challenges. For instance, in the case of Brownian motion, the distribution at $t=\\infty$ assumes a uniform distribution over the entire ${\\mathbb{R}^n}$, which is not well-defined. \n\nAdditionally, the probability flow of Brownian motion at $t = 0$ also experiences a singularity. By taking the limit of the right-hand side of Eq. (19) and letting $d_i = x_t - x_i$, we obtain:\n\n$$\\begin{aligned} \\lim_{t \\to 0^+} -\\frac{1}{2} \\nabla_{x_t} p_B (x_t, t) = & \\lim_{t \\to 0^+} \\frac{\\sum_i exp(-\\frac{d_i^2}{2t}) \\frac{d_i}{2t}}{\\sum_j exp(-\\frac{d_j^2}{2t})} \\\\\\\\ = & \\lim_{z \\to +\\infty} \\sum_i \\frac{d_i}{\\sum_j exp((d_i^2 - d_j^2) z) / z}.\\end{aligned}$$\nSince\n\n$$\\lim_{z \\to +\\infty} exp((d_i^2 - d_j^2)z) / z =\\begin{cases}+\\infty, & \\text{if}~~d\\_i^2 \\> d\\_j^2,\\\\\\\\ 0, & \\text{if}~d_i^2 \\leq d_j^2, \\end{cases}$$\nwe have\n\n$$\\lim_{t \\to 0^+} - \\frac{1}{2} \\nabla_{x_t} p_B (x_t, t) = \\lim_{z \\to +\\infty} \\frac{d_{i_{min}}}{\\sum_j exp((d_{i_{min}}^2 - d_j^2) z) / z},$$\nwhere $i_{min} = \\mathop{\\arg\\min}\\limits_{i} d_i^2 $. (noting that $i_{min}$ may not be unique, but we exclude this scenario as it does not significantly affect our analysis). Consequently, we obtain:\n\n$$\\lim_{t \\to 0^+}-\\frac{1}{2} \\nabla_{x_t} p_B (x_t, t)=\\begin{cases}0, & \\text{if}~x_{t=0}=x_{i_{min}},\\\\\\\\ d_{i_{min}}*\\infty, & \\text{else}, \\end{cases}$$\nwhere $d_{i_{min}} * \\infty$ indicates that the vector is oriented in the direction of $d_{i_{min}}$ and has an infinite norm. Consequently, the right-hand side of Eq.(19) lacks Lipschitz continuity, leading to non-unique solutions. Actually, if Eq.(19) has a unique solution near $t=0$, the distribution $p_B(x,t)$ will always be a summation of Dirac deltas, which contradicts Eq.(18). Due to our reliance on the solution of ODE, we are unable to analyze the behavior in the vicinity of $t=0$. Consequently, we have limited our study to finite intervals. However, we remain committed to continuing our efforts toward resolving this problem and arriving at a comprehensive conclusion.\n\n---\n**Reference:**\n\n[1] Gretton, A. et al.. A kernel two-sample test. JMLR, 2012.\n\n[2] Huang C. W. et al.. A variational perspective on diffusion-based generative models and score matching. NeurIPS, 2021."}, {"review_id": "XS8EhYQc3N", "reviewer": "Reviewer_WcKm", "summary": "The 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.", "questions": "The 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?", "limitations": "It is unclear how the method scales to high dimensional / large number of state size datasets.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "I 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **The 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.**\n\nThank you for your comments. In [1], Khrulkov et al. have demonstrated the results on continuous state spaces through numerical experiments, and our contribution is to provide a theoretical basis for their conclusions. More importantly, we focus on analogizing this conclusion to discrete states and designing our discrete probability flows. Therefore, our experiments are primarily conducted around the study of discrete probability flow.\n\n> **The 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\nPlease refer to our general response.\n\n> **The presentation of the paper is poor in places although the overall flow of the paper is good.**\n\nThanks for your advice, we will improve our presentation.\n\n> **Definition of \"neighborhood\" in Eq. (16), and \"mutual flow\" in line 178.**\n\nThanks for your advice. We will clarify the two definitions in lines 107 and 178 as follows.\n\n  (1) The discrete model is defined on a graph $G = (V, E)$, where $V = \\\\{a_1,..., a_N\\\\}$ is the set of vertices, and $E$ is the set of edges. Therefore, the neighborhood of $i$ in the graph can be defined as:   $N(i) = \\\\{j \\in \\\\{1, 2,..., N\\\\}|\\\\{a_i, a_j\\\\} \\in E\\\\}$. \n\n (2) Given two states $i$, $j$, if it has $Q^{i}\\_{j} > 0$ and  $Q^{j}\\_{i} > 0$, it indicates the presence of mutual flows between states $i$ and $j$.\n\n> **Typos.**\n\nThanks for your suggestion. We will correct these typo errors.\n\n> **The 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\nThanks for your comment. This paper aims to propose a discrete probability flow to reduce the sampling uncertainty. Optimal transmission is an approach to studying DPF, which can better explain and characterize the nature and sampling process of DPF, and unify the concept of discrete and continuous probability flows. Therefore, even without the description of the metric, the definition of DPF still applies, and the results can still be explained in terms of reduced sampling uncertainty.\n\n> **In 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.**\n\nOur method of reducing uncertainty is independent of the number of states and dimensions. However, since Figure 9 and 10 are visualized in the 'float space' instead of the 'encoding space', there is an illusion that the difference between SDDM and PDF decreases as the state increases. Specifically, our synthetic data with a state size of $S$ and a dimension size of $K$ is established by encoding the x and y coordinates of the toy dataset (normalized to [0, 1]) to K/2-bit S-ary respectively. In this encoding, the first dimension of the encoding has the greatest impact on the data position. For example, in binary encoding (state size = 2), the first bit divides the data space into two parts, and determines the part in which it resides. However, as the number of states increases, the space is divided into more parts, and the small change of the first bit can not significantly change the position of the number it represents. This will lead to a narrowing of the gap between our DPF and SDDM in the visualization. Therefore, in such situations, the quantitative results in Table 2 are more appropriate for verifying the reduction of uncertainty in the encoding space.\n\n>**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? Scales to high dimensional / large number of state size datasets?**\n\nPlease refer to our general response.\n\n---\n**Reference:**\n\n[1] Khrulkov, V. et al.. Understanding ddpm latent codes through optimal transport. ICLR, 2023."}]

Author rebuttal:
# General Response to All Reviewers 

We thank all the reviewers for their time, insightful suggestions, and valuable comments. We are delighted that they find our approach interesting and theoretically sound (**Reviewer Rnpe, wEPe**), establishes a clear link between diffusion processes and optimal transport (**Reviewer v2T5**), and will have large significance for the community (**Reviewer WcKm**). We have also been inspired by the reviewers to further explore the discrete probability flow and attempt to apply the model to various application scenarios in our future study. 

First of all, we summarize some common questions and provide our response as follows.

>**Scalability and the selection of synthetic datasets.**

As Proposition 5 established without any assumption of the number of states and dimensions, our method is with good scalability, and guaranteed to scale to any size as well as practical datasets. It is worth noting that any model can use our sampling method as long as its generator follows Eq. (25), including SDDM, as we have the same training procedure. To further verify our method is still applicable in higher dimension or state scenarios, we increased the number of states $S$ and dimension $K$ to 50 for experiments. Specifically, we set $S=50, K=20$ and $S=5, K=50$ to avoid dimension redundancy in the K/2-bit S-ary encoding for the toy dataset (float64) coordinates (i.e., $50^{20/2} < 2^{64}$ and $5^{50/2} < 2^{64}$). The results of this experiment are shown in Table 1\* and Table 2\*, which clearly demonstrate that our method can significantly reduce sampling uncertainty even with larger state and dimension sizes.

In this paper, we focus on establishing the fundamental theory for the probability flow of discrete diffusion models, and verifying its related properties. Currently, we have selected synthetic dataset for validation for two reasons: (1) it allows a more intuitive representation of sampling properties for easy analysis, and (2) complete code of SDDM on synthetic data is publicly available for easy reproduction. Through various metrics and visualizations constructed on the synthetic dataset, the properties of discrete probability flow in terms of optimal transport and uncertainty reduction have been explicitly verified. Thanks for your suggestion, in our future study we will continue to explore this topic and make further attempts in real-world applications.

*Table 1\*. Comparison of certainty for SDDM and DPF on the synthetic dataset with $K = 20$ and  $S = 50$, in terms of $ECV$ on 4,000 initial points, each of which has 10 generated samples. Lower values indicate superior certainty.*
|         | 2spirals | 8gaussians | checkerboard | circles | moons | pinwheel | swissroll |
| :----:| :----: |:----:  | :----: |:----: |:----: |:----: |:----: |
| **SDDM** | 25.8777 | 26.5288 | 25.5106 | 25.7398 | 25.6984 | 25.6984 | 26.4767 |
| **DPF (ours)** | 2.7113 | 4.4274  | 2.6217 | 3.7554 | 3.0774  | 3.3054 | 3.8183 



*Table 2\*. Comparison of certainty for SDDM and DPF on the synthetic dataset with $K = 50$ and  $S = 5$, in terms of $ECV$ on 4,000 initial points, each of which has 10 generated samples. Lower values indicate superior certainty.*
|         | 2spirals | 8gaussians | checkerboard | circles | moons | pinwheel | swissroll |
| :----:| :----: |:----:  | :----: |:----: |:----: |:----: |:----: |
| **SDDM** | 47.2706 | 47.5810  | 47.4964 | 47.2733 | 47.0047  | 46.9103 | 46.9819 |
| **DPF (ours)** | 2.0335 | 1.8134  | 0.7418 | 1.7143 | 1.2840 | 1.4245 | 1.5720 |

>**Possible practical applications of the proposed sampling method.**

Thanks for the suggestion, we will supplement the discussion on the practical relevance of the proposed sampling framework. Analogously to the effect of continuous probability flow on the continuous diffusion model, we believe that reducing uncertainty can also bring many benefits to discrete diffusion models. For instance, by selecting appropriate initial data, we can generate results that are pertinent to the initial data to attain controllable generation. Additionally, due to the excellent property of sampling certainty reduction, we can perform operations such as interpolation in latent code to complete data editing.

---

Next, we will provide detailed responses to each reviewer's comments.

Meta-review: This paper presents a theoretical work that defines the probability flow for continuous-time diffusion models with discrete state-space
All reviewers agree that the paper presents some  novelties and a relevant analysis of diffusion processes from the
perspective of OT and probability flow. They all agree that the paper deserves to be published.
Decision: Accept (poster)