A STOCHASTIC CONTROL APPROACH FOR SAMPLING
Year: 2021

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ödinger 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ödinger 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 time-discretization. We experimentally demonstrate the advantages of PIS compared with other startof-the-art sampling methods on a variety of tasks.

heading: 1 INTRODUCTION
text: We are interested in drawing samples from a target density µ̂ = Zµ known up to a normalizing constant Z. Although it has been widely studied in machine learning and statistics, generating asymptotically unbiased samples from such unnormalized distribution can still be challenging (Talwar, 2019). In practice, variational inference (VI) and Monte Carlo (MC) methods are two popular frameworks 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 µ̂ to identify good q 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 q, 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 et al., 2020a; Nicoli et 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 to be bijective and its Jacobian needs to be easy-to-compute (Wu et al., 2020; Cornish et al., 2020; Grathwohl et al., 2018).
Most MC methods generate samples by iteratively simulating a 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 (Chopin & Papaspiliopoulos, 2020). 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 are some recent works exploring the possibility of incorporating neural networks to improve MCMC (Spanbauer et al., 2020; Li et al., 2020b). However, evaluating existing MCMC empirically, not to say designing an objective loss 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 propos-
als (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 is built on the Schrödinger bridge problem (Pavon, 1989; Dai Pra, 1991; Léonard, 2014; Chen et al., 2021) whose original goal was to infer the most likely evolution of a diffusion given its marginal distributions at two time points. With a proper prior diffusion model, this Schrödinger bridge framework can be adopted for the sampling task. Moreover, it can be reformulated as a stochastic control problem (Chen et al., 2016) whose terminal cost depends on the target density µ̂ so that the diffusion under optimal control has terminal distribution µ̂. We model the control policy with a network and develop a method to 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 models have weaker expressive power and flexibility compared with implicit models, both theoretically and empirically (Cornish et al., 2020; Chen et al., 2019; Kingma & Welling, 2013; Mohamed & Lakshminarayanan, 2016). 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 show guaranteed sampling quality in terms of Wasserstein distance from the target density for any given sub-optimal policy.
Our algorithm is based on Tzen & Raginsky (2019), where the authors establish the connections between 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 Tzen & Raginsky (2019). Another closely related work is Wu et al. (2020); Arbel et al. (2021), which extends Sequential Monte Carlo (SMC) by combining deterministic normalizing flow blocks with stochastic MCMC blocks. To be able to evaluate the importance weights efficiently, MCMC blocks need to be chosen based on annealed target distributions carefully. 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, annealing scheduling. An illustration of the advantages of PIS is presented in Fig 1. We summarize our contributions as follows.
1. We propose Path Integral Sampler (PIS), a generic sampler that generates samples through simulating a target-dependent SDE which can be trained with free-form architecture 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 is 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 distributions. 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-the-art sampling performance in several sampling tasks.

heading: 2 SAMPLING AND STOCHASTIC CONTROL PROBLEMS
text: We begin with a brief introduction to the sampling problem and stochastic control problem. Throughout, we denote by τ = {xt, 0 ≤ t ≤ T} a continuous-time stochastic trajectory.

heading: 2.1 SAMPLING PROBLEM
text: We are interested in drawing samples from a target distribution µ(x) = µ̂(x)/Z in Rd where Z is the normalization constant. Many sampling algorithms rely on constructing a stochastic process that drives the random particles from an initial distribution ν that is easy to sample from, to the target distribution µ.
In the variational inference framework, one seeks to construct a parameterized stochastic process to achieve this goal. Denote by Ω = C([0, T ];Rd) the path space consisting of all possible trajectories and by P the measure over Ω induced by a stochastic process with terminal distribution µ at time T . Let Q be the measure induced by a parameterized stochastic and denote its marginal distribution at
T by µQ. Then, by the data processing inequality, the Kullback-Leibler divergence (KL) between marginal distributions µQ and µ can be bounded by
DKL(µ Q‖µ) ≤ DKL(Q‖P) := ∫ Ω dQ log dQ dP . (1)
Thus, DKL(Q‖P) serves as a performance metric for the sampler, and a small DKL(Q‖P) value corresponds to a good sampler.

heading: 2.2 STOCHASTIC CONTROL
text: Consider a model characterized by the stochastic differential equation (SDE) (Särkkä & Solin, 2019)
dxt = f(t,xt)dt+ g(t,xt)(utdt+ dwt), x0 ∼ ν (2) where the drift term f : Rd → Rd is a vector-valued function, and the diffusion coeffient g is a matrix-valued function, xt,ut denote state and control input respectively, and wt denotes standard Brownian motion. In stochastic control, the goal is to find an feedback control strategy that minimizes a certain given cost function.
The standard stochastic control problem can be associated with any cost. In this paper, we only consider the cost of the form
E[ ∫ T
0
1 2 ‖ut‖2 dt+ Ψ(xT ) | x0 ∼ ν], (3)
where Ψ represents the terminal cost. The corresponding optimal control problem can be solved via dynamic programming (Bertsekas et al., 2000), which amounts to solving the Hamilton-JacobiBellman (HJB) equation (Evans, 1998)
∂Vt ∂t + f · ∇Vt − 1 2 ∇V ′t gg′∇Vt + 1 2 Tr(gg′∇2Vt) = 0, VT (·) = Ψ(·). (4)
The space-time function Vt(x) is known as cost-to-go function or value function. The optimal policy can be inferred from Vt(x) as (Pavon, 1989)
u∗t (x) = −g(t,x)′∇Vt(x). (5)

heading: 3 PATH INTEGRAL SAMPLER
text: It turns out that, with a proper choice of initial distribution ν and terminal loss function Ψ, the stochastic control problem coincides with sampling problem, and the optimal policy drives samples from ν to µ perfectly. The process under optimal control can be viewed as the posterior of uncontrolled dynamics conditioned on target distribution as illustrated in Fig 2. Throughout, we denote by Qu the path measure associated with control policy u. We also denote by µ0 the terminal distribution of the uncontrolled process Q0. For the ease of presentation, we begin with sampling from a normalized density µ, and then generalize the results to unnormalized µ̂ in Section 3.4.

heading: 3.1 PATH INTEGRAL AND VALUE FUNCTION
text: Due to the special cost structure, the nonlinear HJB eq (4) can be transformed into a linear partial differential equation (PDE)
∂φt ∂t + f · ∇φt + 1 2 Tr(gg′∇2φt) = 0, φT (·) = exp{−Ψ(·)} (6)
by logarithmic transformation (Särkkä & Solin, 2019) Vt(x) = − log φt(x). By the celebrated Feynman-Kac formula (Øksendal, 2003), the above has solution
φt(x) = EQ0 [exp(−Ψ(xT ))|xt = x]. (7) We remark that eq (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 Q0. This is exactly the Path Integral control theory (Thijssen & Kappen, 2015). Furthermore, the optimal control at (t,x) is
u∗t (x) = g(t,x) ′∇ log φt(x) = lim
s↘t
EQ0{exp{−Ψ(xT )} ∫ s t dwt | xt = x} (s− t)EQ0{exp{−Ψ(xT )} | xt = x} , (8)
meaning that u∗t (x) can also be estimated by uncontrolled trajectories.

heading: 3.2 SAMPLING AS A STOCHASTIC OPTIMAL CONTROL PROBLEM
text: For a given f ,g, there are infinite choices of u such that eq (2) has terminal distribution µ. We are interested in the one that minimizes the KL divergence to the prior uncontrolled process. This is exactly the Schrödinger bridge problem (Pavon, 1989; Dai Pra, 1991; Chen et al., 2016; 2021), which has been shown to have a stochastic control formulation with cost being control efforts. In cases where ν is a Dirac distribution, it is the same as the stochastic control problem in Section 2.2 with a proper terminal cost as characterized in the following result (Tzen & Raginsky, 2019). Theorem 1 (Proof in appendix A). When ν is a Dirac distribution and terminal loss is chosen as Ψ(xT ) = log µ0(xT ) µ(xT ) , the distribution Q∗ induced by the optimal control policy is
Q∗(τ) = Q0(τ |xT )µ(xT ). (9) Moreover, Q∗(xT ) = µ(xT ).
To gain more insight, consider the KL divergence
DKL(Qu(τ)‖Q0(τ |xT )µ(xT )) = DKL(Qu(τ)‖Q0(τ) µ(xT )
µ0(xT ) ) = DKL(Qu‖Q0) + EQu [log
µ0
µ ]. (10)
Thanks to the Girsanov theorem (Särkkä & Solin, 2019),
dQu dQ0 = exp( ∫ T 0 1 2 ‖ut‖2 dt+ u′tdwt). (11)
It follows that
DKL(Qu‖Q0) = EQu [ ∫ T
0
1 2 ‖ut‖2 dt]. (12)
Plugging eq (12) into eq (10) yields DKL(Qu(τ)‖Q0(τ |xT )µ(xT )) = EQu [ ∫ T
0
1 2 ‖ut‖2 dt+ log
µ0(xT )
µ(xT ) ], (13)
which is exactly the cost defined in eq (3) with Ψ = log µ 0
µ . Theorem 1 implies that once the optimal control policy that minimizes this cost is found, it can also drive particles from x0 ∼ ν to xT ∼ µ.

heading: 3.3 OPTIMAL CONTROL POLICY AND SAMPLER
text: Optimal Policy Representation: Consider the sampling strategy from a given target density by simulating SDE in eq (2) under optimal control. Even though the optimal policy is given in eq (8), only in rare case it has an analytic closed-form.
For more general target distributions, we can instead evaluate the value function eq (7) via empirical samples using Monte Carlo. The approach is 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 problem (Au & Beck, 2003).
To overcome the above shortcomings, we parameterize the control policy with a neural network uθ. We seek a control policy that minimizes the cost
u∗ = arg min u
EQu [ ∫ T
0
1 2 ‖ut‖2 dt+ log
µ0(xT )
µ(xT ) ]. (14)
In the space of policies, eq (14) also serves as distance metric between uθ and u∗ as in eq (13).
Gradient-informed Policy Representation: It is believed that proper prior information can significantly boost the performance of neural network (Goodfellow et al., 2016). The score∇ logµ(x) has been used widely to improve the proposal distribution in MCMC (Li et al., 2020b; Hoffman & Gelman, 2014) and often leads to better results compared with proposals without gradient information. In the same spirit, we incorporate∇ logµ(x) and parameterize the policy as
ut(x) = NN1(t,x)− NN2(t)×∇ logµ(x), (15) where NN1 and NN2 are two neural networks. Empirically, we also found that the gradient information leads to faster convergence and smaller discrepancy DKL(Qu‖Q∗). We remark that PIS with policy eq (15) can be viewed as a modulated Langevin dynamics (MacKay, 2003) that achieves µ within finite time T instead of infinite time.
Evaluating Loss: To optimize uθ with objective loss eq (14), we need to evaluate µ0(xT ) numerically. For simplicity, we use x0 ∼ δ0 and f = 0,g = I so that µ0(xT ) has a simple closed-form. Empirically it works well. We include more discussion of f ,g and implementation in appendix F.1.
Optimizing uθ requires the gradient of loss in eq (14), which involves ut and the terminal state xT . To calculate gradients, we rely on backpropagation through trajectory. We train the control policy with recent techniques of Neural SDEs (Li et al., 2020a), which greatly reduce memory consumption during training. We augment the origin SDE with state ∫ t 0 1 2 ‖u‖ 2 ds such that the whole training can be conducted end to end. The full training procedure in provided in Algorithm 1.
Algorithm 1 Training Define: faug(t, [xt, yt]) = [f(t,xt) + g(t,xt)uθt(xt), 12 ‖uθt(xt)‖ 2 ]
Define: gaug(t, [xt, yt]) = [g(t,xt), 0] repeat
Vector: x0 = 0,Scalar: y0 = 0 xT , yT = sdeint(faug,gaug, [x0, y0], [0, T ]) Gradient descent step∇θ[yT + log µ 0(xT ) µ(xT ) ]
until converged
Wasserstein distance bound: The PIS trained by Algorithm 1 can not generate unbiased samples from the target distribution µ for two reasons. First, due to 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 error in simulating SDEs is inevitable. Fortunately, the following theorem quantifies the Wasserstein distance between the sampler and the target density. (More details can be found in appendix C)
Theorem 2. Under Condition 1, with sampling step size ∆t, if ‖u∗t − ut‖2 ≤ d for any t, then W2(Qu(xT ), µ(xT )) = O( √ Td(∆t+ )). (16)

heading: 3.4 IMPORTANCE SAMPLING
text: 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 in appendix B)
wu(τ) = dQu(τ) dQ∗(τ) = exp( ∫ T 0 1 2 ‖ut‖2 dt+ u′tdwt + Ψ(xT )). (17)
Algorithm 2 Sampling Vector: x0 = 0,Scalar: w0 = 0 for i← 1 to N do
∆t = ti − ti−1,∆w ∼ N (0,∆tI), xi = xi−1 + f∆t+ g(u + ∆w) wi = wi−1 + u ′∆w + 12 ‖u‖ 2
end for Outputs: xN , exp(wN + log µ 0(xN ) µ(xN ) )
We note eq (17) resembles training objective eq (14). Indeed, eq (14) is the average of logarithm of eq (17). If the trained policy is optimal, that is, Qu = Q∗, all the particles share the same weight. We summarize the sampling algorithm in Algorithm 2.
Effective Sample Size: The Effective Sample Size (ESS), ESSu = 1EQu [(wu)2] , is a popular metric to measure the variance of importance weights. ESS is often accompanied by resampling trick (Tokdar & Kass, 2010) to mitigate deterioration of sample quality. ESS is also regarded as a metric for quantifying goodness of sampler based on importance sampling. Low ESS means that estimation or downstream tasks based on such sampling methods may suffer from a high variance. ESS of most importance samplers is decreasing along the time. However, thanks to the adaptive control policy in PIS, we can quantify the ESS of PIS based on the optimality of learned policy with proof in appendix D.
Theorem 3 (Corollary 7 (Thijssen & Kappen, 2015)). If maxt,x ‖ut(x)− u∗t (x)‖2 ≤ T , then 1
EQu [(wu)2] ≥ 1− .
Estimation of normalization constants: In most sampling problems we only have access to the target density up to a normalization constant, denoted by µ̂ = Zµ. PIS can still generate samples following the same protocol with new terminal cost Ψ̂ = log µ 0
µ̂ = Ψ − logZ. The additional constant − logZ does not affect the optimal policy and the optimization of uθ. As a byproduct, we can estimate the normalization constants as follows (more details in appendix E). Theorem 4. For any given policy u, the logarithm of normalization constant is bounded below by
Eτ∼Qu [−Ŝu(τ)] ≤ logZ (18)
where Ŝu(τ) = ∫ T
0 1 2 ‖ut(xt)‖ 2 dt+ u′t(xt)dwt + Ψ̂(xT ). The equality holds only when u = u ∗. For any sub-optimal policy, an unbiased estimation of Z based on importance sampling is
Z = Eτ∼Qu [exp(−Ŝu(τ))]. (19)

heading: 4 EXPERIMENTS
text: We present empirical evaluations of PIS and several baselines as well as details of practical implementation in this section. Inspired by Arbel et al. (2021), we conduct experiments for tasks of estimating normalization constant and Bayesian inference.
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 No-U-Turn Sampler (NUTS) (Hoffman & Gelman, 2014). The second is Sequential Monte Carlo with annealing trick (SMC), which is regarded as state-of-the-art sampling algorithm (Chopin & Papaspiliopoulos, 2020). We consider a standard instance of SMC samplers and the recently proposed Annealed Flow Transport Monte Carlo (AFT) (Arbel et al., 2021). Both use a default 10 temperature levels with a linear annealing scheme. Lastly, variational normalizing flow (VI-NF) (Rezende & Mohamed, 2015) is also included for comparison.
The number of steps N of MCMC algorithm and the number of SDE time-discretization steps for PIS work as a proxy for benchmarking computation times. The time complexity of evaluating drift and control function in PIS is comparable with proposal calculation in MCMC. Since we focus on evaluating policy after training, training consumption is not included in the comparison. The training time of PIS highly depends on efficiency of training NeuralSDEs. One future direction is to investigate additional regularizations and structured SDE to speed up the training.
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, denoted as PIS-Grad. When we have an analytical form for the ground truth optimal policy, the policy is denoted as PIS-GT. We remark that PIS-GT serves the purpose of validating the connections between the optimality of the policy and the effectiveness of Path Integral Sampler. The subscript RW is to distinguish PIS with path integral importance weights eq (17) that uses eq (19) to estimate normalization constants from the ones without importance weights that use the bound in eq (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 use 10 times more steps compared with HMC and PIS. We include more details about hyperparameters, training, discussions of NUTS, and experiments with large N in appendices F and G.

heading: 4.1 PIS-GRAD VS PIS-NN: IMPORTANCE OF GRADIENT GUIDANCE
text: We observed that the advantage of PIS-Grad over PIS-NN is more clearer when the target density has multiple modes as it in a toy example shown in Fig 3. The objective DKL(Q‖Q∗) is known to have zero forcing. When the modes of the density are well separated andQ is not expressive enough, minimizing DKL(Q‖Q∗) can drive Q(τ) to zero on some area, even if Q∗(τ) > 0 (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 it is unlikely such noise only can overcome the local minima. On the other hand, the gradient information can help avoid local minima and provide hopeful exploring directions.

heading: 4.2 BENCHMARKING DATASETS
text: Mode-separated mixture of Gaussian: We consider the mixture of Gaussian in 2D dimension. We notice when Gaussian modes are not far from each other, all methods work well. However, when we reduce the variance of Gaussian distribution and separate the modes of Gaussian, the advantage of PIS becomes clear even in the low dimension tasks. We generate 2000 samples from each method and kernel density estimate (KDE) plot are provided in Fig 1. PIS generates samples that are visually indistinguishable from the given density.
Funnel distribution: We consider the popular testing distribution in MCMC literature (Hoffman & Gelman, 2014; Hoffman et al., 2019), 10-dimensional Funnel distribution:
x0 ∼ N (0, 9), x1:9|x0 ∼ N (0, exp(x0)I). The distribution can be pictured like a funnel - with x0 wide at the mouth of funnel, getting smaller as the funnel narrows.
Log Gaussian Cox Process: We further investigate estimating the normalization constant of the challenging log Gaussian Cox process (LGCP), which is designed for modeling the positions of Finland pine saplings. In LGCP (Salvatier et al., 2016), an underlying field λ of positive real values is modeled using an exponentially-transformed Gaussian process. Then λ is used to parameterized Poisson points process to model locations of pine saplings. The posterior density of λ is
λ(x) ∼ exp(− (x− µ) TK−1(x− µ) 2 ) ∏ i∈d exp(xiyi − α expxi), (20)
where d denotes the size of discretized grid and yi denotes observation information. Modeling parameters, including normal distribution and α follow Arbel et al. (2021) (See appendix F).
From the Tab 1, it clearly shows the advantages of PIS among the above three datasets. And importance weight also helps improve the estimation of log normalization constants based on the comparison between PISRW and PIS. We found PIS-Grad that trained with gradient information outperforms PIS-NN. The difference is more obvious in datasets that have well-separated modes, such as MG and LGCP, and less obvious on one mode tasks like Funnel.
In all cases, PISRW -Grad is better than AFT and SMC. Interestingly, even without annealing and gradient information of target density, PISRW -NN can outperform SMC with annealing trick and HMC kernel for the Funnel distribution.
4.3 ALANINE DIPEPTIDE
Thanks to the great success achieved by flow models in the generation of asymptotically unbiased samples from physics models (LeCun, 1998), we explore the applications in the sampling of molecular structure from a simulation of Alanine dipeptide as introduced in Wu et al. (2020). The target density of molecule is µ̂ = exp(−E(x[0:65]) − 12 ∥∥x[66:131]∥∥2). More details about the energy function E are included in appendix F.
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) (Wu et al., 2020). SNF is very close to AFT (Arbel et al., 2021). Both of them couple deterministic normalizing flow layers and MCMC blocks except SNF uses an amortized structure. We include more details of MCMC kernel and modification in appendix F. We show the generated molecular in Fig 4 and quantitative comparison of KL divergence in Tab 2, including overall atom states distribution and five multimodal torsion angles (backbone angles φ, ψ and methyl rotation angles η1, η2, η3). 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.

heading: 4.4 SAMPLING IN VARIATIONAL AUTOENCODER LATENT SPACE
text: Table 3: Estimation of log pθ(x) of a trained VAE.
B S
√
B2 + S2
VI-NF -2.3 0.76 2.42 AFT -1.7 0.95 1.96 SMC -10.6 2.01 10.79
PISRW -NN -1.9 0.81 2.06 PISRW -Grad -0.87 0.31 0.92
In this experiment we investigate sampling in the latent space of a trained Variational Autoencoder (VAE). VAE aims to minimize
DKL(q(x)qφ(z|x)‖p(z)pθ(x|z)), where qφ(z|x) represents encoder and pθ for a decoder with latent variable z and data x. We are interested in posterior distribution z ∼ p(z)pθ(x|z). (21) The normalization constant of such target unnormalized density function p(z)pθ(x|z) is exactly the likelihood of data points pθ(x), which serves as an evaluation metric for the trained VAE.
We investigate a vanilla VAE model trained with plateau loss on the binary MNIST (LeCun, 1998) dataset. For each distribution, we regard the average estimation from 10 long-run SMC with 1000 temperature levels as the ground truth normalization constant. We choose 100 images randomly and run the various approaches on estimating normalization of those posterior distributions in eq (21) and report the average performance in Tab 3. PIS has a lower bias and variance.

heading: 5 CONCLUSION
text: 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 an initial distribution to a target density perfectly when the policy is optimal for a target-dependent optimal control problem. Furthermore, we provide a calibration based on importance weights, ensuring sampling quality even with sub-optimal policies. In the future, we plan to apply this new sampling algorithm to energy-based models learning and higher dimension tasks.

heading: 6 REPRODUCIBILITY STATEMENT
text: The detailed discussion on assumptions and proof of theorems presented in the main paper is included in appendices A and C to E. For empirical experiments, the training settings and implementation tips are included in appendices F and G. We also include an implementation based on PyTorch (Paszke et al., 2019) in the supplementary material.

heading: A PROOF OF THEOREM 1
text: Before proving our main theorem, we introduce the following important lemma. Lemma 4.1 (Dai Pra (1991); Pavon (1989)). The transition density associated with optimal control policy u∗ for eq (2) and eq (3) follows
Q∗s,t(x,y) = Q0s,t(x,y) φt(y)
φs(x) , (22)
where Qus,t(x,y) denote the transition probability from state x at time s to state y at time t.
Proof of Theorem 1: We denote initial Dirac distribution by ν = δx̄0 . Combining eq (7) and Vt(x) = − log φt(x) we obtain
V0(x̄0) = − logEQ0 [exp(−Ψ(xT ))] = − log( ∫ µ
µ0 dµ0) = 0.
Therefore, we can evaluate the KL divergence between Q∗ and Q0(τ |xT )µ(xT ) as DKL(Q∗(τ)‖Q0(τ |xT )µ(xT )) = Eτ∼Q∗ [ ∫ T
0
1 2 ‖u∗‖2 dt+ Ψ(xT )] = V0(x̄0) = 0.
The first equality is based on eq (13). Next, we show that x∗T ∼ µ. The above equations imply Q∗0,T (x̄0, y)dy = exp (−Ψ(y))µ0(dy). It follows that
P[x∗T ∈ A] = ∫ A Q∗0,T (x̄0, y)dy = ∫ A exp(−Ψ(y))µ0(dy) = µ(A).

heading: B PROOF OF IMPORTANCE WEIGHTS
text: By definition
wu(τ) = dQu(τ) dQ∗(τ) = dQu(τ) dQ0(τ) dQ0(τ) dQ∗(τ) . (23)
Plugging eq (11) and eq (22) into the above we obtain
wu(τ) = exp( ∫ T 0 1 2 ‖ut‖2 dt+u′tdwt+ log µ0(xT ) µ(xT ) ) = exp( ∫ T 0 1 2 ‖ut‖2 dt+u′tdwt+ Ψ(xT )).

heading: C PROOF OF THEOREM 2
text: C.1 LIPCHITZ CONDITON AND PRELIMINARY LEMMA
To ease the burden of notations, we assume x0 ∼ δ0. 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 xT ∼ µ. The conditions and properties are studied in Schrödinger-Föllmer process, which dates back to the Schrödinger problem. For proof of existence of a unique strong solution and detailed discussion, we refer the reader to Dai Pra (1991); ?); Eldan et al. (2020); Huang et al. (2021). Condition 1. ||u∗t (x)||22 ≤ C0(1 + ‖x‖2) (24) and ∥∥u∗t1(x1)− u∗t2(x2)∥∥ ≤ C1(‖x1 − x2‖+ |t1 − t2| 12 ). (25) We introduce the following lemma before stating the proof for Theorem 2. Lemma 4.2. (Huang et al., 2021, Lemma A.2.) Assume Condition 1 holds, then the following inequality holds for xt generated from u∗,
E[‖xt2 − xt1‖2] ≤ 4C0 exp(2C0)(C0 + d)(t2− t1)2 + 2C0(t2− t1)2 + 2d|t2− t1|, t1, t2 ∈ [0, T ]. (26)
C.2 PROOF OF THEOREM 2
We denote by x∗(0:T ) the trajectory controlled by the optimal policy u ∗, and by {xt0 ,xt1 , · · · ,xtN } the discrete time process with sub-optimal policy u over discrete time {tk} such that t0 = 0, tN = T and tk − tk−1 = ∆t. The process {xtk} can be extended to continuous time setting as
xtk = xtk−1 + ∫ tk tk−1 utk−1(xtk−1)ds+ dws.
The key of our proof is the bound of ∥∥xtk − x∗tk∥∥2.
∥∥xtk − x∗tk∥∥2 = ∥∥∥∥∥xtk−1 + ∫ tk tk−1 utk−1(xtk−1)ds+ dws − [x∗tk−1 + ∫ tk tk−1 u∗s(x ∗ s)ds+ dws] ∥∥∥∥∥ 2
≤ ∥∥∥xtk−1 − x∗tk−1∥∥∥2 + ∥∥∥∥∥ ∫ tk tk−1 [u∗s(x ∗ s)− utk−1(xtk−1)]ds ∥∥∥∥∥ 2
+ 2 ∥∥∥xtk−1 − x∗tk−1∥∥∥ ∥∥∥∥∥ ∫ tk tk−1 [u∗s(x ∗ s)− utk−1(xtk−1)]ds ∥∥∥∥∥ ≤ ∥∥∥xtk−1 − x∗tk−1∥∥∥2 + (∫ tk
tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥ ds)2 + 2
∥∥∥xtk−1 − x∗tk−1∥∥∥ (∫ tk tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥ ds) ≤ (1 + α)
∥∥∥xtk−1 − x∗tk−1∥∥∥2 + (1 + 1α )( ∫ tk tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥ ds)2 ≤ (1 + α)
∥∥∥xtk−1 − x∗tk−1∥∥∥2 + (1 + 1
α )(tk − tk−1)( ∫ tk tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥2 ds), (27) where the first and second inequality is based on the triangle inequality, the third inequality is based 2ab ≤ αa2+ 1αb2 for any α > 0, and the forth inequality is based on the Cauchy-Schwarz inequality. In the following we bound the second term in eq (27) as∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥2
= ∥∥∥u∗s(x∗s)− u∗tk−1(xtk−1) + u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2
≤ (1 + β) ∥∥∥u∗s(x∗s)− u∗tk−1(xtk−1)∥∥∥2 + (1 + 1β )∥∥∥u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2
≤ 2C21 (1 + β)[ ∥∥x∗s − xtk−1∥∥2 + |s− tk−1|] + (1 + 1β )∥∥∥u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2 ,
where the first inequality uses 2ab ≤ βa2 + 1β b2 for an arbitrary β > 0 and the second one is based on eq (25). It follows that∫ tk
tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥2 ds ≤ ∫ tk tk−1 2C21 (1 + β) ∥∥x∗s − xtk−1∥∥2 ds+ ∫ tk tk−1 2C21 (1 + β)(s− tk−1)ds
+ ∫ tk tk−1 (1 + 1 β ) ∥∥∥u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2 ds.
Thus, for stepsize ∆t = tk − tk−1, we establish∫ tk tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥2 ds ≤ ∫ tk tk−1 2C21 (1 + β) ∥∥x∗s − xtk−1∥∥2 ds
+ C21 (1 + β)∆t 2 + (1 +
1 β ) ∥∥∥u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2 ∆t. (28)
Next we bound ∥∥x∗s − xtk−1∥∥2 in eq (28) as∥∥x∗s − xtk−1∥∥2 ≤ (1 + η)∥∥∥x∗s − x∗tk−1∥∥∥2 + (1 + 1η )∥∥∥x∗tk−1 − xtk−1∥∥∥2 , (29)
where the inequality is based on (a+ b)2 ≤ (1 + η)a2 + (1 + 1η )b2 for an arbitrary η > 0. Plugging eq (29) into eq (28) yields∫ tk
tk−1 ∥∥u∗s(x∗s)− utk−1(xtk−1)∥∥2 ds ≤ (1 + 1 β ) ∥∥∥u∗tk−1(xtk−1)− utk−1(xtk−1)∥∥∥2 ∆t+ C21 (1 + β)∆t2
+ 2C21 (1 + β)(1 + 1 η ) ∥∥∥x∗tk−1 − xtk−1∥∥∥2 ∆t+ 2C21 (1 + β)(1 + η)∫ tk
tk−1 ∥∥∥x∗s − x∗tk−1∥∥∥2 ds. (30)
Plugging eq (28) and (30) into eq (27) yields
LHS ≤[1 + α+ 2C21 (1 + 1
α )(1 + β)(1 +
1 η )∆t2] ∥∥∥x∗tk−1 − xtk−1∥∥∥2 + 2C21 (1 + 1
α )(1 + β)(1 + η)∆t ∫ tk tk−1 ∥∥∥x∗s − x∗tk−1∥∥∥2 ds + (1 + 1
α )(1 +
1 β ) ∥∥∥u∗tk−1(x∗tk−1)− utk−1(xtk−1)∥∥∥2 ∆t2 + (1 + 1α )C21 (1 + β)∆t3.
(31)
Invoking lemma 4.2, we obtain E[ ∫ tk tk−1 ∥∥∥x∗s − x∗tk−1∥∥∥2 ds] ≤ 4C0 exp(2C0)(C0 + d)∆t3 + 2C0∆t3 + 2d∆t2. Taking the expectation of eq (31), in view of the above and the assumption on control, we establish
E[ ∥∥xtk − x∗tk∥∥2] ≤ C3E[∥∥∥xtk−1 − x∗tk−1∥∥∥2] + C4,
where C3 = [1 + α+ 2C21 (1 + 1 α )(1 + β)(1 + 1 η )∆t 2], and
C4 =(1 + 1
α )(1 +
1 β )d ∆t2 + (1 + 1 α )C21 (1 + β)∆t 3
+ 2C21 (1 + 1
α )(1 + β)(1 + η)[4C0 exp(2C0)(C0 + d)∆t
3 + 2C0∆t 3 + 2d∆t2]∆t.
Finally, in view of the fact x0 = x∗0 and fixed step size ∆t, we conclude that by the choice α = C1∆t, β = η = 1
E[‖xT − x∗T ‖2] ≤ C T ∆t 3 − 1 C3 − 1 C4 = O(dT (∆t+ )), (32)
where the last inequality based on ignoring the high order terms, C T ∆t 3 − 1 ≤ O(T ) and C4C3−1 ≤ O(d(∆t+ )).

heading: D PROOF OF THEOREM 3
text: The proof is a natural extension of Corollary 7 in Thijssen & Kappen (2015).
We define random variable
Su(t) = ∫ t 0 ‖us‖2 2 ds+ u′sdws + Ψ(xT ), (33)
and Φ(t) = exp(−Su(0) + Su(t)). (34)
Lemma 4.3. (Thijssen & Kappen, 2015, Lemma 4) For any feasible control policy u for stochastic optimal control problem,
Φ(T )φt(xT )− Φ(t)φt(xt) = ∫ T t Φ(s)φt(xs)(u ∗ s − us)′dws. (35)
Corollary 1. φt(x) = EQ0 [exp(−Ψ(xT ))|xt = x] = EQu [exp(−Ψ(xT )− ∫ T t ‖us‖2 2 ds)|xt = x] (36)
Proof. This follows importance sampling with density ratio from eq (11).
Proof of Theorem 3: We denote the important weight bywu; note it is a random variable. It follows that
wu = exp(−Ψ(xT )−
∫ T 0 (‖u‖ 2 2 dt+ u ′dw))
EQu [exp(−Ψ(xT )− ∫ T 0 (‖u‖ 2 2 dt+ u ′dw))] = exp(−Su(t)) EQu [exp(−Su(t))] . (37)
Dividing the LHS of eq (35) by φ0(x0) we obtain
Φ(T )φt(xT )− Φ(t)φt(xt) φ0(x0) = Φ(T )φt(xT )− φ0(x0) EQu [exp(−Su(0))] = wu − EQu [wu]. (38)
Therefore, the variance of wu equals EQu [(wu − EQu [wu])2] = EQu [( ∫ T
0
Φ(s)φs(xs)
φ0(x0) (u∗s − us)′dw)2]
= EQu [ ∫ T
0
Φ2(s)φ2s(xs)
φ20(x0) (u∗s − us)′(u∗s − us)ds]
= EQu [ ∫ T
0
(wuφs(xs) exp(S u(s)))2(u∗s − us)′(u∗s − us)ds]. (39)
By Jensen’s inequality
φ(s,xs) 2 = (EQu [exp(−Su(s))|xs])2 ≤ EQu [exp(−2Su(s))|xs].
Plugging the above inequality into eq (39), we reach the upper bound of variance EQu [(wu − EQu [wu])2] ≤ ∫ T
0
EQu [(u∗s − us)′(u∗s − us)(wu)2]ds. (40)
In view of the fact Var(wu) + 1 = EQu [(wu)2], we arrive at 1 + EQu [(wu)2] ≤ EQu [(wu)2] ∫ T
0
EQu [(u∗s − us)′(u∗s − us)]ds.
If we consider the near optimal policy such that maxt,x ‖ut(x)− u∗t (x)‖2 ≤ T , then it follows that 1
EQu [(wu)2] ≥ 1− . (41)

heading: E PROOF OF THEOREM 4
text: We consider the KL divergence between trajectory distribution resulted from the policy u and the one from optimal policy u∗:
DKL(Qu(τ)‖Q∗(τ)) = DKL(Qu(τ)‖Q0(τ) µ(xT )
µ0(xT ) ) = Eτ∼Qu [ ∫ T
0
1 2 ‖ut‖2 dt+ u′tdwt + Ψ(xT )]
= Eτ∼Qu [ ∫ T
0
1 2 ‖ut‖2 dt+ u′tdwt + Ψ̂(xT ) + logZ]
= Eτ∼Qu [Ŝu(τ) + logZ] ≥ 0.
The last inequality is based on the fact DKL(Qu(τ)‖Q∗(τ)) ≥ 0 and the equality holds only when u = u∗, pointing to
0 = Eτ∼Q∗ [Ŝu(τ) + logZ].
Therefore, we can estimate the normalization constant by
Z ≥ exp(−Eτ∼Qu [Ŝu(τ)]), Z = exp(−Eτ∼Q∗ [Ŝu(τ)]). (42)
Next we provide an unbiased estimation with sub-optimal policy u based on importance sampling as
1 = Eτ∼Qu [ Q∗(τ) Qu(τ) ]
= Eτ∼Qu [exp(log Q∗(τ) Qu(τ) )]
= Eτ∼Qu [exp(−Ŝu(τ)− logZ)].
The last equality is based on the fact Eτ∼Qu [Q ∗(τ) Qu(τ) ] = ∫ τ Q∗(τ)dτ = 1. Hence, we obtain an unbiased estimation of the normalization constant as
Z = Eτ∼Qu [exp(−Ŝu(τ))].

heading: F EXPERIMENT DETAILS AND DISCUSSIONS
text: The PIS algorithm is implemented in PyTorch (Paszke et al., 2019). We use Adam optimizer (Kingma & Ba, 2014) in all experiments to learn optimal policy with learning rates 5× 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 event 1 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 7 minutes and 9 minutes on LGCP (d=1600).
For all trained PIS and its variants, we use 100 time-discretization steps for the SDEs. Gradient clipping with value 1 is used. A Fourier feature augmentation (Tancik et al., 2020) is employed for time condition. For HMC, we uses 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 of SMC and AFT follow the official implementation (Arbel et al., 2021) in the released codebase 1.
1https://github.com/deepmind/annealed_flow_transport
F.1 CHOICE OF f ,g:
As discuss in Section 3.3, we prefer use a linear function for f ,g to promise a closed-form µ0. Choice of f ,g encodes prior knowledge into dynamics without control and Q∗ is determined based on the prior Q0. Intuitively, the ideal Q0 should drive particles from ν to Q0(xT ) that is close to µ. In PIS, our training objective is to fit Q∗ with parameterized Qu. Thus training can be easier and faster ifQ∗ andQ0 are close since we use zero control as initialization for training policy. However, there is no general approach to choose f ,g such thatQ0(xT ) is close to µ andQ0(xT ) has a closed form. In this work, we adopt the general form with f = 0,g = I. It is interesting to explore other prior dynamics or data-variant f ,g in the future work.
F.2 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 temperature levels. One corresponds to a standard Gaussian distribution and the other one corresponds to the target density. Since the initial distribution ν for SMC and NUTs is chosen as standard Gaussian, we can omit the MCMC steps for it and the total computation efforts required for the specific SMC are for the transitions in HMC and NUTS.
For VI-NF, we use importance sampling∫ µ̂(x)dx = ∫ q(x) µ̂(x)
q(x) dx = Eq[
µ̂(x) q(x) ],
where q is the normalized distribution represented by normalizing flows, to provide an unbiased estimation of normalization constants. We use the ELBO in eq (18) for PIS and the unbiased estimation eq (19) for PISRW .
F.3 2 DIMENSIONAL RINGS EXAMPLE
The ring-shape density function
log µ̂ = −min((‖x‖ − 1) 2, (‖x‖ − 3)2, (‖x‖ − 5)2)
100 .
Consider the special case of gradient informed SDE, which can be viewed as PIS-Grad with a specific group of parameters,
dx = −d log µ̂(x) dx
dt+ √ 2dw.
This is exactly the Langevin dynamics used widely in sampling (MacKay, 2003). As a special case of MCMC, Langevin sampling can generate high quality samples given large enough time interval (MacKay, 2003). From this perspective, PIS-Grad can be viewed as a modulated Langevin dynamics that is adjusted and represented by neural networks.
F.4 BENCHMARKING DATASETS
For mixture of Gaussian, we choose nine centers over the grid {−5, 0, 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 use 2000 samples for estimating the log normalization constant Z. We use the standard MLP network to parameterize the control drift ut(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 2 layer MLP to map the summation of features to 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 of steps, PIS achieves the best performance. We include more results for long-run MCMC methods in Tab 4.
In the experiment with Funnel distribution, Arbel et al. (2021) suggests to use a slice sampler kernel for AFT and SMC, which includes 1000 steps of slice sampling per temperature. In Tab 1, we still use HMC for comparing performance with the same number of integral steps. We also include the results with slice sampler in Tab 5. 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 as
K(u, v) = 1.91× exp(−‖u− v‖2 Mβ ),
and the mean vector equals log(126) − σ2 and α = 1/M2. We note that this setting follows Arbel et al. (2021); Møller et al. (1998). 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 temperatures 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 is 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 is no prior information that can be used to shrink the search space, which makes the training of PIS
with MLP more difficult. We use 2000 particles for estimation of log normalization constants. We also include more experiment results in Tab 6.
F.5 ALANINE DIPEPTIDE
The setup for target density distribution and the comparison method are adopted from Wu et al. (2020). Following Noé et al. (2019), an 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 in vacuum using OpenMMTools (Eastman et al., 2017) 2. Total 105 atoms data points are generated as training data. Situation parameters, including time-step and temperature setting are the same as Wu et al. (2020). We refer the reader to official codebase for more details of setting target density function 3.
Following the setup in Wu et al. (2020), we use unweighted samples to compute the metrics. KL divergence of VI-NF on µ is calculated based on ELBO instead of importance sampling as in normalizing constants tasks. We use Metropolis random walk MCMC block for SMC, SNF and AFT and RealNVP blocks for SNF and AFT (Dinh et al., 2016). For a 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 is originally trained with maximizing data likelihood where some empirical samples are assumed to be available. We modify the training objective function by reversing the original KL divergence as in eq (1).
F.6 MORE DETAILS ON SAMPLING IN VARIATIONAL AUTOENCODER LATENT SPACE
We use a vanilla VAE architecture to train on binary MNIST data. The encoder uses a standard 3 layer MLP networks with 1024 hidden neurons, and maps an image to the mean and standard deviation of 50 dimension diagonal Normal distribution. The decoder employs 3 layer MLP networks to decode images from latent states. ReLU nonlinearity is used for hidden layers. For training, we use the Adam optimizer with learning rate 5 × 10−4, batch size 128. With reparameterization trick (Kingma & Welling, 2013) and closed-form KL divergence between approximated normal distribution and standard normalization distribution, we train networks for totally 100 epochs. We show performance of vanilla in Fig 5.
2https://github.com/choderalab/openmmtools 3https://github.com/noegroup/stochastic_normalizing_flows
We parameterize distribution of decoder pθ(x|z) as log pθ(x|z) = log p(x|Dθ(z)) = x logDθ(z) + (1− x) log(1−Dθ(z)).
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.

heading: G TECHNICAL DETAILS AND TIPS
text: Here we provide a 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 Lipchastiz constants, which is within expectation since large control is required to drive particles to destination with less amount of time. It is reported that it is more difficult to approximate large Lipchastiz functions with neural networks (Jacot et al., 2018; Tancik et al., 2020). We thus recommend to increase T or constraint the magnitude of u to 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 ut(x) has large Lipchastiz constants in time dimension, the discretized error would also increase. For calculating the weights based on path integral, we suggest to 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 formula (Efron, 2011), but we found no obvious improvements on testing the trained sampler or estimating normalization constants.
• Regardless the accuracy and memory advantages of Reversible Heun claimed by torchsde (Li et al., 2020a; Kidger et al., 2021) , we found this integration approach is less stable compared with simple Euler integration without adjoint and results in numerical issues occasionally. We recommend readers to use Ito Euler integration when memory permits or conduct training with a small ∆t.

Review summaries:
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.

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.

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.