Path Integral Sampler: A Stochastic Control Approach For Sampling

1 INTRODUCTION . 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 . 2 SAMPLING AND STOCHASTIC CONTROL PROBLEMS . 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 . 2.1 SAMPLING PROBLEM . 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 . 2.2 STOCHASTIC CONTROL . 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 ) 3 PATH INTEGRAL SAMPLER . 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 . 3.1 PATH INTEGRAL AND VALUE FUNCTION . 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 .

Summary: 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.