Path Integral Sampler: A Stochastic Control Approach For Sampling
Authors: Qinsheng Zhang
Yongxin Chen
Venue: ICLR 2022 Poster (2022)
Keywords: Sampling
Path Integral
Stochastic Differential Equation
MCMC

Abstract: We present Path Integral Sampler~(PIS), a novel algorithm to draw samples from unnormalized probability density functions. The PIS is built on the Schr\"odinger bridge problem which aims to recover the most likely evolution of a diffusion process given its initial distribution and terminal distribution. The PIS draws samples from the initial distribution and then propagates the samples through the Schr\"odinger bridge to reach the terminal distribution. Applying the Girsanov theorem, with a simple prior diffusion, we formulate the PIS as a stochastic optimal control problem whose running cost is the control energy and terminal cost is chosen according to the target distribution. By modeling the control as a neural network, we establish a sampling algorithm that can be trained end-to-end. We provide theoretical justification of the sampling quality of PIS in terms of Wasserstein distance when sub-optimal control is used. Moreover, the path integrals theory is used to compute importance weights of the samples to compensate for the bias induced by the sub-optimality of the controller and the time-discretization. We experimentally demonstrate the advantages of PIS compared with other start-of-the-art sampling methods on a variety of tasks.

Published as a conference paper at ICLR 2022
PATH INTEGRAL SAMPLER: A STOCHASTIC CONTROL APPROACH FOR SAMPLING
Qinsheng Zhang
Yongxin Chen
Center for Machine Learning Georgia Institute of Technology qzhang4 90 gatech... edu
School of Aerospace Engineering Georgia Institute of Technology yongchenngatecch. edu
ABSTRACT
We present Path Integral Sampler (PIS), novel algorithm to draw samples from unnormalized probability density functions. The PIS is built on the Schrödinger bridge problem which aims to recover the most likely evolution of diffusion process given its initial distribution and terminal distribution. The PIS draws samples from the initial distribution and then propagates the samples through the Schrödinger bridge to reach the terminal distribution. Applying the Girsanov the- orem, with simple prior diffusion, we formulate the PIS as stochastic optimal control problem whose running cost is the control energy and terminal cost is chosen according to the target distribution. By modeling the control as a neural network, we establish a sampling algorithm that can be trained end-to-end. We provide theoretical justification of the sampling quality of PIS in terms of Wasser- stein distance when sub-optimal control is used. Moreover, the path integrals theory is used compute importance weights of the samples to compensate for the bias induced by the sub-optimality of the controller and time-discretization. We experimentally demonstrate the advantages of PIS compared with other start- of-the-art art sampling methods on a variety of tasks.
INTRODUCTION
We are interested in drawing samples from target density j Zμ known up to a normalizing constant Z. Although has been widely studied in machine learning and statistics, generating asymptotically unbiased samples from such unnormalized distribution can still be challenging (Tal- war 2019). In practice, variational inference (VI) and Monte Carlo (MC) methods are two popular Trameworks for sampling.
Variational inference employs a density model q, from which samples are easy and efficient to draw, to approximate the target density (Rezende & Mohamed] 2015 Wu et al. 2020). Two important ingredients for variational inference sampling include a distance metric between q and fi to identify good and the importance weight to account for the mismatch between the two distributions. Thus, in variational inference, one needs to access the explicit density of which restricts the possible parameterization of q. Indeed, explicit density models that provide samples and probability density such as Autoregressive models and normalizing flow are widely used in density estimation (Gao) al. 2020a Nicoli al. 2020). However, such models impose special structural constraints on the representation of q. For instance, the expressive power of normalizing flows (Rezende & Mohamed, 2015) is constrained by the requirements that the induced map has be bijective and its Jacobian needs to be easy-to-compute (Cornish et al. 2020 Grathwohl et al. 2018 Zhang & Chen 2021).
well-designed Markov chain (MCMC) or sampling ancestrally (MacKay) 2003). Among them, Sequential Monte Carlo and its variants augmented with annealing trick are regarded as state-of-the-art in certain sampling tasks (Del Moral et al. 2006). Despite its popularity, MCMC methods may suffer from long mixing time. The short-run performance of MCMC can be difficult to analyze and samples often get stuck in local minima (Nijkamp et al. 2019 Gao et al. (2020b). There some recent works exploring the possibility of incorporating neural networks to improve MCMC (Spanbauer et al., 22200; Li et al.) 2020b). However, evaluating existing MCMC empirically, not to say designing an objective loss
Most MC methods generate samples by iteratively simulating

---

Published as a conference paper at ICLR 2022
Uncontrolled Q"
Terminal Cost v log L
Optimal control _- Q
H
Figure 1: Illustration of Path Integral Sampler (PIS). The optimal policy of a specific stochastic control problem where a terminal cost function chosen according the given target density p. can generate unbiased samples over a finite time horizon.
function to train network-powered MCMC, is difficult (Liu et al. 2016 Gorham & Mackey 2017) Most existing works in this direction focus only on designing data-aware proposals (Song et al., 2017 Titsias & Dellaportas 2019) and training such networks can be challenging without expertise knowledge in sampling.
In this work, we propose an efficient sampler termed Path Integral Sampler (PIS) to generate samples by simulating a stochastic differential equation (SDE) in finite steps. Our algorithm built on the
Schrödinger bridge problem (Pavon 1989 Dai Pra 1991 Léonard 2014 Chen et al. [02T]) whose original goal was to infer the most likely evolution of a diffusion given marginal distributions at two time points. With proper prior diffusion model, this Schrödinger bridge framework can be adopted for the sampling task. Moreover, it can be reformulated as a stochastic control problem (Chnn et al.) 2016) whose terminal cost depends the target density so that the diffusion under optimal control has terminal distribution ft. We model the control policy with a network and de- velop method train it gradually and efficiently. The discrepancy of the learned policy from the optimal policy also provides an evaluation metric for sampling performance. Furthermore, PIS can be made unbiased even with sub-optimal control policy via the path integral theorem to compute the importance weights of samples. Compared with VI that uses explicit density models, PIS uses an implicit model and has the advantage of free-form network design. The explicit density mod- els have weaker expressive power and ficxibility compared with implicit models, both theoretically and empirically (Cornish al., 2020 Chen al., 2019 Kingma & Welling 2013 Mohamed & Lakshminarayanan (0166). Compared with MCMC, PIS is more efficient and is able to generate high-quality samples with fewer steps. Besides, the behavior of MCMC over finite steps can be analyzed and quantified. We provide explicit sampling quality guarantee in terms of Wasserstein distance to the target density for any given sub-optimal policy.
Our algorithm is based on Tzen & Raginsky (2019), where the authors establish the connections be- tween generative models with latent diffusion and stochastic control and justify the expressiveness of such models theoretically. How to realize this model with networks and how the method performs on real datasets are unclear in TTen & Raginsky (2019). Another closely related work is Wu et al. 2020): Arbel et al. (202T), which extends Sequential Monte Carlo (SMC) by combining determin- istic normalizing flow blocks with stochastic MCMC blocks. To be able to evaluate the importance weights efficiently, MCMC blocks need to be chosen based annealed target distributions care- fully. In contrast, in PIS one can design expressive architecture freely and train the model end-to-end without the burden of tuning MCMC kernels, resampling or annealing scheduling. We summarize our contributions as follows. 1). We propose Path Integral Sampler, a generic sampler that generates samples through simulating a target-dependent SDE which can be trained with free-form architec- ture network design. We derive performance guarantee in terms of the Wasserstein distance to the target density based on the optimality of the learned SDE, 2). An evaluation metric provided to quantify the performance of learned PIS. By minimizing such evaluation metric, PIS can be trained end-to-end. This metric also provides an estimation of the normalization constants of target dis- tributions. 3). PIS can generate samples without bias even with sub-optimal SDEs by assigning importance weights using path integral theory. 4). Empirically, PIS achieves the state-of of-the-art art sampling performance in several sampling tasks.
2

---

Published as a conference paper at ICLR 2022
2 SAMPLING AND STOCHASTIC CONTROL PROBLEMS
We begin with a brief introduction to the sampling problem and the stochastic control problem. Throughout, we denote by r {x,, 0 t T) continuous-tiie stochastic trajectory.
2.1 SAMPLING PROBLEMS
We are interested in drawing samples from a target distribution μux) = / in Rs where Z is the normalization constant. Many sampling algorithms rely on constructing stochastic process that drives the random particles from an initial distribution that is easy to sample from, to the target distribution jr.
In the variational inference framework, one seeks to construct parameterized stochastic process to achieve this goal. Denote by = C((O, 7]; Rd) the path space consisting of possible trajectories and by PP the measure over 12 induced by stochastic process with terminal distribution Jr at time T. Let Q be the measure induced by a parameterized stochastic and denote its marginal distribution at T by n Then, by the data processing inequality, the Kullback-Leibe divergence (KL) between marginal distributions o and μ can be bounded by
J dQ dQ log
dQ dp
DKI( s DKI.
(1)
Thus, DKL (P) serves as a performance metric for the sampler, and a small (KL(P)P) value corresponds to good sampler.
2.2 STOCHASTIC CONTROL.
Consider model characterized by special stochastic differential equation (SDE) Särkkää & Solin 2019)
dx, = u,dt + dw,,
where X,, u, denote state and control input respectively, and w, denotes standard Brownian motion. In stochastic control, the goal is to find feedback control strategy that minimizes a certain given cost function.
The standard stochastic control problem be associated with any cost and any dynamics. In this work, we only consider cost of the form
E
uell de V(XT) x
where represents the terminal cost. The corresponding optimal control problem can be solved via dynamic programming (Bertsckas et 2000), which amounts solving the Hamilton-Jacobe Bellman (HJB) equation (Evan 1998)
aV, Øt
1 AV = Vi(() = vT(-).
VV 2
(4)
The space-time function (x(x) is known cost-to- to function value function. The optimal policy can be computed from (x(x) as (Pavon 1989)
= (
3 PATH INTEGRAL SAMPLER
It turns out that, with a proper choice of initial distribution v and terminal loss function V the stochastic control problem coincides with sampling problem, and the optimal policy drives samples from perfectly. The process under optimal control can be viewed the posterior of uncon- trolled dynamics conditioned on target distribution as illustrated Fig I Throughout, we denote by Q" the path measure associated with control policy u. We also denote by 4t the terminal distri- bution of the uncontrolled process Q". For the case of presentation, we begin with sampling from a normalized density n, and then generalize the results unnormalized jı in Section 3.4
3

---

Published as a conference paper at ICLR 2022
3.1 PATH INTEGRAL AND VALUE FUNCTION
Thanks to the special cost structure, the nonlinear HJB eq (4) can be transformed into linear partial differential equation (PDE)
Jé. Ōt 2 A0: = =0 Φr(() = exp(-v(.))
(6)
by logarithmic transformation (Särkkä Solin 2019 Vi(x) = - log (x)).. By the celebrated
Feynman-Kaa Kac formula @ksendal 2003). the above has solution
Φe (x) [exp( v (xr)) x = x].
We remark that leq (7) implies that the optimal value function can be evaluated without knowing the optimal policy since the above expectation is with respect to the uncontrolled process Q This
is exactly the Path Integral control theory (Theodorou et al. 2010: Theodorou & Todorov 2012
Thijssen Kappen 2015). Furthermore, the optimal control (t, x))
e (expl (xxp(- Jt dw, x]
un (x) = (x) lim
s (s gg (cxpl -x(())) X; =
meaning that u can also be estimated by uncontrolled trajectories.
3.2 SAMPLING AS A STOCHASTIC OPTIMAL CONTROL PROBLEM
There are infinite choices of control strategy u such that eq (2) has terminal distribution μ.. We
interested in the that minimizes the KL divergence to the prior uncontrolled process. This
exactly the Schrödinger bridge problem (Pavon 1989 Dai Pra 1991] Chen et al., [016 [021])
which has been shown have stochastic control formulation with cost being control efforts. In
cases where v a Dirac distribution, it is the same as the stochastic control problem in /Section 2.2
with proper terminal cost as characterized in the following result (Tzen & Raginsky 2019).
Theorem (Proof in lappendix A). When v is a Dirac distribution and terminal loss chosen as
(xT) the distribution Q induced by the optimal control policy is (xx) log p(x))
(*(r)= = 2"(x=))(()))(x
(9)
Moreover, Q" (xTx = p(xT).
To gain more insight, consider the KI. divergence
m(xT), ( (x(x)) = DKL. (2" + Ez [log
= (V
(10)
Thanks to the Girsanov theorem
Särkkä & Solin 2019).
dQu dQ
expl (" |ell dt u.dw;).
(11)
follows that
11" oo ||uu|12 dt).
(KL(O"1Q0)) EQu
(12)
Plugging eg (12)
into eq (10)
vields
DKL (T))
= EQu 1//0
2 | 1|2 dt + log
1 (xx)) μ μ(XT)
(13)
which is exactly the cost defined in (3) with V log _ Theorem implies that once the optimal
control policy that minimizes this cost found, also drive particles from Xo to XT p.
4

---

Published as a conference paper at ICLR 2022
Algorithm 1 Training
Input: Vector: Xo 0, Scalar: Yo = 0
Output: uz(x) parameterized by 0
Define:
SDE drift [tt, [x,, yel) = [ooe (x,), 2
1 [x+) diffusion g(t, , [x,, Mi]) = [1,0]
loop epoches
NTT = sfeint(f. 8: [xo, yol, [0. 2]
# Integrate SDE from 0 to T with Neural SDE
Gradient descent step Volyr + log p(xT)
(xxr)
Optimize control policy
done
3.3
OPTIMAL CONTROL POLICY AND SAMPLER
Consider the sampling strategy from given target density by simulating SDE in eq (2) under optimal control. Even though the optimal policy is characterized by only rare case (Gaussian target distribution) it has an analytic closed-form.
Optimal Policy Representation:
eq (8)
For more general target distributions, we can instead evaluate the value function jeg (7) via empir- ical samples using Monte Carlo. The approach essentially importance sampling whose proposal distribution is the uncontrolled dynamics. However, this approach has two drawbacks. First, it is known that the estimation variance can be intolerably high when the proposal distribution is not close enough to the target distribution (MacKay] 2003). Second, even if the variance is acceptable, without a good proposal, the required samples size increases exponentially with dimension, which prevents the algorithm from being used in high or even medium dimension settings (Neal) 2001).
To overcome the above shortcomings, we parameterize the control policy with neural network ug. We seek control policy that minimizes the cost
u arg min eo= [ |lull dt log (xx)
1 (XT)
(14)
2
The formula eq (14) also serves as distance metric between ug and u" as in eq (13)
Gradient-informed Policy Representation: It believed that proper prior information can signifi- cantly boost the performance of neural network (Goodfellow et al. 2016). The score V log (xx) has been used widely to improve the proposal distribution in MCMC (Li al. [2020b] Hoffman Gel- man (214)) and often leads to better results compared with proposals without gradient information. In the same spirit, we incorporate log (xx) and parameterize the policy as
u, (x) NN, (t, x) NN2(t) x V log ((x),
(15)
where NN, and NN₂ are two neural networks, Empirically, we also found that the gradient informa- tion leads to faster convergence and smaller discrepancy DKL (Q" Q*).. We remark that PIS with
policy leq (15) can be viewed as modulated Langevin dynamics (MacKay 2003) that achieves within finite time instead of infinite time.
Optimize Policy: Optimizing Ug requires the gradient of loss in (14) which involves u, and the terminal state xy- To calculate gradients, we rely on backpropagation through trajectories. We train the control policy with recent techniques of Neural SDEs eLi al. 2020a Kidger et al., (02I), which greatly reduce memory consumption during training. The gradient computation for Neural SDE is based stochastic adjoint sensitivity, which generalizes the adjoint sensitivity method for Neural ODE (Chen et al. 2018). Therefore, the backpropagation in Neural SDE is another SDE associated with adjoint states. Unlike the training of traditional deep MLPs which often runs into gradient vanishing/explodin
issues, the training of Neural SDE/ODE is more stable and not sensitive the number of discretization steps (Chen et al., 2018 Kidger et al. 2021). We augment the origin SDE with state Lo ||u,l] 2| ds such that the whole training can be conducted end to end. The full training procedure in provided in Algorithm
Wasserstein distance bound: The PIS trained by Algorithm I can not generate unbiased samples from the target distribution μ for two reasons. First, due the non-convexity of networks and randomness of stochastic gradient descent, there is no guarantee that the learned policy is optimal. Second, even if the learned policy is optimal, the time-discretization emor in simulating SDEs is in-
5

---

Published as a conference paper at ICLR 2022
evitable. Fortunately, the following theorem quantifies the Wasserstein distance between the sampler and the target density. (More details and a formal statement can found in lappendix C)
Theorem (Informal), Under mild condition, with sampling step size At. if I|u - uel S de for
any I, then
(16)
3.4 IMPORTANCE SAMPLING
The training procedure for PIS does not guarantee its optimality. To compensate for the mismatch between the trained policy and the optimal policy, we introduce importance weight to calibrate generated samples. The importance weight can be calculated by (more details infppenddix B)
(T) exp( h w t (r) (Q
|u, dt - u/dw, ((xT))
(17)
Algorithm 2 Sampling
We note eq (17) resembles training objective (14)
Input: Vector:
O, Scalar: yo = 0 Output: Samples with weights
Indeed, eq (14) is the average of logarithm of eq (17) If the trained policy is optimal, that is, Qa Q all
for i < 1 to N do
the particles share the same weight. We summarize the
Al = 1, 4--11 1, Aw ~ NO, AIT,,
sampling algorithm in in/llgorihmm 2
x,- u Aw
Effective Sample Size: The Effective Sample Size (ESS), ESS [[=](r is popular metric to measure the variance of importance weights, ESS is of- ten accompanied by resampling trick (Tokdar & Kass 2010) to mitigate deterioration of sample quality. ESS also regarded as a metric for quantifying goodness of sampler based importance sampling. Low ESS means that estimation or downstream tasks based on such sampling methods may suffer from high variance. ESS of most importance samplers is decreasing along the time. Thanks to the adaptive control policy in PIS, we can quantify the ESS of PIS based on the optimality of learned policy. For the sake of completeness, the proof in provided in lappendix D
Vi- Vi-1 + u A' 2 |ull2 At end for Outputs:
XN XN, exp(-yN log
((xx)) p(xN)
Theorem (Corollary 7 (Thijssen & Kappen 2015 If max,, [x); ut (x)]1 S 1. then
Estimation of normalization constants: In most sampling problems we only have access to the target density up normalization constant, denoted by jr Zμ. PIS can still generate samples following the same protocol with new terminal cost V = log n0 V log Z. The additional con- stant- log Z is independent of xy and thus does not affect the optimal policy and the optimization of ug. As byproduct, we can estimate the normalization constants (more details in lappendix E}.
Theorem 4. For any given policy the logarithm of normalization constant is bounded below by
E (-)) log Z,
(18)
where (-) Moreover, for any sub-optimal policy, unbiased estimation of Z using importance sampling is
= fo [x:))x) dt + u/(x)) (x, dw, + ). The equality holds only when u = u
7 Er-ou (exp(-S""()
(19)
4 EXPERIMENTS
In this section we present empirical evaluations of PIS and the comparisons several baselines. We also provide details of practical implementations. Inspired by Arbel et al. (2021), ), we conduct experiments for tasks of Bayesian inference and normalization constant estimation.
We consider three types of relevant methods. The first category is gradient- guided MCMC methods without the annealing trick. It includes the Hamiltonian Monte Carlo (HMC) (MacKay) 2003) and

---

Published as a conference paper at ICLR 2022
O
PIS-Grai
AFT
SMC
NUTS
HMC
VI-NF
PIS-NN
Figure 2: Sampling performance on rings-shape density function with 100 steps. The gradient information can help PIS-Grad and MCMC algorithm improve sampling performance.
No-U-Turn Sampler (NUTS) (Hoffman & Gelman 2014). The second is Sequential Monte Carlo with annealing trick (SMC), which is regarded state-of-the-art sampling algorithm (Del Moral it al. (0066) in terms of sampling quality. We choose a standard instance of SMC samplers and the recently proposed Annealed Flow Transport Monte Carlo (AFT) (Arbel et al. 202I). Both use
default 10 temperature levels with a linear annealing scheme. We note that there are optimized SMC variants that achieve better performance (Del Moral et al.) | Chopin Papaspiliopoulos 2020 Zhou et al.] 2016). Since the introduction of advanced tricks, we exclude the comparison with those variants for fair comparison purpose. We note PIS can also be augmented with annealing trick, pos- sible improvement for PIS can be explored in the future. Last, the variational normalizing flow (VI- NF) (Rezende & Mohamed] 2015) is also included for comparison. We note that another popular line of sampling algorithms use Stein Siii--aaiaaionnll Gradient Descent (SVGD) or other particle- based variational inference approaches (Liu & Wang 2016 Liu et al. (019). We include the comparison and more discussions on SGVD in inlppeendix F.3 due to significant difference. In our experiments, the number of steps N of MCMC algorithms and the number of SDE time-discretizatio steps for PIS work as proxy for benchmarking computation times.
We also investigate the effects of two different network architectures for Path Integral Sampler. The first one is a time-conditioned neural network without any prior information, which we denote as PIS-NN, while the second one incorporates the gradient information of the given energy function as in eq (15) denoted as PIS- AradG When we have an analytical form for the ground truth optimal policy, the policy is denoted as PIS-GT. The subscript RW is to distinguish PIS with path integral importance weights (17) that usefeq (19)to estimate normalization constants from the ones with- out importance weights that the bound in ey (18) to estimate Z. For approaches without the annealing trick, we take default N = 100 unless otherwise stated. With annealing, N steps are the default for each temperature level, thus AFT and SMC rougly use 10 times more steps com- pared with HMC and PIS. We include more details about hyperparameters, training time, sampling efficiency, and more experiments with large N in lappendices F and G
4.1 PIS-GRAD PIS-NN: IMPORTANCE OF GRADIENT GUIDANCE
We observed that the advantage of PIS-Grad over PIS-NN is clearer when the target density has multiple modes in the toy example shown in Fig 2 The objective (xl( Q*) is known have zero forcing. In particular, when the modes of the density are well separated and Q is not expressive enough, minimizing DKL ( (QII*) can drive ((7) to zero on some area, even if (* > [Fox & Roberts 2012). PIS-NN and VI-NF generate very similar samples that almost cover half the inner ring. The training objective function of VI-NF can also be viewed as minimizing KL divergence between two trajectory distributions (Wu et al. 2020). The added noise during the process can encourage exploration but unlikely such noise only can overcome the local minima. On the other hand, the gradient information can help cover more modes and provide exploring directions.
4.2 BENCHMARKING DATASETS
Mode-separated mixture of Gaussian: We consider the mixture of Gaussian in 2-dimension. We notice that when the Gaussian modes are not far away from each other, all methods work well. However, when we reduce the variances of the Gaussian distributions and separate the modes of Gaussian, the advantage of PIS becomes clear even in this low dimension task. We generate 2000 samples from each method and plot their kernel density estimate (KDE) in Fig 4 PIS generates samples that are visually indistinguishable from the target density.
7

---

Published as a conference paper at ICLR 2022
Funnel distribution: We consider the popular testing distribution in MCMC literature (Hoffman &
Gelman 2014 Hoffman et al. 2019), the 10-dimensional Funnel distribution charaterized by
To ~ (0,9). T1:9220 N NO, exp(2x)I).
This distribution can be pictured as funnel with vp wide at the mouth of funnel, getting smaller the funnel narrows.
MG (d =
2)
Funnel (d
10)
LGCP (d 1600)
B
S
A
B
S
A
B
S
A
PIS Rw-
-0.012
0.013
0.018
PIS-NN
-1.691
0.370
1.731
-0.098
5e-3
0.098
-92.4
6.4
92.62
PIS-Grad
-0.440
0.024
0.441
-0.103
9e-3
0.104
-13.2
3,21
13.58
PIS RW-NN
-1.192
0.482
1.285
-0.018
7c-3
0.02
-60.8
4.81
60.99
PIS awwdGrrd
-0.021
0.030
0.037
-0.008
9e-3
0.012
-1.94
0.91
2.14
AFT
-0.509
0.24
0.562
-0.208
0.193
0.284
-3.08
1.59
3.46
SMC
-0.362
0.293
0.466
-0.216
0.157
0.267
-435
14.7
436
NUTS
-1.871
0.527
1.943
-0.835
0.257
0.874
-1.ee
8.01
1.3e3
HMC
-1.876
0.527
1.948
-0.835
0.257
0.874
-1.ee3
8.01
1.3e3
VI-NF -1.632
0.965
1.896
-0.236
0.0591
0.243
-77.9
5.6
78.2
Table I: Benchmarking on mode separated mixture Gaussian (MG), Funnel distribution and Log Gaussian Cox Process (LGCP) for estimation log normalization constants. B and S stand for estimation bias and standard deviation among 100 runs and A²= B2 S2.
We further investigate the normalization constant estimation problem for the challenging log Gaussian Cox process (LGCP), which is designed for modeling the posi- tions of Finland pine saplings. In LGCP (Salvatier al. 2016), underlying field X of positive real values is modeled using an exponentially seieenttilly-tranffommeee Gaussian process. Then used to parameterize Poisson points process to model locations of pine saplings. The posterior density is
Log Gaussian Cox Process:
X(x) ~ exp(
(x 11) POKK-p
p)
II iEd
exp(x.1: a 2x).;;),
(20)
where d denotes the size of discretized grid and y/4 denotes observation information. The modeling parameters, including normal distribution and a, follow Arbel et al. (202I) (See lappendix F).
Tab clearly shows the advantages of PIS for the above three datasets, and supports the claim that importance weight helps improve the estimation of log normalization constants, based on the comparison between PIS and PIS. We also found that PIS-Grad trained with gradient information outperforms PIS-NN. The difference more obvious in datasets that have well-separated modes, such MG and LGCP, and less obvious unimodal distributions like Funnel.
In all cases, PIS nw- -Gad is better than AFT and SMC. Interestingly, even without annealing and gradient information of target density, PISnw-Nh -NN outperform SMC with annealing trick and HMC kernel for the Funnel distribution.
4.3 ADVANTAGE OF THE SPECIALIZED SAMPLING ALGORITHM
From the perspective particles dynamics, most existing MCMC algorithms invariant to the target distribution. Therefore, particles are driven by gradient and random noise in a way that is independent of the given target distribution. In contrast, PIS learns different strategies to combine gradient information and noise for different target densities. The specialized sampling algorithm can generate samples more efficiently and shows better performance empirically in our experiments. The advantage can be showed in various datasets, from unimodal distributions like the Funnel dis- tribution to multimodal distributions. The benefits and efficiency of PIS are more obvious high dimensional settings as we have shown.
4.4 ALANINE DIPEPTIDE
Building on the success achieved by flow models the generation of asymptotically unbiased sam- ples from physics models (LeCCnn 1998), investigate the applications in the sampling of molec-
8

---

Published as a conference paper at ICLR 2022
KI.
1
Φ
1)1
1/2
1/3
VI-NF
175.6
+4.5
24.2 4.1
3.1 + 0.05
14.6 6.4
7e-2#5e-3
8.55-2+3.5e-3
SMC
183.3
-2.3
18.3
0.32 + 0.08
9.6 1.2
0.12=0005
0.15 = 9e-3
SNF
181.8
+0.75
6.34 0.71
0.17100.05
1.58 + 0.36
0.11 0.03
8.8e-2 488-3 8.7e-2--3e3
PIS-NN
171.3
4.061
5.24 0.35
0.32+40.03
1.03 + 0.23
5e-2--
Table KL-divergences comparison among variational approaches of generated density with target density in overall atom states distribution and five multimodal torsion angles. We emphasis Kl.* denote the KL divergence between unnormalized distribution due to lack of ground truth normaliza- tion constants. Mean and standard deviation are conducted with five different random seeds.
ular structure from simulation of Alanine dipeptide introduced in Wu et al. (2020). The target density of molecule
exp( (
X106:131
We compare PIS with popular variational approaches used in generating samples from the above model. More specifically, we consider VI-NF, and Stochastic Normalizing Flow (SNF) (Wul et al.) (020). SNF is very close to AFT (Arbel et al.) (02I). Both of them cou- ple deterministic normalizing flow layers and MCMC blocks except SNF uses an amor- tized structure. We include more details of MCMC kernel and We show generated molecular in Fig 3 and quantitative comparison in terms of KL divergence in Tab 2 including overall atom states dis- tribution and five multimodal torsion angles (backbone angles o, t and methyl rotation angles 1/, 722 73)) We remark that unweighted samples are used to approximate the density of torsion angles and all approaches do not use gradient information. Clearly, PIS gives lower divergence.
modification in lappendix F Figure 3: Sampled Alanine dipeptide molecules
4.5 SAMPLING IN VARIATIONAL AUTOENCODER LATENT SPACE
In this experiment we investigate sampling in the latent space of a trained Variational Autoen- coder (VAE). VAE aims to minimize (KL)g(()44)) (z|xx) (1)o)))) where (z(z(x) represents encoder and for decoder with latent variable z and data x. We investigate the posterior distri- bution
p(z)po
(21)
The normalization constant of such target unnormalized density function p(z) (x is exactly the likelihood of data points po (x), which serves as an evaluation metric for the trained VAE.
We investigate vanilla VAE model trained with plateau loss on the binary MNIST (LeCun) 1998) dataset. For each distribution, age estimation from 10 long-run SMC with 1000 temperature levels as the ground truth normal- ization constant. We choose 100 images ran- domly and run the various approaches on estimat- ing normalization of those posterior distributions in eq (21) and report the average performance Tab 3 PIS has lower bias and variance.
Table 3: Estimation of log (o(po()) of trained VAE.
regard the aver-
B
S
v B S
VI-NF
-2.3
0.76
2.42
AFT
-1.7
0.95
1.96
SMC
-10.6
2.01
10.79
PIS RNN
-1.9
0.81
2.06
PISRW-Grre
-0.87
0.31
0.92
CONCLUSION
Contributions. In this work, we proposed a new sampling algorithm, Path Integral Sampler, based on the connections between sampling and stochastic control. The control can drive particles from a simple initial distribution to a target density perfectly when the policy is optimal for an optimal control problem whose terminal cost depends the target distribution. Furthermore, provide calibration based on importance weights, ensuring sampling quality even with sub-optimal policies.
Limitations. Compared with most popular non-learnable MCMC algorithms, PIS requires training neural networks for the given distributions, which adds additional computational overhead, though this can be mitigated with amortization. Besides, the sampling quality of PIS in finite steps depends on the optimality of trained network. Improper choices of hyperparameters may lead to numerical issues and failure modes as discussed in inlapeendix G.T

---

Published as a conference paper at ICLR 2022
6 REPRODUCIBILITY STATEMENT
The detailed discussions on assumptions and proofs of theorems presented in the main paper are included infappendices A and C to The training settings and implementation tips of the algorithms are included in appendices F and G An implementation based on Py PyTorch (Paszke et al. 2019) of PIS can be found https: /ggttttt........ com/ qah- com/gsh-2///spi..
ACKNOWLEDGMENTS
The authors would like to thank the anonymous reviewers for useful comments, This work is par- tially supported by NSF ECCS-1942524 and NSF CCF-2008513.
REFERENCES
Michael Arbel, Alexander GDG Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. arXiv preprint arXivv2102.07501, 2021.
Dimitri P Bertsekas et al. Dynamic programming and optimal control: Vol. 1. Athena scientific
Belmont, 2000.
Ricky T. Q. Chen, Jens Behrmann, David Duvenaud, and Jörn-Henrik Jacobsen. Residual flows for
invertible generative modeling. In Advances in Neural Information Processing Systems, 2019.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differ- ential equations. arXiv preprint arXiv:1806.07366. 2018.
Yongxin Chen, Tryphon T Georgiou. and Michele Pavon. On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169(2):671-691 2016.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge Schrödinger bridge. SIAM Review, 63(2):249-313, 2021.
Nicolas Chopin and Omiros Papaspiliopoulos. An introduction to sequential Monte Carlo. Springer, 2020.
Rob Cornish, Anthony Caterini, George Deligiannidis, and Arnaud Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference on Machine Learning, pp.
PMLR, 2020.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1)333-332, 1991.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(3):411-433 2006.
Laurent Dinh, Jascha -ickssick,,,,,, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 16:-16005 00800008803, 2016.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C. Simmonett, Matthew P Harrigan, Chaya Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 13(7);e 1005659, 2017.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 1062496) 2011.
Ronen Eldan, Joseph Lehec, and Yair Shenfeld. Stability of the logarithmic Sobolev inequality via the Föllmer process. In Annales de I'Institut Henri Poincaré, Probabilités et Statistiques, volume 56, pp. 2253-2269 Institut Henri Poincaré, 2020.
Lawrence C Evans. Partial differential equations. Graduate studies in mathematics, 19(4):7, 1998.
10

---

Published as a conference paper at ICLR 2022
Charles W Fox and Stephen J Roberts. A tutorial variational Bayesian inference. intelligence review, 38(2):85-99) 2012.
Artificial
Christina Gao, Joshua Isaacson, and Claudius Krause. i-flow:
High-dimensional integration and sampling with normalizing flows. Machine Learning: Science and Technology, 1(4):045023, 2020a.
Ruiqi Gao. Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy-based models by diffusion recovery likelihood. arXiv preprint arXiv:2012.08) 125, 2020b.
lan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1.
MIT Press, 2016.
Jackson Gorham and Lester Mackey. Measuring sample quality with kernels. In International Conference on Machine Learning, pp. 1292-1301. PMLR, 2017.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367, 2018.
Matthew Hoffman, Pavel Sountsov, Joshua V Dillon, lan Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lizing bad geometry in Hamiltonian Monte Carlo using neural transport. arXiv preprint arXiv: 1at1::9003.0 907003004 2019.
Matthew D Hoffman and Andrew Gelman. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1):1593-16) 1623, 2014.
Jian Huang, Yuling Jiao, Lican Kang, Xu Liao, Jin Liu, and Yanyan Liu. Schrödinger-Föllmer sampler: Sampling without ergodicity. arXiv preprint arXiv:2 106. 10880, 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. arXiv preprint 1rXiv:1806.000 806.07572, 2018.
Patrick Kidger, James Foster, Xuechen Li, and Terry Lyons. Efficient and accurate gradients for neural sdes. arXiv preprint arXiv:2105.13492, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint MrXiv:1412.698 2014.
Diederik Kingma and Max Welling. Auto-encoding variational Bayes.
arviv preprint
arXiv:1312.6114, 2013.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/cxdb/mnistl, 1998.
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014.
Xuechen Li, Ting-Kam Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statistics, pp. 3870-3882, PMLR, 2020a.
Zengyi Li, Yubei Chen, and Friedrich T Sommer. neural network meme sampler that maximizes proposal entropy. arXiv preprint arXiv:2010.03582 2020b.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and acceler- ating particle-basee variational inference. In International Conference on Machine Learning, pp. 4082-4092. PMLR, 2019.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. arXiv preprint www.:1608.0471... 2016.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized Stein discrepancy for goodness-of tests. In International conference machine learning, pp. 276-284. PMLR, 2016.

---

Published as a conference paper at ICLR 2022
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
Shakir Mohamed and Balaji Lakshminarayannn. Learning in implicit generative models. arXiv preprint arXiv.1610003483 2016.
Jesper Møller, Anne Randi Syversveen, and Rasmus Plenge Waagepetersen. Log Gaussian cox processes. Scandinavian journal of statistics, 25(3):451-482, 1998.
Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2)125-139, 2001.
Kim Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojeiech Samek, Klaus-Robert Müller, and Pan Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, 101 (2):0)433344, 2020.
Erik Nijkamp. Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent non- persistent short-run meme toward energy-based model. arXiv preprint 19Xiv:1904.09770, 2019.
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.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65-84. Springer, 2003.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Michele Pavon. Stochastic control and nonequilibrium thermodynnaical systems. Applied Mathematics and Optimization, 191):1877-222 1989.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows In International Conference Machine Learning, pp. 1530-1538. PMLR, 2015.
John Salvatier, Thomas V Wiecki, and Christopher Fonnesbeck. Probabilistic programming in python using pymc3. PeerJ Computer Science, 2:e55, 2016.
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Jiaming Song, Shengjia Zhao and Stefano Ermon. A-nice-MC: Adversarial training for MCMC. arXiv preprint 170,00761, 2017.
Span Spanbauer, Cameron Freer, and Vikash Mansinghka. Deep involutive generative models for neural meme. arXiv preprint 15116 2020.
Kunal Talwar. Computational separations between sampling and optimization. arXiv preprint arXiv: 1119911020444 2019.
Matthew Tancik, Pratul P Srinivasan, Ben Mildenhall, Sara -riil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks leamn high frequency functions in low dimensional domains. arXiv:2006.10733 2020.
arXiv preprint
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A generalized path integral control ap- proach to reinforcement learning. The Journal of Machine Learning Research, 11:3137-3181, 2010.
Evangelos A Theodorou and Emanuel Todorov. Relative entropy and free energy dualities: Connec- tions to path integral and KL control. In 2012 icee 51 iece conference decision and control (cde), pp. 1466-1477 IEEE, 2012.
Sep Thijssen and HJ Kappen. Path integral control and state-dependent feedback. Physical Review E, 9133))032104 2015.
12

---

Published as a conference paper at ICLR 2022
Michalis Titsias and Petros Dellaportas. Gradient-basee adaptive Markov chain Monte Carlo. Advances in Neural Information Processing Systems, 32:15730-15739, 2019.
Surya T Tokdar and Robert E Kass, Importance sampling: a review. Wiley Interdisciplinary Reviews: Computational Statistics, 2(1):54-60, 2010.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pp. 3084-33 3114. PMLR, 2019,
Hao Wu, Jonas Köhler, and Frank Noc. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances Neural Information Processing Systems, Curran Associates, Inc., 2020.
volume 33, pp. 5933-5944,
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34, 2021.
Yan Zhou, Adam M Johansen, and John AD Aston. Toward automatic model comparison: an adaptive sequential Monte Carlo approach. Journal of Computational and Graphical Statistics, 25 (3):701-726, 2016.
13

---

Published as a conference paper at ICLR 2022
A PROOF OF THEOREM
Before proving our main theorem, we introduce the following important lemma.
Lemma 4.1 (Dai Pra
1991); I; Pavon 1989)). The transition density associated with optimal control
policy u* for eq (2) and eq (3) is
(y(y) (x,( (x,
(22)
0, (x(x))
where (x, y) denotes the transition probability from state x at time to state at time t.
Proof Theorem We denote the initial Dirac distribution by v òxx- Combining eq (7) and
Vx(x) = log (x) we obtain
1 dp O. log Eg /exp(- Vr (xr))) =-:0=-loggg log(] /
Vo(xo)
w10
Therefore, we can evaluate the KL divergence between Q" and Q (((x)) as
By Ry L/ | de + ((xT)] = Vo(xo) 0.
The first equality is based eq (13) Next, we show that xi p. The above equations imply
25r (X0, exp (-V (dy). ). It follows that
P E A] / 16, / exp(- = m(A).
B PROOF OF IMPORTANCE WEIGHTS
By definition
dQ' (T) (Q)
(
(23)
(Q) (T) dQ)(T)
Plugging eq (11) and eq (22) into the above we obtain
(" T
T 1r (x(xx), = exp( /"
w²(r) = exp( /
2 || di (.--u,dw--- log
||uell dt (0--u,dw:- (x()))
u(xT)
C PROOF OF THEOREM
C.1 LIPCHITZ CONDITON AND PRELIMINARY LEMMA
To case the burden of notations, we assume 5g. Our conclusion and proof can be generalized to other Dirac distributions easily. We start by assuming some conditions on the Lipschitz of optimal policy u*. It promises the existence of a unique strong solution with Xy ~ μ.. The conditions and properties are studied in Schrödinggr-Fööllee process, which dates back to the Schrödinger problem. For proof of existence of unique strong solution and detailed discussion, we refer the reader to Dai Pra 1991): Léonard (014); Eldan et al. (2020): Huang et al. (202I).
Condition 1.
[x)] < C1 + [xx1
(24)
and
| (x1) už, (xx) < C1([1x - xzll + lt1 - =z1²²)
(25)
Below the formal statement of Thhoremm
Theorem 5. Under Condition with sampling step size At, if |u; -u,l? < de for any then
(xx(" 0 O(y
(26)
14

---

Published as a conference paper at ICLR 2022
We introduce the following lemma before stating the proof for Theorem 2
Lemma 5.1. HHunng et al. 2021 Lemma A.2.) Assume Condition inequality holds for x, generated from u'.
holds, then the following
Bl|xx;2 x1,1 < 4C eCo (l2 4)) 2Col22 -4)² +2)2+22.44 411, 41 l2 E [0,
(27)
C.2 PROOF OF THEOREM 2
We denote by ((:T) the trajectory controlled by the optimal policy u", and by (xt, Xt,; the discrete time process with sub-optimal policy over discrete time {4x such that to = 0,'N = T and La 1--1 = At. The process (xx,) can be extended to continuous time setting as
... xtx)
X12 X12-1 + S
Uee- (Xtz-1) + dw
The key of our proof is the bound of ||xe xtx11 112 as
f
f pl
||x - x/, 112 = Xlu-1 +
Use (Nta- )ds + dw, - -ixia- +
(x + dw.]
12-1
? +
$ X1--1 x(en
(u((x))
U (Xex- ]dds
th
1 [x](x)) ur (x), ]ss
+ 2
x
x
he-1
x ll2 ( 11t de
<
X1
[|u5(x)) u -u.-, (Xte ds)
12-1
22
X1
xin_|
(/
[|,((()) - Uu (x1e-))]l 1)l
f4-1
(1+a)
X12-1 X l'
2
+
11 fK 1/-1
[|u](x)) uu (xx, ds}
((1 Xin-1 Xtx ||
+ (1+ = (tx - tu tx-1)(
|| - -u (x
(28)
where the first and second inequalities are based on the triangle inequality, the third inequality is based 2ab & aaa2 2 for any a 0, and the forth inequality is based on the Cauchy-Schwee
inequality.
In the following we bound the second term in eq (28) as
(x,)) ut- (x >/1
= =||u;(x;) -uj (Xte-.) + ui (Xte-i) - uce, (x
S (1 + (8) |1x(x;) - uin-, (X6e_1) I|| + (1 + }) | (Xee.>) - usa-1 (xte-.)
< 201(1 + =B]1|>> xxm_,l" 1| 2 Is t--ill (1+ 1 B) ||| piin (Xee-s) - Usy- (X(o_)) )
15

---

Published as a conference paper at ICLR 2022
where the first inequality uses 2ab 3a² + b2 for arbitrary 8 > 0 and the second one is based on eq (25) It follows that
[x((x)) 11 (xt=-1) ds
th 20 +,3)|12² xte-1 11 ds L 2C9(1 + 3)(s tx-1)ds
1
I (1 1 ( ule-s (x8e)11) /| ds.
14-1
Thus, for stepsize Al = 4x - 1-11, we establish
[[((x)) ue, (xxe- =)1 ds
20 + +13) [x; -xax 112 ds
16-1
1 } | (Xte-1) - ue (Xt. 1|² At.
Cf(1 + 8) A) + +(1+ +
(29)
Next we bound || - X1x 11 in (29) as
||=; X 1| s (1+ x xix .|| + (1+ + 1) | X (30)
1)
where the inequality is based (a 8)2 < (1 n)a? + (1 1) for arbitrary 1) 0.
Plugging eq (30) into eq (29) yields
[x{t) - ure (x1x 1)1 ds
< (1 + 3) Ilui (Xtr ,) - usy 18te }l( Al Al Cf(1 +
+ +20 + 3)(1 ) / Xlz-1
At + 20+20((1+000 8)( 7)
f/
AM
|x]
x(-_ ds.
(31)
Plugging eq (29) and (31) into eq (28) yields
LHS 5[1+ a 20 + - (1 + 3)(1 + - )t?t
Xin_1 - X(4-1
+ +2C + (1 () + n) 1"h
xx² - Xtx 112cs ds
12-1
(1 (1+ } ui (x/
u (x(e..)) ) /| 2 At + (1+ 1 +1 3)At3.
(32)
Invoking lemma 5.1
we obtain
E
x ds) < 4Co exp(2C)) + d)Ar + 2C,At + 2dA
Taking the expectation of lqq (32) in view the above and the assumption on control, we establish
B X1x x7, 11 $
C Xta..
-xi, 1 Ct,
16

---

Published as a conference paper at ICLR 2022
where C3 = 20 + () + 8)(1 1)) and
C c1=1+ [1 2 jddA + (1+ cf 3)At"
a
+ +2 + ) 3)(1 + (4C exp((Co) + d) + 20, + 2d
Finally, in view of the fact Xo = x and fixed step size Åt, we conclude that by the choice a =
C1 C1 = =
C - 1 - - - 6))
E Biixx] x
(33)
C3
where the last inequality is based on ignoring the high order terms, Cg² T Odd(At+))) +)).
- O(T) and
C2
S
C3
D PROOF OF THEOREM
The proof is a natural extension Corollary
Thijssen Kappen (0155)
We define random variable
uu ds + u'dw, + ((x)),
(1) = 0 2
(34)
and
(t) exp(- + (t))
(35)
(Tijssee & Kappen optimal control problem,
Lemma 5.2.
2015 Lemma 4) For any feasible control policy for stochastic
((T)))(x))- - ((t)) = =/[ (())d((x.((u; - u,,'dw..
(36)
Corollary 1.
|u, 2
(x) = EEga
=x] =Bqu [oop[eep--Vxxx
[s)[xx
x
(37)
2
Proof. This follows importance sampling with density ratio from eq (11)
Proof of Theorem 3: We denote the important weight by w"; note is random variable. It follows
that
w (Tp(-V((x))) 1-' ulll di + u'dw))
exp(-S"(())
2
(38)
EQ- (exp(- (x() 1 ull dt + u'dw))]
E (exp((
Dividing the LHS of ef (36) by (xx) we obtain
@ (TT):((x)) (()) (XT) (0 w" [o-
(39)
(xo)
Eo
Therefore, the variance of "h equals
((s)),
Equ I(w" EQ (10
EQ 1(
(x((x) (u [
" (x))0
Egu!
(u: - uz)(u, - u.)ds]
(x())
EEou 1//2/
(w (w" exp(S" (s))) u.,'(u, -u)ds]. (40)
17

---

Published as a conference paper at ICLR 2022
By Jensen's inequality Φ(8, x2)² (E (exp(([XXXXXX))))
E (exp(_25"" (s){x
Plugging the above inequality into eq (40) we reach the upper bound of variance
Eo~((w0 [eg[[w"]
Eg u.,'(u; u.)(w²)) Jdds.
(41)
In view of the fact Var(w") + +1= +1=Be- -===((w")) [() we arrive at
+ (< < E [[ /' Be Be-l(u; u,)(u,; u.)lds.
we consider the near optimal policy such that max,, |l,(()) - ux(x)]l S 7.. then follows that
2 -1
(42)
o- [2xn22]
PROOF OF THEOREM 4
We consider the KL divergence between trajectory distribution resulted from the policy u and the one from optimal policy u²:
DEL DxL(Q"(t)ll* (T) p(xT)
p (xT)
ETwey!
|rl| dt + u/dw, + (xT))
E
|ue 2 di n,dw, + (((T)) log %
0
E- (5) log Z]
2 O.
The last inequality is based on the fact Dxu(Q" 2 0 and the equality holds only when u2, pointing
0 E. 1 (T) + log Z). Therefore, we can estimate the normalization constant by
Z exp( -E (T)() Z exwppeppppccce ($)
(43)
Next provide unbiased estimation with sub-optimal policy based importance sampling as
2¹ (T) Q"(T)
(*(
Q"(()
= En~g' ( - log 2)].
The last equality is based on the fact Ex [-(G]] J, @*(d) 1. Hence, we obtain an
unbiased estimation of the normalization constant
Z Ey-gu (exp_
F EXPERIMENT DETAILS AND DISCUSSIONS
18

---

Published as a conference paper at ICLR 2022
1
PIS
AFT
SMC
NUTS
HMC
VI-NF
Figure Sampling performance on challenging 2D unnormalized density model with well- separated modes. Kernel density estimation plots compared with 2k samples. AFT and SMC use anncaling trick with 10 decreasing temperate levels and HMC kernel following (Arbel et al, 2021]. Even without annealing trick and resampling. Path Integral Sampler (PIS) generates visually indistinguishable samples from target density with 100 steps. PIS starts Xg from origin point while others start from standard Gaussian. The underlying distribution is chosen deliherately to distin- guish the performance of different methods. In particular, 100 steps are not sufficient for general MCMC converge the stationary distribution. We also note performance of compared methods can be further improved with tuning temperature scheduling, samples initialization. Our generic algorithm can explore more modes with similar initialization and less tuning parameters.
F.1 TRAINING TIME, MEMORY REQUIREMENTS AND SAMPLING EFFICIENCY
The PIS can be trained once before sampling and did not contribute to the runtime in sampling computation. Most MCMC methods do not have learnable parameters. However, PIS policy is trained once and used everywhere. Thus, the training time can be amortized when PIS policy is deployed in generating large number of samples. The training time of PIS highly depends on efficiency of training NeuralSDEs. One future direction is to investigate additional regularizations and structured SDEs speed up the training. The PIS algorithm is implemented in PyTorch (Paszke et al. (219). We use Adam optimizer (Kingma & Ba, 2014) in all experiments to learn optimal policy with learning rates 5 x 10-3 and other default hyperparameters. All experiments are trained with 30 epochs and 15000 points datasets. Loss in most experiments plateau after 3 epochs, some even epoch. Experiments are conducted using an NVIDIA A6000 GPU. Training one epoch on 2d example takes around 15 seconds for PIS-NN and 30 seconds for PIS- Grad, 1.6 minutes and 1.8 minutes respectively on Funnel (d = 10), and minutes and 9 minutes on LGCP (d = 1600). We note both PIS and AFT can be trained once and used everywhere. Therefore, the training time be amortized when the PIS policy is deployed in generating a large number of samples. We further compare empirical sampling time for various approaches in Tab
In the high dimensional data, thanks to the efficient adjoint SDE solver, we do not need to cache the whole computational graph and the required memory is approximately the cost associated with one forward and backward pass of (x) network. In our experiments, the total consumed memory is around 1.5GB for the toy and Funnel example, and around 5GB for the LGCP.
method AFT
Sampling Time
Training Time 711.2ms
352.2ms
SMC
110.8 ms
PIS-NN
16.8 ms
30.3 ms
PIS-Grad
34.3 ms
61.2
SNF
130.6 ms
256.1 ms
Table 4: Sampling and training efficiency comparison. For sampling, we measure the consumed time for generating 2000 particles with each method for 100 times in lower dimensional data, 2-D points datasets, and report average time for each method. We also include one batch training time for AFT and PIS.
F.2 NETWORK INITIALIZATION
19

---

Published as a conference paper at ICLR 2022
In most experiments in the paper, found PIS with the default setting, T = 1 and zero control initialization, gives reasonable performance. One exception is LGCP, where we found training with T = 1 sometimes suffers from numerical issues and gives NAN loss. We sweep T 1, 2, 5 and found T = 5 gives encouraging results. As we discussed in appendix G] the optimal policy u² depends T and large 7' not only results in large cover area but also "smoother" u². For neural network weight initialization, we use the zero initialization for the last layer of parameterized policy (x) and other weights follow default initialization PyTorch. We note there no guarantee that PIS can cover all modes initially with only the given unnormalized distribution. However, the initialization problem exists for most sampling algorithms and MCMC algorithms suffer longer mixing times compared with the ones with proper initialization (Chopin & Papaspilloppolos! PIS also suffers from improper initialization and we report some failure mode appendix
2020).
F.3
DETAILS FOR COMPARED METHODS
For all trained PIS and its variants, we use uniform 100 time- discretirrtiiiiionn steps for the SDEs, Gra- dient clipping with value is used. A Fourier feature augmentation (Tancik et al. (2020) is employed for time condition. Throughout all experiments, we use the same network with the modified first layer and last layer for different input data shape. In one pass (x)x), we augmented scalar t with Fourier feature to 128 dimension, followed by a linear layers to extract 64 dimension signal fea- ture. We use 2 layer linear layers to extract 64 dimension feature for æ and the concatenate æ, t features before feeding into another 3 linear layer with 64 hidden neurons to get (x). We found that the training stable with our simple MLP parametrization our experiments.
For HMC, we use 10 iterations of Hamiltonian Monte Carlo with 10 leapfrog steps per iterations, totaling 100 leapfrog steps. For NUTS, we set the maximum depth of the tree built as 5. Note that samples of HMC and NUTS used in our experiments are from separate trajectories instead of from one trajectory at different timestamps. We observed that the latter is more likely to generate samples that concentrate on one single mode. For SMC and AFT, we use 10 transitions with each transition using the same amount computation as HMC. The settings SMC and AFT follow the official implementation (Arbel et al. 2021 in the released codebase | In the Alanine Dipeptide experiments, for AFT we sweep adam learning rate Ir = 1 x 10 4 5 x 10 10-4, 11xx x 10 5 10 1x 10 and select 5 x 10-3, Other settings follow the default setup from the official codebase.
F.4 COMPARISON WITH SVGD ALGORITHM
In this section, we present some comparison between celebrated SVGD sampling algorithm (Lu
al., 2016) 2019) and PIS sampler. SVGD is based on collections of interacting particle, and we found its performance increase with more samples in a batch as it shown in Fig 5) Even with 5000 particles a bach, sample quality of SVGD is still worse than PIS-Grad. Mcanwhile, the kernel- based approach pays much more computation results the number of active particles increases as Tab 5 compared with PIS. In this perspective, PIS enjoy much better scalibility.
we show in
method
Time (mean + std) 19.3 ms + 1.2 ms
SVGD 100 particles
SVGD 1000 particles
29.4 ms + 1.8 ms
SVGD 5000 particles
434 ms + 2.43 ms
PIS-NN 5000 particles
223 us + 14.2 js
PIS- FIS-Grad 5000 particles
412 + 15.1 μs
Table 5: Consumed time for one iteration of SVGD and one step in PIS in 2D points example. Sampling with SVGD is less efficient compared with PIS.
20

---

Published as a conference paper at ICLR 2022
100 particles
1000 particles
5000 particles
Figure 5: Generated samples from SVGD (Liu & Wang 2016 with 100 steps. We generated samples with batch size 100, 1000, 5000. We find with more particles, samples generated are more closed to the ground truth data.
F.5 CHOICE OF PRIOR SDE:
We use general SDE
dx; f(t, x, ct + g(t, x;) dw.), Xo ~
(44)
instead of eq (2) for sampling. As discuss in Section 3.3 we prefer use a linear function for f. g.g promise closed-form μr. The choice of f, g encodes prior knowledge into dynamics without control and Q* is determined based on the prior -" Intuitively, the ideal @ should drive particles from (xT) that is close to jr. In PIS, Q". Thus training can be easier and faster if Q" and Q" are close since we use zero control as initialization for training policy. However, there
training objective is fit Q" with parameterized
general approach to choose f,g such that (x is close to μ and (x(x)) has closed form. In this work, we adopt the general form with would be interesting explore other prior dynamics or data-vart variant f. g in the future work, e.g., underdamped Langevin.
f 0,g
F.6 ESTIMATION OF NORMALIZATION CONSTANTS
As discussed in Chopin Papaspiliopoulos 2020). normalization constants estimation of SMC and its variants AFT can be achieved with incremental importance sampling weights.
In our experiments we treat HMC and NUTS as special cases of SMC with only two different tem- perature levels. One corresponds to a standard Gaussian distribution and the other one corresponds to the target density. Since the initial distribution v for SMC and NUTs is chosen as standard Gaus- sian, we can omit the MCMC steps for it and the total computation efforts required for the specific SMC are for the transitions HMC and NUTS.
For VI-NE, importance sampling
/ [(x)dx= / q(xx) ((xx) (x))
[(xx) (xx)
dx Eg
where q is the normalized distribution represented by normalizing flows, to provide an unbiased esti- mation of normalization constants. We use the ELBO inleq (18) for PIS and the unbiased estimation eq (19) for PISaw-
F.7 2 DIMENSIONAL RINGS EXAMPLE
The ring-shape density function
log A = min([[|x|| 1)2,(([x]) ([|x]| 3)) (|/x|| 5)2)
100
21

---

Consider the special case of gradient informed SDE, which can be viewed as PIS-Grad with spe- cific group parameters, Published as a conference paper at ICLR 2022
dx.- V log ((x))dt + 22dw..
This is exactly the Langevin dynamics used widely in sampling MacKay 2003). As special case of MCMC, Langevin sampling can generate high quality samples given large enough time interval (MacKay 2003). From this perspective, PIS- PIS-Grad can be viewed modulated Langevin dynamics that is adjusted and represented by neural networks.
F.8 BENCHMARKING DATASETS
MG(d-2)
Funnel(d=10)
LGCP(d=1600)
A
B
B
A
AFT-10
-0.509
0.24
0.562
-0.249
0.0758
0.261
-3.08
1.59
3.46
SMC-103
-0.362
0.293
0.466
-0.338
0.136
0.364
-440
14.7
441
AFT-2 x 10a
-0.371
0.477
0.604
-0.249
0.0758
0.261
-1.23
0.826
1.48
SMC-2 102
-0.398
0.198
0.444
-0.338
0.136
0.364
-197
5.21
197
AFT-3 X 10a
-0,316
0.365
0.483
-0.281
0.0839
0.293
-1.05
0.514
1.17
SMC-3 10
-0.137
0.62
0.635
-0.323
0.064
0.329
-109
5.58
109
AFT-5 10a
-0.194
0.319
0.373
-0,253
0.0397
0.256
-0.949
0.439
1.05
SMC-5 1
-0.129
0.246
0.278
-0.298
0.0564
0.303
-37.5
5.04
37.8
AFT-10
-0.03
0,515
0.515
-0.194
0.0554
0.202
-0.277
0.356
0.901
SMC-104
-0.171
0.446
0.477
-0.239
0.0412
0.243
-6.47
1.95
6.76
PIS-10
-0.021
0.03
0.037
-0.008
9e-3
0.012
-1.94
0.91
2.14
Table 6: Long-run MCMC mode separated mixture of Gaussian (MG), Funel distribution and Log Gaussian Cox Process (LGCP) for estimating log normalization constants. The suffix denotes the total number of discrete-time steps for each method, which equals the number of layers multiply steps per layer. We experiments 10, 20, 30, 50, 100 layers for annealing and 100 leapfrog steps per layer. As the number of steps increases, the performance of AFT and SMC gradually improves. PIS denotes the PIS (ISpw-Grade B and S stand for estimation bias and standard deviation among 100 runs and A B² 52.
For mixture of Gaussian, choose nine centers the grid {-5, (-5,0,5), and each Gaussian has variance 0.3. The small variance is selected deliberately to distinguish the performance of the different methods. We 2000 samples for estimating the log normalization constant Z. We use the standard MLP network to parameterize the control drift u, (x), where the time signal is augmented by Fourier feature using 64 different frequencies. We use 2 layer (64 hidden neurons in each layer) MLP to extract features from the augmented time signal and x separately, and another layer MI.P to map the summation features the policy command. We note that all these methods for comparison, including HMC, NUTS, SMC, AFT, can reach reasonably good results given large enough iterations. However, with small finite number steps, PIS achieves the best performance. We include more results for long-run MCMC methods in Tab 6
the experiment with Funnel distribution, Arbel al. (2021) suggests to use slice sampler kernel for AFT and SMC, which includes 1000 steps of slice sampling per temperature. In Tab still use HMC for comparing performance with the same number integral steps. We also include the results with slice sampler in Tab 7 We use 6000 particles for the estimation of log normalization constants. The network architecture of PIS is exactly the same as that in the experiments with mixture of Gaussian.
In the example with Cox process, the covariance K is chosen
||u - v/]
K(u, v) = 1.91 exp(
MB
22

---

Published as a conference paper at ICLR 2022
Funnel(d=10)
B
S
A
AFT-10
0.128
0.376
0.398
SMC-10
-0.193
0.067
0.204
AFT-20
0.0134
0.173
0.174
SMC-20
-0.113
0.0878
0.143
AFT-30
0.074
0,309
0.318
SMC-30
-0.006
0.188
0.188
PrSaw-Grad
-0,008
0.009
0.012
Table 7: AFT and SMC with slice sampler kernel. The suffix denotes the number of temperature levels for annealing. 1000 slicing sampling steps are used for each temperature. Though there no annealing and only 100 steps are used, the performance of PIS is competitive.
and the mean vector equals log (126) o² and a 1/M². We note that this setting follows Arbel et al. (2021); y); Møller et al. (998). Totally 2000 samples are used to evaluate the log normalization constant. We treat the mean of estimation results from 100 repetitions of SMC with 1000 tempera- tures as ground truth normalization constants. In this experiment, we found clipping gradient from target density function help stabilize and speed up the training of PIS-Grad. This example the most challenging task among the three. One major reason is the high dimensionality of the task; the PIS needs to find optimal policy u (t, Rd) + Rd in high dimensional space. In addition, there no prior information that can be used to shrink the search space, which makes the training of PIS with MLP more difficult. use 2000 particles for estimation of log normalization constants. We also include more experiment results in Tab 8
LGCP(d=1600)
B
A
AFT-10
-0.827
0.356
0.901
PIS nw-Grad-1 x 10
-1.94
0.91
2.14
PIS nw-GGrd-5 x 102
-1.25
0.57
1.373
PIS RW- Grad-10 x 102
-0.832
0.214
0.859
Table 8: PIS with large number of integral step. The suffix number is the total integral steps. For AFT-104, we use 100 annealing layers and run 100 leapfrog steps per each annealing layer.
F.9 ALANINE DIPEPTIDE
The setup for target density distribution and the comparison method are adopted from Wu et al. (2020). Following Noč et al. (2019), invertible transformation between Cartesian coordinates and the internal coordinates is deployed before output the final samples. Then we normalize the coordinates by removing means and dividing them by the standard deviation of train data. To setup the target distribution, we simulate Alanine dipeptide al. 2017 Total 10a atoms data points are generated as training data. Situation parameters, including time-step and temperature setting are the same as Wu et all (2020). We refer the reader to official codebase for more details setting target density function
vacuum using OpenMMTools (Eastman)
Following the setup in Wu et al. (2020), we use unweighted samples to compute the metrics. KL divergence of VI-NF on μ calculated based on ELBO instead of importance sampling as in nor- malizing constants tasks. We use Metropolis random walk MCMC block for SMC, SNF and AFT and RealNVP blocks for SNF and AFT (Dinh et 2016). For fair comparison, we use PIS-NN instead of PIS-Grad since none of the approaches in this example uses the gradient information. We note that SNF originally trained with maximizing data likelihood where some empirical samples
23

---

Published as a conference paper at ICLR 2022
assumed to be available. We modify the training objective function by reversing the original KL divergence as in eq (1)
F.10
MORE DETAILS ON SAMPLING IN VARIATIONAL AUTOENCODER LATENT SPACE
8 8 0
9
4
2
8 2
8 8
9
4
2 8
2
Figure 6: Origin data images and their reconstructions from trained vanilla VAE. It can be seen that reconstruction images smoother compared with the original images.
We use a vanilla VAE architecture to train on binary MNIST data. The encoder uses a standard layer MLP networks with 1024 hidden neurons, and maps image to the mean and standard devi- ation of 50 dimension diagonal Normal distribution. The decoder employs 3 layer ML.P networks decode images from latent states. ReLU nonlinearity is used for hidden layers. For training, we use the Adam optimizer with learning rate 5 x 10-4, batch size 128. With reparameterization trick (Kingma & Welling 2013) and closed-form KL divergence between approximated normal dis- tribution and standard normalization distribution, we train networks for totally 100 epochs. We show performance of vanilla Fig 6
We parameterize distribution of decoder pe (x) as
log z) = lo(p(x(Do(z)) = x log Dz(z) + (1 log(1 Do(2).
For PIS, we use the same network and training protocol as that in the experiment for mixture of Gaussian and Funnel distributions. We also use gradient clip to prevent the magnitude of control drift from being too large.
G TECHNICAL DETAILS AND FAILURE MODES
G.1 TIPS
Here provide list of observations and failure cases we encountered when we trained PIS. We found such evidences through some experiments, though there is no way we are certain the following claims are correct and general for different target densities.
We notice that smaller T may result in control strategy with large Lipehastiz constants, which is within expectation since large control is required to drive partieles to destina- tion with less amount of time. It is reported that more difficult approximate large Lipchastiz functions with neural networks (Jacot et al., 2018 Tancik et al. (2020). We thus recommend increase T' or constraint the magnitude of stablize training when encountering numeric issue or when results are not satisfactory.
We found batch normalization can help stablize and speed up training, and the choice of nonlinear activation (ReLU and its variants) does not make much difference.
We also notice that if the control (x))) has large Lipchastiz constants in time dimension, the discretized error would also increase. For calculating the weights based on path integral, we suggest decrease time stepsize and increase N when the number of integral steps is small and discretization error is high.
We obtained more stable and smaller training loss when training with Tweedie's for- mula (Efron 2011), but we found obvious improvements testing the trained sampler or estimating normalization constants.
Regardless the accuracy and memory advantages of Reversible Heun claimed by torchsde (Li al. 2020a Kidger al. 2021) found this integration approach is
24

---

Published as a conference paper at ICLR 2022
Figure 7: Generated 5000 uncurated samples with 0.1, 0.2, 0.4, 0.6, 0.6,0.8,,8888. 0.8, 1.0, 2.0, 4.0, 5.0. PIS with small T may miss some modes.
less stable compared with simple Euler integration without adjoint and results in numer- ical issues occasionally. We empirically found that methods without adjoint more stable and lower loss compared with adjoint ones, even in low dimensional data like 2d points. We use Muler-Maruyama without adjoint for low dimension data and Reversible Heun method (Kidger et al. 2021) in datasets when required memory is overwhelming to cache the whole computational graph. We recommend readers
Ito Euler integra- tion when memory permits or conduct training with small At. The more steps in PIS and smaller At not only reduce policy discretization error Theorem 2] but also reduce the Fuler-Maruyama SDE integration errors.
G.2 FAILURE MODES
In this section, we show some failure modes PIS. With an improper initialization and an extremely small 7, PIS suffers from missing mode. The failure can be resulted from following factors. First, the untrained and initial p(xr) = (0, 7T) may be far away from the target the distribution modes. Thus it is extremely difficult for training PIS to cover region of high probability mass under u. Second, small results in policies with large Lipchastic constants that are challenging to train networks
discussed in appendix G.I Third, the policy with limited representation power tends to miss modes due to the proposed KL training scheme. We show some uncurated samples trained with different T and other failure cases.
Thoughout our experiments, we also find that large T leads to more stable training and PIS sampler covers more modes. However large T with large öt deteriorates sample quality due to discretization error but it can be eased with increasing number steps.
We note PIS is not perfect sampler and failure modes exists for all compared methods, experiments results in Section we reported are based on uncurated experiments.
25

---

Published as a conference paper at ICLR 2022
(a) T 10, A = 0.1, 100 (b) T = 10, At = 0.01, N = 10
Figure 8: Large T with large öt deteriorates sample quality due discretization error in Fig 8a but it can be eased with increasing number of steps in Fig 8b
Figure 9: Another failure case with T 2, At 0.02, N 100 due to randomness of training networks.
26


Reviews:
[{"review_id": "zyOZIKcIPHg", "reviewer": "Reviewer_dEAt", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper is excellently written and great care is taken\nto highlight exactly what properties of each formalism are\nused to get the nice closed-form equations. The work is\nhighly novel and one that would be greatly valued by the\ncommunity. The proofs are clear and easy to follow as well.\n\nThe experiments are rigorous and persuasive. One thing I\ndon't understand is why was the work not compared against\nneural density estimators? There is a rich literature of\nflow models that is referenced in the paper but aside from\nSNF not really compared against. Some of it even uses\nneural networks that backpropagate through a SDE solver!\n", "summary_of_the_paper": "This paper shows that sampling can be treated as a stochastic\noptimal control program and introduces a novel sampling algorithm\nthat works by simulating a stochastic differential equation under\noptimal control with appropriately chosen dynamics and cost function.\n", "summary_of_the_review": "This work ntroduces a novel sampling algorithm with lots of interesting\nproperties. This paper makes a solid contribution with the strengths and\nlimitations of the approach very clearly articulated.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SNINq3da6FI", "reviewer": "Reviewer_ZcgD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strengths: clear exposition, theoretical guarantees provided, informative experiments and ablations\n\nWeaknesses:\n0) It would be informative to see a comparison to Stein-Variational Gradient Descent (SVGD) or other ParVI methods [1]\n1) The paper says \"To calculate gradients, we rely on backpropagation through trajectory.\" Are there any issues associated with that? E.g., exploding/vanishing gradient. How long can trajectories be? How much memory/time is required?\n2) In Algorithm 1, specify what is input and output of the algorithm\n3) In Algorithm 2, it is a bit confusing that 'w' is used as the argument of the exponential, which is different from Eq. (17), where it is used as the exponential itself\n4) In Theorem 2, Condition 1 is only defined in the Appendix. Either state the condition before the theorem or expand it in the statement of the theorem\n5) The paper says \"training consumption is not included in the comparison\". Could the authors state at least the order of magnitude of how much time the proposed method takes compared to the baselines?\n6) In Table 1, highlighting PIS_RW-GT may be seen as a distraction, since it is using the ground truth. Therefore, perhaps one could either separate that line or set those numbers in italic and use bold on PIS_RW-Grad for MG(d=2) columns\n7) \"more clearer\" -> \"clearer\"; \"as it in a toy example\" -> \"as in a toy example\"\n8) What are the limitations? Adding a paragraph on drawbacks and limitations would be helpful.\n\n[1] Liu, C., Zhuo, J., Cheng, P., Zhang, R., & Zhu, J. (2019, May). Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning (pp. 4082-4092). PMLR.", "summary_of_the_paper": "The paper proposes an algorithm called Path Integral Sampler (PIS) for sampling from unnormalized distributions by parameterizing the control policy in the Schrödinger bridge problem with neural networks and supplying it with the gradient information from the target density. Performance guarantees are provided, and experiments on several illustrative toy problems as well as on sampling molecular structures and latent space of VAEs are conducted. The method is shown to produce high-quality samples from the posterior, demonstrating state-of-the-art performance.", "summary_of_the_review": "The paper is overall good. The method is based on taking the established framework of inference as control via Schrödinger bridge problem and utilizing NNs as a parameterization for the policy. Exposition is clear, convergence is characterized, the properties of the algorithm are highlighted on relatively low-dimensional problems.\n\n====\n\nAfter rebuttal\n\nI thank the authors for addressing my comments. I trust that the authors will include the promised changes and additions regarding the training time and the comparison to ParVI into the final version of the paper. I maintain my score \"8: accept, good paper\".", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HTyYC8Eiy_g", "reviewer": "Reviewer_1wV8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The idea of using stochastic control for sampling is very interesting and the empirical performance seems improved. I have the below questions.\n\n1.\tCompared with previous MCMC methods, one benefit claimed in the paper is that PIS can generate samples with fewer steps. It is not clear to me why this is the case. T seems a hyperparameter in the algorithm and there are no theoretical or empirical supports to show that T is less than the number of steps in MCMC methods. Besides, what does “until converged” mean specifically in Algorithm 1? How do the authors check convergence? From Figures 1 and 2, the main advantage of PIS seems the ability of sampling from multimodal distributions. Why is it the case?\n\n2.\tIn Theorem 2, the Wasserstein distance between the sampler and the target density increases as T increases which seems weird to me. Should the sampler be closer to the target distribution as the number of SDE time-discretization steps increases?\n\n3.\tSince the authors use a neural network to parameterize control policy, how does the expressiveness of neural networks affect the performance? When do the users need to include the importance sampling step? There seems no discussion about how to choose the neural networks.\n\n4.\tThe empirical comparison between PIS and previous methods generally lacks interpretation. Surely the evaluation metric of PIS seems better, but it will be much helpful to provide an explanation on why it is the case. For example, is it because the target distributions are multimodal or PIS converges faster? This kind of interpretation of empirical results sheds more light on when and how to use the proposed method for users. \n\n5.\tThis paper seems closely related to [Deep Generative Learning via Schrodinger Bridge, ICML 2021] which also uses Schrodinger Bridge algorithm for sampling. A discussion with this paper would be helpful.\n\n", "summary_of_the_paper": "The authors propose a Path Integral Sampler based on Schrodinger bridge. They turn the sampling problem into a stochastic control problem and use a neural network to parameterize the control policy. In order to find the optimal control policy, several approximation techniques are introduced. The authors further propose an important sampling scheme to handle the case of a suboptimal policy. They provide convergence guarantees under Wasserstein distance and demonstrate the proposed method on synthetic datasets, alanine dipeptide and binary MNIST.", "summary_of_the_review": "In summary I think the proposed method is interesting and technically sound. But the benefits over previous methods are not clear both in theory and practice.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2WmXr7jPA4M", "reviewer": "Reviewer_d7Mk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The theoretical result giving the KL expression is known as acknowledged by the authors. However, to the best of my knowledge, this has never been exploited to solve sampling problems. Instead it has been proposed to instead directly approximate the Schrodinger-F\\\"ollmer drift. It is an interesting idea to attempt to exploit the KL expression to obtain a practical sampling algorithm. \nThe authors propose to parameterize the control term by some neural networks which involves the gradient of the log-target; see e.g. equation (15). The gradient of the loss is trained using recent neural SDEs techniques. The method is then demonstrated on various examples where it appears to outperform solid baselines.\n\nI find the idea proposed in this paper interesting but I also think that a significant part of the paper was spent discussing unnecessary material and that important details were omitted. The limitations of the methodology should also be better spelled out. \n\nFirst, the authors discussed a general class of controlled diffusions, review HJB for such models and Fleming logarithmic transform. It is shown that the the logarithm of the value function satisfies a simple linear PDE which can be solved using Feynman-Kac. However, these results are never used in the paper. The paper only deals with scenarios where we initialize the diffusion at the origin and consider $dX_t=u_t(X_t)dt+dW_t$ as one needs to know the marginal density of the uncontrolled process at time T. So I think it's really unnecessary to spend so much time on the general case, you might want to include this material in the supplementary and instead concentrates on more important algorithmic details. Similarly please simplify the presentation of Algorithms 1 and 2.\n\nA key component of the paper is the parameterization $u_t(x)=NN_1(t)-NN_2(t)\\nabla \\log \\mu(x)$ in eq. (15) of the control policy using neural networks. This should be better motivated. The use of the gradient of the log-target seems essential for the method to work well but I find it hard to believe it can solve all issues. Clearly if one has a multimodal target, gradient information will not be sufficient to explore all the modes; i.e. Langevin does not mix well in multimodal scenarios. Obviously one has also another component $NN_1(t,x)$ in the control which could deal with scenarios. However, if we are interested in sampling a high dimensional multimodal target, how can the initial control ever visit the region of high probability mass under $\\mu$ and learn anything. The forward KL is the criterion been minimized so I would expect you will miss modes in some examples. In this respect, it would be good to discuss how the method is initialized (i.e. what is $u_t(x)$ at the first iteration, how is $T$ selected? etc). \n\nThe failure modes of the proposed method are never mentioned, was it necessary to train multiple times the model with various initializations to obtain good results etc, did you find examples where the method fails miserably? (this would be fine, I don't expect any method to solve all sampling problems). I think these are really important points which are unfortunately not discussed. \n\nSimilarly, I was somewhat frustrated by the lack of details on the gradient computation. Instead of just citing a paper, I would encourage the authors to discuss these details (instead of having a fairly long introduction on stochastic control). It would also be good to explain if there are any gradient issue (gradient vanishing/exploding) as the number of discretization steps increases (you seem to be using 100 in simulations, are there numerical problems when you increase the number of steps) etc.\n\nOther questions:\n\n- As in the AFT paper, the training time is not taking into account when evaluating the performance of the algorithm. \nHow large is the training time? Is it better/worse than AFT? If the training time is included, is the path integral sampler competitive to an SMC algorithm using a large number of intermediate temperatures?\n\n- How do you select T and $\\delta t$? If you initialize using $u_t(x)=0$, then at time $T$ you will have a state $X_T \\sim N(0,T*I)$ but if the mode of the target is really in the tails of this distribution, would not this be a problem?\n\n- Shouldn't the weight of interest be  $dQ^{\\star}(\\tau)/dQ^u(\\tau)$ and not its inverse given in (17). You are sampling from $Q^u(d\\tau)$ and are targetting $Q^{d\\star}(\\tau)$. The ESS is $\\mathbb{E}_{Q^u}[{dQ^{\\star}(\\tau)/dQ^u(\\tau)}^2]$ and not what is written in Theorem 3 given in eq. (17).\n\n- It is claimed that a single step of the sampler in Algorithm 2 is similar to a single step of leap frog step of HMC. As the control is computed by a neural network, this really depends on the size of the network. I could not find this point discussed in the paper. \n\n- There is quite scant information on the tuning of SMC and AFT (what flow was used? how was HMC tuned? was the Adam step size reasonable?)\n\n- 10 temperatures as a default for SMC is certainly not standard practice. In practice, the temperature schedule is also typically selected adaptively.\n\n- In the 2D example (Figure 1), the scale of the target is set to be larger than the initial distribution so no surprise SMC will struggle. I find the example quite misleading. The comment \"The small variance is selected deliberately to distinguish the performance of the different method\" in F4 should be included earlier. \n\n- I couldn't find any details on the neural nets used in the experiment. Did you use the same architecture for all examples?\n\nMinor comments\n\n- SMC with annealing trick is due to Del Moral et al. (2006) not Chopin & Papaspiliopoulos (2020).\n\n- page 13: Dai Pra (1991); ?)\n\n- Table 1: How do you estimate the log normalizing constant for LGCP?\n\n- Table 2: caption indicates KL-divergences but you cannot compute it as you don't have the normalizing constant.\n\nPost-rebuttal comment: The authors didn't submit a revised version of their manuscript on time so I revised my score accordingly. I was eventually asked by the AC to consider the revised version that was sent after the deadline. As some of my comments have been taken into account, I have re-adjusted my score. In particular, I do appreciate that some of the limitations of PIS are now detailed. However, a few important points need to be addressed seriously by the authors in their final version of the paper. \n\n- \"Our approach avoid long mixing time theoretically and is more efficient\": this claim is not justified and should be removed from the paper. You have no guarantee to be able to compute anything close to the optimal control and, even if you were, the Lipschitz constant of the control might be very high which would require a very fine discretization of the SDEs for numerical stability and hence a long run time. This type of bold claims is not necessary.\n\n- Figure 1 is an unncessary ``clickbait\" figure which, if kept in the manuscript, will just convince any AIS/SMC expert that you haven't run serious comparisons. The caption mentions that \"The underlying distribution is chosen deliberately to distinguish performance of different methods\". This is not quite correct, the problem is not the \"underlying\" (i.e. target) distribution but the fact that, instead of picking a diffuse initial distribution compared to the target as recommended in the SMC iterature, you pick a standard normal distribution whereas the modes of the target are at (-5,0), (-5,-5) etc. If you were using an initial distribution N(0,\\sigma^2 I) with sigma=5, I am sure it'd work very well. Also the selection of 10 intermediate distributions is just arbitrary and, contrary to what the authors state, is not a recommendation made in [Arbel, 2021]. I recommend the authors to read and cite Zhou, Yan and Johansen, Adam M and Aston, John AD, Toward automatic model comparison: an adaptive sequential Monte Carlo approach, JCGS 2016. Ideally you'd run simulations with this method.\n\n- I still believe that having a general drift and volatility in eq (2) is just unncessary when you use f=0 and g=1. Please simplify the presentation. Algorithm 1 and Algorithm 2 could be similarly simplified. It's not like you can use very general f and g in the first place as you need to know the marginal distribution of X_T for the uncontrolled SDE.\n\n- Please cite standard references in importance sampling instead of Au & Beck, 2003. Also please do cite Neal, Annealed Importance Sampling, 2001.\n\n- Figure 3 is another unnecessary \"our method works best\" figure which tells us nothing given no detail is provided (i could not find any) in the main text.\n\n- You mentioned it in your rebuttal that \"We admit a better temperature schedule may improve the performance of SMC\": this is certainly true and should be stated explicitly in the paper and you should include here the reference to  Zhou, Yan and Johansen, Adam M and Aston, John AD  which addresses this problem. This is followed by a claim that you follow \"the setup for SMC and AFT recommended by AFT and its official codebase\". I checked the code (which is a codebase for SMC and not the official codebase) and the paper [Arbel, 2021]. You should use obviously the same settings as in this paper when you present similar examples but, if 10 temperatures is written in a code, this is not a recommendation for all examples. I think including more thorough comparisons (i.e. more temperatures, optimized schedules) would make the paper better.  You compute the \"ground thruth\" for the LGCP using SMC (something which should be added in the paper too and not only mentioned in the rebuttal), so it must work well sometimes.\n\n-Section F.1 didn't address my concerns and has to be rewritten/completed. \n\"We note PIS training algorithm is different from AFT, where the AFT algorithm requires optimization during generating samples and\nretraining the network for the next batch of samples. In contrast, the PIS control policy is trained once and used everywhere.\"\nOnce you run AFT to learn the flows, then you can rerun it with the trained flows (this is what the authors appear to be doing from my reading - \"We concentrate our empirical value evaluation on the learnt flow, which is equivalent to using the test set particles.\" ), so it's not different from PIS. \nI want to see a table with two columns, one for training where the training times for all methods are given (what is diplayed currently is the execution time for PIS and the training time for AFT) and another column where the execution times for all methods are presented  (AFT will correspond here to SMC+learned flows).\n\n\n\n", "summary_of_the_paper": "The paper presents a control approach to sampling. It is proposed to use a controlled diffusion initialized at time t=0 to sample at a given time t=T from a given unnormalized target distribution. This is achieved by minimizing a forward KL between two suitable diffusions. The resulting expression is simple as it is given by the integrated squared control plus a final cost term. The authors propose to parameterize the control by a neural network and minimize the KL using stochastic gradient techniques. In experiments, this novel method appears to outperform some recent alternatives.", "summary_of_the_review": "Overall, this is an interesting approach and the authors have demonstrated empirically on a variety of models that it appears to perform well. \n\nTwo points should be addressed:\n\n- I believe that it would be good to give much more details about neural network architectures, initialization, gradient computation etc. \n\n- The limitations of the methodology should be better spelled out. The authors should for example illustrate the failure modes of the algorithm.  \n\nMore generally, it remains unclear whether any of the recent proposed Monte Carlo methods relying on training flows/neural networks by minimizing KL/maximizing ELBO are competitive with AIS/SMC when taking into account the training time.  \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a control-based approach to sampling. All of the reviewers found the idea interesting. There were serious concerns by some of the reviewers regarding how the paper positioned itself relative to the literature, how it designed baselines for experiments, and how it compared itself to existing methods. There was vigorous rebuttal phase. The authors submitted a slightly late revision based on a procedural misunderstanding, and I decided to incorporate their late revision.

Based on the revision, the majority of reviewers felt that the paper was at least above the bar for acceptance and some of the more positive reviewers stood strongly by the paper. I believe that this paper is of value to the community, so I will recommend that it is accepted, but I want to be very clear about something: the authors **must** incorporate the late revision as the basis of their camera ready and I **strongly recommend** that they address the concerns of reviewer d7Mk, including but not limited to:

- "Our approach avoid long mixing time theoretically and is more efficient": this claim is too strong.
- Figure 1 is not particularly informative and the authors should reconsider it.
- The presentation of Eq (2), Algorithm 1, and Algorithm 2 should be simplified. 
- Section F.1 should incorporate the comments of Reviewer d7Mk.
- Citing standard references mentioned by Reviewer d7Mk.

The reason I highlight these recommendations is that I believe they will greatly improve the quality, longevity, and impact of this paper. Slightly overselling ideas feels like a good strategy, but it is a bad long-term strategy. I believe addressing these points is in the interest of the authors.
Decision: Accept (Poster)