Score-Based Diffusion meets Annealed Importance Sampling
Authors: Arnaud Doucet
Will Sussman Grathwohl
Alexander G. D. G. Matthews
Heiko Strathmann
Venue: NeurIPS 2022 Accept (2022)
Keywords: importance sampling
mcmc
diffusion

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 tractable initial distribution and the target distribution of interest which we simulate from approximately using a non-homogeneous Markov chain. To obtain an importance sampling estimate of the marginal likelihood, AIS introduces an extended target distribution to reweight the Markov chain proposal. While much effort has been devoted to improving the proposal distribution used by AIS, by changing the intermediate distributions and corresponding Markov kernels, an underappreciated issue is that AIS uses a convenient but suboptimal extended target distribution. This can hinder its performance. We here leverage recent progress in score-based generative modeling (SGM) to approximate the optimal extended target distribution for AIS proposals corresponding to the discretization of Langevin and Hamiltonian dynamics. We demonstrate these novel, differentiable, AIS procedures on a number of synthetic benchmark distributions and variational auto-encoders.

Score-Based Diffusion meets Annealed Importance Sampling
Arnaud Doucet, Will Grathwohl, Alexander G. D. G. Matthews Heiko Strathmann DeepMind
(arnauddoucet, sg
Abstract
More than twenty years after its introduction, Annealed Importance Sampling (AIS) remains one of the most effective methods for marginal likelihood estima- tion. It relies on sequence of distributions interpolating between tractable initial distribution and the target distribution of interest which we simulate from approximately using a ho-hoooggeeous Markov chain. To obtain an importance sampling estimate of the marginal likelihood, AIS introduces an extended target distribution to reweight the Markov chain proposal. While much effort has been devoted to improving the proposal distribution used by AIS, an underappreciated issue is that AIS uses convenient but suboptimal extended target distribution. We here leverage recent progress in score-based generative modeling (SGM) to approximate the optimal extended target distribution minimizing the variance of the marginal likelihood estimate for AIS proposals corresponding to the discretization Langevin and Hamiltonian dynamics. We demonstrate these novel, differentiable, AIS procedures number of synthetic benchmark distributions and variational auto-encoders.
Introduction
Evaluating the marginal likelihood, also known as evidence, is of key interest in Bayesian statistics as it allows not only model comparison but also often used select hyperparameters. large variety of Monte Carlo methods have been proposed to address this problem, including path sampling [19], AIS [37] and related Sequential Monte Carlo methods [13]. An appealing feature of AIS is that provides an unbiased estimate of the marginal likelihood and can thus be used to define an evidence lower bound (ELBO) or mutual information bounds; see e.g. [53, 51, 7].
AIS builds a proposal distribution using a Markov chain (= initialized at an sasy-l distribution followed by sequence of Markov chain Monte Carlo (MCMC) transitions targeting typically annealed versions of the posterior. By proceeding this way, we obtain a proposal ææ whose distribution is expected to be reasonable approximation to the target posterior. However, this distribution is intractable as it requires integrating the joint proposal distribution over previous states AIS bypasses this issue by instead using Importance Sampling (IS) on the whole path (xx)) k= through the introduction of an artificial extended target distribution whose marginal at time K coincides with the posterior.
(xx) k=
(
There has been much work devoted to improving AIS in machine learning and statistics but also in physics where it was introduced independently in [29, 9]. A standard approach to improve AIS to modify the intermediate distributions [45, 21, 34] and corresponding transition kernels of proposal [10, 53, 18, 51, 54]. We here address distinct problem. For given proposal, it was shown in [13] that the extended target distribution minimizing the variance the evidence estimate is not
alphabeticaa order, equal contribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Xo ~ n0
dx, = V log (X, + V2dB,
XT ~M
Initial distribution
Target distribution
XT K ax,
V log ( + log (r-) 22B,
Xo ~ n
Figure 1: Top: Samples X, from AIS proposal (red) obtained by sampling initially from Gaussian at f and diffusing through Langevin dynamics on intermediate targets πn (white). The intermediate marginals of the proposal, g.. approximated by the samples are such that @r for reasonably fast mixing diffusion. Bottom: Computing importance weights. The optimal extended target used compute the weights is the distribution obtained by initializing Xo exactly from æ and then following the reverse- time dynamics the forward AIS proposal. This requires access to score vectors of the marginals gr.
the one used by AIS but is instead defined through the time-reversal of the proposal. However, this result difficult 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 obtain improved, lower variance, AIS estimates. We concentrate on scenarios where we use unadjusted overdamped Langevin [23, 53, 51] and unadjusted Hamiltonian proposals with partial momentum refreshment (i.e. underdamped Langevin) [10, 53, 18, 54, 28] which correspond time-discretized diffusion processes. The first benefit of using such proposals is that, by omitting Metropolis Hastings steps, one obtains differentiable versions of the Evidence Lower Bound (ELBO) amenable to the reparameterization trick. The second benefit of these proposals is that their time-reversal can be approximated by adapting techniques developed for SGM [24, 48, 15] to our setup. We derive principled parameterization for approximation of their time-reversal which we learn by maximizing the ELBO. As for SGM, this ELBO coincides with a denoising score matching loss [27, 52, 24, 48]. This provides novel, optimized and differentiable, AIS estimators which refer as Monte Carlo Diffusion (MCD). We demonstrate the benefits of this approach on synthetic benchmark distributions and variational auto-encoders (VAEs) [31]. All proofs can be found in the Appendix. A preliminary version of this work appeared [16].
2
Annealed Importance Sampling
2.1 Setup and algorithm
Consider probability density v on Rd of the form
7(x)
(σ)
2= 7x))dx
(1)
Z
Jud
where (xx) can be evaluated pointwise. We want to approximate the intractable normalizing constant
Z. Bayesian framework, (xx = p(x) x)) the joint density of parameter and data D,
(x) = p(x D) the corresponding posterior and Z p(D) the evidence.
To estimate AIS introduces the intermediate distributions ( bridging smoothly from tractable distribution To to the target distribution π π of interest. One typically uses T& (x) oc
(x(x) with 7x (x) = To Tx(x)) (x:) (x) 8 for 0 = 3o < 31 < ... < 3K = but other choices are possible [21]. The IS proposal used by AIS is then obtained by running Markov chain (
such that 20 ~ no (-), and then Iy ~ F& (-|æx-1) for k 2 1 where F& is a MCMC kernel invariant
The The proposal is thus given by
Q(æo.k) To(o)) IT Fa Fx(x//0-111
(2)

---

Denote by Qh the marginal distribution of 26 under Q satisfying 9 (x(2)) 1 Fx (2
( for k 2 1 and qo up. it is typically intractable for 2 1, As чк cannot be evaluated in complex scenarios, the marginal IS estimate Wmar (x (( k//kk//0 //1x)//x(/2) (TK) of Z is intractable.
One can bypass this issue by introducing an extended target distribution
(( P(xo:x)= = 7
K T( = >(xx) II
Be(
(3)
k=0
where (=) are backward Markov transition kernels, i.e. I Be ( ddr 1 for any 2h+1+ so that by construction æ ~ under P. For any selection of backward kernels such that the ratio T/Q well-defined, we then have
[(
Bo(w[(Fok]] Z, for w(xox)) Q(Tu:K)²
i.e. w<(o:k) is an unbiased estimate of Z for TO:K Q.
The AIS estimate of the evidence is specific instance of the estimator (4) relying on the back- ward kernels [xx| = Th++1 T FR+1 1 [x))/ TR////TTIff (#++1). This yields the following expression for log (xo:k)):
log wais (TO:K) = Ex=n log (7/
(5)
2.2 Limitations of AIS
While designing P in (3) by using the backward Markov kernels (BT*)) k== is convenient, is also sub-
R-iisyK-1 K-1
optimal in terms of variance. For example, consider the ideal scenario where Fs (x = n (xx). This scenario has been used many times in the literature to provide some guidelines AIS, see
e.g. [37, 21]. In this case, varqllog Wais (Io.K) while gx [maar [xx)] = var, [m (Ix) = O.
K var
Another illustration of the suboptimality of AIS is to con- sider a scenario where the proposal is a homogeneous MCMC chain, i.e. 3o To and Ji& F(-[**-1) for F a m-invariant MCMC kernel; i.e. use F& = F and Tk = for k 1....y K. If F is reasonably well-
x(0,1)
logz estiante
-25
mixing, then s т for K large enough and the evi- dence estimate Wmar (Ix) = (x //)//)/////))))))
should have
-75
AIS
MCD
small variance. However, it is easy to check that we have Was (Fo: ( Fu (xro) for the exact same proposal; i.e. the AIS estimate does not depend on the MCMC sam- ples æ1:K and boils down the estimate of Z using the proposal Wo-
-100
50
100
150
200
250
steps
Figure 2: Comparing log Z estimates as a function of using AIS and MCD. Both estimates use the same forward kernels but reweight samples in a differ- ent way using distinct backward kernels. Initial distribution Tn is 20-dimensional ((0, I) and progressively shifts to the target = т = (10, I). The MCD estimate much closer the ground truth (log Z 0) than AIS.
These examples illustrate that it would be preferable
to use (T rather than Wai (Fo:k). In Appendix A, provide detailed comparison of both estimates in
scenario where their variance can be computed analytically, We propose in the next section an unbiased estimate of the evidence (MCD) approximating Wmar (m based on a different choice of backward kernels. As illustrated in Figure 2, significant gains can be achieved.
3 Optimized Annealed Importance Sampling
We show here that the optimal extended target distribution P minimizing the variance of the evidence estimate (4) is defined through the time-reversal of the proposal Q. By exploiting a connection to SGM, we can approximate this reversal using score matching when the proposal is obtained through an unadjusted overdamped or underdamped Langevin algorithm.
3

---

3.1 Optimal Extended Target Distribution via Time Reversal
We summarize here Proposition 1 of [13]; see also [46, 5}.
Proposition 1. For proposal Q form (2), the extended target P of the form (3) minimizing both the Kullback-Leite divergence ( and the variance of the evidence estimate w w(ook) I ( (/Q( for 2o.K is given by (TM(k) (Fo: / where
K-1
BR (xx| 9b (xxx)) (
(ra:k)) ((xx)
II
BR (x=())))),
(6)
9k+1 (x++1)
In particular, one has
7(ax)
Dopt (2o.k) and DKI[Q] = DKI((KK|TT)
Wmar(Tr)
(7)
qK (xxx
(xTo.))
This result follows simply from the chain rule and the law of total variance which yield
(Ql(P) = (gxl() + Eyg (Q(- [[(x((()==x))))((((())))) varq [w][Foo ) = var 2x (Mma(((x)) )] + Eg [varo (
(8)
(9)
Both quantities are clearly minimized by selecting P(mo:x-1/xx) = ((Fo
We emphasize that Proposition 1 applies any forward kernels ( including MCMC kernels, unadjusted Langevin kernels or even deterministic maps?. It shows that popt is the distribution of a backward process initialized at æ which then follows the time-reversed dynamics of the forward process Q. If we had 9K x, then we would have pope Q as then popt would correspond to the backward decomposition of Q.
3.2 Time reversal, Score matching and ELBO for unadjusted overdamped Langevin
We concentrate here on the case where (Fi) correspond to a time-inoomogeneous unadjusted (overdamped) Langevin algorithm (ULA) as used in [23, 53, 51]; that is we consider F (1 (xx-1) = N(zx; -1 + 8V log ThK (x) 2 where >> 0 is a stepsize. Let 8 == T/K K then, as K + oc, the proposal Q converges to the path measure of the following inhomogeneous Langevin diffusion (xe ( defined by the stochastic differential equation (SDE)
da.
V log n. (x;)dt + V2dB.,
J0 ~ To
(10)
where (Bi)te(o,T is standard multivariate Brownian motion and are slightly abusing notation from now on as Te fort t = t = kō corresponds to Te in discrete-time. Many quantitative results measuring the discrepancy between the law of or and रт for such annealed diffusions have been obtained;
see e.g. [17, 50). From [22], is known that the time-reversed process (ã+) = (xT-+))e(o.)) is also
diffusion given by
dz = Vlog (3t) + 2V 2V log (E.)) + Vdd
Zg qr
(11)
where (Bo)eeto.2) another multivariate Brownian motion. The continuous-time version of popt the path measure popt defined by the diffusion (11) but initialized at Tp 7 rather than @r as
noted in [5]; see Figure/ for an illustration. This shows that approximating (B )=== requires K-1
approximating the so-called scores (V log (x))))=((or) This can be derived heuristically through the fact that a Taylor expansion yields the following approximation of the optimal backward kernels, B (2x (0)((=++1) s (Zx;;; #x+1 8V log T (#x+1) + 28 log qk+ 2881), which indeed
corresponds to a Euler discretization of (11); see c.g. [12, Section 2.
In SGM [48], one gradually adds noise to data using
diffusion to transform the complex data distribution into a Gaussian distribution and the generative model is obtained 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 using denoising
"The case of deterministic maps corresponds to normalizing flow components, where the inverse flow is the optimal and only valid reversal, e.g. [2] which includes detailed literature review.

---

score matching [52]. While in our setup, the diffusion (10) instead goes from a simple distribution to a complex one (see Appendix D for discussion), we can still use score matching ideas. We define path measure Po approximating poot using a neural network so(T I, Tt) in place of V log (
(11), i.e. we consider
dē = { V lo + (T t, + v2dB,
To
(12)
We would like learn by minimizing (el(P)) over 8, i.e. equivalently we maximize continuous-time ELBO. Note that it is neither easily feasible to minimize DKI. (PoM())) (as one cannot sample from x) nor it is desirable as the evidence estimate is computed using samples from Q.
Hence we want the scores to be well-approximated in regions of high- mass under Q.
In practice, the diffusions corresponding to and Pg have be discretized, so more direct route adopted here is to simply take inspiration from (11) and to consider the parameterized backward kernels Bº (zz (2 N (z . 8V log Th+ (I+++1 + 20sg (k + L, (.++1),, 1), 261) to obtain parameterized extended target Po and corresponding unnormalized target rg. We then learn @ by minimizing PQ where
Q(Fo:)) (x)) I FR+1 (F++( (o = オ (Ix) I (((
This is obviously equivalent to maximizing the ELBO Eg llo wy (2):k))) where wy(2r.k) I' I'o ) //((zox)).. We note that it has previously been proposed to learn parameterized back- ward kernels for general AIS proposals [44, 26]. However, the parameterization adopted therein, (x (=+(22++) ((x no (zx+1)), Eg (2)++))), does not leverage the structure of the true reversal and performs poorly experimentally [51, Section 4.2].
As established in the next proposition, the continuous and discrete time approaches coincide for 5 < 1. Once 0 is learned, we then obtain unbiased estimate of Z through we (2(. for Fo:K Q.
Proposition 2. Under regularity conditions, we have
[ [1s)(t, - Vlog + Ci
DRt(21/P)) -Ee
(13)
K
Be [1lss(t, - V log atgg ( (0 dt + +22
(14)
k=1
where t, = kō, = T/5, 9//s (1)=(=((z)) (z) is the density of xt = x' given I, under Q and Ct Cz constants independent of Let 6 EQ [ks xx.) - V log Fs F= denote a discrete-time approximation of this loss. We have DKL(QP))P)) = V (0) + = for some function satisfying lim (0) O.
Equation (13) shows that (KL(2117) corresponds score matching loss as for SGM [47]. It is possible to rewrite this loss as (14) so as to replace the intractable score term V log 9: (xx) by the easy to approximate gradients of the log- -ra-iiannitions log qe (2x (x.,) [52]. In practice, as mentioned above, we simply learn f by minimizing the discrete-ting time KL. discrepancy D
This formulation is also very convenient as we can additionally learn potential parameters 6 of
Qo using the same criterion. Note from equation (8) that the KL divergence decomposes as firstly term penalizing the difference the approximating measure and the fixed target at the final time and secondly another term which can be reduced by optimization of both Qo and Po conditioned on 2T; see c.g. [1].
Pseudo-code for our approach (with comparison to the AIS algorithm proposed in [23, 53, 51J) can be found in Algorithm
3.3 Incorporating Hamiltonian dynamics via the underdamped Langevin equation
We now consider a proposal defined on an extended space which arises from the cime-discretizatiin of tim--inhomogeneuus underdamped Langevin dynamics; e.g. [33, Chapter 6]. In this scenario, we first focus on continuous-time as the development of suitable numerical integrators is much more involved than for overdamped Langevin diffusions. We consider the diffusion (Ii, Pe (e) where
Pe E Rd is momentum variable
dz,- . -1 p,,,, -M'pp,
dp: V
Sp,dt + 2/2
(15)
5

---

Algorithm Unadjusted Langevin AIS/MCD - red instructions for AIS and blue for MCD
Require:
Unnormalized target 7(z), initial state proposal To(z), number steps K, stepsize an- nealing schedule {B3 } score model sø(k,z)
Sample 2g (T))
Set log = - log (Too
for
K do
Defin log 7)c (-) 3R. log (.) + B&) log To(-)
Define (r [K 1) = (Tx 1k-1 8 YI: (Tx- (Ix-1),281) 251)
Sample n&
F-
Define Bx (Tk- [xk) = (IR- (xx)
AIS MCD
Define Ba (TR- (xx) = N (xx Ih 8V log (xx) 205a (ki 261)
Set log w log w + log B (T (TR) log (2
end for
Set log w = log w + log (xKK
initialized at 20 no. Po (O, M) defining the path measure Q. Here M is positive definite mass matrix,
a friction coefficient and (Br) a multivariate Brownian motion. m was not time-varying, e.g. %t. ग,, the invariant distribution of this diffusion would be given by T (2, = ग (x) 0, M). Intuitively, in the time varying case the SDE will have enough time to approximate each intermediate कt (2. p) (p) 0, M) if we change the target sufficiently slowly. We can think of underdamped Langevin as a continuous-tin dynamics with continuous stochastic partial momentum refreshment [25].
version of Hamiltonian
From [22], the time-reversal of the diffusion (15) also a diffusion process (Tx, Pc)et.)
(wT- PT (o,) given by (Zo. Po) 1/ and
dæ. = M pidt,
(16)
dø. = V log (Z))( + (P:dt + 2ÇM V,,, log (.) /2CM
where 7/ denotes the density (xt, Pb) under (15). In this case, the continuous-time version of Popt is the path measure Pon defined by the diffusion (16) but initialized To t,, NN M) rather than 7T. We will approximate it by the path measure Pa using a neural network 8q (T 1, 1,Tr, in place of V log 1/r -1 (T,,Pu in (16), i.e. we consider
dø, -1 p,dl,
(17)
dp: = - v log nт (Z: (p.tt + 2çMso (T - t,,,, Fi V2CM2/2AB
As for overdamped Langevin, we could also learn 0 by minimizing ( PP) over 0. This again corresponds to minimizing
matching loss albeit of form slightly different from (13).
Proposition 3. Under regularity conditions, we have
(KI = <EE [ [|s,((, Vn logne(s,, Pl) C₁
(18)
=
Ee [[lso(t, Vp. log ( Pe Pel.teeerPP..... P13 dt + Cz,
where ||]]| == Mu, 4 = kō, K = T/8, C1, C2 are constants independent of 0 and (as p p) is the density of (3t, ,',,') given (x., Ps) = (x.p) under Q
While the continuous-time perspective shed light on how to parameterize an approximation to the time- reversal, this does not lead directly to 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ø obtain some discrete-tima forward and backward kernels and then compute wy (To:K Po:k) ) = (o Po:k) not well-defined due to the lack of noise the position component in both (15) and (17).
) (z. Po.k). ). This is because this ratio is
The integrator we use for the forward equation (15), consists in alternating partial momentum refreshments and deterministic leapfrog steps (see e.g [33, 38]) giving
P148 Whpp, (1 M)²)M), (I Pt48)) = ( Pc-s),
(19)

---

Algorithm 2 Unadjusted Hamiltonian AIS/MCD red instructions for AIS and blue for MCD
Require:
Unnormalized target 7(z), initial state proposal (o( number steps K, stepsize 77 an- nealing schedule {B5 } J damping coefficient h, mass matrix M, score model sy (k, x, x.
Sample 2g (Tu) and pu (pu : 0, MM
Set log = - log (To) - log. Log.Von (pu: 0, MM)
for
to K do
Defie log 7x (.) 3r log 7.) + (1 - Be) log (-
Sample fi (Nhhp-- 1, M)
Set Myg = Pk-1
Set μ = Pk
B UHA reversal mean
Set M4 = Pk 2 (
Pil
MCD reversal mean
Set log log w + log. ((p hypp, (1 M) ) log (p hyg: (1 - M)
Run leapfrog integrator on 7h and set (2k. = ((xx-1,Pk)
end for
Set log w = log w + log (Fx) + log
with h = - and dg is the leapfrog integrator for тr. The resulting forward sampler is similar to the one proposed by [18], except we do not flip the momentum after the leapfrog step'. This integrator may be interpreted as a splitting method for equation (15); see e.g [33, Chapter 7].
We need the integrator for the reversal fulfill criteria. First, by definition, as the time step - O must recover the SDE (17). Second, the importance weight of the forward sampler to the reversal must be well defined. Since the leapfrog integrator is diffeomorphism (or flow) the only possible way to get well defined reversal for these steps is to take the inverse Φ, ' As the transformation is also volume preserving, the contribution from the deterministic forward and reverse terms will then exactly cancel in the importance weight. The required form of the reverse integrator is
(It, = o,- (II+S Pt+55,
Pr (h + d, 2r, Pe+s), (1 1²)
(20)
where So(l + 8, Fe, P1+6) := pe+s + 82C[M (1, wy: Pt+6) + Pe+5]. Appendix C we show that as 5 - 0 this can indeed be interpreted point of algorithmic difference from [10, 18, 54] arises from our necessary form for the mean of the reverse momentum refreshment. These works use hpi+5 the mean instead of h + d, It, Pc-5). We again transition to discretized notation with 5 = T/K, and = O...., K. In this case, the log importance weight, corresponding the log evidence estimate, satisfies
valid split integrator for the reverse SDE (17). The crucial
log wa (x. p) = llog ( (PK: 0, M) To (zo) (po:: M)
log N(Pk-1::fool 1. (1-9)M)) (21)
W(pu: (1 12)
where (x, p) denote all the variables introduced by our integration scheme. We can show informally that minimizing (xI(([|PP)) i.c. maximizing the ELBO given EqQ [log we (x p)], again corresponds to minimizing score matching type loss (18) when 8 < 1; see Appendix C. Pseudo-code for our approach can be found in Algorithm 2.
4 Experiments
We run number of experiments estimating normalizing constants to validate our approach, MCD and compare to differentiable AIS with ULA [53, 51] and Unadjusted Hamiltonian Annealing (UHA) [18, 54]. We first investigate the performance of these approaches on static target distributions using the same, fixed initial distribution 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 method compared to baselines. Additional results on a Normalizing Flow target can be found in Appendix F.2. Full experimental details, chosen prprr-parametes, and model architectures can be found in Appendix E.
Our score model is parameterized by an MLP with residual connections that is conditioned on integration time f, and on the momentum term for the Hamiltonian case (see Algorithm 2). For an ablation on various network architectures we refer the reader to Appendix F.3.
[[18] were attempting discretize underdamped Langevin dynamics.
7

---

4.1 Static Targets
We estimate the normalizing constants of five simple distributions with known normalizing constants equal to log Z = 0: (10,, T). 0.17), Gaussian mixture with 8 components whose means are drawn from . I) where each component has variance 1, standard Laplace distribution and a Student's T distribution with 3 degrees of freedom. We use NN, () our initial distribution for all targets except for the Gaussian mixture and (0, 0.1f) which both use (0, 3²1). We run each method using K E {64, 256} steps and use fixed, linear annealing schedule. For all methods, sampling step-sizes per-timestep are tuned to via gradient descent to maximize the ELBO and the diagonal mass matrix is learned for the Hamiltonian samplers. Gaussian Mixture and Student-T results can be found in Tables and 2, respectively. Additional results can be found in Appendix F.1.
Sampler # steps
ULA
UHA
ULA-MCD
UHA-MCD
64
256
64
256
64
256
64
256
-0.47
-0.13
-0.03
0.015
0.01
0.01
0.01
-0.01
Dim-20
- 0.22
+ 0.07
+ 0.18
= 0.03
0.02
= 0.01
= =0002
- 0.01
-85.62
-21.98
-8.20
-1.26
-0.25
0.08
0.20
-0.05
Dim-200
+ 2.01
+ 1.35
+ 1.84
0.35
+ 0.03
+ 0.122
0.49
+ 0.04
-304.06
-83.50
-44.45
-8.60
-2.61
1.01
-1.74
-1.02
Dim-500
+ 5.48
+ 7.88
+ 3.24
+ 1.69
+ 1.26
0.99
+ 1.25
+ 0.03
Table 1: log Z estimates for a Gaussian mixture target. Averages and standard errors over 3 seeds.
On average and as expected, UHA outperforms ULA on most targets. Further our approximation to the optimal backward kernels can yield considerable improvements enabling the ULA sampler to produce better results than UHA (using the standard AIS backward kernels). Additionally, we see that our approximation 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-par with 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 our method are only roughly twice that of the baselines.
Sampler # steps
ULA
UHA
ULA-MCD
UHA-MCD
64
256
64
256
64
256
64
256
-0.09
-0.02
-0.04
-0.02
-0.06
-0.00
-0.03
-0.01
Dim-20
+ 0.02
- 0.01
- 0.02
0.01
+ 0.02
- 0.01
0.04
- 0.01
-1.63
-0.88
-0.36
-0.10
-0.82
-0.30
-0.48
-0.10
Dim-200
0.23
0.15
- 0.34
0.10
0.28
+ 0.38
0.20
- 0.07
-5.43
-3.10
-2.86
-1.03
-4.00
-2.23
-2.03
0.06
Dim-500
0.78
0.07
+ 0.36
+ 0.13
+ 0.51
- 0.18
0.18
0.30
Table 2: log Z estimates for Student-T T' target. Averages and standard errors over 3 seeds.
4.2
Application to Amortized Inference
Next, we explore the application of our method to amortized inference, in the context VAEs [31, 41]. These models are trained to infer latent representations using an inference neural network that consumes an input and produces parameters an approximation 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 [51, 18, 54], the output of the inference network parameterizes the initial distribution of the annealing sequence to the true posterior. By training this end-to-end, we effectively learn the initial distribution for the diffusion process. Consequently, the diffusion marginals and their score vectors V log 9: different for every input, and need to condition score model $a on the inputs to reflect that. We achieve this by projecting the last hidden layer of the inference network into a summary vector that is concatenated the other conditioning inputs of so-
We train VAE on the binarized MNIST dataset [43], re-using architectures proposed in [18, 8] (two- layer MLP encoder/decooder, Bernoull likelihood). All generative models use the same architecture
8

---

and hyper-ppaaeee parameters. We compare standard amortized variational inference with annealed ULA and UHA with standard AIS backward transition kernels, well ULA and UHA with our MCD transition kernels. We match the number of sampler steps between TLA//LLAMMC) and UHA/UHA- MCD to 64 and 32 respectively. ELBO and log-likelihood values on the test set are presented in Table 3.
Sampler
VI
ULA
UHA
ULA-MMD
UHA-MCD
ELBO
-9.32 + 0.40
-90.41 = 0.17
-88.58 = 0.51
-90.10 = 0.10
-88.08 + 0.07
Log-lik.
-69.35 + 0.36
-62.43
0.25
-61.05
1.84
-61.83
0.42
-58.58 0.34
Iteration time
0.024s
0.055s
0.050s
0.101s 11304.36s
0.098s
Total time
3263.40s
4694.52s
5072.19s
11345.59s
Table 3: Test set performance for MNIST VAE. Averages and standard errors over 5 seeds. We additionally report runtimes for single training iteration, and total experiment time.
We see that, as reported in prior works [18, 51, 54], Monte Carlo based inference methods provide significant benefit over standard amortized variational inference. In addition, our learned backward kernels lead to improved performance of both ULA and UHA, with the difference being more distinct for the latter. We also note that UHA has significantly larger error bars than other methods, but UHA-MCD does not appear to inherit this. We further provide a runtime comparison in Table 3, both for a single training iteration, and total experiment time (including evaluation). We find that fitting the score models of our method results in roughly twice as much runtime (recall fixed the number of sampler steps).
Limitations
While we have demonstrated that our optimized backward kernels can lead large improvements over standard AIS backward kernels, our proposed approach has a number of limitations. First, since our method relies upon unadjusted Langevin Hamiltonian sampling, we inherit many of the issues with these approaches. Using unadjusted samplers enables fully-differentiable ELBO estimate which can in theory be used to tune the many parameters of the sampler. However, these samplers require repeated gradient steps as an inner-loop which can lead to divergence dynamics and numerical instability, making optimization difficult. Next, our proposed training procedure has notable increase in memory consumption as we store the entire sampling trajectory in memory to train our neural network. We further note that possible to compute our importance weights online using O(1) memory. We could utilize this with Monte Carlo sub- sampling of the timestep (as with SGM (24)) to derive a constant memory training procedure for our approximate time-reversal at the cost of increased variance but we leave exploring this to future work. Finally, we observed that, without mitigation, parameterized forward samplers Qo sometimes dropped modes during training. Since this is a known challenge of using reverse KL, this might be improved by using alternative optimization objectives [35, 36].
Discussion
In this work have explored AIS using unadjusted Langevin and Hamiltonian dynamics. We have demonstrated that the backward transition kernels typically used are suboptimal and we have backward kernels. We have further shown
presented the form of the optimal how an approximation to these kernels can be learned using score matching and that this objective corresponds maximizing the ELBO in the limit of infinitesimally small dime-discretization step- sizes. We have illustrated the benefit of using our proposed optimized backward kernels on a number of inference problems including fixed targets and amortized tasks with model learning.

---

References
[1] Agakov, F. V. and Barber, D. (2004). An auxiliary variational method. In Advances in Neural Information Processing Systems.
[2] Arbel, M., Matthews, A. G., and Doucet, A. (2021). Annealed flow transport Monte Carlo. In International Conference on Machine Learning.
[3] Ba, J. L., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. arXiv preprint urXiv:1607.0666
[4] Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Fantacei, C., Godwin, J., Jones, C., Hennigan, T., Hessel, M., Kapturowski, S., Keck, T., Kemaev, I., King, M., Martens, L., Mikulik, V., Norman, T., Quan, J., Papamakarios, G., Ring, R., Ruiz, F., Sanchez, A., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stokowiee, W., and Viola, F. (2020). The DeepMind JAX Ecosystem.
[5] Bernton, E., Heng. J., Doucet, A., and Jacob, P. E. (2019). Schrödinger bridge samplers, arXiv preprint 1rXiv:19/2.13170.
[6] Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman- Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.
[7] Brekelmans, R., Huang. S., Ghassemi, M., Ver Steeg, G., Grosse, R. B., and Makhzani, A. (2022). Improving mutual information estimation with annealed and energy-basee -based bounds. In International Conference on Learning Representations.
[8] Burda, Y., Grosse, R. B., and Salakhutdinov, R. (2016). Importance weighted autoencoders. In International Conference on Learning Representations.
[9] Crooks, G. E. (1998). Nonequilibrium measurements of free energy differences for microscopi- cally reversible Markovian systems. Journal of Statistical Physics, 144--)):4411-1487
[10] Dai, C.,, Heng, J., Jacob, P. E., and Whiteley, (2020). An invitation to sequential Monte Carlo samplers. arXiv preprint arXiv:2007.11953
[11] De Bortoli, V., Doucet, A., Heng, J., and Thornton, J. (2021a). Simulating diffusion bridges with score matching. arXiv preprint 1
[12] De Bortoli, V., Thornton, J., Heng. J., and Doucet, A. (2021b). Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems.
[13] Del Moral, P., Doucet, A., and Jasra, A. (2006). Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series B, 68(3):411 1--466.
[14] Dinh, L., Krueger, D., and Bengio, Y. (2014). Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516.
[15] Dockhorn, T., Vahdat, A., and Kreis, K. (2022). Score-based generative modeling with critically- damped Langevin diffusion. In International Conference on Learning Representations.
[16] Doucet, A., Grathwohl, W. S., Matthews, A. G. D. G., and Strathmann, H. (2022). Annealed importance sampling meets score matching. In ICLR Workshop on Deep Generative Models for Highly Structured Data.
[17] Fournier, N. and Tardif, C. (2021). On the simulated annealing in Rd. Journal of Functional Analysis, 28 281(5):009086.
[18] Geffner, T. and Domke, J. (2021). MCMC variational inference via uncorrected Hamiltonian annealing. In Advances Neural Information Processing Systems.
[19] Gelman, A. and Meng, X.-I.. (1998). Simulating normalizing constants: From importance sampling to bridge sampling to path sampling. Statistical Science, 183
10

---

[20] Grillo, G. (1994). Logarithmic Sobolev inequalities and Langevin algorithms in R". Stochastic Analysis and Applications, 12(3):30--328.
[21] Grosse, R. B., Maddison, C. J., and Salakhutdinov, R. R. (2013). Annealing between distribu- tions by averaging moments. In Advances in Neural Information Processing Systems.
[22] Haussmann, U. G. and Pardoux, E. (1986), Time reversal of diffusions. The Annals of Probability, 14(3):1188 188-12088
[23] Heng. J., Bishop, A. N., Deligiannidis, G., and Doucet, A. (2020), Controlled sequential Monte Carlo. The Annals of Statistics, 4865):2944-
[24] Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances Neural Information Processing Systems.
[25] Horowitz, A. M. (1991). A generalized guided Monte Carlo algorithm. Physics Letters B, 268(2)247-255
[26] Huang, C.-W., Tan, S., Lacoste, A., and Courville, A. C. (2018). Improving explorability in variational inference with annealed variational objectives. In Advances in Neural Information Processing Systems.
[27] Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. The Journal of Machine Learning Research, 6:695-709.
[28] Jankowiak, M. and Phan, D. (2022). Surrogate likelihoods for variational annealed importance sampling. In International Conference on Machine Learning.
[29] Jarzynski, C. (1997). Nonequilibrium equality for free energy differences. Physical Review Letters, 78614)32600-2963
[30] Kingma, D. P. and Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint 0rXiv:1412.6900
[31] Kingma, D. P. and Welling, M. (2013). Auto-encoding variational Bayes. arXiv preprint 112.
[32] Klebaner, F. C. (2012). Introduction to Stochastic Calculus with Applications. Imperial College Press.
[33] Leimkuhler, B. and Matthews, C. (2016). Molecular Dynamics with Deterministic and Stochas- tic Numerical Methods. Springer.
[34] Masrani, V., Brekelmans, R., Bui, T., Nielsen, F., Galstyan, A., Ver Steeg, G., and Wood, F. (2021). q-paths: Generalizing the geometric annealing path using power means, In Uncertainty in Artificial Intelligence.
[35] Masrani, V., Le, T. A., and Wood, F. (2019). The thermodynamic variational objective. In Advances in Neural Information Processing Systems.
[36] Matthews, A. G. D. G., Arbel, M., Rezende, J., and Doucet, A. (2022). Continual repeated annealed flow transport Monte Carlo. In International Conference on Machine Learning.
[37] Neal, R. M. (2001). Annealed importance sampling. Statistics and Computing, 1 11(2):12--139
[38] Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2(11):2.
[39] Opper, M. (2019). Variational inference for stochastic differential equations. Annalen der Physik, 531:1800233.
[40] Ramachandran, P., Zoph, B., and Le, Q. V. (2017). Searching for activation functions. arXiv preprint 1rXiv:1710.05941.

---

[41] Rezende, D. J., Mohamed, S., and Wierstra, D. (2014). Stochastic backpropagation and approximate inference deep generative models. In International Conference on Machine Learning.
[42] Roberts, G. O. and Tweedie, R. L. (1996). Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, pages 341-363.
[43] Salakhutdinov, R. and Murray, I. (2008). On the quantitative analysis of deep belief networks. In International Conference Machine Learning.
[44] Salimans, T., Kingma, D., and Welling, M. (2015). Markov chain Monte Carlo and variational inference: Bridging the gap. In International Conference on Machine Learning.
[45] Schmiedl, T. and Seifert, U. (2007). Optimal finite-timg processes in stochastic thermodynamies Physical Review Letters, 98110)110830.
[46] Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsuper- vised learning using nonequilibrium thermodynamics, In International Conference Machine Learning.
[47] Song, Y., Durkan, C., Murray, I., and Ermon, S. (2021a). Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems.
[48] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021b). Score- based generative modeling through stochastic differential equations. In International Conference on Learning Representations.
[49] Sottinen, and Särkkä, (2008). Application of Girsanov theorem particle filtering of discretely observed continuous-time non-linear systems. Bayesian Analysis, 3(3):555-555 584.
[50] Tang, W. and Zhou, X. Y. (2021). Simulated annealing from continuum to discretization: convergence analysis via the Eyring-Kramers law. arXiv preprint (rXiv:2102.0233)
[51] Thin, A., Kotelevskii, N., Durmus, A., Panov, M., Moulines, E., and Doucet, A. (2021). Monte Carlo variational auto-encoders. In International Conference on Machine Learning.
[52] Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation, 2367):1661-1674
[53] Wu, H., Köhler, J., and Noé, F. (2020). Stochastic normalizing flows. In Advances in Neural Information Processing Systems.
[54] Zhang, G., Hsu, K., Li, J., Finn, C., and Grosse, R. (2021). Differentiable annealed importance sampling and the perils of gradient noise. In Advances in Neural Information Processing Systems.
Checklist
1. 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 5
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Paper is methodological paper approximate inference methods with only standard datasets. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See proofs of main results in Appendix.
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix.
3. If you experiments..
12

---

(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or URL)? [Yes] See Appendix for 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 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., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. 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 the supplemental material or as 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]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "sB6CykdwO-b", "reviewer": "Reviewer_aT45", "summary": "The key problem this work addresses is the suboptimality of the extended target distribution that is widely used in current annealed importance sampling (AIS) methods. The authors, building on previous literature, show that the custom choice of the backward Markov kernel of the unnormalized extended target is comfortable but suboptimal in terms of variance.\nTo overcome such limitations, the authors extend some key observations that have been made in the literature of diffusion processes used for generative modeling (e.g. Sohl-Dickstein et al., 2015), and draw a connection between forward/backward diffusion and forward/backward transition kernels used in importance sampling. Armed with this realization, the authors show that AIS can be studied through the lenses of forward and backward diffusion SDEs: the forward diffusion begins in a simple to sample from distribution, and gradually transforms it into the target distribution (or an extension thereof), while the backward diffusion corresponds to the (hypothetical) process of transforming the target to a simple distribution.\nIn addition, and as an extension to a previous paper that appeared in a recent workshop, the authors consider Hamiltoinan dynamics, and derive a similar connection to AIS, but with improved performance due to a higher fidelity representation of the paths taken by the Hamiltonian SDEs, when compared to first order dynamics.\nFor practical reasons, the authors produce discretized versions of the dynamics they study, and implement numerical integration schemes for the simulation of the SDEs, where necessary.\nNumerical results on both static distributions and on an instance of a variational autoencoder indicate dramatic performance improvements of the proposed approach when compared to vanilla and Hamiltonian vanilla versions of AIS.", "questions": "Here are some questions, but these are somehow superficial as I didn't find any major flaws in the paper.\n\n1- I think there is a typo in eq 12: the term $ 2 \\nabla \\log s_{\\theta} ( \\cdot ) $ should read $ 2 s_{\\theta} ( \\cdot ) $\n\n2- I may have made mistakes in my own derivations, but isn't there a typo in eq 6 for the backward transition kernel? In the numerator you should have $q_{k+1}(x_k)$ and not $q_{k}(x_k)$? I'm not 100% sure here\n\n3- May I suggest to improve the notation? I find it difficult to follow when $x, x'$ are used in place of $x_{k+1}, x_k$\n\n4- Would it make sense to have an additional section in the appendix to clarify not only the similarities but also the differences w.r.t. score-based diffusion models? See my comment on the weaknesses above.\n\n== Post Rebuttal == \nThanks for your answers to the questions and comments above. I have modified my score, raising it to an accept/7. ", "limitations": "I think they did address limitations appropriately.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* This is a very important topic, and the performance gains from the proposed method are substantial\n* I liked the idea of connecting the dots between AIS and (score-based) diffusion models\n* The extension to Hamiltoinan dynamics is interesting, and a nice complement to [1]\n\nWeaknesses:\n* I think the exposition of this paper should be improved: the connection with diffusion models is somehow reversed, as in that literature the forward process is going from the target distribution to a simple one, and the backward process is doing the opposite, whereas in this paper it is the other way around. This could be discussed in the appendix. Also, I would study the \"steady state\" distribution of the forward process defined in eq 10, by means of Fokker-Plank equations, to show that indeed the forward process moves an \"easy\" distribution to the target one. This can be helpful for a reader familiar with diffusion models, whereby you can demonstrate that for variance preserving methods, the process converges to isotropic Gaussian.\nTo put things differently: a reader with familiarity in diffusion models knows that the forward process \"destroys\" data, whereas in this work it is the other way around.\n* experiments only compare the proposed method to baseline AIS and to Hamiltonian AIS (which I appreciated a lot). However, this work could be cast under the umbrella of energy based models and density estimation. Then, would it make sense to compare it to such line of work?\n* there exist a workshop paper presenting preliminary ideas that are extended in this work. The main difference here is the treatment of Hamiltonian dynamics, and the experiments with the variational autoencoder. If accepted, this paper could maybe give more space to the Hamiltonian part, and most importantly, to the excellent results obtained. This is the main reason why, in its current form, the key contribution of this paper is somehow diluted, and seem a repetition: I like this work, and I would be happy to raise the score I gave which is only tainted by this discussed weakness\n\n[1] Tim Dockhorn and Arash Vahdat and Karsten Kreis, Score-Based Generative Modeling with Critically-Damped Langevin Diffusion, ICLR 2022\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "igrOTHHpRGR", "reviewer": "Reviewer_xpLj", "summary": "Annealed importance sampling is one of the most effective methods for marginal likelihood estimation. The commonly used extended target distribution is simple but sub-optimal. The paper leverage the recent progress in score-based generative modeling to approximate the optimal extended target distribution for annealed importance sampling and give the discretization of Langevin and Hamiltonian dynamics. On a number of synthetic distributions and variational auto-encoders, the proposed method obtains better performance compared to existing methods.", "questions": "* Considering unadjusted algorithms are used, will the error in score estimation brings bias in the sampling? If so, is it able to quantify the bias in terms of the estimation error?\n\n* What is the training time of the score model? How is it compared to the sampling time?", "limitations": "The paper mentions some limitations of the proposed methods. I think it misses two important problems:\n1. The method requires learning a score model, which is time consuming and hard. This factor reduces the efficiency or even correctness of the proposed method.\n2. The estimation error in the score model may lead to bias in sampling, and the paper lacks a quantification of the potential bias.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "* The writing is of high quality. The paper gives a concise review of the annealed importance sampling and illustrates the limitations clearly via two examples. Then, the paper provides necessary details in deriving the Langevin and Hamiltonian algorithms and makes them easy to read.\n\n* As a sampling algorithm, the effectiveness of the proposed method needs further consideration. In particular, the effectiveness of the proposed method relies on the qualities of the learned score function. However, learning the score models is not easier than estimating the partition functions. Whether the proposed method can apply well to hard distributions is still doubtful. It will be helpful if the paper can provide an ablation study that investigates the relationship between the quality of the learned score functions and the MCD samplers.\n\n* Although I have some reservations about the effectiveness of the current algorithm due to the reason mentioned above, its contribution to further studies on extended target distributions should be recognized. For example, when the energy function is learned via a score-based model and is associated with a score function in nature, the proposed method could be very helpful.\n\n* The evaluation in the experiment is not fair. We can run ULA and UHA directly, but we need to learn the score functions to run ULA-MCD and UHA-MCD. The paper should also report the learning time and the sampling time such that the readers can have better judgments.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GeLoADmowWw", "reviewer": "Reviewer_KR2v", "summary": "This paper proposes an improved Annealed Importance Sampling (AIS) for score-based diffusion models. The author argues that past papers on AIS have predominantly focused on modifying the intermediate distributions or the Markov kernel but have overlooked the suboptimal nature of the extended target distribution. Thus, the paper integrates a new AIS procedure in combination with recent works to approximate the extended target distribution.", "questions": "- Were there any experiments done on ULA-MCD and UHA-MCD on more complex distributions? I understand estimating the normalizing constant is very difficult, but a single experiment would be nice to see.\n- Why do we need applications for amortized inference? Also, why not experiment on more datasets?\n- I am not very familiar with AIS, are differentiable AIS and ULA and UHA the current SOTA benchmarks for likelihood estimation?\n- How would you provide a solution for the numerical instability in optimizing due to the unadjusted Langevin and Hamiltonian sampler? This seems like an imperative issue that needs to be addressed.", "limitations": "The authors transparently addresses the limitations of their work and have a section dedicated to limitations. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well motivated and theoretically proves the correctness of the proposed methods.\n- Overall, the paper is well written and explained useful insight on the limitations on the Markov kernel:\n  - E.g. the illustration of the suboptimality of a homogeneous MCMC chain give good intuition on the current problem of AIS.\n\nWeaknesses:\n- While experimental results show improvement, it would be more convincing if there were more real-world dataset experiments (other than binarized MNIST).\n\nMinor Suggestions:\n- Why not bold the best log Z estimate on Table 1 & 2?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept