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: This paper proposes a new MCMC transition kernel. This kernel is parameterized by neural networks and is optimized through an objective maximizing the proposal entropy. Specifically, the authors use a combination of a flow model and non-volume preserving flow in [Dinh et al., 2016] as the neural network parameterized kernel. Then they use the objective in [Titsias & Dellaportas, 2019] which maximizes the proposal entropy to optimize the kernel. The proposed method is tested on synthetic datasets, Bayesian logistic regression and a deep energy-based model.