First Hitting Diffusion Models for Generating Manifold, Graph and Categorical Data
Abstract
We propose a family of First Hitting Diffusion Models (FHDM), deep generative models that generate data with a diffusion process that terminates at a random first hitting time.
This yields an extension of the standard fixed-time diffusion models that terminate at a pre-specified deterministic time.
Although standard diffusion models are designed for continuous unconstrained data, FHDM is naturally designed to learn distributions on continuous as well as a range of discrete and structure domains.
Introduction
Standard diffusion processes used in ML can be classified into two categories: 1) infinite (or mixing) time diffusion processes such as Langevin dynamics, which requires the process to run sufficiently long to converge to the invariant distribution , whose property is leveraged for the purpose of learning and inference; and 2) fixed time diffusion processes such as DDPM, SMLD, and Schrodinger bridges [De Bortoli et al., 2021], which are designed to output the desirable results at a pre-fixed time.
Although fixed-time diffusion has been show to surpass infinite time diffusion on both speed and quality, it still yield slow speed for modern applications due to the need of a pre-specified time and the incapability to adapt the time based on the difficulty of instances and problems.
Moreover, standard diffusion models are naturally designed on R d , and can not work for discrete and structured data without special cares.
In this work, we study and explore a different first hitting time diffusion model that terminates at the first time as it hits a given domain, and leverages the distribution of the exit location (known as exit distribution, or harmonic measure [Oksendal, 2013]) as a tool for learning and inference.
We provide the basic framework and tools for first hitting diffusion models.
We leverage our framework to develop a general approach for learning deep generative models based on first hitting diffusion.
This approach generalizes SMLD and its SDE extensions but can be attractively applied to a range of discrete and structured domains.
This contrasts with the standard diffusion models, which are restricted to continuous R d data.
In particular, we instantiate our framework to three cases, yielding new diffusion models for learning 1) spherical, 2) binary and 3) categorical data.
Main Framework 2.1 First Hitting Diffusion Processes
Let Π ∗ be a distribution of interest on a domain Ω ⊂ R d .
The goal is to construct a first hitting stochastic process , which starts from a point outside of Ω and returns a sample drawn from Π ∗ when it first hits set Ω .
We start with introducing the new first hitting model.
[Figure] Let Z := { Z t : t ∈ [0 , + ∞ ) } be a continuous-time Markov process with probability law Q taking value in a set V that contains Ω as a subset.
Here Q is a probability measure defined on the space of all continuous trajectories C ([0 , + ∞ ) , R d ) .
We use Q t to denote the marginal distribution of Z t at time t .
We assume that the process is initialized from a point Z 0 outside of Ω .
Denote by τ the first hitting time of Z t on Ω , that is, τ = inf t { t ≥ 0: Z t ∈ Ω } .
We call that Z t is absorbing to set Ω if
i)
The process enters Ω in finite time almost surely when initialized from anywhere in V , that is, Q ( τ < + ∞ | Z {+
+}
= z ) = 1 , ∀ z ∈ V .
ii)
The process stops to move once it arrives at Ω , that is, Q ( Z t + s =
Z t | Z t ∈ Ω)
= 1 , ∀ s, t ≥ 0 .
We define the Poisson kernel of Q as the conditional distribution of Z τ given Z t =
z , denoted by Q Ω (d x | Z t = z ) :
=
Q ( Z τ =
d x | Z t = z ) .
The marginal distribution of Z τ , which we write as Q Ω (d x ) =
Q ( Z τ = d x ) , is called the exit distribution, or harmonic measure .
Note that Q Ω (d x ) =
(cid:82) V Q Ω (d x | Z 0 =
z ) Q 0 (d z ) .
The crux of our framework is to leverage the exit distribution Q Ω as a tool for statistical learning and inference, which is different from traditional frameworks that exploit the properties of the distributions at a fixed time or at convergence.
Example 2.1 (Sphere Hitting) .
As shown in Figure 1-B, let V = { x ∈ R d : ∥ x ∥ ≤ 1 } be the unit ball and Ω =
S d := ∂V the unit sphere.
Let Z be a Brownian motion starting from z ∈ V and stopped once it hits the boundary Ω .
It is written as
[Equation]
where W t is a Wiener process; the indicator function I ( ∥ Z t ∥ < 1) sets the velocity to zero and hence stops the process once Z t hits Ω .
The Poisson kernel in this case is a textbook result:
[Equation]
where µ Ω is the surface measure on Ω = S d .
[Equation]
Example 2.2 (Boolean Hitting) .
As shown in Figure 1-C, let V =
[0 , 1] d be the unit cube and Ω
= B d := { 0 , 1 } d the Boolean cube.
Let Z be a Brownian motion starting from Z 0 ∈ V and confined inside the cube V in the following way:
[Equation]
where Z t,i is the i -th element of Z .
Here, each coordinate Z t,i stops to move once it hits one of the end points ( 0 or 1 ).
It can be viewed as a particle flying in a room that sticks on a wall once it hits it.
Proposition 2.3.
The Poisson kernel of Q B d is a simple product of Bernoulli distributions:
[Equation]
z
[Equation]
:= d
[Equation]
Ber( x i |
z i ) , where Ber( x i |
z i ) =
x i z i + (1 − x i )(1 − z i );
Ber( x i |
z i ) is the likelihood function of observing x i ∈ { 0 , 1 } under Bernoulli ( z i ) with z i ∈ [0 , 1] .
Example 2.4 (Fixed Time Hitting) .
Our first hitting framework includes the more standard models with fixed terminal time.
To see this, let ¯ Z t = ( t, Z t ) be a stochastic process Z t with law Q augmented with time t as one of its coordinates.
Let V =
[0 , t ] × R d and Ω = { t } × R d , where Ω is a vertical plane on the augmented space.
Then the hitting time τ equals t deterministically, and the exit distribution equals the marginal distribution of Z t at time t .
See Figure 1-A, for illustration.
2.2 Diffusion Process Tools: Conditioning and h -transform
We introduce some basic tools for diffusion processes, including how to conduct conditioning, and exponential tilting (via h -transform) on diffusion processes.
We apply these tools to the first hitting models we have.
The readers can find related background in Oksendal [2013], Särkkä and Solin [2019].
Assume Z is a general Ito diffusion process in V that is absorbed to Ω , denoted as Ito Ω ( b, σ ) ,
Q ∼ Ito Ω ( b, σ ) :
[Equation]
where b t ( x ) ∈ R
d is the drift term and σ t ( x )
∈ R
d × d is a positive definite diffusion matrix.
We always assume that b and σ are sufficiently regular to yield a unique weak solution of (3).
Conditioning A step in our work is to find the distribution of the trajectories of a process Q conditioned on a future event, e.g., the event of hitting a particular value x at exit, that is, { Z τ = x } .
A notable result is that the conditioned diffusion processes are also diffusion processes.
Given a point x ∈ Ω on the exit surface, the process of Q ( · | Z τ = x ) can be shown to be the law of the following diffusion process [Doob and Doob, 1984, Särkkä and Solin, 2019]:
[Equation]
where q Ω ( x | z ) is the density function of the Poisson kernel Q Ω (d x | Z t = z ) w.r.t.
a reference measure µ Ω on Ω , and σ 2 is the matrix square of σ , and the conditional initial distribution µ 0 | x = Q 0 ( · |
Z τ = x ) is the posterior probability of Z 0 given Z τ = x .
Intuitively, the additional drift term ∇ Z t log p Ω ( x | Z t ) plays the role of steering the process towards the target x , with an increasing magnitude as Z t approaches Ω (because P Ω ( · | Z t = z ) converges to a delta measure centered at x when z approaches Ω ).
This process is known as a diffusion bridge , because it is guaranteed to achieve Z τ = x at the first hitting time with probability one.
Proposition 2.5.
For Q S d , the process conditioned on Z τ = x ∈ S d at exit is
[Equation]
Here the additional drift term (colored in blue) grows to infinity if ∥ Z t ∥ → 1 but ∥
Z t − x ∥ is large, and hence enforces that Z τ
= x when we exit the unit ball.
Proposition 2.6.
For Q B d , the process conditioned on Z τ = x ∈ { 0 , 1 } d at exit is Q B
[Equation]
The additional drift term (colored in blue) enforces that Z τ,i = x
i at the exit time as the drift would be infinite if z i is still far from x i when z i is close to { 0 , 1 } .
Proposition 2.7.
For the fixed time diffusion in Example 2.4, let Q T be the standard Brownian motion d Z t =
d W t stopped at a fixed time t =
T , then Q conditioned on Q T ( Z | Z T = x ) is Q T ( ·| Z τ = x ) :
[Equation]
The additional drift (colored in blue) forces Z T
= x as it grows to infinity if Z t ̸
= x while t → T .
h -Transform Assume we want to modify the Markov process Z such that its exit distribution Q Ω matches the desirable target distribution Π ∗ .
Doob’s h -transform
Doob and Doob [1984] provides a simple general procedure to do so.
Note that by disintegration theorem, we have Q (d Z )
= (cid:82) Q Ω (d x ) Q (
d Z | Z τ = x ) , which factorizes Q into the product of the exit distribution and the conditional process given a fixed exit location Z τ = x .
To modify the exit distribution of Q to Π ∗ , we can simply replace Q Ω with Π ∗ in the disintegration theorem, yielding
[Equation]
where π ∗ =
dΠ ∗ d Q Ω is the Radon–Nikodym derivative (or density ratio) between Π ∗ and Q Ω , and Q Π ∗ is called an h -transform of Q .
Intuitively, Q Π ∗ is the distribution of trajectories Z ∼ Q ( ·| Z τ = x ) when the exit location x is randomly drawn from x ∼ Π ∗ .
We can also view π ∗ ( Z τ ) as an importance score of each trajectory Z based on its terminal state Z τ , and Q Π ∗ is obtained by reweighing (or tilting) the probability of each trajectory based on its score.
If Q is a diffusion process, then Q Π ∗ is also a diffusion process.
In addition, Q Π ∗ is the law of the following diffusion process:
[Equation]
where the initial distribution Q Π ∗ 0 and h Π ∗ in the drift term are defined as
[Equation]
It is clear that h coincides with π ∗ on the boundary, that is, h π ∗ ( x, t ) =
π ∗ ( x ) for all x ∈ Ω , t ≥ 0 .
The name of h -transform comes from the fact that h Π ∗ is a (space-time) harmonic function w.r.t.
Q in the light of a mean value property: h Π ∗ t ( z ) =
E Q [ h Π ∗ t + s ( Z t + s ) | Z t =
z ] , ∀ s, t > 0 .
Q Π ∗ yields a simple variational representation in terms of Kullback–Leibler (KL) divergence.
2.3 Learning First Hitting Diffusion Models
Assume Π ∗ is unknown and we observe it through an i.i.d.
sample { x ( i ) } ni
=1 drawn from Π ∗ .
We want to fit the data with a parametric diffusion process Ito Ω ( s θ , σ ) in V that is absorbing to Ω ,
[Equation]
such that the exit distribution P θ Ω matches the unknown Π ∗ .
Here s θt ( z ) is a deep neural network with input ( z, t ) and parameters θ .
We should design s θ and σ properly to ensure the absorbing property.
The standard approach to estimate Π ∗ is maximum likelihood estimation, which can be viewed as approximately solving min θ KL (Π ∗ || P θ Ω ) .
However, calculating the likelihood of the exit distribution
P θ Ω of a general general diffusion process is computationally intractable.
To address this problem, we fix Q as a “prior” process, and augment the data distribution Π ∗ to the h -transform Q Π ∗ , whose exit distribution Q Π ∗ Ω matches Π ∗ by definition.
Note that we can draw i.i.d.
sample from Q Π ∗ in a “backwar” way: first drawing an exit location x ∼ Π ∗ from the data, and then draw the trajectory Z from Q ( ·| Z τ = x ) with the fixed exit point.
To train a generative model, we train P θ to fit it with the data drawn from Q Π ∗ by maximum likelihood estimation:
[Equation]
where p θ =
d P θ d Q Π ∗ is Radon–Nikodym density function of P θ relative to Q Π ∗ .
By the chain rule of
KL divergence ( ?? )
, we have KL (Π ∗ || P θ Ω ) ≤ KL ( Q Π ∗
|| P θ ) .
Therefore, if minimizing the KL divergence allows us to achieve P θ ≈ Q Π ∗ , we should also have P θ Ω ≈ Q Π ∗ Ω = Π ∗ .
Using Girsanov theorem [Liptser and Shiriaev, 1977], we can calculate the density function p θ and hence the loss function.
Proposition 2.8.
Assume Q in (3) , and P θ in (12) are absorbing to Ω .
We have
[Equation]
where b t ( z | x ) :
= b t ( z ) + σ 2 t ( z ) ∇ z log p Ω ( x | z ) is the drift of the conditioned process Q ( ·| Z τ = x ) in (4) , and p θ 0 is the probability density function of the initial distribution P θ 0 .
In addition, θ ∗ achieves the global minimum of L ( θ ) if
[Equation]
Therefore, the optimal drift term s θ ∗ t should match the conditional expectation of b t ( z | x ) with x ∼ Q Ω ( ·| Z t = z ) , which coincides with the drift of Q Π ∗ in (9).
The initial distribution of P θ should obviously match the initial distribution of Q Π ∗ .
In practice, we recommend eliminating the need of estimating P θ 0 by starting Q from a deterministic point Z
=
z 0 , in which case P θ should initialize from the same deterministic point.