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