L EARNING D IFFUSION B RIDGES ON C ONSTRAINED D OMAINS
A BSTRACT
Diffusion models have achieved promising results on generative learning recently.
However, because diffusion processes are most naturally applied on the unconstrained Euclidean space R d , key challenges arise for developing diffusion based models for learning data on constrained and structured domains.
We present a simple and unified framework to achieve this that can be easily adopted to various types of domains, including product spaces of any type (be it bounded/unbounded, continuous/discrete, categorical/ordinal or their mix).
In our model, the diffusion process is driven by a drift force that is a sum of two terms: one singular force designed by Doob’s h-transform that ensures all outcomes of the process to belong to the desirable domain, and one non-singular neural force field that is trained to make sure the outcome follows the data distribution statistically.
Experiments show that our methods perform superbly on generating tabular data, images, semantic segments and 3D point clouds.
Code is available at https: //github.com/gnobitab/ConstrainedDiffusionBridge .
I NTRODUCTION
Diffusion-based deep generative models, notably score matching with Langevin dynamics (SMLD)
(Song & Ermon, 2019, 2020), denoising diffusion probabilistic models (DDPM)
(Ho et al., 2020), and their variants (e.g., Song et al., 2020b,a; Kong & Ping, 2021; Song et al., 2021; Nichol &
Dhariwal, 2021), have shown to achieve new state of the art results for image synthesis (Dhariwal & Nichol, 2021;
Ramesh et al., 2022; Ho et al., 2022; Liu et al., 2021), audio synthesis (Chen et al., 2020; Kong et al., 2020), point cloud synthesis (Luo & Hu, 2021a,b; Zhou et al., 2021), and many other AI tasks.
These methods train a deep neural network to drive as drift force a diffusion process to generate data, and are shown to outperform competitors, mainly GANs and VAEs, on stability and sample diversity (Xiao et al., 2021; Ho et al., 2020; Song et al., 2020b).
However, due to the continuous nature of diffusion processes, the standard approaches are restricted to generating unconstrained continuous data in R d .
For generating data constrained on special structured domains, such as discrete, bound data or mixes of them, special techniques , e.g., dequantization (Uria et al., 2013; Ho et al., 2019) and multinomial diffusion (Hoogeboom et al., 2021; Austin et al., 2021), need to be developed case by case and the results still tend to be unsatisfying despite promising recent advances (Hoogeboom et al., 2021; Austin et al., 2021).
⇧ ⇤ <latexit sha1_base64="baXQBVQzdlQ1vBlkbVvdqNyT5g8=">AAAB7HicdVDLSgMxFM3UV62vqks3wSKIiyEzHVu7K7pxWcFpC+1YMmmmDc1khiQjlNJvcONCEbd+kDv/xvQhqOiBC4dz7uXee8KUM6UR+rByK6tr6xv5zcLW9s7uXnH/oKmSTBLqk4Qnsh1iRTkT1NdMc9pOJcVxyGkrHF3N/NY9lYol4laPUxrEeCBYxAjWRvK7DXZ31iuWkO26FVSuQGR7VQ+5niG185rjOtCx0RwlsESjV3zv9hOSxVRowrFSHQelOphgqRnhdFroZoqmmIzwgHYMFTimKpjMj53CE6P0YZRIU0LDufp9YoJjpcZxaDpjrIfqtzcT//I6mY4uggkTaaapIItFUcahTuDsc9hnkhLNx4ZgIpm5FZIhlphok0/BhPD1KfyfNF3b8ezyjVeqXy7jyIMjcAxOgQOqoA6uQQP4gAAGHsATeLaE9Wi9WK+L1py1nDkEP2C9fQKcZ46S</latexit> Time 𝑡 ℙ 𝜽
This work proposes a simple and unified framework for learning diffusion models on general constrained domains Ω embedded in the Euclidean space R d .
The idea is to learn a continuous R d -valued diffusion process Z t on time interval t ∈ [0 , T ] , with a carefully designed force field, such that the final state Z T guarantees to 1) fall into the desirable domain Ω , and 2) follows the data distribution asymptotically.
We achieve both steps by leveraging a key tool in stochastic calculus called Doob’s h -transform (Doob, 1984), which provides formula for deriving the diffusion processes whose final states are guaranteed to fall into a specific set or equal a specific value.
Input
Learning : Solve the optimization below with stochastic gradient descent (or other optimizers)
[Equation]
where
[Equation]
with x drawn from the dataset D , ξ ∼ N (0 , I ) , and x 0 any initial point.
Sampling :
Generate sample Z T from
[Equation]
Remark When the domain Ω is a manifold (e.g., line or surface) in R d , simply replace the sum (cid:80) e ∈ Ω with the corresponding line or surface (in general Hausdorff) integration (cid:82) Ω on Ω .
Our simple procedure can be applied to any domain Ω once a properly defined summation (for discrete sets) or integration (for continuous domains) can be evaluated. To give a quick overview on the practical intuition without invoking the mathematical theory, we show in Algorithm 1 a simple instance of the framework when the domain is a discrete set Ω = { e 1 , . .
, e K } .
The idea is to set up the diffusion model to have a form of
[Equation]
where the drift is a sum of a non-singular (e.g., bounded) term f θ ( z, t ) which is a trainable neural force field with parameter θ , and a singular term ∇ z ψ Ω ( z, t ) , which drives Z t towards set Ω as a gradient ascent on ψ Ω ( z, t ) .
The ψ Ω ( z, t ) measures the closeness of z to set Ω , as the log-likelihood of a Gaussian mixture model (GMM) centered on the elements in Ω with variance T − t .
When t approaches to the terminal time T , the variance T − t of the GMM goes to zero, and the magnitude of ∇ z ψ Ω ( z, t ) grows to infinity, hence ensuring that Z T must belong to Ω .
In particular, note that
[Equation]
which increases with an O (1 / ( T − t )) rate as t → T ;
here ω Ω ( e | z, t ) is the softmax probability measuring the relative closeness of z to the elements e in Ω (see also Eq (2)).
As we show in Section 2.3, once f θ is non-singular in the sense of the mild condition of (cid:82) T 0 E [ (cid:13)(cid:13)
f θ ( Z t , t ) (cid:13)(cid:13) 2 ]
d t < + ∞ , the diffusion model in (3) guarantees to yield a final state Z T that belongs to Ω , and hence provides a flexible model family on Ω .
Moreover, as shown in Eq 1 in Algorithm 1, the neural field f θ can be simply trained to approximate ∇ log ω Ω ( e | z, t ) with e plugged as the data point that we expect to achieve when starting from z at time t .
Intuitively, such fitted f θ increases the relative probability of the observed data points and hence allows us to fit the data distribution.
Empirically, diffusion models learned through Ω -bridge achieves favorable results in generating mixed discrete/continuous tabular data, point clouds on grids, categorical semantic segments and discrete CIFAR10 images.
Outline The rest of the paper is organized as follows.
Section 2 introduces h -transform, which allows us to derive bridge processes that are guaranteed to enter specific sets at the terminal time, and Section 2.3 specifies the parametric diffusion models for Ω -bridges.
Then, with the learnable diffusion models, Section 3 introduces the general learning framework along with the loss function.
B ACKGROUND : D IFFUSION P ROCESSES AND h -T RANSFORM
A diffusion process Z = { Z t : t ∈ [0 , T ] } on R d follows a stochastic differential equation of form
[Equation]
(4) where W t is a Wiener process, and σ : [0 , T ] × R d
→ R is a positive diffusion coefficient, and b ∈ [0 , T ] × R d → R d is a drift function.
We use Q (or P ) to denote the path measure of stochastic processes Z , which are probability measures on the space of continuous paths.
Let Q t be the marginal distribution of Z t at time t under Q .
Our framework heavily relies on the bridge processes , special stochastic processes that guarantee to achieve a deterministic value or fall into a given set at the final state T .
☞ For a set Ω ⊆ R d , a process Z in R d with law Q is called an Ω -bridge if Q ( Z T ∈ Ω) = 1 .
One natural approach to constructing bridge processes is to derive the conditioned process of a general unconstrained process given that the desirable bridge constraint happens.
Specifically, assume that Q is the law of a general unconstrained diffusion process of form (4), and denote by
Q Ω ( · ) = Q ( · | Z T ∈ Ω) the conditioned distribution given that the event of Z T ∈ Ω happens.
Then Q Ω is guaranteed to be an Ω -bridge by definition.
Importantly, a remarkable result from Doob (Doob, 1984), now known as h -transform, shows that Q Ω is the law of a diffusion process with a properly modified drift term.
Below, we introduce this results, first for the case x -bridge when Ω = { x } includes a single point, and then for more general sets Ω .
For simplicity, we only state the formula from h -transform that are useful for us without proofs.
See e.g., Oksendal (2013); Rogers & Williams (2000) for more background on h -transform.
2. x -B RIDGES
Let us first consider the x -bridge Q x ( · ) :
= Q ( · | Z T = x ) , the process Q pinned at a deterministic terminal point Z
T = x .
By the result from h -transform (see e.g., Oksendal (2013)), the conditioned process Q x ( · ) :
= Q ( · | Z T = x ) , if it exists, can be shown to be the law of
[Equation]
where q T | t ( x | z ) is the density function of the transition probability
[Equation]
where d x denotes an infinitesimal volume centering around x .
Compared with the diffusion process (4) of Q , the main difference is that the conditioned process has an additional drift force σ 2 ( z, t ) ∇ z log q T | t ( x | z ) which plays the role of steering Z t towards the target Z T =
x ; this is a singular force whose magnitude increases to infinity as t → T , because q T | t ( · | z ) is a delta measure centered at z when t = T .
In addition, by Bayes rule, the distribution of the initial state Z 0 should be given by
[Equation]
Example 2.1.
If Q is the law of d Z t = σ t d W t , we have Q T | t ( ·| z ) =
N ( z, β T − β t ) , where β t = (cid:82) t 0 σ 2 s d s .
Hence, following the formula in (5) , Q x := Q ( ·| Z T = x ) is the law of
[Equation]
and Z 0 ∼ Q 0 | T (d z ) ∝ Q 0 (d z ) ϕ
( x | z, β T − β t ) , and ϕ ( ·| µ, σ 2 ) is the density function of N ( µ, σ 2 ) .
The process in (7) is known as a (time-scaled) Brownian bridge.
Note that the drift in (7) grows to infinity in magnitude with a rate of O (1 / ( β T − β t )) as t → T , which ensures that Z t
= x with probability one.
Arbitrary initialization To make (5) the conditioned process of (4), the initial distribution must follow the Bayes rule in (6).
However, thanks to the singular force ∇ z log q T | t ( x | z ) , the process (5) can guarantee Z xt = x from an arbitrary initialization once the process is well defined.
When the initialization is different from (6), the process in (5) is no longer the conditioned process of (4), but it remains to be an x -bridge in that Z T
= x is still guaranteed.
To see why this is the case, assume that Q is initialized from a deterministic point Z
= x 0 .
Then we would still have Z 0 =
xwhen conditioned on Z T ∈ Ω by Bayes rule.
This suggests that (5) starting from any deterministic initialization is the condition process of Q with the same deterministic initialization, and is hence an x -bridge.
As a result, (5) from any stochastic initialization is also an x -bridge because it can be viewed as the mixture of the processes with different deterministic initialization, all of which are x -bridges.
See Appendix A.4 for a detailed analysis, in which it is shown that (5) with an arbitrary initialization can be viewed as the conditioned process of a special class of non-Markov processes called reciprocal process.
2. Ω -B RIDGES
More generally, for the law Q of (4) and a set Ω ∈ R d , the Ω -bridge Q Ω := Q ( · | Z T ∈ Ω) follows Q
[Equation] [Equation]
(8) with
[Equation]
where drift force η Ω is similar to that of the x -bridge in (5), except that the final state x is now randomly drawn from an Ω -truncated (or Ω -conditioned) transition probability:
[Equation]
which is the transition probability from Z t to Z T , conditioned on that Z T ∈ Ω .
In practice, its form can be derived using Bayes rule.
Example 2.2.
Assume Q follows d Z t = σ t d W t .
Then Q Ω yields the following Ω -bridge:
[Equation]
where N Ω ( z, σ 2 )
=
Law( Z | Z ∈ Ω) with Z ∼ N ( µ, σ 2 ) , which is an Ω -truncated Gaussian distribution
N ( µ, σ 2 ) , whose density function is ϕ Ω ( x ) ∝ I ( x ∈ Ω) ϕ
( x | µ, σ ) with ϕ ( x | µ, σ ) the density function of
[Equation]
Note that it is tractable to calculate η Ω once we can evaluate the expectation of N Ω ( z, β T − β t ) .
A general case is when Ω = I 1 × · · · I d , for which the expectation reduces to one dimensional Gaussian integrals.
See Appendix A.6 and A.7 for details and examples of η Ω .
As in the x -bridge, we can set the initialization to be any distribution supported on the set of points that can reach Ω following Q (precisely, points z 0 that satisfy Ω ∩ supp( Q T ( ·| Z 0 = z 0 ))
̸ = ∅ ) using the mixture of initialization argument.
2. A P ARAMETRIC F AMILY OF Ω -B RIDGES
The formula in (8) only provides a fixed process for a given Q .
For the purpose of learning generative models, however, we need a rich family of Ω -bridges within which we can search for a best one to fit with the data distribution.
It turns out we can achieve this by simply adding an extra non-singular drift force, which can be a trainable neural network, on top of the Ω -bridge in (8).
Specifically, we construct the following parametric diffusion model P θ :
[Equation]
where f θ ( z, t ) is a neural network with input ( z, t ) and parameter θ , which will be trained based on the empirical observations.
Adding the neural drift σ ( Z t , t )
f θ ( Z t , t ) term does not break the Ω -bridge condition, once it satisfies a very mild regularization condition:
Proposition 2.3.
For any Q Ω following d Z t = η Ω ( Z t , t )d t + σ ( Z t , t )d W t that is an Ω -bridge, the P θ in (10) is also an Ω -bridge if E Z ∼ Q Ω [ (cid:82) T 0 (cid:13)(cid:13)
f θ ( Z t , t ) (cid:13)(cid:13) 2 2 d t ] < + ∞ and KL ( Q Ω0 ||
P θ 0 ) < + ∞ .
The condition on f θ is very mild, and it is satisfied if f θ is bounded. Moreover, it can easily hold even when f θ is not bounded.
For example, assuming that f θ ( x ) ≤ a ∥ x ∥ β + b , which holds for
{+ReLU network with β = 1 , we just need to require that the underlying process has a bounded moment E Z ∼ Q Ω [ (cid:82) T 
+}
∥
Z t ∥ 2 β d t ] < + ∞ , which is a typical regularity condition to expect.
L EARNING Ω -B RIDGE M ODELS
Let { x ( i ) } ni =1 be an i.i.d.
sample from an unknown distribution Π ∗ on a domain Ω ⊆ R d .
Our goal is to learn the parameter θ for the Ω -bridge model P θ in (10) such that the terminal distribution
Z T ∼ P θT matches the data X ∼ Π ∗ .
We should distinguish P θ , which is the trainable generative model, and Q , which is a fixed “baseline process” that helps us to derive methods for constructing and learning the model.
Q can be the simple Brownian motion in Example 2.1 and 2.2.
P ✓T <latexit sha1_base64="UqCQhsPGq/gANGjB2hp14D3cFfI=">AAAB/HicdVDLSgMxFM3UV62vapdugkVwNWScUttd0Y3LCn1BW0smTdvQzIPkjlCG+ituXCji1g9x59+YaSuo6IHA4Zx7uSfHi6TQQMiHlVlb39jcym7ndnb39g/yh0ctHcaK8SYLZag6HtVcioA3QYDknUhx6nuSt73pVeq377jSIgwaMIt436fjQIwEo2CkQb7Q8ylMPC+pz297MOFAB41BvkjsKnHLpIqJXaq6rpsSp+wQUsGOTRYoohXqg/x7bxiy2OcBMEm17jokgn5CFQgm+TzXizWPKJvSMe8aGlCf636yCD
[Table]
/AEnrVUe9RetNflaE5b7RyBH9DePgFig5Su</latexit>
Time 𝑡 training
As the case of other diffusion models, P θ can be viewed as a model with an infinite dimensional latent variable of the intermediate trajectories of Z .
Hence, a canonical learning approach is expectation maximization (EM), which alternates between
1) E-step: estimating the posterior P θ,x :=
P θ ( Z | Z T = x ) of the latent trajectories Z given the observation Z T = x ; 2) M-step: estimating the parameter θ with Z imputed from P θ,x .
A key challenge, however, is that the posterior distribution P θ,x is difficult to calculate due to the presence of neural force field in P θ (as the h -transform formula would have no closed form), and it need to be iteratively updated as θ changes.
Following DDPM (Ho et al., 2020), we consider a simpler approach that replaces the posterior P θ,x with an arbitrary x -bridge, denoted by Q x .
This yields a simplified EM algorithm without the expensive posterior inference in the E-step.
A natural choice is the conditioned process
Q x :=
Q ( Z | Z T = x ) , but the method works for a general x -bridge.
Specifically, let Q Π ∗ ( · ) =
(cid:82) Q x ( · )Π ∗ (d x ) be the mixture of the x -bridges whose end point x is randomly drawn from the data distribution x ∼ Π ∗ .
The trajectories from Q Π ∗ can be generated in the following “backward” way: first drawing a data point x ∼ Π ∗ , and then Z ∼ Q x conditioned on the end point x .
Obviously, by construction, the terminal distribution of Q Π ∗ equals Π ∗ , that is,
Q Π ∗ T =
Π ∗ .
Then, the model P θ can be estimated by fitting data drawn from Q Π ∗ using maximum likelihood estimation:
[Equation]
The classical (variational) EM would alternatively update θ (M-step) and Q x (E-step) to make Q x ≈
P θ,x .
Why is it OK to simply drop the E-step?
At the high level, it is the benefit from using universal approximators like deep neural networks: if the model space of P θ is sufficiently rich, by minimizing the KL divergence in (11), P θ can approximate the given Q Π ∗ well enough (in a way that is made precise in the Appendix A) such that their terminal distributions are close: P θT ≈ Q Π ∗ T
= Π ∗ .
☞ Learning latent variable models require no E-step if the model space is sufficiently rich.
We should see that in this case the latent variables Z in the learned model P θ is dictated by the choice of the imputation distribution Q since we have P θ,x = Q x when the KL divergence in (11) is fully minimized to zero; EM also achieves P θ,x = Q x but has the imputation distribution Q x determined by the model P θ , not the other way.
Loss Function
In its general form, the x -bridge Q x that we use can be a non-Markov diffusion process Q x :
[Equation]
which has the same diffusion coefficient σ ( Z t , t ) as P θ in (10), and any x -dependent η x and initialization µ x once the x -bridge condition is ensured.
As the general framework, we assume that η x can depend on the whole trajectory Z .
Input : A dataset D :=
{ x ( i ) } drawn from distribution Π ∗ on a domain Ω . Setup : Specify an x -bridge Q x and an Ω -bridge Q Ω
[Equation]
Specify the generative model P θ based on Q Ω and a neural network f θ :
[Equation]
Default: let Q be the law of d Z t =
σ t d W t and derive the bridges by h -transform as Q x =
Q ( ·| Z T = x ) in Eq (7) and Q Ω = Q ( ·| Z T ∈ Ω) in Eq (9) .
Training :
Estimating θ by minimizing the loss function (13) using any off-the-shelf optimizer.
Sampling :
Generate sample Z T from P θ with the trained parameter θ .
Using Girsanov theorem (e.g., Oksendal, 2013), with P θ in (10) and Q x in (12), the KL divergence in (11) can be shown to equal to
L ( θ ) =
[Equation]
score matching
[Equation]
, (13) where we write s θ as the overall drift force of P θ in (10), that is, s [Equation] [Equation] , t ) =
[Equation]
[Equation]
[Equation]
[Equation]
+ η Ω ( z, [Equation] ) , and p θ 0 is the probability density function (PDF) of the initial distribution P θ 0 .
Therefore, L ( θ ) is a sum of the negative log-likelihood of the initial distribution that encourages P θ 0 ≈ Q Π ∗ 0 , and a least squares loss between s θ and η x .
In practice, we simply fix the initial distribution P θ 0 to be a delta measure on a fixed point (say x 0 = 0 ), so we only need to train the drift function f θ .
Algorithmshows a simple instance of the framework when the baseline process Q is the standard Brownian motion d Z t =
d W t , Q x = Q ( · | Z T = x ) and σ ( z, t )
= 1 .
Note that the least squares term in (13) can be viewed as enforcing f θ ≈ σ − 1 ( η x − η Ω ) , which reduces to f θ ≈ ∇ log ω Ω in the case of
Algorithm 1.
Related Works Bridge processes provide a simple and flexible approach to learning diffusion generative models, which was explored in Peluchetti (2021); Ye et al.
; Wu et al. De Bortoli et al. Heng et al. (2021) investigates the orthogonal problem of simulating from the bridge Q x for a given Q .
In comparison, our method learns diffusion models on general domains Ω on which an Ω -bridge can be derived (using h -transform or any other method), and hence provides a highly flexible framework for learning with structured data (including discrete, continuous, and their mixes).
This distinguishes it with existing approaches that are designed for special types of data (e.g., Ho et al.
Hoogeboom et al.
; Austin et al.
; Li et al.
; Dieleman et al.
(2022) for discrete data).
De Bortoli et al.
(2022) discusses how to learn score-based generative models on general Riemannian manifolds.
Another highly related work is Ye et al. (2022), which proposes to learn first hitting diffusion models for generating data on both discrete sets and spheres.
The advantage of our approach is that it is simpler and easier to derive for more complex types of domains.
E XPERIMENTS
We evaluate our algorithms for generating mixed-typed tabular data, grid-valued point clouds, categorical semantic segmentation maps, discrete CIFAR10 images.
We observe that Ω -bridge provides a particularly attractive and superb approach to generating data from various constrained domains.
Algorithm Overview
For all experiments, we use Algorithm 2 with the default choice of Q x in (7) and Q Ω in (9).
The specific form of η Ω is derived based on the specific choice of the domain Ω .
By default, we set the initialization Z
= 0 and the optimizer Adam. [Table]
4. G ENERATING M IXED -T YPE T ABULAR D ATA
Learning to generate tabular data is challenging, because tabular data usually contains a mixture of discrete and continuous attributes (Xu et al., 2019;
Park et al., 2018).
Unlike carefully designing special GANs as in previous works (Xu et al., 2019; Srivastava et al., 2017)
[Table]
Ω -bridge can be seamlessly applied to mixed-typed tabular data generation without any further modification.
In contrast, diffusion processes that solely work on discrete domain (Austin et al., 2021; Hoogeboom et al., 2021) cannot be applied to this task.
In this experiment
[Table]
we use the Adult Income dataset (Kohavi, 1996) [Table] which contains 30,162 training samples.
The data points are described by a series of attributes, including continuous ( age , capital-gain , etc.) and discrete ( sex , race , etc.).
We compare with conditional tabular
GAN (Xu et al., 2019) (CTGAN), CopulaGAN (Patki et al., 2016), and Table VAE (Xu et al., 2019) (TVAE), which are state-of-the-art GAN-based and VAE-based generative models for mixed-typed tabular data.
Following previous works (Xu et al., 2019; Patki et al., 2016), we measure the classification accuracy on the real data of logistic regression, AdaBoost classifier and MLP classifier when trained on the generated data.
In the Ω -bridge model, we set σ t = 3 exp( − 3 t ) and f θ a 3-layer MLP.
In this case, Ω =
I 1 × · · · ×
I 15 , where I 1 to I 9 are discrete domains and I 10 to I 15 are non-negative continuous domains.
For discrete domains, I = { e 1 , . . .
, e d } , we have, η I ( z, t )
=
σ 2 t ∇ z log (cid:80) [Equation]
− ∥ [Equation]
) (cid:17) ; for nonnegative continuous domains, I = [0 , + ∞ ) , derivation shows η I ( z, t ) = σ [Equation]
∇ [Equation]
log (cid:16) [Equation] [Equation] (cid:17)
[Equation]
where F is the standard Gaussian CDF.
Finally, we have η Ω ( z, t ) =
(cid:80) 15 i =
η
I
i
( z i , t ) for the whole domain Ω .
We set the number of diffusion steps to K = 2000 .
Results are shown in Table 1.
Result All the three different classifiers yield the highest accuracy when trained on the data generated by our method, referred to as Mixed-Bridge in this case.
The result reflects that the data generated by Mixed-Bridge is closer to the real distribution than the baseline methods.
Method
MMD ↓ COV ↑
1-NNA ↓
PCD (Luo & Hu, 2021a) 13.
R d -Bridge 13.
Grid-Bridge 12.
4. G ENERATING I NTEGER - VALUED P OINT C LOUDS
A feature of point clouds in 3D objects in graphics is that they tend to distribute evenly, especially if they are discretized from a mesh.
This aspect is omitted in most existing works on point cloud generation.
As a result they tend to generate non-uniform points that are unsuitable for real applications, which often involve converting back to meshes with procedures like Ball-Pivoting (Bernardini
[Figure]
Methods
[Table]
Uniform Dequantization (Uria et al., 2013)
[Table]
Variational Dequantization (Ho et al., 2019)
[Table]
Argmax Flow (Softplus thres.)
(Hoogeboom et al., 2021) [Table]
[Table]
Argmax Flow (Gumbel distr.)
(Hoogeboom et al., 2021)
[Table] [Table]
Argmax Flow (Gumbel thres.)
(Hoogeboom et al., 2021)
[Table]
Multinomial Diffusion (Hoogeboom et al., 2021)
[Table]
Cat.-Bridge (Constant Noise)
[Table]
Cat.-Bridge (Noise Decay A)
[Table]
Cat.-Bridge (Noise Decay B)
[Table]
Cat [Table] Bridge (Noise Decay C)
[Table]
et al., 1999).
We apply our method to generate point clouds that constrained on a integer grid which we show yields much more uniformly distributed points.
To the best of our knowledge, we are the first work on integer-valued 3D point cloud generation.
A point cloud is a set of points { x i } mi =1 , x i ∈ R 3 in the 3D space, where m refers to the number of points.
We apply two variants of our method: R d -Bridge and Grid-Bridge.
R d -Bridge generates points in the continuous 3D space, i.e., Ω R = R 3 m .
Grid-Bridge generate points that on integer grids, Ω Grid = { 1 , . .
, 128 } 3 m .
We fix the diffusion coefficient σ t = 1 .
The number of diffusion steps K is set to 1000 .
We test our method on ShapeNet (Chang et al., 2015) chair models, and compare it with Point Cloud Diffusion (PCD)
(Luo & Hu, 2021a), a state-of-the-art continuous diffusion-based generative model for point clouds.
The neural network f θ in our methods are the same as that of PCD for fair comparison.
Qualitative results and quantitative results are shown in
Figure 2 and Table 2.
As common practice (Luo & Hu, 2021a,b), we measure minimum matching distance (MMD), coverage score (COV) and 1-NN accuracy (1-NNA) using Chamfer Distance (CD) with the test dataset.
Result Both R d -Bridge and Grid-Bridge get better MMD, COV, and 1-NNA than PCD.
Moreover, by constraining the domain of interest to the integer grids, Grid-Bridge yields even better performance than R d -Bridge.
In Figure 2, since the point clouds generated by Grid-Bridge are limited to integer grids, the reconstructed meshes from Ball-Pivoting clearly have higher quality than
R d -Bridge and PCD.
4. G ENERATING S EMANTIC S EGMENTATION M APS ON C ITY S CAPES
We consider unconditionally generating categorical semantic segmentation maps.
We represent each pixels as a one-hot categorical vector.
Hence the data domain is Ω = { e 1 , . . .
, e c } h × w , where c is the number of classes and e i is the i -th c -dimensional one-hot vector, and h, w represent the height and width of the image.
In CityScapes (Cordts et al., 2016), h = 32 , w = 64 , c = 8 .
In this experiment, we test different schedule of the diffusion coefficient σ t , including (Constant Noise) :
σ t = 1 ; (Noise Decay A) : σ t = a exp( − bt ) ; (Noise Decay B) : σ t = a (1 − t ) ; (Noise Decay C) σ t = a (1 − exp( − b (1 − t ))) .
Here a and b are hyper-parameters.
The number of diffusion steps K is set to 500 .
We measure the negative log-likelihood (NLL) of the test set using the learned models.
The
NLL (bits-per-dimension) is estimated with evidence lower bound (ELBO) and importance weighted bound (IWBO)
(Burda et al., 2016), respectively, as in (Hoogeboom et al., 2021).
We compare Ω -Bridge with a state-of-the-art categorical diffusion algorithm, Argmax Flow (and Multinomial
Diffusion)
(Hoogeboom et al., 2021), and the traditional methods, uniform dequantization (Uria
[Table]
et al., 2013) and variational dequantization (Ho et al., 2019). The numerical results of the baselines are directly adopted from (Hoogeboom et al., 2021), and experiment configuration is kept the same for fair comparison.
The neural network f θ is the same as (Hoogeboom et al., 2021).
The results are shown in Figure 3 and Table 3.
Our Ω -bridge is named Categorical-Bridge (Cat.-Bridge) in this experiment.
Result We observe that all the four kinds of Cat.-Bridge can successfully generate categorical semantic segments, and different noise schedules result in different empirical performance.
Among the four variants of Cat.-Bridge, Cat.-Bridge with Noise Decay A yields the best ELBO and IWBO, surpassing all the other algorithms in comparison.
4. G ENERATING
[Figure]
In this experiment, we apply three types of bridges.
All of these bridges use the same output domain Ω = { 0 , . .
, 255 } h × w × c , where h, w, c are the height, width and number of channels of the images, respectively.
We set σ t = 3 exp( − 3 t ) .
We consider different initial distributions: (Init.
A)
Z 0 = 128 ; (Init.
B)
Z
=
ˆ µ 0 , (Init.
C)
Z 0 ∼ N (ˆ µ 0 , ˆ σ 0 ) , where ˆ µ 0 and ˆ σ 0 are the empirical mean and variance of pixels in the CIFAR10 training set.
The number of diffusion steps K is set to 1000 .
We compare with the variants of a state-of-the-art discrete diffusion model, D3PM (Austin et al., 2021).
For fair comparison, we use the DDPM backbone (Ho et al., 2020) as the neural drift f θ in our method, similar to D3PM.
We report the Inception Score (IS) Salimans et al.
(2016), Fr ´ echet Inception Distance (FID)
Heusel et al.
(2017) and negative log-likelihood (NLL) of the test dataset.
We call our method Integer-Bridge in this case.
The results are shown in Table 4 and Figure 4.
Result
In Table 4, Integer-Bridge with Initialization A,B,C can all get lower FIDs ( ≤ 7) than the variants of D3PM.
Among the three kinds of Integer-bridges, Integer-Bridge (Init.
B) obtains the lowest NLL (3.35).
It also beats D3PM Gauss + logistic (3.435) on NLL, which has the best NLL in the variants of D3PM.
C ONCLUSION AND L IMITATIONS
We present a framework for learning diffusion generative models on constrained data domains.
It leaves a number of directions for further explorations and improvement.
For example, the practical impact of the choices of the bridges Q , in terms of initialization, dynamics, and noise schedule, are still not well understood and need more systematical studies.
Besides, our current method is limited to Ω that are factorizable and integrable.
Moreover, application of Ω -bridge to many other practical fields also needs investigation in the future. A CKNOWLEDGEMENTS This research is supported by NSF CAREER1846421, SenSE2037267, EAGER-2041327, Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML)
A T HEORETICAL A NALYSIS ON B RIDGES A. D ERIVATION OF THE MAIN LOSS IN E QUATION (13)
[Proof of Equation (13)]
Denote by Q x = Q ( ·| Z T = x ) .
Note that
[Equation]
where const denotes a constant that is independent of θ .
Recall that Q x follows d Z t =
η x ( Z [0 ,t ] , t )d t + σ ( Z t , t )d W t , and P θ follows d Z t =
s θ ( Z t , t )d t + σ ( Z t , t )d W t .
By Girsanov theorem (e.g., Lejay, 2018),
[Equation]
x (cid:34)(cid:90) T (cid:13)(cid:13) s θ ( Z t , t ) − η x ( Z [0 ,t ] , t ) (cid:13)(cid:13) [-
-]2 2 d t (cid:35) =
E Z ∼ Q x (cid:34) − log p θ
( Z 0 )
+ 1(cid:90) T (cid:13)(cid:13)
s θ ( Z t , t )
− η x ( Z [0 ,t ] , t ) (cid:13)(cid:13) 2 2 d t (cid:35) + const.
Hence
L ( θ ) =
E x ∼ Π ∗ ,Z ∼ Q x (cid:34) − log p θ
( Z 0 )
+ 1(cid:90) T (cid:13)(cid:13)
s θ ( Z t , t )
− η x ( Z [0 ,t ] , t ) (cid:13)(cid:13) 2 2 d t (cid:35) + const =
E Z ∼ Q Π ∗ (cid:34) − log p θ
( Z 0 )
+ 1(cid:90) T (cid:13)(cid:13)
s θ ( Z t , t )
− η Z T
( Z [0 ,t ] , t ) (cid:13)(cid:13) 2 2 d t (cid:35) + const.
A. D ERIVATION OF THE DRIFT η Π ∗ OF Q Π ∗
Lemma A.1.
Let Q x is the law of d Z xt = η x ( Z x [0 ,t ] , t )d t + σ ( Z xt , t )d W t , Z 0 ∼ Q x 0 , and Q Π ∗ := (cid:82) Q x ( Z )Π ∗ (d x ) for a distribution Π ∗ on R d . Then Q Π ∗ is the law of d Z t = η Π ∗ ( Z [0 ,t ] , t )d t + σ ( Z t , t )d W t , Z 0 ∼ Q Π ∗ 0 , where
[Equation]
[Proof] Q Π ∗ is the solution of the following optimization problem:
[Equation]
By Girsanov’s Theorem (e.g., Lejay, 2018), any stochastic process P that has KL ( Q x || P ) <
+ ∞ (and hence is equivalent to Q x ) has a form of d Z t =
η Π ∗ ( Z [0 ,t ] , t )d t + σ ( Z t , t )d W t for some measurable function η Π ∗ , and E x ∼ Π ∗ [ KL ( Q x || P )]
[Equation]
It is clear that to achieve the minimum, we need to take P 0 ( · ) = E x ∼ Π ∗ [ Q x 0 ( · )] and η Π ∗ ( z [0 ,t ] , t ) =
E x ∼ Π ∗ ,Z ∼ Q x [ η x ( Z [0 ,t ] , t ) | Z [0 ,t ] =
z [0 ,t ] ] , which yields the desirable form of Q Π ∗ . A. P θ ∗ Y IELDS A M ARKOVIZATION OF Q Π ∗ As P θ is Markov by the model assumption, it can not perfectly fit Q Π ∗ which is non-Markov in general. This is a substantial problem because Q Π ∗ can be non-Markov even if Q x is Markov for all x ∈ Ω (see Section A.4).
In fact, using Doob’s h -transform method (Doob, 1984), Q Π ∗ can be shown to be the law of a diffusion process
[Equation]
where η Π ∗ is the expectation of η x when x = Z T is drawn from Q conditioned on Z [0 ,t ] .
We resolve this by observing that it is not necessary to match the whole path measure ( P θ ≈ Q Π ∗ ) to match the terminal ( P θT ≈ Q Π ∗ T
= Π ∗ ).
It is enough for P θ to be the best Markov approximation (a.k.a.
Markovization) of Q Π ∗ , which matches all (hence terminal) fixed-time marginals with Q Π ∗ :
[Equation]
the set of all Markov processes on [0 , T ] .
Proposition A.2.
The global optimum of L ( θ ) in (11) and (13) is achieved by θ ∗ if
[Equation]
In this case, P θ ∗ =
Proj( Q Π ∗ , M ) is the Markovization of Q Π ∗ , with which it matches all timemarginals: P θ ∗ t = Q Π ∗ t for all time t ∈ [0 , T ] .
In addition,
[Equation] Note that s θ ∗ is a conditional expectation of η Π ∗ : [Equation] θ
[Equation]
Theorem 1 of Peluchetti (2021) gives a related result that the marginals of mixtures of Markov diffusion processes can be matched by another Markov diffusion process, but does not discuss the issue of Markovization nor connect to KL divergence.
Theorem 1 of Song et al. (2021) is the special case of (15) when Q Π ∗ is Markov. [Proof of Proposition A.2] It is the combined result of Lemma A.3 and Lemma A.4 below.
Lemma A.3.
Let Q be a non-Markov diffusion process on [0 , T ] of form
[Equation] and M = arg min P ∈M KL ( Q || P ) be the Markovization of Q , where M is the set of all Markov processes on [0 , T ] . Then Q is the law of [Equation]
where [Equation] In addition, we have Q t = M t for all time t ∈ [0 , T ] .
[Proof] By Girsanov’s Theorem (e.g., Lejay, 2018), any process that has KL ( Q || M ) < + ∞ (and hence is equivalent to Q ) has a form of d Z t = m ( Z [0 ,t ] , t )d t + σ ( Z t , t )d W t , where m is a measurable function. Since M is Markov, we have m ( Z [0 ,t ] , t ) = m
( Z t , t ) .
Then
[Equation]
It is clear that to achieve the minimum, we need to take M= Q 0 and m ( z, t ) =
E Z ∼ Q [ η ( Z [0 ,t ] , t ) | Z t = z ] . To prove Q t =
M t , note that by the chain rule of KL divergence:
[Equation]
As the second term P ( ·| Z t ) is independent of the choice of the marginal P t at time t ∈ [0 , T ] , the optimum should be achieved by M only if M t = Q t .
Lemma A.4.
Let
[Equation]
where M is the Markovization of Q (see Lemma A.3).
Then
[Equation]
Hence, assume there exists θ ∗ such that P θ ∗ =
M and write L ( θ ) : = KL ( Q || P θ ) . We have KL
[Equation]
Q [Equation]
[Proof] Note that [Equation] where we define ∥ f ∥ 2 Q ,
[Equation]
= E Z [Equation] [Equation] Z
[Equation]
, t
[Equation]
Z
[Equation] t [Equation] 2 2 d t (cid:105) .
On the other hand,
[Equation] Using Lemma A.5 with a ( z ) = σ ( z, t ) − 1 s θ ( z, t ) , and b ( z [0 ,t ] )
= σ ( z, t ) − 1 η ( z [0 ,t ] , t ) , we have the following bias-variance decomposition: [Equation] Hence, KL ( Q || P
[Equation] = KL [Equation]
KL [Equation] Finally, KL ( M T || P θT ) ≤ KL ( M || P θ ) is the direct result of the following factorization of KL divergence: [Equation]
Lemma A.5.
Let ( X, Y ) be a random variable and a ( x ) , b ( x, y ) are square integral functions.
Let m ( x ) = E [ b ( X, Y ) | X = x ] . We have [Equation]
[Proof]
[Equation]
where
[Equation]
A. M ARKOV AND R ECIPROCAL P ROPERTIES OF Q Π ∗
Mixture of Bridges and Initialization It is an immediate observation that the mixtures of a set of bridges are also bridges: let Q z,A be a set of A -bridges indexed by a variable z ,
then Q A := (cid:82)
Q z,A µ (d z ) is an x -bridge for any distribution µ on z .
A special case is to take the mixture of the conditional bridges in (5) starting from different deterministic initialization, which shows that we can obtain a valid x -bridge by equipping the same drift in (5) with essentially any initialization.
Hence, the choices of the drift force and initialization in
Q x can be completely decouple. Proposition A.6. Let ˜ Q is a path measure and Ω x is the set of z for which ˜ Q z 0 ,x ( · ) : = ˜ Q ( ·| Z T = x, Z 0 =
z 0 ) exists.
Then Q x := (cid:82) ˜ Q z 0 ,x µ (d z 0 | x ) is an x -bridge, for any distribution µ on Ω × Ω . [Proof of Proposition A.6] This is an obvious result. We have Q z 0 ,x ( Z T = x ) = 1 by the definition of conditioned processes. Hence Q x ( Z T = x ) = (cid:82) Q z 0 ,x ( Z T = x ) µ (d z 0 | x ) = (cid:82) µ (d z 0 | x ) = 1 .
Markov and Reciprocal Properties of Q Π ∗ If Q x is constructed as Q x = Q ( ·| Z T = x ) , it is easy to see that Q Π ∗ := (cid:82) Q x ( · )Π ∗ (d x ) is Markov iff Q is Markov.
Proposition A.7.
Assume Q x = Q ( · | Z T = x ) and π ∗ ( z ) : = dΠ ∗ d Q T ( z ) exists and is positive everywhere. Then Q Π ∗ is Markov, iff Q is Markov.
[Proof of Proposition A.7] If Q x = Q ( · | Z T = x ) , we have from the definition of Q Π ∗ :
[Equation]
where π ∗ ( Z T ) = dΠ ∗ d Q T ( Z T ) . Therefore, Q Π ∗ is obtained by multiplying a positive factor π ∗ ( Z T ) on the terminal state Z T of Q . Hence Q Π ∗ has the same Markov structure as that of Q .
If Q x is constructed from mixtures of bridges as above, the resulting Q Π ∗ is more complex.
In fact, simply varying the initialization µ in Proposition (A.6) can change the Markov structure of Q Π ∗ .
Proposition A.8. Take Q x to be the dynamics in (7) initialized from Z 0 ∼ N (0 , v 0 ) . Assume σ t > 0 , ∀ t ∈ [0 , T ] . Then Q Π ∗ is Markov only when v 0 = 0 , or v 0 = + ∞ .
[Proof of Proposition A.8] When taking Q x to be the dynamics (7) initialized from Z 0 ∼ µ 0 = N (0 , v 0 ) , we have Q x = (cid:82) µ 0 (d z 0 ) ˜ Q z 0 ,x , where ˜ Q z 0 ,x = ˜ Q ( ·| Z 0 = z 0 , Z T = x ) with ˜ Q following Brownian motion d Z t = d W t .
Hence, we can write Q Π ∗ (d Z ) = ˜ Q (d Z ) r ( Z 0 , Z T ) , where r ( z 0 , z T ) = d µ 0 ⊗ Π ∗ d˜ Q 0 ,T ( z 0 , z T ) . From L´eonard et al.
, Q Π ∗ is Markov iff r ( x, z 0 ) = f ( x ) g ( z 0 ) for some f and g , which is not the case except the degenerated case ( v 0 = 0 and v 0 = + ∞ ) because ˜
Q 0 , 1 is not factorized. On the other hand, when v 0 = 0 , we have that Q x = Q ( ·| Z T = x ) is the standard Brownian bridge and hence Q Π ∗ is Markov following Proposition A.7. When v 0 = + ∞ , as the case of SMLD, Q Π ∗ is the law of Z t = ˜ Z T − t with d ˜ Z t =
d W t and ˜ Z 0 ∼ Π ∗ , which is also Markov. The right characterization of Q Π ∗ from Proposition (A.6) involves reciprocal processes (L ´ eonard et al., 2014). Definition A.9. A process Z with law Q on [0 , T ] is said to be reciporcal if it can be written into
Q = (cid:82) ˜ Q z 0 ,z T µ (d z 0 , d z T ) , where ˜ Q is a Markov process and ˜ Q z 0 ,z T = ˜ Q ( ·| Z 0 = z 0 , Z T = z T ) , and µ is a probability measure on Ω × Ω . Proposition A.10.
Q Π ∗ is reciprocal iff Q x = (cid:82) ˜ Q z 0 ,x µ (d z 0 | x ) for a Markov ˜ Q and distribution µ .
[Proof of Proposition A.10] Note that
[Equation]
Hence if ˜ Q is Markov, Q Π ∗ is reciprocal by Definition A.9.
On the other hand, if Q Π ∗ is reciprocal, we have Q Π ∗ ( · ) =
(cid:82) M z 0 ,x ( · ) µ (d z 0 , d x ) for some Markov process M and probability measure µ on Ω × Ω .
In this case, we have Q x ( · ) = Q Π ∗ ( ·| Z T = x ) = (cid:82) M z 0 ,x ( · ) µ (d z 0 | x ) , assuming it exits.
Intuitively, a reciprocal process can be viewed as connecting the head and tail of a Markov chain, yielding a single loop structure.
A characteristic property is Q ( X [ s,t ] ∈
A | Z [0 ,s ] , Z [ t,T ] ) =
Q ( X [ s,t ] ∈
A | Z s , Z t ) , where A is any event that occur between time s and t .
Solutions of the Schrodinger bridge problems are reciprocal processes (L´eonard et al., 2014).
A. P RACTICAL A LGORITHM AND E RROR A NALYSIS
In practice, we need to introduce empirical and numerical approximations in both training and inference phases.
Denote by τ = { τ i } K +1 i =1 a grid of time points with 0 = τ 1 < τ 2 . . .
<
τ K +1 = T .
During training, we minimize an empirical and time-discretized surrogate of L ( θ ) as follows
[Equation] where ∆( θ ; Z, t ) : = (cid:13)(cid:13) σ − 1 ( Z t , t )( s θ ( Z t , t ) − η x ( Z [0 ,t ] , t ))
(cid:13)(cid:13) 2 , and { Z ( i ) } is drawn from Q Π ∗ , and τ ( i ) can be either a deterministic uniform grid of [0 , T ] , i.e., τ ( i ) =
{ i/K } Ki =0 , or drawn i.i.d.
uniformly on [0 , T ] (see e.g.,Song et al.
(2020b)
; Ho et al.
A subtle problem here is that the variance of ∆( θ ; Z, t ) grows to infinite as t ↑ T .
Hence, we should not include ∆( θ ; Z, T ) at the end point
τ K +1 =
T into the sum in the loss ℓ ( θ, Z, τ ) to avoid variance exploding.
In the sampling phase, the continuous-time model P θ should be approximated numerically.
A standard approach is the Euler-Maruyama method, which simulates the trajectory on a time grid τ by
[Equation]
The final output is ˆ Z T .
The following result shows the KL divergence between Π ∗ and the distribution of ˆ Z T can be bounded by the sum of the step size and
the expected optimality gap
E [ ˆ L ( θ ) − ˆ L ( θ ∗ )] of the time-discretized loss in (16).
A.5. T IME -D ISCRETIZATION E RROR A NALYSIS (P ROPOSITION A.11)
Proposition A.11.
Assume Ω = R d and σ ( z, t ) = σ ( t ) is state-independent.
Take the uniform time grid τ unif := { iϵ } Ki =0 with step size ϵ = T/K in the sampling step (17) . Assume σ ( t ) > c > 0 , ∀ t and σ ( t ) is piecewise constant w.r.t.
time grid τ unif . Let L ϵ ( θ ) = E Z ∼ Q Π ∗ [ ℓ ( θ ; Z, τ unif )] .
Let
P θ,ϵT be the distribution of the resulting sample ˆ Z T .
Let θ ∗ be an optimal parameter satisfying (14) .
Assume C 0 :
= sup z,t (cid:16)(cid:13)(cid:13) s θ ∗ ( z, t ) (cid:13)(cid:13) 2 / (1 + ∥ z ∥ 2 ) , tr( σ 2 ( z, t )) , E P θ ∗ [ ∥ Z 0 ∥ 2 ] (cid:17) < + ∞ , and (cid:13)(cid:13) s θ ∗ ( z, t ) − s θ ∗ ( z ′ , t ′ ) (cid:13)(cid:13) 2 2 ≤ L (cid:16) ∥
z − z ′ ∥ 2 + | t − t ′ | (cid:17) for ∀ z, z ′ ∈ R d and t, t ′ ∈ [0 , T ] .
Then (cid:113)
[Equation]
We provide the analysis and proof for Proposition A.11 in the following text.
Proposition A.12.
Assume Ω = R d and σ ( z, t ) =
σ ( t ) is state-independent and σ ( t ) >
c > 0 , ∀ t ∈ [0 , T ] .
Take the uniform time grid τ unif := { iϵ } Ki =0 with step size ϵ = T/K in the sampling step (17) .
Let
[Equation]
Z
[Equation] ] with [Equation] where ϵ > 0 is a step size with T =
Kϵ and t k = ( k − 1) ϵ , and σ 2 k : = ( t k +1 − t k ) − 1 (cid:82) t k +1 t k σ ( t ) 2 d t . Let P θ,ϵT be the distribution of the sample ˆ Z T resulting from the following Euler method:
[Equation]
where ξ k ∼ N (0 , I d ) is the standard Gaussian noise in R d .
Let θ ∗ be an optimal parameter satisfying (14) . Assume C 0 : = sup z,t
[Equation]
s θ ∗ ( z, t ) [Equation] / (1 + ∥ z ∥ [Equation]
[Equation]
tr
[Equation] σ [Equation] ( z, t ))
[Equation]
E P θ ∗ [ ∥ Z 0 ∥ 2 ] (cid:17) <
+ ∞ , and s θ ∗ satisfies (cid:13)(cid:13)
s θ ∗ ( z, t ) − s θ ∗ ( z ′ , t ′ )
(cid:13)(cid:13) 2 2 ≤ L (cid:16) ∥
z − z ′ ∥ 2 + | t − t ′ | (cid:17) for ∀ z, z ′ ∈ R d and t, t ′ ∈ [0 , T ] .
Then we have
[Equation]
[Proof of Proposition A.11]
This is the result of Lemma A.13 below by noting that the ˆ P θ there is equivalent to the Euler method above, and L ϵ ( θ ) − L ϵ ( θ ∗ ) ≤ ˜ L ϵ ( θ ) − ˜
L ϵ ( θ ∗ ) (because σ − 2 k = (( t k +1 − t k ) − 1 (cid:82) t k +1 t k [Equation]
σ
[Equation]
Lemma A.13.
Let h be a step size and ϵ = T/K for a positive integer K . For each t ∈ [0 , ∞ ) , denote by ⌊ t ⌋ ϵ = max( { kϵ : k ∈ N } ∩ [0 , t ] ). Assume [Equation]
P
[Equation]
where P θ ∗ is the Markovianization of Q Π ∗ , and ˆ P θ is a discretized version of P θ . Define [Equation] Assume the conditions of Lemma A.17 holds for P θ ∗ , and σ ( z, t ) ≥ c > 0 for all z, t , and s θ ∗ satisfies
[Equation]
∀ z
[Equation]
z
[Equation]
∈
[Equation]
and t, t
[Equation]
∈ [0 , T [Equation] Then (cid:113)
[Equation] where C is a constant that is independent of ϵ . [Proof] Define ∥ f ∥ 2 Q ,σ = {+E Z ∼ Q [ (cid:82) T 
+}
∥ σ ( Z, t ) f ( Z, t ) ∥ 2 ] for convenient notation. Let s θϵ ( Z, t )
=
s θ ( Z ⌊ t ⌋ ϵ , ⌊ t ⌋ ϵ ) , and η ϵ = η Z T ( Z [0 , ⌊ t ⌋ ϵ ] , ⌊ t ⌋ ϵ ) . [Equation] where ω > 0 is any positive number and
[Equation]
A. [Equation] and
[Equation]
A
[Equation] [Equation] where C P θ ∗ is a constant depending on P θ ∗ that comes from Lemma A.17. Hence
[Equation]
This completes the proof.
Lemma A.14.
For any a, b ∈ R d , and ω ≥ 0 , [Equation] [Proof] [Equation]
Lemma A.15. Assume T ≥ 0 , ϵ ≥ 0 and T/ϵ ∈ N .
We have
[Equation] [Proof] [Equation]
Lemma A.16 ( Gr ¨ onwall’s inequality) . Let I denote an interval of the real line of the form [ a, ∞ ) or [ a, b ] or [ a, b ) with a < b . Let α, β and u be real-valued functions defined on I .
Assume that β and u are continuous and that the negative part of α is integrable on every closed and bounded subinterval of I . (a) If β is non-negative and if u satisfies the integral inequality
[Equation]
which is true if
[Equation]
then
[Equation]
(b) If, in addition, the function α
is non-decreasing, then
[Equation]
Lemma A.17.
Consider
[Equation]
Assume there exists a finite constant C 0 , such that
[Equation]
and E [ ∥ Z 0 ∥ 22 ] ≤ C 0 .
Then for any 0 ≤ s ≤ t ≤ T , we have
[Equation]
where K C 0 ,T is a finite constant that depends on C 0 and T .
[Proof] Let η =
sup
x
,t tr( σσ ⊤ ( x, t )) .
We have by Ito Lemma,
[Equation]
Using Gronwall’s inequality,
[Equation]
Taking s = 0 yields that
[Equation]
Hence
[Equation]
Therefore,
[Equation]
≤ C
[Equation]
where
[Equation]
A.5. S TATISTICAL E RROR A NALYSIS (P ROPOSITION A.18)
To provide a simple analysis of the statistical error, we assume that ˆ θ n = arg min θ ˆ L ( θ ) is an asymptotically normal M-estimator of θ ∗ following classical asymptotic statistics (Van der Vaart, 2000), with which we can estimate the rate of the excess risk L ϵ (ˆ θ n )
− L ϵ ( θ ∗ ) and hence the KL divergence.
Proposition A.18.
Assume the conditions in Proposition (A.11) .
Assume ˆ θ n = arg min θ ˆ L ϵ ( θ ) with ˆ L ϵ ( θ )
=
(cid:80) ni
=1 ℓ ( θ ; Z ( i ) , τ unif ) /n , Z ( i ) ∼ Q Π ∗ .
Take Q x to be the standard Brownian bridge d Z xt = x − Z xt T − t d t + d W t with Z 0 ∼ N (0 , v 0 ) and v 0 > 0 .
Assume √ n (ˆ θ n − θ ∗ )
d −→ N (0 , Σ ∗ ) as n → + ∞ , where Σ ∗ is the asymptotic covariance matrix of the
M estimator ˆ θ n .
Assume L ϵ ( θ ) is second order continuously differentiable and strongly convex at θ ∗ .
Assume Π ∗ has a finite covariance and admits a density function π that satisfies sup
[Equation]
∇ θ
{+s θ ∗ ( Z t , t ) (cid:13)(cid:13) 
+}
(1 + ∥∇ log π ( Z T ) ∥ 2 + tr( ∇ 2 log π ( Z T ))) (cid:105) <
+ ∞ .
We have
E
[Equation]
log
[Equation]
[Equation]
The expectation in Eq.
(18) is w.r.t.
the randomness of ˆ θ n .
The log(1 /ϵ ) factor shows up as the sum of a harmonic series as the variance of ∆( θ ; Z, t ) grows with O (1 / ( T − t )) when t ↑ T .
Taking ϵ
= 1 /n yields KL (Π ∗
||
P ˆ θ n ,ϵ T ) =
O (log n/n ) .
If we want to achieve KL (Π ∗ || P ˆ θ n ,ϵ T )
=
O ( η ) , it is sufficient to take K =
T/ϵ =
O (1 /η ) steps and n =
Ø(log(1 /η ) /η )
data points.
[Proof of Proposition A.18] Let
[Equation]
where { Z ( i ) } ni =1 is drawn i.i.d. from Q Π ∗ .
We assume that ˆ θ n is an asymptotically normal Mestimator, in which case we have [Equation]
where [Equation] and
[Equation] where f ≍ g denotes that f − g = o (1) . We now need to bound tr
( H − 1 ∗ V ∗ ) . Combining the results in Lemma A.19 and Lemma A.23, we have when t k = ( k [Equation]
ϵ and T = Kϵ , [Equation] Hence,
E [ [Equation]
Lemma A.19. Assume the conditions in Proposition A.18.
Define
[Equation] for ∀ k = 1 , . . .
K .
Then [Equation] [Proof] From Lemma A.21, tr( H − 1 ∗ V ∗ ) ≤ ( λ min ( H ∗ )) − 1 tr( V ∗ ) . Hence we just need to bound tr( V ∗ ) .
[Table] 1 /= I 1 / 2+
K
K (cid:88) k = [Equation] I 1 /
[Equation] ≤ I 1 / [Equation]
+ (cid:118)(cid:117)
[Equation]
(cid:116) [Equation] K K (cid:88) k
=
[Equation] I k . Lemma A.20. Assume the results in Lemma A.19 and Lemma A.23 hold.
Assume max [Equation] Then for k = 1 , . . , K , we have I k
= O (cid:16) 1 T − t k + 1 (cid:17) .
[Proof] It is a direction application of (20).
Lemma A.21.
Let A and B be two d × d positive semi-definite matrices. Then tr( AB ) ≤ λ max ( A )tr( B ) . [Proof] Write A into A = (cid:80) di =1 λ i u
i u ⊤ i where λ i and u i is the i -th eigenvalue and eigenvectors of
A , respectively.
Then
[Equation]
Controlling the Conditional Variance of the Regression Problem Assume Q x is the standard Brownian bridge: [Equation] In this case, the (ideal) loss function is
[Equation] , where Y t = X − Z t− t .
The second part of the loss is a least square regression for predicting Y t = η X ( Z t , t ) with s θ ( Z t , t ) . The conditioned variance cov( Y t | Z t ) is an important factor that influences the error of the regression problem. We now show that tr(cov( Y t | Z t ))
=
O (1 /T − t ) which means that it explodes to infinity when t ↑ T . First, note that tr(cov( Y t | Z t )) = [Equation] 2 tr(cov( X | Z t )) .
Using the estimate in Lemma A.23, we have [Equation]
Lemma A.22.
For the standard Brownian bridge in (19) , we have
[Equation] [Proof] Let Z z 0 ,x t be the same process that is initialized from Z z 0 ,x= z
We have from the textbook result regarding Brownian bridge that we can write Z z
[Equation] x t [Equation]
where ξ t is some standard Gaussian random variable.
The result follows directly as Z xt = Z Z
[Equation]
x t with Z 0 ∼ N (0 , v 0 ) .
Lemma A.23.
Let π ∗ be the density function Π ∗ on R d whose covariance matrix exists.
When X ∼ Π ∗ and Z ∼ Q X from (19) with v 0 > 0 . Then the density function ρ t ( x | z ) of X | Z t = z t satisfies
[Equation]
In addition, there exists positive constants c < + ∞ and τ ∈ (0 , T ) , such that [Equation] t ≤ T c [Equation]
when 0 ≤ t ≥ τ, where w t = T ( T − t ) t + [Equation] − t ) 2 t.
So tr(cov ρ t ( x | z )) is bounded and decay to zero with rate O ( T − t ) as t ↑ T . [Proof] We know that X ∼ Π ∗ and Z Xt | X ∼ N ( t/TX, w t ) . Hence, (21) is a direct result of Bayes rule. Then Lemma A.25 gives
[Equation] On the other hand, [Equation] When t → 0 , we have 1 /w t → 0 and T/ ( tw t ) → 0 .
Hence, ρ t ( x | z ) converges to π ∗ ( x ) as t → 0 , as a result, tr(cov ρ t ( x | z )) → tr(cov π ∗ ( x )) < + ∞ . Therefore, for any c > 0 , there exists t 0 > 0 , such that tr(cov ρ t ( x | z )) ≤ tr(cov π ∗ ( x ))
+ c when 0 ≤ t ≤ t 0 .
Remark A.24.
We need to have v 0 > 0 to ensure that T/ ( tw t ) → 0 in the proof of Lemma A.23.
This is purely a technical reason, for yielding a finite bound of the conditioned variance when t is close to 0 . We can establish the same result when v 0 = 0 by adding the assumption that max [Equation] ∇
[Equation] cov
[Equation] ∗ Ztk [Equation]
Z [Equation] )) (cid:105) < + ∞ , where Π ∗ z is the distribution with density π ∗ z ( x ) ∝ π ∗ ( x ) exp( z ⊤ x/T ) .
Lemma A.25.
Let p ( x ) ∝ π ( x ) exp (cid:16) − α ∥ x − b ∥ 22(cid:17) be a positive probability density function on R d , where α > 0 , b ∈ R and log π is continuously second order differentiable. Then
[Equation] [Proof] Let us focus on the case when b = 0 first. Stein’s identity says that
[Equation]
for a general continuously differentiable function ϕ when the integrals above are finite. Taking ϕ = x yields that [Equation]
which gives [Equation] On the other hand, taking ϕ ( x ) = ∇ x log π ( x ) yields
[Equation]
which gives [Equation] This gives
[Equation] For b ̸ =
[Equation]
define ˜ p ( x ) ∝ π ( x + b ) exp (cid:16) − α 2 ∥ x
∥ 2 (cid:17) , which is the distribution of ˜ x
= x − b when x ∼ p . Then applying the result above to ˜ p yields
[Equation]
A. C ONDITION FOR Ω - BRIDGES
We provide the proof for Proposition 2.3.
[Proof of Proposition 2.3] By the formula of KL divergence between two diffusion processes, we have
[Equation]
This means that Q Ω and P θ are absolutely continuous to each other, and hence have the same support.
Therefore, Q Ω ( Z T ∈ Ω)
= 1 implies that P θ ( Z T ∈ Ω) = 1 .
A. E XAMPLES OF Ω -B RIDGES
If Ω is a product space, the integration can be factorized into one-dimensional integrals.
Specifically, assume Ω
=
I 1 × · · · I d , then
[Equation]
where η I i is the drift fore of the
I i -bridge, and z i is the i -th element of z =
[ z i ] .
Therefore, it is sufficient to focus on 1D case below.
Consider the bridge process constructed from the Brownian motion in (9).
If Ω is a discrete set, say Ω = { e 1 . .
, e K } , we have
[Equation]
where
[Equation]
If Ω =
[ a, b ] , we have
[Equation]
where F is the standard Gaussian CDF.
B A DDITIONAL M ATERIALS OF THE E XPERIMENTS
In our experiments, T = 1 and ϵ = T/K = 1 /K .
Moreover, we take the time grid by randomly sampling from { i/K } K − 1 i =0 for the training objective Eq.
For evaluation, we calculate the standard evidence lower bound (ELBO) by viewing the resulting time-discretized model as a latent variable model:
[Equation]
where t k = ( k − 1) ϵ , and ˆ p θ is the density function of the time-discretized version of P θ , and q is the density function of Q .
We adopt Monte-Carlo sampling to estimate the log-likelihood.
As in (Song et al., 2020b), we repeat 5 times in the test set for the estimation.
For categorical/integer/grid generation, the likelihood of the last step should take the rounding into account: in practice, we have ˆ Z T = rounding( ˆ Z t K + ϵs θ ( ˆ Z t K , t K )+ √ ϵσ
( Z t k , t K ) ξ K , Ω) , where rounding( x, Ω) denotes finding the nearest element of x on Ω , and hence the likelihood ˆ p θT | t K of the last step should incorporate the rounding operator as a part of the model.
B. G ENERATING M IXED -T YPE T ABULAR D ATA
In this experiment, the metrics are measured by the implementation from Synthetic Data Vault (SDV) (Patki et al., 2016).
For baseline methods, we adopt their open-sourced official implementation 1 .
For the machine learning models adopted for evaluation, logistic regerssion, AdaBoost and
MLP, we directly use their default configuration in SDV.
For the results in Table 1, we repeat the experiments with 5 different random seeds and report their standard deviation.
We provide additional generated samples from Mixed-Bridge in Figure 5.
[Figure]
In this experiment, we need to process point cloud data on integer grid.
To prepare the data, we firstly sample 2048 points from the ground truth mesh.
Then, we normalize all the point clouds to a unit bounding box.
After this, we simply project the points onto grid point by rounding the coordinate to integer.
The metrics in the main text, MMD, COV and 1-NNA are computed with respect to the post-processed integer-valued training point clouds.
For the results in Table 2, we repeat the experiments for 3 times and report the mean of the experiments.
B. G ENERATING S EMANTIC S EGMENTATION M APS ON C ITY S CAPES
In this experiment, we set (Noise Decay A) : σ 2 t = 3 exp( − 3 t ) ; (Noise Decay B) : σ 2 t = 3(1 − t ) ; (Noise Decay C) σ 2 t = 3 − 3 exp( − 3(1 − t )) .
We visualize the noise schedule in Figure 6.
Note that, except for Constant Noise, all the other three processes gradually decrease the magnitude of the noise as t → 1 .
For fair comparison, we use the same neural network as in Hoogeboom et al.
The network is optimized with Adam optimizer with a learning rate of 0 . 0002 .
The model is trained for 500 epochs.
The CityScapes dataset (Cordts et al., 2016) contains photos captured by the cameras on the driving cars.
A pixel-wise semantic segmentation map is labeled for each photo.
As in (Hoogeboom et al., 2021), we rescale the segmentation maps from cityscapes to 32 × 64 images using nearest neighbour interpolation.
Our training set and test set is exactly the same as that of (Hoogeboom et al., 2021) for fair comparison.
For the results in Table 3,we repeat the experiments for 3 times and report the mean of the experiments.
We provide more samples in Figure 9.
B. D ISCRETE CIFAR10 GENERATION
The model is trained using the same training strategy as DDPM
(Ho et al., 2020) with the code base provided in Song et al.
(2020b).
Specifically, the neural network is the same U-Net structure as the implementation in Song et al.
(2020b).
The optimizer is Adam with a learning rate of 0 . 0002 .
According to common practice (Song & Ermon, 2020; Song et al., 2020b), the training is smoothed by exponential moving average (EMA) with a factor of 0 . 999 .
We use K = 1000 and d t = 0 .
001 for discretizing the SDE.
To account for the discretization error, after the final step, we apply rounding to the generated images to get real integer-valued images.
We compare the value distribution of the generated images in Figure 8.
[Figure] [Figure]