PaperlayerView exact revision on arXiv

Source-linked AI summary

Formulating Discrete Probability Flow Through Optimal Transport

Pengze Zhang, Hubery Yin, Chen Li, Xiaohua Xie

arXiv:2311.03886v1cs.LGcs.AImath.OC

TL;DR

Discrete diffusion models lack an established deterministic probability flow, and a deterministic transport map may not exist for discrete distributions. This paper develops an optimal-transport formulation for discrete probability flow and reports reduced sampling uncertainty on synthetic toy and CIFAR-10 datasets.

  • Problem

    Discrete diffusion models have received less attention regarding sampling certainty, and existing literature has not addressed their probability flow.

    Source: §1 · pp. 1–2

  • Method

    The paper proves a conditional Monge-map result for continuous probability flow and modifies a discrete diffusion model to produce a Kantorovich plan that defines discrete probability flow.

    Source: §1 · p. 2, §4 · p. 6, §4.1 · p. 7

  • Results

    The proposed sampling method significantly reduces sampling uncertainty, with effectiveness reported on synthetic toy and CIFAR-10 datasets.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    An optimal-transport formulation provides a discrete probability-flow definition even when a deterministic map between discrete distributions may not exist.

    Source: §1 · p. 1, §4 · p. 6, §4.1 · p. 7

    The OU–Brownian equivalence result assumes θ_t ≥ 0, σ_t > 0, and independent driving Brownian motions.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

Continuous diffusion models are commonly acknowledged to display a deterministic probability flow, whereas discrete diffusion models do not. In this paper, we aim to establish the fundamental theory for the probability flow of discrete diffusion models. Specifically, we first prove that the continuous probability flow is the Monge optimal transport map under certain conditions, and also present an equivalent evidence for discrete cases. In view of these findings, we are then able to define the discrete probability flow in line with the principles of optimal transport. Finally, drawing upon our newly established definitions, we propose a novel sampling method that surpasses previous discrete diffusion models in its ability to generate more certain outcomes. Extensive experiments on the synthetic toy dataset and the CIFAR-10 dataset have validated the effectiveness of our proposed discrete probability flow. Code is released at: https://github.com/PangzeCheung/Discrete-Probability-Flow.

1 Introduction

Continuous diffusion models provide a deterministic probability flow, but discrete diffusion models lack an established counterpart and deterministic transport may not exist for discrete distributions. This paper develops an optimal-transport foundation for discrete probability flow and proposes a sampling method aimed at greater certainty.

1 Introduction

  • Discrete diffusion research has focused mainly on sampling quality and efficiency, while sampling certainty and probability flow have received little attention.For discrete distributions, a deterministic map between distributions may not exist.Source: §1 · pp. 1–2
  • The paper proves that, under certain conditions, continuous probability flow coincides with the Monge optimal transport map over finite time intervals within (0, ∞).It also introduces a discrete probability flow under optimal transport and identifies properties shared by the continuous and discrete formulations.Source: §1 · p. 2
  • The proposed sampling method significantly improves sampling certainty on a synthetic toy dataset and CIFAR-10.The authors report experiments on both datasets as evidence of effectiveness.Source: §1 · p. 2

2 Background on Diffusion Models and Optimal Transport

This section reviews diffusion processes, optimal transport, and gradient-flow concepts used in the paper. It distinguishes continuous and discrete diffusion formulations and the transport frameworks relevant to them.

2.1 Continuous state diffusion models

  • Continuous diffusion models use a forward process to transform data into a tractable reference distribution and a reverse process to generate samples from that reference.The forward process is described as a time-dependent Ornstein–Uhlenbeck process, and its reverse is a reverse-time process.Source: §2.1 · p. 2
  • The continuous probability flow is a deterministic process sharing the diffusion process’s single-time marginals, and its flow map is a diffeomorphism for t ≥ s > 0.Training minimizes discrepancy between the model output and the Stein score, equivalently matching the conditional score function.Source: §2.1 · p. 2

2.2 Discrete state diffusion models

  • Discrete diffusion models include DTDS models based on discrete-time Markov chains and CTDS models based on continuous-time Markov chains.Reverse-generator factorization permits parallel dimension-wise sampling only when each time step is small.Source: §2.2 · p. 3

2.3 Optimal transport

  • The Monge formulation transports all mass at each point to one destination, whereas the Kantorovich formulation allows more general transport plans.The formulations are equivalent for measures absolutely continuous with respect to Lebesgue measure, but can differ for discrete measures because the Monge constraint may not be satisfied.Source: §2.3 · p. 3

2.4 Fokker-Planck equation by gradient flow

  • The Fokker–Planck equation is presented as a gradient flow, with Brownian motion’s equation corresponding to the negative Gibbs–Boltzmann entropy under the 2-Wasserstein metric.A discrete analogue uses discrete Gibbs–Boltzmann entropy and a newly defined metric on a graph.Source: §2.4 · p. 4

3 Continuous probability flow

The paper relates time-dependent Ornstein–Uhlenbeck processes to Brownian motion and establishes conditions under which continuous probability-flow maps are Monge optimal transports. The results are bounded by assumptions on the initial distribution and time range.

3.1 The equivalence of Ornstein-Uhlenbeck processes and Brownian motion

  • The paper focuses its analysis on standard Brownian motion after showing that a time-dependent Ornstein–Uhlenbeck process can be deterministically converted to it.The proposition’s setup specifies θ_t ≥ 0, σ_t > 0, and independent Brownian motions.Source: §3.1 · p. 4

3.2 Probability flow is a Monge map

  • For Brownian motion with the specified initial distribution and all initial sample locations on one line, the probability-flow map is Monge optimal between pB(x, t) and pB(x, t + s) for t > 0 and s ≥ 0.The paper also states that the optimal-transport conclusion applies only to its specified cost function.Source: §3.2 · pp. 4–5
  • The paper distinguishes its result from a counterexample involving the limit at infinite time and an initial distribution outside the specified form.Those differences are why the cited counterexample does not apply to the paper’s stated setting.Source: §3.2 · p. 5
  • An infinitesimal-transport flow initialized at the probability-flow state is Monge optimal for sufficiently small positive time increments.For each t > 0, the stated result guarantees a δ_t such that the map is optimal when 0 < s < δ_t.Source: §3.2 · p. 5
  • The corresponding Ornstein–Uhlenbeck probability-flow map is Monge optimal under the stated initial condition and collinearity of the initial locations.This result holds for all t > 0 and s ≥ 0.Source: §3.2 · p. 5

4 Discrete probability flow

The section defines a discrete probability flow using optimal transport, then introduces a sampling method intended to reduce uncertainty while retaining the discrete diffusion model’s marginal distributions.

4 Discrete probability flow

  • The proposed discrete probability flow uses optimal transport to address the absence of a deterministic probability flow for discrete diffusions.For discrete distributions, a deterministic map from one distribution to the next may not exist.Source: §4 · p. 6

4.1 Constructing discrete probability flow

  • Because QD can include mutual flows between states, the authors modify the forward process to obtain an optimal transport plan and define QD’s process as the discrete probability flow.The modified process is described as a Kantorovich plan between the time marginals.Source: §4.1 · pp. 6–7
  • The process generated by QD has the same single-time marginals as Q, and for sufficiently small time increments Q provides an optimal transport map under cost dD.The map applies from PD(t) to PD(t + s), for 0 < s < δt.Source: §4.1 · p. 6

4.2 Sampling by discrete probability flow

  • The modified model uses the same training process as SDDM, while reverse sampling uses Euler steps and sequential draws from the reverse conditional distribution.The training method directly learns conditional probabilities; the equality of the relevant marginals makes training identical to [50].Source: §4.2 · p. 7
  • DPF reduces conditional standard deviation, and its reverse transitions occur only toward higher-probability states, unlike the original reverse processes.The paper uses conditional standard deviation to measure sampling certainty.Source: §4.2 · p. 7

5 Related Work

Related work connects probability flow to deterministic alternatives to diffusion and to optimal transport, while discrete-diffusion research has largely emphasized forward-process design. The section also situates the paper among continuous-time discrete-state models.

5 Related Work

  • Probability flow was introduced as a deterministic alternative to Itô diffusion [49], and DDIM was shown equivalent to it [47].Work on optimal transport proposed a Monge-map interpretation, but a later result found an initial distribution for which probability flow is not optimal [30].Source: §5 · p. 8
  • Discrete diffusion began with a binary model [46], and much subsequent work focused on designing the forward process [25] [24] [4] [27] [17] [10].Source: §5 · p. 8
  • Continuous-time discrete-state models were introduced by [7] and subsequently developed by [50].Source: §5 · p. 8

6 Experiments

Experiments compare DPF with SDDM on synthetic data and with τLDR-0 on CIFAR-10. The reported results show comparable binary generation quality and reduced sampling uncertainty, including for images drawn from the same initial point.

6 Experiments

  • DPF’s MMD is slightly higher than SDDM’s, but the reported difference is minimal and does not significantly affect generated-sample quality.The authors attribute the increase potentially to approximation of an additional term in the reverse generator.Source: §6 · pp. 8–9
  • Synthetic tests use state sizes 2, 5, and 10, with 40,000 binary samples and 4,000 samples for other state sizes.The experiments use the same setup as SDDM except for replacing generator Q with Equation (26).Source: §6 · p. 8
  • DPF produces binary samples comparable in quality to SDDM and reduces sampling uncertainty on synthetic data and same-start CIFAR-10 image generation.The binary quality comparison is visualized in Figure 1, and the CIFAR-10 comparison in Figure 4.Source: §6 · pp. 9–10, Figure 1 · p. 9, Figure 4 · p. 10
  • For different initial points, DPF samples cluster by starting-point color and tend to remain closer to their respective initial points, unlike the disorganized SDDM samples.The paper reports similar behavior across different state sizes.Source: §6 · pp. 9–10, Figure 3 · p. 9

