Semi-Empirical Objective Functions for Neural MCMC Proposal Optimization
Authors: Chris Cannella
Vahid Tarokh
Venue: ICLR 2022 Submitted (2022)
Keywords: Markov Chain Monte Carlo
Neural MCMC
Generative Models
Deep Generative Models
Normalizing Flows

Abstract: 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 architectures.  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 efficiency 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 find that Ab Initio objective functions are sufficiently robust to enable the confident optimization of neural proposal distributions parameterized by deep generative networks extending beyond the regimes of traditional MCMC schemes.

Under review as a conference paper at ICLR 2022
SEMI-EMPIRICAL OBJECTIVE FUNCTIONS FOR NEURAL MCMC PROPOSAL OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Current objective functions used for training neural MCMC proposal distributions implicitly rely on architectural restrictions 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 efficiency for use in proposal op- timization. Our experimental results demonstrate that Ab Initio objective functions maintain favorable performance and preferable optimization behavior compared existing objective functions for neural MCMC optimization. We find that Ab Initio objective functions are sufficiently robust to enable the confident optimization of neural proposal distributions parameterized by deep generative networks extending beyond the regimes of traditional MCMC schemes.
INTRODUCTION
The development of efficient Markov Chain Monte Carlo (MCMC) proposal distributions 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 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 al. (001 Roberts and Rosenthal (007)) focus 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 the problem of MCMC proposal optimization. Recent research regarding applications of deep learning to proposal optimization has yielded new MCMC schemes (Song
al. 2017, Spanbauer et al., 2020; Hoffman al., 2019) and extensions of existing schemes (de Freitas et al. [201]] 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 eurrently used for MCMC proposal optimization rely on model class restrictions imposed on the proposal architecture and 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 remain compatible with any proposal distribution defined using deep generative models.
Presumably, there exists some "ground truth" objective function underlying notion of MCMC sampling performance. However, even after the decades of rescarch regarding MCMC methods, there appears be no universal definition for what mean by MCMC efficiency. Metrics like effective

---

Under review as a conference paper at ICLR 2022
sample size (ESS) (Gelman 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 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 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 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 pa- rameterized by deep generative models. Research into the applications of deep learning for MCMC proposal optimization (Song et al. [ et al.] 2020 Hoffman et al. 2019 de Freitas et al., 2001 Habib and Barber 2018 Levy et al. 2018 Li et al. (021) has yielded 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, 20188) Neklyudov 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, 2 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 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 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
2

---

Under review as a conference paper at ICLR 2022
from the target their unique global minima. Another potential advantage of the Ab Initio objective of Equation 2 over MSJD based objective functions is that does not rely on the notion metric of the underlying space of the proposal distribution, and may be more suited optimizing proposals in settings where the space equipped with only 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 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 given optimization problem, which will generally be unknown when using very general neural architectures. Our example Ab Initio objective function of Equation ① takes inspiration from the components of the GSM, while also ensuring the functional limitations argued for in Section 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
R") ) to sample from target density π (X) for MCMC task. Proposals are accepted with rate O'ge, (x" (x) that ensures the resulting Markov Chain converges towards π (X). For this work, we restrict ourselves to Metropolis-Hasting type schemes, wherein Cg x min (X
In this work, we consider the task of optimizing a proposal density go(X² (X
((x x') ). For optimization,
we utilize some measure of sampling performance define an objective function Lig: x] that imposes an ordering over proposal distributions by defining 91 < 92 exactly when Llgi ; L'g2; T.. We do not seek to provide an argument for the application of deep learning methods to MCMC proposal optimization, we will simply assume that all objective functions of interest 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 allowed proposals to consider during optimization and let D be the group of almost sure diffeomorphisms
the space of our data. To perform optimization, we first select some T E DD that provides us with the coordinate system we will use for optimization. Defining:
8T
To To 1 ((T-
Oz
We finally optimize find Jopt argmin CIT g: o m..
gEG
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 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 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 number of first principles properties:
L is Proper Define an objective function to be proper if it attains a unique global minimum at x(X'|xx) (xX') and az (x' x) 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.
3

---

Under review as a conference paper at ICLR 2022
L is Representation Independent Define objective function, L, to be representation inde- pendent over group of almost sure diffeomorphisms T if, for all E T, and for all proposal distributions 91 92. : L(o2; n|] if and only if g o > L 92: T o rr. Similarly, Llg; и for all g. If L were not D2 then our definition of sampling performance would depend on which I' E D/H. is used when computing the optimization, where H. is the maximal subgroup of D over which C 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.
we will say & is representation invariant over T if L o g: 7 representation independent
L Yields Established Optimal Results Prior theoretical analysis has established certain proper- ties regarding optimal proposal distributions for a number of MCMC schemes under diffusionary limits. If C' 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.
5 CONSTRUCTING AB INITIO OBJECTIVE FUNCTIONS
Without knowing the exact form of L" must resort to finding 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 simplified by restricting our consideration to representation invariant objective functions. As shown in Appendix the positive weighted combination of proper and representation invariant objective functions is itself proper and representation invariant objective function.
This allows to construct potential approximations of L by the combination of simpler objective functions as weighted by hyperparameter coefficients. These hyperparameter coefficients can then be fit 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:
Cig: π Ex)n(s)) DKL (x | (x)))
A dd ~ x) llog [x(
~g(x
(1)
Where A is hyperparameter coefficient for fitting and d is the dimensionality of the target. In Appendix C we demonstrate that this functional class is proper and representation invariant.
FITTING AND VERIFYING AB INITIO OBJECTIVE FUNCTIONS
There are many possible approaches to fitting the coefficients of an Ab Initio objective function to recover optimal results over reference problems. General and principled methodologies for this fitting are provided by procedures like stochastic Levenberg-Maun Marquardt optimization (Bergou et al.) 2018).
For simplicity, we determined A by manual approximate Newton Raphson result that RWM has an optimal acceptance rate of 0.234 when targeting a multivariate gaussian of within the diffusionary limit
Newton-Raphson to match the theoretical
zero mean and identity covariance. As the theoretical results of d - 00, we must match the theoretical result a problem with sufficiently large dimension in practice. Here, perform this match to theoretical reference on 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:
Cig: m) (x)1 DKL (x)
0.18125 Ex (xx [lo (x'
By fitting the parameters of Ab Initio objective function, hope approximate general notion of MCMC efficiency 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) perform MCMC
4

---

Under review as a conference paper at ICLR 2022
Table 1: Verification results from applying Ab Initio objective function of Equation 2 multiple MCMC optimization tasks with analytically known optimal properties. The reference problem used for coefficient fitting is shown in bold. Value means are reported at most the first significant digit of standard error (reported in parentheses).
d 100
d 1000
10000
Initio
Analytic Simulated
Initio
Analytic Simulated
Initio
Analytic Simuluted
Analytic Exact
Eq. 2
Eq. 2
Eq. 2
Reference Problem
Ace. Rate
0.246(2)
0.236(1)
0.233(5)
0.234(1)
0.228(6)
0.235(2)
0.234 M
Gaussian
RWM
MSJD 1.32(2)
1.32(1)
1.32(1)
1.324(4)
1.33(1)
1.33(1)
Scheme Robustness
Ace. Rate
0,546((5
0.572(2)
0,503(5)
0.572(2)
0.495(9)
0.573(3)
0.574
Gaussian
MALA
MSID
38.6(2)
38.4(2)
166(1)
167(1)
738(6)
748(3)
Target Rohustness
Uniform
RWM
Ace. Rate
0.146(9)
0.135(2)
0.146(1)
0.134(2)
0.144(9)
0.137(2)
0.135
MSJD
8.2(1))-3
8.2(1)e-3
8.5(1))-4
8.4(1)e-4
8.5(2)e-5
8.6(1)e-5
Ace. Rate MSID
0,231(4)
0.234(3)
0,226(4)
0.234(2)
0,229(2)
0.235(2)
0.234
Laplace
RWM
1.47(2)
1.48())
1,37())
1,37(1)
1.33(2)
1,34(1)
Acc. Rute 0,237(3)
0.235(2)
0.236(5)
0.234(1)
0,233(3)
0.233(2)
0.234
Cauchy
RWM
MSJD 2.72(3)
2.72(3)
2.67())
2.64(1)
2.67(3)
2.61(2)
Dimensional Robustness
optimization in a number of tasks where the optimal acceptance rate 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 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 efficiency 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 verification optimization replicated a total 5 times. These results are listed in Table Additional experimental details and results provided in Appendix D.
We find 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 dimensional Gaussian target, the Ab Initio objective 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. 0.14 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 statistically insignificant.
tuned to replicate optimal behavior when optimizing RWM for 1000
These results demonstrate that the Ab Initio objective remains useful approximation of our notion of MCMC efficiency across a wide range of optization problems within the diffusionary limit for which we have analytically known solutions. At the same time, defining the Ab Initio objective to be proper ensures, in principle, valid optimization behaviour the opposite asymptotic limit of model complexity wherein i.i.d. resampling from the target lies within the proposal distribution's 's model class. The Ab Initio objective is a sufficiently robust approximation to enable confident optimization of highly expressive neural MCMC proposal distributions.
Roberts 2001)
Gelman et al. (1977)
Roherts et al.
2001
Roberts and Rosenthal 1998
Roherts
200T
Neal ct al. 2012)
5

---

Under review as a conference paper at ICLR 2022
7 EXPERIMENTAL RESULTS
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 first and second moments of the target distribution's i marginals (further explained in Appendix D). reporting the minimum obtained across all dimensions. For comparisons, we consider MSJD optimization (MSJD Opt.) Pasarriaa 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 and K
7.1 OPTIMIZING MUL.TI-SCHEME PROPOSALS
We now introduce a normalizing flow Tabak and Turner, 2013 Rezende and Mohamed 2015
Papamakarios et al., 2019) based proposal distribution. With To denoting a normalizing flow's parameterized transformation, the functions fL,0, (D,,, Ex,o, and ED specified by neural networks, and @ denoting element wise multiplication, our proposal distribution is defined via:
ñ ~ NO; 1)
و = ML.0 (T (x)) X' ( + e + Tz
(x) + E2.0 EL. on
This implementation is fundamentally conditional normalizing flow (Winkler et al. 2019) and is similar to NeuTra-HMC (Hoffmaan et al. 2019). which employs HMC within the latent space of a pre-fit normalizing flow to improve MCMC efficiency. This multi-schene distribution differs from NeuTra-HMC by allowing both the flow's transformation and the latent distribution to be optimized to improve MCMC efficiency, rather than being fixed 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 cross formation with a maximal distance of 8 between component centers. We use the NICE architecture (Dinh et al. 2014) to define the normalizing flow and vary the number of coupling layers (using either 8 or layers), which influences how accurately the flow is able to approximate the target distribution. After each optimization, we determine efficiency measures on the basis of 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 replications for each considered objective function. The results of these optimizations are provided in Table Comparisons of the computational costs for optimizing each objective are provided in Table
Density plots of multi-scheme proposal distributions optimized using an MSJD objective and our Ab Initio objective are provided in Figure T. Position dependence within the proposal distribution lead MSJD and L2HMC's modification to optimize towards arbitrarily non-ergodic proposal distributions, sampling only either the horizontal or vertical components of the target distribution.
This experiment demonstrates fundamental difficulty in applying the GSM objective to multi-scheme proposals. Here, the flow's 's depth determines it's maximal accuracy in approximating independent resampling of the target distribution, creating maximum acceptance rate that can be obtained while still approximating global resampling. Attaining greater acceptance rates is possible, but forces inefficient R-ike behavior. In this case, 8 coupling layers within the flow permits maximal resampling acceptance rate near 0.70, while 3 coupling layers lowers this maximal rate below 0.60. A
*Calculations performed using 2080 RTX GPU and Xeon Gold 61552 CPU.

