Formulating Discrete Probability Flow Through Optimal Transport

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.

Paper

References (58)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer WcKm7/10 · confidence 4/52023-07-04

Summary

The paper answers two important questions for the diffusion modelling community: (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. (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. The low stochasticity sampler is interpreted on 2d datasets.

Strengths

I like this paper as it very satisfyingly answers the above two important research questions in a very neat way using optimal transport. Regarding 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. Regarding 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. The 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. 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. The strength of the theoretical results stands alone. Note that I have not the expertise in this area to check the proofs. 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. The presentation of the paper is poor in places although the overall flow of the paper is good. In eq(16) you should define what a neighborhood is properly. On line 178 you should define what you mean by mutual flow. Typos: eq (2) dt is in the brackets as well for drift line 72 x_{t+t} should be x_{t+1} again line 74 k+t should be k+1 eq 17 extra bracket eq (26) s=0 line 174 I assume you mean the sum is to K not S? Edit after rebuttal: I have read the author response and the new experiments alleviated my main concern and I have raised my score to 7.

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. 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. 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

It is unclear how the method scales to high dimensional / large number of state size datasets.

Reviewer v2T57/10 · confidence 3/52023-07-05

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.

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. To 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.

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. It 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. **Experiments** The 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. I 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? **Miscellaneous** To 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. Further, some **typos**: - In Table 1, given that the MMD value for the SDDM simulation on the checkerboard dataset has the wrong sign. - In line 72: $(x_{t+t})$. - In Eq. 19: missing a ½ factor? - In the Discussion section.

Questions

What are possible practical applications of the proposed sampling method? What are possible alternative approximations of the quantity $Q_t$, which is likely the cause of the inferior quality of reconstructed marginals? Why do points in Figure 2 appear to form rectangles, e.g., in the 2spirals, pinwheel, and swissroll datasets? Even 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

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.

Reviewer wEPe7/10 · confidence 2/52023-07-06

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: - 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]. - 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. - 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. [1] Understanding ddpm latent codes through optimal transport, Khrulkov et al., 2022. [2] Score-based Continuous-time Discrete Diffusion Models, Sun et al., 2022.

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. - Paper very well-written, easy to follow.

Weaknesses

The scalability of the experiments.

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 ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The notion of probability flow is only defined for the noising process corresponding to the Brownian motion.

Reviewer Rnpe8/10 · confidence 4/52023-07-07

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.

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. Detailed comments: - 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? - Line 105: "using a newly defined metric" what is this metric? - Line 107, a bit more detail about the neighborhood and eq 16 seems helpful. - Line 121: "This paper" --> "We" would be clearer. (Also "This paper" in other places.) - 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.

Questions

None.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

N/A

Reviewer Rnpe2023-08-10

Re: author response

Thanks for the response and clarifications! I'll highlight that several of my confusions arise from use of the referent "this", which was ambiguous (to me). I would suggest going through the paper and considering whether uses of "this" could be replaced with a more explicit word/phrase. Nice paper, thanks!

Authorsrebuttal2023-08-12

Thank you very much for the reply

We are extremely thankful for your appreciation of our work! We greatly appreciate the suggestions in the comments, and we will keep striving to perfect the writing in our paper.

Reviewer v2T52023-08-10

Thank You for the Rebuttal

I am still at unease about the choice of experiments and feel that at least some empirical evaluation of the method is necessary. The current selection of toy data represents settings where discrete probability flow are not even necessary, and are of a complexity that does not substantiate the paper's claims. I am not sure I understand the point of the argumentation " complete code of SDDM on synthetic data is publicly available for easy reproduction" in the General Response. What speaks against evaluating it at least against some of the "real-world" settings explored in Sun et al. (2023) [47]?

Authorsrebuttal2023-08-12

Thank you very much for the reply

Thank you very much for the reply! Since this paper has a sufficient theoretical derivation of the proposed discrete diffusion flow, and its sampling properties have been clearly verified on the synthetic dataset, we apologize for previously neglecting the experiments in the application domain. We greatly appreciate your suggestion and have already started the experiments on the real-world dataset. However, since the code of the real-world dataset needs to be reproduced by ourselves and requires training, it indeed takes some time. We will supplement this experiment in the camera-ready version and discuss the results.

Authorsrebuttal2023-08-15

Thank you for waiting

Thank you for waiting. We have validated our DPF on the CIFAR-10 dataset, using the pre-trained discrete diffusion model provided by the paper [1]. Firstly, we selected 1,000 initial points, and sampled 10 images from each initial point. To measure the sampling certainty on the image data, we used a pre-trained CIFAR-10 classifier to classify the image, and introduce two new metrics, i.e., class-std and class-entropy. The class-std calculates the standard deviation of the categories of the images sampled from the same initial point. While the class-entropy calculates the entropy of the category distribution of the images sampled from the same initial point. Lower class-std and class-entropy indicate better sampling certainty. The experimental results, shown in Table 3*, demonstrate that our method can significantly reduce the sample uncertainty compared to the $\tau$LDR-0 method [1]. Additionally, we visualized the sampled images, and it was clear that from an initial point, our method samples almost the same images, while the original sampling method obtains totally different images. Thank you again for your suggestion, we will supplement these results and discuss them in the camera-ready. *Table 3\* Comparison of certainty for $\tau$LDR-0 and DPF on the Cifar-10 dataset. Here, $ECV$, class-std, and class-entropy are calculated on 1,000 initial points, each of which has 10 generated images. Lower values indicate superior certainty.* | | $ECV$ | class-std | class-entropy| | :----: | :----: | :----: | :----: | | **$\tau$LDR-0** [1] | 57.6898 | 2.6628 | 1.7703 | | **DPF (ours)** | 9.4420 | 1.1819 | 0.5291 | --- **Reference** [1] Campbell, A. et al.. A continuous time framework for discrete denoising models. NeurIPS, 2022.

Reviewer v2T52023-08-17

Updated Score

Thank you for the continued efforts and running these experiments. I have updated my score accordingly.

Reviewer WcKm2023-08-11

Thank you for the response. I like the new experiment scaling to high dimensionality and showing the minimized variance in sampling effect is still preserved. This has alleviated my main worry and I will increase my score accordingly.

Authorsrebuttal2023-08-12

Thank you very much for the reply

We are extremely thankful for your recognition of our work, as well as the time and effort you have dedicated to our paper.

Reviewer wEPe2023-08-12

Thank you for the response. I find the new experiments with higher dimensionality pretty convincing, and I think that it definitely makes the contribution stronger on the experimental side (which was the main weakness for me). Besides this, my theoretical question was well answered and I suggest the authors to include this explanation in the paper. Overall, it is very nice !

Authorsrebuttal2023-08-12

Thank you very much for the reply

We are thankful for your appreciation of our work! We will carefully review our paper and incorporate your valuable feedback into the final version of our paper to make it even stronger.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC