Learning Optimal Flows for Non-Equilibrium Importance Sampling
Authors: Yu Cao
Eric Vanden-Eijnden
Venue: NeurIPS 2022 Accept (2022)
Keywords: non-equilibrium importance sampling
machine learning

Abstract: Many applications in computational sciences and statistical inference require the computation of expectations with respect to complex high-dimensional distributions with unknown normalization constants, as well as the estimation of these constants. Here we develop a method to perform these calculations based on generating samples from a simple base distribution, transporting them by the flow generated by a velocity field, and performing averages along these flowlines. This non-equilibrium importance sampling (NEIS) strategy is straightforward to implement and can be used for calculations with arbitrary target distributions. On the theory side, we discuss how to tailor the velocity field to the target and establish general conditions under which the proposed estimator is a perfect estimator with zero-variance. We also draw connections between NEIS and approaches based on mapping a base distribution onto a target via a transport map. On the computational side, we show how to use deep learning to represent the velocity field by a neural network and train it towards the zero variance optimum. These results are illustrated numerically on benchmark examples (with dimension up to $10$), where after training the velocity field, the variance of the NEIS estimator is reduced by up to $6$ orders of magnitude than that of a vanilla estimator. We also compare the performances of NEIS with those of Neal's annealed importance sampling (AIS).

Learning Optimal Flows for Non-Equilibrium Importance Sampling
Yu Cao
Courant Institute of Mathematical Sciences, New York University
yucaoycloutlook.com
Eric Vanden-Eijnden
Courant Institute of Mathematical Sciences, New York University eve20cims nyu. edu
Abstract
Many applications in computational sciences and statistical inference require the computation of expectations with respect to complex oigh-dimensional distributions with unknown normalization constants, as well as the estimation these constants. Here we develop method to perform these calculations based on generating samples from a simple base distribution, transporting them by the flow generated by velocity field, and performing averages along these flowlines. This non- equilibrium importance sampling (NEIS) strategy straightforward to implement and can be used for calculations with arbitrary target distributions. On the theory side, we discuss how to tailor the velocity field the target and establish general conditions under which the proposed estimator is a perfect estimator with zero- variance. We also draw connections between NEIS and approaches based on mapping base distribution onto target via transport map. On the computational side, we show how to use deep learning to represent the velocity field by a neural network and train towards the zero variance optimum. These results illustrated numerically on benchmark examples (with dimension up to 10), where after training the velocity field, the variance of the NEIS estimator is reduced by up to 6 orders of magnitude than that vanilla estimator. We also compare the performances of NEIS with those of Neal's annealed importance sampling (AIS).
Introduction
Given a potential function U1 Ω + R on the domain Ω C Rd, the main goal of this paper is to evaluate
Z / (x) dr.
The calculations of such integrals arise in many applications from several scientific fields. For instance, Zi is known as the partition function in statistical physics (2I), where it is used characterize the thermodynamic properties of system with energy U,, and as the evidence in Bayesian statistics, where it used for model selection [2]
When the dimension cable to ( and the method of choice estimate Z is Monte-Carhe sampling 222 8 This requires expressing Z, as an expectation, which can be done, e.g., by realizing that
the domain Q is large, standard numerical quadrature methods are inappli-
Z Eole U1 /po],
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

where Ey denotes the expectation with respect to the probability density function Po > 0. If pu is both known (i.e., we can evaluate pointwise in Q2 normalization factor included) and simple to sample from, we can build an estimator for 21 by replacing the expectation on the right hand side of (2) by the empirical average of e /Po over samples drawn from Po. Unfortunately, finding a density Po that has the two properties above is hard: unless Po is well-adapted to e e-Vi, the estimator based on (2) is terrible in general, with a standard deviation that is typically much larger than its mean or even infinite. A similar issue arises if we want to estimate the expectation Eoo of some function f: Q + R, and the two problems are in fact connected when f > 0, since the second reduces to for U1 lo)
These difficulties have prompted the development of importance sampling strategies 14 whose aim is to produce estimators with reasonably low variance for 2 or Eof. These include for example umbrella sampling
replica exchange (aka parallel tempering) [4 23 nested sampling [37
38) in which the estimation of 2 is factorized into the calculation of several expectations of the
type (2) but with better properties, that can then be recombined using thermodynamic integration [7
or Bennett acceptance ratio method [6].
Complementary to these equilibrium techniques, non-equilibrium sampling strategies have also been introduced for the calculation of (Tp. For example, Neal's annealed importance sampling (AIS) [26] based on the Jarzynski equality [[S] I6 calculates 21 using properly weighted averages over sequences of states evolving from samples from Pu. without requiring that the kernel used to generate these states be in belailed-balance with respect to either Po. or P1 == e U1 /1, or any density interpolating between these two. Instead, the weight factors are based on the probability distribution of the sequence of states in the path space. Other non-equilibrium sampling strategies in this vein
include bridge and path sampling [13,, and sequential Monte Carlo (SMC) sampling [24 2).
In this paper, we analyze another non-equilibrium importance sampling (NEIS) method, originally introduced in [33]. NEIS is based on generating samples from a simple base density po. then propagating them forward and backward in time along the flowlines of velocity field, and computing averages along these trajectories- the basic idea of the method is to use the flow induced by this velocity field to sweep samples from Pa through regions in S that contribute most to the expectation. As shown in [33] and recalled below, this procedure leads to consistent estimators for the calculation of 2, or f via generalization of ( One advantage of the method, which a rare feature among importance sampling strategies. is that it leads to estimators that always have lower variance than the vanilla estimator based on (2) (3 The question we investigate in this paper is how low their variance can be made, both in theory and in practice. Our main contributions are:
Under mild assumptions on U1 and Pu, we show that the NEIS velocity field is the gradient of a potential that satisfies Poisson equation, the NEIS estimator for Z₁ has zero variance.
Under the same assumptions, we show that this optimal flow can be used to construct a perfect transport map from Po to Pl. This allows us to compare NEIS with importance sampling strategies involving transport maps like normalizing flows (NF) that have recently gained popularity [32 19 30, and highlight some potential advantages of the former over the latter.
On the practical side, derive variational problems for the optimal velocity field in NEIS, and show how to solve these problems by approximating the velocity by a neural network and optimizing its parameters using deep learning training strategies, similar to what is done with neural ODE 9
We illustrate the feasibility and usefulness of this approach by testing it on numerical examples. First we consider Gaussian mixtures in up to 10 dimensions. In this context, we show that training the velocity used in NEIS allows to reduce the variance of a vanilla estimator using standard Gaussian distribution as on by up to orders of magnitude. Second we study Neal's 10-dimensional funnel distribution [27 2) for which the variance of the vanilla importance sampling method is infinity; training a linear dynamics with 2 parameters in NEIS can lead to an estimator with more accurate estimate of Z,. In these examples we also show that after training, NEIS leads to estimators with lower variance than AIS (26).
Related works. The idea of transporting samples from Po to lower the variance of the vanilla estimator based on (2) is also at the core of importance sampling strategies using normalizing flows (NF) 40
29 25 The type of transport used in NF-based method is however different nature from the one used NEIS. With NE, one tries to construct a map that transforms
39
32
19 30
45
each sample from Po into a sample from the target P1 = e U, U1/Z /Z, In contrast, NEIS uses samples
2

---

from Po as initial conditions to generate trajectories, and uses the data along these entire trajectories to build an estimator. Intuitively, this means that samples likely on Po must become likely on p sometime along these trajectories rather than at given time specified beforehand, which is easier to enforce.
NEIS bears similarities with Neal's AIS [26]. except that in NEIS the sampling is done once from Po generate deterministic trajectories to gather data for the estimator, whereas AIS uses random trajectories. There are some methods based on AIS that optimize the transition kemel: for instance, stochastic normalizing flows (SNF) proposed in [66] incorporate NF between annealing steps; and annealed flow transport (AFT) in (2 combines NF with the sequential Monte Carlo method to provide optimized flow transport. These approaches require learning several maps along the annealed transition, whereas the NEIS herein only needs to learn a single flow dynamics.
A time-discrete version of NEIS, termed NEO, was proposed in 422. The current implementation of NEO iterates on a map that needs to be prescribed beforehand, but this map could perhaps be optimized using strategy similar to the one proposed here.
From a practical standpoint, the idea of optimizing the velocity field in NEIS using a neural network approximation for this field can be viewed as application of neural ODEs [9] that uses the variance of the NEIS estimator as the objective function to minimize. The nature of this objective poses specific challenges in the training procedure, which investigate here.
Notations. For symmetry, we denote po( co((2) with Uo = -og Pu Ω + Zo S e (Uo(()) dr = 1. We denote d-dimensional vector filled with zeros as O₂ and the d x d identity matrix as La. ( is the Euclidean inner product in Rd. We assume that the domain Q is either an open and connected subset of Rd with smooth boundary or d-dimensional torus (without boundary). We denote by ./ [K,, Σ) the multivariate Gaussian density with mean μ and covariance matrix Σ. For two functions f.g D -> I where D is domain of interest, the notation f S g means that there exists a constant >>> 0 such that ((x) S Cg(x) for any x E D. Suppose T Ω - Ted is map and is a distribution, then the pushforward distribution of p by the map T is denoted as ##p. The notation |-| is the usual &g norm for vectors and ||-I| is the matrix norm or functional norm.
IR and
Flow-based NEIS method
Here we recall the main ingredients of the non-equilibrium importance sampling (NEIS) method proposed in 333 Let Ω - TR. velocity field which we assume belongs to the vector space
B = {b be C (I,
b nlan = 0, sup /bb(x) < oo).
(3)
se
where n is the normal vector at the boundary 0Ω. Define the associated flow map X; Ω - Ω via
d =
Xo(x)=
di
and let [xr)) be the Jacobian of this map:
J((x) := |det (V (x)( = exp (% (f'v.((x.x)) ds) ds)
(5)
Finally, let denote
(x) (ix J((x)
(6)
for k E {0, 10.1), x E Ω and & E R. NEIS is based on the following result, proven in Appendix B.
Proposition 2.1. b E B, then for any -0 < t- <t4 < 0o, we have
Z = EgAr
where
A Ae_ (2) = 1 F()((r) (x) dt.
(8)
fi-t- F! (x)
11_
3

---

1.0
0..2
0.8
0.01
0
0.01
0.2
0.02
0..3
0.0
0.2
0.4
0.6
0.8
1.0
Figure 1: Contour plot of and flowlines of b VV, where V solves the Poisson's equation [I]
with D = 1, assuming that = and is mixture density with modes; see (4)). With this VV. we have A(x) 21 for almost all [0, 1]3.
In addition, if
J Fi (x) dt
lim
Ae_ A(x)
S& F (z) dt
t -000
exists for almost all Po. then
21 = EsA.
(10)
When b = Og, or t_ 0 and 4 + 0, reduces to (2) The aim, however, is to choose so that the
estimator based on 7 has lower variance than the one based (2): we will show below that this can
indeed be done. For now, note that Jensen's 's inequality implies that an estimator based on (10) for
any b has lower variance than the one based 2 see [33 Proposition H.3 below for details.
Note also that, if one allows the magnitude of the flow b to be arbitrarily large, the finite-time NEIS (8) will behave like the infinite-time NEIS such relation will be discussed and elaborated in Appendix B.4.
Finally, note that the estimator 100) based on invariant with respect the parameterization of the flowlines generated by the dynamics b, as shown by the following result proved in Appendix B.2
Proposition 2.2 (An invariance property). Suppose ab B, where 8 (St, R) satisfies inf (xx)) > O. Then the fields b and ab generate the same flowlines, and A Aa where A
and Abb are the function defined in (9) using b and ob, respectively.
3 Optimal NEIS
The NEIS estimator for 10 is unbiased no matter what b is. However, its performance relies on the choice of b. Therefore, natural question is to find the field b that achieves the largest variance reduction. The next result shows that an optimal exists that leads to a zero-variance estimator:
Proposition 3.1 (Existence of zero-variance dynamics). Assume that Ω [0, 1] is torus and Un. U, E (S, RR. Let D : Ω > (0, 20) be some smooth positive function with infzen D(x)) > 0,
and suppose that V C (Ω) solves the following Poisson's 's equation Ω
[ v(x) =
v (DVV) = po+
with
(11)
If the solution V is Morse function, then b VV is a zero-variance velocity field: that is, if we use define (9) we have
[ (x) / Pi (x)
for almost all x. ~ Po>
(12)

---

and as result
A(x) 2 for almost all Pa-
(13)
This proposition is proven in Appendix E where we also make connection between 11 and Beckmann's 's transportation problem. We stress that the optimal b specified in Proposition 3.1 is not unique (see Proposition E.9): however, we show below in Proposition [4. that, under certain conditions, all local minima of the variance (viewed a functional of b) are global minima. We also note that the assumption that the solution is a Morse function is mostly a technicality, as discussed in Appendix [E.3] Similarly, we consider the torus in Proposition 3.1 for simplicity mainly; we expect that the proposition will hold in general when Ω has compact closure or even when s2 R see Appendix E for examples including that of Gaussian mixture distributions.
For illustration, the contour plot of V and the flowlines of b = VV are shown in Figure ! in simple example two-dimensional torus where pa and #1 a mixture density with modes; their explicit expressions are given in (54); in this example, we solved (II) numerically with D = 1, see Appendix E..5 for more details. Some other examples where the zero-variance dynamies is explicit are discussed Appendix F
Connection to transport maps and normalizing flows. transport map T from Po to Pt, as shown in:
The zero-variance dynamics provides a
Proposition 3.2 (Existence of a perfect generator). Suppose D = 1 for simplicity. Under the same assumption as in Proposition 3.7] lel V be the Morse function solving (TT and b VV the associated zero-variance dynamics. Then there exists a continuously differentiable function x (defined almost everywhere on (2) such that
do 1 (x(())
Po (X, (x)x))))(())))
PI (
(14)
Furthermore, the map T(x): = X (x) is a transport map from Pa to Pl. i.e. ## P1.
The proof is given in Appendix G.1 Note that we consider again b VV on the torus for technical simplicity: the statement of the proposition should hold in general for vero-variance dynamics The solution of (14) is particularly simple in one-dimension, where we can take b(x) = 1, and
straightforwardyy verify that x(x) ((r) with
where R((x) = S p:(y) dy, i == 0,
FF
(15)
We also illustrate the statement of Proposition 3.2 via numerical examples in Appendix G.2
To avoid confusion, we stress that we will not use the transport map T of Proposition 3.2 in the examples below. Indeed, using this map would require identifying x. which introduces an unnecessary additional calculation which we can avoid using the NEIS estimator directly. In addition, the NEIS estimator will likely have better properties than those based on transport maps, as we can think of NEIS as using time-paun parameterized family of transport maps rather than single one. In particular, the variance of the NEIS estimator will be small if samples likely on Po become likely on P1 sometime along the NEIS trajectories, rather than at the same fixed time for all samples. The former seems easier to fulfill than the latter. For example, in one-dimension, the NEIS estimator has zero variance for any b bounded away from zero, whereas building transport map from Pu to P1 is already nontrivial in that simple case since it requires solving [5
4 Variational formulations
The Poisson equation (1I) admits a variational formulation:
min / JIVV/D+ - -P).
(16)
If 'D is chosen to be probability density function (for example D Po D = (p + po)), the
two terms in the objective in (T6) are expectations which can be estimated via sampling (using, e.g.. direct sampling for the expectation with respect to Po and NEIS for the one with respect to P1). This
5

---

means that we can in principle use an MCMC estimator of (16) as empirical loss, and minimize it over all V in some parametric class, Here however, we will follow a different strategy that allows us to directly parametrize b instead of V (i.e., relax the requirement that b = VV) and simply use the variance of the estimator as objective function.
Specifically, since we quantify the performance of the estimators based on 7 and (10) by their variance, we can view these quantities as functionals of b that we wish minimize. Since the estimators are unbiased, these objectives are
Var, .14 (b) = MM_ + (b) - -2,, (finite-time);
(17)
Var(b) .M(b) - -zi:
(infinite-tmme)
where we defined the second moments M,_,t (b) := E. [A. and Eo [[43]] With the finite-time objective, we know that with b Od, 7 reduces to (2) Therefore, minimizing Mr_.t. (b) over b by gradient descent starting from b near 0g will necessarily produce a better estimator: while we cannot guarantee that the variance of this optimized estimator will be zero, the experiments conducted below indicate that it can be a several order of magnitude below that of the vanilla estimator.
For the infinite-time objective, know that for any b, ( leads to an estimator with lower variance than the one based on (2 [33] Minimizing Var(b) over b using gradient descent leads to a local minimum: the next result shows that all such local minima are global:
Proposition 4.1 (Global minimum). Under some technical assumptions listed in Proposition H.]
i. E B is a local minimum of Var(.) where the functional derivative of Var(b) with respect to vanishes, ie., (Var(b.)/bb Oa on S2 then b. is global minimum and Var(b.) 0.
The expression of the functional derivative d'Var(b (b.) )6b is given in Proposition D.1 assumptions under which Proposition 4.1 holds are explained in Appendix A and the proof given in Appendix
The technical
5 Training towards the optimal b
Here we discuss how to use deep learning techniques to find the optimal b; these techniques will be illustrated on numerical examples in Section 6 Some technical details are deferred to Appendix
= 1- +. Two with L- E 4- natural choices are (_ = 0 and (_ 1122, which will be used below in the numerical experiments. This leads to no loss of generality a priori since in the training scheme we put no restriction on the magnitude that b can reach, and with large b the flow line can travel a large distance even during t E [-1,1] (the range of integration in s, in (8); see the discussion in Appendix B.4 for more details. In practice, use the standard Runge- Kutta scheme of order (RK4) to integrate the ODE (4) over E [-1, Kunge-Kuttta using uniform time step (At 1/N,). We note that this numerical discretization introduces bias. However, this bias can be systematically controlled by changing the time step or using higher order integrators. In our experiments, we observed that the RK4 integrator led negligible errors, see Table 5
Objective. We use the finite-time objective Mr .t (b)
dine-discretized version of (8) with 2N, discretization points, and
Neural architecture. In our experiments, we either parameterize b by neural network directly, or we assume that b is a gradient field,
b VV
(gradient form),
and parameterize the potential V by neural network. We always use an l-layer neural network with width m for all inner layers: therefore, from now on, we simply refer the neural network structure by a pair (f, m); see Appendix I.2 for more details. When we parametrize the potential function V instead of b, the only difference is that the output dimension of the neural network becomes 1 instead of d. The activation function is chosen as the softplus function (a smooth version of ReL.U) that gave better empirical results compared to the sigmoid function. At initialization the neural parameters were randomly generated. Theoretical results about the gradient of M, .t, (b) with respect to parameters given in Appendix C and corresponding numerical implementations explained in Appendix [

---

Direct training method. We minimize Mr_. (b) with respect to the parameters in the neural network using stochastic gradient descent (SGD) in which we evaluate the loss and its gradient empirically using mini-batches of data drawn from Po at every iteration step. For simplicity, we choose Po as the standard Gaussian in Section 6 below.
Assisted training method. When local minima of U1 are far away from the local minimum of Ug. the direct training method by sampling data from pa and minimizing Mi_. (b) fails, because the flowlines may not reach the importance region of f1 due to poor initializations of b. More specifically, if along almost all trajectories, e-Ui(X. (t)) R O for E [-1, 1], then with large probability Ar_. t.. (x) Re O where x' ~ Pos as result, the empirical variance of the estimator can be extremely small if the number of samples is small, while the true variance could be extremely large.
Such mode collapse phenomenon is common in rare event simulations.
To get around this difficulty, recall that ideally we would like to find a dynamics such that A is approximately constant function in the infinite-time time case. That is, if b is zero-variance dynamics, then A is a constant function and Ep [(A - (E}] = O for any distribution constrained to use the base distribution po and minimize the functional b + Var(b) by this idea, we use an assisted training scheme in which, iteration of SGD, the loss function is
and we are not
(b) in (7). Motivated
Ep. (AL_+ - Ep, AL..4)
(18)
Here Ep. denotes expectation with respect to the probability density p. defined as
P: = (1- p1 + #Po.
G = max {c - uL 0).
(19)
where E (0,1) controls the number of steps during which the training is assisted, c E (0, 1). is the total number of training steps, and := Z==1 is the time-1 map of the ODE (x(x) = -VV (z)) with (z= z and < > 0 is a parameter. In essence, using (18) means that, for the first J. training steps, there is a probability e; that the data x ~ Po are replaced by Z(x). that the training method can better explore important regions near local minima of U1. Subsequently, the assistance is turned off so that some subtle adjustment can be made. If some samples from were available beforehand, we could equivalently replace Z 4o in (19) by the
P1 = e e- empirical distribution over these samples. We emphasize that the assisted training method is only used to guide the training initially and the NEIS estimator for Z, is unbiased.
6 Numerical experiments
We consider three benchmark examples to illustrate the effectiveness of NEIS assisted with training. The first two examples involve Gaussian mixtures, for which we use NEIS with f- = 0; the third example is Neal's funnel distribution, for which we use NEIS with t_ -1/2. In all examples, we compare the performance of NEIS with those of annealed importance sampling (AIS) [26]: the number of transition steps in AIS is denoted as K and we refer to this method as AIS-K below; for more details see Appendix [.] For the comparison, we choose to record the query costs to Ui and VU, as the measurement computational complexity, which connects to the framework in the theory of information complexity (see e.g.. (28). The runtime could depend on coding, machine condition, etc., whereas query complexity more or less only depends on the computational problem (Uo. U1 and b) itself; for most examples of interest, Uo is simple whereas U₁ and its derivatives will be expensive to compute; as a remark, VU, is almost always more expensive to compute than U1.
For simplicity, we always use as base density po (z) = (2T)) -02e- ln1 We remark that a better choice of po (i.c., more adapted to #1) can significantly improve the sampling performance; our po is precisely used to validate the performance of NEIS in situations where po is not well chosen. would be interesting to study how to adapt the choice of Po for easier training in NEIS, but this is left for future investigations.
When presenting results, we rescale the estimates so that the exact value is Z = 1 for all examples. More implementation details about training are deferred to Appendix [.2] All trainings and estimates of 21 are conducted on a laptop with CPU 17-1270 we use 15 threads at maximum. The runtime of training ranges approximately between 45 ~ 76 seconds for Gaussian mixture (2D), 9.5w112 minutes for Gaussian mixture (10D); for the Funnel distribution (10D), the runtime is around 25 minutes
7

---

for generic linear ansatz and around 2 minutes for two wwo-parametric ansatz. Appendix inneludes additional figures about training. When computing the gradient of the variance with respect to parameters, we use an eengraation-assed method when L_ -1/2 (for the convenience of numerical implementation) and use an ODE-based method when f._ = 0 (for higher accuracy); details about these two approximation methods are given in Appendix 1.3 and Appendix I.4 respectively. The codes are accessible on htttps://githi wwt.giittu..ooomeeeeee woom/yuccoyc/N/...
iincluess
An asymmetric 2-mode Gaussian mixture in 2D. 2-mode Gaussian mixture
As first illustration, we consider an asymmetric
1 Us x 5 (xte,,, ofTa) (-Ne, ofTa).
4
(20)
where e = [b) e2 11 01 = vO.i, A = 5. With this choice of parameters, the variance of the vanilla estimator based on 2 is approximately 1.85 106, We use NEIS with f_ 0 and set the time step to At = 1/50 for ODE discretization during both training and estimation of Z₁. We train over L = 50 SGD steps using the loss (18) by imposing bias in the first 60% of the training period only (i.e., with e 0.6). The evolution of the variance during the training is shown in Figure| 7 in Appendix: the best optimized flow has variance of about 1, as opposed to 10 for the vanilla estimator. Since p; in (19) is quite different from Po during the assisted learning period, it may happen that the empirical variance significantly exceeds the variance of the vanilla importance sampling; this does not contradict with Proposition H.33 below. As seen in Figure 7 at the end of the assisted period, the variance is already quite small and most cases, the variance continues to reduce as b gets further optimized.
bebllo....
After training, we estimate 2 using NEIS with the optimized flow and compare its performance with AIS-10 and AIS-100. We first record the query cost for training and then set a total number of queries to U1, VU, as budgets. Given the query budget, we estimate Σ, using each method 10 times, leading to the results given in (Jelow. When we determine the estimation cost of NEIS, we deduct the query cost of training from the total query budget for fairer comparison. Note that NEIS uses less queries to produce more accurate estimate of Z1: in particular. the standard deviation of estimating Z1 by NEIS method is around I magnitude smaller than AIS-100. Moreover, the bias from ODE discretization appears to be negligible. Figure shows an optimized flow and also provides a visual comparison of NEIS with AIS under fixed query budget; more comparisons using various ansatzes or architectures can be found in Figures 11 and 12
shows
symmetric 4-mode Gaussian mixture in 10D. mixture in d = 10 dimension with
Next we consider a symmetric 4-mode Gaussian
e (.()
Σ N M E),
(21)
where the vector μ4 [AAcos ()) Asin( ( 00 u| and Σ = Diagl of of 03 o3 ... a3] is diagonal matrix. The parameters
d 10, 01 V0.1, 02 v0.5 and X 5. With this choice of parameters, the variance of the vanilla estimator based on 2 is approximately 2.15 x 108. We use NEIS with t_ 0 and the time step At = 1/60 is used for ODE discretization during both training and estimation of Z,. We show the training result in Figure 8 The variance reduces to about 10 after 60 SGD steps for the gradient ansatz (here we only considered this ansatz it produces more promising empirical results).
Similar to the last example, compare NEIS using the optimized b with AIS, under fixed query budgets; see T The best result from NEIS has an estimator with the standard deviation less than 1/3 of the one from AIS-100. This comparison suggests that AIS- needs more than 9 times more resources than NEIS with optimized flow in order to achieve similar precision and the cost spent on training indeed pays off if we require an accurate estimate of 2, (meaning less fluctuation for Monte Carlo estimates). Moreover, this table also shows that the bias from ODE discretization is negligible.
Figure 2 shows a particular optimized flow: can be seen, the mass near the origin flows towards four local minima of U,, as we would intuitively expect. More optimized flows and comparisons can be found in Figure 13 in Appendix.
8

---

VU::
VUu:
VUx: 0.0 B
1.25
trial
1.00
0.50
(a) Asymmetric Gaussian mixture (2D): gradient ansatz, 2, 20
Us: 48.6 MB
Us: 48.6 MB
Us: 72.9 MB
10
1.3
VU: 48.6 MB
VU: 48.6 MB
vU 0.0
1.1
tal
1.0
0.9
-5
0.8
AIS-10
A/S-100
NEIS
-10
10
(b) Symmetric Gaussian mixture (10D): gradient ansatz, = 2, me 30
Us: 178.3 MB
U:: 178.3 MB
Us: 267.5 MB
1.2
TU: 178.3 MB
VU:: 178.3 MB
VU: 0.0
5.0
2.5
1.0
tiil
0.0
55.0
AS:10
AI5-100
NEIS
-10
-6
2
(c) Funnel distribution (10D): generic linear ansatz
U.: 80.8 MB
VU:: 80.8 MB
VU:: 0.0 B
2.5
0.0
-5
0.6
Funnel distribution (10D): two-parametric
Figure 2: Selective comparison results for various models. Left panels: estimates of 2 by AIS and NEIS with optimized flow under the fixed query budget shown above the panels; we repeat these calculations 10 times and show boxplots of these 10 estimates for each method. The trial number refers to the index for randomly chosen initialization. The query numbers refer to the queries used for each estimate of Z, for each method; 1MB 106. The dashed red lines show the exact value Z 1. Right panels: streamlines of optimized flows atop the contours of U1, both projected into the .1-022 plane for the two 1OD examples. Full comparison and figures can be found in Figures 12 and 14 in Appendix.

---

A funnel distribution in 10D. We consider the following 10D funnel distribution studied in 2 27 for the state x [T1, T2. .2.............................. ........
210 E R.
21 (0,9),
I2 ~ (0,,
2 2sisd. d.
For numerical stability, we consider the above funnel distribution restricted to unit ball centered at the origin with radius 25. Instead of heuristically parameterizing the dynamics via networks, consider generic linear ansatz and wwopparametric linear ansatz:
b(r) W1z + b1.
W1 E I 1
BER
(22a)
b(x)) [13, a.2 aT3, .T,uu a.iu),
a
(22b)
The generic linear ansatz can be regarded
neural network without inner layers. With (2a).
drawn the entries in the matrix W1 randomly and we set bj = 10 initially, and we use the assisted training method; with (22b). we set = 8 = 2 initiallly, and we use the direct training method. In both cases, we choose the finite-NEIS scheme with l_ = -1/2. We notice that the asymmetric choice L_ 0 can also leads into more optimal dynamics, but its performance is not as competitive as the symmetric case t_ 1/2. It is very likely that such a difference is due to the structure of funnel distribution: each coordinate z; (1 S < 10) has mean 0 and therefore, a symmetric version can probably better weight the contribution from both forward and backward flowlines.
The training results are shown in Figures and 10 Appendix. In Figure TO
can observe that both error and variance are overall decreasing during the training and the parameters a, 3 tend to increase with similar speed. We use the same protocol in the two previous examples to compare NEIS with AIS. As can be observed in Table 5 the two-parametric ansatz (22b) gives the best estimate; the generic linear ansatz (22a) is not as competitive as the two-parametric ansatz (probably due to wver-parameterization)) but it still outperforms the AIS-100, under fixed query budget. Figure 2 shows these optimized flows; more results can be found in Figure 14jin Appendix. The apparent gap between estimates and the ground truth in Figure 2 (or see Table 5) comes from insufficient sample size.
7 Conclusion and outlook
In this work, we revisited the NEIS strategy proposed [33 and analyzed its capabilities, both from theoretical and computational standpoints. Regarding the former, we showed that NEIS leads to a zero-variance estimator for a velocity field b = VV with a potential V that satisfies a certain Poisson equation with the difference between the target and the base density source. Moreover, a zero-variance dynamics can be used to construct a transport map from Po to P1. In turn, we highlighted the connection and difference between NEIS and importance sampling strategies based on the normalizing flows (NF).
On the computational side, we showed that the variance of the NEIS estimator can be used as objective function to train the velocity field b. This training procedure can be performed in practice by approximating the velocity field by neural network, and optimizing the parameters in this network using SGD, similar to what done in the context of neural ODE but with a different objective. Our numerical experiments showed that this strategy is effective and can lower the variance of vanilla estimator for Z1 by several orders of magnitude.
While the numerical examples we used in the present paper are somewhat academic, the results suggest that NEIS has potential more realistic settings. In order to explore other applications, would be interesting to investigate how to best parametrize b (e.g.. less parameters and non-stiff energy landscape with respect to these parameters) and how to best initiate the training procedure. would also be interesting ask whether we can improve the performance of NEIS by optimizing certain parameters in the base density Po in concert with b. The answers to these questions are probably model specific and are left for future work.
Acknowledgment
We would like to thank Jonathan Weare and Fang-Hua Lin for helpful discussions, and the anonymous referees for their useful comments and suggestions. The work of EVE is supported by the National Science Foundation under awards DMRR-420077 DMS-2012510. and DMS-2134216 by the Simons Collaboration Wave Turbulence, Grant No. 617006, and by Vannevar Bush Faculty Fellowship.
10

---

References
[1] Christophe Andrieu, James Ridgway, and Nick Whiteley. Sampling normalizing constants in high dimensions using inhomogeneous diffusions, 2018. arXiv: 1612.07583.
[2] Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 Proceedings of Machine Learning Research, pages 318-330. PMLR, 2021.
[3] Larry Armijo. Minimization of functions having Lipschitz continuous first partial derivatives. Pac. J. Math., 16(1):1-3, 1966.
[4] Michèle Audin and Mihai Damian. Morse Theory and Floer Homology. Springer London, 2014. doi: 10. 10.1007/978- 10.10079978-14477-5566 1-4471-5496-9.
[5] Martin Beckmann. A continuous model of transportation. Econometrica, 20(4):643 660, 1952.
[6] Charles H Bennett. Efficient estimation of free energy differences from Monte Carlo data. J.
Compul. Phys., 22(2):245-268, 1976.
[7] O. Bräunling. Fourier series the n-dimensional torus, 2004. https://www wnw.uuniimate http:////caiictttaacc///i.ooooo pdf.
[8] Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Li Meng. Handbook of Markov Chain Monte Carlo. CRC Press LLC, 2011.
[9] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[10] Weinan E, Chao Ma, and Wu. The Barron space and the flow-induced function spaces for neu- ral network models. Constr. Approx., 55(1):369 406, 2022. doi: 10. 10.10077:00665-01--9599,,9
[11] Lawrence C. Evans. Partial Differential Equations. American Mathematical Society, 2010.
[12] F. Feroz and M. P. Hobson. Multimodal nested sampling: an efficient and robust alternative to Markov Chain Monte Carlo methods for astronomical data analyses. Mon. Not. Astron. Soc., 384(2):449-446 463, 2008. doi: 11..111/..165-296666.
[13] Andrew Gelman and Xiao-Li Meng. Simulating normalizing constants: from importance sampling to bridge sampling to path sampling. Statist. Sci., 13(2):163-185, 1998. doi: 10. 1214/ ss/1028905934.
Number:
[14] Charles J. Geyer. Importance sampling, simulated tempering, and umbrella sampling. In Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng, editors, Handbook of Markov Chain Monte Carlo. CRC press, 2011.
[15] Jarzyssk... Equilibrium free-energy differences from nonequilibrium measurements: A master- equation approach. Phys. Rev. E, 5655):0188855055 1997. doi: 10.1 1.110PhyyeReee.55.50001
[16] C. Jarzynski. Nonequilibrium equality for free energy differences. Phys. Rev. Lett., 78(14): 2690-28 26930, 1997. doi: 10. 10.1103PPhy3ReeeLett.....55
[17] John G Kirkwood. Statistical mechanics of fluid mixtures. Chem. Phys., 3(5):300-313, 1935.
[18] Achim Klenke. Probability Theory: A Comprehensive Course. 2014.
[19] Ivan Kobyzev, Simon J. D. Prince, and Marcus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Trans. Pattern Anal. Mach. Intell., pages 1-1, 2020. doi: 10.1 2..1009/PAMM..22002292222
[20] Peter Lax. Change of variables in multiple integrals II. Am. Math. Mon., 108(2):115-11) 2001. doi: 10.2307/269552..

---

[21] Evgenii Mikhailovich Lifshitz and Lev Petrovich Pitaevskii. Statistical Physics: Theory of the Condensed State, volume 9. Elsevier, 2013.
[22] Jun Liu, Monte Carlo Strategies in Scientific Computing, volume 10. Springer, 2001.
[23]
Jianfeng Lu and Eric Vanden- Vanden-Eijnden.. Methodological and computationnl aspects parallel tempering methods in the infinite swapping limit. J. Stat. Phys., 174(3): 2019. doi:
10.1007/s10955-008 010-555-01-220000 2-220--..
[24] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. J. R. Statist. Soc. B, 68(3):41 1-436, 2006. doi: 10.1111/j; 1467-9886..20066005555
[25] Thomas Müller, Brian Mewilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural importance sampling. ACM Trans. Graph., 38(5):145:1-145: 2019. doi: 10.1 10.1145/3341155
[26] Radford M. Neal. Annealed importance sampling. Stat. Comput., 1 11(2): 10.1023/A:1000992215522
139, 2001. doi:
[27] Radford M. Neal. Slice sampling. Ann. Stat., 31(3), 2003, doi: 10.1214/awx/1555524461
[28] Erich Novak. Deterministic and Stochastic Error Bounds in Numerical Analysis. Lecture Notes in Mathematics. Beringer-Verlagee 1988. doi: 10.1007/BF600099992.
[29] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457), 2019. doi: 10.1 10.1126/science..a 147.
[30] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res., 22(57):1-64, 2021.
[31] Philipp Petersen. Neural network theory, 2022.
http://ppc-etereen.ee/eeeer
Network_Theee Theory. pdd
[32] Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on International Conference on Machine Learning- Volume 37, ICML' 15, pages 1530-1538, Lille, France, 2015.
[33] Grant M. Rotskoff and Eric Canden-Eijndee Dynamical computation of the density of states and Bayes factors using nonequilibrium importance sampling. Phys. Rev. Lett., 122(15):150602. 2019. doi: 10. 103/Physseev ett. 122.150660
[34] Walter Rudin. Principles of Mathematical Analysis. 1976.
[35] Filippo Santambrogio. dacorogna- moserr approach to flow decomposition and minimal flow problems. ESAIM: ProcS, 45:265 274, 2014. doi: 10. /555//occ/014450222
[36] Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser, Cham, 2015.
[37] John Skilling. Nested sampling. AIP Conf. Proc., 735(1):395 405, 2004. doi: 10.1063/1. 1835238.
[38] John Skilling. Nested sampling for general Bayesian computation. Bayesian Anal., 1(4): 833-859, 2006. doi: 10. 1211406-BA122
[39] E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Commun. Pure Appl. Math., 16(2):145-164, 2013. doi: 10.1002/cpa.21423.
[40] Esteban G. Tabak and Eric Vanden -ijnden. Density estimation by dual ascent of the log- likelihood. Commun. Math. Sci., 8(1):217-233: 2010.
[41] Erik H. Thiede, Brian Van Koten, Jonathan Weare, and Aaron R. Dinner. Eigenvector method for umbrella sampling enables error analysis. Chem. Phys., 145(8):084115, 15, 2016. doi: 19.1069//.496069.
12

---

[42] Achille Thin, Yazid Janati EI Idrissi, Sylvain Le Corff, Charles Ollion, Eric Moulines, Arnaud Doucet, Alain Durmus, and Christian X Robert. NEO: Non equilibrium sampling on the orbits of deterministic transform. In Advances in Neural Information Processing Systems, volume 34, pages 17060-17071. Curran Associates, Inc., 2021.
[43] Glenn M Torrie and John P Valleau. Nomphysical sampling distributions in Monte Carlo free-energy estimation: Umbrella sampling. J. Comput. Phys., 23(2):187-199, 1977.
[44] Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste-Julien. Painless stochastic gradient: Interpolation, line-search, and convergence rates. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.,
2019.
[45] Peter Wirnsberger, Andrew J. Ballard, George Papamakarios, Stuart Abererombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. J. Chem. Phys., 153(14):144112, 2020. doi: 10. 10.1063/ 5.0018903.
[46] Hao Wu, Jonas Köhler, and Frank Noe. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 9934-9944,20 2020.
13


Reviews:
[{"review_id": "jDNFKg6ZnwG", "reviewer": "Reviewer_LyJz", "summary": "The paper approaches the problem of normalization constant estimation for the case when the density of the target distribution is given as an unnormalized function. To be precise, the paper builds upon the Non-Equilibrium Importance Sampling (NEIS) [1]. Unlike the conventional Importance Sampling, NEIS takes the samples from the reference density (known normalized density) and propagates them along some vector field, which we are free to choose. One can obtain the resulting estimator of the normalization constant as the properly weighted average of the target density (unnormalized) along the sampled trajectories.\n\nThe main contribution of this paper is the following. The authors use the fact that NEIS yields an unbiased estimate for any choice of the vector field. Using this degree of freedom, they learn a parametric model of the vector field that minimizes the variance of the resulting estimator.\n\nThe main theoretical result of the paper is the achievability of the zero-variance estimate of the normalization constant. Namely, when we take the infinite time for the trajectory integration, there exists a vector field that yields a zero-variance NEIS estimator. Furthermore, the authors demonstrate that every local minimum for this problem is actually global. In practice, the authors consider only finite-time estimators, though.\n\nFor the empirical evaluation, the authors consider 2d, 10d mixtures of Gaussians, and 10d Funnel distribution. They demonstrate that the proposed method outperforms AIS and the original NEIS in a reasonable setting. Moreover, for the Funnel distribution, the authors parameterize the vector field as a linear model, significantly reducing the computational efforts.\n\n1. Grant M. Rotskoff and Eric Vanden-Eijnden. Dynamical computation of the density of states and Bayes factors using nonequilibrium importance sampling. Phys. Rev. Lett., 122(15):150602, 2019. doi: 10.1103/PhysRevLett.122.150602.\n", "questions": "In the section \"Strengths and Weaknesses,\" I gave my suggestion for the improvements of the experimental section.\n\nRegarding the approach, I have the following question for the authors. Assuming the network can generate arbitrary vector fields, does it make sense to change the time interval $[t_-, t_+]$?", "limitations": "As I've mentioned in \"Strengths and Weaknesses,\"  the computational side of the proposed method is not fully addressed, and I would expect a comparison of the methods with the fixed computational budget.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, the paper does a decent job on every single aspect of a NeurIPS publication. It's clearly presented. It focuses on the intuition of the method in the main body and exhaustively describes all the details in the appendix. The theory is novel and immediately results in a motivated application. One might say that the empirical evaluation is not extensive. However, it's neat (properly compared with relevant algorithms) and demonstrates the clear superiority of the approach.\n\nI feel that the weakest point of the method is its computational burden. Namely, backpropagation through the trajectories of ODEs is computationally expensive. Also, the unbiased estimate is formulated as a continuous integral, which, in practice, has to be discretized. While the latter seems to be not an issue since the integral is 1d, I could imagine that the computational efforts spent on learning might overweight the performance gains during the evaluation. I think that a comparison of the methods with the fixed computational budget is essential here (I haven't found the corresponding numbers in Appendix I).", "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": "dH8vh2e7CU", "reviewer": "Reviewer_vi5E", "summary": "This paper considers non-equilibrium importance sampling (NEIS), wherein samples from a base distribution are evolved along flow lines and a particular quantity is averaged along these flow lines.  This results in a theoretically unbiased estimate of a normalizing constant (like in standard importance sampling), and since standard importance sampling can be derived as a special case of NEIS, this approach should allow one to construct lower variance estimators.  The authors derive a PDE which is satisfied by vector fields that produce optimal (zero variance) importance schemes, but then use a neural ODE approach to minimize the empirical variances of samples.  Finally, a few problems are considered where the NEIS method greatly reduces the variance of normalizing constant estimators compared to adaptive importance sampling.", "questions": "* How were the flow fields chosen for the direct GD method for Figures 2, 3 and 4?  in Figure 2 it is $-\\nabla U_1$, in Figure 3 it is $-5\\nabla U_1$, in Figure 4 it is $-2\\nabla U_1$.\n* Why was a neural ODE not used in the funnel toy model?  In general one does not know a priori what parameterization of the vector field will be good, and so using the neural ODE on the funnel model would be a good test-case for the behavior of the neural ODE + NEIS approach to pathological distributions that may arise in practice.\n* In figures 2, 3, and 4 the NEIS-based estimators (especially Direct GD in figures 2 and 3, and the Linear (before model  and even Linear (after) model to some extent in Figure 4) show substantial bias.  This is at odds with the theoretical result that regardless of the vector field NEIS should be unbiased.  Is this due to the discretization of the ODE?  Or are there just rare events where a very large normalizing constant is estimated and the runs presented here got \"unlucky\" in having exclusively underestimated normalizing constants?  In any case, this disconnect between the theoretical result that all of these estimates should be unbiased and the observed substantial bias should be discussed / addressed.", "limitations": "In terms of limitations, please see the discussion above in \"Strengths and Weaknesses\" about issues surrounding discretization.  I agree with the authors that there are no obvious negative societal impacts of the work.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* Overall I really enjoyed the paper.  The main results were technical, but clearly presented, and the approach seems to be promising.\n\n* The theoretical results in the appendix were very thorough and I appreciated the examples of explicitly solvable dynamics and other concrete examples.\n\nWeaknesses:\n\n* A major weakness of the present paper, in my opinion, is the general lack of discussion of the impact of discretizing the flow along the vector field.  All of the theoretical results are about the idealized flow, but in practice the ODEs have to be solved numerically.  Presumably this introduces complicated biases and if the vector field is learned using a neural net, there is no obvious reason that the learned ODE needs to be amenable to easy and accurate numerical integration.  More discussion of and investigation of the impact of discretization would strengthen the paper, in my opinion.\n\n* A minor point -- 3 different methods of using the neural ODE framing (directly estimating an arbitrary vector field, estimating a scalar potential and then taking gradients, or the divergence-free form) but then results were only presented from one of these for each toy problem (other than training curves in Figure 9).\n\n* The training of the neural ODE seems difficult -- the two trials seem to rarely converge to similar optima in Figure 9, and the runtime of the training seems quite long -- in the hours that it takes to train the neural ODE to construct the vector field, many more samples could be drawn from the proposal distribution and used in a less fancy importance sampling scheme.  It would be good to consider comparisons of the various estimators in terms of accuracy per fixed amount of wall time (e.g., how does AIS or vanilla IS do if it just samples continuously for the amount of time it takes to train the NEIS vector field?).\n\n* One other aspect that I felt was neglected was the importance of the proposal distribution, $\\rho_0$.  In standard importance sampling, how to choose $\\rho_0$ (from some family) is well-understood -- we want it to be as close to $\\rho_1$ as possible, where closeness is measured by the chi divergence.  In the NEIS case, the optimal $\\rho_0$ is less obvious -- presumably choosing it close to $\\rho_1$ is useful, but one also wants the vector field to be \"nice\" and easy to learn.\n\nMinor comments and Typos:\n\n* Line 16: \"order\" --> \"orders\"\n* Line 45: \"based on Jarzynski\" --> \"based on the Jarzynski\"\n* Line 75: \"order\" --> \"orders\n* Line 77: \"lead into\" --> \"lead to\"\n* Line 83: \"one use in\" --> \"one used in\"\n* Line 168: \"sampling procedure\" --> \"a sampling procedure\"\n* Line 188: \"will necessarily produces\" --> \"will necessarily produce\"\n* Line 211: I believe the plural of \"anstaz\" is \"anstazes\"\n* Line 227: \"minimize\" --> \"minimizing\"\n* Figure 2b caption: \"descend\" --> \"descent\"\n* Line 289: \"descend\" --> \"descent\"\n* The notation $|\\cdot|$ is used throughout.  I believe that this is the $\\ell_2$ norm, but that should be made explicit.\n* It should be made explicit that the minimization of equation (15) is over $V$ ", "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": "TlcemdLzyIF", "reviewer": "Reviewer_Hz5v", "summary": "The paper analyzes a recently proposed non-equilibrium importance sampling approach to integrating / estimating the partition function of intractable densities. This approach provides an unbiased estimator with bounded variances typically much better than the variances provided by other estimators (which are in some cases infinite). Some experiments on basic multivariate densities including asymmetric and symmetric MoG, and Neal's funnel in 10d, provide empirical support.", "questions": "Why were these particular densities chosen for numerical examples?\n\nIf we ran AIS for 1h, could we get much more exhaustive exploration and tighter variance?\n\nHow does the training approach compare/contrast with diffusion models, where we are training a map that pushes a perturbed sample back toward the original sample?\n\nIs the bias on Neal's funnel attributable to the variational approximation, or some other approximation in the procedure?", "limitations": "The academic nature of the examples is noted. Perhaps the expense of training ought to be more explicitly called out.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality\nThe NEIS method seems to be new to the ML/stats community. The predecessor work in physics is only a couple years old. Connecting this work to Neural ODEs and neutral transport maps is a novel contribution, as are the variational approximations enabling these connections.\n\nQuality\nThe math is developed reasonably well and connections to existing work are shown (might also think about connections to diffusion models, neural transport HMC, neural optimal transport). A possible complaint might be that the math is developed fairly quickly/compactly in the body of the paper, with >30pg of detail reserved to supplementaries (which I've only browsed).\nA weaker portion of the paper, in my view, is the experiments. On the positive side, the experiments demonstrate fairly well the unbiased and low variance nature of this estimation approach. However, I would like to see some richer densities than [hierarchical/mixture] gaussians. Multiple libraries exist that could provide other interesting densities (inference gym, ppl bench, posterior db). One reason the authors might have chosen these particular densities is for analytically computable Z, which enables scaling all plots to a 0-1 size. But I do see a log-Gaussian cox process in the anon github, which suggests other densities might have been considered? The supplementary indicates that the method takes 1-2h to fit/optimize on these densities, which calls into question the fairness of comparison with 100 transitions of AIS (which should take only seconds unless I've misunderstood something). If we ran AIS for 1h, could we get much more exhaustive exploration and tighter variance? I might like to see plots of Z bias against number of density evaluations for each of the methods considered.\n\nClarity\nI would have been helped by adding some prose about usage of $\\nabla$ to the notation section. Presumably it is used throughout as an operator, but some notations like $\\nabla \\cdot$ (e.g. eq 5) are unfamiliar.\n\nSignificance\nBecause the approach seems to be fairly expensive and is only compared to much less expensive AIS, I am somewhat unclear about its practical significance. I do think the approach seems academically interesting to the NeurIPS audience and the numerical experiments indicate that, at least in cases where we are willing to pay compute, this provides an interesting new point on the frontier of approximate inference fidelity vs cost. I wonder whether comparisons with a trainable-RealNVP-reparameterized AIS/SMC might give a more fair comparison in terms of cost.", "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": "SpmyArYIzid6", "reviewer": "Reviewer_QJgR", "summary": "The paper has two main contributions. First, a theoretical analysis of Non-Equilibrium Importance Sampling (NEIS), a technique used to get unbiased estimates (in theory) of expectations and normalizing constants under intractable distributions, based on performing averages averages along the flow generated by the evolution of a dynamical system (given as a first order ODE). The main theoretical results proved in the paper involved showing the existence (under certain assumptions) of an optimal velocity field that yields an estimator with zero variance, and a connection with optimal transport maps. And second, the paper explains how NEIS with a learnable velocity field parameterized by a NN can be used in practice by training the method's parameters to minimize the estimator's variance, and shows an empirical evaluation on three low-dimensional models.", "questions": "- Have you tried using smaller batch-sizes?\n\n- If I understand correctly, theoretical results (prop 3.1, prop 4.1) involve the formulation for equation 10, which integrates time from $-\\infty$ to $\\infty$. Do the results hold for the case where finite integration limits are used (i.e. eq. 8), which is the one used in practice?", "limitations": "See above regarding limitations.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The paper is written in a clear way, and presents original theoretical findings. The theoretical analysis answers relevant questions about NEIS, including conditions for the existence of the optimal dynamics that lead to an estimator with zero variance. All statements are clearly stated including all assumptions, and most assumptions are adequately discussed in the main text.\n\nWeaknesses: In my opinion, the bigger weaknesses involve the empirical evaluation, the claims made around it, and the method's practicality. For instance:\n\n- The conclusion claims that the empirical evaluation shows the method's effectiveness in high dimensions, when the evaluation is done on three models of dimensionalities 2, 10 and 10.\n\n- The method appears to be extremely slow in practice, even for the simple 2d model, but this is not mentioned in the main text. It takes 1.5 hours to train for 200 SGD steps on the two dimensional model (mixture of two Gaussians) and 5 hours on one of the ten dimensional model (mixture of 4 Gaussians). I understand that the method uses neural networks, and simulations were done using CPUs, but all NNs are fairly small, of around 2 layers with 10 neurons each, so I don't think that should be a big bottleneck even when using CPUs. (Are training times on the funnel model, for which NNs are not used, reported?).\n\n- The training uses minibatches of thousands of samples per training step, which may indicate that there are some issues with the variance of the objective and its gradients. This affects the cost of the method. Is it necessary to use such large minibatches?\n\n- The vanilla estimator is taken to be importance sampling with a standard Gaussian as proposal, which is expected to work very poorly, specially since it is an extremely bad approximation of the targets used. I think that replacing this vanilla baseline with a diagonal Gaussian trained with VI, for instance, would provide a more interesting comparison.\n\nOverall, I feel the empirical evaluation does not really show the practicality of the method, somewhat contradicting the claims made in the paper (if any of the statements above is wrong, please let me know). I think that more training details should be included in the main text explicitly (training times, sizes of the networks, steps used to numerically evaluate integrals and simulate dynamics, etc), as they are extremely relevant and currently only present in the Appendix. Also, I think that a more explicit discussion regarding the method's cost should be included, specially since the objective is somewhat complex, involving two nested integrals whose integrands depend on the values taken by a dynamical system at different times.", "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