---

Under review as a conference paper at ICLR 2022
Table Comparison of sampling performance obtained optimizing the multi-scheme proposals of Equation using various objective functions. Results are reported for two choices for the number of coupling layers used the proposal's 's normalizing flow (L and L 3). Value means are reported to most the first significant digit of standard error (reported in parentheses).
1.=3
Acc. Rate
ESS per Proposal
Acc. Rate
ESS per Proposal
MSJD
MSJD
24 0.33(5)
x; 0.41(5)
x²
Ab Initio, Eq.
0.67(3)
22(()
0.55(4)
17(1)
0.18(8)
0.2(1)
MSID Opt.
0,975(6)
66.2(9)
1.10())
6(1)e-4
0.978(4)
67(1)
1.10()
7(1)--4
L2HMC Ohj.
0.974(3)
65.7(5)
1.12(2)
6(1)e-4
0.974(5)
66(1)
1.1((1)
6.1(3)e-4
GSM-90
0,92(1)
1.86(5)
5.6(4)e-4
7(1)e-
0.90(2)
1.80(4)
5.3(4)e-4
7(1)-4
GSM-60
0,65(6)
21(2)
0.3(2)
0.3(2)
0.60(1)
1.41(1)
6(1)e-4
1(1)e-3
GSM-30
0.31(1)
10.7(4)
0.14(2)
0.17(1)
0.302(6)
10.5(4)
0.13(2)
0.18(2)
L.D., Resample 1.00
34.4(3)
0.98(5)
0.99(6)
1.00
33.9(3)
0.90(7)
1.03(5)
Table 3: Comparison of number of gradient evaluations computed per second when optimizing the multi-scheme proposals of Equation (3) using various objective functions. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
Initio, Eq.
MSID Opt.
L2HMC Obj.
GSM
L=8
7.6(1)
9,9(3)
10.3(3)
7,33(7)
L=3
10.6(1)
12.4(6)
12.6(4)
9.4(4)
y(o.y
Contours of Target Density, n(x y)
StartingPoint, (X)
aa) Multi-scheme proposals obtained by maximizing MSJD. g(x', yoo,
Contours Target Density, n(x.) StartingPoint, (X).
(b) Multi-scheme proposals obtained by optimizing Ab Initio Objective.
Figure 1: Comparison multi-scheme mixture proposals (L=8) Equation (3) obtained by max- imizing MSJD (a) and optimizing Ab Initio objective (b). Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
user with 8 coupling layers who uses the GSM targeting an acceptance rate of 60% to optimize their distribution would fortuitously arrive at efficient proposal distribution. However, if they had used
7

---

Under review as a conference paper at ICLR 2022
Table 4: Comparison of sampling performance obtained optimizing augmented multi-scheme propos- als of Equation (3) using various objective functions targeting for logistic regression of UCI datasets. Value means are reported at most the first significant digit of standard error (reported in parentheses).
posterior distributions of parameters
German Credit (d=21)
Heart Disease (d=14)
Acc. Rate
ESS per Proposal
Acc. Rate
ESS per Proposal
MSJD
MSJD
32 0.63(6)
x²²
Ab Initio, Eq. 2
0.81(1)
0.30(1)
0.56(3)
0,49(7)
0.86(2)
1.03(4)
0,58(7)
MSJD Opt.
0.77(2)
0.32(1)
0.52(7)
0.4(1)
0.81(1)
1.40K3)
0.69(3)
0.32(6)
1.2MCC Ohj.
0.76(2)
0.32(1)
0.54(4)
0.48(5)
0.82(1)
1.43(4)
0.73(4)
0.28(7)
GSM-90
0.90(1)
0.012(1)
0.006(1)
0.005(1)
0,900(3)
0.138(3)
0.034(2)
0.035(1)
GSM-60
0.61(3)
0.24(1)
0.40(3)
0.40(3)
0,59(2)
0.72(2)
0.34(3)
0.35(3)
GSM-30
0.29(5)
0.13(2)
0.13(4)
0.13(4)
0.29(3)
0.42(4)
0.13()
0.14(3)
L.I.D. Resample 1.00
0.3700(3)
0.98(1)
0.983(3)
1.00
1.222(1)
0.97(1) 0.97(2)
Table 5: Comparison of number of gradient evaluations computed per second when optimizing the multi-scheme proposals of Equation (3) using various objective functions targeting posterior distributions of parameters for logistic regression UCI dataset. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
Ab Initio, Eq. 2
MSJD Opt.
L2HMC Obj. 14.8(9)
GSM
German Credit
1 11.5(6)
1.((1)
10.1(3)
Heart Disease
11.7(4)
15.0(3)
14.7(7)
10.8(3)
model with coupling layers instead and targeted the same 60% acceptance rate, the result would be proposal distribution far more inefficient than the actual capabilities of their proposal model class.
In both cases, the Ab Initio objective is able to optimize towards efficient approximations of i.i.d. resampling. An intuitive explanation for the results of this experiment is provided in Appendix H
As this experiment utilizes a complicated proposal distribution, it is reasonable to question whether the problems exhibited here arise from some peculiarities of the model class of Equation (3). For this reason, Appendix Drppeats a similar experiment with simpler proposal distribution and recovers similar results. We conclude that the effects demonstrated within Table may be traced back the choice of objective function used for optimization.
7.2 OPTIMIZING AUGMENTED MULTI-SCHEME PROPOSALS
Augmentation with auxiliary variables (e.g. the momenta of HMC) is common feature of MCMC schemes (Levy et al., 2018; Habib and Barber 2018) and, in this experiment, we augment the original distributions with number of independent auxiliary variables equal to the distribution's original dimensionality. In Appendix E. 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 multi- scheme proposals defined in Equation
3 for MCMC sampling of regression weights in bayesian logisite regression for various UCI datasets (Bache and Lichman 2013). After each optimization, we determine efficiency measures on the basis of Markov Chains of 20000 total proposals starting from samples drawn from 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 coste5 for optimizing each objective are provided in Table Under these expanded conditions, we find the objective functions exhibiting much of the same behavior as seen in Section 7.I in particular with the Heart Disease dataset.
Calculations performed using 2080 RTX GPU and Xeon Gold 61552 CPU.
8

---

Under review as a conference paper at ICLR 2022
Table 6: Comparison of sampling performances and Ab Initio losses (lower is better) obtained from various MCMC schemes for the logistic regression of MNIST digits (d=785). Value means are reported to most the first significant digit of standard error (reported in parentheses).
Acc.
ESS Per Proposal
Ab Initio
Samples
Model
Obj. Func.
MSJD
Rate
Eq. 2
Per Sec.
x
x²²
Precond. MALA
Eq.
0,501(7)
129(2)
4(1)e-6
5(1)--6
1.644())3
115(4)
Precond. MALA
MSJD
0.543(3)
241(1)
1(1)e-10
2(2)e-10
2.59(3)e3
113(4)
Precond. RWM
Eq.
2
0.239(4)
1.36(2)
7(2)e-
1.1(3)e-6
2.653(1)e3
205(4)
Multi-Scheme,, Eq.
Eq.
0.24(4)
1.2(1)
5(2)e-
6(3)e-7
2.83(7)e3
47.5(4)
Resampling Norm. Flow
NLL
1e-3-1e-6
<4.73
5,2(3)65
68.2(2)
7.3 APPLICATION TO MCMC SCHEME SELECTION
As a general measure of MCMC efficiency, the Ab Initio objective funtion of Equation ( 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 logisite regression between two MNIST (LeCun al.) 1998) digits. For each evaluated scheme, after optimization using the respective objective function, we determine efficiency 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 The resampling normalizing flow (an unconditional normalizing flow optimized to directly approximate the target distribution) did not produce single accepted proposal within these evaluation Markov Chains, hence we unable to estimate ESS and our estimates of acceptance rates and MSJD for this proposal distribution should be viewed 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 sacrificed 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 find the example Ab Initio objective function of Equation to be robust measure of MCMC efficiency that may be used to evaluate and compare the efficiency of multiple MCMC schemes in addition to its utility as an objective function for optimizing the individual parameters of highly expressive MCMC schemes.
8 CONCLUSION AND FUTURE WORK
In this work, we have shown how to construct Ab Initio objective functions that suited for the optimization of highly expressive proposal distributions. We find 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 existing objective functions (Levy al. 2018 Titsias and Dellaportas) 2019; Pasarica and Gelman 2010). By design, Ab Initio objective functions are approximations of our notion of MCMC efficiency 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 ( our proposed Ab Initio procedure allows for further improvements by considering alternative component functions and coefficient fitting 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 efficient samples, and how to handle scaling into higher dimensions remain important practical considerations for MCMC optimization. We argue that these considerations selection. We will therefore investigate these questions in a future work focusing on comparisons among specific proposal architectures and training procedures.
primarily influenced by proposal architecture and training procedure

---

Under review as a conference paper at ICLR 2022
REFERENCES
Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation State Calculations by Fast Computing Machines. The Journal of Chemical Physics, 21(6): 108--1997 1953.
Julian Besag and Peter J Green. Spatial Statistics and Bayesian Computation. Journal of the Royal Statistical Society: Series B (Methodological), 55(1):25-37, 1993.
Simon Duane, Anthony D Kennedy, Brian J Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216-222, 1987.
Gareth O Roberts and Jeffrey S Rosenthal. Examples of Adaptive MCMC. Journal of Computational and Graphical Statisticx, 18())3399-367, 2009.
Dino Sejdinovic, Heiko Strathmann, Maria Lomeli Gareia, Christophe Andrieu, and Arthur Gretton. Kernel Adaptive Metropolli-Hastings.. In International Conference on Machine Learning, pages 1665-1673. PMLR, 2014.
Heikki Haario, Eero Saksman, Johanna Tamminen, et al. An Adaptive Metropolis Algorithm. Bernoulli, 722:2233-22, 2001.
Gareth Roberts and Jeffrey Rosenthal. Coupling and Ergodicity of Adaptive Markov Chain Monte Carlo Algorithms. Journal of Applied Probability, 44(2):458-455, 2007.
Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-NICE-MC: Adversarial Training for MCMC. In Proceedings of the 3 International Conference on Neural Information Processing Systems, pages 5146-5156, 2017
Span Spanbauer, Cameron Freer, and Vikash Mansinghka. Deep Involutive Generative Models for Neural MCMC. arXiv preprint 201.277000.1166, 2020.
Matthew Hoffman, Pavel Sountsov, Joshua V Dillon, Ian Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport arXiv preprint 1wXXiv:9003.030000 2019.
Nando de Freitas, Pedro Hojen-Sprensen, Michael I Jordan, and Stuart Russell. Variational MCMC. In Proceedings of the Seventeenth Conference Uncertainty in Artificial Intelligence, pages 120-127, 2001.
Raza Habib and David Barber. Auxiliary Variational MCMC. In International Conference on Learning Representations, 2018.
Daniel Levy, Matthew D Hoffman, and Jascha Sohl-Dickstein. Generalizing Hamiltonian Monte
Carlo with Neural Networks. International Conference on Learning Representations, 2018.
Zengyi Li, Yubei Chen, and Friedrich T Sommer. Neural Network MCMC Sampler that Maximizes Proposal Entropy. Entropy, 23(3):269, 2021.
Andrew Gelman, John B Carlin, Hal S Stern, David B Dunson, Aki Vehtari, and Donald B Rubin. Bayesian Data Analysis. CRC press, 2013.
Dootika Vats, James M Flegal, and Galin Jones. Multivariate Output Analysis for Markov Chain Monte Carlo. Biometrika, 106(2):321-3) 2019.
Gareth O Roberts, Jeffrey S Rosenthal, et al. Optimal Scaling for Various Hatropolis-Hastinn Algorithms. Statistical Science, 16(4):351-377, 2001.
Andrew Gelman, Walter Gilks, and Gareth O Roberts. Weak Convergence and Optimal Scaling of Random Walk Metropolis Algorithms. The Annals of Applied Probability, 7(1): 1997.
Gareth O Roberts and Jeffrey S Rosenthal. Optimal Scaling of Discrete Approximations Langevin Diffusions. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 60(1): 255-268, 1998.
10

