Enhanced gradient-based MCMC in discrete spaces
Authors: Benjamin Rhodes
Michael U. Gutmann
Venue: Accepted by TMLR ()

Abstract: The recent introduction of gradient-based Markov chain Monte Carlo (MCMC) for discrete spaces holds great promise, and comes with the tantalising possibility of new discrete counterparts to celebrated continuous methods such as the Metropolis-adjusted Langevin algorithm (MALA). Towards this goal, we introduce several discrete Metropolis-Hastings samplers that are conceptually inspired by MALA, and demonstrate their strong empirical performance across a range of challenging sampling problems in Bayesian inference and energy-based modelling. Methodologically, we identify why discrete analogues to \emph{preconditioned} MALA are generally intractable, motivating us to introduce a new kind of preconditioning based on auxiliary variables and the `Gaussian integral trick'.

Published in Transactions on Machine Learning Research (10/2022)
Enhanced gradient-based MCMC in discrete spaces
Benjamin Rhodes University of Edinburgh
ben.rhodes Ged. ac. ac.uk
Michael U. Gutmann University of Edinburgh
michael.gatmannneeeee uk
Reviewed on OpenReview: https: openreview. formm? id- j d5NFUJ
Abstract
The recent introduction of gradient-based Markov chain Monte Carlo (MCMC) for discrete spaces holds great promise, and comes with the tantalising possibility of new discrete counterparts to celebrated continuous methods such the Metropolis-adjusted Langevin algorithm (MALA). To- wards this goal, we introduce several discrete MetropolisHeestingss samplers that conceptually inspired by MALA, and demonstrate their strong empirical performance across range of challeng- ing sampling problems in Bayesian inference and energy-based modelling. Methodologically, we identify why discrete analogucs to preconditioned MALA are generally intractable, motivating us to int roduce new kind of preconditioning based on auxiliary variables and the 'Gaussian integral trick'.
1 Introduction
Gradient-based Markov Chain Monte Carlo (MCMC) offers an efficient, robust way to sample from a wide-class of probability distributions. The gradient serves concise descriptor of local geometry, which can be exploited when designing MCMC transition operators. In continuous spaces, such operators are often based on the Langevin diffusion (Roberts Tweedie, 1996; Roberts & Rosenthal, 1998) or Hamiltonian Monte Carlo (Duane et al., 1987; Neal et al., 2011), which can be unified under a single complete framework (Ma al., 2015). Until recently, gradient-base operators were only viable for continuous distributions after all, the standard gradient is undefined
Sibbs-with-Gradieen
Nomm-Costrrined Gradient
Preconditioned AVG
39.50
28.71
15.50
8.64
Figure 1: Metropolis-Hastingg proposal distributions for an existing gradient-based sampler (left) and our proposed ones (centre & right). The blue contours delineate a continuous function that, when restricted to the discrete lattice, equals the target log-probability function (up to a constant).. The white circles represent the proposal distribution given the current state of the Markov chain (black X). Larger circles mean higher probability mass. The new samplers can update multiple dimensions time, which can lead more efficient mixing of the Markov chains. In addition, preconditioned AVG (right) accounts for the strong positive correlation between the two dimensions.

---

Published in Transactions on Machine Learning Research (10/2022)
for discrete domains. However, Grathwohl et al. (2021) made the important observation that many probability mass functions are naturally thought of as the restriction of a continuous function (defined in e.g. Rd) to a discrete subset (e.g. N Gradients in this ambient space can thus inform the design of transition operator in the restricted discrete space. They demonstrate this via promising new method, Bilbss-wthh-Gradiee (GWG), that can be seen as gradient-based version of the 'locally informed proposals' introduced by Zanella (2020).
Whilst the use of pradient-informatioo makes GWG appealing, the method has multiple unresolved limitations. It:
(L.1) cannot update more than one dimension a time.
(L.2) does not exploit second-order interactions in the target distribution to define the proposal distribution.
In this paper, we will address these shortcomings via the introduction of several new discrete gradient- based sam- plers. Our strategy will be define discrete analogues to continuous gradient-based samplers that do not have these limitations. Specifically, we consider the celebrated Netropolis-Adjustee Langevin Algorithm (MALA), which is free from limitation 1.,1, and its preconditioned variant. (PMALA), which is free from T.2. 'Discretising' these samplers is not a straightforward task: (P)MALA is typically viewed through the lens of stochastic differential equations (SDEs), which are not casily translated into discrete state-spaces. Fortunately, we identify two non-SDE characterisations of (P)MALA that are more readily imported into the discrete setting, resulting in the following methods (illustrated in Figure 1):
The Conn-Consta Gradient (NCG) sampler is constructed by viewing MALA as a locally- informed proposal (Zanella, 2020) whose domain is restricted to a discrete space. NCG replaces the Hamming-ball constraint in Gible-witt--Gradientts with a soft norm constraint, enabling multiple dimen- sions to be updated a time, thereby addressing limitation L.1. Unfortunately, preconditioning this sampler in a manner analogous to preconditioned MALA does not appear possible, leaving limitation L.2 unresolved.
The Mexiliary-Variable Gradient (AVG) sampler is constructed by viewing MALA as a marginalised auxiliary sampler (Titsias & Papaspiliopoulos, 2018) whose domain is restricted to discrete space. The resulting sampler is conceptually and algebraically similar to NCG, but incurs extra computations that make it less appealing. However, AVG has the major advantage that it can be extended to include a kind of preconditioning, thereby addressing limitation L.2. Intriguingly, this Preconditioned AVG (PAVG) generalises an existing auxiliary-variable scheme (Martens & Sutskever, 2010) that is only applicable to pairwise Markov Random fields (a.k.a. Boltzmann machines).
The value of these new methods is two-fold: 1) Empirically, they show superior sampling efficiency compared to several important basclines across a range of problems. In particular, they outperform Gibbs- Goobb--wth--Grddenttt validating our hypothesis that L.1 and T.2 are indeed limitations that should be addressed. 2) Methodologically, they demonstrate how new discrete gradient-basse samplers can be derived via the frameworks of Zanella (2020) and Titsias & Papaspiliopoulos (2018). To our knowledge, we are the first paper to use the latter framework for discrete problems, and our derivation of PAVG shows that auxiliary variables open up fundamentally new possibilities for constructing gradient-based samplers.
2 Background
Throughout this work, we will be dealing with distributions p(s) that are only available up to proportionality ie.
logp(s) = -f(s) log Z
(1)
where is known function defined on set Sd and the value of Z is presumed unknown. Depending the context these distributions may be continuous, with S S= R and Z fge exp(J(s))dss or discrete, with S C IR and Z = Egest (s)) In the discrete setting, our exposition focuses on two important cases: binary vectors for which S = (0, and finite ordinal data for which S = (st..... **) is ordered set of real numbers. However, in Appendix A we show that our methods also easily extend to one-hot vectorised categorical data.
'Tee NCG method was independently discovered by Zhang et al. (2022) (first made public on Arxiv on the 20th June 2022); see related work discussion Section 4.
2

---

Published in Transactions on Machine Learning Research (10/2022)
For continuous distributions, it common to assume f differentiable with derivative V... discrete settings, the standard definition of a derivative inapplicable and the notation Vf is undefied. However, the recent work of Grathwohl et al. (2021) provide useful definition of V in discrete settings that applies to many (but not all) functions f. Their key observation is that many functions that accept discrete imputs also accept continuous inputs: for instance an Ising model has the form f(s) = b²ss s Js, where the elements of s are usually constrained to {-1,-1}, but if we instead assumed them to be real-valued, f would still be a valid function. So many discrete functions can be uniquely extended to R by leaving their mathematical form intact, and only altering the domain. If we further assume that this unique extension differentiable with derivative Vʃ, then we can use this as the definition of the 'gradient' of discrete-input function f.
2.1 Metropolis- Hastingss
A common approach to sampling such distributions is to use Metropolis- HaropoolisHHastigsss (MH) Metropolis et al. (1953); Hastings (1970), which evolves a chain of samples by iteratively sampling proposal S1+1 ~> g(S++1 [sc) and accepting it with probability
min {1, exp(f(S(+1)) g(se |s441)
(2)
exp(J(s:)) ((s++1 [s))
The key challenge here is the choice of the proposal distribution q. Roughly speaking, there are four desiderata (Robert et al., 1999) i) tractable to sample ii) tractable to evaluate iii) yields reasonable acceptance rates and iv) the resulting chains have weak dependencies between successive states. Satisfying all four criteria is difficult, especially in high dimensions. Most 'obvious' approaches fail meet at least one criterion; for instance a uniform distribution over neighbouring states of 8, can satisfy the first three criteria, but fail dramatically on the last.
2.2 MALA and its preconditioned variant
Mddjopolis-AAjusted Langevin Algorithm (MALA) (Roberts & Tweedie, 1996; Dwivedi et al., 2018) is an effective comlinuons-space MH sampler that uses gradient-based proposal distribution of the form
- N + - e.)
(3)
where is tunable step-size. If we were to skip the MH accept accept/reject step, then iteratively sampling this proposal distribution corresponds to a discrete-time simulation of an SDE whose stationary distribution is the larget p(s). Such time-discretisations induce errors which the MH step corrects for.
Many target distributions exhibit strong second-order interactions between dimensions. A natural way to handle this via a suitable linear change of coordinates; a technique known as preconditioning Preconditioned MALA (PMALA) (Roberts & Stramer, 2002) has a proposal distribution of the form
(s st) = Ns; + c (Vs),,
(4)
where M is a user-specified symmetric pesitive-defintt matrix. This proposal also corresponds to a discretised SDE whose stationary distribution is p(s).
2.3 Locally-informed proposals
Zanella (2020) proposed different framework for incorporating local geometry into an MH proposal distribution, with the advantage of being applicable to discrete spaces. They define pointuiis--iformeed proposals of the form
q(s 8²) g(exp([(s) - ((s)))
(5)
(st)
where g is a non-negative 'balancing' function, K2 symmetric kernel whose 'width' is controlled by a (e.g. a uniform distribution over a local ball of radius σ) and (s() is a normalising constant.

---

Published in Transactions on Machine Learning Research (10/2022)
The intuition here that we take an 'uninformed" kernel and re-weight according to the target distribution, biasing our proposal in favour of higher density regions. Zanella (2020) identity class of balancing functions g that are optimal when the kernel is sufficiently local (i.e. is small); an important member of that class is the squarc-root function g(x) = VT.
2.4 Gibbs-witt--Gradients GWG
A major challenge when using locally-informed schemes in discrete spaces is the cost of normalising and sampling the proposal distribution. Even for the most simple choice of kernel uniform distribution over a Hamming ball of radius 1-the cost is (dd) evaluations of f for d-dimensional problem. To reduce this cost, Grathwohl et al. (2021) use the innovative trick of treating f as a continuous function and approximating it with a first-order Taylor expansion about the current state St, giving a first-order informed proposal of the form
exp 2 VS(ss) st) T(s (s:))
q (x)
(6)
Z(s:)
where (sx)) = E (s.) exp (1 ( - s:)) and H (st) is a Hamming ball of radius 1 around Sy (i.e. those points that differ from s, in single dimension). This proposal distribution is generally inexpensive to normalise and sample as the dominant cost comes from single gradient computation f(()).. Once we have computed this gradient, the normaliser is just a sum over d inexpensive terms. These d terms are then reused for sampling: we simply sample from a categorical distribution with d states (representing which dimension of S, we will change) with probabilities given by the d possible values of exp f(.) (s)ssss))(sss))) Zss) for E H1{))..
There two significant limitations of this proposal distribution: (L.1) For each gradient computation, we only update one dimension. Updating one dimension at a time can be slow, and a straightforward approach to 'broaden' the GWG proposal by using larger Hamming ball is prohibitively expensive- see Appendix C for more details (L.2) Neither the linear approximation of the target distribution nor the Hamming ball kernel account for second-order interactions between dimensions. That is, equation contains no cross-terms S;8y.
3 Discrete analogues to MALA
Our goal in this section to design new discrete gradient-based MCMC schemes that overcome the aforementioned limitations. We achieve this by 'importing" two different characterisations of the cetropolis-adjustee Langevin Algorithm (MALA) into discrete setting.
3.1 Norm-constrained gradient sampler: NCG
As discussed, the use of Hamming balls in dimension time. Instead, propose to soft norm 'constraint', transforming equation 6 into
makes it challenging to update more than
(s(s st) x exp ( ( (ss
- si)) si))ex))) exp ( (- --11 -=13 2c
(7)
Such first-order informed proposal was already discussed by Zanella (2020) in the context of continuous spaces, where, after normalising, it corresponds exactly to the proposal distribution used by Metropolis-adjustel Langevin Algorithm (MALA) in equation 3. Our key insight is that the functional form of this proposal is still valid in discrete spaces. Moreover, an important benefit of this functional form is that it can be factorised across all dimensions. After factorising and rearranging" terms, we 3
(s sx) = II ( -s.. f=1
exp(x) Es exp(x)
(x)
(8)
2e
Due its factorised structure, this proposal distribution is efficient to evaluate and sample, making it straightforward to use in Metropolis- Hastiiss scheme; see Algorithm B.1 for pseudocode and Appendix C for further discussion of
When rearranging, we can ignore any terms the exponential that not depend on s, since they are absorbed into the normaliser. 3 The definition (x) abuses notation by assuming x a function of E S, and that we sum over all values of s the denominator.

---

Published in Transactions on Machine Learning Research (10/2022)
computational cost. Like MALA, has a step-size parameter that controls the 'width' the distribution. Larger step-sizes enable proposals Sc+1 that differ from s in multiple dimensions; this is a key property that GWG lacks. Finally, we note that NCG should be viewed as an alternative to (and not generalisation of) GWG, since there is setting of the step-size that recovers GWG.
3.1.1 The intractability of preconditioning NCG
It's natural to wonder: if discretising MALA yields NCG, what does discretising preconditioned MALA yield? Unfortunately, find that this question leads to dead end. To see why, we first note that PMALA, just like MALA, can be viewed as first-order informed proposal by replacing the Euclidean norm ||s sell² in equation 7 with a Mahalanobis norm, giving
(e(s) si) & exp ( ()V)(s))) -s-s(s))s))))))))))))))))))) si)) exp (- 1 2 (-- (s- (s-ss) si))
(9)
Specifically, we obtain PMALA by noting that, for E Rd, the hight-hand side can be rearranged and normalised to obtain the Gaussian distribution in equation 4. However, when we restrict to discrete domain, this proposal distribution is a pairwise Markov random field that, in general, is intractable to normalise and sample from, making it unusable within a Metropolis-Hastinnee scheme. The key point here is that log-quadratic distributions are tractable in Rd, but this tractability vanishes as soon as we alter the domain.
One could avoid this intractability by choosing a highly restricted, sparse form for M-1 (e.g. diagonal*). We prefer to instead focus the more general problem of constructing a gradient-based sampler that can work with arbitrary M. To solve this challenge, we depart from the bocally-ifformed framework of Zanella (2020) and introduce an alternative MH-framework that operates in an extended state-space. This framework also contains a discrete analogue to MALA that resembles NCG, with the key distinction that it admits kind of tractable preconditioning.
3.2 Auxiliary Variable Gradient sampler: AVG
We follow Titsias & Papaspiliopoulos (2018) in showing how MALA can be derived as part of a continuous--pace auxiliary variable framework. Subsequently, we show how to translate the method into a discrete space.
First, our ionuu state s E Rd augmented with Gaussian auxiliary variables z Rd, to give an unnormalised target density x(s, z) = exp(f(s)) (/
In theory, this distribution could be sampled in a Block-Gibbe fashion via alternate sampling of z ~ N(z:s /e/2,, I), and S ~ a(s 2i) x T(s, 2:). However, for general functions ʃ this second sampling step is intractable, so replaced with an MH accept-reject step using the proposal distribution:
(s(s| s:) x exp(f(s.) + vf(s)) s:)) ///2,)
(10)
where we have approximated r(s, z1) via a Taylor expansion of f(s). After rearranging and normalising, we obtain (s Ze, mp) = Nss Ve/2z, + (s))
(11)
Equipped with this proposal distribution, block-wise MH sampling can then be performed in the extended state space. Alternatively, the latent variables could be marginalised out, which yields the MALA proposal in equation 3
(s st) f N(z:s:/Ve/22 T) zı, dz N(s: + cf(
(12)
We propose to apply the auxiliary variable procedure just described to discrete state-spaces. Our starting point is equation 10, which is still well-defined if we assume s belongs to Sd C Rd and the gradient operator is defined as in Section 2 i.e. using the discrete Taylor approximation 'trick' of Grathwohl et al. (2021). However, if we now rearrange and normalise equation 10, we no longer obtain a factorised Gaussian distribution, but a different kind of factorised distribution II (s. St) where each factor has the form
(5: Ve,,, 81) = (<((((()) -
+ 2
exp(x)
ox)) = Es Exp(xx)
(13)
4This possihility was explored by the concurrent work of Zhang et al. (2022).

---

Published in Transactions on Machine Learning Research (10/2022)
We can easily evaluate and sample this proposal distribution, enabling us to perform block-wise MH sampling as summarised in Algorithm B.2. We refer this method as the Auxiliary Variable Gradient (AVG) sampler. In appendix D we provide some intuition how to interpret the auxiliary variables used in AVG via toy example.
Instead of using bise MH, one might try to marginalise out Zt. giving (e 5²) = II Bz.. ge, | s1),
However, we do not use such a marginalised proposal distribution in this paper, since the required expectations have no closed-form solution, necessitating careful numerical approximations. More fundamentally, we focus on the block-wise scheme as it admits an effective kind of preconditioning.
3.3 Preconditioned Auxiliary Variable Gradient sampler: PAVG
If the auxiliary-variabbl view of MALA yields AVG, what does this view imply for preconditioned MALA? Unfor- tunately, much like in Section 3.1.1, this route of enquiry encounters difficulties. As we explain in Appendix E, we can obtain PMALA by changing the choice of conditional Gaussian in our derivation of MALA above, but doing so sabotages the tractability of the corresponding discrete proposal distribution.
Fortunately, we now show that a different kind of 'preconditioning" is viable. We start by noting that there were essentially two ingredients in our derivation of AVG: i) a local approximation of f(s) and ii) choice of conditional auxiliary distribution. We will alter both of these choices in the subsequent derivation. First, we make a second-order approximation
f(s) ~ f(s)) + fss - s:) + (1/2)(s - s.) - St),
(14)
where the second-order term uses a global (independent of t) symmetric positive definite matrix M. The nature of this approximation, and how to choose M, is discussed below in Section 3.3.1.
We then define an unnormalised joint distribution T(s, z) exp(f(s) ((2; M. I), where M, := M M2(2/c)I Just like in our derivation of AVG, Gibbs sampling of z and is prevented by the intractable x zz) x ats 2,) and so we replace this intractable sampling step with an MH accept/reject step. Leveraging equation 14, we approximate rs|| ze) with the following proposal distribution
(s 24, 8t) xx exp(s VS(s)) s;) + (1/2) - -x) M(s (s)) Nix M I)
(15)
exp ([Vf(s;) Ms, + +M// M1/221]² -(1/c)s'ss. s).
(16)
II (V](st); (Ms.): + (M) 2z); (2) 1
exp(x)
where x(x) Es exp(x))
(17)
which is fully-factorised and thus tractable evaluate and sample. The resulting block-wise MH sampling scheme is called Preconditioned AVG (PAVG) and is summarised in Algorithm B.3. By comparing equation 13 and equation 17, we see that setting M 0 recovers AVG.
The key step in obtaining a factorised proposal distribution occurs when going from equation 15 to equation 16, since the cross-terms (s18,, i * j) cancel out. Using Gaussian auxiliary variables to induce such cancellations has a long history in statistical physics where it is known as the Bulbbard-Stratonoviie the 'Gaussian integral trick" (Hertz al., 1991). The trick has also boen used in the Machine learning literature; first by Martens & Sutskever (2010) and then extended by Zhang et al. (2012). In both these prior works however, the target function was exactly quadratic i.c. f(s) bTs s" Ms, in which case the proposal distribution above is 'exact' and an MH accept/reject step is unnecessary. Thus, the proposed PAVG approach can be seen as an extension of Martens & Sutskever (2010) to more general target distributions."
transform (Hubbard, 1959) or
3.3.1 Choice of preconditioning matrix M
The above derivation assumed an unusual kind of quadratic approximation of [() in equation 14 that mixed a local first-order Taylor approximation with a global second-order term. A more standard approximation would use
5 'Important caveat: Martens Sutskever (2010) allow M to have negative cigenvalucs. We can also allow this, with small modifications to our proposal distributions as described in Appendix F.

---

Published in Transactions on Machine Learning Research (10/2022)
a second-order Taylor approximation. Our foremost reason for avoiding this option is that Hessian computations typically cost (dd evaluations of f. which would greatly increase of the cost of one iteration of PAVG, which, with fixed matrix M, only costs (1) evaluations of f.
Intuitively, one may hope that the approxination in equation 14 is reasonable whenever there are global second- order interactions in the target distribution. We propose to capture these second-order interactions by building on ideas from adaptive contimuous-space MCMC (Rosenthal et al., 2011), where one can, for instance, estimate an empirical covariance or procision matrix Memp from a 'dataset' of samples 'D obtained during a burn-in period. Given this malrix, we then use a common procedure (e.g. see Algorithm of Andrieu & Thoms (2008)) of re-scaling it by an adaptively learned parameter 7. We give full details in Appendix F. A benefit of this re-scaling is that we can automatically 'fallback' to AVG by learning 7 = O. However, in practice, we find that values of y far from 0 are learned. As final remark, we would like to emphasise that we have no strong apriori reason for thinking that scaled covariance/ ppreciiion matrices should satisfy equation 14; nevertheless, our empirical results show that PAVG works robustly across a range of scenarios and always outperforms the special case of AVG. Future work may well benefit from finding more theoretically principled choices of M that (unlike Hessians) are cheap to compute.
4 Related Work
Most closely related to our work is the concurrent paper of Zhang et al. (2022), who introduce a discrete "Tangevin- like' sampler called DMALA that coincides exactly with NCG. Their empirical results agree with ours in showing that NCG/DMALA is an effective sampler that robustly outperforms GWG. Beyond this shared method, our two works diverge: Zhang et al. (2022) explore variants of DMALA that drop the Metropolis accept-reject step, or use diogonal preconditioning; in contrast, we explore a. new MALA-inspired auxiliary variable framework that admits a generic non-diagonal type of preconditioning (PAVG). In addition, our work is the first to identify the intractability of general preconditioning matrices within the DMALA/N /NCG framework, as explained in Section 3.1.1.
A range of prior works have 'mapped' discrete spaces to contimuou ones, thereby enabling MALA /HMC in the new continuous space. These methods are generally specialised to particular forms of the target distribution (Zhang et al., 2012), or to certain data types, e.g. binary data (Pakman & Paninski, 2013) or trees (Dinh et al., 2017). More problematically, the embedded distributions are typically piece-wise continuous (Nishimura et al., 2017) and highly multi-modal, which may explain their limited empirical success (Zanella, 2020; Grathwohl et al., 2021). One int interesting recent advance that may address the aforementioned problems is given by Jaini et al. (2021) who use SurVAE flows to transform a complex embedded distribution into an approximately Gaussian one. The drawback of such an approach is the requirement of training a potentially complex normalising flow model prior to running an MCMC sampler.
5 Experiments
We evaluate the newly proposed methods NCG, AVG & PAVG on four problem types: 1) sampling from highly correlated ordinal mixture distributions 2) a sparse Bayesian variable selection problem 3) estimation of Ising models and 4) sampling deep energy- basod model parameterised by a convolutional neural network.
Our key baselines are Gibbs-with-Graddeee & Geman, 1984). For our ordinal experiments, we also compare to Metropolis-Hastingg sampler with a uniform proposal over a local ball of radius r and a simple effective extension of GWG for ordinal data that we introduce: ordinal-GWG. In Appendices G & H, we provide details on these baselines and our methodology for tuning step-size parameters to maximise ||s - sellz between successive MCMC states, similar to e.g. Levy et al. (2018).
(GWG) Grathwohl et al. (2021) and standard Gibbs sampler (Geman

---

Published in Transactions on Machine Learning Research (10/2022)
Marginal estimation error
Covariance estimation error
ESS per second
0.14
Average (log-scale)
PAVS
Frobenius nom
Gibbs
CWG
0.10
CING-ondinl
306
Mit uniform
per
PAY
ESS
Gioes
0.04
DING
GNG-ordial
Mt. un-usfomm
0.00
NCG
AVG
NWVG
Gibba Method
Minutes
Minutes
Marginal estimation error
Covariance estimation error
ESS second
0.05
Awerage (log-scale)
107
Frobenius norm
pir second
NCG AVG
10.
NGG
FAVG
AVG
HWVG
0.02
30*
0.01
Miturifomm
MHunilom
NCG
AVG
Globs
GWS
GNG-ord Unl
Minutes
Minutes
Methad
Figure 20D mixture--f-polyoomia results. Top row: results for 2nd order polynomial in equation 19. Bottom row: results for 4th order polynomial in equation 20. Left: KL divergence between true and estimated marginals, averaged over all dimensions. Middle: Estimation error of the empirical covariance matrix (lower is better). Right: Effective sample size per second (higher is better). All error bars computed across 100 parallel chains.
5.1
20D ordinal mixture-of-polyyomiie
We define highly-corrclated ordinal target distributions logp(s) f(s) log Z over 20- dimensional lattices 220 c R220 where S contains 50 equally spaced points in the interval [-1.5,3.0] We construct these target distributions from mixtures of fully-factorised distribu- tions, enabling exact sampling and evaluation of the normaliser Z. f equals
50
20
f(s) log
Σ exp
Σ
94 (s:))
(18)
where k indexes a component of the mixture distribution, and 9k S + R is polynomial that we allow to take one of two forms:
(2nd order)
Ih (u) 1.5 21 61 go (u) = -t + t22 th
1 u + k/25
(19)
h4th order)
ty = (2u 1 3k/50).
(20)
Figure 3: 2D ordinal illustrations; equa- tion 19 & equation 20 correspoud top & bottom, respectively.
We visualise 2D versions of the resulting target distributions in Figure 3. The two dimensions are highly correlated with probability mass concentrating along the main diagonal. Similarly, in 20 dimensions, mass concentrates along the diagonal of a hypercube, meaning all dimensions are positively correlated.
We track the similarity between the empirical distribution of MCMC samples g(s) and the targel distribution p(s) using the following metrics i) marginal estimation error: the average KL-divergence between marginals (1/d) E Dku(q, p:) and fi) covariance estimation error: the difference, in Frobenius norm, between the empir- ical covariance matrices estimated under both distributions. We also track the Effective Sample Size (ESS) each sampler. Full details of these evaluation metrics are provided Appendix J.

---

Published in Transactions on Machine Learning Research (10/2022)
5.1.1 Results
Figure 2 shows the results, with the top & bottom rows corresponding to the 2nd & 4th order polynomials, respec- tively. In the 2nd order case, the evaluation metrics imply the following ranking:
PAVG NCG AVG Gibbs MHH-uiform > GWG-ordinal GWG
(21)
This ranking shows that all newly proposed samplers are either competitive or superior to baselines, with the preconditioned sampler, PAVG, showing especially strong performance. It is slightly surprising that GWG-based methods would perform worse than a standard Gibbs sampler, however this partly due higher wall-clock costs; GWG-ordinal actually matches Gibbs per-iteration; see Appendix I.
For the 4th order polynomial (bottom row), the estimation error metrics (columns 1 & 2) imply similar ranking as before. The main difference is that PAVG and NOG are now matched in performance. This relative reduction in the performance of PAVG is not entirely surprising: the derivation PAVG relied on a kind of global quadratic approximation of (s) (equation 14). Intuitively, it makes sense that our approximation degraded with the addition of 3rd and 4th order terms in the polynomial of equation 20. The ESS metric in the third column provides less specialised indicator of performance that doesn't use knowledge of the ground truth target distribution. Nevertheless,
that it provides similar ranking of methods in the top row, and ranks all methods more- or-less equal in the bottom row.
5.2 100D Sparse Bayesian linear regression
A common use-case of MCMC is sampling posterior distributions that arise in Bayesian analysis. We consider Bayesian treatment of sparse variable solection in linear regression models, using an experimental setup inspired by that of Titsias & Yau (2017). Given an n x d design matrix X, the response E R" is modelled as
y X(s av
N/ I.),
(22)
where w E IRd is vector of weights and s E Sd is binary random vector that masks out certain covariates. We place a conjugate marginalised out. Combined with a sparsity-promotiy posterior over binary masks p(s] X, y)-see see Appendix K for a complete description. This posterior tells us which covariates are 'relevant" for predicting the response y, and which are 'irrelevant'. Thus, it can be viewed a kind of sparse variable selection procedure.
prior over weights and noise-variance (w.o²), which can then be analytically
prior over s, we obtain an unnormalised expression for the
This posterior is 20D distribution over binary vectors, implying ~ Im possible states, which we can sum over (enabling normalisation). To make the sampling problem more challenging, whilst keeping normalisation tractable, we append additional 'irrelevant' dimensions by multiplying the posterior with 11 B(00001), where B is the Bernoulli distribution. The fact that we can normalise (and sample) the resulting 100D target distri- bution, p(s), enables us to compare it to the empirical distribution of MCMC samples g(8) using the following
Marginal estimation error
estimation error
ESS per second
NGS
N0S
A0G RANG
AMG
2,0
absolure error (log)
S1005
ESS second
BIG
DWG
1.5
20-3
(Tog) absolute
1.0
0.5
2
10
Minutes
Minutes
Method
Figure 4: 100D Sparse Bayesian linear regression results. The proposed samplers (NCG, AVG PAVG) converge faster as measured by bivariate marginals (middle) and are competitive better as measured by univariate marginals (left) and Effective Sample Size (ESS) (right). Smaller errors and larger ESS indicate better performance.

---

Published in Transactions on Machine Learning Research (10/2022)
Table Estimation error of an Ising Lattice matrix learned with persistent contrastive divergence, with standard deviations across 5 runs in parentheses. Each run uses a different seed for all sources of randomness.
K-1*
K-5
K=10
K=15
K=20
NCG
0.837 (t.05)
0.117
(#.0003)
0.117
(1.0004)
0.117 (#.0005)
0.117
(f.0005)
AVG
5.535 (#.002)
5.452 (=.008)
5.273
(=.007)
4.814 (=.029)
0.120 (#.0006)
PAVG model-specific"
0.120 (+.003)
0.118 (+.0003)
0.118
(#.0003)
0.118 (+.0004)
0.118 (#.0005)
PAVG model-agnostic
5.521 (#.002)
0.124 (=0.0005)
0.121
(=0.0004)
0.119 (=0.0004)
0.119
(+:0.0007)
GWG
5.271 (#t0038)
0.163 (#.0009)
0.138 (=.0004)
0.132 (=00))
0.128
(#-00088)
Gibbs
4.825 (#.002)
0.805 (=.003)
0.167 (=0000))
0.136 (003))
0.132 (=.0005)
K refers to the number of MCMC steps used by AVG; suitable multipliers ensure each method has the same budget.
This corresponds precisely to the Block-Gibbs sampler introduced by Martens & Sutskever (2010).
two metrics i) marginal estimation error: (I/d) Σx (s: = 1) - pi(s, = 1) and ii) pairwise estimation error: Exg Emte(o.) (s, k, S, =() Pt.j(s: k.s 1)). These metrics are described more fully in Appendix J.
5.2.1 Results
Figure 4 shows the results. Under all evaluation metrics, the methods rank as follows: PAVG 2 NCG 2 AVG Z GWG 2 Gibbs, with the inequalities being strict. for the pairwise-error metrie. This ranking matches our findings from the previous experiment, except for the swapping of GWG and Gibbs. Most interestingly, we continue to see clear benefit from preconditioning, even though the target log- log-probability is far from quadratic.
5.3 Estimation of Ising models via persistent-contraasine divergence
Pairwise undirocted graphical models are an important class of distributions used in physics, protcomics and eco- nomics (MacKay, 2003; Lapedes el al., 1999; Sornelle, 2014). Here, we consider the Ising model
1 logp(s) b + 5 Js log Z,
S f-1,1)
(23)
where J binary matrix multiplied by constant and Z is the (generally intractable) normaliser.
Following Grathwohl et al. (2021), define 100-dimensional Ising model with ground-truth parameters ba 0 and J is set to cyclic lattice as depicted in Appendix M. We then obtain 'ground-truth' samples from this model by running a Gibbs sampler for one million iterations. These samples are then used as dataset from which we re- estimate the Ising model using Persistent Contrastive Divergence (PCD) (Neal, 1992; Younes, 1999; Tieleman, 2008; Du Mordatch, 2019), which is an approximation to gradient-based maximum likelihood learning that requires an MCMC sampler; see details in Appendix and pscudocode in Algorithm L.6.
PCD has two key free-parameters: the MCMC sampler itself and the number of sampling steps K per parameter update. Better samplers enable lower values of K to obtain desired level of estimation error. We assess the performance of the samplers in terms of the estimation error of the estinated matrix J, as measured by the Frobenius norm || - [[ The model's bias b is fixed at the ground-truth value.
We compare two types of PAVG: model-agnostic and model-specific. The former is the approach we used in previous experiments as discussed in Section 3.3.1. The model-specifie version uses J as the preconditioning matrix. As described the end of 3.3, this renders PAVG identical to the -og-quadratic Block-Gibbs sampler of Martens & Sutskever (2010).
5.3.1 Results
Table shows the results. The model-specifie PAVG method performs best, achieving low estimation error even when K- ==1. 1. This not entirely surprising given the correspondence the method by Martens Sutskever (2010)
10

---

Published in Transactions on Machine Learning Research (10/2022)
Table 2: Estimation error of quadratic term when sampling from a particular kind of deep energy-based model, with standard deviations across runs in parentheses. Each run uses a different seed for all sources of randomness.
K-5"
K=10
K=15
K=20
NCG
0.128 (1.001)
0.116
(f.0002)
0.112 (#.0009)
0.112
(#.0004)
AVG
3.310 (#.078)
2.599 (+.051)
0.496 (#.119)
0.114
(=.000))
PAVG
3.212 (#.079)
0.117 (+.0004)
0.114 (# (#.0000)
0.112 (=.0007)
GWG
0.732 (4.013)
0.156 (4.003)
0.120 (+.001)
0.114 (#.0005)
Gibbs
3.121 (#.146)
2.270 (+.067)
1.637 (#.007)
1.204 (#.008)
= number of MCMC steps for AVG; suitable multipliers ensure each method has the same budget.
that was developed for log-quadratic target distributions. However, the superiority of this model-specific sampler over GWG is a new finding; indeed Grathwohl et al. (2021) perform multiple experiments with log-quadratic target distributions but do not compare to Martens Sutskever (2010).
Amongst the non model-specific methods, the results imply the following ranking: NCG PAVG (model-agnostic) > GWG > Gibbs AVG. In contrast to previous experiments, we now see a modest but clear advantage for NCG over PAVG, and that AVG underperforms compared to all other samplers unless K is sufficiently large. It's important to note that, throughout learning, the e step-size parameter used by NCG and (P)AVG is held fixed. This means that the same step-size must be effective across range of target distributions (since the target changes every time we update the model's parameters). The results show that NCG and PAVG work robustly with a fixed step-size, whilst AVG is less robust. We suspect that all three methods would benefit from adaptively-tuning the step-size throughout learning, but leave this possibility to future work.
5.4 Sampling deep convolutional energy- based models
Deep energy-based models (EBMs) take the form log ((s) - log Z, where f is a deep neural network. Such models have attracted significant attention for continuous data (Du & Mordatch, 2019; Arbel al., 2020; Qin al., 2022) where Langevin-based samplers the default approach. Less progress has been made in the discrete setting, with GWG (Grathwohl et al., 2021) being the first paper to showcase the potential of deep EBMs here.
A major challenge in comparing the efficacy of different samplers for deep EBMs is the lack of an evaluation metric. Unlike the Ising model, parameter estimation error is not meaningful since the models are not identifiable. We propose a novel strategy for dealing with this problem consisting of the following steps:
i) Given a real dataset, fit a ground-trur and J is symmetric. This fitting can be done via PCD with any reference MCMC sampler and a large number of sampling steps K (we use GWG and K 50). Denote the final learned model by 1 J's + (s).
energy-based model of the form 1 Js f(s) where f is neural network
ii) Sample dataset. D from the learned model by running the reference sampler for many (e.g. 50k) iterations.
iii) Fit model of the form 1 Js + (s) to the dataset D, where / is fixed, and only the symmetric matrix J is estimated. This fitting is done with PCD, using the MCMC sampler we wish to assoss. This model is identifiable, and hence the estimation error, || - F is a valid performance metric.
We apply this methodology to the USPS 256-dimensional image dalaset of binarised handwritten digils (Hull, 1994) and parameterise fe as a 7-layer convolutional network as detailed in Appendix O.
5.4.1 Results
Table 2 shows the results. Roughly speaking, we can rank the methods as NCG > PAVG 2 GWG > AVG > Gibbs. This ranking is similar that obtained in the previous experiment, and demonstrates that the newly proposed samplees are still effective for rather high-dimensional and non-linear target distributions. We note that the estimation errors in Table correlate well with sample- based metrics like MMD and visual sample quality; see Figures 13 and 15 the appendix.
11

---

Published in Transactions on Machine Learning Research (10/2022)
6 Discussion
We have presented multiple discrete gradient-based MCMC samplers that show strong performance across range of problem types in Bayesian inference and energy-based modelling. In particular, obtained the NCG sampler by viewing the Metropolis-adjuslee Langevin Algorithm through the lens of locally-informed proposals (Zanella, 2020) and the PAVG sampler through the lens of gradient- based auxiliary samplers Titsias & Papaspiliopoulos (2018).
Depending on the task, we saw that either NCG or PAVG show the strongest performance, and both generally outperform withb--wth--rrdientt (GWG) by a clear margin. This demonstrates the value of using proposal dis- tributions that update multiple dimensions at once, and do so in a correlated way. However, it is important note that these advantages do not come for free: GWG requires no tuning, whereas both NCG and PAVG have step-size parameters. In practice, we would recommend rumning GWG alongside the new methods when facing new sampling problem, as it provides a strong and reliable baseline.
To our knowledge, we are the first to adapt the auxiliary variable framework of Titsias & Papaspiliopoulos (2018) to discrete state-spaces. We believe there is significant scope to build on this idea, by investigating different choices of continuous conditional distributions, and understanding how /oo/wwen marginalised auxiliary proposals can be leveraged. In the context of PAVG, a natural question is what happens if one replaces the global preconditioning matrix with a state-specific matrix such the Hessian, in similar vein to Manifold MALA (Girolami & Calderhead, 2011).
References
Christophe Andrieu and Johannes Thoms. A tutorial adaptive MCMC. Statistics and computing, 18(4):343 373, 2008.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. ar arXiv:2003.05000 2020.
arXiv preprint
Vu Dinh, Arman Bilge, Cheng Zhang, and Frederick A Matsen IV. Probabilistic path Hamiltonian Monte Carlo. In International Conference on Machine Learning, pp. 1009 1018. PMLR, 2017.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
Simon Duane, Authony D Kennedy, Brian J Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics letters B, 195(2):216-22) 1987.
Raaz Dwivedi, Yuansi Chen, Martin J Wainwright, and Bin Yu. Log-concave sampling: Metropolis- Hastingg algorithms are fast! In Conference on learning theory, pp. 793 797. PMLR, 2018
Stuart Geman and Donald Geman. Stochastic relaxation, Gibbs distributions, and the Bayesian restoration of images. IEEE Transactions on pattern analysis and machine intelligence, (6):721 741, 1984.
Mark Girolami and Ben Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo methods. Journal of the Royal Statistical Society: Series (Statistical Methodology), (3(2):123-214, 2011.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops took a gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning, pp. 3831 3841.
PMLR, 2021.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-777 -73,, 2012. URL http://jmlr. .aaier.rg//peee org/papers/ v13/ h/s/greetton2a.hhre
W. K. Hastings. Monte Carlo sampling methods using Markov chains and their applications. Biometrika, 57(1):
97 109, 1970. ISSN 00063444. URL http://www jstor. org/stable//334900
John Hertz, Anders Krogh, Richard G Palmer, and Heinz Horner. Introduction to the theory of neural computation. Physics Today, 44(12):70, 1991.
12

---

Published in Transactions on Machine Learning Research (10/2022)
John Hubbard. Calculation of partition functions. Physical Review Letters, 3(2):77, 1959.
Jonathan J. Hull. A database for handwritten text recognition research. IEEE Transactions on pattern analysis and machine intelligence, 16(5):550 554, 1994.
Priyank Jaini, Didrik Niclsen, and Max Welling. Sampling in combinatorial spaces with survac flow augmented meme. In International Conference on Arlificial Intelligence and Slatislics, pp. 3349 3357. PMLR, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint orXiv: arXiv:1412.6980, 2014.
Alan S Lapedes, Bertrand G Giraud, LonChang Liu, and Gary Stormo. Correlated iutations in models of protein sequences: phylogenetic and structural effects. Lecture Monogrppog Series, pp.
1999.
Daniel Levy, Malt D HofIman, and Jascha Sohl-Dickstein. Generalizing Hamiltonian Monte Carlo with neural networks. In International Conference on Learning Representations, 2018.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient MCMC. Advances in neural information processing systems, 28, 2015.
David JC MacKay. Information theory, inference and learning algorithmss. Cambridge university press, 2003.
James Marlens and Πya Sutskever. Parallelizable sampling of Markov random fields. Proceedings of the Thirleenth International Conference on Artificial Intelligence and Statistics, pp. 517-524. JMLR Workshop and Conference Proceedings, 2010.
Nicholas Mctropolis, Arianna W Rosenbluth, Marshall Rosenbluth, Augusta Teller, and Edward Teller. Equa- tion of state calculations by fast computing machines. The journal of chemical physics, 21(6):1087 1092, 1953.
Radford M Neal. Connectionist learning of belief networks. Artificial intelligence, 56(1):71 113, 1992.
Radford M Neal et al. MCMC using Hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Akihiko Nishimura, David Dunson, and Jianfeng Lu, Discontimuous Hamiltonian Monte Carlo for sampling discrete parameters. arXiv preprint Xix:1705.08514 2, 2017.
Ari Pakman and Liam Paninski. Auxiliary-variable exact Hamiltonian Monte Carlo samplers for binary distribu- tions. Advances in neural information processing systems, 26, 2013.
Lianhui Qin, Sean Welleck, Daniel Khashabi, and Yejin Choi. COLD decoding: Energy-based constrained text generation with Langevin dynamics. arXiv preprint ar arXi::220.. 11705, 2022.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint ar XXK::1710.05941 2017.
Christian P Robert, George Casella, and George Casolla. Monte Carlo stafistical methods, volume Springer, 1999.
Gareth O Roberts and Jeffrey S Rosenthal. Optimal scaling of discrete approximations to Langevin diffusions.
Journal of the Royal Statislical Sociely: Series B (Stalistical Methodology), 60(1):255 268, 1998.
Gareth O Roberts and Osnat Stramer. Langevin diffusions and Eatropolis-Hastiine algorithms. Methodology and computing in applied probability, 4(4):337 337, 2002.
Gareth O Roberts and Richard L Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, pp. 341-363 1996.
Jeffrey S Rosenthal et al. Optimal proposal distributions and adaptive MCMC. Handbook of Markov Chain Monte Carlo, 1(10,1201), 2011.
Didier Sornette. Physics and financial economics (1776-2014) 2014): puzzles, Ising and agent-based models. Reports on progress in physics, 0666000 2014.
13

---

Published in Transactions on Machine Learning Research (10/2022)
Tijmen Tieleman. Training restricted Boltzmann machines using approximations to the likelihood gradient. Proceedings of the 25th international conference on Mochine learning, pp. 1064-1071, 2008.
Michalis K Titsias and Omiros Papaspiliopoulos. Auxiliary gradient-based based sampling algorithms. Journal of the Royal Stalistical Sociely: Series B (Stalistical Methodologu), 80(())7449 767, 2018.
Michalis K Titsias and Christopher Yau. The Hamming ball sampler. Journal of the American Statisticol Associa- tion, (2(6520):199810000 2017.
Laurent Younes. On the convergence of Markovian stochastic algorithms with rapidly decreasing ergodicity rates. Stochastics: An International Journal of Probability and Stochostic Processes, 65(3-4):177 228, 1999.
Giacomo Zanella. Informed proposals for local MCMC in discrete spaces. Journal of the American Statistical Association, 115(530):852 865, 2020.
Arnold Zellner. On assessing prior distributions and Bayesian regression analysis with g-prior distributions. Bayesion inference and decision techniques, 1986.
Ruqi Zhang, Xingchao Liu, and Qiang Liu. A Langevin-like sampler for discrete distributions. In International Conference Machine Learning, pp. 26375 26396.. PMLR, 2022.
Yichuan Zhang, Zoubin Ghahramani, Amos J Storkey, and Charles Sutton. Continuous relaxations for discrete Hamiltonian Monte Carlo. Advances in Neural Information Processing Systems, 25, 2012.
14

---

Published in Transactions on Machine Learning Research (10/2022)
A One-hot categorical random variables
A categorical random variable with possible values can be represented as one-hot vector belonging the
h S = fs [0,1),
Σ 8 = =1 10.
where |S| = k
(24)
1=1
A d-length vector of k-valued categorical random variables thus belongs the 8 CR
The methods introduced in this paper (NCG, AVG and PAVG) all use proposal distributions that, before nor- malisation, have the form exp(a's + (s o s) for some choice of and b. Such formulas remain valid in the categorical case. However, the normalised versions of these distributions have slightly different form. Specifically, to normalise this distribution, we note that there are d-independent groups of dimensions, and that within each group, there are only possible scttings, yiclding
exp(x)
II aa + bj. (s, O8y))
o(x) o(x)= Es exp(x))
(25)
where Sj == 11j-kkks aj == 1ojーl)ust and b, = b 1)bilh
B Sampling algorithms
Algorithm B.1 NCG step
Require: Step-size e. Unnormalised log prob function f(-). Triple (st f(t),, f(())). Sample S1+1 ~ (s(()) (t) as in Eq. (binary/ordinal) or the categorical equivalent implied by Eq. 25
Compute f(4++) & f((+1)
Accept 8441 with probability min (x) f(s)) (e s1+1)
(26)
qe (S 8s))
Algorithm B.2 AVG step
Require: Step-size c. Unnormalised log prob function f(.). Triple (st, ((t), V ((:)).
Sample 2, N(z; 22/e s:)
Sample $411 ~ (s z;, 8t) as in Eq. 13 (binary/ordinal) or the categorical equivalent implied by Eq. 25
Compute f($++1) & f((+1)
Accept 8++1 with probability
min (1, ex,) f(s))) N(z: /2/c S1+1, 1) (. Ze, 81+1) N(z; /2/c St. I) ge(se+1 Ze St)
(27)
Algorithm B.3 PAVG step
Require: Step-size e. Preconditioner M. Unnormalised log prob function f()). Triple (s,, f(s)), (s()))
Sample 2t ~ N(z M2 St, I) for Me defined in Eq. 39.
Sample $t+1 (s 24, s.) in Eq. 40 (binary/ordinal) or the categorical equivalent implied by Eq. 25
Compute f($++1 & f((++)
Accept 8 with probability
N(z; / S I) g (s 2 (1,8(+1)) N(z; !
min (1, (1,cxp(f(5441) f(s:))
st, I) qo (St+1 Z1, 8:)
(28)
15

---

Published in Transactions on Machine Learning Research (10/2022)
Algorithms B.1, B.2 and B.3 show how to perform a single step of NCG, AVG and PAVG, respectively. For simplicity of presentation, we write these algorithms for single MCMC chain. However, is straightforward use vectorised implementation for batch of MCMC chains, which is what we do in practice since it enables efficient GPU acceleration. This vectorised PyTorch code will be accessible upon publication.
C The computational benefit of NCG over GWG with large Hamming balls
We analyse the computational cost of GWG for Hammiing balls of radius r (henceforth GWG-r) and contrast this with the cost of NCG. For simplicity, we assume we are working in a binary space, S = {0, 1]" but similar arguments hold in the /caterrcatao cases. The proposal distribution for GWG-r is
exp V - s.) ) I(s (s:))
( sx)
(29)
((s))
where ((s)) == (e) exp Vf(; - s:) and (s:) contains all those points in St that differ from St in most r dimensions. A naive computation of (s:) involves enumcrating every point in the Hamming ball (s:);
the number of terms in the sum is Ei=o (i), which grows extremely quickly as a function of as illustrated in Table 3. It is for this reason that we claim in the main text that "straightforward approach' to using large values of is prohibitively expensive.
Table 3: Growth of size of Hamming ball as function of dimension and radius r
r r=1 1
r=
==3
r=
r' 5
r 6
50
51
1276
20876
251176
2369936
18260636
100
101
5051
166751
4087976
79375496
1271427896
d 200
201
20101
1333501
66018451
2601668491
85010294791
We believe it should be possible to compute (s:) more efficiently through suitable rearrangements and caching of computations in the manner of dynamic programming, however this is far from straightforward to imple- ment, especially when working with vectorised batches of MCMC chains. More fundamentally, such a method for computing Z(s,) does not immediately solve the issue of sampling the proposal in Equation 29, The naive way for sampling this proposal is to enumerate every point in the Hamming ball along with their probabilities exp (s)) (s-sssss) s:)) /Z(s). and then use the fact these probabilities define categorical distribution which can be sampled in the usual manner (i.e. inverse transform sampling). Enumerating all Ex (+) probabilities is prohibitively expensive for large values of r, and (to the best of our knowledge) it is open problem if if/how the dynamic programming approach mentioned above can be leveraged to reduce this sampling cost.
In contrast, NCG uses the fully-factorised proposal distribution that can be written as
d II xm1 (s(s) exp ((s)) - [ 2 s
(s
(30)
where Z1(8)) == == ( (s))) + t8.. 1 Each per-dimension normaliser Z;(ss) is just a sum over [S] terms (i.e. 2 terms for binary variables), and since there are only d such normalisers, the total number of terms that need computing is only d|S). Moreover, sampling the NCG proposal distribution is also easy since it fully factorised, and hence we can sample each dimension independently. For instance, when working with binary vectors, all need to do is sample d independent Bernoulli distributions. We note that the other proposed samplers in this work (AVG& & PAVG) also use fully-factorised proposal distributions, and hence the same arguments we just provided for NCG apply to them as well.
Finally, we note that GWG (with radius 1), NCG, AVG and PAVG all have very similar run-times in practice (sce Section T) since the predominant cost in all cases a single function and gradient evaluation (s(s)), Vs(s)) per iteration.
16

---

Published in Transactions on Machine Learning Research (10/2022)
D Interpretation of auxiliary variables
We aim here to shed some light the role of the auxiliary variables used in Sections 3.2 and 3.3. To keep things simple, we illustrate how AVG works for single binary variables with target distribution p(s) ex)(f(s))/Z In this simple scenario, Z can be calculated as exp(f(0)) + exp(f and hence p(s) is a known function. Recall that the first step in AVG is augmenting this binary variable with Gaussian auxiliary variables distributed according to p(= (8) = (22 / 8, 1). Figure 5 plots these conditional Gaussian distributions for two different values of e.
8= 0.02
== 0.2
12-
v
10
10
Auxiliary state z
8-
8
6-
6-
4-
2
2
0
-2
Discrete state S
Discrete state S
Figure 5: Conditional Gaussian distributions used in AVG for two different values of E.
The key feature of these plots that e controls the amount of separation between the two Gaussian distributions (for s = 0 and 8 = 1). Greater separation means that a block-wise MCMC sampler less likely to move between the two binary states. Why is this the case? An informal argument goes as follows: imagine we are using = 0.02 (left plot)) and shown at the bottom of the y-axis. Suppose we sample
current state is 0. Conditional on O, we sample 2, Lee. we sample from the Gaussian 2. Now, we would like to sample
conditional 2, i.e. would like to sample value of that could have plausibly generated = 2. It is very unlikely that = could have generated = 2 since that would be an 8-sigma event (the Gaussian at the top of y-axis has mean 10 and standard deviation 1). With s = excluded, we therefore end up sampling = 0, and remain stuck there.
The above analysis is imprecise in a couple of ways. We were intentionally vague on how to sample 8 given z. As explained in the main paper, AVG uses an MH accept-reject step with particular proposal distribution. Here, we will simplify things and use the exact conditional p(s) 2). which is given by Bayes rule
p(s) N(: 22
p(s (0) N(= 0,1) + P(1) N(;; V2.11
(31)
With this distribution at hand, we can be more precise about what happens when we sample from 2 2 2). Probability mass will concentrate on s = 0 (as we claimed before) unless the 'prior probability' p(1) is extremely close to 1. Thus, for 'most' target distributions (i.c. those that are not nearly point-masscs), our block-wise MH algorithn will be "sticky', remaining at s = for an extraordinarily long period of time.
The analysis changes however when we use the larger step-size of 0.2 (right subplot). Now, if we repeat our argument, and consider starting at 0 and sampling
2, there is much greater chance that the next state in the chain is s = 1. This is because z = 2 is plausible value under the Gaussian distribution corresponding to = 1 (in fact, more plausible than the Gaussian for s = 0). Our block-wise MH algorithm will be less sticky than before, and we will see faster convergence to the target distribution.
If larger values of are advantageous, why not use an extremely large value? This actually is the correct thing do if we can perform exact block-Gii Gibbs sampling and we don't require an MH step. This because as -> oc,
17

---

Published in Transactions on Machine Learning Research (10/2022)
8 and z become independent (all values of map to the same Gaussian distribution (o 1)) and hence we have p(s [] = p(s). Thus, our exact block-sampler alternates sampling irrelevant Gaussian noise and exact samples from p(s)- However, this extreme scenario illustrates precisely why exact block-Gibbs sampling is generally infeasible since special case of it (c- oo) assumes we evaluate and sample p(s), which is precisely the problem we are trying to solve in the first place! In practice, we cannol evaluate or sample p(s [)), and we need to approximate it using proposal distribution. Our choice of proposal distribution for AVG (based on a local first-order Taylor approximation) is often reasonable for sufficiently small values of c, but may be poor for larger values, resulting in unacceptably high rejection rales in the MH accept-reject step.
E
Preconditioned MALA as an auxiliary variable scheme
Here, we explain how to frame PMALA (equation 4) as an auxiliary variable sampler, and why its 'discrete analogue' is intractable. Our derivation here has the same structure as that of AVG in Section 3.2, except we now use a different form for the conditional Gaussian auxiliary variables.
For a continuous state s E Rd, consider the unnormalised target density =ts, 2) exp(I(s))N M-1/3s, 1). In theory, this distribution could be sampled in a block-Gibbs fashion via alternate sampling of z, ~ (2z M-1/2,, I), and 1 (s Ze) x (s, ze). However, for general functions f this second sampling step is intractable, so it is replaced with an MH rcjept-reject step using the proposal distribution
(s 22, 8:) x exp(J(s:) VS(ss) Sx)).N(n;, M-,0;M-11 M-1/2s,
(32)
x exp -
2 - M- 11)
where p== M2/22, + MV)(s;)
(33)
N(ss M2/22 + MV f(s:,, M)
(34)
where equation 32 approximates x(, via a Taylor expansion of ((s). If we now marginalise out the latents, we obtain
f N M ts DW(s; 112 + MV M)dz, N(s: st + MV)(ss)
(35)
which is equal to the PMALA proposal in equation 4 after redefining M as M.
If we now restrict e 8d c Rd to be a discrete random variable in equation 33, then we no longer obtain Gaussian proposal distribution as in equation 34, but rather discrete pairwise Markov random field that is generally intractable to normalise and sample.
F
Choice of preconditioning matrix for PAVG
We use matrices of the form 7,, where is an adaptively learned scaling paramcter. In general scttings, when we have no domain-kknwwledee to help us select M, we choose it to be an empirical covariance precision matrix (see below for how to choose) computed from a set of initial samples obtained during the burn-in phase. The algorithm for this general case is presented as algorithm F.4.
avarrance/precision
However, when sampling from energy-basee models (EBMs), we usually have access to a real-world dataset that the EBM is modelling. In this case, we can compute the empirical covariance/ precision of this dataset, and use that our M. This simplifies algorithm F.4: we drop lines 5-8, and no longer need a vero-initialisation of M in line 2.
covariance/preccsion
F.1 Covariance or precision? Automatically picking M from list of options.
If we were sampling continuous variables drawn from Gaussian distribution, then the correct choice of M, based the approxination we made in equation 14, would be the (negative) precision matrix. This choice also works best in our ordinal experiments in Section 5.1. However, in all of our binary experiments, the covariance matrix, not precision, works better (often substantially so).
The choice between covariance and precision (or, more generally, a list of candidate matrices) is a hyperparameter. To avoid manual tuning, we propose simple hcuristic to automatically select the 'best' matrix from list of
18

---

Published in Transactions on Machine Learning Research (10/2022)
options. We define 'best' as the matrix that (after rescaling) minimises the approximation error in equation 14 across pairs of consecutive samples accumulated during burn-in phase. Specifically, do the following:
1. During the burn-in period, collect (st, ((s,,, (s())) into a 'dataset' DD
2. Assemble list of candidates [M1, M2... Mm] In experiments, consider only two candidates: the empirical covariance & precision matrix of collected samples {st).
For each M our list, solve the least-squares linear regression problem for the scalar To
arg min Σ lIve 702re1
(36)
where y = f(:)) (s++1 st), 24 := (1/2(s[s---11- M(s41 st)
(37)
Return the re-scaled matrix M which obtained the lowest least-squared loss.
Across many of our experiments, the returned matrix M performed very well without any additional alterations. However, to maximise the performance of PAVG, we found it beneficial to only use the scaling factor To as an 'iitialisationn and continue lo update it adaptively using AdaptGamma as shown Algorithms F.4 and F.5.
F.2 Preconditioners with negative eigenvalues
When deriving PAVG in section 3.3, we defined a conditional Gaussian distribution of the form
N(z; ,,,
M, = M (2/c
(38)
We note that there are multiple types of matrix square-root, and any of them is valid here. However, such square- roots only exist M, is semi-positive definite (i.c. all eigenvalues are non-negative). This only guaranteed to be the case (for any value <) if M also semi-positive. This is because the eigenvalues of M, are of the form + (2/c), where x; is an eigenvalue of M.
We can specify alternative definition of M, that ensures semi-positive definiteness
M, M [max((, -mmin) +2/cc]
(39)
de
where Amin S for all i. This definition of M, means that PAVG uses different conditional Gaussian, and the MH-proposal distribution in equation 17 now becomes
(s Zt, St)
II σ ([((s)) (Mst), + (M)
de (22)) 2 5/
where o(x) = Exes exp(x))
cxp(x)
(40)
What is the consoquence of this new definition M.? Previously, we diagonally perturbed M by some value (0, o0), where larger step-sizes e meant smaller perturbations. Now, we diagonally perturb M by some value in (d=o, oo), where do max(i0,Amin) 2 0. Thus compared to the old regime, the new regime enforces kind of maximum step-size (minimum perturbation).
G Baseline ordinal samplers
Our ordinal experiments two additional baselines: Ordinal-GWG and MH-uniform. Ordinal-GWG is newly introduced in this paper as it is arguably an 'obvious' way to improve GWG when dealing with ordinal data, and incurs very little additional implementation or computational complexity compared to standard GWG. We provide visual illustrations of the proposal distributions of these Metropolis- -astiigg samplers Figure 6. As can be seen from the figure, Ordinal-GWG, like GWG, can only update one dimension a time. However, can jump multiple states along any given dimension, with the maximum number of moves controlled by a radius parameter r. Thus,
19

---

Published in Transactions on Machine Learning Research (10/2022)
Algorithm F.4 Adaptive learning of preconditioning matrix. (Default values in brackets are used across all experiments)
Require: Integers Niters; Nehains (100), NM (1000) and Nadapt (100).
Require: Initial adaptation rate 5 (0.25) and decay factor p (0.99). 1: Initialise chains S [s), ..... Neetiii and history D = [[
Initialise M 0, ? Vold 1.0
3: for EE (I.... Niters do
4: Compute S141 from St using one step of PAVG (Algorithm B.3) with preconditioner yM
5: < NM then
6:
Append S411 to history D
7: else if t= NM then
8:
Define new M described in Section (F.1) using history D
9:
else if t mod Nadapt 0 then
10:
Y: 7 AdaptGamma(?; Yold: 6, Nadapt; D)
Update scaling factor
11:
8 p8
& exponential decay of step-size
12: end
13: end for
Algorithm F.5 Adapt y to maximise 'jump' distance ||s 8:-1lli
function ADAPTGAMMA(3, Told; 8, Nadapt. D)
Onew f Average value of ||s, - St-1 || computed over all chains in D(- N Nadapt
Oold + Average value of |ss s computed over all chains in D[- 2.n - -N Nadapt)
Increased +- 2 Told
Improved +- Onew 2 Gold
if (Increased and Improved) (not Increased and not Improved) then
8
Positive adjustment
else
-
Negative adjustment
end if
f- Told
if hy| 1 then
7 7 (1
Multiplicative adjustment
else
7 + 8
> Additive adjustment (allows 7 to change sign)
end if
return 7, 7old
end function
20

---

Published in Transactions on Machine Learning Research (10/2022)
GWG
Ordinal-GWG
MH Uniform
7.922
5.800
3.313
3
0.529
0.000
Figure 6: Illustration of the different proposal distributions used by various baseline MH samplers. The darker blue contours represent the target distribution, the white dots represent the proposal distribution around the current state (black X). Left: GWG Middle: Ordinal-GWG with radius 3 Right: Uniform MH with radius 2.
the support set of this proposal distribution is (at most) of size (2r - 1)d, where is the dimensionality. The Ordinal-GWG proposal has the same functional form as GWG equation 6, except that the summation is no longer over a Hamming ball of radius 1, but instead over the aforementioned support set.
MH-uniform (final column of Figure 6) also has radius paramcter r. This radius controls the size of the hypercube over which the proposal is uniform.
H Tuning step-size parameters
NCG, AVG & PAVG have step-size parameters Our grid-search based tuning procedure involves rumning each sampler for a short amount of time (1000 iterations, which takes S 1 minute in most of our experiments) with different step-sizes, and selecting the step-size that maximises the average L1-distance ||s++1-8:]| between successive states (averaged over all time-steps and parallel chains). For NCG, AVG & PAVG we first first identify the best orden-of-maanittdee by searching, in parallel, over the values in the set (0.05, 0.5, 5.0, 50.0, 500.0). We did not find it necessary to consider other orders of magnitudes. Indeed, our analysis in Section D indicates that (at least in the context of AVG and binary variables) step-sizes below 0.01 are unlikely work. After identifying the best order-of-mmanittute c.g. f0.1, . 0.9}. Thus, in total, scarch procedure tests 15 values of and, due to the parallelism, only lakes a couple of minutes for our experiments. Admittedly, for more complex experiments with expensive target distributions, the cost of this search procedure will increase.
we then search, in parallel, each decile within that particular order of magnitude
Ordinal-GWG and MH-uniform also have step-sizo-like parameters r and we tune these in a similar way as described above. We first search r E (1,5, 5, 10, 15, 20} and then search within the best interval e.g. r (15, 16, 17, 18, 19] (note: the maximum possible value of the radius parameter is 50 these experiments). The final selected step-size parameters are shown in Table 4.
It is natural to wonder how selection of in PAVG interacts with the choice of preconditioning matrix. As explained in Appendix F, PAVG's preconditioning matrix is automatically tuned during each run. Hence, when we perform the above grid-search, we are fixing a particular at the beginning of a run, and then automatically learning the preconditioning matrix given that fixed e.
H.1 Step-size sensitivity analysis
It is important to have a sense of how robust NCG, AVG and PAVG are to the choice of step-size. We investigate this in the context of the Ising model experiment in Section 5.3. We fix the number of sampling steps K 10, and consider how the errors reported in Table vary as we multiplicatively perturb the step-sizes used to generate those results. Doing so yields Figure The main finding here that all samplers are relatively stable over a wide range of step-sizes. For NCG and PAVG, we can multiply our original step-size by any number between 0.3 and
21

---

Published in Transactions on Machine Learning Research (10/2022)
Experiment
NCG
AVG
PAVG
Ordinal-GWG
MH-Uniform
Ordinal (poly2)
0.05
0.02
1000.0
16
2
Ordinal (poly4)
0.05
0.02
0.06
Bayesian regression
0.03
1000.0
1000.0
Ising lattice
0.5
0.2
0.2
ConvEBM
0.2
0.1
0.1
Table 4: Final step-sizes used across different experiments after tuning.
10.0, and the resulting is within 15% of the original error and remains lower than GWG's error. Finally, note that chosen step-sizes (i.e. relative step-size 1) are not optimal, which is to be expected since search procedure described above does not have access to our error metric (which is important, because the error metric relies on ground-trun truth information that is not typically available for real-world problems).
Effect of step-size on estimation error
NCG
AVG
PAVG model-agnood
NCG
AVG
PAVG model-agnostic
GWG
GWG
GWG
1.00
Relative error
0.10
1.0
Relative step-size
Relative step-size
Relative step-size
Figure 7: Relative error multiplicatively perturb the step-sizes used by NCG, AVG and PAVG in Ising model experiments of 5.3. Relative error is defined as the new error divided by the 'original' error reported in Table 1. Similarly, the dashed line for GWG is also relative to the errors reported in Table 1.
Wall-clock versus per-iteration results
In all the experiments presented in the main text, we allotted equal wall-clock time to each competing MCMC method. Whilst this the fairest form comparison, is implementation and system-dependent. To give a sense of how the different methods compare on a per-iteration basis, see Figures 8 and for 20D ordinal and 100D) Bayesian regression results, respectively. Methods with lower wall-clock costs per-iteration are run for more iterations in total. In particular, the Gibbs sampler is run for significantly more iterations than other samplers due its low cost in these experiments. However, note that Gibbs is not guaranteed to be the cheapest method universally; it can be very expensive for categorical distributions with large state-spaces (Grathwohl et al., 2021).
Figure 8, GWG & GWG-ordinal run for the shortest number of iterations. This is potentially surprising since one may expect GWG to cost the same amount. per-iteration as NCG. This occurs because GWG-based proposal distributions have irregularly shaped support sets when the boundary of the state-spaco- the centre panel
22

---

Published in Transactions on Machine Learning Research (10/2022)
Covarlance estimation error
Marginal estimation error
NCG
NCG
Average KL (log-scale)
PAVG
6-
PAVG
Gibbs
Gibbs
GWG
5
GWG
GWG-ordinal
GWG-ordinal
10a
MH-uniform
Frobenius
4
MHH-uiffrm
3
10-1
5000
10000
15000
20000
5.000
10000
15000
20000
Iterations
Iterations
Marginal estimation error
Covariance estimation error
16
NCG
NCG
AVG
Average KL (log-scale)
PAVG
Gibbs
Frobenius norm
Gibbs
GWG
GWG
GWG-ordinal
GWG-ordingg PVG-oordinal
MHH-uniform
MH-uniform
8
10a
2000
4000
6000 8000 10000 12000 14000 16000
2000 4000 6000 8000
10000 12000 14000 16000 Iterations
Iterations
Figure 20D mixturc--ff results. These figures are the per-iteration versions those Figure 2.
Marginal estimation error
Pairwise estimation error
NCG
NCG
PNVG
PAVG
absolute error (log)
Gibbs
(log) absolute error
Gibbs
10-2-
10-1
20000
40000
60000
a0000
100000
20000
40000 Iteration
60000
80000
100000
Iteration
Figure 9: 100D Sparse Bayesian linear regression results. This Figure is the per-iteration version of Figure 4.
of Figure 6 for an illustration. Normalising the proposal over such irregularly shaped support sets is difficult to implement efficiently in a vectorised manner for a batch of MCMC chains (since each chain can have a differently- shaped support set, and complex masking operations are seemingly required to handle this). Even in an idcalised scenario where this computational overhead zero and so GWG matches NCG in per-iteration cost our rankings reported in the main text would not be impacted significantly.
23

---

Published in Transactions on Machine Learning Research (10/2022)
J
Evaluation metrics for ordinal & Bayesian regression experiments
The experiments in sections 5.1 and 5.2 use the same methodology for evaluating performance. In both cases we:
Run 100 parallel chains for 10 minutes with a burn-in period of 1 minute. The resulting number of iterations per-method & experiment can be seen in Figures and 9.
After the burn-in period, we begin saving the history of cach chain.
Every minute after the burn-in, we use the chain histories to compute estimation errors (Table 5) for each chain separately. We then compute a mean and standard error across our 100 parallel chains, and use these to construct Figures 2 and 4.
At the end of the run, we estimate the Effective Sample Size (ESS) of each chain. Following Zanella (2020) and Grathwohl et al. (2021), we map every state our chains to test statistic, namely ||s Srandom where Srandom is randomly selected point in the state-space. The resulting test statistics can be stored in matrix S of shape num_iters x 100, and the ESS-per-chain is then estimated with tfp.mmeree mcmc wwfeeeiveeeseaeeeeeeeeeeee sample_sizede densorflox--probbllit Finally, to construct the ESS plots in Figures 2 and 4, we compute box-plot statistics (median quartiles) across the 100 ESS estimates.
filter_ _ertte_beeodd__oe peaeeddpssiiieeeeeeeeeeeeeeeeeee _ using version 0.14.1 1 of
Experiment
marginal error
covariance/pairwiss error
Ordinal
(1/d) DK |l p.)
||M MpllF
Bayesian regression
(1 (a:s; 1) - (sss 1)1 # Els ExIEO..) (a k, =) - Pis(s: k,s; 0)
Table 5: qi refers to the 1th univariate marginal distribution of the empirical samples accumulated across a single chain. Similarly, 4. refers to a bivariate marginal of such empirical samples. P, and Pt, refer to marginals of the target distribution, and are computed exactly (i.e. they are not empirical distributions). My and Mp are both empirical covariance matrices; the former is computed using samples accumulated across single chain, whilst the latter is computed using 100, 000 samples drawn from the target distribution.
K
Posterior distribution for sparse Bayesian linear regression
We define Bayesian regression model using a similar procedure to Titsias & Yau (2017). The regression takes the form
y X(s +o ov
Y.VRR S
(41)
The quantitics in this equation random variables described by the following generative process
Place a sparsity-promoting prior on s
D 5t
p(s) = 1 Σ 8 TDD
s + - 1=1
3x),
(42)
where I is the gamma function (not distribution) and (ar, Br) are hyperparameters with default values (,001,10.0), This prior can itself be viewed as the marginal of Bayesian model specified by Ih Bernoullis;; m) and Reta(T;;az,90
2. Let -~ No IN).
Define X, Xdiag(s) and note that ((s o w) Xyw.
24

---

Published in Transactions on Machine Learning Research (10/2022)
Place a conjugate normall-inverse-aammaa prior over weights and noise-variance (w, o2) Namely,
p(w, x) = NVw 0, go (x² + AID) IneGamma( a [a,30))
(43)
We note that this choice normal distribution is a kind of perturbed g g-prior (Zellner, 1986), with hyperpa- rameters (g. A) with default values (20, 0.001). The gammma hyperparameters (aa: 9o) have default values (0.1.0.1)
5. As implied by the regression formula in equation 41, our model of follows
p s, X,W, = Ny Xgw, INI
(44)
6. Putting the previous sleps together, we arrive at the joint distribution
p(s, o w X) p(s)p(w, o² |s, X)((() |s, ,
(45)
The variables (o², w) can be analytically integrated out, resulting in the posterior distribution
X X, + MDD
2 iN
p(s 1), x p(s)
(2 + yu -9"" Xx + g) Xa + Xm
(46)
+ g) XX AT
Finally, we use the following steps to construct 'observed data' (X,y) that we then plug into our posterior
Let o be ii.d random variables drawn from a uniform distribution over the discrete set {0, 1,,1,2).
Define the observed response Yobs Ei= Tz..
'Dpllcatee' the covariales times i.e. X, == Xj (mod 5)+1 for j E 76,7....,00
Repeat the above steps N = 20 times to obtain a 'design matrix' X E 20x20 and response E R This design matrix and response can be obtained under the regression model in equation 41 by setting w to a vector ones, and a = 0 (i.c. noiseless regime).
The duplication of covariates induces multi-modality in the posterior p(s Yobs), since masking out æ and leaving its copy unmasked is equivalent to masking 26 and leaving unmasked.
L Persistent contrastive divergence (PCD)
The gold-standard approach for the estimation of statistical models is maximum-likellhood estimation (MLE). Unfortunately, for unnormalised models
logp(s: 0) f(s: 0) log Z(0),
z(0) = Σ exp(f(s: 0))
(49)
the normaliser Z(0) is presumed intractable, and so we cannot actually compute log 0), which is required for MLE. However, can conveniently express the gradient as
Vo ps: 0) Vofo(s) Epts: o) [)[os))
(50)
and then use Monte-Carke estimate of the second term, where approximate samples are drawn from p(s; 0) using MCMC sampler.
Running an MCMC sampler afresh every time we wish to perform gradient-based parameter update quickly becomes prohibitively expensive. Persistent contrastive divergence provides a solution to this problem, by not starting afresh each time, but by persisting MCMC chains across parameter updates. One version of PCD given in Algorithm L.6. This implementation adopts vanilla stochastic gradient descent to update the model parameters, but alternative optimisers (c.g. Adam (Kingma Ba, 2014)) can be substituted.
25

---

Published in Transactions on Machine Learning Research (10/2022)
Algorithm L.6 Persistent contrastive divergence with buffer
Require: Discrete dataset D. Integers Niters: Natch, Nbufeer.
Require: Unnormalised log probability function f(:: Q). Step-size c. Optional regulariser h(0).
Require: MCMC transition operator T(1-.). Number of MCMC steps K.
B + "st,....s:Vuuiii
Initialise buffer of persistent chains
for E (1...., NitEr) do Sample minibatch B of size Naach from buffer B for jE E K,...,,K do
B~ T(- B)
Parallelised update to minibatch
end for Update buffer B with the new values in B Sample minibatch X of size Nbatch from dataset D
Di i.e. persist the chains
g
1
Vof(s: 0) Vof(s: 0)] Voh(O)
(47)
Nbath
SEX
SEB
t A + +g
(48)
end for
00
103
Figure 10: Left: ground truth lattice Ising model with connection strength 0.2. Right: samples from this model generated with 1, 1,000,00 000 steps of Gibbs sampling.
M Estimation of Ising models
The 100 100 lattice Ising matrix used our experiments was generated via. igraph, specifically, we call agraph.Graph. Lattice (dim=[10, 10], circular=True) which returns binary adjacency matrix, and then multiply this by 'connection strength' parameter 8 0.2. The resulting matrix, along with samples from the Ising model, is shown in Figure 10.
For learning, we use PCD as described in Algorithm L.6, replacing vanilla SGD with Adam. The dataset D consists of 10, 000 samples gemerated via 1, 000, 000 steps of Gibbs sampling. We set Niters 2, 000, Nbatch = 50, Nhuffer 5000, e = 0.0003. Our model takes the form f(s; J) = =T Js, where J is a symmetric matrix. We use the regulariser h(J) 0.01 Ev, [L,p]-
We use version 0.9.8 the igraph package.
26

---

Published in Transactions on Machine Learning Research (10/2022)
N Sampling Lattice Ising models with controllable higher order interactions
The motivation for PAVG was incorporate socond-order interactions into the sampling process. We expect this to be of clear benefit when the target distribution is log-quadratic, and that is indeed the case, as shown in the Ising model experiments of Section 5.3. It is interesting consider how the performance of PAVG changes we gradually add higher-order terms to such Ising model. As simple experiment, we consider sampling from low-dimensional Lattice Ising target distribution (d 16) defined by
f(s) = s' +a Σ
St8ySk:
(51)
(iJAA)ZZ
where I is an index set containing 50 randomly selected triples and the matrix J is as described in Appendix M. We use low-dimensional model so that, similar to our Bayesian regression experiment, we can analytically compute all pairwise marginals and use the estimation error of such marginals (as defined in the bottom right quadrant of table 5) for evaluation.
We compare PAVG (with preconditioning matrix J) with AVG for different settings of a. AVG is natural comparison method, as it is a special case of PAVG that does not have access to the preconditioning matrix. The results are shown Figure 11.
Effect of higher order terms on estimation error
PAVG
0.25
AVC
0.20
Estimation atror
0.15
0.10
coefficient third-order terms a
Figure 11:
Estimation error for different values of a when sampling from a target distribution given by equation 51.
When the coefficient is 0, we simply have an Ising model, and PAVG outperforms AVG as expected. For small a (in absolute value), PAVG and AVG perform similarly, whilst for larger values of a, PAVG again outperforms AVG. This U-shaped performance gap is somewhat surprising, and we do not claim to have a full explanation of why it occurs. However, we believe it might be understood in terms the relative speed of PAVG versus AVG during the two phases of MCMC sampling: the burn-in period and, afterwards, the exploration of the typical set.
0.5
PAG
8
Avorag ditince Jarpud
0.5
0.2
0.1-
245
30K
12x
24h
3.0
400
0.3
(b)
0.5
Figure 12: Average L1 distance moved (across 100 parallel chains) for two different values of a.
27

---

Published in Transactions on Machine Learning Research (10/2022)
Our conjecture is that, in the specific context of Ising models with non-negligible third-order terms, the precondi- tioning in PAVG isn't very useful for faster exploration of the typical set, but it is useful during the burn-in period. Some evidence for this claim can be found in Figure 12, where we plot the average 'jump distance' (l|s, 8 s(-1ll)) between successive states in the MCMC chains for 0.3 and 0.5. We see that, for both values of a, PAVG and AVG eventually converge to steady-state where the jump distance is low, which may indicate that the typical set has been reached and is now being explored. In both cases, PAVG is faster to reach this steady-state suggesting that it has a smaller burn-in period. However, PAVG is only marginally faster than AVG when a = 0.3, perhaps because the initial distribution of MCMC chains is already reasonably close to the target distribution, so neither method has difficulty. In contrast, PAVG is substantially faster when a = 0.5.
One corollary of our conjecture is that, for 0.5, the estimation error for AVG ought to eventually come down if we run it for longer. Indeed, running it for twice as long (120k iterations) does reduce the error to the much lower value of 0.08 (although this is still more than twice the error of PAVG). We believe that further analysis of these results is beyond the scope of this appendix, and leave it to future work. Regardless of whether our above conjecture holds, a key take-away message from Figure 11 is that preconditioning with second-order interactions can be highly beneficial in the presence of strong higher-order interactions. Indeed our other experiments on sparse Bayesian regression and deep energy-based models also validate this claim.
Sampling convolutional energy-based models
The architecture of the convolutional neural network used in our experiments is shown in Table 14a, alongside images from the USPS dataset and our ground- trrth EBM. To learn the ground-truth EBM model 3,5 J f*), we use PCD as shown in A Algorithm L.6 with GWG as our sampler and K 50. We set Niters 10, 000, Nbathh 50, Nbuffer = 5, 000, 0.0003. We use weight decay of 0.0001 on the neural net weights.
After learning this ground-trute truth EBM, we rerun PCD using the model 3s f*),,,,, where J is symmetric matrix of parameters. During this re-estimation phase, the settings of PCD remain the same, except for Niters = 2, 000.
0,0035
R*2 0.7151
R*2 0.6983
R*2 0.9530
0.0030-
5.025
0,0025
maiggnnl error
coveriance error
0,000
0.0020
0.0010
2.005
0.0035
0.0
0.5
1.2
155
200
2.5
3.0
0.0
0.5
1.0
1.5
2.0
255
3.0
0.0
0.5
1.0
1.5
2,0 25
3.0
Parameter error
Parameter error
Parameter error
Figure 13: Comparison of the parameter estimation error metric ||J- J"I (as reported in Table 2) with 3 different hample-based metrics. Each point in a figure corresponds to a cell in Table 2 i.e. it corresponds to a particular choice of MCMC method and value of K (thus, there are 20 points per subplot). These sample-based metrics all compare the similarity of two sets of samples. Here, we compare the 5K buffer samples produced by a particular MCMC method against 5k samples from the ground-truth EBM. Maximum mean discrepancy (MMD) (Gretton et al., 2012) is computed using the same code as Grathwohl et al. (2021). 'marginal error' is the absolute difference in empirical means (averaged over all dimensions). 'covariance error' is the Frobenius norm of the difference between the empirical covariance matrices computed using each set of samples.
28

---

Published in Transactions on Machine Learning Research (10/2022)
8
6 1
0
5
0
5 32 404
Conv2d(1, 16, 3, 1, 1) SiLU()
2
6
6
1
3
1
1 9 166 4 823
9 2
(16, 32, 4, 2, 1) SiLU()
1 5 12107163>0011
O 3479350048
275
Conv2d(32, 32, 3, 1) SiLU()
0 7 21]22069612900
3 7 7744333536104a)
C3nv2d(32, 64, 4, 2, 1) SiLU()
76 6 9 4380841800411
Conv2d(64, 64, 3, 1, 1) SiLU()
J (
6 9 / 3
Conv2d(64, 128, 4, 2, 1) SiLU()
A 1 1111033314117
(21403307771
0 1107001511300009
Conv2d(128, 128, 2, 1, 0)
SiLU()
R00309100--1133
77864583328.
(a) 7-layer convolutional neural net- work used an EBM model USPS digits. The SiLU() activation function (Ramachandran et al., 2017) also called the 'swish' activation.
5 210200813217? 133 3
(b) Top row: real images from the USPS dataset. Bottom samples from ground-truth quadratic-EBB model trained with GWG and K = 50.
29

---

Published in Transactions on Machine Learning Research (10/2022)
K=5
K=10
K=15
K=20
0.128
0.116
0.112
0.112
4
3
2
8
f
30
| & 7 126
2
4
0
6
5 529
3 7
7
(
1
i
6 05
1
1
1
NCG
1
5,
8
9
2
07660630
79
0
e
3
O
6
0
1
17 C 1-8811 1.7
7 1
770900
20.30
9 17300
130065
1 4100]
7 14309
906
7 7 1230600
7336000
0000
0
7 1.62
7h
7
4
709 8 1
O
0
9 6 543
3 56002010
6 G 1000TT
3 Q
B
7 1
1000
Q f se
2
6
O
9
3.310
2.599
0.496
0.114
R
7
33
5 4
73 13
3
A
I
8
72 S
3.5
7
1
9
A9
24
19
(
5 r
1
0
Q 22500
0
AVG
9
.2.
210 9
BTC99
3
4:
2
R
4
4
48
6
7 1R1020
Y 1 1110800
F7
: 4 7799
7 7 X x 7
3.
3 8478 9
& c 1
0F 104
1 0
3.212
0.117
0.114
0.112
F
4
Y
2
5 Y
3
9
0
3
0
J
K 1
4
3
7
1
0
1 0
6
PAVG
0
()
7 %
M.
1006
0
8
00
5
3
.)
5
6
9
3
3
1
0
RRS 3 2
79
9
3
4
Y 2 0
C
699 19104 ?
7
0
8
0
2
7 6
7 + 99
" < J
39
' D
L
7 22
O
0
b
N
6 1
3 3 4r
J
O
Y
L011
0.732
0.156
0.120
0.114
53 3
10
6 10
C 3
9 9
Y
2
6
2
4 1
8 9 T 070 0 " 1 4 G 0000000
3 ( 0 / 1 2
$ 4
3
1
1
3J
1
4
0
/
9
O056
9
7
7
C 30 J 6 Q
0
GWG
1 6
J253 0
393 00.4
4 1:099
0
+130091
A
X 6669
K
O
(
4 222
0
5 3 2
Q
GO 0 000001
060933>2
O6 0
Q
3
3 71001
D
5 7 7077
3
2
2
2
9
0
3.121
2.270
1.637
1.204
$
5
0
1
Q
$ 3
7 4
7
1
7
7
1
1
2
1
3
9 <
3
7 700
4
3
7
4
0
5
1
6
Gibbs
4. 2.
0 110
7
1
0
1
1
6
1
3
9
1
100
7
<2
$
000
6 5 1001
4 7
8
8 0
9 0
30
U 6
00
1 0 1000 0
5 7
0
3 Q
8
100 ?
1
J.
2
7
Figure 15: Samples for the methods described in Table 2. Above each image is the corresponding estimation error.
30


Reviews:
[{"review_id": "6ueaYD0QQ3", "reviewer": "Reviewer_QtjV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Novelty: The methods introduced in the paper are appear to be a straightforward \ncombination and/or extension of existing ideas. As TMLR prioritizes correctness over \nnovelty, this is not a significant issue in and of itself.\n \nExperiments: The comparison with other methods are thorough, although I would have liked\nto see the recent method of Zhang et al '22. Any reason why this was not included in the\nevaluation?\n\nMy main criticism of the paper in its current form is the lack of sensible defaults for \nthe hyperparameters and an analysis of sensitivity to hyperparameters for the proposed \nmethods and the baselines. I believe this is important to justify the authors' claims tha\nt\ntheir method performs better than prior methods on the problem.\nIn my experience using probabilistic programming tools, a fair amount of tuning is\nrequired to get MCMC methods working efficiently. The authors have specified the specific\nhyperparameter choices in Table 3, but have not commented on how much effort was spent to\ntune each method for each problem. Ideally, the authors would show results for all method\ns\non several hyperparameter choices. \n\n\nOther questions: \n\nEquation 8: It appears sampling from q requires a summation over all points in the\ndiscrete space. This can be computationally prohibitive in high dimensional settings.\nIs this avoidable in an implementation? If not, how come your method does well in \npractice? ", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents three methods for sampling from intractable distributions in discrete\nspaces. The problem is well-motivated: such situations are ubiquitous and sampling from \ndiscrete distributions is significantly more difficult than sampling from Euclidean \nspaces. The authors develop principled MCMC-based methods adopted from similar methods \ndeveloped in prior work in different settings. \n\nWhile I am not an expert on this area, I am cautiously optimistic about the paper. I \nbelieve the paper can be accepted after some suggested revisions below. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "See aboveNo", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "D18fBBEDfc", "reviewer": "Reviewer_MRuB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n- The ideas to extend GwG to MALA type samplers to overcome the drawbacks of GwG are interesting and clever. I also appreciated the authors discussion on challenges faced by these discrete samplers when incorporating preconditioning thereby showing that it is non-trivial to achieve in the discrete space.\n- The algorithms have been tested on a wide variety of tasks (albeit simpler and more synthetic in nature) demonstrating its performance on a diverse set of problems. \n- The paper is also well written and easy to follow. Sufficient explanation have been provided at appropriate places to make the arguments clear. I'd also commend the authors on providing intuitive explanations for most ideas (I wish more ML papers did this).\n\n**Weakness:**\n- I think a weakness of the paper is the restriction of experiments on relatively easier tasks on either synthetic or simpler datasets. As I mentioned earlier, the authors did a good job to show the performance of the algorithms on diverse tasks. I also understand that these tasks have become commonplace in the literature. However, I believe the paper would have benefitted greatly from considering at-least one task closer to a real-world application or based on a more complicated dataset. For example, Zhang et.al 2022 who concurrently proposed the NCG sampler considered the task of text infilling. Currently, I find the experiments to be limited to relatively lower dimensions and simpler tasks to be interesting beyond as just proofs-of-concept.\n- A motivation for authors with P-AVG was to incorporate higher-order interaction terms. While, the experiments and the performance of P-AVG *can imply* that incorporating such information is useful to the performance of the sampler, I wonder if it might have been useful to construct a more controlled experiment to develop this point. For example, the authors can extend their Ising model experiment to create a frustrated system on a square lattice where the nodes can have both long-range interactions as well as second-order couplings. Sampling from such a system can be beneficial to more directly isolate the benefit of incorporating second-order interactions.\n- *Eq.8*: the normalization term here still seems expensive to me to compute or am I missing something here? How well will it scale to increasing dimensions? In any case, I think it might be useful to explicitly give the computational costs for the methods to keep things transparent and for future works in the community.\n\n*minor points and comments*:\n- It would be easier to understand the gains made the proposed samplers if the code was available through some anonymous link. \n- Through pre-conditioning the authors were able to incorporate second-order interactions. I am curious if the authors have any intuition to extend this to even higher order interactions which might be useful for many problems in Physics like in spin-glass systems? Specifically, in problems which has such higher-order interactions, do the authors have a sense of how well the proposed method will work or how badly they will fail?\n- In Section 3.1.1, the authors provide a discussion on why preconditioning is difficult for NCG. I think the subsection may benefit from a more detailed exposition on how the preconditioning lead to a pairwise MRF and show its intractability. This will help make the manuscript more accessible.\n- In section 4, the authors discuss methods where discrete space is embedded in a continuous one and highlight their drawbacks. The authors miss the ref \"Sampling in Combinatorial Spaces with SurVAE Flow Augmented MCMC, AISTATS 2021\" which addresses some of the drawbacks mentioned in the section. \n- Another small question: in the algorithms box you mention the step of sampling from the proposal distributions you have created. However, it is not immediately clear to me how you sample from say Eq.8? Is it similar to what is done in GwG? If not, can you please add more details around this\n", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors consider the problem of sampling in combinatorial spaces using gradient-based methods that have seen a recent surge in interest due to works of Zanella (2020) on locally-balanced proposals and Grathwohl et.al (2021) on Gibbs with Gradients (GwG). \n\nThe authors identify two main problems with GwG, namely, inability to update multiple dimensions concurrently and reliance on only first order structure of the problem ignoring higher-order interactions. The main contributions of the paper are to address these shortcomings. \n\nThe authors achieve this by first extending the ideas of Grathwohl et.al (2021) to other continuous gradient based samplers like MALA and Preconditioned-MALA to propose Norm-constrained Gradient sampler which can update multiple dimensions at once and Auxiliary Variable Gradient + Preconditioned AVG samplers based on the auxiliary sampler proposed by (Titsias & Papaspiliopoulos) which is capable of incorporating second-order information. They show that the proposed samplers have superior performance to GwG on simple experiments.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Please see my comments above for the details of my points below.\n- Adding relevant citations\n- adding details in section 3.1.1\n- discuss computational complexity of the methods\n- (good to have but I won't consider it necessary at this stage) additional experiment on a real-world application or more complex dataset", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AMuNmsweUG", "reviewer": "Reviewer_UqhQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written, generally easy to follow, and shows expertise from the authors. The two new samplers are generally well derived and presented, and the experiments clearly demonstrate the benefits of the new methods in consistent way. The paper suffers from some (minor) clarity issues, listed below. After these have been adressed, I'm happy to see this accepted.\n\n", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes gradient-based MCMC methods to discrete problems. The paper is based on the idea of interpreting discrete spaces as restrictions of an underlying continuous field, and utilising their gradients. The paper then extends such gibbs-with-gradient MCMC method by applying similar treatment to MALA and its preconditioned version.\n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "\n* Please specify if “f” (or s) is known, unknown, or possible to evaluate. Define what it means that “f is restriction of a differentiable function [which one?]”. Define what it means to take a gradient of a \"restriction\". The background gives the impression that the paper is written as a companion to Gratwohl’21, while it should be self-contained. Similarly in eq 6 the discrete f changes into continuous f, and gradient is taken, but with no explanation what this means mathematically. The paper needs to give a precise and formal presentation of f and nabla f.\n* The discussion of cost of eq 6 is unclaer, and I’m not sure what is costly and what is cheap about eq 6. I assume that entire eq 6 is cheap to evaluate, but doing search around the q(s|..) is expensive. Please elaborate.\n* Eq 6 is in multivariate form where it seems to directly support multivariate updates to s_t. Yet, the text claims that one can only update one dimension at a time. Please explain.\n* The sec 2.4. last phrase implies that linearising the target AND using symmetric kernels can’t account for non-linear target interactions. Why do we need the “and” here? Surely linearising something alone removes non-linear effects? Does the symmetric kernels also pose problems?\n* “as discussed, use of hamming balls in gwg is challenging..” I don’t recognize this discussion. I assume it refers to post-eq-6 paragraph, but it does not explain why increasing the radius is a problem. If you have a linear approximation around center point s_t, surely one can continue the linear plane for as far away as one wants with little expense. Surely one can also make it multidimensional for cheap as well? Please make it explicit where is the bottleneck here. Furthermore, I really don’t see how changing a hamming ball to a gaussian ball changes the cost of anything: surely both are equally tractable and fast to compute. \n* “if we restrict s to discrete space” (before eq8) implies that eq 7 did not restrict f(s) to discrete space. This makes me very confused: did we have a fully continuous f(s) in eq 7? Is this a typo? Please explain what is the nature of f and s.\n* I have hard time understanding eq 8. I don’t understand how we go from 7 to 8. I don’t get why restriction leads to taking a product over dimensions. In eq 7 it was a sum.. I assume that we are completing the square here, but it seems that the s_t^2 term is missing. Please make this understandable.\n* It seems that eq 8 is a MALA-version of eq 6. The main bottleneck was that you couldn’t broaden the hamming ball. I don’t see what eq8 improves. Please explain why the eq8 solves the problem(s) of eq6.\n* 3.2.: I don’t think f(s) has been defined for s \\in R. Eq1 only defines f(s) for discrete s, so thus the pi(s,z) is undefined. The fluid nature of {s,p,f} as continuous or discrete things is very confusing. I would suggest separating the disc/cont versions notationally, eg. by using \\tilde. \n* Can you motivate or explain the N(z | s/sqrt,I) term? I can’t follow why are we dividing the mean with square root. It would also be good to explain what z is or tries to model or conceptualise. It seems that z is some kind of scaled version of s, and s is then following z’s. Perhaps this introduces some lag or momentum to the system? Not sure.. Please explain.\n* In eq 9 the taylor expansion over f(s) leads to the eq 9, but previously also eq 6 was a taylor expansion of f(s) as well, and it lead to a different form without the lone f(s_t) term. Are these different kinds of Taylors? Please explain. \n* eq 11 is missing \\cal{N}\n* eq 13: this looks like second order taylor, but is not. Can we make this kind of approximation? Surely this only applies if M is carefully chosen. However, there is no discussion of what M is. Please elaborate. Later it is claimed that we can use empirical covariance or precision as M. I don’t understand this. M should be a hessian matrix or second derivatives. Surely covariates have little to do with higher order gradients. One should also explain what happens if the eq13 is not satisfied: isn’t this a problem?\n* fig 1 right seems odd. Why do we increase the sampling probably towards corners? Is this intentional? Does the sampling probably keep increasing as a function of distance? The fig1 right needs to zoom out and also show regions where the probably is decreasing or vanishing.\n", "broader_impact_concerns": "No issues", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All three reviewers were satisfied with the efforts and responses put forward by the authors during the response phase, which cleared up some minor outstanding issues raised by the reviewers during the initial review phase. There is consensus that this is a strong contribution and universal support for the publication of the revised manuscript.
Decision: Accept (poster)