SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pengze Zhang *
zhangpz3@mail2.edu.cn
Hubery Yin*
WeChat, Tencent Inc.
hubery@tencent.com
Chen Li
chaselli@tencent.com
Xiaohua Xie
xiexiaoh6@mail.edu.cn
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.
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 = \frac{1}{2}, P_1 = \frac{1}{2})$  and  $(P_0 = \frac{1}{3}, P_1 = \frac{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
*Equal contribution. This work was done when Pengze Zhang was an intern at WeChat.
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,\infty)$ . 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  $\hat{X}_t$  and  $\tilde{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:
$$
d X _ {t} = - \theta_ {t} X _ {t} d t + \sigma_ {t} d B _ {t}, \tag {1}
$$
where  $\theta_t \geq 0, \sigma_t > 0, \forall t \geq 0$  and  $B_t$  is the Brownian Motion (BM). The backward process is the reverse-time process of the forward process [2]:
$$
d X _ {t} = \left[ - \theta_ {t} X _ {t} - \sigma_ {t} ^ {2} \nabla_ {X _ {t}} \log p (X _ {t}, t) \right] d t + \sigma_ {t} d \tilde {B} _ {t}, \tag {2}
$$
where  $\tilde{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]:
$$
\frac {\partial}{\partial t} p (x, t) = \theta_ {t} \nabla_ {x} (x p (x, t)) + \frac {1}{2} \sigma_ {t} ^ {2} \Delta_ {x} p (x, t). \tag {3}
$$
In order to train a diffusion model, the primary objective is to minimize the discrepancy between the model output  $s_{\theta}(x_t,t)$  and the Stein score function  $s(x_{t},t) = \nabla_{x_{t}}\log p(x_{t},t)$  [25]. Song et al. [45] demonstrate that, it is equivalent to match  $s_{\theta}(x_t,t)$  with the conditional score function:
$$
\theta^ {*} = \underset {\theta} {\arg \min } \mathbb {E} _ {t} \left\{\lambda_ {t} \mathbb {E} _ {x _ {0}, x _ {t}} \left[ \| s _ {\theta} (x _ {t}, t) - \nabla_ {x _ {t}} \log p (x _ {t}, t | x _ {0}, 0) \| ^ {2} \right] \right\}, \tag {4}
$$
where  $\lambda_{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):
$$
d \hat {X} _ {t} = \left[ - \theta_ {t} \hat {X} _ {t} - \frac {1}{2} \sigma_ {t} ^ {2} s (\hat {X} _ {t}, t) \right] d t. \tag {5}
$$
In accordance with the global version of Picard-Lindelöf theorem [1] and the adjoint method[36, 7], the map
$$
T _ {s, t}: \mathbb {R} ^ {n} \longrightarrow \mathbb {R} ^ {n}, \tag {6}
$$
$$
\hat {X} _ {s} \longmapsto \hat {X} _ {t}.
$$
is a diffeomorphism  $\forall t\geq 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_{t + 1}) = \frac{q_{t + 1|t}(x_{t + 1}|x_t)q_t(x_t)}{q_{t + 1}(x_{t + 1})}$ . In practical applications, it is preferable to parameterize this model using  $p_{0|t + 1}^{\theta}$  [24, 4] and obtain  $p_{k|k + 1}^{\theta}$  through
$$
\begin{array}{l} p _ {k | k + 1} ^ {\theta} (x _ {k} | x _ {k + 1}) = \sum_ {x _ {0}} q _ {k | k + 1, 0} (x _ {k} | x _ {k + 1}, x _ {0}) p _ {0 | k + 1} ^ {\theta} (x _ {0} | x _ {k + 1}) \\ = \sum_ {x _ {0}} q _ {k + 1 | k} \left(x _ {k + 1} \mid x _ {k}\right) \frac {q _ {k \mid 0} \left(x _ {k} \mid x _ {0}\right)}{q _ {k + 1 \mid 0} \left(x _ {k + 1} \mid x _ {0}\right)} p _ {0 \mid k + 1} ^ {\theta} \left(x _ {0} \mid x _ {k + 1}\right). \tag {7} \\ \end{array}
$$
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:
$$
\frac {\partial}{\partial t} q _ {t | s} \left(x _ {t} \mid x _ {s}\right) = \sum_ {y} q _ {t | s} \left(y \mid x _ {s}\right) Q _ {t} \left(y, x _ {t}\right). \tag {8}
$$
The reverse process is:
$$
\frac {\partial}{\partial s} q _ {s \mid t} \left(x _ {s} \mid x _ {t}\right) = \sum_ {y} q _ {s \mid t} \left(y \mid x _ {t}\right) R _ {t} \left(y, x _ {s}\right). \tag {9}
$$
The generator of the reverse process can be written by [6, 47]:
$$
R _ {t} (y, x) = \frac {q _ {t} (x)}{q _ {t} (y)} Q _ {t} (x, y) = \sum_ {y _ {0}} \frac {q _ {t | 0} (x | y _ {0})}{q _ {t | 0} (y | y _ {0})} q _ {0 | t} (y _ {0} | y) Q _ {t} (x, y). \tag {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, 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  $\mu$  and  $\nu$  on  $(\mathbb{R}^n,\mathcal{B})$  and a cost function  $c:\mathbb{R}^n\times \mathbb{R}^n\to [0, + \infty ]$ . The Monge problem is
$$
\left(\mathrm {M P}\right) \inf  _ {\mathrm {T}} \left\{\int c (x, \mathrm {T} (x)) \mathrm {d} \mu (x): \mathrm {T} _ {\#} \mu = \nu \right\}. \tag {11}
$$
The measure  $\mathrm{T}_{\#}\mu$  is defined through  $\mathrm{T}_{\#}\mu(A) = \mu(\mathrm{T}^{-1}(A))$  for every  $A \in \mathcal{B}$  and is called the pushforward of  $\mu$  through  $\mathrm{T}$ .
It is evident that the Monge Problem (MP) transports the entire mass from a particular point, denoted as  $x$ , to a single point  $\mathrm{T}(x)$ . In contrast, Kantorovich provided a more general formulation, referred to as the Kantorovich problem:
$$
\left(\mathrm {K P}\right) \inf  _ {\gamma} \left\{\int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} c \mathrm {d} \gamma : \gamma \in \Pi (\mu , \nu) \right\}, \tag {12}
$$
where  $\varPi(\mu,\nu)$  is the set of transport plans, i.e.,
$$
\varPi (\mu , \nu) = \left\{\gamma \in \mathcal {P} \left(\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}\right): \left(\pi_ {x}\right) _ {\#} \gamma = \mu , \left(\pi_ {y}\right) _ {\#} \gamma = \nu \right\}, \tag {13}
$$
where  $\pi_x$  and  $\pi_y$  are the two projections of  $\mathbb{R}^n\times \mathbb{R}^n$  onto  $\mathbb{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.
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:
$$
\frac {\partial}{\partial t} p (x, t) = \frac {1}{2} \Delta p (x, t), \tag {14}
$$
and it represents the gradient flow of the Gibbs-Boltzmann entropy multiplied by  $-\frac{1}{2}$ :
$$
- \frac {1}{2} S (p) = \frac {1}{2} \int_ {\mathbb {R} ^ {n}} p (x) \log p (x) \mathrm {d} x. \tag {15}
$$
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:
$$
S (p) = \sum_ {i} p _ {i} \log p _ {i}, \tag {16}
$$
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, \dots, 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:
$$
\frac {d}{d t} p _ {i} = \sum_ {j \in N (i)} p _ {j} - p _ {i}, \tag {17}
$$
where  $N(i) = \{j\in \{1,2,\dots,N\} |\{a_i,a_j\} \in 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 Ornstein-Uhlenbeck 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} = -\theta_{t}X_{t}dt + \sigma_{t}dB_{t}^{(1)}$ ,  $dY_{t} = dB_{t}^{(2)}$ , where  $B_{t}^{(1)}$  and  $B_{t}^{(2)}$  are two independent Brownian motions and  $\theta_t \geq 0, \sigma_t > 0, \forall t \geq 0$ . Let  $\phi_t = \exp(\int_0^t \theta_\tau \, \mathrm{d}\tau)$ ,  $\beta_t = \int_0^t (\sigma_\tau \phi_\tau)^2 \, \mathrm{d}\tau$ . Then  $X_{t}$  coincides in law with  $\phi_t^{-1}Y_{\beta_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) = \frac{1}{2} |x - y|^2$ .
Within the context of generative models, a collection of training samples denoted as  $\{x_{i}\}_{i = 1}^{N}$  is typically provided, and these samples are intrinsically defined by a distribution:
$$
p (x, 0) = \frac {1}{N} \sum_ {i = 1} ^ {N} \delta \left(x - x _ {i}\right), \tag {18}
$$
where  $\delta(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]
$$
p _ {B} (x, t) = \frac {1}{N} \sum_ {i = 1} ^ {N} (2 \pi t) ^ {- \frac {n}{2}} \exp \left(- \frac {\left| x - x _ {i} \right| ^ {2}}{2 t}\right). \tag {19}
$$
The probability flow is defined as [46]:
$$
d \hat {Y} _ {t} = - \frac {1}{2} \nabla_ {\hat {Y} _ {t}} \log p _ {B} (\hat {Y} _ {t}, t) d t. \tag {20}
$$
According to [1, 36, 7], the solution exists for all  $t > 0$  and the map  $\hat{Y}_{t + s}(\hat{Y}_t)$  is a diffeomorphism for all  $t > 0, s \geq 0$ . We have discovered that  $\hat{Y}_{t + s}(\hat{Y}_t)$  is the Monge map under some conditions and the time does not reach 0 or  $+\infty$ .
Proposition 2. Given that  $Y_0$  follows the initial condition (18), and all  $x_i$  lie on the same line, the diffeomorphism  $\hat{Y}_{t + s}(\hat{Y}_t)$  is the Monge optimal transport map between  $p_B(x,t)$  and  $p_B(x,t + s)$ ,  $\forall t > 0, s \geq 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  $\hat{Y}_{+\infty}(\hat{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 Gibbs-Boltzmann entropy concerning the  $W_{2}$  metric [27]. As  $W_{2}$  is associated with optimal transport, it is reasonable to anticipate that the "infinitesimal transport"  $\hat{Y}_{t + dt}(\hat{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_c^2 (\mathbb{R}^n)$  denote the set of twice continuously differentiable functions on  $\mathbb{R}^n$  with compact support. The generator  $A_{t}$  is defined as follows [35]:
$$
\hat {A} _ {t} f = \lim  _ {\Delta t \rightarrow 0 ^ {+}} \frac {f \left(\hat {Y} _ {t + \Delta t}\right) - f \left(\hat {Y} _ {t}\right)}{\Delta t}, \forall f \in C _ {c} ^ {2} \left(\mathbb {R} ^ {n}\right). \tag {21}
$$
It is straightforward to verify that
$$
\hat {A} _ {t} = - \frac {1}{2} \nabla_ {x} \log p _ {B} (x, t) ^ {T} \nabla_ {x}. \tag {22}
$$
We define the "infinitesimal transport" to be the diffeomorphism  $\tilde{Y}_{t+s}(\tilde{Y}_t)$  where  $\tilde{Y}_{t+s}$  evolves according to the following equation
$$
d \tilde {Y} _ {t + s} = - \frac {1}{2} \nabla_ {\tilde {Y} _ {t}} \log p _ {B} \left(\tilde {Y} _ {t} \left(\tilde {Y} _ {t + s}\right), t\right) d s, \tag {23}
$$
with the initial condition  $\hat{Y}_t = \hat{Y}_t$ . The generator of  $\hat{Y}_{t + s}$  is
$$
\tilde {A} _ {t + s} = - \frac {1}{2} \nabla_ {\tilde {Y} _ {t}} \log p _ {B} \left(\tilde {Y} _ {t} \left(\tilde {Y} _ {t + s}\right), t\right) \nabla_ {x}. \tag {24}
$$
Proposition 3. Given any  $t > 0$ , there exists a  $\delta_t > 0$  s.t.  $\forall 0 < s < \delta_t$ , the diffeomorphism  $\hat{Y}_{t + s}(\hat{Y}_t)$  with the initial condition  $\hat{Y}_t = \hat{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  $\hat{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$  lie on the same line, the diffeomorphism  $\hat{X}_{t + s}(\hat{X}_t)$  is the Monge optimal transport map for all  $t > 0, s \geq 0$ .
4 Discrete probability flow
The continuous probability flow is deterministic, which means the "mass" at  $\hat{Y}_t$  is entirely transported to  $\hat{Y}_{t + s}$  during the time interval  $[t,t + s]$ . However, it is widely acknowledged that for discrete distributions  $\mu$  and  $\nu$ , there may not exist a T such that  $\mathrm{T_{\#}}\mu = \nu$ . 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, \ldots, i_K)$ , where  $i_j \in \{1, 2, \ldots, S\}$ . The Kolmogorov forward equation for this process is
$$
\frac {d}{d t} P _ {j} ^ {i} (t | s) = \sum_ {j ^ {\prime}} P _ {j ^ {\prime}} ^ {i} (t | s) Q _ {D j} ^ {j ^ {\prime}} (t), \tag {25}
$$
where  $P_{j}^{i}(t|s)$  means  $P(x_{t} = j|x_{s} = i)$  and  $Q_{D}$  is defined as:
$$
Q _ {D j} ^ {i} = \left\{ \begin{array}{l l} 1, & d _ {D} (i, j) = 1, \\ - \sum_ {j ^ {\prime} \in \{k: d _ {D} (i, k) = 1 \}} Q _ {D j ^ {\prime}} ^ {i}, & d _ {D} (i, j) = 0, \\ 0, & \text {o t h e r w i s e ,} \end{array} \right. \tag {26}
$$
where  $d_{D}(i,j) = \sum_{l = 1}^{K}|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:
$$
P _ {D i} (t) = \sum_ {j} P _ {0 j} Q _ {D i} ^ {j} (t | 0). \tag {27}
$$
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_{j}^{i} > 0$  and  $Q_{i}^{j} > 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$ :
$$
Q _ {j} ^ {i} (t) = \left\{ \begin{array}{l l} \frac {\operatorname {R e L U} \left(P _ {D _ {i}} (t) - P _ {D _ {j}} (t)\right)}{P _ {D _ {i}} (t)}, & d _ {D} (i, j) = 1, \\ - \sum_ {j ^ {\prime} \in \{k: d _ {D} (i, k) = 1 \}} Q _ {j ^ {\prime}} ^ {i} (t), & d _ {D} (i, j) = 0, \\ 0, & \text {o t h e r w i s e .} \end{array} \right. \tag {28}
$$
$$
R e L U (x) = \left\{ \begin{array}{l l} x, & x > 0, \\ 0, & x \leq 0. \end{array} \right. \tag {29}
$$
In order to avoid singular cases, We define  $Q_{j}^{i}(t)$  to be 0 when  $P_{D_i}(t) = 0$ . In fact, it is easy to verify that  $P_{D_i}(t) > 0$  for all  $t > 0$ ,  $i \in \{1,2,\dots,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  $\forall t > 0$ .
Proposition 6. Given any  $t > 0$ , there exists a  $\delta_t > 0$  s.t.  $\forall 0 < s < \delta_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  $\frac{d}{ds}\hat{P} (t + s) = \hat{P} (t + s)Q(t)$ , also provides an optimal transport map.
Proposition 7. Given any  $t > 0$ , there exists a  $\delta_t > 0$  s.t.  $\forall 0 < s < \delta_t$ , the process above provides an optimal transport map from  $\hat{P}(t)$  to  $\hat{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^{\theta}(i_l(t)|\{i_1,\ldots ,i_{l - 1},i_{l + 1},\ldots ,i_K\} (t))$ . According to proposition 5, it follows that  $P^{\theta} = P^{\theta}_{D}$ , and consequently, the training process is identical to that of [47]. For the sake of brevity, we will employ the notation  $P_{i_i|i\backslash i_l}^{\theta}(t)$  to replace  $P^{\theta}(i_l(t)|\{i_1,\dots ,i_{l - 1},i_{l + 1},\dots ,i_K\} (t))$ .
The generator of the reverse process is
$$
R _ {j} ^ {i} (t) = \left\{ \begin{array}{l l} R e L U \left(\frac {P _ {D j _ {l} | i \backslash i _ {l}} ^ {\theta} (t)}{P _ {D i _ {l} | i \backslash i _ {l}} ^ {\theta} (t)} - 1\right), & d _ {D} (i, j) = 1 \text {a n d} i _ {l} \neq j _ {l}, \\ - \sum_ {j ^ {\prime} \in \{k: d _ {D} (i, k) = 1 \}} R _ {j ^ {\prime}} ^ {i} (t), & d _ {D} (i, j) = 0, \\ 0, & o t h e r w i s e. \end{array} \right. \tag {30}
$$
We use the Euler's method to generate samples. Given the time step length  $\epsilon$ , the transition probabilities for dimension  $l$  is: