D ENOISING D IFFUSION S AMPLERS
A BSTRACT
Denoising diffusion models are a popular class of generative models providing state-of-the-art results in many domains.
One adds gradually noise to data using a diffusion to transform the data distribution into a Gaussian distribution.
Samples from the generative model are then obtained by simulating an approximation of the time-reversal of this diffusion initialized by Gaussian samples.
Practically, the intractable score terms appearing in the time-reversed process are approximated using score matching techniques.
We explore here a similar idea to sample approximately from unnormalized probability density functions and estimate their normalizing constants.
We consider a process where the target density diffuses towards a Gaussian.
Denoising Diffusion Samplers (DDS) are obtained by approximating the corresponding time-reversal.
While score matching is not applicable in this context, we can leverage many of the ideas introduced in generative modeling for Monte Carlo sampling.
Existing theoretical results from denoising diffusion models also provide theoretical guarantees for DDS.
We discuss the connections between DDS, optimal control and Schr ¨ odinger bridges and finally demonstrate
DDS experimentally on a variety of challenging sampling tasks.
I NTRODUCTION
Let π be a probability density on R d of the form
[Equation]
where γ : R d → R + can be evaluated pointwise but the normalizing constant Z is intractable.
We are here interested in both estimating Z and obtaining approximate samples from π .
A large variety of Monte Carlo techniques has been developed to address this problem.
In particular Annealed Importance Sampling (AIS) (Neal, 2001) and its Sequential Monte Carlo (SMC) extensions (Del Moral et al., 2006) are often regarded as the gold standard to compute normalizing constants.
Variational techniques are a popular alternative to Markov Chain Monte Carlo (MCMC) and SMC where one considers a flexible family of easy-to-sample distributions q θ whose parameters are optimized by minimizing a suitable metric, typically the reverse Kullback–Leibler discrepancy
KL( q θ || π ) .
Typical choices for q θ include mean-field approximation (Wainwright & Jordan, 2008) or normalizing flows (Papamakarios et al., 2021).
To be able to model complex variational distributions, it is often useful to model q θ ( x ) as the marginal of an auxiliary extended distribution; i.e. q θ ( x ) = (cid:82) q θ ( x, u )d u .
As this marginal is typically intractable, θ is then learned by minimizing a discrepancy measure between q θ ( x, u ) and an extended target p θ ( x, u )
=
π ( x ) p θ ( u | x ) where p θ ( u | x ) is an auxiliary conditional distribution (Agakov & Barber, 2004).
Over recent years, Monte Carlo techniques have also been fruitfully combined to variational techniques.
For example, AIS can be thought of a procedure where q θ ( x, u ) is the joint distribution of a Markov chain defined by a sequence of MCMC kernels whose final state is x while p θ ( x, u ) is the corresponding AIS extended target (Neal, 2001).
The parameters θ of these kernels can then be optimized by minimizing KL( q θ || p θ ) using stochastic gradient descent (Wu et al., 2020; Geffner &
Domke, 2021; Thin et al., 2021; Zhang et al., 2021; Doucet et al., 2022; Geffner & Domke, 2022).
Instead of following an AIS-type approach to define a flexible variational family, we follow here an approach inspired by Denoising Diffusion Probabilistic Models (DDPM), a powerful class of generative models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021c).
In this context, one adds noise progressively to data using a diffusion to transform the complex data distribution into a Gaussian distribution.
The time-reversal of this diffusion can then be used to transform a
Gaussian sample into a sample from the target.
While superficially similar to Langevin dynamics, this process mixes fast even in high dimensions as it inherits the mixing properties of the forward diffusion (De Bortoli et al., 2021, Theorem 1).
However, as the time-reversal involves the derivatives of the logarithms of the intractable marginal densities of the forward diffusion, these so-called scores are practically approximated using score matching techniques.
If the score estimation error is small, the approximate time-reversal still enjoys remarkable theoretical properties (De Bortoli, 2022; Chen et al., 2022; Lee et al., 2022).
These results motivate us to introduce Denoising Diffusion Samplers (DDS).
Like DDPM, we consider a forward diffusion which progressively transforms the target π into a Gaussian distribution.
This defines an extended target distribution p ( x, u ) =
π ( x ) p ( u | x ) .
DDS are obtained by approximating the time-reversal of this diffusion using a process of distribution q θ ( x, u ) .
What distinguishes DDS from DDPM is that we cannot simulate sample paths from the diffusion we want to time-reverse, as we cannot sample its initial state x from π .
Hence score matching ideas cannot be used to approximate the score terms.
We focus on minimizing KL( q θ || p ) , equivalently maximizing an Evidence Lower Bound (ELBO), as in variational inference.
We leverage a representation of this KL discrepancy based on the introduction of a suitable auxiliary reference process that provides low variance estimate of this objective and its gradient.
We can exploit the many similarities between DDS and DDPM to leverage some of the ideas developed in generative modeling for Monte Carlo sampling.
This includes using the probability flow ordinary differential equation (ODE) (Song et al., 2021c) to derive novel normalizing flows and the use of underdamped Langevin diffusions as a forward noising diffusion (Dockhorn et al., 2022).
The implementation of these samplers requires designing numerical integrators for the resulting stochastic differential equations (SDE) and ODE.
However, simple integrators such as the standard Euler–Maryuama scheme do not yield a valid ELBO in discrete-time.
So as to guarantee one obtains a valid ELBO, DDS relies instead on an integrator for an auxiliary stationary reference process which preserves its invariant distribution as well as an integrator for the approximate time-reversal inducing a distribution absolutely continuous w.r.t.
the distribution of the discretized reference process.
Finally we compare experimentally DDS to AIS, SMC and other state-of-the-art
Monte Carlo methods on a variety of sampling tasks.
D ENOISING D IFFUSION S AMPLERS : C ONTINUOUS T IME
We start here by formulating DDS in continuous-time to gain insight on the structure of the timereversal we want to approximate.
We introduce C = C
([0 , T ] , R d ) the space of continuous functions from [0 , T ] to R d and B ( C ) the Borel sets on C .
We will consider in this section path measures which are probability measures on ( C , B ( C )) , see L´eonard (2014a) for a formal definition.
Numerical integrators are discussed in the following section.
2. F ORWARD DIFFUSION AND ITS TIME - REVERSAL
Consider the forward noising diffusion given by an Ornstein–Uhlenbeck (OU) process
[Equation]
where ( B t ) t ∈ [0 ,T ] is a d -dimension Brownian motion and t → β t is a non-decreasing positive function.
This diffusion induces the path measure P on the time interval [0 , T ] and the marginal density of x t is denoted p t .
The transition density of this diffusion is given by p t | 0 ( x t | x 0 ) =
N ( x t ; √ 1 − λ t x 0 , σ 2 λ t I ) where λ t = 1 − exp( − 2 (cid:82) t 0 β s d s ) .
From now on, we will always consider a scenario where (cid:82) T
β s d s ≫ 1 so that p T ( x T ) ≈
N ( x T ; 0 , σ 2 I ) .
We can thus think of (2) as transporting approximately the target density π to this
Gaussian density.
From (Haussmann & Pardoux, 1986), its time-reversal ( y t ) t ∈ [0 ,T ] =
( x T − t ) t ∈ [0 ,T ] , where equality is here in distribution, satisfies
[Equation]
where ( W t ) t ∈ [0 ,T ] is another d -dimensional Brownian motion.
By definition this time-reversal starts from y 0 ∼ p T ( y 0 )
≈ N ( y 0 ; 0 , σ 2 I ) and is such that y T ∼ π .
This suggests that if we could approximately simulate the diffusion (3), then we would obtain approximate samples from π .
However, putting this idea in practice requires being able to approximate the intractable scores ( ∇ ln p t ( x ))
t ∈ [0 ,T ] .
To achieve this, DDPM rely on score matching techniques (Hyv¨arinen, 2005;
Vincent, 2011) as it is easy to sample from (2).
This is impossible in our scenario as sampling from (2) requires sampling x 0 ∼ π which is impossible by assumption.
2. R EFERENCE DIFFUSION AND VALUE FUNCTION
In our context, it is useful to introduce a reference process defined by the diffusion following (2) but initialized at p ref 0 ( x 0 )
=
N ( x 0 ; 0 , σ 2 I ) rather than π ( x 0 ) thus ensuring that the marginals of the resulting path measure P ref all satisfy p ref t ( x t )
=
N ( x t ; 0 , σ 2 I ) .
From the chain rule for KL for path measures (see e.g. (L´eonard, 2014b, Theorem 2.4) and (De Bortoli et al., 2021, Proposition 24)), one has KL( Q||P ref ) =
KL( q 0 || p ref 0 )
+ E x 0 ∼ q
[KL( Q ( ·| x 0 ) ||P ref ( ·| x 0 ))] .
Thus P can be identified as the path measure minimizing the KL discrepancy w.r.t.
P ref over the set of path measures Q with marginal q 0 ( x 0 ) =
π ( x 0 ) at time t = 0 , i.e. P = arg min Q { KL( Q||P ref ) :
q 0 = π } .
A time-reversal representation of P ref is given by ( y t ) t ∈ [0 ,T ] =
( x T − t ) t ∈ [0 ,T ] satisfying
[Equation]
As β T − t y t + 2 σ 2 ∇ ln p ref T − t ( y t ) =
− β T − t y t , we can rewrite the time-reversal (3) of P as
[Equation]
where ϕ t ( x ) =
p t ( x ) /p ref t ( x ) is a so-called value function which can be shown to satisfy a Kolmogorov equation such that ϕ t ( x t ) =
E P ref [ ϕ 0 ( x 0 ) | x t ] , the expectation being w.r.t.
P ref .
2. L EARNING THE TIME - REVERSAL
To approximate the time-reversal (3) of P , consider a path measure Q θ whose time-reversal is induced by
[Equation]
so that y t ∼ q θT − t .
To obtain s θ ( t, x ) ≈ ∇ ln p t ( x ) , we parameterize s θ ( t, x ) by a neural network whose parameters are obtained by minimizing
[Equation]
where we have used the chain rule for KL then Girsanov’s theorem (see e.g. (Klebaner, 2012)).
This expression of the KL is reminiscent of the expression obtained in (Song et al., 2021b, Theorem 1) in the context of DDPM.
However, the expectation appearing in (7) is here w.r.t.
Q θ and not w.r.t.
P and we cannot get rid of the intractable scores ( ∇ ln p t ( x ))
t ∈ [0 ,T ] using score matching ideas.
Instead, taking inspiration from (5), we reparameterize the time-reversal of Q θ using
[Equation]
i.e. f θ
[Equation]
t, x ) =
s θ
[Equation]
t, x ) − ∇ ln p ref
[Equation]
x ) =
s θ
[Equation]
t, x ) + x/σ 2 .
This reparameterization allows us to express KL( Q θ ||P ) in a compact form.
Proposition 1.
The Radon–Nikodym derivative d Q θ d P ref ( y [0 ,T ] ) satisfies under Q θ
[Equation]
From the identity KL( Q θ ||P ) =
KL( Q θ ||P ref )
+ E y T
[Equation]
[ln
[Equation]
y T
[Equation]
T )
[Equation]
] , it follows that
[Equation]
For θ minimizing (10), approximate samples from π can be obtained by simulating (8) and returning y T ∼ q θ 0 .
We can obtain an unbiased estimate of Z via the following importance sampling identity
[Equation]
where the second term can be computed directly from (9) and y [0 ,T ] is obtained by solving (8).
2. C ONTINUOUS - TIME NORMALIZING FLOW
To approximate the log likelihood of DDPM, it was proposed by Song et al.
(2021c) to rely on the probability flow ODE.
This is an ODE admitting the same marginal distributions ( p t ) t ∈ [0 ,T ] as the noising diffusion given by d x t = − β t (cid:8) x t
+ σ 2 ∇ ln p t ( x t ) (cid:9) d t .
We use here this ODE to design a continuous-time normalizing flow to sample from π .
For θ such that f θ ( t, x ) ≈ ∇ ln ϕ t ( x ) (obtained by minimization of the KL in (10)), we have x t + σ 2 ∇ ln p t ( x ) ≈ σ 2 f θ ( t, x ) .
So it is possible to sample approximately from π by integrating the following ODE over [0 , T ]
[Equation]
If denote by ¯ q θT − t the distribution of y t then y T ∼ ¯ q θ 0 is an approximate sample from π .
We can use this normalizing flow to obtain an unbiased estimate of Z using importance sampling ˆ
Z = γ ( y T ) / ¯ q θ
( y T ) for y T ∼ ¯ q θ 0 .
Indeed, contrary to the proposal q θ 0 induced by (8), we can compute pointwise ¯ q θ 0 using the instantaneous change of variables formula (Chen et al., 2018)
such that ln ¯ q θ
( y T ) =
ln N ( y 0 ; 0 , σ 2 I ) − σ 2 (cid:82) T 0 β T − t ∇· f θ ( T − t, y t )d t and where ( y t ) t ∈ [0 ,T ] arises from the ODE (12).
The expensive divergence term can be estimated using the Hutchinson estimator (Grathwohl et al., 2018; Song et al., 2021c).
2. E XTENSION TO U NDERDAMPED L ANGEVIN D YNAMICS
For DDPM, it has been proposed to extend the original state x ∈ R d by a momentum variable p ∈ R d .
One then diffuses the data distribution augmented by a Gaussian distribution on the momentum using an underdamped Langevin dynamics targeting N ( x ; 0 , σ 2 I ) N ( m ; 0 , M ) where M is a positive definite mass matrix (Dockhorn et al., 2022).
It was demonstrated empirically that the resulting scores are smoother, hence easier to estimate and this leads to improved performance.
We adapt here this approach to Monte Carlo sampling; see Section B for more details.
We diffuse ¯ π ( x, m ) =
π ( x ) N ( m ; 0 , M ) using an underdamped Langevin dynamics, i.e.
[Equation]
The resulting path measure on [0 , T ] is denoted P and the marginal of ( x t , m t ) is denoted η t .
Here, the reference process P ref is defined by the diffusion (13) initialized using x 0 ∼
N (0 , σ 2 I ) , m 0 ∼ N (0 , M ) .
This ensures that η ref t ( x t , m t ) =
N ( x t ; 0 , σ 2 I ) N ( m t ; 0 , M ) for all t and the time-reversal process ( y t , n t ) t ∈ [0 ,T ] =
( x T − t , m T − t ) t ∈
[0 ,T ] (where equality is in distribu- tion) of this stationary diffusion satisfies
[Equation]
Using manipulations identical to Section 2.2, the time-reversal of P can be also be written for ϕ t ( x, m ) :
=
η t ( x, m ) /η ref t ( x, m ) as d y t = − M − 1 n t d t and
[Equation]
[Equation]
To approximate P , we will consider a parameterized path measure Q θ whose time reversal is defined for ( y 0 , n 0 ) ∼ N ( y 0 ; 0 , σ 2 I ) N ( n 0 ; 0 , M ) by d y t = − M − 1 n t d t and
[Equation] [Equation]
We can then express the KL of interest in a compact way similar to Proposition
1: Input: σ > 0 , γ : R d → R + , ( β k )
Kk =1 ∈ ( R + ) K , θ ∈ R p , λ >2: for i
. , training iterations do 3: for n
, N do # Iterate over batch size.
Sample y 0 ,n ∼ N (0 , σ 2 I ) and set r
=5: for k = 0 , . .
, K − 1 do # Iterate over integration steps.
6: λ
[Equation]
:=
[Equation]
− α K − k 7: y k
[Equation]
1 ,n =
[Equation]
1 − α K
[Equation]
k y k,
[Equation]
σ
[Equation]
k f
[Equation]
( K
[Equation]
k, y
[Equation]
,n )
[Equation]
K
[Equation]
k,n , ε k,n i.i.d.
[Equation]
(0 , I ) 8: r k
[Table]
1 ,n = r k,
[Equation]
[Table] [Table]
[Equation]
f
[Table]
( K
[Table]
[Equation]
, y
[Table]
,n ) ||9: end for 10: θ ← θ − λ ∇ θ 1 N (cid:80) Nn =
[Table]
(cid:16) r K,n
[Table]
ln (cid:16) N ( y K,n ;0 ,
[Table]
I )
[Table]
( y K,n ) (cid:17)(cid:17) 11: end for 12: end for 13: Return: θ
D ENOISING D IFFUSION S AMPLERS : D ISCRETE -T IME
We introduce here discrete-time integrators for the SDEs and ODEs introduced in Section 2.
Contrary to DDPM, we not only need an integrator for Q θ but also for P ref to be able to compute an approximation of the Radon–Nikodym derivative d Q θ / d P ref .
Additionally this integrator needs to be carefully designed as explained below to preserve an ELBO.
For sake of simplicity, we consider a constant discretization step δ > 0 such that K = T/δ is an integer.
In the indices, we write k for t k = kδ to simplify notation; e.g. x k for x kδ .
3. I NTEGRATOR FOR P ref
Proposition 1 shows that learning the time reversal in continuous-time can be achieved by minimising the objective KL( Q θ ||P ) given in (10).
This expression is obtained by using the identity
KL( Q θ ||P ) =
KL( Q θ ||P ref )
+ E y T ∼ q θ
[ln (cid:0) p ref 0 ( y T ) /p
( y T ) (cid:1) ] .
This is also equivalent to maximizing the corresponding ELBO, E Q θ
[ln ˆ Z ] , for the unbiased estimate ˆ Z of Z given in (11)
An Euler–Maruyama (EM) discretisation of the corresponding SDEs is obviously applicable.
However, it is problematic as established below.
Proposition 2.
Consider integrators of P ref and Q θ leading to approximations p ref ( x 0: K ) and q θ ( x 0: K ) .
Then KL( q θ || p ref ) + E y K ∼ q θ
[ln (cid:0) p ref 0 ( y K ) /π
( y K ) (cid:1) ] is guaranteed to be non-negative and
E Q θ [ln ˆ Z ] is an ELBO for ˆ Z
=
( γ ( y K ) p ref ( y 0: K )) /
( N ( y K ; 0 , σ 2 ) q θ ( y 0: K ))
if one has p ref K ( y K ) =
p ref0 ( y K ) .
The EM discretisation of P ref does not satisfy this condition.
A direct consequence of this result is that the estimator for ln Z that uses the EM discretisation can be such that E Q
θ [ln ˆ Z ] ≥ ln Z as observed empirically in Table 5.
To resolve this issue, we derive in the next section an integrator for P ref which ensures that p ref k ( y ) =
p ref0 ( y ) for all k .
3. O RNSTEIN –U HLENBECK
We can integrate exactly (4).
This is given by y 0 ∼ N (0 , σ 2 I ) and
[Equation]
for α k = 1 − exp (cid:16) − 2 (cid:82) kδ
( k − 1)
δ β s d s (cid:17) 2 .
This defines the discrete-time reference process.
We propose the following exponential type integrator (De Bortoli, 2022) for (8) initialized using y 0 ∼ N (0 , σ 2 I )
[Equation]
[Figure]
where ε k i.i.d.
∼
N (0 , I ) .
Equations (17) and (18) define the time reversals of the reference process p ref ( x 0: K ) and proposal q θ ( x 0: K ) .
For its time reversal, we write q θ ( y 0: K ) =
N ( y 0 ; 0 , σ 2 I ) (cid:81)
Kk =1 q θk − 1 | k
( y K − k +1 | y K − k ) abusing slightly notation, and similarly for p ref ( y 0: K ) .
We will be relying on the discrete-time counterpart of (9).
Proposition 3.
The log density ratio between q θ ( y 0: K ) and p ref ( y 0: K ) satisfies for y 0:
K ∼ q θ ( y 0: K ) ln (cid:18) q θ ( y 0: K ) p
[Equation] [Equation]
(19) where λ k :=
1 − √ 1 − α k and ε k defined through (18) is such that ε k i.i.d.
∼ N (0 , I ) .
In particular, one obtains from KL( q θ || p ) =
KL( q θ |
[Equation]
We compute an unbiased gradient of this objective using the reparameterization trick and the JAX software package (Bradbury et al., 2018).
The training procedure is summarized in Algorithm 1 in
Appendix D. Unfortunately, contrary to DDPM, q θk | K is not available in closed form for k < K −so we can neither mini-batches over the time index k without having to simulate the process until the minimum sampled time nor reparameterize x k as in Ho et al.
Once we obtain the parameter θ minimizing (20), DDS samples from q θ using (18).
The final sample y K has a distribution q θapproximating π by design.
By using importance sampling, we obtain an unbiased estimate of the normalizing constant ˆ Z =
( γ ( y K ) p ref ( y 0: K )) /
( N ( y K ; 0 , σ 2 ) q θ ( y 0: K ))
for y 0: K ∼ q θ ( y 0: K ) .
Finally Appendices B.4 and B.6 extend this approach to provide a similar approach to discretize the underdamped dynamics proposed in Section 2.5.
In this context, the proposed integrators rely on a leapfrog scheme (Leimkuhler & Matthews, 2016).
3. T HEORETICAL G UARANTEES
Motivated by DDPM, bounds on the total variation between the target distribution and the distribution of the samples generated by a time-discretization of an approximate time reversal of a forward noising diffusion have been first obtained in (De Bortoli et al., 2021) then refined in (De Bortoli, 2022; Chen et al., 2022; Lee et al., 2022) and extended to the Wasserstein metric.
These results are directly applicable to DDS because their proofs rely on assumptions on the score approximation error but not on the way these approximations are learned, i.e. via score matching for DDPM and reverse KL for DDS.
For example, the main result in Chen et al.
(2022) shows that if the true scores are L -Lipschitz, the L 2 ( p t ) error on the scores is bounded and other mild integrability assumptions then DDS outputs samples ϵ -close in total variation to π in O
( L 2 d/ϵ 2 ) time steps.
As pointed out by the authors, this matches state-of-the-art complexity bounds for Langevin Monte Carlo algorithm for sampling targets satisfying a log-Sobolev inequality without having to make any log-concavity assumption on π .
However, the assumption on the approximation error for score estimates is less realistic for DDS than DDPM as we do not observe realizations of the forward diffusion.
[Figure]
3. I NTEPRETATION , R ELATED W ORK AND E XTENSIONS
DDS as KL and path integral control.
The reverse KL we minimize can be expressed as
[Equation]
This objective is a specific example of KL control problem (Kappen et al., 2012).
In continuoustime, (10) corresponds to a path integral control problem; see e.g. (Kappen & Ruiz, 2016).
Heng et al.
(2020) use KL control ideas so as to sample from a target π and estimate Z .
However, their algorithm relies on p ref ( x 0: K ) being defined by a discretized non-homogeneous Langevin dynamics such that p K ( x K ) is typically not approximating a known distribution.
Additionally it approximates the value functions ( ϕ k )
Kk =0 using regression using simple linear/quadratic functions.
Finally it relies on a good initial estimate of p ( x 0: K ) obtained through SMC.
This limits the applicability of this methodology to a restricted class of models.
Connections to Schr¨odinger Bridges.
The Schr ¨ odinger Bridge (SB) problem (L ´ eonard, 2014a;
De Bortoli et al., 2021) takes the following form in discrete time.
Given a reference density p ref ( x 0: K ) , we want to find the density p sb ( x 0: K ) s.t.
p
sb = arg min q { KL( q || p ref ) :
q 0 = µ 0 , q K = µ K } where µ 0 , µ K are prescribed distributions.
This problem can be solved using iterative proportional fitting (IPF) which is defined by the following recursion with initialization p 1 = p ref p
[Equation]
=
[Equation]
Consider the SB problem where µ 0 ( x 0 )
=
π ( x 0 ) , µ K ( x K ) =
N ( x K ; 0 , σ 2 I ) and the time-reversal of p ref ( x 0: K ) is defined through (17).
In this case, p 2 = p corresponds to the discrete-time version of the noising process and p 3 to the time-reversal of p but initialized at µ K instead of p K .
This is the process DDS is approximating.
As p K ≈ µ K for K large enough, we have approximately p sb
≈ p 3 ≈ p 2 .
We can thus think of DDS as approximating the solution to this SB problem.
Consider now another SB problem where µ 0 ( x 0 )
=
π ( x 0 ) , µ K ( x K ) =
δ
( x K ) and p ref ( x 0: K ) =
δ
( x K ) (cid:81) K − 1 k =
N ( x k ; x k +1 , δσ 2 I ) , i.e. p ref is a pinned Brownian motion running backwards in time.
This SB problem was discussed in discrete-time in (Beghi, 1996) and in continuous-time in (F¨ollmer, 1984; Dai Pra, 1991; Tzen & Raginsky, 2019).
In this case, p 2 ( x 0: K ) =
π ( x 0 ) p ref ( x 1: K | x 0 ) is a modified “noising” process that transports π to the degenerate measure
δ
and it is easily shown that p sb = p 2 .
Sampling from the time-reversal of this measure would generate samples from π starting from δ 0 .
Algorithmically, Zhang et al.
(2021) proposed approximating this time-reversal by some projection on some eigenfunctions.
In parallel, Barr et al.
(2020), Vargas et al.
and Zhang &
Chen (2022) approximated this SB by using a neural network parameterization of the gradient of the logarithm of the corresponding value function trained by minimizing a reverse KL.
We will adopt the Path Integral Sampler (PIS) terminology proposed by Zhang & Chen (2022) for this approach.
DDS can thus be seen as an alternative to PIS which relies on a reference dynamics corresponding to an overdamped or underdamped OU process instead of a pinned Brownian motion.
Theoretically the drift of the resulting time-reversal for DDS is not as a steep as for PIS (see Appendix A.2) and empirically this significantly improves numerical stability of the training procedure; see Figure 1.
The use of a pinned Brownian motion is also not amenable to the construction of normalizing flows.
[Figure]
Forward KL minimization.
In Jing et al.
(2022), diffusion models ideas are also use to sample unnormalized probability densities.
The criterion being minimized therein is the forward KL as in
DDPM, that is KL( p || q θ ) .
As samples from p are not available, an importance sampling approximation of p based on samples from q θ is used to obtain an estimate of this KL and its gradient w.r.t.
θ .
The method is shown to perform well in low dimensional examples but is expected to degrade significantly in high dimension as the importance sampling approximation will be typically poor in the first training iterations.
E XPERIMENTS
We present here experiments for Algorithm 1.
In our implementation, f θ follows the PIS-GRAD network proposed in (Zhang & Chen, 2022): f θ ( k, x ) =
NN 1 ( k, x ; θ ) + NN 2 ( k ; θ ) ⊙ ∇ ln π ( x ) .
Across all experiments we use a two layer architecture with 64 hidden units each (for both networks), as in Zhang & Chen (2022), with the exception of the NICE (Dinh et al., 2014) target where we use 3 layers with 512, 256, and 64 units respectively.
The final layers are initialised to 0 in order to make the path regularisation term null.
We use α 1 / 2 k ∝ α 1 / 2 max cos 2 (cid:16) π 2 1 − k/K
+ s 1+ s (cid:17) with s = 0 .
as in (Nichol & Dhariwal, 2021).
We found that detaching the target score stabilised optimization in both approaches without affecting the final result.
We adopt this across experiments, an ablation of this feature can be seen in Appendix C.9.1.
Across all tasks we compare DDS to SMC (Del Moral et al., 2006;
Zhou et al., 2016), PIS (Barr et al., 2020; Vargas et al., 2021; Zhang & Chen, 2022), and Mean Field-VI (MF-VI) with a Gaussian variational distribution.
We also compare DDS to AIS (Neal, 2001) and optimized variants of
AIS using score matching (MCD) (Doucet et al., 2022; Geffner & Domke, 2022) for two standard
Bayesian models.
Finally we explore a task introduced in (Doucet et al., 2022) that uses a pre-trained normalising flow as a target.
Within this setting we propose a benchmarking criterion that allows us to assess mode collapse in high dimensions and explore the benefits of incorporating inductive biases into f θ .
We carefully tuned the hyper-parameters of all algorithms (e.g. step size, diffusion coefficient, and such), details can be found in Appendix C.2.
Finally training time can be found in
Appendix C.5.
Additional experiments for the normalizing flows are presented in Appendix C. and for the underdamped approach in Appendix C.12.
We note that these extensions did not bring any benefit compared to Algorithm 1.
4. B ENCHMARKING T ARGETS
We first discuss two standard target distributions which are often used to benchmark methods; see e.g. (Neal, 2003; Arbel et al., 2021; Heng et al., 2020; Zhang & Chen, 2022).
Results are presented in Figure 2.
Funnel Distribution: This 10-dimensional challenging distribution is given by γ ( x 1:10 ) =
N ( x 1 ; 0 , σ 2 f ) N ( x 2:10 ; 0 , exp( x 1 ) I ) where σ
= 9 (Neal, 2003)
Log Gaussian Cox process: This model arises in spatial statistics (Møller et al., 1998).
We use a d = M × M = 1600 grid, resulting in the unnormalized target density γ ( x ) =
N ( x ; µ, K ) (cid:81) i ∈
[1: M ] 2 exp( x i y
i − a exp( x i )) .
4. B AYESIAN M ODELS
We explore two standard Bayesian models and compare them with standard AIS and MCD benchmarks (See Appendix C.8) in addition to the SMC and VI benchmarks presented so far.
Results for
Ionosphere can be found in the 3rd pane of Figure 2 whilst SONAR and Brownian are in Figure 3.
Logistic Regression: We set x ∼ N (0 , σ 2 w I ) , y i ∼ Bernoulli(sigmoid( x ⊤ u i )) .
This Bayesian logistic model is evaluated on two datasets, Ionosphere ( d = 32 ) and Sonar ( d = 61 ).
Brownian Motion: We consider a discretised Brownian motion with a Gaussian observation model and a latent volatility as a time series model, d = 32 .
This model, proposed in the software package developed by Sountsov et al.
(2020), is specified in more detail in Appendix C.3.
4. M ODE C OLLAPSE I N H IGH D IMENSIONS
Normalizing Flow Evaluation: Following Doucet et al.
we train NICE
(Dinh et al., 2014) on a down-sampled d = 14 × 14 variant of MNIST (LeCun & Cortes, 2010) and use the trained model as our target.
As we can generate samples from our target, we evaluate the methods samples by measuring the Sinkhorn distance between true and estimated samples.
This evaluation criteria allows to asses mode collapse for samplers in high dimensional settings.
Results can be seen in
Table 1 and the third pane of Figure 3.
[Table]
D ISCUSSION
We have explored the use of DDPM ideas to sample unnormalized probability distributions and estimate their normalizing constants.
The DDS in Algorithm 1 is empirically competitive with state-of-the-art SMC and numerically more stable than PIS.
This comes at the cost of a non-negligible training time compared to SMC.
When accounting for it, SMC often provide better performance on simple targets.
However, in the challenging multimodal NICE example, even a carefully tuned SMC sampler using Hamiltonian
Monte Carlo transitions was not competitive to DDS.
This is despite the fact that DDS (and PIS) are prone to mode dropping as any method relying on the reverse KL.
We have also investigated normalizing flows based on the probability flow ODE as well as DDS based on an underdamped dynamics.
Our experimental results were disappointing in both cases in high dimensional scenarios.
We conjecture that more sophisticated numerical integrators need to be developed for the normalizing flows to be competitive and that the neural network parameterization used in the underdamped scenario should be improved and better leverage the structure of the logarithmic derivative of the value functions.
Overall DDS are a class of algorithms worth investigating and further developing.
There has much work devoted to improving successfully DDPM over the past two years including, among many others, modified forward noising mechanisms (Hoogeboom & Salimans, 2022), denoising diffusion implicit models (Song et al., 2021a) and sophisticated numerical integrators (Karras et al., 2022).
It is highly likely that some of these techniques will lead to more powerful DDS.
Advances in KL and path integral control might also be adapted to DDS to provide better training procedures; see e.g.
Thalmeier et al.
A A PPENDIX A. P ROOF OF P ROPOSITION
The Radon–Nikodym derivative expression (9) directly follows from an application of Girsanov theorem (see e.g. Klebaner (2012)).
The path measures P and P ref are induced by two diffusions following the same dynamics but with different initial conditions so that
[Equation]
So it follows directly that KL( Q θ |
[Equation]
P
[Equation]
Note we apply Girsanov theorem to the time reversals of the path measures Q θ and P using that the Radon–
Nikodym between two path measures and their respective reversals are the same.
As W t is a Brownian motion under Q θ , the final expression (10) for KL( Q θ ||P ) follows.
A. C OMPARING DDS AND PIS D RIFTS FOR G AUSSIAN TARGETS
The optimal drift for both PIS and DDS can be expressed in terms of logarithmic derivative of the value function plus additional terms:
[Equation]
where ϕ t is the corresponding value function for DDS or PIS respectively.
Recall that ln ϕ t ( x ) =
ln p t ( x ) − ln p ref t ( x ) .
For a target π ( x ) =
N ( x ; µ, Σ) , p ref t ( x ) =
N ( x t ; µ ref t , b ref t I ) and p t | 0 ( x t | x 0 ) =
N ( x t ; a t x 0 , b t I ) , we obtain
[Equation]
so
[Equation]
Corollary 1.
For DDS, we have µ ref t = 0 , b ref t = σ 2 I , a t = √ 1 − λ t , b t = σ 2 λ t
I with λ t = 1 − exp( − 2 (cid:82) t 0 β
s d s ) .
So, for example, for σ = β t = 1 and Σ
=
I , we have
[Equation]
Corollary 2.
For PIS, we have a reference process which is the time-reversal of a pinned Brownian motion with µ ref t = 0 , b ref t = σ 2 ( T − t )
I , a t = T − t T , b t = σ 2 t ( T − t ) T .
For σ = 1 , Σ
= I , we obtain
[Equation]
In particular b PIS ( x, t ) ≈ xt − ( x − µ ) as t → 0 .
Hence for PIS the drift function explodes close to the origin compared to DDS.
This explosion holds for any target π as it is only related to the fact that p ref t concentrates to δ 0 as t → T .
This makes it harder to approximate for a neural network.
Additionally when discretizing the resulting diffusion, this means that smaller discretization steps should be used close to t = 0 .
A. P ROOF OF P ROPOSITION
Proof.
Let us express the discrete time version of the reference process as
[Equation]
and denote by p ref k the corresponding marginal density of y k which satisfies
[Equation]
The backward decomposition of this joint distribution is given by
[Equation]
where p ref
[Equation]
If our chosen integrator induces a transition kernel p ref k +1 | k
( y k +1 | y k ) which is such that p ref K ( y K ) =
p ref0 ( y K )
[Equation]
then
[Equation]
31) is a valid (normalised) probability density.
Hence it follows that
[Equation]
If the integrator does not preserve the marginals we have that
[Equation]
This is not a probability density and thus the objective is no longer guaranteed to be positive and consequently the expectation of our estimator of ln Z will not be necessarily a lower bound for ln Z .
Finally simple calculations show that the Euler discretisation does not preserve the invariant distribution of P ref in DDS.
A. P ROOF OF P ROPOSITION
We have ln
[Equation]
Now by construction, we have from (17) that p ref
[Equation]
−
[Equation]
y K
[Equation]
) =
N ( y K − k +1 ; √ 1 − α k y K − k , σ 2 α k I ) and from (18) we obtain q θk − 1 | k
( y K − k +1 | y K − k ) =
N ( y K − k +1 ; √ 1 − α k y K − k + 2 σ 2 (1 − √ 1 − α k ) f
θ ( k, y K − k ) , σ 2 α k I ) .
It follows that
[Equation]
where we have exploited the fact that q K ( y 0 )
=
p ref K ( y 0 ) =
N ( y
[Equation] [Equation]
σ 2 I
[Equation]
Now using ϵ k :=
[Equation] [Equation]
we can rewrite (35) as
[Equation]
Now (19) follows directly from (34) and (66).
Note we have for δ ≪ 1 that 2 σ 2 (1 −√ 1 − α k ) 2 α k ≈
σ 2 β k δ and 2 σ (1 −√
− α k ) √
α k ≈ √
β
kδ δ as expected from (9).
Finally the final expression (20) of the KL follows now from the fact that E q θ [ f θ ( k, y K − k ) ⊤ ϵ k ] =
[Equation]
A. A LTERNATIVE K ULLBACK –L EIBLER DECOMPOSITION
A KL decomposition similar in spirit to the one developed for DDPM (Ho et al., 2020) can be be derived.
It leverages the fact that
[Equation]
for
[Equation]
=
[Equation]
Proposition 4.
The reverse Kullback–Leibler discrepancy KL( q θ || p ) satisfies
[Equation]
+ KL( q θ 0 | 1 ( x 0 | x 1 )
||
π ( x 0 ))] .
So for q θ k − 1 | k ( x k − 1 | x k ) =
N ( x k − 1 ; √ 1 − β k x k + σ 2 β k f θ
( k, x k ) , σ 2 β k I ) , the terms
KL( q θk − 1 | k ( x k − 1 | x k ) ||
p k − 1 | 0 ,k ( x k − 1 | x 0 , x k )) are KL between two Gaussian distributions and can be calculated analytically.
We found this decomposition to be numerically unstable and prone to diverging in our reverse KL setting.
Proof.
The reverse KL can be decomposed as follows
[Equation]
where
[Equation]
Now using the identity for k ≥
[Equation]
we can rewrite L ( θ ) as
[Equation]
|
[Equation]
(cid:17)
[Equation]
| 1 ( x 0 | x 1 ) (cid:105)
[Equation]
The result now follows directly.
B U NDERDAMPED L ANGEVIN D YNAMICS
In the generative modeling context, it has been proposed to extend the original state x ∈ R d by a momentum variable m ∈ R d .
One then diffuses in this extended space using an underdamped Langevin dynamics (Dockhorn et al., 2022) targeting N ( x ; 0 , σ 2 I ) N ( m ; 0 , M ) .
It was demonstrated empirically that the resulting scores one needs to estimate are smoother and this led to improved performance.
We adapt this approach to Monte Carlo sampling.
This adaptation is non trivial and in particular requires to design carefully numerical integrators.
B. C ONTINUOUS TIME
We now consider an augmented target distribution π ( x ) N ( m ; 0 , M ) where M is a positive definite mass matrix.
We then diffuse this extended target using the following underdamped Langevin dynamics, i.e.
[Equation] [Equation]
where x 0 ∼ π, m 0 ∼ N (0 , M ) .
The resulting path measure on [0 , T ] is denoted again P .
From (Haussmann & Pardoux, 1986), the time-reversal process is also a diffusion satisfying
[Equation] [Equation]
for ( y 0 , n 0 ) ∼ η T where η t denotes the density of ( x t , m t ) under (38).
Now consider a reference process P ref on [0 , T ] defined by the forward process (38) initialized using x 0 ∼ N (0 , σ 2 ) , m 0 ∼ N (0 , M ) .
In this case one can check that η ref t ( x t , m t ) =
N ( x t ; 0 , σ 2 I ) N ( m t ; 0 , M ) and the time-reversal process of P ref satisfies
[Equation]
as ∇ n ln η ref
[Equation]
y,
[Equation]
Hence it follows that the time-reversal (39) of P can be also be written as
[Equation] [Equation]
where ϕ t ( x, m ) :
=
η t ( x, m ) /η ref t ( x, m ) .
To approximate P , we consider a parameterized path measure Q θ whose time reversal is defined for ( y 0 , n 0 ) ∼ N ( y 0 ; 0 , I ) N ( n 0 ; 0 , M ) by d y t = − M − 1 n t d t and d n t =
y
[Equation]
B. L EARNING THE TIME - REVERSAL THROUGH KL MINIMIZATION
To approximate P
[Equation]
we will consider a parameterized diffusion whose time reversal is defined by
[Equation]
for ( y 0 , n 0 ) ∼ N ( y 0 ; 0 , I ) N ( n 0 ; 0 , M ) inducing a path measure Q θ on the time interval [0 , T ] .
We can now compute the Radon-Nikodym derivative between Q θ and P ref using an extension of
Girsanov theorem (Theorem A.3.
in Sottinen & S¨arkk¨a (2008))
[Equation]
To summarize
[Equation]
we have the following proposition
[Equation]
Proposition 5.
The Radon-Nikodym derivative d Q θ d P ref ( y [0 ,T ] , n [0 ,T ] ) satisfies under Q θ
[Equation]
From KL( Q θ ||P )
=
KL( Q θ ||P ref )
+ E Q θ (cid:104) ln (cid:16) p ref 0 ( y T ,n T ) η 0 ( y T ,n T ) (cid:17)(cid:105) , it follows that
[Equation]
The second term on the r.h.s.
of (46) follows from the fact that ln( p ref 0 ( y T , n T ) /p
( y T , n T ))
=
ln( N ( y T ; 0 , σ 2 I ) /π
( y T )) .
B. N ORMALIZING FLOW THROUGH ORDINARY DIFFERENTIAL EQUATION
The following ODE gives exactly the same marginals η T as the SDE (38) defining P
[Equation]
Thus if we integrate its time reversal from 0 to T starting from ( y 0 , q 0 ) ∼ η T
[Equation]
[Equation]
then we would obtain at time T a sample ( y T , n T ) ∼ π ( y T ) N ( n T ; 0 , M ) .
In practice, this suggests that once we have learned an approximation f θ ∗ ( t, x, m ) ≈ ∇ m ln ϕ t ( x, m ) by minimization of the ELBO then we can construct a proposal using
[Equation]
[Equation]
for ( y 0 , n 0 ) ∼ N ( y 0 ; 0 , σ 2 I ) N ( n 0 ; 0 , M ) .
The resulting sample ( y T , n T ) ∼ ¯ η 0 will have distribution close to π × N (0 , M ) .
Again it is possible to compute pointwise the distribution of this sample to perform an importance sampling correction.
B. F ROM CONTINUOUS TIME TO DISCRETE TIME
We now need to come up with discrete-time integrator for the time-reversal of P ref given by (40) and the time-reversal of Q θ given by (43).
Let us start with (40).
We split it into the two components
[Equation]
and
[Equation] [Equation]
We will compose these transitions.
To obtain ( y k +1 , n k +1 ) from ( y k , n k ) , we first integrate (50).
To do this, consider the Hamiltonian equation d y t =
M − 1 n t d t, d n t = − y t σ 2 d t which preserves N ( y ; 0 , σ 2 I ) N ( n ; 0 , M ) as invariant distribution.
We can integrate this ODE exactly over an interval of length δ and denote its solution Φ( y, n ) ; see Section B.5 for details.
We use its inverse Φ − 1 ( y, n ) =
Φ flip ( y, n ) ◦ Φ( y, n ) ◦ Φ flip ( y, n ) where Φ flip ( y, n ) =
( y, − n ) so that ( y k +1 , n ′ k ) =
Φ − 1 ( y k , n k ) .
Then we integrate exactly (51) using
[Equation]
We have thus design a transition of the form
[Equation] [Equation]
Now to integrate (43), we split it in three parts.
We first integrate (50) using exactly the same integrator ( y k +1 , n ′ k ) =
Φ − 1 ( y k , n k ) .
Then we integrate
[Equation]
using
[Equation]
where we abuse notation and write f θ ( K − k, y k +1 , n ′ k ) instead of f θ (( K − k ) δ, y k +1 , n ′ k ) .
Finally, we integrate the OU part using (52) but replacing n ′ k by with n ′′ k .
So the final transition we get is
[Equation]
However
[Equation]
we can integrate n ′′ k analytically to obtain
[Equation]
B. E XACT SOLUTION OF H AMILTONIAN EQUATION
Consider the Hamiltonian equation and M
=
I , the solution of
[Equation] [Equation]
can be exactly written as ( x t , m t ) =
Φ t ( x 0 , m 0 ) defined through
[Equation]
This is the so-called harmonic oscillator.
Now the inverse of the Hamiltonian flow satisfies Φ
[Equation] [Equation]
) (see e.g. Leimkuhler & Matthews (2016))
so we can simply integrate
[Equation]
using ( y k +1 , n ′ k )
=
Φ − 1 τ ( y k , n k ) .
This gives
[Equation]
Obviously, if we have M = αI then
[Equation]
then we use a reparameterization ˜ m t
= m t /α and
[Equation]
and the solution is as above with σ being replaced by σ √ α .
So for ˜ n k = n k /α , we have
[Equation]
so as n k = α ˜ n k and similarly n ′ k
= α ˜ n ′ k , we finally obtain
[Equation]
B. D ERIVATION OF U NDERDAMPENED KL IN D ISCRETE T IME
In this section we derive the discrete-time KL for the underdamped noising dynamics.
Proposition 6.
The log density ratio lr = ln
( q θ ( y 0: K , n 0: K , n ′
: K ) /p ref ( y 0: K , n 0: K , n ′ 0: K )) for y
: K , n 0: K , n ′ 0:
K ∼ q θ ( · ) equals
[Equation] [Equation]
where κ k := √ 1 − α k
(1 − √ 1 − α k ) and ε k is obtained as a function of n k +1 and n ′′ k when describing the integrator for Q θ .
In particular, we have ε k i.i.d.
∼ N (0 , I ) and one obtains
[Equation]
The integrator has been designed so that the ratio between the transitions of the proposal (56) and the reference process (53) is well-defined as the deterministic parts are identical in the two transitions so cancel and
[Equation]
The calculations to compute the Radon–Nikodym derivative are now very similar to what we did in the proof of Proposition
[Equation]
We have
[Equation]
where we have exploited the fact that q K ( y 0 , n 0 )
=
p ref K ( y 0 , n 0 ) =
N ( y 0 ; 0 , σ 2 I ) N ( n 0 ; 0 , M ) .
Now let us introduce
[Equation]
hence
[Equation]
We can rewrite (69) as
[Equation]
where we note that √ 1 − α k (1 −√ 1 − α k ) √ α k
[Equation]
k /
[Equation]
δ .
C E XPERIMENTS
In this section we incorporate additional experiments and ablations as well as further experimental detail.
For both PIS and DDS we created a grid with δ = KT = 0 . 05 and values of
T ∈ { 3 . 4 , 6 . 4 , 12 . 8 , 25 . 6 } and corresponding number of steps K ∈ { 64 , 128 , 128 , 256 } .
When tuning PIS we explore different values of σ which in practice, when discretised, amounts to the same effect as changing δ .
For DDS we apply the cosine schedule directly on the number of steps however we ensure that (cid:80)
k
α k
=
α max T such that we have a similar scaling as in PIS.
Finally both
PIS and DDS where trained with Adam (Kingma & Ba, 2015) to at most 11000 iterations, although in most cases most both converged in less than 6000 iterations.
Across all experiments used a sample size of 300 to estimate the ELBO at train time and of 2000 to for the reported IS estimator of Z for all methods.
Finally for the ground truth where available we use the true normalizing constant as is the case with the Funnel distribution and the Normalising Flows used in the NICE Dinh et al.
(2014) task.
For the rest of the tasks we follow prior work (Zhang & Chen, 2022; Arbel et al., 2021) use a long run SMC chain (1000 temperatures, 30 seeds with 2000 samples each).
C. N ETWORK P ARAMETRISATION
In order to improve numerical stability we re-scale the neural network drift parametrisation as follows:
[Equation]
This was done such that the reciprocal term α − 1 K − k in the DDS objective is cancelled as the term reaches very small values in the boundaries causing the overall objective to be large.
Finally as
[Equation]
k 2 it follows that our proposed re-parametrisation converges to the same
SDE as before, but now with stabler and simpler updates:
[Equation]
C. T UNING H YPER - PARAMETERS
For both DDS and PIS we explore a grid of 25 hyper-parameter values.
For PIS we search for the best performing value of the volatility coefficient over 25 values, depending on the task we vary the end points of the drift however we noticed that PIS leads to numerical instabilities for γ > 4 thus we never search for values larger than this.
For DDS we searched across 5 values for σ and 5 values for α max , this led to a total of 25 combinations we explored for each experiment.
Finally for SMC we searched over 3 of its different step sizes leading to a total of 5 3 = 125 .
C. B ROWNIAN M OTION M ODEL W ITH G AUSSIAN O BSERVATION N OISE
The model is given by
[Equation]
y
[Equation]
The goal is to perform inference over the variables α inn , α obs and { x i } 30 i =1 given the observations { y i } 10 i =1 ∪ { y i } 30 i =20 .
[Figure]
[Table]
C. D RIFT M AGNITUDE
In Figure 4 we compare the magnitudes of the NN drifts learned by PIS and DDS on a scalar target N (6 , σ 2 ) target.
The drifts where sampled on randomly evaluated trajectories from each of the learned samplers.
We observe the PIS drift to have large magnitudes with higher variance more prone to numerical instabilities, whilst the OU drift has a notably lower magnitude with significantly less variance.
We conjecture this is the reason that PIS becomes unstable at training time for a large number of steps across several of our PIS experiments, many of which did not converge due numerical instabilities.
For simplicity in this experiment we set α k to be uniform.
C. T RAINING T IME
We evaluate the training times of DDS and PIS on an 8-chip TPU circuit and average over 99 runs.
Table 2 shows running times per training iteration.
In total we trained for 11000 iterations thus the total training times are 19 minutes; 37 minutes; 1 hour and 14 minutes; 2 hours and 26 minutes for k = 64 , 128 , 256 , 512 respectively.
C.
V ARIATIONAL E NCODER IN L ATENT S PACE
Following Arbel et al.
(2021) we explore estimating the posterior of a pre-trained Variational Auto
Encoder (VAE) (Kingma & Welling, 2013).
We found all approaches to reach a very small error for ln Z very quickly however as seen with PIS across experiments results became unstable for large T and thus many hyper-parameters when tuning the volatility for PIS led to the loss exploding.
C. G ENERATED I MAGES
In this section we provide some of the generated samples for the normalising flow evaluation task.
In Figure 6 we can observe how both PIS and DDS are able to generate more image classes than
SMC due to mode collapse.
Out of the 3 approaches we can see that DDS mode collapses the least whilst SMC generates the higher quality images.
Using a neural network with inductive biases such
[Figure]
as a convolutional neural network (LeCun et al., 1998) should improve the image quality over the simple feed-forward networks we have used, we believe this is the reason behind the low quality in the images as both minimising the path KL and sampling high quality images is a difficult task for a small feed-forward network.
C. C OMPARISON TO AIS AND MCD
We compare our results to the AIS and MCD baselines presented in Geffner & Domke (2022).
Results can be seen in Tables 4 and 3, we can see that both DDS and PIS outperform LDVI across all number of steps and tasks with DDS outperforming PIS on higher values of k due to inconssitent and unstable behaviour presented by PIS.
C. E ULER –M AYURAMA VS E XPONENTIAL I NTEGRATOR
In this section we demonstrate empirically how using the Euler Mayurama integrator directly on the
DDS objective leads to overestimating ln Z .
We compare on the Funnel dataset for which we know ln Z to be 0 and the Ionosphere dataset which is small and for which we have a reliable estimate of ln Z via a long run SMC.
Results in Table 5 show case how for lower values of k = 64 ,
[Figure]
DDS
PIS
ULA
MCD
[Figure]
LDVI
DDS
PIS k
=-113.
[Table]
Brownian Motion
ULA
MCD
UHA
LDVI
DDS
PIS k
=-0.
k
=-0.
k
=-0.
[Table]
the Euler–Maruyama (EM) based approach significantly overestimates ln Z and overall does not perform well.
C.9. D ETACHED G RADIENT A BLATIONS
In this section we perform an ablation over our proposed modification of the PIS-Grad network architecture.
We compare the same architecture with and without the gradient attached.
We found that detaching gradient led to a favourable performance as well as more stable training.
Results are presented in table 7 and were carried out for K = 128 and using the best found diffusion coefficient from each task.
We chose to explore this feature as it is known that optimization through unrolled computation graphs can be chaotic and introduce exploding/vanishing gradients (Parmas et al., 2018;
Metz et al., 2020) which can numerically introduce bias.
This has been successfully applied in related prior work (Greff et al., 2019) where detached scores are provided as features to the recognition networks.
Alternative approaches exist based on smoothing the chaotic loss-landscape (Vicol et al., 2021) but we leave this for future work.
C. C OSINE BASED DECAY FOR PIS
Unlike DDS, with PIS it is not immediately clear how to schedule the step-sizes in a similar manner to Nichol & Dhariwal (2021).
The simplest approach would be to schedule δ (the discretisation step).
We perform an ablation with the cosine schedule on the discretisation step with PIS.
Results are displayed in Table C.10, we can see that overall the cosine discretisation decreases or preserves the performance of PIS with the exception of the Sonar data-set in which it increases.
Whilst there may be a way to improve PIS with a bespoke discretisation we were unable to find such in this work.
Additionally we can see from Table C.10 how DDS improves significantly across every task when using the cosine squared schedule compared to uniform.
C. P ROBABILITY F LOW ODE
Figures 7 and 8 show samples obtained from the probability flow ODE of a trained DDS model.
We can see that the probability flow ODE is able to perfectly sample from a uni-modal Gaussian, matters became more challenging with multi-modal distributions.
Initially we discretised the ODE (12) using the same type of integrators as the SDEs to obtain y k +
=
y k
+ δσ 2 (1 −√ 1 − α K − k )
f θ ( K − k, y k ) for y 0 ∼ N (0; σ 2 I ) .
Unfortunately under this discretisation we found the probability flow
ODE to become stiff with more complex distributions such as the mixture of Gaussians, resulting in strange effects in the samples (matching the modes but completely wrong shapes).
By using a Heun
[Table] [Table]
integrator as proposed in Karras et al.
we were able to obtain better results, however we also found we needed to increase the network size to improve the results as well as train with a learning rate decay of 0 . 99 .
With these extra features we were able to simulate a probability flow ODE that roughly matched the marginal densities of the SDE.
However even with these fixes we still found the probability flow based estimator of ln Z to drastically overestimate, this is not entirely surprising as in discrete time the expectation of this estimator is not guaranteed to be an ELBO.
C. U NDERDAMPED OU R ESULTS
We perform some additional experiments with the underdamped OU reference process.
Similar to the damped setting we parametrise the network as:
[Equation]
Unfortunately, unlike in DDS and PIS we cannot directly aid the update/proposal for x t with ∇ ln π ( x ) , thus this naive parametrisation is not the ideal inductive bias.
Results in Figures 9 and 10 show some experiments for this approach.
The results perform worse than DDS for a standard OU reference process and PIS.
However they are still better than VI.
We believe that future work exploring better inductive biases for the network f θ ( k, x, p ) could narrow the performance gap in this approach.
Additionally we highlight that in Figure 9 (Funnel Distribution) we can see the underdamped approach overestimates ln Z for k = 64 .
We believe this may be due to numerical error when generating a sample.
C. B OX P LOTS FOR M AIN R ESULTS
For completeness in this section we also present our results via box plots as can be seen in Figures 11 and 12.
D D ENOISING D IFFUSION S AMPLERS IN D ISCRETE - TIME : F URTHER DETAILS D. F ORWARD P ROCESS AND I TS T IME R EVERSAL
In discrete-time, we consider the following discrete-time “forward” Markov process
[Equation]
Following (17), we use p k | k − 1 ( x k | x k − 1 ) =
N ( x k ; √ 1 − α k x k − 1 , σ 2 α k I ) .
The parameters ( α k ) Kk =1 are such
that p K ( x K ) ≈ N ( x K ; 0 , σ 2 I ) .
We propose here to sample approximately from π by approximating the ancestral sampling scheme for (75) corresponding to the backward decomposition p ( x 0: K ) =
p K ( x K )
K (cid:89) k
= p k − 1 | k ( x k − 1 | x k ) , with p k − 1 | k ( x k − 1 | x k ) =
p k − 1 ( x k − 1 ) p k | k − 1 ( x k | x k − 1 ) p k
( x k ) , (76) where p 0 ( x 0 )
=
π ( x 0 ) .
If we could thus sample x K ∼ p K ( · ) then x k − 1 ∼ p k − 1 | k ( ·| x k ) for k = K, ..., 1 then x 0 would indeed be a sample from π .
However as the marginal densities ( p k )
Kk = are not available in closed-form, we cannot implement exactly this ancestral sampling procedure.
As we have by design p K ( x K ) ≈ N ( x K ; 0 , σ 2 I ) , we can simply initialize the ancestral sampling procedure by a Gaussian sample.
The approximation of the backward Markov kernels p k − 1 | k
( x k − 1 | x k ) is however more involved.
D. R EFERENCE P ROCESS AND B ELLMAN R ECURSION
We introduce the “simple” reference process p ref ( x 0: K ) defined by
[Equation]
which is designed by construction to admits the marginal distributions p ref k ( x k )
=
N ( x k
[Equation]
0 , σ 2 I ) for all k .
It can be easily verified using the chain rule for KL that the extended target process p is the distribution minimizing the reverse (or forward) KL discrepancy w.r.t.
p ref over the set of path measures q ( x 0: K ) with marginal q 0 ( x 0 ) =
π ( x 0 ) at the initial time, i.e.
[Equation]
The intractable backward Markov densities p k − 1 | k one wants to approximate can also be rewritten as twisted versions of the tractable backward Markov densities p ref k − 1 | k by some value functions ( ϕ k ) K k =0 .
[Figure] [Figure]
Proposition 7.
We have for k = 1 , ..., K
[Equation]
The value functions ( ϕ k )
Kk =1 satisfy a forward Bellman type equation
[Equation]
It follows that
[Equation]
Proof.
To establish (78), we use Bayes’ rule
[Equation]
where we have used the fact that p k ( x k ) =
ϕ k ( x k ) p ref k ( x k ) , p ref k | k −
[Equation]
x
[Equation]
x k −
[Equation]
p k | k
[Equation]
x k | x k
[Equation]
1 ) and Bayes’ rule again.
Now we have (cid:82) p k − 1 | k
( x k − 1 | x k )d x k − 1 = 1 for any x k , so it follows directly from the expression of this transition kernel that the value function satisfies
[Equation]
By iterating this recursion, it follows that
[Equation]
[Figure] [Figure]
D. A PPROXIMATING THE B ACKWARD KERNELS
We want to approximate the backward Markov transitions p k − 1 | k .
From (78), this would require not only to approximate the value functions - Monte Carlo estimates based on (80) are high variancebut also to sample from the resulting twisted kernel which is difficult.
However, if we select β k ≈ 0 , then we have ϕ k − 1 ( x ) ≈
ϕ k ( x ) and by a Taylor expansion of ∇ ln ϕ k around x
k we obtain p
[Equation]
By differentiating the identity (80) w.r.t.
x k , we obtain the following expression for ϕ k
[Equation]
Alternatively, we can also use ∇ ln ϕ k
( x k )
= ∇ ln p k
[Equation]
x k )
+ xσ 2 where ∇
ln p k
( x k ) can be easily shown to satisfy
[Equation]
For DDPM, the conditional expectation in (85) is reformulated as the solution to a regression problem, leveraging the fact that we can easily obtain samples from π ( x 0 ) p k
| 0 ( x k | x 0 ) as π is the data distribution in this context.
In the Monte Carlo sampling context considered here, we could sample approximately from p 0 | k
( x 0 | x k ) ∝ π ( x 0 ) p k | 0 ( x k | x 0 ) using MCMC so as to approximate the conditional expectations in (84) and (85) but this would defeat the purpose of the proposed methodology.
We propose instead to approximate ∇ ln ϕ k by minimizing a suitable criterion.
In practice, we will consider a distribution q θ ( x 0: K ) approximating p ( x 0: K ) of the form
[Equation]
[Equation]
i.e. we select q θ K ( x K ) =
q K ( x K ) =
N ( x K ; 0 , σ 2 I ) as p K ( x K )
≈ N ( x K ; 0 , σ 2 I ) .
We want q θk − 1 | k
( x k − 1 | x k ) to approximate p k − 1 | k ( x k − 1 | x k ) , i.e. f θ ( k, x k )
≈ ∇ ln ϕ k ( x k ) .
Inspired by (83), we will consider an approximation of the form
[Equation]
[Equation]
where f θ ( k, x k ) is differentiable w.r.t.
θ .
D. H YPERPARAMETERS D.4. F ITTED H YPERPARAMETERS
To aid reproducibility we report all fitted hyperparameters for each of our methods and PIS across all experiments in Tables 8 and 9.
D.4. O PTIMISATION H YPERPARAMETERS
As mentioned in the experimental section across all experiments modulo the Funnel we use the
Adam optimiser with a learning rate of 0 . 0001 with no learning decay and 11000 training iterations, for the rest of the optimisation parameters use the default settings as provided by the Optax library (Hessel et al., 2020) which are b 1 = 0 .
9 , b 2 = 0 .
999 , ϵ = 10 − 8 naming as per Kingma & Ba (2015).
From the github repository of Zhang & Chen (2022) we were only able to find hyperparameters reported for the Funnel distribution.
In order to first reproduce their results we used the a learning rate of 0 .
005 and a learning rate decay of 0 .
95 as per their implementation, their results were initially not reproducible due to a bug in setting σ f = 1 despite comparing to methods at the less favourable values of σ f = 3 .
For σ f = 1 we were able to reproduce their results.
However we report results at σ f = 3 as this is the traditional value used for this loss.
As no other optimisation configuration files were reported we used the more conservative learning rate of 0 . 0001 since PIS was very unstable for 0 .
005 with decay 0 .
95 across many of our tasks.
Finally we would like to clarify that the exact same optimiser settings where used for both PIS and DDS in order to ensure a fair comparison.
D.4. D RIFT AND G RADIENT C LIPPING
We follow the same gradient clipping as in Zhang & Chen (2022) that is :
[Equation]
[Table] [Table]
0 .σ
= 0 .
, α = 3 . 7 , m =
0 .σ
= 2 .
1 , α = 1 .
σ = 0 .σ
= 1 . 7 , α = 4 . 5 , m = 0 .σ
= 0 .
, α = 0 .σ
= 0 .σ
= 0 .
6 , α = 3 ., m = 1 .
T a b l e 8 : F itt e d hyp e r p a r a m e t e r s f o r f unn e l , l g c p a nd i on e xp e r i m e n t s .
l r s on a r v ae b r o w n i a n n i ce
DD S
P I S
UD M P
DD S
P I S
UD M P
DD S
P I S
UD M P
DD S
P I S
UD M P
K = σ = 0 .
3 , α = 1 .
σ = 0 .
σ = 1 .
15 , α = 1 . 7 , m = 2 .σ = 0 .
61 , α = 2 .σ
= 0 .
NA σ = 0 .
1 , α = 2 .
σ = 0 . σ = 0 .
115 , α = 4 . 8 , m = 2 .σ
= 1 . 5 , α = 2 .
σ = 0 .
σ = 1 . 2 , α = 1 . 0 , m = 0 .
K = σ = 0 .
3 , α = 1 .σ
= 0 .
σ = 0 .
55 , α = 1 . 7 , m = 3 .σ 0 .
61 , α = 1 .
σ = 0 .
NA σ = 0 .
1 , α = 1 .σ
= 0 .
σ = 0 .
115 , α
[Table] [Table]
[Table]
K = σ = 0 .
3 , α = 0 .
σ = 0 . σ = 0 .
55 , α = 2 .
9 , m = 2 .σ
= 0 .
61 , α = 1 .
[Table] [Table]
[Table] [Table]
[Table]
K = σ = 0 .
3 , α = 0 .
σ = 0 . σ = 0 .
55 , α = 2 .
9 , m = 3 .σ
= 0 .
61 , α = 1 .σ = 0 .
NA σ = 0 .
1 , α =
[Table] [Table]
[Table] [Table]
[Table]
T a b l e 9 : F itt e d hyp e r p a r a m e t e r s f o r s on a r , b r o w n i a n , v ae a nd n i ce
e xp e r
i m e n t s .