Score-Based Diffusion meets Annealed Importance Sampling
Anonymous Author(s)
AffiliationAddress email
Abstract
More than twenty years after its introduction, Annealed Importance Sampling(AIS) remains one of the most effective methods for marginal likelihood estimation.
It relies on a sequence of distributions interpolating between a tractableinitial distribution and the target distribution of interest which we simulate fromapproximately using a non-homogeneous Markov chain.
To obtain an importancesampling estimate of the marginal likelihood, AIS introduces an extended targetdistribution to reweight the Markov chain proposal.
While much effort has beendevoted to improving the proposal distribution used by AIS, an underappreciatedissue is that AIS uses a convenient but suboptimal extended target distribution.
This can hinder its performance.
We here leverage recent progress in score-basedgenerative modeling (SGM) to approximate the optimal extended target distributionfor AIS proposals corresponding to the discretization of Langevin and Hamiltoniandynamics.
We demonstrate these novel, differentiable, AIS procedures on a numberof synthetic benchmark distributions and variational auto-encoders.
Introduction
Evaluating the marginal likelihood, also known as evidence, is of key interest in Bayesian statisticsas it allows not only model comparison but is also often used to select hyperparameters.
A largevariety of Monte Carlo methods have been proposed to address this problem, including path sampling[Gelman and Meng, 1998], AIS [Neal, 2001] and related Sequential Monte Carlo methods [Del Moralet al., 2006].
An appealing feature of AIS is that it provides an unbiased estimate of the marginallikelihood and can thus be used to define an evidence lower bound (ELBO) or mutual informationbounds; see e.g. [Wu et al., 2020; Thin et al., 2021;
Brekelmans et al., 2022].
AIS builds a proposal distribution using a Markov chain ( x k )
Kk =0 initialized at an easy-to-sampledistribution followed by a sequence of Markov chain Monte Carlo (MCMC) transitions targetingtypically annealed versions of the posterior.
By proceeding this way, we obtain a proposal x K whosedistribution is expected to be a reasonable approximation to the target posterior.
However, thisdistribution is intractable as it requires integrating the joint proposal distribution over previous states( x k ) K − 1 k =0 .
AIS bypasses this issue by instead using Importance Sampling (IS) on the whole path( x k )
Kk =0 through the introduction of an artificial extended target distribution whose marginal at time
K coincides with the posterior.
There has been much work devoted to improving AIS in machine learning and statistics but alsoin physics where it was introduced independently in [Jarzynski, 1997; Crooks, 1998].
A standardapproach to improve AIS is to modify the intermediate distributions [Schmiedl and Seifert, 2007;
Grosse et al., 2013; Masrani et al., 2021] and corresponding transition kernels of the proposal [Daiet al., 2020; Wu et al., 2020; Geffner and Domke, 2021; Thin et al., 2021; Zhang et al., 2021].
We here address a distinct issue: it was shown in [Del Moral et al., 2006] that the extended target distribution used by AIS is suboptimal and that the optimal extended target minimizing the varianceof the evidence estimate is defined through the time-reversal of the proposal.
However, this result isdifficult to exploit algorithmically as the time-reversal is intractable for useful proposals.
In this paper, we show how one can combine this result with recent advances in SGM to obtainimproved, lower variance, AIS estimates.
We concentrate on scenarios where we use unadjusted overdamped Langevin
[Heng et al., 2020; Wu et al., 2020;
Thin et al., 2021] and unadjusted Hamiltonianproposals with partial momentum refreshment (i.e. underdamped Langevin)
[Dai et al., 2020; Wuet al., 2020; Geffner and Domke, 2021; Zhang et al., 2021] which correspond to time-discretizeddiffusion processes.
The first benefit of using such proposals is that, by omitting Metropolis–Hastingssteps, one obtains differentiable versions of the ELBO amenable to the reparameterization trick.
Thesecond benefit of these proposals is that their time-reversal can be approximated by adapting techniques developed for SGM [Ho et al., 2020; Song et al., 2021b; Dockhorn et al., 2022] to our setup.
We derive a principled parameterization for an approximation of their time-reversal which we learnby maximizing the ELBO.
As for SGM, this ELBO coincides with a denoising score matching loss[Hyvärinen, 2005; Vincent, 2011; Ho et al., 2020; Song et al., 2021b].
This provides novel, optimizedand differentiable, AIS estimators which we refer to as Monte Carlo Diffusion (MCD).
We demonstrate the benefits of this approach on synthetic benchmark distributions and variationalauto-encoders (VAEs)
[Kingma and Welling, 2013].
All proofs can be found in the Appendix.
Annealed Importance Sampling
2.1 Setup and algorithm.
Consider a probability density π on R d of the form
[Equation]
where γ ( x ) can be evaluated pointwise
and we want to approximate the intractable normalizingconstant Z .
In a Bayesian framework, γ ( x ) = p ( x ) p ( D| x ) is the joint density of parameter x anddata D , π ( x ) =
p ( x |D ) the corresponding posterior and Z = p ( D ) the evidence.
To estimate Z , AIS introduces the intermediate distributions ( π k )
Kk =1 bridging smoothly from atractable distribution π 0 to the target distribution π K
= π of interest.
One typically uses π k ( x ) ∝γ k ( x ) with γ k ( x ) =
π 0 ( x ) 1 − β k γ K ( x ) β k for 0 =
β 0 < β 1 < · · · < β K = 1 but other choices arepossible [Grosse et al., 2013].
The IS proposal used by AIS is then obtained by running a Markov chain ( x k ) Kk =0 such that x 0 ∼ π 0 ( · ) , and then x k ∼ F k ( ·| x k − 1 ) for k ≥ 1 where F k is a MCMCkernel invariant w.r.t.
π k .
The proposal is thus given by
[Equation]
Denote by q k the intractable marginal distribution of x k under Q satisfying q k ( x k ) =
67 (cid:82) q k − 1 ( x k − 1 ) F k
( x k | x k − 1 )d x k − 1 for k ≥ 1 and q 0 = π 0 .
As q K cannot be evaluated in complex scenarios, the marginal IS estimate w mar ( x K ) =
γ K ( x K )
/q
K ( x K ) of Z is intractable.
One can bypass this issue by introducing an extended target
[Equation]
where ( B k ) K − 1 k =0 are backward Markov transition kernels.
For any selection of backward kernels suchthat the ratio Γ /Q is well-defined, we have
[Equation]
for
[Equation] [Equation]
i.e. w ( x 0: K ) is an unbiased estimate of Z for x 0: K ∼ Q .
The AIS estimate of the evidence is a specific instance of the estimator (4) relying on the backward kernels B ais k ( x ′ | x ) =
π k +1 ( x ′ )
F k +1 ( x | x ′ )
/π
k +1 ( x ) .
This yields the following expression forlog w ( x 0: K ) :
[Equation]
2.2 Limitations of AIS.
While designing P in (3) by using the backward Markov kernels ( B ais k ) K − 1 k =0 is convenient, it is also suboptimal.
For example, consider the ideal scenario where F k ( x k | x k − 1 ) =
π k ( x k ) .
This scenario hasbeen used many times in the literature to provide some guidelines on AIS, see e.g. Neal [2001]; Grosseet al.
In this case, var Q [log w ais ( x 0: K )]
= (cid:80)
[Equation]
var π
[Equation] [Equation]
[Equation] [Equation]
))] > 0while var q K [ w mar ( x K )]
= var π K [ w mar ( x K )]
= 0 .
Both estimates use the same forward kernels but reweight samples in a different way using distinct backward kernels.
Initial distribution π 0 is 20-dimensional
N (0 , I ) and progressively shifts to the target π = π K = N (10 , I ) .
The MCD estimate is much closer to the ground truth ( log Z = 0 ) than AIS.
Another illustration of the suboptimality of AIS is to consider a scenario where the proposal is a homogeneous
MCMC chain, i.e. x 0 ∼ π 0 and x k ∼ F ( ·| x k − 1 ) for
F a π -invariant MCMC kernel; i.e. use F k = F andπ k = π for k = 1 , ..., K .
If F is reasonably wellmixing, then q K ≈ π for K large enough and the evidence estimate w mar ( x K ) =
γ K ( x K )
/q
K ( x K ) shouldhave small variance.
However, it is easy to check thatwe have w ais ( x 0: K ) =
γ ( x 0 ) /π
( x 0 ) for the exact sameproposal; i.e. the AIS estimate does not depend on the
MCMC samples x 1: K and boils down to the IS estimateof Z using the proposal π 0 .
These two examples illustrate that it would be preferableto use w mar ( x K ) rather than w ais ( x 0: K ) .
In Appendix A,we provide a detailed comparison of both estimates in ascenario where their variance can be computed analytically.
We propose in the next section an unbiased estimate ofthe evidence (MCD) approximating w mar ( x K )
based ona different choice of backward kernels.
As illustrated in
Figure 2, significant gains can be achieved.
Optimized Annealed Importance Sampling
We show here that the optimal extended target distribution P is defined through the time-reversal ofthe proposal Q .
By exploiting a connection to SGM, we can approximate this reversal using scorematching when the proposal is obtained through an unadjusted overdamped/underdamped Langevinalgorithm.
3.1 Optimal Extended Target Distribution via Time Reversal
We summarize here Proposition 1 of Del Moral et al.
[2006]; see also [Sohl-Dickstein et al., 2015;
Bernton et al., 2019].
Proposition 1.
For a proposal Q of the form (2), the extended target P of the form (3) minimizing boththe Kullback–Leibler divergence D KL ( Q || P ) and the variance of the evidence estimate w ( x 0: K ) =
Γ( x 0: K ) /Q
( x 0: K ) for x 0: K ∼ Q is given by P opt ( x 0: K ) =
Γ opt ( x 0: K ) /Z
where
[Equation]
In particular, one has
[Equation]
and D KL ( Q || P opt ) =
D KL
[Equation]
q K || π K ) .
This result follows simply from the chain rule and the law of total variance which yield
[Equation] [Equation]
Both quantities are clearly minimized by selecting P ( x 0: K − 1 | x K ) =
Q ( x
[Equation]
: K − 1 | x K ) .
We emphasize that Proposition 1 applies to any forward kernels ( F k )
Kk =1 including MCMC kernels,unadjusted Langevin kernels or even deterministic maps 1 .
It shows that P opt is the distribution of abackward process initialized at π K which then follows the time-reversed dynamics of the forwardprocess Q .
If we had π K = π , then we would have P opt = Q as then P opt corresponds to thebackward decomposition of Q .
3.2 Time reversal, Score matching and ELBO for unadjusted overdamped Langevin
We concentrate here on the case where ( F k ) Kk =1 correspond to a time-inhomogeneous unadjusted(overdamped)
Langevin algorithm (ULA) as used in [Heng et al., 2020; Wu et al., 2020; Thin et al.
,2021]; that is we consider F k ( x k | x k − 1 ) =
N ( x k ; x k − 1 + δ ∇ log π k ( x k − 1 ) , 2 δI ) where δ > 0 isa stepsize.
Let δ :=
T/K then, as K → ∞ , the proposal Q converges to the path measure Q ofthe following inhomogeneous Langevin diffusion ( x t ) t ∈ [0 ,T ] defined by the stochastic differential 127 equation (SDE)
[Equation]
where ( B t ) t ∈ [0 ,T ] is standard multivariate Brownian motion and we are slightly abusing notationfrom now on as π t for t
= t k = kδ corresponds to π k in discrete-time.
Many quantitative resultsmeasuring the discrepancy between the law of x T and π T =
π for such annealed diffusions have beenobtained; see e.g. [Fournier and Tardif, 2021; Tang and Zhou, 2021].
It can be shown [Haussmannand Pardoux, 1986] that the time-reversed process (¯ x t ) =
( x T − t ) t ∈ [0 ,T ] is also a diffusion given by
[Equation]
where ( ¯ B t ) t ∈ [0 ,T ] is another multivariate Brownian motion.
The continuous-time version of P opt isthe path measure P opt defined by the diffusion (11) but initialized at ¯ x 0 ∼ π T rather than q T as notedin Bernton et al.
[2019]; see Figure 1 for an illustration.
This shows that approximating ( B opt k ) K − 1 k =0requires approximating the so-called scores ( ∇ log q t ( x )) t ∈ [0 ,T ] .
This can be derived heuristicallythrough the fact that a Taylor expansion yields the following approximation of the optimal backwardkernels, B opt k ( x k | x k +1 )
≈ N ( x k ; x k +1 − δ ∇ log π k +1 ( x k +1 )
+ 2 δ ∇ log q k +
x k +1 ) , 2 δI ) , whichindeed corresponds to a Euler discretization of (11); see e.g. [De Bortoli et al., 2021b, Section 2.1].
In SGM [Song et al., 2021b], one gradually adds noise to data using an Ornstein–Ulhenbeck diffusionto transform the complex data distribution into a Gaussian distribution and the generative model isobtained by approximating the time-reversal of this diffusion initialized by Gaussian noise.
Practically, the time-reversal approximation is obtained by estimating the scores of the noising diffusion usingdenoising score matching [Vincent, 2011].
While in our setup, the diffusion (10) instead goes froma simple distribution to a complex one (see Appendix D for a discussion), we can still adapt scorematching ideas to our setup.
We define a path measure P θ approximating P opt using a neural networks θ ( T − t, ¯ x t ) in place of ∇ log q T − t (¯ x t ) in (11), i.e. we consider
[Equation]
We would like to learn θ by minimizing D KL ( Q||P θ ) over θ , i.e. equivalently we maximize acontinuous-time ELBO.
Note that it is neither feasible to minimize D KL ( P opt ||P θ ) as one cannotsample from π T nor it is desirable as the evidence estimate is computed using samples from Q .
Hencewe want the scores to be well-approximated in regions of high-probability mass under Q .
In practice, the diffusions corresponding to Q and P θ have to be discretized, so a more direct routeadopted here is to simply take inspiration from (11) and to consider the parameterized backwardkernels B θk ( x k | x k +1 )
=
N ( x k ; x k +1 − δ ∇ log π k +1 ( x k +1 )
+ 2 δs θ
( k + 1 , x k +1 ) , 2 δI ) to obtain aparameterized extended target P θ and corresponding unnormalized target Γ θ .
We then learn θ byminimizing D KL ( Q || P θ ) where
[Equation]
This is obviously equivalent to maximizing the ELBO E Q [log w θ ( x 0: K )] where w θ ( x 0: K ) =
Γ θ ( x 0: K ) /Q
( x 0: K ) .
We note that it has previously been proposed to learn parameterized backward kernels for general AIS proposals [Salimans et al., 2015; Huang et al., 2018].
However, theparameterization adopted therein, B θk ( x k | x k +1 ) =
N ( x k ; µ θ ( x k +1 ) , Σ θ ( x k +1 )) , does not leveragethe structure of the true reversal and performs poorly experimentally [Thin et al., 2021, Section 4.2].
As established in the next proposition, the continuous and discrete time approaches coincide forδ ≪ 1 .
Once θ is learned, we then obtain an unbiased estimate of Z through w θ ( x 0: K ) for x 0: K ∼ Q .
Proposition 2.
Under regularity conditions, we have
[Equation]
(14) where t k = kδ , K = T/δ , q t | s ( x ′ | x ) is the density of x t = x ′ given x s = x under Q and C 1 , C 2constants independent of θ .
Let L ( θ )
=
δ (cid:80)
Kk =1 E Q (cid:2)
|| s θ
( k, x k ) − ∇ log F k ( x k | x k − 1 ) || 2 (cid:3) denote adiscrete-time approximation of this loss.
We have ∇ D KL ( Q || P θ ) =
∇L ( θ ) + ϵ ( θ ) for some functionϵ satisfying lim K →∞ ϵ ( θ ) = 0 .
Equation (13) shows that D KL ( Q||P θ ) corresponds to a score matching loss as for SGM [Songet al., 2021a].
It is possible to rewrite this loss as (14) so as to replace the intractable score term∇ log q t ( x t ) by the easy to approximate gradients of the log-transitions ∇ log q t ( x t | x t k − 1 )
[Vincent,2011].
In practice, as mentioned above, we simply learn θ by minimizing the discrete-time KLdiscrepancy D KL ( Q || P θ ) .
This formulation is also very convenient as we can additionally learnpotential parameters ϕ of a Q ϕ using the same criterion.
Note from equation (8) that the KLdivergence decomposes as firstly a term penalizing the difference in the approximating measure andthe fixed target at the final time and secondly another term which can be reduced by optimization ofboth Q ϕ and P θ conditioned on x T ; see e.g. [Agakov and Barber, 2004].
Pseudo-code for our approach (with a comparison to the AIS algorithm proposed in Heng et al.[2020]; Wu et al.
[2020]; Thin et al.
[2021]) can be found in Algorithm 1.
3.3 Incorporating Hamiltonian dynamics via the underdamped Langevin equation
We now consider a proposal defined on an extended space which arises from the time-discretizationof a time-inhomogeneous underdamped Langevin dynamics; see e.g. [Leimkuhler and Matthews,2016, Chapter 6].
In this scenario, we first focus on continuous-time as the development of suitablenumerical integrators is much more involved than for overdamped Langevin diffusions.
We considerthe diffusion ( x t , p t ) t ∈ [0 ,T ] where p t ∈ R
d is a momentum variable
[Equation]
Require:
Unnormalized target γ ( x ) , initial state proposal π 0 ( x ) , number steps K , stepsize δ , annealing schedule { β k } Kk =0 , score model s θ ( k, x )
Sample x 0 ∼ π 0 ( x 0 )
Set log w = − log π 0 ( x 0 ) for k = 1 to K do
Define log γ k ( · ) =
β k log γ ( · ) + (1 − β k ) log π 0 ( · )
Define F k ( x k | x k − 1 ) =
N ( x k ; x k − 1 + δ ∇ log γ k ( x k − 1 ) , 2 δI )
Sample x k ∼ F k ( x k | x k − 1 )
Define B k − 1 ( x k − 1 | x k )
=
F k ( x k − 1 | x k ) ▷ AIS
Define B k − 1 ( x k − 1 | x k ) =
N ( x k − 1 ; x k − δ ∇ log γ k ( x k ) + 2 δs θ ( k, x k ) , 2 δI )
▷ MCD
Set log w = log w
+
log
B
k − 1 ( x k − 1 | x k ) − log F k ( x k | x k − 1 ) end for
Set log w = log w + log γ ( x K ) initialized at x 0 ∼ π 0 , p 0 ∼ N (0 , M ) defining the path measure Q .
Here M is a positive definite massmatrix, γ > 0 a friction coefficient and ( B t ) t ∈ [0 ,T ] a multivariate Brownian motion.
If π t was not timevarying, the invariant distribution of this diffusion would be given by ¯ π t ( x, p ) =
π t ( x ) N ( p ; 0 , M ) .
Intuitively, in the time varying case the SDE will have enough time to approximate each intermediate¯ π t ( x, p ) if we change the target sufficiently slowly.
We can think of underdamped Langevin as acontinuous-time version of Hamiltonian dynamics with continuous stochastic partial momentumrefreshment [Horowitz, 1991].
From [Haussmann and Pardoux, 1986], the time-reversal of the diffusion (15) is also a diffusionprocess (¯ x t , ¯ p t ) t ∈ [0 ,T ] =
( x T − t , p T − t ) t ∈ [0 ,T ] given by (¯ x 0 , ¯ p 0 ) ∼ η T and
[Equation]
[Equation]
where η t denotes the density ( x t , p t ) under (15).
In this case, the continuous-time version of P opt isthe path measure P opt defined by the diffusion (16) but initialized at ¯ x 0 ∼ π T , ¯ p 0 ∼ N (0 , M ) ratherthan η T .
We will approximate it by the path measure P θ using a neural network s θ ( T − t, ¯ x t , ¯ p t ) inplace of ∇ log η T − t (¯ x t , ¯ p t ) in (16), i.e. we consider
[Equation]
[Equation]
As for overdamped Langevin, we could also learn θ by minimizing D KL ( Q||P θ ) over θ .
This againcorresponds to minimizing a score matching loss albeit of a form slightly different from (13).
Proposition 3.
Under regularity conditions, we have
[Equation]
where || x || M :=
u T Mu , t k = kδ , K = T/δ , C 1 , C 2 are constants independent of θ andη t | s ( x ′ , p ′
| x, p ) is the density of ( x t , p t ) =
( x ′ , p ′ ) given ( x s , p s ) =
( x, p ) under Q .
While the continuous-time perspective shed light on how to parameterize an approximation to the timereversal, this does not lead directly to an implementable discrete-time algorithm for underdamped
Langevin.
Contrary to overdamped Langevin, we cannot indeed simply use an Euler discretization of(15) defining Q and (17) defining P θ to obtain some discrete-time forward and backward kernelsand then compute w θ ( x 0: K , p 0: K )
=
Γ θ ( x 0: K , p 0: K ) /Q
( x 0: K , p 0: K ) .
This is because this ratio isnot well-defined due to the lack of noise on the position component in both (15) and (17).
The integrator we use for the forward equation (15), consists in alternating partial momentumrefreshments and deterministic leapfrog steps (see e.g [Leimkuhler and Matthews, 2016; Neal, 2011])giving
[Equation]
Require:
Unnormalized target γ ( x ) , initial state proposal π 0 ( x ) , number steps K , stepsize η , annealing schedule { β k } Kk =0 , damping coefficient h , mass matrix M , score model s θ ( k, x, p )
Sample x 0 ∼ π 0 ( x 0 ) , and p 0 ∼ N ( p 0 ; 0 , M )
Set log w = − log π 0 ( x 0 ) − log N ( p 0 ; 0 , M ) for k = 1 to K do
Define log γ k ( · ) =
β k log γ ( · ) + (1 − β k ) log π 0 ( · )
Sample ˜ p
k ∼ N (˜ p k ; hp k − 1 , (1 − h 2 ) M )
Set µ q = p k − 1 Set µ p = ˜
p
k ▷ UHA reversal mean
Set µ p = ˜
p k − 2 log( h )
[ Ms θ ( k, x k − 1 , ˜ p k ) + ˜ p k ] ▷ MCD reversal mean
Set log w = log w + log N
( p k − 1 ; hµ p , (1 − h ) 2 M ) − log N (˜ p k ; hµ q , (1 − h ) 2 M )
Run leapfrog integrator on γ k and set ( x k , p k ) =
Φ( x k − 1 , ˜ p k ) end for
Set log w = log w + log γ ( x K ) + log N ( p K ; 0 , M ) with h = exp {− γδ } and Φ t is the leapfrog integrator for π t .
The resulting forward sampler is similarto the one proposed by Geffner and Domke [2021], except we do not flip the momentum after theleapfrog step 2 .
This integrator may be interpreted as a splitting method for equation (15); see e.g[Leimkuhler and Matthews, 2016, Chapter 7].
We need the integrator for the reversal to fulfill two criteria.
First, by definition, as the time stepδ → 0 it must recover the SDE (17).
Second, the importance weight of the forward sampler tothe reversal must be well defined.
Since the leapfrog integrator is a diffeomorphism (or flow) theonly possible way to get a well defined reversal for these steps is to take the inverse Φ − 1 t .
Since thetransformation is also volume preserving, the contribution from the deterministic forward and reverseterms will then exactly cancel in the importance weight.
The required form of the reverse integrator is(
[Equation]
t ,
[Equation]
where f θt + δ ( x t , ˜ p
[Equation]
+ δ ) :
= (cid:101) p t + δ + δ 2 γ
[ Ms θ ( t, x t , (cid:101) p t + δ ) + (cid:101) p t + δ ] .
In Appendix C we show that as 224 δ → 0 this can indeed be interpreted as a valid split integrator for the reverse SDE (17).
The crucialpoint of algorithmic difference from [Dai et al., 2020; Geffner and Domke, 2021; Zhang et al., 2021]arises from our necessary form for the mean of the reverse momentum refreshment.
These worksuse h ˜
p
t
+ δ as the mean instead of hf θt + δ ( x t , ˜ p t + δ ) .
We again transition to discretized notation withδ :
= T/K , and k = 0 , ..., K .
In this case, the log importance weight, corresponding to the logevidence estimate, satisfies
[Equation] [Equation]
where ( x , p ) denote all the variables introduced by our integration scheme.
We can show informallythat minimizing D KL ( Q || P θ ) , i.e. maximizing the ELBO given E Q [log w θ ( x , p )] , again correspondsto minimizing the score matching loss (18) when δ ≪ 1 ; see Appendix C. Pseudo-code for ourapproach can be found in Algorithm 2.
Experiments
We run a number of experiments estimating normalizing constants to validate our approach, MCDand compare to differentiable AIS with ULA [Wu et al., 2020; Thin et al., 2021] and Unadjusted
Hamiltonian Annealing (UHA) [Geffner and Domke, 2021; Zhang et al., 2021].
We first investigatethe performance of these approaches on static target distributions using the same, fixed initialdistribution and annealing schedule.
Finally, we explore the performance of the methods for VAEs.
Here, being the most expensive of our experiments, we include runtime comparisons of our methodcompared to baselines.
Additional results on a Normalizing Flow target can be found in Appendix
F.2.
Full experimental details, chosen hyper-parameters, and model architectures can be found in
Appendix E.
Our score model is parameterized by an MLP with residual connections that is conditioned onintegration time t , and on the momentum term for the Hamiltonian case (see Algorithm 2).
For anablation on various network architectures we refer the reader to Appendiex F.3.
4.1 Static Targets
We estimate the normalizing constants of five simple distributions with known normalizing constantsequal to log Z = 0 ; N (10 , I ) , N (0 , 0 .
1 I ) , a Gaussian mixture with 8 components whose means aredrawn from N (3 , I ) where each component has variance 1, a standard Laplace distribution and a
Student’s T distribution with 3 degrees of freedom.
We use N (0 , I ) as our initial distribution forall targets except for the Gaussian mixture and N (0 , 0 .
1 I ) which both use N (0 , 3 2 I ) .
We run eachmethod using K ∈ { 64 , 256 } steps and use a fixed, linear annealing schedule.
For all methods,sampling step-sizes per-timestep are tuned to via gradient descent to maximize the ELBO and thediagonal mass matrix is learned for the Hamiltonian samplers.
Gaussian Mixture and Student-Tresults can be found in Tables 1 and 2, respectively.
Additional results can be found in Appendix F.1.
[Table] [Table]
[Table] [Table]
[Table] [Table]
On average and as expected, UHA outperforms ULA on most targets.
Further our approximation tothe optimal backward kernels can yield considerable improvements – enabling the ULA sampler toproduce better results than UHA (using the standard AIS backward kernels).
Further, we see that ourapproximation to the optimal backward kernels improves the performance of UHA as well.
We further emphasize that often ULA-MCD and UHA-MCD with 64 steps outperform or is on-parwith ULA and UHA with 256 steps, a difference of factor 4 in terms of target gradient evaluations.
As we show below in Table 3, the additional computational costs of fitting the score model of ourmethod are only roughly twice that of the baselines.
[Table] [Table]
[Table] [Table]
4.2 Application to Amortized Inference
Next, we explore the application of our method to amortized inference, in the context of VAEs[Kingma and Welling, 2013; Rezende et al., 2014].
These models are trained to infer latent representations using an inference neural network that consumes an input and produces parameters of anapproximation to the true posterior distribution of the underlying generative model.
In particular,this posterior distribution is different for each input.
When applying AIS to VAE inference [Thinet al., 2021; Geffner and Domke, 2021; Zhang et al., 2021], the output of the inference networkparameterizes the initial distribution of the annealing sequence to the true posterior.
By training thisend-to-end, we effectively learn the initial distribution for the diffusion process.
Consequently, thediffusion marginals and their score vectors ∇ log q t are different for every input, and we need to condition our score model s θ on the inputs to reflect that.
We achieve this by projecting the last hiddenlayer of the inference network into a summary vector that is concatenated to the other conditioninginputs of s θ .
We train a VAE on the binarized MNIST dataset [Salakhutdinov and Murray, 2008], re-using architectures proposed in [Geffner and Domke, 2021; Burda et al., 2016] (two-layer MLP encoder/decoder,
Bernoulli likelihood).
All generative models use the same architecture and hyper-parameters.
Wecompare standard amortized variational inference with annealed ULA and UHA with standard AISbackward transition kernels, as well as ULA and UHA with our MCD transition kernels.
We match thenumber of sampler steps between ULA/ULA-MCD and UHA/UHA-MCD to 64 and 32 respectively.
ELBO and log-likelihood on the test set can be found in Table 3.
[Table]
We see that, as reported in prior works [Geffner and Domke, 2021; Thin et al., 2021; Zhang et al.
,2021], Monte Carlo based inference methods provide a significant benefit over standard amortizedvariational inference.
In addition, our learned backward kernels lead to improved performance ofboth ULA and UHA, with the difference being more distinct for the latter.
We also note that UHAhas significantly larger error bars than all other methods, but UHA-MCD does not appear to inheritthis.
We further provide a runtime comparison in Table 3, both for a single training iteration, and totalexperiment time (including evaluation).
We find that fitting the score models of our method results inroughly twice as much runtime (recall we fixed the number of sampler steps).
Limitations
While we have demonstrated that our optimized backward kernels can lead to large improvementsover standard AIS backward kernels, our proposed approach has a number of limitations.
First,since our method relies upon unadjusted Langevin or Hamiltonian sampling, we inherit many ofthe issues with these approaches.
Using unadjusted samplers enables a fully-differentiable ELBOestimate which can in theory be used to tune the many parameters of the sampler.
However, thesesamplers require repeated gradient steps as an inner-loop which can lead to divergence dynamicsand numerical instability, making optimization difficult.
Next, our proposed training procedure hasa notable increase in memory consumption as we store the entire sampling trajectory in memoryto train our neural network.
We further note that it is possible to compute our importance weightsonline using O (1) memory.
We could utilize this with Monte Carlo sub-sampling of the timestep (aswith SGM [Ho et al., 2020]) to derive a constant memory training procedure for our approximatetime-reversal at the cost of increased variance but we leave exploring this to future work.
Finally, weobserved that, without mitigation, parameterized forward samplers Q ϕ sometimes dropped modesduring training.
Since this is a known challenge of using reverse KL, this might be improved by usingalternative optimization objectives [Masrani et al., 2019; Matthews et al., 2022].
Discussion
In this work we have explored AIS using unadjusted Langevin and Hamiltonian dynamics.
Wehave demonstrated that the backward transition kernels typically used are suboptimal and we havepresented the form of the optimal variance-minimizing backward kernels.
We have further shownhow an approximation to these kernels can be learned using score matching and that this objectiveconverges to maximizing the ELBO in the limit of infinitesimally small time-discretization step-sizes.
We have illustrated the benefit of using our proposed optimized backward kernels on a number ofinference problems including fixed targets and amortized tasks with model learning.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] Paper is
methodological paper on approximate inference methods with only standard datasets.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See proofs ofmain results in Appendix.(b)
Did you include complete proofs of all theoretical results?
[Yes] See Appendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Appendixfor architectures, Algorithm 2 for pseudo-code.
We used standard datasets from the
literature.
We did not publish code.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See experimental section and Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] Where appropriate and feasible.
(d) Did you include the total amount of compute and the type of resources
used (e.g., typeof GPUs cluster, cloud provider)
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]