7 Discussion

Discrete diffusion sampling remains stochastic, motivating a probability flow grounded in optimal transport. The paper connects this construction to continuous probability flow and identifies extensions still open for study.

7 Discussion

  • The discussion summarizes the continuous Monge-map result, the discrete probability-flow construction, shared properties, and a sampling method that reduces uncertainty.Open directions include general initial conditions, extending a small-time Kantorovich-plan existence result globally, and applying the definition to broader discrete diffusion models.Source: §7 · p. 10

A Overview of our DPF

  • Continuous probability flow makes sampling from an initial point deterministic, addressing the stochasticity that can complicate defining an encoder to latent space.The overview notes that stochastic samples from one initial point may span the data space, while probability flow follows a deterministic path.Source: A Overview of our DPF · p. 16
  • Figure 5 contrasts continuous and discrete diffusion sampling and illustrates the proposed optimal-transport-based discrete probability flow.The overview describes stochastic trajectories and a deterministic probability-flow path for continuous models, then proposes a discrete counterpart to reduce uncertainty.Source: A Overview of our DPF · pp. 15–16, Figure 5 · p. 16

B Definitions and Theorems Employed in this Appendix

The appendix collects mathematical and optimization results used in the paper, then applies a time-change argument to relate the Ornstein–Uhlenbeck process to other diffusion formulations.

B Definitions and Theorems Employed in this Appendix

  • The appendix states that its definitions and theorems are presented for the case within R^n and includes a quadratic-cost optimal-transport result for a gradient map.The cited result assumes finite second moments and a convex potential differentiable almost everywhere.Source: B Definitions and Theorems Employed in… · p. 16
  • LICQ is defined through linear independence of active constraint gradients, and the appendix states that KKT conditions are sufficient for linear programming.The KKT conditions are presented as necessary under LICQ for a local solution; the linear-programming case is also sufficient.Source: B Definitions and Theorems Employed in… · p. 17

C.1 Proof of Proposition 1

  • The Proposition 1 proof uses a time-change result to show equality in law for the process, and its remark treats VE, VP, sub-VP, and cited models as equivalent.The remark characterizes the Ornstein–Uhlenbeck process as a scaling of Brownian motion with a change in time.Source: §C.1 · p. 18

C.2 Proof of Proposition 2

These appendix proofs establish optimal-transport properties for continuous and discrete processes, followed by the reverse-sampling procedure and synthetic-data setup used for validation.

C.2 Proof of Proposition 2; C.3 Proof of Proposition 3

  • The map Ŷ_T(Ŷ_t) is optimal for quadratic transport cost; for small s, Ỹ_{t+s}(Ỹ_t) is also optimal, and particles travel at a uniform rate.Source: §C.2 · p. 19, §C.3 · pp. 19–20

C.4 Proof of Proposition 4

  • For the Ornstein–Uhlenbeck process, the probability-flow path and comparison process solve the same ODE with identical initial conditions, yielding a quadratic-cost optimal map.Source: §C.4 · p. 20

C.5 Proof of Proposition 5

  • The modified discrete process has the same time-dependent state probabilities as the original process, as both solve the same forward equation.Source: §C.5 · pp. 20–21

C.6 Proof of Proposition 6

  • When adjacent-state flux directions remain fixed over the time interval, the constructed transport plan satisfies the optimal-transport conditions and solves the problem.Source: §C.6 · pp. 21, 23

C.7 Proof of Proposition 7

D.1 Algorithm

  • The reverse sampler computes transitions for staying, adding, or subtracting a state, then iterates until t = 0.Source: §D.1 · p. 24

D.2 Synthetic Dataset; D.3 Experiment Details

  • Synthetic validation uses seven distributions encoded at state sizes 2, 5, and 10; training uses a three-layer, 256-channel MLP for 300,000 iterations.The state-size-5 and state-size-10 datasets use 8-bit 5-base and 6-bit decimal codes, respectively.Source: §D.2 · pp. 24–25, §D.3 · p. 25

D.4 Quality of Generated Samples

On synthetic data, DPF’s MMD is slightly higher than SDDM’s, while the reported sample visualizations show comparable generation quality. The authors attribute the slight quality difference to neural-network approximation in the transition-rate calculation.

