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 proposes four desirable properties that should be satisfied by objective functions used to tune MCMC proposals. These are based on first principles, such as that the objective must have a unique global minimum when the proposal is equal to the target, must be representation invariant, and must recover known optimal results in well-studied settings (e.g. optimal acceptance rate for Langevin dynamics with a Gaussian target). Using these principles it proposes a new objective based on GSM (Titsias and Dellaportas, 2019) that satisfies all proposed principles.