---

Under review as a conference paper at ICLR 2022
Peter Neal, Gareth Roberts, Wai Kong Yuen, al. Optimal Scaling of Random Walk Metropolis Algorithms with Discontinuous Target Densities. Annals of Applied Probability, 22(5):1880-1977, 2012.
Alexandros Beskos, Natesh Pillai, Gareth Roberts, Jesus-Maria Sanz-Serna, Andrew Stuart, et al. Optimal Tuning of the Hybrid Monte Carlo Algorithm. Bernoulli, : 2013.
Warren J Hehre. Ab Initio Molecular Orbital Theory. Accounts of Chemical Research, 9(1))3999-466, 1976.
MT Yin and Marvin I. Cohen. Theory of Ab Initio Pseudopotential Calculations. Physical Review B, 25(12):7403, 1982.
Dominik Marx and Jürg Hutter. Ab Initio Molecular Dynamics: Basic Theory and Advanced Methods. Cambridge University Press, 2009.
Kirill Neklyudov, Evgenii Egorov, Pavel Shvechikov, and Dmitry Vetrov. Metropolis-Hastinng View on Variational Inference and Adversarial Training. arXiv preprint arXiv: 2018.
Tongzhou Wang, Yi Wu, David A Moore, and Stuart J Russell. Meta-Learning MCMC Proposals. Proceedings of the 32nd International Conference on Neural Information Processing Systems, 2018.
Michalis Titsias and Petros Dellaportas. Gradient-Based Adaptive Markov Chain Monte Carlo. Advances in Neural Information Processing Systems, 32:15730-15739 2019.
Cristian Pasarica and Andrew Gelman. Adaptively Scaling the Metropolis Algorithm Using Expected Squared Jumped Distance. Statistica Sinica, pages 343-364, 2010.
E Bergou, Y Diouane, V Kungurtsev, and CW Royer. A Stochastic evrenberg-Marquaad Method Using Random Models with Application to Data Assimilation. arXiv preprint arXiv: 107.0000101111000 02176, 2018.
Esteban G Tabak and Cristina V Turner. A Family of Nonparametric Density Estimation Algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
Danilo Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows. In International Conference on Machine Learning, pages 1530-1538. PMLR, 2015.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXiv preprint arXiv: 2019.
Christina Winkler, Daniel Worrall, Emiel Hoogeboom, and Max Welling. Learning Likelihoods with Conditional Normalizing Flows. arXiv preprint arXiv:1912.00042, 2019.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear Independent Components Estimation. arXiv preprint arXiv: arXiv:1410.8511 2014.
Kevin Bache and Moshe Lichman. UCI Machine Learning Repository, 2013.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. mrudient-based Learning Applied to Document Recognition. Proceedings of the IEEE, 86(1 11:22788-2333 22224,, 1998.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastie Optimization. In ICLR (Poster), 2015.
Brian Ripley. Pattern Recognition and Neural Nenworks. Cambridge University Press, 2007.
Adam D Cobb, Atilmm Güneş Baydin, Andrew Markham, and Stephen J Roberts. Introducing Explicit Symplectic Integration Scheme for Riemannian Manifold Hamiltonian Monte Carlo, arXiv preprint arXiv: 1910.06243.33.. 2019.
11

---

Under review as a conference paper at ICLR 2022
Aki Vehtari, Andrew Gelman, Daniel Simpson, Bob Carpenter, and Caul-Christian Bürkner. Rank- normalization, folding, and localization: An improved R for assessing convergence of meme. arXiv preprint arXiv: 08008, 2019.
Anand Patil, David Huard, and Christopher J Fonnesbeck. PyMC: Bayesian Stochastic Modelling in Python. Journal of Statistical Software, 35(4):1, 2010.
A COMBINING PROPER AND REPRESENTATION INVARIANT FUNCTIONS
As part of our Ab Initio approach, we rely on the following proposition:
Every positive weighted combination of proper and representation invariant objective functions is itself a proper and representation invariant objective function.
Proposition:
To prove this, assume that we have two proper and representation invariant (over group of almost sure diffeomorphisms D) objective functions f and with target æ and allowed set proposals G. Let us select arbitrary constants C1 C2 >> 0 and consider the properties of the objective function cif + c2h. From the invariance f and h we may immediately conclude:
(af + c2h) e T a g: o c2 o 9; To
= c m| + c2h(g, (caf + c2
Hence f + egh is itself representation invariant. Finally, for all g(>1lx)) # m(²)):
+ faf+cch]lasn = + cal[T;
< ci flg.g o (cif c2h)lg=m]
Hence / + oah is also proper. This concludes the proof of the above proposition.
This offers a straightforward approach to approximating our ground truth objective 4 Although we do have a guarantee that C* 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 fitting the coefficients of Ab Initio objective function to match analytical results is more easily accomplished than deriving particular objective function that happens to match analytical results, we argue that our Ab Initio approach is 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 of the main paper. Below, we start with a definition.
Definition:
An objective function is said to be nearly proper if it attains (not necessarily unique) global minimum at (x'(( = π|( rTx')) and 0g (x')x)) ) = for all x', x.
We may consider minor generalization of the previous proposition:
Every positive weighted combination of at least one proper and representation invari- ant objective function with any number of nearly proper and representation independent functions itself proper and representation invariant objective function.
Proposition:
The proof that the resulting combination 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 be proper and let be nearly proper. Let us select arbitrary constants C₁ C2 > and consider the properties of the objective function c + eah. For all g(rl(x)) o (<x)::
12

---

Under review as a conference paper at ICLR 2022
(eif + =2h)a; = c =
S cohlg
<
(ar c
Hence of + c2h is also proper.
B DERIVATION OF OUR EXAMPLE AB INITIO OBJECTIVE FUNCTION
We begin with the definition of the (maximized) lower bound of the GSM, where is hyper- parameter to be adapted during optimization to match the user specified acceptance rate:
f m] (x)+( (1))) (X'
xx) + B B (x log a (x'
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 in- clusion into an Ab Initio objective function. The addition of an expected cross entropy term, Eg (x) [ (g (z) [log क (X and factoring out an overall negative sign yields an objective function to be minimized that follows the functional constraints listed in Section
Clgs Eg (xxx((("1())]((( B Eg (ggg" 101 [x]]
However, encounter additional complication. If [x)+(x)] (xz(x(X'xx) [x(x)) and Ex (x) (g (2) [log [XX|x)) have differing dependencies on dimensionality as d + co, 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 Ab Initio function of the form:
C[g: Ex))(R)) ( C(d) Ex Dog
Future research may asymptotic analysis derive functional forms for C(d) that are best suited for exceedingly high dimensions. In this work, we followed a procedure manual trial and error and found that C(d) Ad yields useful results least up 10, 000. This completes our adaptation of the GSM into the Ab Initio objective function listed in Equation (I:
Clg: n] E (g(x)()))))ll) (x())) 2 loo (x*(x))
As the Ab Initio objective function of Equations and (2) functionally differs from the GSM by the addition a cross-entropy term and by dimensional scaling of the multiplicative coefficient, we should expect that their computational costs and complexities to be comparable. This is supported by our results Tables 3 5 12 and 14 18 While the GSM is often slightly more computationally costly,
we believe this is more due to particular implementation of the online update of its hyperparameter for tuning to match the targeted acceptance rate as opposed to reflecting a fundamentally different complexity scaling of the GSM compared the Ab Initio objective function.
13

---

Under review as a conference paper at ICLR 2022
C
PROPERTIES OF OUR EXAMPLE AB INITIO OBJECTIVE FUNCTION
The functional form our example Ab Initio objective function is given by:
fig: π| = B ( C(d) (x) (x)
Where ((d) is dimension dependent positive constant. We will now consider the following proposition:
Proposition: The example Ab Initio objective function. Lig: x]. defined above is a proper and representation independent objective function for all C(d) 2 0.
Following the general restrictions we' 've used to define our optimization problem, let g and be positive and smooth and let T denote any almost sure diffeomorphism. Let Xr and Zy denote the regions of the domain and range T over which is differentiable and invertible. From the properties of KL-Divergence, we know that (x)n(x) (x | (x))) 2 O and is equal to 0 if and only if gx x xxx') ) for all X, X² From the definition of KL-Divergence and the law of the unconscious statistician, have:
E / g(X'( log
(x)l
dx')
x(x²)
-EZTOOTEE / g(8'(z) log
g(X' [
dx
π ((X')
Nr
T g ( dz")
-EZTTO
To log
T o π (z')
2r
DKL(T g(<|x>]|[ [
{ o ा (x'))]
Hence, [x)) (x(1 is a proper and representation invariant objective function.
acceptance rate is guaranteed to fall in the range 0 S (x* S 1. Additionally, if gxX') (x) = x (x'), ), (x'(xx) x = 1. This guarantees
Next, for all proposals g and targets æ, the Metropolis-Hastoee
that (x))(s) (LE= ~ Hastings acceptance rate and the law of the unconscious statistician, we have:
log [x]]] is nearly proper. And from the definition of the Metropolis-
Ex) [ log (x
[x]] [x [gg [x)]
π x g X')
-Ex-x(8) Ex ( log min (1.
x H]
T (7)-
= B log min (1,
T(T (Z)
To x g(7
( log min
To ग (Z) g
Ex Bx (Ty log (x*(x]]
And so (gg-T()) [x> log (x'' [x]] is a nearly proper and coordinate invariant objective function. With C(d) 2 0, our example Ab Initio objective function is therefore the positive weighted combination of proper and representation invariant objective function with 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 also representation independent and thus belongs to the same functional class as approximated "ground truth" objective.
14

---

Under review as a conference paper at ICLR 2022
D ADDITIONAL DATA FOR VERIFICATION TASKS
For verification 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 mixture distribution of a standard uniform distribution and standard multivariate gaussian (with probabilities of I+e- and I+c. respec- tively, 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 go (x)) are generated to empirically estimate expected loss on the basis of 50 total samples. The gradient update steps take the form of Algorithm [ with M 1 and N 50. Optimization is performed using the Adam optimizer (Kingma and Ba 2013) with stepsize of 3e-4 (reduced to 3e-5, 3e-6, and 3 for uniform targets of dimensions 100,,000, and 1 10000). For optimizing GSM (Titsias and Dellaportas 2019). 3 is initially set to d'1 and p is set to 0.02. 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 verification optimization is replicated a total 5 times. Tables 7 I 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 fixed step size determined to yield the analytically known optimal acceptance rates.
I0 and
We were unable to perform gradient based optimization of MSJI) and 1.2HMC's objective when targeting the uniform distribution, and resorted to estimating their optima by fitting 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 fitting a parabola near the objective function's 's maxima, with uncertainties reported following the typical propagation of error. Even following this procedure, L2HMC's 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 failure case for L2HMC's objective.
Table 7: Verification results from applying MSJD maximization to the MCMC optimization tasks with analytically known optimal properties. As these verification tasks are set within the diffusionary limit, MSJD recovers known optimal results. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
d = 100
d= 1000
d = 10000
MSID Opl.
Analytic Simulated
MSJD
Analytic Simulated
MSJD
Analytic Simulated
Analytic Exact
Opt.
Opt.
Aoc. Rate MSJD
0.237(4)
0,235(1)
0.232(3)
0.234(3)
0.233(2)
0.234(2)
0.234 9
Gaussian
RWM
1.32(2)
1.32(1)
1.32(1)
1.32(3)
1.33(2)
1.33(2)
Acc. Rate
0,538(6)
0,574(1)
0.559(5)
0.573(1)
0.56(1)
0.574(2)
0.574 F
Gaussian
MALA
MSJD
38 38.6(1)
38.4(1)
167(1)
166(1)
749(2)
747(5)
Acc. Rate MSID
0.15(2)
0.135(55
0.15(2)
0.137(4)
0.134(2)
0.132(3)
0.135
Uniform
RWM
8.2(7)e-3
8.2(3)e-3
8.4(5)e-4
8.6(2)e-4
8.4(4)e-5
8.2(2)e-5
Acc. Rate MSID
0,21(2)
0.234(1)
0,23(1)
0.233(1)
0.24(1)
0.234(3)
0.234 6
Laplace
RWM
1.4(())
1.48(1)
1.37()
(,37(0)
1.34(1)
1,34(2)
Auc. Rate
0.223(5)
0.235(3)
0.230(8)
0.235(3)
0.23(2)
0.235(2)
0.234
Cauchy
RWM
MSJD
2.72(3)
2.74(5)
2.65(1)
2.67(3)
2.64(2)
2.65(1)
Roberts
(2001):
Gelman et al. (1997)
Roherts et al.
2001
Roberts and Rosenthal 1998
Roberts
200T
Neal ct al. 2012)
15

---

Under review as a conference paper at ICLR 2022
Table 8: Verification results from applying L2HMC's objective to the MCMC optimization tasks with analytically known optimal properties. L2HMC's objective only recovers optimal behavior for MALA proposals and cannot be used to optimize RWM proposals with uniform target. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
100
1000
10000
L2HMC
Analytic Simulated
L2HMC
Analytic Simulated
L2HMC
Analytic Simulated
Analytic Exact
Obj.
Obj.
Obj.
Acc. Rate 0.587(3)
0,233(3)
0,580(4)
0.235(1)
0.56(1)
0.233(1)
0.234
Gaussian
RWM
MSJD
0.69(1)
1.31())
0.71(1)
1.34(1)
0.76(4)
1.32(1)
Acc. Rate
0.546(5)
0.575(2)
0.553(6)
0.575(3)
0.562(6)
0.574(1)
0.574
Gaussian
MALA
MSJD 38.6(2)
38.5(2)
167(1)
167(1)
750(3)
749(2)
Acc. Rate
0.135
Uniform
RWM
MSJD
Acc. Rate
0.57(1)
0,235(2)
0,59(1)
0,233(2)
0.59(1)
0.234(2)
0.234 9
Laplace
RWM
MSID 0.78(4)
1.48(1)
0.71())
1.37(1)
0.69(2)
1.340)
Acc. Rate
0.514(4)
0,235(2)
0,516(5)
0,237(1)
0.52(1)
0,236(1)
0.234
Cauchy
RWM
MSJD
1.77(3)
2.71(2)
1.77())
2.69(1)
1.74(5)
2.66(4)
Table 9: Verification results from applying GSM targeting acceptance rates of 0.3 to the MCMC optimization tasks with analytically known optimal properties. Optimal behavior is approximately recovered only for RWM proposals at this acceptance rate. Value means are reported to at most the first significant digit standard error (reported in parentheses).
100
d 1000
10000
Analytic Simulated
Analytic Simulated
Analytic Simulated
Analytic Exact
GSM-30
GSM-30
GSM-30
Acc. Rate
0.298(4)
0.235(2)
0.295(4)
0,234(2)
0.299(9)
0,234(2)
0,234 9
Gaussian
RWM
MSJD
1,28(1)
1,33(2)
1,29(1)
1.32(1)
1,30(1)
1.33(1)
Acc. Rate MSJD
0.297(6)
0.575(1)
0.295(8)
0,575(2)
0.297(6)
0.573(3)
0.574
Gaussian
MALA
31.8(5)
38.4(3)
134(2)
167(1)
595(7)
748(3)
Acc. Rate MSJD
0.302(6)
0.135(1)
0.30(1)
0.136(2)
0,30(1)
0.135(2)
0.135
Uniform
RWM
6.6(2)e-3
8.18(7)e-3
6.8(2)e-4
8.6(2)e-4
6.9(1)e-5
8.4(1)e-5
Acc. Rate
0.303(6)
0.234(3)
0.29(1)
0.235(3)
0.32(1)
0.236(1)
0.234 9
Laplace
RWM
MSJD
1.43(1)
1,48(1)
1.34(2)
1.38(2)
1.29(2)
1.34(1)
Acc. Rate
0,297(2)
0,234(3)
0,291(3)
0.237(1)
0.30(1)
0.234(2)
0.234
Cauchy
RWM
MSJD
2.62(3)
2.71(4)
2.58(3)
2.69(1)
2.57(3)
2.65(2)
Roberts 2001):
Gelman et al. (1977
Roherts et al.
2001
Roberts and Rosenthal 1998
Roherts
200T
Neal al. 2012
16

---

Under review as a conference paper at ICLR 2022
Table 10: Verification results from applying GSM targeting acceptance rates of 0.6 the MCMC optimization tasks with analytically known optimal properties. Optimal behavior is approximately recovered only for MALA proposals at this acceptance rate. Value means are reported most the first significant digit of standard emror (reported in parentheses).
100
1000
10000
Analytic Simulated
Analytic Simulated
Analytic Simulated
Analytic Exact
GSM-60
GSM-60
GSM-60
Ace. Rate
0,602(2)
0,234(1)
0.60(1)
0.234(2)
0,60(1)
0,233(2)
0.234
Gaussian
RWM
MSJD
0.66(0)
1.31(1)
0.66(2)
1.33(1)
0.65(2)
1.32(1)
Ace. Rate 0.607(7)
0,575(1)
0.60(1)
0.575(1)
0.596(5)
0.576(2)
0.574
Gaussiun
MALA
MSJD (2)
38.5(2)
166(0)
167(0)
743(4)
751(4)
Ace. Rate
0,61(1)
0.135(2)
0.58(1)
0.135(1)
0.57(4)
0.138(1)
0.135
Uniform
RWM
MSJD 2.4(1)e-3
8,20(9)e-3
2.7(1)e-4
8.52(8)e-4
2.8(4)e-5
8.64(4)e-5
Luplace
RWM
Ace. Rate
0.60(1)
0,234(2)
0.61(3)
0,233(2)
0.63(4)
0,237(1)
0.234 12
MSJD
0.69(3)
1.48(1)
0.64(6)
1.37()
0.6(1)
1.35(1)
Acc. Rate
0,599(6)
0.235(2)
0.603(5)
0.236(3)
0.60(1)
0.234(2)
0.234
Cauchy
RWM
MSJD
1.33(3)
2.72(2)
1.31(3)
2.69(3)
1.31(4)
2.64(3)
Table 11: Verification results from applying GSM targeting acceptance rates of 0.9 the MCMC optimization tasks with analytically known optimal properties. Results are significantly suboptimal. Value means are reported most the first significant digit of standard error (reported in parentheses).
d 100
1000
10000
Analytic Simulated
Analytic Simulated
Analytic Simulated
Analytic Exact
GSM-90
GSM-90
GSM-90
Acc. Rate
0.900(3)
0.234(2)
0.900(2)
0.232(2)
0.910(6)
0,233(2)
0.234 12
Gaussian
RWM
MSID
0.06(0)
1.3((1)
0.06(2)
1.32(1)
0.05(1)
1.32(1)
Acc. Rate MSJD
0.899(2)
0.575(1)
0.901(2)
0.576(2)
0.902(2)
0.574(1)
0.574
Gaussian
MALA
20.6(3)
38.5(1)
92(1)
167(1)
421(4)
750(2)
Acc. Rate MSJD
0.90(1)
0.136(2)
0.90(2)
0,133(3)
0.95(5)
0,135(2)
0.135
Uniform
RWM
1.5(3)e-4
8.2(1)e-3
1.8(7)-5
8.4(2)e-4
1(1)o-6
8.5(1)e-5
Ace.
0.905(5)
0.234(2)
0.91(1)
0.235(2)
0.90(2)
0.234(1)
0.234 12
Laplace
RWM
MSJD
0.05(1)
1.47())
0.05(1)
1.38(2)
0.06(2)
1.38(2)
Acc. Rate
0,901(4)
0.234(1)
0.901(2)
0.235(2)
0.898(5)
0.233(2)
0.234
Cauchy
RWM
MSJD
0.11(1)
2.73(1)
0.11(0)
2.65(2)
0.12(1)
2.65(3)
Roberts 2001):
Gelman et al. (1977
Roherts et al.
2001
Roberts and Rosenthal 1998
Roherts
200T
Neal al. 2012
17

---

Under review as a conference paper at ICLR 2022
E APPROXIMATING MCMC SCHEMES WITH MULTI-SCHEME PROPOSALS
As part of our experiments, we utilize normalizing flow based multi-scheme proposal distribution, parameterized as:
n ~ N1O; 1}
z' x, ñ = HL. (x)) + E1 (T (x)) (x) o = )D,0 + ED,0(x) @ Tz(2")
Where To denotes a normalizing flow's parameterized transformation, the functions PL.9- J.D.. EL.. and ED specified by neural networks, and @ 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, go 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: HTTONRTT.HL. 1DD00 Ō, and EL,0 Ep,o then:
go ( X
(xx)
Depending the accuracy of the flow's approximation, the multi-scheme distribution approximately recovers i.i.d. resampling from the target.
Isotropic RWM: Te = I, AL.e = a X, E1 I, and Ep then:
X ~ N (x:
And the multi-scheme distribution recovers isotropic RWM with step-size T.
Isotropic MALA: If = -, FL.0 = 0. JD.0 R x+ V log π (x), Et,0 = T, and Ep, = V2r,
then:
x x ~ N +
log ग
2 1}
Depending on the accuracy of the approximation of μ.D ()), the multi-scheme distribution recovers isotropic MALA with step-size т.
Preconditioned Diffusions: If To(z) s A- where A is an invertible preconditioning matrix, then appropriate choices uL,8. PD.. DD and EL will approximately recover preconditioned RWM and MALA.
Latent MALA: ON M...R.PP..... (X) R (1+ pn. Er,0 = 2T, and E, = I, then the multi-scheme distribution will mimic the normalizing flow based proposal scheme Hoffman et al. (211)). with MALA used within the flow's 's latent space rather than HMC.
F. AUGMENTING VARIABLES
One method of greatly expanding the model class of proposal distribution is to augment the distri- bution with auxiliary variables. This is most notably accomplished with the canonical momenta in HMC. We will denote augmentation with a semicolon, x; indicates that x has been augmented with the auxiliary variables a.
If E..o & (Z(Z') R Ō and (x))(x) R Xint Pint, where Xinti Pent are the results of applying any desired integrator to the system then the flow based proposal may approximate HMC by iteratively resampling p from normal distribution, generating a new proposal, and then performing a Metropolis-Hastime approximation should be expected to be accurate, merely that it mathematically possible within the flow proposal's '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.
Approximate HMC:
correction. Do note that
not claiming that this
18

---

Under review as a conference paper at ICLR 2022
These examples are not exhaustive. Depending implementation, these flw-based 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 DETAILS OF COMPARED OBJECTIVE FUNCTIONS
Throughout this work, we focus on comparing the properties four objective functions for MCMC proposal optimization.
Example Ab Initio
From Equation 2 to be minimized:
f[g: π|
(x)x [ 0.18125 E 10 (xg
MSJD As used by Pasarica and Gelman 2010), to be maximized:
[g T] = [x)+(x] Ex -((()))))
[| x||
L2HMC's Modification
As proposed by Levy et al. 2018), to be minimized:
||x - x|
Clg: x| = Ex))(x)) Ex
o
min
Where oi 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:
Gig n| = (x) [xx-0((()))) [-8 log x) + log (xx(
Where 3 a hyper We follow the suggested method for updating / in online manner. Following our discussion regarding the asymptotic behaviors of the two components with respect to dimension, we found beneficial initialize 3 to 5, rather than the default value of After every parameter update, the 3 is updated following:
3 + 3 (1 Pa (Atarget - Qcurrentt)
Where the suggested value of p3 0.02 throughout.
G
FURTHER DETAILS FOR MULTI-SCHEME PROPOSAL OPTIMIZATION
Our experiments with multi-scheme proposals utilize distribution specified as follows:
ñ ~ N1 1}
z' x n = (1. (x) Σr.o (T, (x)) o X (D) + ED) C (z
The target distribution an equal mixture of 4 standard gaussians positioned in a cross formation with maximal distance of 8 between component centers. The additive functions, JLL8 and pp.a. are specified by standard ReLU networks (input dimension 2, hidden dimension 16, 4 hidden layers, output dimension 2). The multiplicative functions, Σ,,0 and EpDo are specified by standard ReLU networks (input dimension 2, hidden dimension 16, 4 hidden layers, output dimension 2) followed by component-wiss exponentiation. The normalizing flow, Ts. follows the NICE architecture (Dinh) et al. 2014), with additive coupling layers specified by standard ReLU networks (input dimension I, 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 taken, with each training batch taken from starting points sampled i..dd from the target distribution from each of which 100 independent proposals from go [x] generated to empirically estimate
19

---

Under review as a conference paper at ICLR 2022
expected loss the basis of 800 total samples. The gradient update steps take the form of Algorithm
I with M and N = 100. Optimization is performed using the Adam optimizer (Kingma and 2015) with 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), 33 is initially set to 0.5 and P3 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 independent Markov Chains 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 3 4 and 3 density plots of the multi-scheme proposals obtained following optimization using L2HMC's objective Levy et al. (2018) and the GSM Titsias Dellaportas/ (2019) with varying target acceptance rates. Figures 6 9 10 show density plots of the multi-scheme proposals obtained following optimization of all considered objective functions.
Ba
g(x', y' yo)
Contours of Target Density, n(x, y)
StartingPoint, (X)
Figure Comparison of sclti-scheme mixture proposals (L=8) of Equation 4 obtained by optimizing L2HMC's objective. Each subplot illustrates the proposal distribution's 's density when sampling from a given starting position (in red).
g(x', y' yo)
Contours of Target Density, n(x, y)
StartingPoint, (Xo.
Figure 3: Comparison of multi-scheme mixture proposals (L=8) of Equation 4 obtained by optimizing GSM targeting an acceptance rate of 0.3. Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
20

---

Under review as a conference paper at ICLR 2022
g(x',' vo)
Contours of Target Density, n(x, y)
StartingPoint, (xo.
Figure Comparison of multi-scheme mixture proposals (L=8) of Equation 4 obtained by optimizing GSM targeting an acceptance rate of 0.6. Each subplot illustrates the proposal distribution's 's density when sampling from given starting position (in red).
g(x', y' yo)
Contours of Target Density, n(x, y)
StartingPoint, (Xo.
Figure Comparison of multi-scheme mixture proposals (L=8) of Equation 4 obtained by optimizing GSM targeting an acceptance rate of 0.9. Each subplot illustrates the proposal distribution's ' density when sampling from given starting position (in red).
y yo)
Contours of Target Density, n(x, y)
StartingPoint, (Xo,
Figure 6: Comparison of multi-scheme mixture proposals (L=3) of Equation obtained by optimizing Ab Initio objective. Each subplot illustrates the proposal distribution's density when sampling from a given starting position (in red).
21

---

Under review as a conference paper at ICLR 2022
g(x'y' yo)
Contours of Target Density, n(x, y)
StartingPoint, (Xo.y)
Figure 7: Comparison of multi-scheme mixture proposals (L=3) of Equation obtained by maximiz- ing MSJD. Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
g(x', y'xo
Contours of Target Density, n(x. y)
StartingPoint, (Xo.
Figure 8: Comparison of multi-scheme mixture proposals (L=3) of Equation 4 obtained by optimizing L2HMC's objective. Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
g(x'y'l y'oo. Vo)
Contours of Target Density, n(x, y)
StartingPoint, (X).
Figure 9: Comparison of multi-scheme mixture proposals (L=3) of Equation obtained by optimizing GSM targeting acceptance rate of 0.3. Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
22

---

Under review as a conference paper at ICLR 2022
g(x' Vo)
Contours of Target Density, n(x, y)
StartingPoint, (Xo.
Figure 10: Comparison of multi-scheme mixture proposals (L=3) of Equation 4 obtained by opti- mizing GSM targeting an acceptance rate of 0.6. Each subplot illustrates the proposal distribution's density when sampling from given starting position (in red).
g(x',y
Contours of Target Density, n(x, y)
StartingPoint, (Xo.
Figure 11: Comparison of multi-scheme mixture proposals (L=3) of Equation 4 obtained by opti- mizing GSM targeting an acceptance rate of 0.9. Each subplot illustrates the proposal distribution's ''s density when sampling from given starting position (in red).
23

---

Under review as a conference paper at ICLR 2022
H FURTHER EXPLANATION OF MULTI-SCHEME OPTIMIZATION RESULTS
To provide an intuitive explanation of the results of Section 7.1 we will start by restating (in a more simplified form from the original in Section 3 the overall aim of MCMC proposal optimization, namely finding a solution to the optimization problem:
9opt argmin Cig: T]
gEG
Where G is the model class of our proposal architecture, π is the target distribution, and C is the objective function chosen for the optimization. With the experimental setup of Section 7.1 we are confident that the results in Table 2/fflect proposals that are very close to the global optima of the compared objective functions. Figure T2 illustrates the placement of the global optima for MSJD within the model class of the multi-scheme proposal distribution of Equation in very abstract and simplified manner. For simplicity, we also plot hypothetical locations for RW sampling and i.i.d. resampling from the target distribution, n.
Grautti scheme Distribttion
9opt, MEID
RW M
Figure 12: Location of the optimum of MSJD when optimizing the multi-scheme architecture of Equation
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 efficiency in model class neighborhoods very close preconditioned RWM, MALA, and HMC, it is clear that MSJD and its derivatives are not robust measures of MCMC efficiency for proposal architectures with model classes like those of our example multi-scheme distribution.
When utilizing the GSM, the optimization procedure complicated slightly by targeting the user- specified acceptance rate. Effectively, this restricts G to those proposal distributions within the model class that attain the user specified acceptance rate. Figures [3 and [4 illlstrate the placement of the global optima for the GSM when targeting acceptance rates of 30% and 60% using the multi-scheme
24

---

Under review as a conference paper at ICLR 2022
Grratti-adent
Stpt. GSM
JWW
RWM
(a) Targeting acceptance rate of 60%.
(b) Targeting acceptance rate 30%.
Figure 13: Locations of the optima of the GSM when optimizing the multis-scheme architecture of Equation 3 with 8 coupling layers targeting acceptance rates (a) 60% and (b) 30%. 10%.rragg lines represent the restrictions of the model classes that attain the desired acceptance rates.
distribution of Equation 3 with 3 and 8 coupling layers. In these figures, the orange lines represent the restrictions of the model classes that attain the desired acceptance rates.
Geeutt--akkeee Invertaaionn
Ganaa-rchhne Dutriiatio
Sgn, GSM-3)
Sep
RWM
RWM
(a) Targeting acceptance rate of 60%.
(b) Targeting acceptance rate of 30%.
Figure 14: Locations of the optima of the GSM when optimizing the multi-scheme architecture of Equation 3 with 3 coupling layers targeting acceptance rates of (a) 60% and (b) 30%. Orange lines represent the restrictions of the model classes that attain the desired acceptance rates.
There two model class regimes that most relevant for the optimization of the multi-scheme architecture with the GSM in Section 7. a neighborhood of primarily inefficient. RWM-like proposals and neighborhood of approximate i.i.d. resampling proposals. With 8 coupling layers, both
25

---

Under review as a conference paper at ICLR 2022
regimes accessible while maintaining acceptance rates of 30% and 60%. In Figures Figures L and (b), this is represented by the presence of two orange surfaces within G sckeme Din one in the lower-left trepresenting the SWM-llike regime) and one in the upper-nigtt (representing the approximate resampling regime). When using 8 coupling layers, specifying target acceptance rates of either 30% 60% will still include some portion of the approximate resampling regime of the multi-scheme architecture's s model class within the restriction over which the GSM optimizes, In the end, with coupling layers, GSM-30 and GSM-60 have global optima that lie within the approximate resampling regime of the multi-schent architecture.
Moving on to the case with coupling layers, the architecture's ability approximate resampling is diminished, which greatly impacts the qualitative behaviour of the restrictions of the model class different target acceptance rates. As depicted in Figure 14 (b), the case of optimizing a multi- scheme architecture with 3 coupling layers with the GSM targeting an acceptance rate of 30% is not significantly changed, there there is portion of the approximate resampling regime that lies within the restriction of attaining acceptance rate of 30% and the global optimum of GSM-30 some point within this restriction of the approximate sampling regime. However, when using coupling layers and targeting sampling regime that falls into the restriction over which the GSM is optimized, leaving only inefficient RWM-like proposals from which to choose the optimum of GSM-60 in this case. Without significant prior knowledge regarding the interaction between the proposal architecture and the target distribution or a significant time investment testing many different acceptance rates, there is little that can be done to avoid the detrimental effect of choosing highly sub-optimal target acceptance rate when optimizing model classes like the multi-scheme distribution using the GSM.
acceptance rate of 60%, there is no member of the approximate
Model A
Model
MCMC ficieney
X 8
X Acceptance Rote
X + 8
Figure 15: Comparison of maximal MCMC efficiency (C*) attained given a target acceptance rate in the neighborhood of accepance rate 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 coupling layers around an acceptance rate of 55-56%.
The "drop-out" of the approximate resampling regime when using coupling layers with the multi- scheme distribution and targeting acceptance rate of 60% is reflective of discontinuity of MCMC
26

---

Under review as a conference paper at ICLR 2022
efficiency 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 efficiency attained at 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 fixed acceptance rate with the model classes of preconditioned RWM and MALA. So when we find RWM optimized to an acceptance rate of 20% MALA optimized to an acceptance rate of 60%, we are reasonably confident that the proposals are close to optimality and are not drastically inefficient. But, when using a proposal architecture like the sclti- distribution, L can behave very sharply with respect to the target acceptance rate.
extrrited
This is illustrated Figure 15 where the maximal MCMC efficiency (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 the multi-scheme architecture with 8 coupling layers above) behaves much and MALA, where minor errors in fixing a target acceptance rate have minor impacts on the end efficiency. 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 efficiency 56%). Here, very minor differences in the target acceptance rate can result in an orders of magnitude change MCMC efficiency for Model B. Models A and B in this example do not need differ by their proposal architecture, but may instead differ by their target distribution, which makes it difficult translate the knowledge gained from successful optimization in one task to the next task with the same proposal architecture.
would expect given experiences with RWM
target acceptance
X (in case, corresponding around 55-
Miutti schemee Distrilnation
9spt, Ab Inibio
RW M
Figure 16: Location of the optimum of an Ab Initio objective when optimizing the multi-scheme architecture of Equation 3
As robust approximations of our notion of MCMC efficiency throughout the expansive model classes of architectures like the multi-scheme distribution, Ab Initio objective functions allow us. in principle, to confidently approach MCMC proposal optimization with highly expressive neural proposals as
27

---

Under review as a conference paper at ICLR 2022
would any other traditional deep learning task. Figure [6] illlstrrtte the placement of the global optima for an Ab Initio objective function within the model class of the multi-scheme proposal distribution of Equation[]] In contrast to MSJD based objectives, the functional constraints of Section and the fitring and verification procedures of Section [6 work to ensure that Ab Initio objective functions remain robust and align with notion of MCMC efficiency throughout the model classes of very expressive architectures. In contrast to the GSM, Ab Initio objectíve functions seek to directly approximate our notion of MCMC efficiency without relying on prior knowledge of optimal acceptance rates, which are very difficult 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 efficiency and are compatible with traditional deep learning approach to optimizing neural architectures.
1 POSITION-DEPENTE
MIXTURE PROPOSAL OPTIMIZATION
To examine general objective function behavior without additional complications arising from particular proposal architectures, we first consider the optimization of proposals consisting of a mixture of gaussians (with position independent means and covariances) with position dependent component weights specified by neural network. This proposal distribution may be specified as follows:
(S (S [M4,0 E,,0) P(S =|x] = fo(
(4)
The target distribution an equal mixture of 4 standard gaussians positioned in a cross formation with maximal distance of 8 between component centers. The proposal distribution consists of gaussian components. The position-dependent weights, (x), are specified by a standard ReLU network (input dimension hidden dimension 16, hidden layers, output dimension 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 single starting point sampled i.i.d. from the target distribution from which 50 independent proposals from go [x) are generated to empirically estimate expected loss on the basis of 50 total samples. The gradient update steps take the form of Algorithm with M = 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), B is initially set to 0.5 and pa is set to 0.002. 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 replicated a total 5 times. Figures 17
20 21 and 22 show density plots of the position- dependent 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
18
19
GSM (Titsias and Dellaportas] 2019) with varying target acceptance rates.
28

---

Under review as a conference paper at ICLR 2022
g1(x'y'
92(x'y')
g3(x',y'
ga(x'y')
Contours of Target Density, m(x.y)
Contours of Component Proposal Densities, gix', ,y
0.0
0.2
0.4
0.6
0.8
10
P(gix.
mixture proposals of Equation 3 obtained by optimizing Equation 2 Each subplot illustrates the contours of a proposal component (in red) alongside density plot of the probability of sampling from that component based starting position.
Figure 17: Comparison of
91(x',' .")
92 y
gx y")
9a(X', y")
Contours Target Density, n(x, y)
Contours of Component Proposal Densities, gi(''
0.0
0.2
0.4
0.6
0.8
1.0
plg](x.)
Figure 18: Comparison of position-dependee mixture proposals of Equation 3 obtained by optimizing MSJD. Each subplot illustrates the contours proposal component (in red) alongside density plot of the probability of sampling from that component based on starting position.
91 y")
92 y")
0a(x'yy
9a(X', yn
Contours Target Density, n(x, y)
Contours of Component Proposal Densities, gix''
0.0
0.2
0.4
0.6
0.8
10
P(g(xxy)
Figure 19: Comparison of vosition-dependent mixture proposals of Equation 3 obtained by optimizing L2HMC's objective. Each subplot illustrates the contours of a proposal component (in red) alongside density plot of the probability of sampling from that component based starting position.
29

---

Under review as a conference paper at ICLR 2022
91(x' yy)
92(x,y)
gs(x'y)
g4(X',y .Y
Contours of Target Density, n(x.
Contours of Component Proposal Densities, g/('',
0.6
0.8
plgix.
Figure 20: Comparison of position-dependee mixture proposals of Equation 3 obtained by optimizing GSM targeting an acceptance rate 0.3. Each subplot illustrates the contours of proposal component (in red) alongside a density plot of the probability of sampling from that component based on starting position.
91 y')
92(x', yy)
ga(x', y")
94(x',' yy)
Contours of Target Density, n(x,) y)
Contours of Component Proposal Densities, g,(x' y")
P(gilx,y)
Figure 21: Comparison of position-dependens mixture proposals of Equation 3 obtained by optimizing GSM targeting an acceptance rate of 0.6. Each subplot illustrates the contours of proposal component (in red) alongside density plot of the probability of sampling from that component based on starting position.
91 y?
g2(x',y
ga(x'.
Dx(x'yy")
Contours Target Density, n(x,) y)
Contours Component Proposal Densities, gi(x', y')
0.0
0.2
0.4
0.6
0.8
10
Pl
Figure 22: Comparison of position-dependern mixture proposals of Equation 3 obtained by optimizing GSM targeting an acceptance rate 0.9. Each subplot illustrates contours of proposal component (in red) alongside density plot of the probability of sampling from that component based starting position.
30

---

Under review as a conference paper at ICLR 2022
The results of these optimizations provided in Table 12
Table 12: Comparison of sampling performance obtained optimizing the position-dependent mixture
proposals of Equation 4 using various objective functions. Value means are reported to most the first significant digit of standard error (reported in parentheses).
Acc.
ESS per Proposal
Grad. Per Sec.
MSJD
Rate
x²²
Ab Initio, Eq. 2
0.998(1)
34.0(3)
1.00(2)
1.05(3)
17.((3)
MSJD Opt.
0.998(1)
(5.9(4)
1.00(1)
6(1)e-4
24(2)
L2HMC Obj.
0.998(1)
66.0(2)
1.01(1)
7(1)e-4
23(1)
GSM-90
0.900(2)
30.8(3)
0.81(2)
0.81(6)
GSM-60
0.601(5)
20.9(2)
0,41(2)
0.40(4)
20(1)
GSM-30
0.297(7)
10.3(4)
0.15(1)
0.16(1)
LLD.. Resample
1.000
34.2(2)
0.96(5)
1.01(3)
We find that i.i.d. resampling remains stable global minimum of our Ab Initio objective function. The GSM objective also optimizes towards constant and equal weighting of the proposals components, but increases the variance of each component match its targeted acceptance rate. Both MSJD and 1,2HMC's modification find 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 reflected by the poor ESS performance for estimating second moments.
FURTHER DETAILS FOR LOGISTIC REGRESSION POSTERIOR SAMPLING
For this experiment, we augment the target distribution with number of auxiliary variables. Similar to the momenta in HMC, these augmenting variables are independent of the original data variables and are distributed following 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 number of augmenting variables, a, equal to twice the dimensionality of the original data, d.
This experiment again uses multi-scheme proposals, specified as follows:
ń ~ N( 1)
x ń = p1.0 ( (x)) EL.O O X' = ( Ep,0(x) 3 (z)
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 independent gaussians with mean of 0 and standard deviation of 10. The additive functions, FLL8 and μp,a. are specified by standard ReLU networks (input dimension a + d, hidden dimension (W) (a. + d), 4 hidden layers (D), output dimension a + d). The multiplicative functions, ΣL,0 and Ep.e. 8. are specified by standard ReLU networks (input dimension a + d, hidden dimension (W) (a + d). 4 hidden layers (D), output dimension a + d) followed by component-wise exponentiation. The normalizing flow, Ta. follows the NICE architecture (Dinh et al. 2014), with additive coupling layers specified by standard ReLU networks (imput dimension a-re 2 hidden dimension (W) 3 (a + d). 4 hidden layers (D), output dimension a+d, ). For this experiment, the same number of coupling layers is used (L 5) for all datasets. Table 13 summarizes the network parameters used this experiment.
Before each optimization, long equilibrated Markov Chain is obtained following 500,000 proposals of HMC tuned to attain an acceptance rate of 0.65 (using hamiltorch's 's (Cobb et al., 2019) implemen- tation of HMC). A burn-in period of 1000 proposals is utilized for this HMC tuning and is discarded.
31

---

Under review as a conference paper at ICLR 2022
For optimization, 40000 gradient steps 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 (>'(az) generated to empirically estimate expected loss the basis of 800 total samples.
The gradient update steps take the form of Algorithm with M = and N 100. Uniform sampling from this Markov Chain approximates i.i.d sampling from the target. Optimization is performed using the Adam optimizer KKngmaa and Baj 2015) with stepsize of 3e-4. For optimizing GSM (Titsias) and Dellaportas 2019), 3 is initially set to d' and Pg 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 standard error of reported variables, each optimization is replicated total 5 times. Tables 14 15
T6
17 and T8 summarize
the sampling performance measures obtained following optimization, including minimum, median, and maximum one-dimennt niee-dmensional ESS obtained across all dimensions of the original data variables. Minimum one-dimensional ESS is most indicative of sampling efficiency across the entire state space of the distribution and is therefore the focus of our comparisons.
Table 13: Specification of network parameters used in multi-scheme proposals for logistic regression sampling tasks.
PL, and JDD.
Er and E
Tg
W
D
W
D
W
D
L
German Credit
21
126
126
126
4
5
Australian Credit
15
15
90
4
90
4
90
4
5
Heart
14
84
84
5
Pima
54
4
54
4
54
4
5
Ripley
27
27
27
4
5
Table 14: Results from optimizing multi-scheme proposal distributions targeting posterior distribution of logistic regression weights for German Credit dataset. Value means are reported at most the first significant digit of standard error (reported in parentheses).
ESS per Proposal
Grad. Per Sec.
Acc. Rate
MSJD
x;
Ab Initio
0.81(1)
0.30(1)
0.56(3)
0.49(7)
11.5(6)
MSJD Opt.
0.77(2)
0.32(1)
0.52(7)
0.4(1)
15(1)
L2HMC Obj.
0.76(2)
0.32(1)
0.54(4)
0.48(5)
14.8(9)
GSM-90
0.90(1)
0.012(1)
0.005(1)
0.005(1)
GSM-60
0.61(3)
0.24(1)
0.40(3)
0.40(3)
]
10.1(3)
GSM-30
0.29(5)
0.13(2)
0.13(4)
0.13(4)
HMC
0.65
0.201(3)
0.006(3)
0.005(3)
L.D.D Resample
1.00
0.3700(3)
0.98(1)
0.983(3)
32

---

Under review as a conference paper at ICLR 2022
Table 15: Results from optimizing selti-schemme proposal distributions targeting posterior distribution of logistic regression weights for Australian Credit dataset. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
ESS per Proposal xf 0.57(8)
Grad. Per Sec.
Acc. Rate
MSJD
2;
Ab Initio
0.85(1)
1.58(4)
0.63(4)
12.2(2)
MSJD Opt.
0.77(2)
2.39(3)
0.55(2)
0.53(3)
14(1)
L2HMC Obj.
0.76(1)
2.39(3)
0.47(4)
0.41(7)
14.5(7)
GSM-90
0.901
0.088(3)
0.0105(2)
0.0105(2)
}
GSM-60
0.61(1)
1,23(2)
0.41(1)
0.42(1)
10.6(2)
GSM-30
0.32(1)
0.71(3)
0.16(1)
0.16(1)
HMC
0.65
1.43(2)
0.02(1)
0.02(1)
I.I.D Resample
1.00
1.901(4)
0.97(1)
0.97(1)
Table 16: Results from optimizing multi-scheme proposal distributions targeting posterior distribution of logistic regression weights for Heart dataset. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
ESS per Proposal xt? 0.58(7)
Grad. Per Sec.
Acc. Rate
MSJD
32
Ab Initio
0.86(2)
1.03(4)
0.63(6)
11.7(4)
MSJD Opt.
0.81(1)
1.40(3)
0.69(3)
0.32(6)
15.0(3)
L2HMC Obj.
0.82(1)
1.43(4)
0.73(4)
0.28(4)
14.7(7)
GSM-90
0,900(3)
0.138(3)
0.034(2)
0.035(1)
GSM-60
0.59(2)
0.72(2)
0.34(3)
0.35(3)
10.8()
GSM-30
0.29(3)
0.42(4)
0.13(2)
0.14(2)
HMC
0.65
0.70(1)
0.02(1)
0.02(1)
LLD. Resample
1.00
1.222(1)
0.97(1)
0.97(2)
Table 17: Results from optimizing multi-scheme proposal distributions targeting posterior distribution of logistic regression weights for Pima dataset. Value means are reported to most the first significant digit of standard error (reported in parentheses).
ESS per Proposal
Grad. Per Sec.
Acc. Rate
MSJD
F; x
Ab Initio
0.88(3)
0.18(1)
0.7(1)
0.7(1)
11.2(8)
MSID Opt.
0.84(2)
0.28(1)
0.79(6)
0.13(6)
15.0(9)
L2HMC Obj.
0,84(4)
0.27(1)
0.78(5)
0.14(8)
15.2(7)
GSM-90
0.91(1)
0.186(3)
0.69(6)
0.68(6)
GSM-60
0.60(3)
0.14(1)
0.43(3)
0.44(3)
11.1(5)
GSM-30
0.30(3)
0.08(1)
0.15(2)
0.15(2)
HMC
0.65
0.139(3)
0.007(2)
0.007(2)
I.D.. Resample
1.00
0.2118(2)
0.98(1)
0.98(1)
33

---

Under review as a conference paper at ICLR 2022
Table 18: Results from optimizing multi-scheme proposal distributions targeting posterior distribution of logistic regression weights for Ripley dataset. Value means are reported to at most the first significant digit of standard error (reported parentheses).
ESS per Proposal x² 0.88(4)
Grad. Per Sec.
Acc. Rate
MSJD
T3 0.90(4)
Ab Initio
0.97(1)
0,48(1)
11.6(7)
MSJD Opt.
0.92(1)
0.87(2)
1.35(6)
1.2(1)
14.3(8)
L2HMC Obj.
0.91(2)
0.87(3)
1.4(1)
1.2(1)
15.1(9)
GSM-90
0,90(1)
0.46(1)
0.82(3)
0.81(3)
GSM-60
0.61
0.35(1)
0,50(1)
0.51
11.3(6)
GSM-30
0.30(1)
0.20(1)
0.20(1)
0.20(1)
HMC
0.65
0.202(5)
0.15(3)
0.18()
I.I.. Resample
1.00
0.498(2)
1.00(1)
0.99(2)
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 L.HMC''s objective yields proposals whose behavior significantly deviates from i.i.d. resampling from the target distribution and result 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 numerically evidenced by severe decrease in ESS relating estimation of the distribution's '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 significantly greater than which is indicative of the proposal distribution yielding greatly negative auto-correlations. These results demonstrate that this behavior is not the result of specific 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 sufficiently complex (relative the target distribution) position dependence.
K FURTHER DETAILS FOR SCHEME SELECTION EXPERIMENT
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 mean of 0 and standard deviation of 10. Our proposal distributions are diagonally preconditioned MALA, diagonally preconditioned RWM, the multi-scheme proposal distribution of Equation and a normalizing flow intended to solely approximate i.i.d. resampling from the target.
Before each optimization, long equilibrated Markov Chain obtained following 500,000 proposals of HMC tuned to attain an acceptance rate of 0.65 (using hamiltorch's (Cobb et al. 2019) implemen- tation 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 proposal go (x) is generated to empirically estimate expected loss. The gradient update steps take the form of Algorithm with M = and = 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 following 10000 proposals from starting points independently sampled from the long equlibrated HMC chain. Our rough, upper-bound bound estimates for the MSJD and acceptance rates of the resampling normalizing were estimated using the same procedure 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
34

---

Under review as a conference paper at ICLR 2022
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, ML,8 and J.D.. are specified by standard ReLU networks (input dimension 785, hidden dimension (W) 2355, 4 hidden layers (D), output dimension 785). The multiplicative functions, ΣL.0 0 and Σp specified by standard ReLU networks (input dimension 785, hidden dimension (W) 2355, hidden layers (D), output dimension 785) followed by component-wise exponentiation. The normalizing flow, Ts. follows the NICE architecture (Dinh et al. 2014). with additive coupling layers specified by standard ReLU networks (input dimension 392, hidden dimension (W) 2355, hidden layers (D), output dimension 392). The flow 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 flow follows the NICE architecture (Dinh et al. 2014), with additive coupling layers specified by standard ReLU networks (input dimension 392, hidden dimension (W) 3140, 5 hidden layers (D), output dimension 392). The flow uses 6 coupling layers A stepsize of 3e-6 used. The resampling normalizing flow 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 total 5 times.
L CALCULATING AND EVALUATING EFFECTIVE SAMPLE SIZE
Throughout our experiments, optimized proposal distributions. From a series of samples drawn sequentially from a Markov Chain, X2, ..... NN). the simplest estimation of one-dimensional ESS can be obtained by the single chain calculation:
ESS means of verifying the sampling performance of
N
ESS ++ 2 2 Pt.
Where Pe is the autocorrelation at lag t within the series X2........ Xx). 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. (201)) 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 often viewed as providing a measure of how many effectively independent samples provided within the series (X1, X2, Xx), it is important to keep in mind that this interpretation applics only for the purpose of estimating the particular expectation [[ In general, calculating ESS from the series (
-., f(xx)) relates to the uncertainty of estimating [X] based on the series (X1, X2, Nn). As these ESS results can vary dramatically depending on the expectation of interest, [X)] 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 estimating first and second moments of the target distribution [E[X] and E[X²).
When evaluating the performance of restricted model class MCMC propsoals like RWM or MALA, ESS relating estimating first moments (E[X) often reliable performance comparison and ESS per proposal results are typically found to be less than 1. Our results in Section 77. show that these expectations will not generally hold true when evaluating the performance very expressive proposal distributions. As seen with proposals optimized with MSJD and L.2HMC's objective, very expressive neural proposals can exhibit significant negative autocorrelations that yield ESS per proposals greater than and can have first moment ESS results that mask significant 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.
35

---

Under review as a conference paper at ICLR 2022
M OPTIMIZING MCMC OBJECTIVE FUNCTIONS AND ONLINE ADAPTATION
Throughout our experiments, we utilize either exact or approximate samples from (xx) to approx- imate the outer expectations (8)) found within all compared objective functions during optimization. As a result, our gradient update steps take the form illustrated in Algorithm (note, when direct sampling from rxx) not available, we approximate this sampling by uniformly sam- pling from long equilibrated HMC chain). Algorithm requires samples from л (xx) to be available during optimization, which is typically not the case during practical applications of MCMC proposal optimization. For practical applications where samples from (xx) are not available, online adaptation procedures are used to approximate the outer expectations (Ex~)(s) [..]). with gradient update steps not advocating for the use of update steps like Algorithm|] I practice. We are instead using Algorithm speeiiiiall because it allows 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 efficiency when optimizing proposal distributions with highly expressive model classes like that of Equation (3).
taking a form like that illustrated in Algorithm 2 We must emphasize that
Algorithm 1: Gradient Update Step with Sampling from (xx).
Input: Target distribution π Txx), proposal distribution go (x" x)) objective function [(x; go] (x, x"), integers M, N 2 1, optimizer ParamOpt(C,0) (C,
for i in range .M) do
Sample X ~ π (X);
for j in range (0, .V) do
Sample 3 ~ go (x);
end
end
C - MN
1 En, g0
0 + Pa
For the purposes of our experiments, we do not employ online adaptation because the adaptation procedures introduce additional confounding factor to the analysis of performance differences between objective functions. By using Algorithm I 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: Gradient Update Step with Online Adaptation.
Input: Target distribution π (x), proposal distribution ge (x" x), objective function [[m; gol (x, X²), integers M, N 2 1, M initial Markov Chain states x,, optimizer ParamOpt( [,, for i in range (0, M) do for in range .) do Sample X ~ go (x(xxxx))
end end
L MN Σns f1x X(); 0 ParamOpt(C,00 0);
for i range (0, M) do
Sample X (x' x:);
a + 11,in(1, xiX" 1ee
=(, ) J:
Sample ~ Uni 00. 1]: if u < a: then
X, < x':
end
36

---

Under review as a conference paper at ICLR 2022
distributions. The use of adaptation procedures like A 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 modification, with neural MCMC proposal model classes like Equation (3).
It is important 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 efficiency, 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 modification to the optimization of proposal model classes like Equation (3) or whether some modification required to maintain compatibility with highly expressive proposals.
M. REPLICATION OF SECTION 7.2 USING ONLINE ADAPTATION
Although the applicability of adaptive procedures like Algorithm 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 [ all experimental details remain unchanged from those listed in Appendix J unless otherwise noted. For optimization, 60000 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 specified in Appendix [ 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 19: Comparison of sampling performance obtained optimizing augmented multi-schent
proposals of Equation 3 using online adaptive procedure of Algorithm 2 and various objective functions targeting posterior distributions of parameters for logistic regression of UCI datasets. Value means are reported to at most the first significant digit of standard error (reported in parentheses).
German Credit (d=21)
Heart Disease (d=14)
Acc. Rate
ESS per Proposal
Acc. Rate
ESS per Proposal
MSJD
MSJD
Z2 0.61(4)
JT
x,
Ab Initio, Eq.
2
0.85(1)
0.305(4)
0.54(7)
0,87(1)
1.05(2)
0.65(4)
0.61(5)
MSJD Opt.
0,82(1)
0,36(1)
0.64(4)
0.5(1)
0.79(1)
1.38(5)
0.66(3)
0.29(5)
L2HMC Obj.
0.84(2)
0.35(1)
0.63(4)
0.54(7)
0.84(1)
1.41(3)
0.76(4)
0.36(4)
GSM-60
0.58(3)
0.23(1)
0.37(3)
0.36(2)
0.60(2)
0,80(2)
0.40(3) 0.42(3)
LLD.. Resample
1.00
0.3700(3)
0,98(1)
0.983(3)
1.00
1.222(1)
0.97(1) 0.97())
These results using online adaptation are consistent with those obtained using Algorithm 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 Tike Algorithm 2 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 the optimization of neural MCMC proposals:
By using Algorithm 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
37

---

Under review as a conference paper at ICLR 2022
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 needed to validate their appli- cability the optimization of highly expressive proposal distributions. Until such validation is performed, it is difficult to interpret the results of experiments using these adaptive procedures (unmodified from their use with very restricted proposal model classes) for the purpose of comparing the performance of MCMC objective functions and expressive proposal architectures.
38


Reviews:
[{"review_id": "xEeCha77_H", "reviewer": "Reviewer_pwUX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n\nI think the problem was amazingly motivated. The description of the objective function was very clear: both in terms of its description as well as its need. I am also fascinated that for a fixed $A$, the proposed approach is able to reproduce classic theory regarding the optimal acceptance probability for random walk metropolis and metropolis adjust langevin diffusion. \n\n# Weaknesses\n\nI first want to note that the length of the weakness section relative to the strength section does not reflect my opinion of the paper. While I'm a big fan of the approach I'm somewhat confused about the actual proposed method. First, the proposed objective function is an expectation with respect to the target density, $\\pi(x)$. While I understand that theoretically, this is important, practically this is *very* restrictive as one either needs to have access to the target density--which if this is the case, why do MCMC--or one needs to run access to approximate samples, but if you have access to these samples why run this method. \n\nThe second thing that isn't clear is the actual implementation: is the proposal optimized before running the MCMC sampler or is it optimized while the MCMC algorithm is running? Also, in Appendix D it says \"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 proposals are generated to calculate expected loss\". I'm confused about what is meant by 50 proposals? Is the gradient estimated using 50 samples? Is the sampler run for 50 time steps, and those samples are used to approximate expectations? This same question is also relevant to the experiment section as well. The confusion around proposals makes it much harder to understand what is actually going on. \n\nLastly, I have some comments about the experiments. I think for all quantities, it would be useful to state whether higher/lower is better.\nIn experiment 7.1, the authors state \"Position dependence within the proposal distribution lead MSJD and L2HMC’s modification to optimize towards arbitrarily non-ergodic proposal distributions, sampling only either the horizontal or vertical components of the target distribution.\". From this description, I would expect the MJSD and L2HMC to not mix well but according to the results in Table 3, they both have very high ESS (also I thought ESS was upper bounded by 1?) and MSJD. I noticed a similar discrepancy in 7.3 where the authors state \"The resampling normalizing flow 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.\" but the resampling normalizing flow has the highest Ab Initio! Also, I'm confused about what resampling normalizing flow is (I also don't think it was described anywhere in the paper).\n\n# Comments\n1. There should be a section on how the ESS was computed.\n2. It would be beneficial to the reader if MSJD was defined.", "summary_of_the_paper": "The authors design an objective function from first principles that can be used to optimize proposal distributions for MCMC. They empirically verify that their objective function with a fixed hyperparameter--tuned a priori on another task--can reproduce theoretical results over a variety of target distributions.", "summary_of_the_review": "While I like the motivation of the paper, I think the actual algorithm and the experimental details are somewhat opaque. This can be easily fixed by fixing the manuscript and making these things very clear. I'm also somewhat concerned about the practicality of the algorithm and I think this needs to be addressed by the authors in the manuscript. For these reasons, I lean towards a weak reject though I am more than happy to increase my scores if the authors address my concerns.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "cRqGS-ROpLM", "reviewer": "Reviewer_Etgp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The paper tackles the important problem of developing objectives to automatically tune MCMC proposals. This is necessary to tune complex and flexible proposals parameterized by neural networks.\n- The first principles provide a sensible guide to design of objectives to tune MCMC proposals. So far different objectives have been proposed without a clear underlying framework. This paper takes a first step in that direction.\n\nWeaknesses:\n- While the formulation with the guiding principles is quite general, the only concrete objective proposed is a modification of GSM. Also, even in this case, some parameters are tuned using some ad-hoc methods. For instance, the value of the constant A is set by targeting some optimally known acceptance ratio for a specific setting, and a factor ``d'' is added to correct for the magnitude miss-match between different terms in the objective. Thus, despite the generic presentation, it is not clear to me whether this framework actually simplifies the task of creating new objectives to tune MCMC proposals.\n- I'm not fully convinced by the empirical evaluation. The models considered, while illustrative, are quite simple (e.g. Gaussians, low dimensional logistic regression). In addition, the method is not tested for online tuning of MCMC kernel parameters, which is how it would be used in practice. All the evaluations are using true samples from the target, or samples obtained running a long HMC chain. While the authors add this as future work, I think including it in this work would be quite good.", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall I think the whole idea of providing a set of guiding principles for the design of objectives to tune MCMC proposals is interesting and relevant. However, I feel the presentation in this paper may be a bit too generic, and the derivation of the one concrete objective requires a decent amount of manual tuning. Thus, it is not extremely clear to me whether this framework, as presented, simplifies the development of new objectives significantly or not. In addition, I think the concrete objective proposed could be tested more thoroughly, specifically in some settings where samples from the target are not available (i.e. tuning its parameters online as most adaptive MCMC methods do).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "DzL5ujK8NK8", "reviewer": "Reviewer_qDJ6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Although the paper approaches an important problem in the MCMC field, its contribution is incremental. One may consider it as a special case of the objective from [Titsias, 2019]. The proposed approach has several major issues:\n- The conjecture of the existence of universal $\\beta$ is not well-supported. The authors motivate the choice of GSM objective by verifying several properties. However, further developments are based on the conjecture that there should be some optimal tradeoff between the acceptance rate and the proposal entropy for all target distribution and all proposals. This conjecture is validated neither theoretically nor empirically.\n- The way of finding $\\beta$ by \"manual trial and error\" is not acceptable. Although the authors mention that it is possible to estimate the constant using some grounded techniques, they describe their procedure as \"trial and error\". This fact cannot persuade practitioners to use the proposed objective. The constant should be estimated at least by using grid search demonstrating the performance of the algorithm for different values. Furthermore, this empirical study has to be extended to different proposals and different targets demonstrating the universality of the constant. Finally, the variance of the estimation should be reported.\n- The evaluation of the proposed approach is not practical. To be more precise, the proposed objective requires samples from the target distribution. Throughout the experiments, the authors collect these samples separately running a long HMC chain. However, obtaining a large set of uncorrelated samples from the target using HMC could be impossible in practice. Not to mention that this is the final goal of any MCCM algorithm. Note that in [Titsias, 2019], this issue is bypassed by collecting the samples online and treating them as samples from the target distribution. This setting is very important for practical applications and should be explored empirically.\n\nOther comments:\n- Section 7.2 is not motivated. Indeed, some MCMC algorithms use augmented space for sampling. However, the reason for using it here is unclear. This section could give more insights to the reader if two algorithms have been compared (with augmented space and without).\n- Section 7.3 is unclear. The authors conclude that the proposed objective yields a robust measure of MCMC efficiency and can be used as a metric for comparison. However, this conclusion is not argued. In all previous examples, the authors rely on ESS to measure the efficiency of chains and motivate the effectiveness of their objective also by ESS. Thus, we have a question why would one use the proposed objective instead of ESS to measure the efficiency?\n", "summary_of_the_paper": "The paper approaches the problem of learning the optimal parameters of the MCMC chain. Namely, it tries to design an objective function that allows for the learning of a perfect proposal distribution in the Metropolis-Hastings algorithm. The main contribution of the current paper is based on the developments from [Titsias, 2019]. [Titsias, 2019] proposes the following objective for the optimization of the proposal\n$$\\text{GSM} = \\text{Acceptance Rate} + \\beta \\cdot \\text{Entropy of the proposal}.$$\nMaximization of this objective results in better mixing properties of the chain. Indeed, the next sample comes from a distribution with high entropy and should be accepted with a high probability, and the proposal should converge to the independent sampler in the perfect scenario. \n\nHowever, the choice of $\\beta$ remains to be a question since it defines the tradeoff between the acceptance rate and the entropy, which is known to be different for different proposals and different targets. The original paper [Titsias, 2019] solves this issue by adaptively setting the acceptance rate to some desired value.\n\nThe current paper then can be summarized as follows. Instead of setting $\\beta$ adaptively based on the acceptance rate, the authors propose to set it constant. The constant is then found by \"manual trial and error\" in such a way that the acceptance rate of Random-Walk Metropolis-Hastings is 0.234 for high-dimensional standard Gaussian, which is known to be optimal for this task. Then this constant is claimed to be universal for all targets and proposals. To verify this empirically, the authors first optimize the parameters of MALA and RW Metropolis-Hastings for several synthetic targets (Uniform, Laplace, Cauchy, Gaussian) and demonstrate that the procedure yields the right parameters.\n\nFinally, the authors perform an empirical study of the proposed procedure. They use the proposed objective to learn a proposal parameterized by the normalizing flow and compare it to the optimization of the original objective from [Titsias, 2019] targeting different acceptance rates. The target distributions for the empirical study are a mixture of four Gaussians (2-D), Bayesian logistic regression (21-D, 14-D).\n\n[Titsias, 2019]: Titsias, Michalis, and Petros Dellaportas. \"Gradient-based adaptive markov chain monte carlo.\" Advances in Neural Information Processing Systems 32 (2019): 15730-15739.", "summary_of_the_review": "The paper proposes and analyses a special case of the objective from [Titsias, 2019]. Unfortunately, the empirical study is not sufficient to make positive conclusions about the proposed technique.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "89H6JKBQwew", "reviewer": "Reviewer_cPU7", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Finding an effective objective function for training neural MCMC proposals is of great importance for MCMC methods. However, the effectiveness of the proposed Ab Initio objective requires further justifications. Below are my major concerns.\n\n1. Although the motivation for the new objective is quite general, the author ends up with a rather heuristic formulation of the Ab Initio objective matching the optimal acceptance rate of RWM for multivariate Gaussian targets. The general applicability hence is quite doubtful given the relatively simple target distributions used in the demonstrations. The author may try out more complicated models. Also, no significant advantages over MSJD and L2HMC had been found in the experiments. \n\n2. The computation of the Ab Initio objective is not clearly stated. Given that \\pi is unknown, the expectation in the first term of L[g;\\pi] seems to be intractable in practice.\n\n3. The author claimed (3) can approximate HMC proposals. This should be verified by comparing with the hand-tuned HMC methods.", "summary_of_the_paper": "This paper proposed an objective function for training neural MCMC proposal, called Ab Initio objective. The new objective seems to improve upon existing alternatives by maintaining property and representation invariance, and reducing the requirement of prior knowledges (e.g., optimal acceptance rate). The author demonstrated the effectiveness of Ab Initio objective on several benchmark models.", "summary_of_the_review": "The proposed Ab Initio objective seems to be an interesting addition to the current candidates of training objectives for MCMC proposals. However, given the heuristic formulation and inadequate empirical evidence, I am cautious about its general applicability.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes guiding principles with which to design objective functions for proposal distributions for MCMC. They design one such objective based on GSM (Titsias and Dellaportas, 2019). The two concerns raised by reviewers that resonated the most with me were:

- it was not clear that the actual proposed objective was the best way to implement these guiding principles
- a weak empirical evaluation that did not consider online tuning and high-dim, highly non-Gaussian targets. 

After rebuttal, revision, and discussion, reviewers felt that the authors did a reasonable job of addressing the issue of online tuning, but very highly non-Gaussian targets were not addressed. There was still a sense that the ultimate instantiation of the design principles was a somewhat adhoc loss. Ultimately, I think that this work is just below the bar for acceptance and it can be improved by clarifying the choices made in implementing the objective and some more ambitious experiments.
Decision: Reject