Semi-Empirical Objective Functions for Neural MCMC Proposal Optimization

1 INTRODUCTION . The development of efficient Markov Chain Monte Carlo ( MCMC ) proposal distributions is vital to enable numerical estimation and statistical inference in increasingly complicated problem domains . If we had an exact definition of the notion of MCMC efficiency , this could be accomplished by straightforward optimization over a set of proposal distributions . Past computational limitations encouraged the use of architecturally limited MCMC schemes that yield useful results for a wide range of target distributions , such as Random Walk Metropolis ( RWM ) ( Metropolis et al. , 1953 ) , Metroplis Adjusted Langevin Diffusion ( MALA ) ( Besag and Green , 1993 ) , and Hamiltonian Monte Carlo ( HMC ) ( Duane et al. , 1987 ) . Although much research has been devoted to adaptive methods to improve MCMC performance , traditional adaptive methods ( Roberts and Rosenthal , 2009 ; Sejdinovic et al. , 2014 ; Haario et al. , 2001 ; Roberts and Rosenthal , 2007 ) focus on optimization across highly restricted proposal distribution model classes . Ideally , we would like to develop a practical means of optimizing MCMC performance over arbitrarily parameterized classes of proposal distributions . Given their demonstrated success in parameterizing expressive distributions , deep generative models are naturally suited to the problem of MCMC proposal optimization . Recent research regarding applications of deep learning to proposal optimization has yielded new MCMC schemes ( Song et al. , 2017 ; Spanbauer et al. , 2020 ; Hoffman et al. , 2019 ) and extensions of existing schemes ( de Freitas et al. , 2001 ; Habib and Barber , 2018 ; Levy et al. , 2018 ; Li et al. , 2021 ) . A variety of objective functions are utilized to optimize these deep learning based approaches , with functional forms generally dependent on the type of MCMC scheme being optimized . As demonstrated by the experiments of this work , the objective functions currently used for MCMC proposal optimization rely on model class restrictions imposed on the proposal architecture and are not suitable for optimizing proposal distributions more expressive than traditional MCMC schemes . In this work , we introduce and demonstrate Ab Initio objective functions for MCMC proposal optimization intended to remain compatible with any proposal distribution defined using deep generative models . Presumably , there exists some `` ground truth '' objective function underlying our notion of MCMC sampling performance . However , even after the decades of research regarding MCMC methods , there appears to be no universal definition for what we mean by MCMC efficiency . Metrics like effective sample size ( ESS ) ( Gelman et al. , 2013 ; Vats et al. , 2019 ) generally coincide with our notion of sampling performance , but no such metric serves as the canonical definition of MCMC efficiency . Theoretical analysis regarding optimal acceptance rates for particular MCMC schemes ( Roberts et al. , 2001 ; Gelman et al. , 1997 ; Roberts and Rosenthal , 1998 ; Neal et al. , 2012 ; Beskos et al. , 2013 ) considers restricted proposal schemes and targets within a continuous diffusionary limit wherein our common performance metrics converge in their definition of optimality ( Roberts et al. , 2001 ) . Within this diffusionary limit , useful properties regarding MCMC optimality ( e.g . the rules of thumb to seek an acceptance rate of 0.234 when using RWM and of 0.574 when using MALA ) may be derived without specifically defining sampling performance . In light of the lack of an exactly specified objective function for MCMC efficiency , our Ab Initio objective functions seek to approximate the ground truth definition of sampling performance by adhering to certain reasonable first principles properties and fitting to reproduce the `` mathematical observations '' provided by existing theoretical analysis of optimal acceptance rates . Our contributions are as follows : • We describe a set of first principles properties that may be reasonably assumed of the ground truth objective function underlying our notion of MCMC efficiency . • We illustrate the construction of an example Ab Initio objective function via the combination of simpler objective functions with coefficients determined to reproduce optimal behavior on reference problems with analytically known solutions . • We verify the generality of the resulting Ab Initio objective function through its ability to closely reproduce analytically known optimal results for a wide range of optimization tasks beyond the reference problem used in its construction . • Through a series of illustrative experiments , we demonstrate the advantages of Ab Initio objective functions for optimizing arbitrarily defined MCMC proposal distributions . 2 RELATION TO PRIOR WORK . Ab Initio techniques ( Hehre , 1976 ; Yin and Cohen , 1982 ; Marx and Hutter , 2009 ) are used in the physical sciences to simulate systems that would otherwise be unobservable in a laboratory setting . These Ab Initio methods are founded on principled approximations of fundamental physical laws with parameters chosen to reproduce the properties of reference systems that can be observed . We take inspiration from these methods in the physical sciences in forming the methodology of this work . The purpose of this work is to introduce a procedure for selecting objective functions for MCMC proposal optimization that are suited to the optimization of proposal model classes arbitrarily parameterized by deep generative models . Research into the applications of deep learning for MCMC proposal optimization ( Song et al. , 2017 ; Spanbauer et al. , 2020 ; Hoffman et al. , 2019 ; de Freitas et al. , 2001 ; Habib and Barber , 2018 ; Levy et al. , 2018 ; Li et al. , 2021 ) has yielded a number of potential candidates for this objective function . We therefore compare our Ab Initio objective functions to these candidates on the basis of their suitability for optimization of arbitrary proposal distributions . Pure KL-divergence based objectives ( de Freitas et al. , 2001 ; Habib and Barber , 2018 ; Neklyudov et al. , 2018 ; Wang et al. , 2018 ) have found some success , particularly when optimizing resampling style schemes . These objectives are unable to properly optimize proposals within a diffusionary limit ( Titsias and Dellaportas , 2019 ) . We therefore omit pure KL-divergence based objectives from the comparisons within this work . Adversarial objectives have been used to optimize proposal distributions , notably with A-NICE-MC ( Song et al. , 2017 ; Spanbauer et al. , 2020 ) . We view adversarial training as approximately optimizing an existing performance measure ( e.g . KL-divergence ) , rather than defining a fundamentally new performance measure . We therefore also omit adversarial objectives from our comparisons . Mean squared jump distance ( Pasarica and Gelman , 2010 ) ( MSJD ) remains a popular objective for optimizing proposals . Notably , L2HMC ( Levy et al. , 2018 ) is optimized using a modification of MSJD with a regularization intended to encourage mixing of the resulting Markov Chain . As shown in our experiments , both MSJD and L2HMC ’ s modification can produce arbitrarily non-ergodic Markov Chains when optimizing proposal distributions with position dependence . Our Ab Initio objective functions avoid this undesirable optimization behavior by maintaining i.i.d . resampling from the target as their unique global minima . Another potential advantage of the Ab Initio objective of Equation ( 2 ) over MSJD based objective functions is that it does not rely on the notion of a metric of the underlying space of the proposal distribution , and so may be more suited to optimizing proposals in settings where the space is equipped with only a measure and not a metric ( e.g . MCMC settings involving discrete or graph based distributions ) . Recently , Titsias and Dellaportas ( 2019 ) introduced the Generalized Speed Measure ( GSM ) as an objective function for MCMC proposal optimization . The GSM amounts to maximizing proposal entropy subject to the constraint of achieving a user specified acceptance rate . Although it is theoretically well motivated , the GSM relies on knowledge of optimal acceptance rates for a given optimization problem , which will generally be unknown when using very general neural architectures . Our example Ab Initio objective function of Equation ( 1 ) takes inspiration from the components of the GSM , while also ensuring the functional limitations argued for in Section 4 . A key advantage of our Ab Initio objective functions over the GSM is that they do not require prior knowledge to recover near optimal MCMC behavior ( e.g . our Ab Initio objective function , whose construction only involved the knowledge of the optimal acceptance rate for RWM , recovers optimized MALA proposals with acceptance rates around 0.5 ) . An additional advantage of our Ab Initio objective functions over the GSM is that they may be used to compare the efficiencies of proposal distributions with differing acceptance rates outside of the context of parameter optimization . 3 MCMC PROPOSAL OPTIMIZATION . In this work , we consider the task of optimizing a proposal density gθ ( ~x′|~x ) ( ~x ∈ Rn ) to sample from a target density π ( ~x ) for a MCMC task . Proposals are accepted with rate αgθ , π ( ~x ′|~x ) that ensures the resulting Markov Chain converges towards π ( ~x ) . For this work , we restrict ourselves to Metropolis-Hastings type schemes , wherein αgθ , π ( ~x ′|~x ) = min { 1 , π ( ~x ′ ) gθ ( ~x|~x′ ) π ( ~x ) gθ ( ~x′|~x ) } . For optimization , we utilize some measure of sampling performance to define an objective function L [ g ; π ] that imposes an ordering over proposal distributions by defining g1 < g2 exactly when L [ g1 ; π ] > L [ g2 ; π ] . We do not seek to provide an argument for the application of deep learning methods to MCMC proposal optimization , so we will simply assume that all objective functions of interest are well-behaved for optimization via deep learning techniques ( i.e . they are continuous and almost surely differentiable ) . We will also assume that all proposal and target densities considered are positive and non-singular . Let G be the set of allowed proposals to consider during optimization and let D be the group of almost sure diffeomorphisms over the space of our data . To perform optimization , we first select some T ∈ D that provides us with the coordinate system we will use for optimization . Defining : T ◦ f ( ~z ) = f ( T−1 ( ~z ) ) |∂T −1 ( ~z ) ∂~z | We finally optimize to find gopt = argmin g∈G L [ T ◦ g ; T ◦ π ] . The focus of this work is to illustrate the construction of objective functions such that , when the model class G is very expansive ( e.g . having been parameterized by a deep generative model ) , the optimized proposal gopt aligns with our notion of an efficient MCMC proposal , as discussed below . 4 PROPERTIES OF THE GROUND TRUTH OBJECTIVE FUNCTION L∗ . For us to sensibly pursue the task of MCMC proposal optimization , we must assume the existence of some ground truth objective , L∗ , that produces our notion of sampling performance . As previously stated , we currently do not know a universal definition for L∗ . We may , however , assume that the ground truth objective satisfies a number of first principles properties : L∗ is Proper Define an objective function to be proper if it attains a unique global minimum at g ( ~x′|~x ) = π ( ~x′ ) and αg , π ( ~x′|~x ) = 1 for all ~x′ , ~x . The overall goal of our MCMC methodology is to approximate perfect i.i.d . sampling from the target , π . We therefore find it uncontroversial to assume that L∗ is proper . L∗ is Representation Independent Define an objective function , L , to be representation independent over a group of almost sure diffeomorphisms T if , for all T ∈ T , and for all proposal distributions g1 , g2 , L [ g1 ; π ] > L [ g2 ; π ] if and only if L [ T ◦ g1 ; T ◦ π ] > L [ T ◦ g2 ; T ◦ π ] . Similarly , we will say L is representation invariant over T if L [ T ◦ g ; T ◦ π ] = L [ g ; π ] for all g. If L∗ were not representation independent over D , then our definition of sampling performance would depend on which T ∈ D/H is used when computing the optimization , whereH is the maximal subgroup of D over which L∗ remains representation independent . To fully justify an ordering , we would need to justify our selection of a particular member from D/H , which we should expect to be exceptionally burdensome . Of course , we usually have little prior justification for selecting a particular T and instead often use the coordinate system in which data was originally collected , perhaps applying some rescaling and recentering . Thus , unless contradicted by future experimental or theoretical results , we should assume that L∗ is at least representation independent . L∗ Yields Established Optimal Results Prior theoretical analysis has established certain properties regarding optimal proposal distributions for a number of MCMC schemes under diffusionary limits . If L∗ is to correspond to the same notion of sampling performance , it must yield the same results . Thus , we should expect that optimization of L∗ will recover the properties established within these theoretical works when applied in the same diffusionary limits .

Summary: The paper approaches the problem of learning the optimal parameters of the MCMC chain. Namely, it tries to design an objective function that allows for the learning of a perfect proposal distribution in the Metropolis-Hastings algorithm. The main contribution of the current paper is based on the developments from [Titsias, 2019]. [Titsias, 2019] proposes the following objective for the optimization of the proposal $$\text{GSM} = \text{Acceptance Rate} + \beta \cdot \text{Entropy of the proposal}.$$ Maximization of this objective results in better mixing properties of the chain. Indeed, the next sample comes from a distribution with high entropy and should be accepted with a high probability, and the proposal should converge to the independent sampler in the perfect scenario.   However, the choice of $\beta$ remains to be a question since it defines the tradeoff between the acceptance rate and the entropy, which is known to be different for different proposals and different targets. The original paper [Titsias, 2019] solves this issue by adaptively setting the acceptance rate to some desired value.  The current paper then can be summarized as follows. Instead of setting $\beta$ adaptively based on the acceptance rate, the authors propose to set it constant. The constant is then found by "manual trial and error" in such a way that the acceptance rate of Random-Walk Metropolis-Hastings is 0.234 for high-dimensional standard Gaussian, which is known to be optimal for this task. Then this constant is claimed to be universal for all targets and proposals. To verify this empirically, the authors first optimize the parameters of MALA and RW Metropolis-Hastings for several synthetic targets (Uniform, Laplace, Cauchy, Gaussian) and demonstrate that the procedure yields the right parameters.  Finally, the authors perform an empirical study of the proposed procedure. They use the proposed objective to learn a proposal parameterized by the normalizing flow and compare it to the optimization of the original objective from [Titsias, 2019] targeting different acceptance rates. The target distributions for the empirical study are a mixture of four Gaussians (2-D), Bayesian logistic regression (21-D, 14-D).  [Titsias, 2019]: Titsias, Michalis, and Petros Dellaportas. "Gradient-based adaptive markov chain monte carlo." Advances in Neural Information Processing Systems 32 (2019): 15730-15739.