S EMI -E MPIRICAL O BJECTIVE F UNCTIONS FOR N EURAL MCMC P ROPOSAL O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Current objective functions used for training neural MCMC proposal distributions implicitly rely on architectural restrictions to yield sensible optimization results which hampers the development of highly expressive neural MCMC proposal archi tectures
In this work, we introduce and demonstrate a semi-empirical procedure for determining approximate objective functions suitable for optimizing arbitrarily parameterized proposal distributions in MCMC methods.
Our proposed Ab Initio objective functions consist of the weighted combination of functions following constraints on their global optima and transformation invariances that we argue should be upheld by general measures of MCMC efﬁciency for use in proposal optimization.
Our experimental results demonstrate that Ab Initio objective functions maintain favorable performance and preferable optimization behavior compared to existing objective functions for neural MCMC optimization
We ﬁnd that Ab Initio objective functions are sufﬁciently robust to enable the conﬁdent optimization of neural proposal distributions parameterized by deep generative networks extending beyond the regimes of traditional MCMC schemes.
I NTRODUCTION
The development of efﬁcient 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 deﬁnition of the notion of MCMC efﬁciency, 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 deﬁned 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 deﬁnition for what we mean by MCMC efﬁciency.
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 deﬁnition of MCMC efﬁciency.
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 deﬁnition 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 speciﬁcally deﬁning sampling performance.
In light of the lack of an exactly speciﬁed objective function for MCMC efﬁciency, our Ab Initio objective functions seek to approximate the ground truth deﬁnition of sampling performance by adhering to certain reasonable ﬁrst principles properties and ﬁtting to reproduce the "mathematical observations" provided by existing theoretical analysis of optimal acceptance rates.
Our contributions are as follows:
• We describe a set of ﬁrst principles properties that may be reasonably assumed of the ground truth objective function underlying our notion of MCMC efﬁciency.
• We illustrate the construction of an example
Ab Initio objective function via the combination of simpler objective functions with coefﬁcients 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 deﬁned MCMC proposal distributions.
R ELATION TO P RIOR W ORK
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 deﬁning 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 modiﬁcation of
MSJD with a regularization intended to encourage mixing of the resulting Markov Chain.
As shown in our experiments, both MSJD and L2HMC’s modiﬁcation 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 speciﬁed 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 efﬁciencies of proposal distributions with differing acceptance rates outside of the context of parameter optimization.
MCMC P ROPOSAL O PTIMIZATION
In this work, we consider the task of optimizing a proposal density g θ ( (cid:126) x
(cid:48) | (cid:126) x ) ( (cid:126) x ∈ R n ) to sample from a target density π
( (cid:126) x ) for a MCMC task.
Proposals are accepted with rate α g θ ,π ( (cid:126) x (cid:48) | (cid:126) x ) that ensures the resulting Markov Chain converges towards π ( (cid:126) x ) .
For this work, we restrict ourselves to
Metropolis-Hastings type schemes, wherein α g θ ,π ( (cid:126) x (cid:48)
| (cid:126) x ) =
min { 1 , π ( (cid:126) x (cid:48) )
g θ
( (cid:126) x | (cid:126) x (cid:48) )
π ( (cid:126) x ) g θ ( (cid:126) x (cid:48) | (cid:126) x ) } .
For optimization, we utilize some measure of sampling performance to deﬁne an objective function L [ g ; π ] that imposes an ordering over proposal distributions by deﬁning g 1 <
g
exactly when L [ g 1 ; π ] > L [ g 2 ; π ] .
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 ﬁrst select some T ∈ D that provides us with the coordinate system we will use for optimization.
Deﬁning:
[Equation]
|
We ﬁnally optimize to ﬁnd g opt = 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 g opt aligns with our notion of an efﬁcient MCMC proposal, as discussed below.
P ROPERTIES OF THE G ROUND T RUTH O BJECTIVE F UNCTION 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 deﬁnition for L ∗ .
We may, however, assume that the ground truth objective satisﬁes a number of ﬁrst principles properties:
L ∗ is Proper Deﬁne an objective function to be proper if it attains a unique global minimum at g ( (cid:126) x (cid:48) | (cid:126) x )
=
π ( (cid:126) x (cid:48) )
and α g,π ( (cid:126) x (cid:48) | (cid:126) x ) = 1 for all (cid:126) x (cid:48) ,(cid:126) x .
The overall goal of our MCMC methodology is to approximate perfect i.i.d.
sampling from the target, π .
We therefore ﬁnd it uncontroversial to assume that L ∗ is proper.
L ∗ is Representation Independent Deﬁne 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 g 1 , g 2 , L [ g 1 ; π ] > L [ g 2 ; π ] if and only if L [ T ◦ g 1 ; T ◦ π ] > L [ T ◦ g 2 ; 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 deﬁnition of sampling performance would depend on which T ∈ D / H is used when computing the optimization, where H 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 justiﬁcation 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.
C ONSTRUCTING A B I NITIO O BJECTIVE F UNCTIONS
Without knowing the exact form of L ∗ , we must resort to ﬁnding a useful approximation.
Although the assumptions of being proper and representation independent limit the functional class to which
L ∗ belongs, this functional class remains quite expansive.
This situation is greatly simpliﬁed by restricting our consideration to representation invariant objective functions.
As shown in Appendix
A, the positive weighted combination of proper and representation invariant objective functions is itself a proper and representation invariant objective function.
This allows us to construct potential approximations of L ∗ by the combination of simpler objective functions as weighted by hyperparameter coefﬁcients.
These hyperparameter coefﬁcients can then be ﬁt so as to recover optimal properties established within existing theoretical works.
We call the resulting objective function an Ab Initio objective function.
For this work, we adapt the GSM into an Ab Initio objective function (details are provided in
Appendix B) to follow the constraints of Section 4 and consider a functional class of the form:
[Equation]
Where A is a hyperparameter coefﬁcient for ﬁtting and d is the dimensionality of the target.
In
Appendix C, we demonstrate that this functional class is proper and representation invariant.
F ITTING AND V ERIFYING A B I NITIO O BJECTIVE F UNCTIONS
There are many possible approaches to ﬁtting the coefﬁcients of an Ab Initio objective function to recover optimal results over reference problems.
General and principled methodologies for this ﬁtting are provided by procedures like stochastic Levenberg-Marquardt optimization (Bergou et al., 2018).
For simplicity, we determined A by manual approximate Newton-Raphson to match the theoretical result that RWM has an optimal acceptance rate of 0.234 when targeting a multivariate gaussian of zero mean and identity covariance.
As the theoretical results are set within the diffusionary limit of d → ∞ , we must match the theoretical result in a problem with sufﬁciently large dimension in practice.
Here, we perform this match to theoretical reference on a problem with 1000 dimensions.
We recovered the desired acceptance rate with A = 0 .
18125 and the Ab Initio objective function considered through the remainder of this work is therefore:
[Equation]
By ﬁtting the parameters of an Ab Initio objective function, we hope to approximate a general notion of MCMC efﬁciency beyond the particularities of the reference problem(s) used in its construction.
To verify this robustness, we utilized the Ab Initio function listed in Equation (2) to perform MCMC
[Table]
Dimensional Robustness
optimization in a number of tasks where the optimal acceptance rate is analytically known, with varying target distribution (independent gaussian, laplace, cauchy, and uniform), dimensionality (100-10,000 dimensions), and MCMC scheme (RWM with gaussian proposals and MALA).
These tasks are all within the diffusionary limit and involve a severely limited model class for proposal distributions (optimizing only proposal step size).
Under these limitations, acceptance rate adequately summarizes differences between optimized model parameters and MSJD adequately summarizes the efﬁciency of the resulting Markov Chains.
These measures are estimated for each optimized model based on 25,000 proposals from starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each veriﬁcation optimization is replicated a total 5 times.
These results are listed in Table 1.
Additional experimental details and results are provided in Appendix D.
We ﬁnd that the Ab Initio objective of Equation (2) exhibits good agreement with known analytical results beyond the reference problem used in its construction.
The best agreement, both in acceptance rate and resulting MSJD, is obtained in the RWM optimization tasks with continuous target densities.
Even though the objective was tuned to replicate optimal behavior when optimizing RWM for adimensional Gaussian target, the Ab Initio objective is also able to reproduce optimal behavior when optimizing a MALA proposal against a continuous target (recovering an acceptance rate around 0.5) and when optimizing RWM proposal against a discontinuous target (recovering an acceptance rate around 0.145).
With one exception (the 10000 dimensional MALA optmization experiment, where the difference in MSJD is less than 2%), the difference in MSJD perfomance between the Ab Initio and analytically optimized proposals is statistically insigniﬁcant.
These results demonstrate that the Ab Initio objective remains a useful approximation of our notion of MCMC efﬁciency across a wide range of optization problems within the diffusionary limit for which we have analytically known solutions.
At the same time, deﬁning the Ab Initio objective to be proper ensures, in principle, valid optimization behaviour in the opposite asymptotic limit of model complexity wherein i.i.d.
resampling from the target lies within the proposal distribution’s model class.
The Ab Initio objective is a sufﬁciently robust approximation to enable conﬁdent optimization of highly expressive neural MCMC proposal distributions.
E XPERIMENTAL R ESULTS
To demonstrate the advantages of Ab Initio objective functions, we present a series of illustrative experiments.
In these experiments, we compare objective functions on the basis of their capabilities for optimizing very general parameterizations of proposal distributions that should be of interest for
MCMC proposal optimization.
Throughout these experiments, we use effective sample size (ESS) as a measure of sampling performance.
We consider one-dimensional ESS relating to uncertainty estimating ﬁrst and second moments of the target distribution’s marginals (further explained in
Appendix L), reporting the minimum obtained across all dimensions.
For comparisons, we consider
MSJD optimization (MSJD Opt.)
(Pasarica and Gelman, 2010), L2HMC’s objective (L2HMC Obj.)
(Levy et al., 2018), and the Generalized Speed Measure (Titsias and Dellaportas, 2019) targeting acceptance rates of 0.9, 0.6, and 0.3 (GSM-90, -60, and -30).
As we are comparing the properties of the global optima of these objective functions, we estimate sampling from the target distribution by either direct sampling or by sampling form a long equilibrated HMC chain prepared before optimization, as opposed to the online, persistent single chain updates common used in many adaptive MCMC algorithms.
Details of the compared objective functions are provided in Appendix F. A comparison of our optimization procedure with traditional online adaptive techniques is provided in
Appendix M. Additional experimental details and results are provided in Appendices G, J and K.
7. O PTIMIZING M ULTI -S CHEME P ROPOSALS
We now introduce a normalizing ﬂow (Tabak and Turner, 2013; Rezende and Mohamed, 2015;
Papamakarios et al., 2019) based proposal distribution.
With T θ denoting a normalizing ﬂow’s parameterized transformation, the functions µ L,θ , µ D,θ , Σ L,θ , and Σ D,θ speciﬁed by neural networks, and (cid:12) denoting element wise multiplication, our proposal distribution is deﬁned via:
[Equation]
This implementation is fundamentally a conditional normalizing ﬂow (Winkler et al., 2019) and is similar to NeuTra-HMC (Hoffman et al., 2019), which employs HMC within the latent space of a pre-ﬁt normalizing ﬂow to improve MCMC efﬁciency.
This multi-scheme distribution differs from
NeuTra-HMC by allowing both the ﬂow’s transformation and the latent distribution to be optimized to improve MCMC efﬁciency, rather than being ﬁxed from the start.
In Appendix E, we demonstrate how this parameterization is able to approximate or recover various existing MCMC schemes.
The target distribution is an equal mixture of 4 standard gaussians positioned in a cross formation with a maximal distance of 8 between component centers.
We use the NICE architecture (Dinh et al., 2014) to deﬁne the normalizing ﬂow and vary the number of coupling layers (using either 8 orlayers), which inﬂuences how accurately the ﬂow is able to approximate the target distribution.
After each optimization, we determine efﬁciency measures on the basis of 5 Markov Chains of 1000 total proposals starting from independent samples from the target distribution.
Statistics regarding the mean and standard error of these measures are collected from 5 replications for each considered objective function.
The results of these optimizations are provided in Table 2.
Comparisons of the computational costs 4 for optimizing each objective are provided in Table 3.
Density plots of multi-scheme proposal distributions optimized using an MSJD objective and our Ab
Initio objective are provided in Figure 1.
Position dependence within the proposal distribution lead
MSJD and L2HMC’s modiﬁcation to optimize towards arbitrarily non-ergodic proposal distributions, sampling only either the horizontal or vertical components of the target distribution.
This experiment demonstrates a fundamental difﬁculty in applying the GSM objective to multi-scheme proposals.
Here, the ﬂow’s depth determines it’s maximal accuracy in approximating independent resampling of the target distribution, creating a maximum acceptance rate that can be obtained while still approximating global resampling.
Attaining greater acceptance rates is possible, but forces inefﬁcient RWM-like behavior.
In this case, 8 coupling layers within the ﬂow permits a maximal resampling acceptance rate near 0.70, while 3 coupling layers lowers this maximal rate below 0.60.
A
[Table]
.I
[Table]
D.
[Table] [Table]
[Figure]
user with 8 coupling layers who uses the GSM targeting an acceptance rate of 60% to optimize their distribution would fortuitously arrive at an efﬁcient proposal distribution. However, if they had used a
[Table] [Table]
model with 3 coupling layers instead and targeted the same 60% acceptance rate, the result would be a proposal distribution far more inefﬁcient than the actual capabilities of their proposal model class.
In both cases
[Table]
the Ab Initio objective is able to optimize towards efﬁcient approximations of i.i.d.
resampling.
An intuitive explanation for the results of this experiment is provided in Appendix H
[Table]
As this experiment utilizes a complicated proposal distribution
[Table]
it is reasonable to question whether the problems exhibited here arise from some peculiarities of the model class of Equation (3).
For this reason [Table] Appendix I repeats a similar experiment with a simpler proposal distribution and recovers similar results.
We conclude that the effects demonstrated within Table 2 may be traced back to the choice of objective function used for optimization.
7. O PTIMIZING A UGMENTED M ULTI -S CHEME P ROPOSALS
Augmentation with auxiliary variables (e.g. the momenta of HMC) is a common feature of MCMC schemes (Levy et al., 2018; Habib and Barber, 2018) and, in this experiment, we augment the original distributions with a number of independent auxiliary variables equal to the distribution’s original dimensionality.
In Appendix E.1, we explain how this augmentation enables the model class of Equation (3) approximate HMC.
To demonstrate the application of Ab Initio objectives to
MCMC optimization tasks of a more practical nature, we consider the optimization of the multischeme proposals deﬁned in Equation (3) for MCMC sampling of regression weights in bayesian logisitc regression for various UCI datasets (Bache and Lichman, 2013).
After each optimization, we determine efﬁciency measures on the basis of 5 Markov Chains of 20000 total proposals starting from samples drawn from a long, equilibrated Markov Chain obtained using tuned HMC.
Statistics regarding the mean and standard error of these measures are collected from 5 replications for each considered objective function.
The results of these optimizations are provided in Table 4.
Comparisons of the computational costs 5 for optimizing each objective are provided in Table 5.
Under these expanded conditions, we ﬁnd the objective functions exhibiting much of the same behavior as seen in Section 7.1, in particular with the Heart Disease dataset.
[Table]
7. A PPLICATION TO MCMC S CHEME S ELECTION
As a general measure of MCMC efﬁciency, the Ab Initio objective funtion of Equation (2) may be used to evaluate the performance of any MCMC proposal distributions.
In this experiment, we optimize a variety of different MCMC schemes for MCMC sampling of regression weights in bayesian logisitc regression between two MNIST (LeCun et al., 1998) digits.
For each evaluated scheme, after optimization using the respective objective function, we determine efﬁciency measures on the basis of 100 Markov Chains of 1000 total proposals starting from samples drawn from a long, equilibrated Markov Chain obtained using tuned HMC.
Statistics regarding the mean and standard error of these measures are collected from 5 replications for each considered objective function.
The results of these optimizations are provided in Table 6.
The resampling normalizing ﬂow (an unconditional normalizing ﬂow optimized to directly approximate the target distribution) did not produce a single accepted proposal within these evaluation Markov Chains, hence we are unable to estimate ESS and our estimates of acceptance rates and MSJD for this proposal distribution should be viewed as highly uncertain.
Although preconditioned MALA proposals have a restricted model class, optimization with MSJD yields exceptionally low minimum ESS results as sampling across certain subspaces of the distribution are sacriﬁced in order to maximize overall MSJD.
This problem is related to the representation dependence of MSJD, which is alleviated by the representation independence of Ab Initio objective functions as discussed in Section 4.
Overall, we ﬁnd the example
Ab Initio objective function of Equation (2) to be a robust measure of MCMC efﬁciency that may be used to evaluate and compare the efﬁciency of multiple MCMC schemes in addition to its utility as an objective function for optimizing the individual parameters of highly expressive MCMC schemes.
C ONCLUSION AND F UTURE W ORK
In this work, we have shown how to construct Ab Initio objective functions that are suited for the optimization of highly expressive proposal distributions.
We ﬁnd that Ab Initio objectives are suitably robust to enable the optimization of neural MCMC proposal distributions.
Our experimental results show that Ab Initio objective functions can maintain favorable performance and preferable optimization behavior compared to existing objective functions (Levy et al., 2018;
Titsias and
Dellaportas, 2019; Pasarica and Gelman, 2010).
By design, Ab Initio objective functions are approximations of our notion of MCMC efﬁciency and we do not presume that the particular example
Ab Initio objective function of Equation (2) will be absolutely universal.
However, should future experimental or theoretical analysis demonstrate some sub-optimal properties of Equation (2), our proposed Ab Initio procedure allows for further improvements by considering alternative component functions and coefﬁcient ﬁtting procedures, which we plan to explore in future work.
The experimental methodology this work is intended to isolate fundamental effects of the choice of objective function on MCMC optimization.
How to best train proposals distributions in an online adaptive setting, how to attain computationally efﬁcient samples, and how to handle scaling into higher dimensions remain important practical considerations for MCMC optimization.
We argue that these considerations are primarily inﬂuenced by proposal architecture and training procedure selection.
We will therefore investigate these questions in a future work focusing on comparisons among speciﬁc proposal architectures and training procedures.
A C OMBINING P ROPER AND R EPRESENTATION I NVARIANT F UNCTIONS
As part of our Ab Initio approach, we rely on the following proposition:
Proposition: Every positive weighted combination of proper and representation invariant objective functions is itself a proper and representation invariant objective function.
To prove this, assume that we have two proper and representation invariant (over a group of almost sure diffeomorphisms D ) objective functions f and h with target π and allowed set of proposals
G .
Let us select arbitrary constants c 1 , c 2 > 0 and consider the properties of the objective function c 1 f + c 2 h .
From the invariance of f and h we may immediately conclude:
[Equation]
Hence c 1 f + c 2 h is itself representation invariant.
Finally, for all g ( x (cid:48) | x ) (cid:54) =
π ( x (cid:48) ) :
[Equation]
Hence c 1 f + c 2 h is also proper.
This concludes the proof of the above proposition.
This offers a straightforward approach to approximating our ground truth objective L ∗ .
Although we do not have a guarantee that L ∗ lies within the class of representation invariant objective functions (we have only assumed representation independence), we also have too few analytical results to empirically conclude that L ∗ is not representation invariant.
As ﬁtting the coefﬁcients of an Ab Initio objective function to match analytical results is more easily accomplished than deriving a particular objective function that happens to match analytical results, we argue that our Ab Initio approach is a practical method for obtaining approximate objective functions suitable for the optimization of arbitrary MCMC proposal distributions.
Next, we slightly extend this result in order to justify the functions composing Equation 1 of the main paper.
Below, we start with a deﬁnition.
Deﬁnition:
An objective function is said to be nearly proper if it attains a (not necessarily unique) global minimum at g ( (cid:126) x (cid:48) | (cid:126) x )
=
π ( (cid:126) x (cid:48) )
and α g,π ( (cid:126) x (cid:48) | (cid:126) x ) = 1 for all (cid:126) x (cid:48) ,(cid:126) x .
We may now consider a minor generalization of the previous proposition:
Proposition: Every positive weighted combination of at least one proper and representation invariant objective function with any number of nearly proper and representation independent functions is itself a proper and representation invariant objective function.
The proof that the resulting combination is representation invariant follows from the representation invariance of the component functions, exactly as in the above.
To show that the combination is proper, assume that we have representation invariant (over a group of almost sure diffeomorphisms
D ) objective functions f and h with target π and allowed set of proposals G .
Let f be proper and let h be nearly proper.
Let us select arbitrary constants c 1 , c 2 > 0 and consider the properties of the objective function c 1 f + c 2 h .
For all g ( x (cid:48) | x ) (cid:54) =
π ( x (cid:48) ) :
( c 1 f + c 2 h )
[ π ; π ]
=
c
f
[ π ; π ] + c 2 h [ π ; π ] ≤ c 1 f [ π ; π ] + c 2 h
[ g ; π ] < c
f
[ g ; π ] + c 2 h [ g ; π ] = ( c 1 f + c 2 h )
[ g ; π ]
Hence c 1 f + c 2 h is also proper.
B D ERIVATION OF O UR E XAMPLE A B I NITIO O BJECTIVE F UNCTION
We begin with the deﬁnition of the (maximized) lower bound of the GSM, where B is a hyperparameter to be adapted during optimization to match the user speciﬁed acceptance rate:
[Equation]
Given the success of the GSM for optimizing proposal distributions when the user knows the optimal acceptance rates for the particular problem, these components are natural choices for inclusion into an Ab Initio objective function.
The addition of an expected cross entropy term, E (cid:126)
[Equation]
∼
[Equation]
( (cid:126) x )
[ E
[Equation]
x (cid:48) ∼
[Equation]
( (cid:126) x (cid:48) | (cid:126) x ) [log π ( (cid:126) x (cid:48) | (cid:126) x )] , and factoring out an overall negative sign yields an objective function to be minimized that follows the functional constraints listed in Section 4:
[Equation]
x ) [ D KL ( g ( (cid:126) x (cid:48) | (cid:126) x
[Equation]
( (cid:126) x (cid:48) )
[Equation]
− B
[Equation]
(cid:126) x (cid:48) ∼ g
( (cid:126) x (cid:48) | (cid:126) x ) [log α
( (cid:126) x (cid:48)
| (cid:126) x )] ]
However, we encounter an additional complication.
If E (cid:126) x ∼ π ( (cid:126) x ) [ D KL ( g ( (cid:126) x (cid:48) | (cid:126) x ) ||
π ( (cid:126) x (cid:48) ))] and E (cid:126) x ∼ π ( (cid:126) x ) [ E (cid:126) x (cid:48) ∼ g
( (cid:126) x (cid:48) | (cid:126) x ) [log α
( (cid:126) x (cid:48)
| (cid:126) x )] ] have differing dependencies on dimensionality as d → ∞ , then optimization with this Ab Initio objective function will incorrectly optimize towards acceptance rates of either 0 or 1 in the diffusionary limit for our chosen reference problem, as either of the two terms will completely dominate the other.
To prevent this, we must introduce a dimension dependent function
A ( d ) that balances the asymptotic behavior of the two components.
This yields an Ab Initio function of the form:
[Equation]
Future research may use asymptotic analysis to derive functional forms for C ( d ) that are best suited for exceedingly high dimensions.
In this work, we followed a procedure of manual trial and error and found that C ( d ) =
A d yields useful results at least up to 10 , 000 .
This completes our adaptation of the GSM into the Ab Initio objective function listed in Equation (1):
[Equation]
As the Ab Initio objective function of Equations (1) and (2) functionally differs from the GSM by the addition of a cross-entropy term and by a dimensional scaling of the multiplicative coefﬁcient, we should expect that their computational costs and complexities to be comparable.
This is supported by our results in Tables 3, 5, 12, and 14-18.
While the GSM is often slightly more computationally costly, we believe this is more due to our particular implementation of the online update of its hyperparameter for tuning to match the targeted acceptance rate as opposed to reﬂecting a fundamentally different complexity scaling of the GSM compared to the Ab Initio objective function.
C P ROPERTIES OF O UR E XAMPLE A B I NITIO O BJECTIVE F UNCTION
The functional form of our example Ab Initio objective function is given by:
[Equation]
Where C ( d ) is a dimension dependent positive constant.
We will now consider the following proposition:
Proposition: The example Ab Initio objective function, L [ g ; π ] , deﬁned above is a proper and representation independent objective function for all C ( d ) ≥ 0 .
Following the general restrictions we’ve used to deﬁne our optimization problem, let g and π be positive and smooth and let T denote any almost sure diffeomorphism.
Let X T and Z T denote the regions of the domain and range of T over which T is differentiable and invertible.
From the properties of KL-Divergence, we know that E (cid:126) x ∼ π
( (cid:126) x ) [ D KL ( g ( (cid:126) x (cid:48)
| (cid:126) x ) ||
π ( (cid:126) x (cid:48) ))] ≥ 0 and is equal to
if and only if g ( (cid:126) x (cid:48) | (cid:126) x ) =
π ( (cid:126) x (cid:48) ) for all (cid:126) x ,(cid:126) x (cid:48) .
From the deﬁnition of KL-Divergence and the law of the unconscious statistician, we have:
[Equation]
g
[Equation]
x
[Equation]
|
[Equation]
x ) || π
[Equation]
x
[Equation]
))] =
[Equation] [Equation]
[Equation]
x
[Equation]
◦
[Equation]
x
[Equation] [Equation]
Hence, E (cid:126) x ∼ π ( (cid:126) x ) [ D KL ( g ( (cid:126) x (cid:48) | (cid:126) x ) ||
π ( (cid:126) x (cid:48) ))] is a proper and representation invariant objective function.
Next, for all proposals g and targets π , the Metropolis-Hastings acceptance rate is guaranteed to fall in the range 0 ≤ α ( (cid:126) x (cid:48) | (cid:126) x ) ≤ 1 .
Additionally, if g ( (cid:126) x (cid:48) | (cid:126) x ) =
π ( (cid:126) x (cid:48) ) , α ( (cid:126) x (cid:48) | (cid:126) x )
= 1 .
This guarantees that E (cid:126) x ∼ π ( (cid:126) x ) [ E (cid:126) x (cid:48)
∼ g
( (cid:126) x (cid:48) | (cid:126) x ) [ − log α
( (cid:126) x (cid:48) | (cid:126) x )]] is nearly proper.
And from the deﬁnition of the MetropolisHastings acceptance rate and the law of the unconscious statistician, we have:
[Equation]
T
[Equation]
] =
[Equation]
=
[Equation]
| (cid:126) x )]]
And so E (cid:126) x ∼ π ( (cid:126) x ) [ E (cid:126) x (cid:48) ∼ g ( (cid:126) x (cid:48) | (cid:126) x ) [ − log α
( (cid:126) x (cid:48) | (cid:126) x )]] is a nearly proper and coordinate invariant objective function.
With C ( d ) ≥ 0 , our example Ab Initio objective function is therefore the positive weighted combination of a proper and representation invariant objective function with a nearly proper and representation invariant function.
From the previous section, we know that this example Ab Initio objective function is itself proper and representation invariant.
As representation invariance implies representation independence, this example objective function is also representation independent and thus belongs to the same functional class as our approximated ”ground truth” objective.
D A DDITIONAL D ATA FOR V ERIFICATION T ASKS
For veriﬁcation tests, we optimize objective functions for MCMC proposals targeting isotropic gaussian, laplace, cauchy, and uniform distributions with dimensionalities between 100 and 10,000.
For numerical stability, our target uniform distribution is a mixture distribution of a standard uniform distribution and a standard multivariate gaussian (with probabilities of 1 1+ e − 100 and 1 1+ e 100 respectively, to provide some non-zero likelihood beyond the support of the uniform distribution).
For optimization, 20000 gradient steps are taken, with each training batch taken from a single starting point sampled i.i.d.
from the target distribution from which 50 independent proposals from g θ ( x (cid:48) | x ) are generated to empirically estimate expected loss on the basis of 50 total samples.
The gradient update steps take the form of Algorithm 1 with M = 1 and N = 50 .
Optimization is performed using the Adam optimizer (Kingma and Ba, 2015) with a stepsize of 3e-4 (reduced to 3e-5, 3e-6, and 3e-7 for uniform targets of dimensions 100,1000, and 10000).
For optimizing GSM (Titsias and Dellaportas, 2019), β is initially set to d − 1 and ρ β is set to 0 .
Acceptance rate and MSJD are estimated for each optimized model based on 25000 proposals from starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each veriﬁcation optimization is replicated a total 5 times.
Tables 7, 8, 9, 10, and 11 report results from optimizing all objective functions we use in the main work for comparison with our example
Ab Initio objective function.
In these tables, simulated analytic results are also reported, which refer to acceptance rate and MSJD as estimated in the same manner for proposals with a ﬁxed step size determined to yield the analytically known optimal acceptance rates.
We were unable to perform gradient based optimization of MSJD and L2HMC’s objective when targeting the uniform distribution, and resorted to estimating their optima by ﬁtting curves of these objective functions with respect to acceptance rate.
For MSJD, the results reported for MSJD targeting uniform distributions are therefore the result of ﬁtting a parabola near the objective function’s maxima, with uncertainties reported following the typical propagation of error.
Even following this procedure,
L2HMC’s objective function is unable to produce non-trivial optima when targeting the uniform distribution and we therefore consider RWM optimization targeting the uniform distribution a failure case for L2HMC’s objective. with analytically known optimal properties.
As these veriﬁcation tasks are set within the diffusionary limit, MSJD recovers known optimal results.
Value means are reported to at most the ﬁrst signiﬁcant digit of standard error (reported in parentheses)
[Table] [Table]
[Table] [Table]
[Table]
E A PPROXIMATING MCMC S CHEMES WITH M ULTI -S CHEME P ROPOSALS
As part of our experiments, we utilize a normalizing ﬂow based multi-scheme proposal distribution, parameterized as:
[Equation]
Where T θ denotes a normalizing ﬂow’s parameterized transformation, the functions µ L,θ , µ D,θ , Σ L,θ , and Σ D,θ are speciﬁed by neural networks, and (cid:12) denotes element wise multiplication.
This section provides examples of how this multi-scheme distribution is able to approximate many different existing MCMC schemes.
Throughout, let π denote the target distribution, g θ denote the density of the resulting proposal distribution, and let n denote the density of a standard gaussian distribution with zero mean and identity covariance.
Independent Resampling: If T θ ◦ n ≈ π , µ L,θ = µ D,θ = (cid:126) 0 , and Σ L,θ = Σ D,θ = (cid:126) 1 , then:
[Equation]
Depending on the accuracy of the ﬂow’s approximation, the multi-scheme distribution approximately recovers i.i.d.
resampling from the target.
Isotropic
[Equation] [Equation]
x , Σ L,θ = (cid:126) 1 , and Σ D,θ = τ , then: (cid:126) x
(cid:48) | (cid:126) x ∼ N{ (cid:126) x ; τ I }
And the multi-scheme distribution recovers isotropic RWM with step-size τ .
Isotropic MALA: If T θ =
I , µ L,θ = (cid:126) 0 , µ D,θ ( (cid:126) x ) ≈ (cid:126) x
+ τ ∇ log π ( (cid:126) x ) , Σ L,θ = (cid:126) 1 , and Σ D,θ = √ 2 τ , then:
[Equation]
Depending on the accuracy of the approximation of µ D,θ ( (cid:126) x ) , the multi-scheme distribution recovers isotropic MALA with step-size τ .
Preconditioned Diffusions: If T θ ( (cid:126) z )
≈ A − 1 (cid:126) z , where A is an invertible preconditioning matrix, then appropriate choices of µ L,θ , µ D,θ , Σ D,θ , and Σ L,θ will approximately recover preconditioned
RWM and MALA.
Latent MALA:
If T θ ◦ n ≈ π , µ L,θ ( (cid:126) x ) ≈ (1 + τ ) (cid:126) x , µ D,θ = (cid:126) 0 , Σ L,θ = √ 2 τ , and Σ D,θ = (cid:126) 1 , then the multi-scheme distribution will mimic the normalizing ﬂow based proposal scheme of Hoffman et al.
(2019), with MALA used within the ﬂow’s latent space rather than HMC.
E. A UGMENTING V ARIABLES
One method of greatly expanding the model class of proposal distribution is to augment the distribution with auxiliary variables.
This is most notably accomplished with the canonical momenta in
HMC.
We will denote augmentation with a semicolon, x ; a indicates that x has been augmented with the auxiliary variables a .
Approximate HMC: If Σ D,θ ( (cid:126) x ; (cid:126) p ) (cid:12) T θ ( (cid:126) z (cid:48) )
≈ (cid:126)
and if µ D,θ ( (cid:126) x ) ≈ (cid:126) x int ; (cid:126) p int , where (cid:126) x int ; (cid:126) p int are the results of applying any desired integrator to the system then the ﬂow based proposal may approximate HMC by iteratively resampling (cid:126) p from a normal distribution, generating a new proposal, and then performing a Metropolis-Hastings correction.
Do note that we are not claiming that this approximation should be expected to be accurate, merely that it is mathematically possible within the ﬂow proposal’s model class.
Intriguingly, the generality of the multi-scheme proposal distribution enables the construction allows for the construction of generalizations of HMC by varying the number of auxiliary variables used.
These examples are not exhaustive.
Depending on implementation, these ﬂow-based multi-scheme distributions could also parameterize useful proposal distributions that are not covered by existing
MCMC schemes.
An appealing aspect multi-scheme distributions is that their optimization, in theory, solves the problem of selecting which MCMC scheme to apply to a particular problem, which otherwise usually requires some intervention from the user.
However, their expressiveness means that we cannot rely on guarantees that would be provided by restricting to individual MCMC schemes.
F D ETAILS OF C OMPARED O BJECTIVE F UNCTIONS
Throughout this work, we focus on comparing the properties of four objective functions for MCMC proposal optimization.
Example Ab
Initio From Equation (2), to be minimized:
[Equation]
MSJD As used by Pasarica and Gelman (2010), to be maximized:
[Equation]
L2HMC’s Modiﬁcation As proposed by Levy et al.
(2018), to be minimized:
[Equation]
Where σ 2 min is the smallest one-dimensional variance of the target distribution.
As we are primarily interested in behavior at stationarity, we omit the additional "burn-in" regularization suggested by.
Generalized Speed Measure As proposed by Titsias and Dellaportas (2019), to be maximized:
[Equation]
Where β is a hyper-parameter.
We follow the suggested method for updating β in an online manner.
Following our discussion regarding the asymptotic behaviors of the two components with respect to dimension, we found it beneﬁcial to initialize β to 1 d , rather than the default value of 1 .
After every parameter update, the β is updated following:
[Equation]
Where we use the suggested value of ρ β = 0 .
02 throughout.
G F URTHER D ETAILS FOR M ULTI -S CHEME P ROPOSAL O PTIMIZATION
Our experiments with multi-scheme proposals utilize a distribution speciﬁed as follows:
[Equation]
( (cid:126) x ) (cid:12) T θ ( (cid:126) z (cid:48) )
The target distribution is an equal mixture of 4 standard gaussians positioned in a cross formation with a maximal distance of 8 between component centers.
The additive functions, µ L,θ and µ D,θ , are speciﬁed by standard ReLU networks (input dimension 2, hidden dimension 16, 4 hidden layers, output dimension 2).
The multiplicative functions, Σ L,θ and Σ D,θ , are speciﬁed by standard ReLU networks (input dimension 2, hidden dimension 16, 4 hidden layers, output dimension 2) followed by component-wise exponentiation.
The normalizing ﬂow, T θ , follows the NICE architecture (Dinh et al., 2014), with additive coupling layers speciﬁed by standard ReLU networks (input dimension 1, hidden dimension 6, 3 hidden layers, output dimension 1).
For this experiment, differing numbers of coupling layers are considered ( L = 8 and L = 3 ).
For optimization, 20000 gradient steps are taken, with each training batch taken from 8 starting points sampled i.i.d.
from the target distribution from each of which 100 independent proposals from g θ ( x (cid:48) | x ) are generated to empirically estimate expected loss on the basis of 800 total samples.
The gradient update steps take the form of Algorithm 1 with M = 8 and N = 100 .
Optimization is performed using the Adam optimizer (Kingma and Ba, 2015) with a stepsize of 3e-4.
Optimized proposals that fall into local optima of the objective functions are discarded.
For optimizing the GSM (Titsias and Dellaportas, 2019), β is initially set to 0 .
5 and ρ β is set to 0 .
02 and initialization is varied to ensure the target acceptance rate is attained.
After optimization, acceptance rate, MSJD, and ESS are estimated for each optimized model based on 5 independent Markov Chains of 1000 proposals with starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each optimization is replicated a total 5 times.
Figures 2, 3, 4, and 5 show density plots of the multi-scheme proposals obtained following optimization using L2HMC’s objective Levy et al.
(2018) and the GSM Titsias and Dellaportas (2019) with varying target acceptance rates.
Figures 6, 7, 8, 9, 10, and 11 show density plots of the multi-scheme proposals obtained following optimization of all considered objective functions.
H F URTHER E XPLANATION OF M ULTI -S CHEME O PTIMIZATION R ESULTS
To provide an intuitive explanation of the results of Section 7.1, we will start by restating (in a more simpliﬁed form from the original in Section 3) the overall aim of MCMC proposal optimization, namely ﬁnding a solution to the optimization problem:
[Equation]
Where G is the model class of our proposal architecture, π is the target distribution, and L is the objective function chosen for the optimization.
With the experimental setup of Section 7.1, we are conﬁdent that the results in Table 2 reﬂect proposals that are very close to the global optima of the compared objective functions.
Figure 12 illustrates the placement of the global optima for MSJD within the model class of the multi-scheme proposal distribution of Equation 3 in a very abstract and simpliﬁed manner.
For simplicity, we also plot hypothetical locations for RWM sampling and i.i.d.
resampling from the target distribution, π .
Fundamentally, the highly expressive model class of the multi-scheme distribution offers many pathways to increase MSJD.
The end result is, unfortunately, a pathologically non-ergodic Markov
Chain.
Even though MSJD is very well aligned with our notion of MCMC efﬁciency in model class neighborhoods very close to preconditioned RWM, MALA, and HMC, it is clear that MSJD and its derivatives are not robust measures of MCMC efﬁciency for proposal architectures with model classes like those of our example multi-scheme distribution.
When utilizing the GSM, the optimization procedure is complicated slightly by targeting the userspeciﬁed acceptance rate.
Effectively, this restricts G to those proposal distributions within the model class that attain the user speciﬁed acceptance rate.
Figures 13 and 14 illustrate the placement of the global optima for the GSM when targeting acceptance rates of 30% and 60% using the multi-scheme
[Figure]
distribution of Equation 3 with 3 and 8 coupling layers.
In these ﬁgures, the orange lines represent the restrictions of the model classes that attain the desired acceptance rates.
[Figure]
There are two model class regimes that are most relevant for the optimization of the multi-scheme architecture with the GSM in Section 7.1: a neighborhood of primarily inefﬁcient, RWM-like proposals and a neighborhood of approximate i.i.d.
resampling proposals.
With 8 coupling layers, both regimes are accessible while maintaining acceptance rates of 30% and 60%.
In Figures Figures 13 (a) and (b), this is represented by the presence of two orange surfaces within G Multi − scheme Distribution , one in the lower-left (representing the RWM-like regime) and one in the upper-right (representing the approximate resampling regime).
When using 8 coupling layers, specifying target acceptance rates of either 30% or 60% will still include some portion of the approximate resampling regime of the multi-scheme architecture’s model class within the restriction over which the GSM optimizes.
In the end, with 8 coupling layers, GSM-30 and GSM-60 have global optima that lie within the approximate resampling regime of the multi-scheme architecture.
Moving on to the case with 3 coupling layers, the architecture’s ability to approximate resampling is diminished, which greatly impacts the qualitative behaviour of the restrictions of the model class at different target acceptance rates.
As depicted in Figure 14 (b), the case of optimizing a multischeme architecture with 3 coupling layers with the GSM targeting an acceptance rate of 30% is not signiﬁcantly changed, as there is there is a portion of the approximate resampling regime that lies within the restriction of attaining an acceptance rate of 30% and the global optimum of GSMis some point within this restriction of the approximate sampling regime.
However, when usingcoupling layers and targeting an acceptance rate of 60%, there is no member of the approximate sampling regime that falls into the restriction over which the GSM is optimized, leaving only inefﬁcient RWM-like proposals from which to choose the optimum of GSM-60 in this case.
Without signiﬁcant prior knowledge regarding the interaction between the proposal architecture and the target distribution or a signiﬁcant time investment testing many different acceptance rates, there is little that can be done to avoid the detrimental effect of choosing a highly sub-optimal target acceptance rate when optimizing model classes like the multi-scheme distribution using the GSM.
X .
Model A exhibits continuous behaviour, as we might expect from restricted architectures like preconditioned RWM.
Model B exhibits discontinuous behaviour, as exhibited by the multi-scheme architecture with 3 coupling layers around an acceptance rate of 55-56%.
The "drop-out" of the approximate resampling regime when using 3 coupling layers with the multischeme distribution and targeting an acceptance rate of 60% is reﬂective of a discontinuity of MCMC efﬁciency with respect to acceptance rate that should be expected of expressive neural MCMC architectures like the multi-scheme proposal distribution.
With very restricted architectures, like preconditioned RWM or MALA, our notion of maximal MCMC efﬁciency attained at a given acceptance rate remains continuous with respect to the acceptance rate.
In the nomenclature of
Section 4, L ∗ is very tame when optimized under the restriction of attaining a ﬁxed acceptance rate with the model classes of preconditioned RWM and MALA.
So when we ﬁnd RWM optimized to an acceptance rate of 20% or MALA optimized to an acceptance rate of 60%, we are reasonably conﬁdent that the proposals are close to optimality and are not drastically inefﬁcient.
But, when using a proposal architecture like the multi-scheme distribution, L ∗ can behave very sharply with respect to the target acceptance rate.
This is illustrated in Figure 15, where the maximal MCMC efﬁciency ( L ∗ ) at a target acceptance rate is plotted against the target acceptance rate for two proposal architectures (A and B).
In the neighborhood of acceptance rate X , Model A (corresponding to the multi-scheme architecture with 8 coupling layers above) behaves much as we would expect given our experiences with RWM and MALA, where minor errors in ﬁxing a target acceptance rate have minor impacts on the end efﬁciency.
For Model B (corresponding to the multi-scheme architecture with 3 coupling layers above), the interaction of the proposal architecture and the target distribution results in a discontinuity in optimized MCMC efﬁciency at a target acceptance rate X (in our case, corresponding to around 5556%).
Here, very minor differences in the target acceptance rate can result in an orders of magnitude change in MCMC efﬁciency for Model B. Models A and B in this example do not need to differ by their proposal architecture, but may instead differ by their target distribution, which makes it difﬁcult to translate the knowledge gained from successful optimization in one task to the next task with the same proposal architecture.
As robust approximations of our notion of MCMC efﬁciency throughout the expansive model classes of architectures like the multi-scheme distribution, Ab Initio objective functions allow us, in principle, to conﬁdently approach MCMC proposal optimization with highly expressive neural proposals as we would any other traditional deep learning task.
Figure 16 illustrates the placement of the global optima for an Ab Initio objective function within the model class of the multi-scheme proposal distribution of Equation 3.
In contrast to MSJD based objectives, the functional constraints of Section 4 and the ﬁtting and veriﬁcation procedures of Section 6 work to ensure that Ab Initio objective functions remain robust and align with our notion of MCMC efﬁciency throughout the model classes of very expressive architectures.
In contrast to the GSM, Ab Initio objective functions seek to directly approximate our notion of MCMC efﬁciency without relying on prior knowledge of optimal acceptance rates, which are very difﬁcult to determine for arbitrary neural proposal architectures and target distributions.
The overall result is that Ab Initio objective functions offer global optima for highly expressive neural proposal architectures that remain well aligned with our notion of
MCMC efﬁciency and are compatible with a traditional deep learning approach to optimizing neural architectures.
I P OSITION -D EPENDENT M IXTURE P ROPOSAL O PTIMIZATION
To examine general objective function behavior without additional complications arising from particular proposal architectures, we ﬁrst consider the optimization of proposals consisting of a mixture of gaussians (with position independent means and covariances) with position dependent component weights speciﬁed by a neural network.
This proposal distribution may be speciﬁed as follows:
[Equation]
The target distribution is an equal mixture of 4 standard gaussians positioned in a cross formation with a maximal distance of 8 between component centers.
The proposal distribution consists ofgaussian components.
The position-dependent weights, f θ ( (cid:126) x ) , are speciﬁed by a standard ReLU network (input dimension 2, hidden dimension 16, 4 hidden layers, output dimension 4) followed by a component-wise softmax layer.
Proposals are initialized near i.i.d.
sampling of the target.
For optimization, 40000 gradient steps are taken, with each training batch taken from a single starting point sampled i.i.d.
from the target distribution from which 50 independent proposals from g θ ( x (cid:48) | x ) are generated to empirically estimate expected loss on the basis of 50 total samples.
The gradient update steps take the form of Algorithm 1 with M = 1 and N = 50 .
Optimization is performed using the Adam optimizer (Kingma and Ba, 2015) with a stepsize of 3e-4.
For optimizing the GSM (Titsias and Dellaportas, 2019), β is initially set to 0 .
5 and ρ β is set to 0 .
After optimization, acceptance rate, MSJD, and ESS are estimated for each optimized model based on 5 independent Markov Chains of 1000 proposals with starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each optimization is replicated a total 5 times.
Figures 17, 18, 19, 20, 21, and 22 show density plots of the positiondependent weighting functions obtained following optimization using the Ab Initio objective function of Equation (2), MSJD (Pasarica and Gelman, 2010), L2HMC’s objective (Levy et al., 2018) and the GSM (Titsias and Dellaportas, 2019) with varying target acceptance rates.
The results of these optimizations are provided in Table 12.
[Table]
We ﬁnd that i
[Table]
i
[Table]
d. resampling remains a stable global minimum of our Ab Initio objective function.
The
GSM objective also optimizes towards a constant and equal weighting of the proposals components, but increases the variance of each component to match its targeted acceptance rate.
Both MSJD and
L2HMC’s modiﬁcation ﬁnd global minima at arbitrarily non-ergodic proposals.
Both of these MSJD based objectives result in Markov Chains that only effectively explore half of the target’s probability mass (either the horizontal or vertical components), which is numerically reﬂected by the poor ESS performance for estimating second moments.
J F URTHER D ETAILS FOR L OGISTIC R EGRESSION P OSTERIOR S AMPLING
For this experiment, we augment the target distribution with a number of auxiliary variables.
Similar to the momenta in HMC, these augmenting variables are independent of the original data variables and are distributed following a standard multivariate gaussian.
Proposal distributions are optimized to target the joint distribution of original and augmenting variables.
For sampling, the augmenting variables are resampled before making each proposal, similar to the procedure employed in HMC.
With one exception (the Ripley dataset) we found best results by using a number of augmenting variables, a , equal to twice the dimensionality of the original data, d .
This experiment again uses multi-scheme proposals, speciﬁed as follows:
[Equation]
The target distribution is the posterior distribution of regression weights for logistic regression of various datasets from the UCI repository (Bache and Lichman, 2013) and from Ripley (2007).
The prior distributions used for regression weights are independent gaussians with a mean of 0 and a standard deviation of 10.
The additive functions, µ L,θ and µ D,θ , are speciﬁed by standard ReLU networks (input dimension a + d , hidden dimension (W) 3 ( a + d ) , 4 hidden layers (D), output dimension a + d ). The multiplicative functions, Σ L,θ and Σ D,θ , are speciﬁed by standard ReLU networks (input dimension a + d , hidden dimension (W) 3 ( a + d ) , 4 hidden layers (D), output dimension a + d ) followed by component-wise exponentiation.
The normalizing ﬂow, T θ , follows the NICE architecture (Dinh et al., 2014), with additive coupling layers speciﬁed by standard ReLU networks (input dimension a + d 2 , hidden dimension (W) 3 ( a + d ) , 4 hidden layers (D), output dimension a + d 2 ).
For this experiment, the same number of coupling layers is used ( L = 5 ) for all datasets.
Table 13 summarizes the network parameters used in this experiment.
Before each optimization, a long equilibrated Markov Chain is obtained following 500,000 proposals of HMC tuned to attain an acceptance rate of 0.65 (using hamiltorch’s (Cobb et al., 2019) implementation of HMC).
A burn-in period of 1000 proposals is utilized for this HMC tuning and is discarded.
For optimization, 40000 gradient steps are taken, with each training batch taken from 8 starting points sampled randomly from HMC Markov Chain and from each starting point 100 independent proposals from g θ ( x (cid:48) | x ) are generated to empirically estimate expected loss on the basis of 800 total samples. The gradient update steps take the form of Algorithm 1 with M = 8 and N = 100 .
Uniform sampling from this Markov Chain approximates i.i.d sampling from the target.
Optimization is performed using the Adam optimizer (Kingma and Ba, 2015) with a stepsize of 3e-4.
For optimizing GSM (Titsias and Dellaportas, 2019), β is initially set to d − 1 and ρ β is set to 0 .
02 and initialization is varied to ensure the target acceptance rate is attained.
After optimization, acceptance rate, MSJD, and ESS are calculated solely based on original data dimensions and are estimated for each optimized model based on 5 independent Markov Chains of 20000 proposals with starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each optimization is replicated a total 5 times.
Tables 14, 15, 16, 17, and 18 summarize the sampling performance measures obtained following optimization, including minimum, median, and maximum one-dimensional ESS obtained across all dimensions of the original data variables.
Minimum one-dimensional ESS is most indicative of sampling efﬁciency across the entire state space of the distribution and is therefore the focus of our comparisons.
[Table] [Table]
[Table] [Table]
[Table] [Table]
The results for the Heart, Pima, and Ripley datasets exhibit the same behavior observed within the experiments involving the gaussian mixture target in Section 7.1.
Optimizing MSJD or L2HMC’s objective yields proposals whose behavior signiﬁcantly deviates from i.i.d.
resampling from the target distribution and result in a proposal MSJD much larger than what would be obtained following i.d.d.
resampling from the target.
As seen in Tables 16 and 17, the target distributions with the Heart and
Pima datasets are similar to those with the gaussian mixture target in that this behavior is numerically evidenced by a severe decrease in ESS relating to estimation of the distribution’s second moments.
From Table 18, we see that the target distribution with the Ripley dataset is such that this behavior is numerically evidenced by ESS calculations signiﬁcantly greater than 1, which is indicative of the proposal distribution yielding greatly negative auto-correlations.
These results demonstrate that this behavior is not the result of speciﬁc choices in proposal architecture or target distribution.
We conclude that this sub-optimal behavior when maximizing MSJD or optimizing L2HMC’s objective occur when the proposal distribution allows sufﬁciently complex (relative to the target distribution) position dependence.
K F URTHER D ETAILS FOR S CHEME S ELECTION E XPERIMENT
For this experiment, the target distribution is the posterior distribution of regression weights for logistic regression of 5’s and 6’s from the training set of MNIST (LeCun et al., 1998) digits (d=785).
The prior distributions used for regression weights are independent gaussians with a mean of 0 and a standard deviation of 10.
Our proposal distributions are diagonally preconditioned MALA, diagonally preconditioned RWM, the multi-scheme proposal distribution of Equation (3), and a normalizing ﬂow intended to solely approximate i.i.d.
resampling from the target.
Before each optimization, a long equilibrated Markov Chain is obtained following 500,000 proposals of HMC tuned to attain an acceptance rate of 0.65 (using hamiltorch’s (Cobb et al., 2019) implementation of HMC).
A burn-in period of 1000 proposals is utilized for this HMC tuning and is discarded.
For optimization, 40000 gradient steps are taken, with each training batch taken from 1 starting point sampled randomly from HMC Markov Chain and from each starting point 1 proposal g θ ( x (cid:48) | x ) is generated to empirically estimate expected loss. The gradient update steps take the form of Algorithm 1 with M = 1 and N = 1 .
Uniform sampling from this Markov Chain approximates i.i.d sampling from the target.
Optimization is performed using the Adam optimizer (Kingma and Ba, 2015).
After the initial 40000 steps, the Ab Initio objective function of Equation (2) is estimated followingproposals from starting points independently sampled from the long equlibrated HMC chain.
Our rough, upper-bound estimates for the MSJD and acceptance rates of the resampling normalizing were estimated using the same procedure as with the Ab Initio objective function.
For optimizing both the preconditioned MALA and RWM proposal distributions, a stepsize of 3e-4 is used.
One of the preconditioned MALA distributions and the preconditioned RWM distribution are optimized using the Ab Initio objective function of Equation (2).
One of the preconditioned MALA distributions and is optimized to maximize MSJD.
For the multi-scheme proposal distribution, the additive functions, µ L,θ and µ D,θ , are speciﬁed by standard ReLU networks (input dimension 785 , hidden dimension (W) 2355 , 4 hidden layers (D), output dimension 785 ).
The multiplicative functions, Σ L,θ and Σ D,θ , are speciﬁed by standard ReLU networks (input dimension 785 , hidden dimension (W) 2355 , 4 hidden layers (D), output dimension 785 ) followed by component-wise exponentiation.
The normalizing ﬂow, T θ , follows the NICE architecture (Dinh et al., 2014), with additive coupling layers speciﬁed by standard ReLU networks (input dimension 392 , hidden dimension (W) 2355 , 4 hidden layers (D), output dimension 392 ).
The ﬂow uses 5 coupling layers A stepsize of 3e-5 is used.
The multi-scheme proposal distribution is optimized using the Ab Initio objective function of Equation (2).
The resampling normalizing ﬂow follows the NICE architecture (Dinh et al., 2014), with additive coupling layers speciﬁed by standard ReLU networks (input dimension 392 , hidden dimension (W) 3140 , 5 hidden layers (D), output dimension 392 ).
The ﬂow uses 6 coupling layers A stepsize of 3eis used.
The resampling normalizing ﬂow is optimized to optimize negative log-likelihood.
After optimization, acceptance rate, MSJD, and ESS are calculated solely based on original data dimensions and are estimated for each optimized model based on 100 independent Markov Chains of 1000 proposals with starting points independently sampled from the target distribution.
To gather statistics regarding the mean and standard error of reported variables, each optimization is replicated a total 5 times.
L C ALCULATING AND E VALUATING E FFECTIVE S AMPLE S IZE
Throughout our experiments, we use ESS as a means of verifying the sampling performance of optimized proposal distributions.
From a series of samples drawn sequentially from a Markov Chain, ( X 1 , X 2 , ..., X N ) , the simplest estimation of one-dimensional ESS can be obtained by the single chain calculation:
[Equation]
Where ρ t is the autocorrelation at lag t within the series ( X 1 , X 2 , ..., X N ) .
For slowly mixing
Markov Chains, more accurate estimations of ESS can be obtained by using multi-chain formulations that utilize multiple series of samples and considers both within-chain and between-chain variance (Vehtari et al.
(2019) provide a summary of these multi-chain methods).
In our experiments, we utilize
PyMC’s (Patil et al., 2010) calculation of ESS, which considers both within-chain and between-chain variance.
Although ESS is often viewed as providing a measure of how many effectively independent samples are provided within the series ( X 1 , X 2 , ..., X N ) , it is important to keep in mind that this interpretation applies only for the purpose of estimating the particular expectation E [ X ] .
In general, calculating
ESS from the series ( f ( X 1 ) , f ( X 2 ) , ..., f ( X N )) relates to the uncertainty of estimating E [ f ( X )] based on the series ( X 1 , X 2 , ..., X N ) .
As these ESS results can vary dramatically depending on the expectation of interest, E [ f ( X )] , a single ESS estimate does not completely summarize the ergodicity of an MCMC sampler for use in general calculations.
For our evaluations, we consider ESS relating to estimating ﬁrst and second moments of the target distribution ( E [ X ] and E [ X 2 ] ).
When evaluating the performance of restricted model class MCMC propsoals like RWM or MALA,
ESS relating to estimating ﬁrst moments ( E [ X ] ) is often a reliable performance comparison and ESS per proposal results are typically found to be less than 1.
Our results in Section 7.1 show that these expectations will not generally hold true when evaluating the performance of very expressive proposal distributions.
As seen with proposals optimized with MSJD and L2HMC’s objective, very expressive neural proposals can exhibit signiﬁcant negative autocorrelations that yield ESS per proposals greater than 1 and can have ﬁrst moment ESS results that mask signiﬁcant non-ergodicities of the evaluated proposal distribution.
When evaluating the performance of highly expressive neural MCMC proposals, additional consideration is required regarding what ESS truly measures and we cannot solely rely on the intuitions built from the applications of ESS to MCMC proposal distributions with very restricted model classes. M O PTIMIZING MCMC O BJECTIVE F UNCTIONS AND O NLINE A DAPTATION
Throughout our experiments, we utilize either exact or approximate samples from π ( (cid:126) x ) to approximate the outer expectations ( E (cid:126) x
∼ π ( (cid:126) x ) [ ... ] ) found within all compared objective functions during optimization.
As a result, our gradient update steps take the form illustrated in Algorithm 1 (note, when direct sampling from π ( (cid:126) x ) is not available, we approximate this sampling by uniformly sampling from a long equilibrated HMC chain).
Algorithm 1 requires samples from π ( (cid:126) x ) to be available during optimization, which is typically not the case during practical applications of MCMC proposal optimization.
For practical applications where samples from π ( (cid:126) x ) are not available, online adaptation procedures are used to approximate the outer expectations ( E (cid:126) x ∼ π
( (cid:126) x ) [ ... ] ), with gradient update steps taking a form like that illustrated in Algorithm 2.
We must emphasize that we are not advocating for the use of update steps like Algorithm 1 in practice.
We are instead using Algorithm 1 speciﬁcally because it allows us to most directly answer the core question of this work: determining whether the compared objective functions have optima that are aligned with our notion of MCMC efﬁciency when optimizing proposal distributions with highly expressive model classes like that of Equation (3).
Input: Target distribution π ( (cid:126) x ) , proposal distribution g θ ( (cid:126) x (cid:48) | (cid:126) x ) , objective function
L [ π ; g θ ]( (cid:126) x ,(cid:126) x (cid:48) ) , integers M, N ≥ 1 , optimizer ParamOpt ( ˆ L , θ ) .
for i in range (0 , M ) do
Sample (cid:126) x i ∼ π
( (cid:126) x ) ; for j in range (0 , N ) do
Sample (cid:126) x (cid:48)
j ∼ g θ ( (cid:126) x (cid:48) | (cid:126) x i ) ; end end ˆ L ← 1 MN (cid:80) i,j L [ π ; g θ ]( (cid:126) x i ,(cid:126) x (cid:48) j ) ; θ ← ParamOpt ( ˆ L , θ ) ;
For the purposes of our experiments, we do not employ online adaptation because the adaptation procedures introduce an additional confounding factor to the analysis of performance differences between objective functions.
By using Algorithm 1, our experiments demonstrate performance differences between the compared objective functions that can be directly attributed to the objective functions themselves and their interactions with the model classes of highly expressive proposal
Algorithm 2:
Input: Target distribution π ( (cid:126) x ) , proposal distribution g θ ( (cid:126) x (cid:48) | (cid:126) x ) , objective function
L [ π ; g θ ]( (cid:126) x ,(cid:126) x (cid:48) ) , integers M, N ≥ 1 , M initial Markov Chain states (cid:126) x i , optimizer
ParamOpt ( ˆ L , θ ) .
for i in range (0 , M ) do for j in range (0 , N ) do
Sample (cid:126) x (cid:48)
j ∼ g θ ( (cid:126) x (cid:48) | (cid:126) x i ) ; end end ˆ L ← 1 MN (cid:80) i,j L [ π ; g θ ]( (cid:126) x i ,(cid:126) x (cid:48) j ) ; θ ← ParamOpt ( ˆ L , θ ) ; for i in range (0 , M ) do
Sample (cid:126) x (cid:48) ∼ g θ ( (cid:126) x (cid:48) | (cid:126) x i ) ; α ← min { 1 , π ( (cid:126) x (cid:48) )
g θ
( (cid:126) x
i | (cid:126) x (cid:48) )
π ( (cid:126) x i )
g θ
( (cid:126) x (cid:48) | (cid:126) x i ) } ;
Sample
; if u < α : then end distributions.
The use of adaptation procedures like Algorithm 2 would introduce uncertainty whether such performance differences were the result of the objective functions themselves and raises the additional question of whether these adaptive procedures are compatible, without modiﬁcation, with neural MCMC proposal model classes like Equation (3).
It is important to note that adaptation procedures like Algorithm 2 have been developed and validated in the context of relatively restricted proposal model classes.
Our results have shown that even objective functions for MCMC efﬁciency, when developed under assumptions of proposal model class restrictions, can be fundamentally incompatible for optimizing proposal distributions with highly expressive model classes.
The problem of model class compatibility also applies to the development of online adaptation procedures.
It is still an open question for future research whether these adaptive procedures, originally designed for use with relatively restricted model class proposals, can be applied without modiﬁcation to the optimization of proposal model classes like Equation (3) or whether some modiﬁcation is required to maintain compatibility with highly expressive proposals.
M. R EPLICATION OF S ECTION 7.2 USING O NLINE A DAPTATION
Although the applicability of adaptive procedures like Algorithm 2 to highly expressive proposal distributions remains an open question, we may still test whether the example Ab Initio objective function of Equation (2) is compatible for optimization using these adaptive procedures.
To this end, we perform a replication of the experiments of Section 7.2 using the adaptive procedure of
Algorithm 2.
Aside from the use of Algorithm 2 instead of Algorithm 1, all experimental details remain unchanged from those listed in Appendix J, unless otherwise noted.
For optimization,gradient steps are taken in all cases and Algorithm 2 is used, with M = 16 for the German Credit dataset and M = 8 for the Heart dataset.
The initial states of the persistent Markov Chains are drawn uniformly from the HMC chains originally obtained as speciﬁed in Appendix J. In practice, these sampled initial states may be easily obtained via bootstrapping (as used in Song et al.
(2017)), either from a less optimized proposal distribution of the same model class or from a traditional MCMC sampling methodology like HMC (as used here).
Our experimental results are listed in Table 19.
[Table]
These results using online adaptation are consistent with those obtained using Algorithm 1 previously listed in Section 7.2.
In this experiment, the example Ab Initio objective function of Equation (2) is compatible for optimization using the online adaptive update of Algorithm 2.
We must still emphasize that future work is required to verify whether online adaptive procedures like Algorithmare
generally compatible with highly expressive proposal model classes
With this in mind, we have two conclusions regarding our results and the employment of adaptive procedures to the optimization of neural MCMC proposals:
• By using Algorithm 1 in lieu of an adaptation procedure and by investigating fundamental properties of MCMC objective functions, our results regarding the compatibility of MCMC objective functions with highly expressive proposal model classes will remain relevant to
future research regarding neural MCMC proposals, regardless of the particular adaptation procedures used in practice or future advancements of the technique of online adaptation.
• The example Ab Initio objective function of Equation (2) is compatible with online adaptive procedures like Algorithm 2, at least for the tasks considered.
• Future research of adaptation procedures like Algorithm 2 is needed to validate their applicability to the optimization of highly expressive proposal distributions.
Until such validation is performed, it is difﬁcult to interpret the results of experiments using these adaptive procedures (unmodiﬁed from their use with very restricted proposal model classes) for the purpose of comparing the performance of MCMC objective functions and expressive proposal architectures.