D.4 Quality of Generated Samples

  • DPF’s MMD is slightly higher than SDDM’s, but the reported gap is small and the generated samples have comparable quality.The authors report that the gap is not enough to affect generation quality, with Figures 1, 6, and 7 supporting the visual comparison.Source: §D.4 · pp. 25–26, Figure 10 · p. 26, Figure 11 · p. 26
  • The evaluation compares generated samples with true synthetic data using Laplace MMD, averaging results across 10 runs.The experiments generate 40,000 binary samples and 4,000 samples for other synthetic data types.Source: §D.4 · p. 25
  • The authors attribute the slightly inferior quality to neural-network approximation of quantities used to evaluate both terms in their transition-rate calculation.They state that neural-network fitting error leaves no feasible alternative approximation for achieving a superior outcome.Source: §D.4 · p. 25

D.5 Standard Deviation of Generated Samples

Across synthetic and CIFAR-10 experiments, DPF reduces sampling uncertainty, with generated samples remaining more concentrated around their initial points. Additional tests report shorter sampling trajectories and higher transport efficiency than SDDM.

D.5 Standard Deviation of Generated Samples

  • On checkerboard (S = 5), DPF’s CSD is 1.4103, 89% lower than SDDM.The experiments fix the same initial point and evaluate the Expectation of the Conditional Standard Deviation.Source: §D.5 · pp. 26–27

D.6 Generated Samples from Different Initial Points

  • Across initial points, DPF samples cluster by their originating point, unlike SDDM samples, which appear mixed together.The visualization uses a 20 × 20 grid of initial points and 10 samples per point.Source: §D.6 · p. 27

D.7 Distance Between the Generated Samples and Initial Points

  • DPF reduces the average L1 distance from generated samples to their corresponding initial points, with outcomes tending to concentrate near high-probability states.The reported pattern is consistent with the paper’s optimal-transport design.Source: §D.7 · p. 27

D.8 Sampling Trajectory Length

  • On checkerboard (S = 5), DPF’s trajectory-length score is 0.7640, 97% lower than SDDM.The paper interprets the lower trajectory consumption as consistent with its optimal-transport design.Source: Table 5 · p. 27, §D.8 · p. 28

D.9 Transport Efficiency

  • SDDM’s transport efficiency is approximately 50%, while DPF’s is close to 100%, indicating that most DPF trajectory steps contribute to the final transition.The paper attributes SDDM’s lower efficiency to duplicate trajectories.Source: Table 6 · p. 28, §D.9 · p. 29

D.10 Visualization of Sampling Trajectory

  • SDDM often repeats sampling trajectories, whereas DPF moves toward high-probability states and avoids duplicate trajectories.Figure 12 visualizes the duplicated SDDM trajectories in the highlighted region.Source: Figure 12 · p. 28, §D.10 · p. 29

D.11 Higher dimension or state scenarios

  • DPF significantly reduces sampling uncertainty in the tested larger settings, including S = 50, K = 20 and S = 5, K = 50.These configurations were selected to avoid dimension redundancy in the toy dataset’s coordinate encoding.Source: §D.11 · p. 29, Table 7 · p. 29

D.12 Image Modeling

  • On CIFAR-10, DPF significantly reduces sample uncertainty versus τLDR-0, producing almost the same images from a given initial point.The comparison uses class-std and class-entropy on 1,000 initial points with 10 generated images each.Source: §D.12 · p. 30, Table 8 · p. 30

E.1 Narrow time interval limited in Proposition 6.

The discrete optimal-transport plan is justified only over narrow time intervals, and the Brownian probability-flow analysis is restricted to finite intervals. The paper also identifies possible applications of reduced sampling uncertainty.

E.1 Narrow time interval limited in Proposition 6.

  • Sign changes in P_i(t)−P_j(t) can make a discrete plan suboptimal over a finite interval; the S=3 example exhibits such a crossing, limiting Proposition 6 to narrow time intervals.The authors state that the optimal plan cannot be extended across the entire interval.Source: §E.1 · p. 30

E.2 Definition of probability flow on universal discrete process.

  • Because discrete stochastic terms have few assumptions and can absorb drift, the authors note that separating drift and stochastic terms may reveal distinctive properties.Source: §E.2 · p. 30

E.3 Practical applications.

  • The authors identify controllable generation from appropriate initial data and latent-code interpolation for data editing as possible uses of reduced sampling uncertainty.Source: §E.3 · pp. 30–31

E.4 Infinite horizon case in Proposition 2.

  • The analysis excludes infinite horizons and neighborhoods of t=0: Brownian distributions are singular at these boundaries, and non-Lipschitz ODE dynamics yield non-unique solutions.The study is consequently limited to finite intervals.Source: §E.4 · p. 31

Related papers

Loading 2311.03886v1…