A Neural Network MCMC sampler that maximizes Proposal Entropy

1 INTRODUCTION . Sampling from unnormalized distributions is important for many applications , including statistics , simulations of physical systems , and machine learning . However , the inefficiency of state-of-the-art sampling methods remains a main bottleneck for many challenging applications , such as protein folding ( Noé et al. , 2019 ) , energy-based model training ( Nijkamp et al. , 2019 ) , etc . A prominent strategy for sampling is the Markov Chain Monte Carlo ( MCMC ) method ( Neal , 1993 ) . In MCMC , one chooses a transition kernel that leaves the target distribution invariant and constructs a Markov Chain by applying the kernel repeatedly . The MCMC method relies only on the ergodicity assumption , other than that it is general . If enough computation is performed , the Markov chain generates correct samples from any target distribution , no matter how complex the distribution is . However , the performance of MCMC depends critically on how well the chosen transition kernel explores the state space of the problem . If exploration is ineffective , samples will be highly correlated and of very limited use for downstream applications . Despite some favorable theoretical argument on the effectiveness of some MCMC algorithms , practical implementation of them may still suffer from inefficiencies . Take , for example , the Hamiltonian Monte Carlo ( HMC ) ( Neal et al. , 2011 ) algorithm , a type of MCMC technique . HMC is regarded state-of-the-art for sampling in continuous spaces Radivojević & Akhmatskaya ( 2020 ) . It uses a set of auxiliary momentum variables and generates new samples by simulating a Hamiltonian dynamics starting from the previous sample . This allows the sample to travel in state space much further than possible with other techniques , most of whom have more pronounced random walk behavior . Theoretical analysis shows that the cost of traversing a d-dimensional state space and generating an uncorrelated proposal is O ( d 1 4 ) for HMC , which is lower than O ( d 1 3 ) for Langevine Monte Carlo , and O ( d ) for random walk . However , unfavorable geometry of a target distribution may still cause HMC to be ineffective because the Hamiltonian dynamics has to be simulated numerically . Numerical errors in the simulation are commonly corrected by a Metropolis-Hastings ( MH ) accept-reject step for a proposed sample . If the the target distribution has unfavorable geometric properties , for example , very different variances along different directions , the numerical integrator in HMC will have high error , leading to a very low accept probability ( Betancourt et al. , 2017 ) . For simple distributions this inefficiency can be mitigated by an adaptive re-scaling matrix ( Neal et al. , 2011 ) . For analytically tractable distributions , one can also use the Riemann manifold HMC method ( Girolami & Calderhead , 2011 ) . But in most other cases , the Hessian required in Riemann manifold HMC algorithm is often intractable or expensive to compute , preventing its application . Recently , approaches have been proposed that possess the exact sampling property of the MCMC method , while potentially mitigating the described issues with unfavorable geometry . Such approaches include MCMC samplers augmented with neural networks ( Song et al. , 2017 ; Levy et al. , 2018 ; Gu et al. , 2019 ) , and neural transport MCMC techniques ( Hoffman et al. , 2019 ; Nijkamp et al. , 2020 ) . A disadvantage of these recent techniques is that their objectives optimize the quality of proposed samples , but do not explicitly encourage exploration speed of the sampler . One notable exception is L2HMC ( Levy et al. , 2018 ) , a method whose objective includes the size of the expected L2 jump , thereby encouraging exploration . But the L2 expected jump objective is not very general , it only works for simple distributions ( see Figure 1 , and below ) . Another recent work ( Titsias & Dellaportas , 2019 ) proposed a quite general objective to encourage exploration speed by maximizing the entropy of the proposal distribution . In continuous space , the entropy of a distribution is essentially the logarithm of its volume in state space . Thus , the entropy objective naturally encourages the proposal distribution to “ fill up ” the target state space as well as possible , independent of the geometry of the target distribution . The authors demonstrated the effectiveness of this objective on samplers with simple linear adaptive parameters . Here we employ the entropy-based objective in a neural network MCMC sampler for optimizing exploration speed . To build the model , we design a flexible proposal distribution for which the optimization of the entropy objective is tractable . Inspired by the HMC algorithm , the proposed sampler uses special architecture that utilizes the gradient of the target distribution to aid sampling . For a 2-D distribution the behavior of the proposed model is illustrated in Figure 1 . The sampler , trained with the entropy-based objective , generates samples that explore the target distribution quite well , while it is simple to construct a proposal with higher L2 expected jump ( right panel ) . Later we show the newly proposed method achieves significant improvement in sampling efficiency compared to previous techniques , we then apply the method to the training of an energy-based image model . 2 PRELIMINARY : MCMC METHODS , FROM VANILLA TO LEARNED . Consider the problem of sampling from a target distribution p ( x ) = e−U ( x ) /Z defined by the energy function U ( x ) in a continuous state space . MCMC methods solve the problem by constructing and running a Markov Chain , with transition probability p ( x′|x ) , that leaves p ( x ) invariant . The most general invariance condition is : p ( x′ ) = ∫ p ( x′|x ) p ( x ) dx for all x′ , which is typically enforced by the simpler but more stringent condition ofdetailed balance : p ( x ) p ( x′|x ) = p ( x′ ) p ( x|x′ ) . For a general distribution p ( x ) it is difficult to directly construct p ( x′|x ) that satisfies detailed balance , but one can easily1 make any transition probability satisfy it by including an additional Metropolis-Hastings accept-reject step ( Hastings , 1970 ) . When we sample x′ at step t from an arbitrary proposal distribution q ( x′|xt ) , the M-H accept-reject process accepts the new sample xt+1 = x′ with probability A ( x′ , x ) = min ( 1 , p ( x ′ ) q ( xt|x′ ) p ( xt ) q ( x′|xt ) ) . If x′ is rejected , the new sample is set to the previous state xt+1 = xt . This transition kernel p ( x′|x ) constructed from q ( x′|x ) and A ( x′ , x ) leaves any target distribution p ( x ) invariant . The most popular MCMC techniques use the described M-H accept-reject step to enforce detailed balance , for example , Random Walk Metropolis ( RWM ) , Metropolis-Adjusted Langevin Algorithm ( MALA ) and Hamiltonian Monte Carlo ( HMC ) . For brevity , we will focus on MCMC methods that use the M-H step , although some alternatives do exist ( Sohl-Dickstein et al. , 2014 ) . All these methods share the requirement that the accept probability in the M-H step must be tractable to compute . For two of the mentioned MCMC methods this is indeed the case . In the Gaussian random-walk sampler , the proposal distribution is a Gaussian around the current position : x′ = x+ ∗N ( 0 , I ) , which has the form x′ = x+z . Thus , forward and reverse proposal probabilities are given by q ( x′|x ) = pN [ ( x′ − x ) / ] and q ( x|x′ ) = pN [ − ( x′ − x ) / ] , where pN denote the density function of Gaussian . The probability ratio q ( x t|x′ ) q ( x′|xt ) used in the M-H step is therefore equal to 1 . In MALA the proposal distribution is a single step of Langevin dynamics with step size : x′ = x+ z with z = − 2 2 ∂xU ( x ) + N ( 0 , I ) . We then have q ( x ′|x ) = pN [ ( x′ − x ) / + 2∂xU ( x ) ] and q ( x|x′ ) = pN [ − ( x′ − x ) / + 2∂x′U ( x ′ ) ] . Both , the forward and reverse proposal probability are tractable since they are the density of Gaussians evaluated at a known location . Next we show how the HMC sampler can also be formulated as a M-H sampler . Basic HMC involves a Gaussian auxiliary variable v of the same dimension as x , which plays the role of the momentum in Physics . HMC sampling consists of two steps : 1 . The momentum is sampled from a normal distribution N ( v ; 0 , I ) . 2 . The Hamiltonian dynamics is simulated for a certain duration with initial condition x and v , typically by running a few steps of the leapfrog integrator . Then , a M-H accept-reject process with accept probability A ( x′ , v′ , x , v ) = min ( 1 , p ( x ′ , v′ ) q ( x , v|x′ , v′ ) p ( x , v ) q ( x′ , v′|x , v ) ) = min ( 1 , p ( x ′ ) pN ( v ′ ) p ( x ) pN ( v ) ) is performed to correct for error in the integration process . We have q ( x , v|x′ , v′ ) q ( x′ , v′|x , v ) = 1 since the Hamiltonian transition is volume-preserving over ( x , v ) . Both HMC steps leave the joint distribution p ( x , v ) invariant , therefore HMC samples from the correct distribution p ( x ) after marginalizing over v. To express basic HMC in the standard M-H scheme , step 1 and 2 can be aggregated into a single proposal distribution on x with the proposal probability : q ( x′|x ) = pN ( v ) and q ( x|x′ ) = pN ( v′ ) . Note , although the probability q ( x′|x ) can be calculated after the Hamiltonian dynamics is simulated , this term is intractable for general x and x′ . The reason is that it is difficult to solve for the v at x to make the transition to x′ using the Hamiltonian dynamics . This issue is absent in RWM and MALA , where q ( x′|x ) is tractable for any x and x′ . Previous work on augmenting MCMC sampler with neural networks also relied on the M-H procedure to ensure asymptotic correctness of the sampling process , for example ( Song et al. , 2017 ) and ( Levy et al. , 2018 ) . They used HMC style accept-reject probabilities that lead to intractable q ( x′|x ) . Here , we strive for a flexible sampler for which q ( x′|x ) is tractable . This maintains the tractable M-H step while allowing us to train this sampler to explores the state space by directly optimizing the proposal entropy objective , which is a function of q ( x′|x ) .

Summary: The paper argues that a better objective to train neural MCMC kernels is to maximize the proposal entropy (Titsias & Dellaportas, 2019) and demonstrate a method on doing so. The method shows improved sampling efficiency compared to previous method, especially one that optimize the alternative L2 expected jump. The novelty is not of the training objective but a neural instantiation with improved sampling efficiency.