Any-scale Balanced Samplers for Discrete Space
Authors: Haoran Sun
Bo Dai
Charles Sutton
Dale Schuurmans
Hanjun Dai
Venue: ICLR 2023 poster (2023)
Keywords: MCMC
Discrete Space Sampling
Locally Balanced Proposal

Abstract: The locally balanced informed proposal has proved to be highly effective for sampling from discrete spaces. However, its success relies on the "local'' factor, which ensures that whenever the proposal distribution is restricted to be near the current state, the locally balanced weight functions are asymptotically optimal and the gradient approximations are accurate.  In seeking a more efficient sampling algorithm, many recent works have considered increasing the scale of the proposal distributions, but this causes the "local'' factor to no longer hold. Instead, we propose any-scale balanced samplers to repair the gap in non-local proposals. In particular, we substitute the locally balanced function with an any-scale balanced function that can self-adjust to achieve better efficiency for proposal distributions at any scale. We also use quadratic approximations to capture curvature of the target distribution and reduce the error in the gradient approximation, while employing a Gaussian integral trick with a special estimated diagonal to efficiently sample from the quadratic proposal distribution. On various synthetic and real distributions, the proposed sampler substantially outperforms existing approaches.

Published as a conference paper at ICLR 2023
ANY-SCALE BALANCED SAMPLERS FOR DISCRETE SPACES
Haoran Sun hsun349 @ gatech.edu
Bo Dai
Charles Sutton
bodai @ @google.com
charlessuttoon @ google.com
Dale Schuurmans $ schuurmans @ google.com
Hanjun Dai hadai @ google.com
ABSTRACT
The locally balanced informed proposal has proved to be highly effective for sam- pling from discrete spaces. However, its success relies the "local" factor, which ensures that whenever the proposal distribution is restricted to be near the current state, the locally balanced weight functions are asymptotically optimal and the gradient approximations are accurate. In seeking a more efficient sampling al- gorithm, many recent works have considered increasing the scale of the proposal distributions, but this causes the "local" factor to no longer hold. Instead, we propose any-scale balanced samplers to repair the gap in non-local proposals. particular, we substitute the locally balanced function with an any-scale balanced function that can self-adjust to achieve better efficiency for proposal distributions at any scale. We also use quadratic approximations to capture curvature of the target distribution and reduce the error in the gradient approximation, while em- ploying Gaussian integral trick with special estimated diagonal efficiently sample from the quadratic proposal distribution. On various synthetic and real dis- tributions, the proposed sampler substantially outperforms existing approaches.
INTRODUCTION
The Markov Chain Monte Carlo (MCMC) algorithm is one of the most widely used methods for sampling from intractable distributions (Robert et al., 1999). Gradient-based samplers that lever- age gradient information to guide the proposal have achieved significant advances in sampling from continuous spaces, demonstrated, for example, by the Metropolis Adjusted Langevin Algorithm (MALA) (Rossky et al., 1978), Hamiltonian Monte Carlo (HMC) (Duane et al., 1987), and re- lated variants (Girolami & Calderhead, 2011; Hoffman et al., 2014). However, for discrete spaces, gradient based samplers remain far less well understood. Recently, a family of locally balanced (LB) samplers (Zanella, 2020; Grathwohl al., 2021; Sun et al., 2021; 2022a; Zhang et al., 2022) have demonstrated promise in sampling from discrete spaces. Such samplers use a locally balanced weight function in an informed proposal Q(z, y) x g (x (zx())
y). such that g R R is weight function that satisfies g() = tg( π is the target distribution, and K, is kernel that determines the scale of the proposal distribution. It is also shown that such locally balanced informed proposal is discrete version of MALA, since they both simulate gradient flows in the Wasserstein manifold (Sun et al., 2022a).
In initial work. Zanella (2020) considered local proposal with kernel K, that restricts next states to lie within a 1-Hamming ball, seeking to capture natural discrete topological structure aris- ing, for example, in spaces of trees, partitions permutations. For more regular discrete spaces, such as lattices, Grathwohl et al. (2021) introduce a gradient approximation for the probability ra- tio ग (yy/m(x) & exp((y x, V log π (I)) to make the locally balanced proposal more scalable.
Work done during an internship at Google.
1 Georgia Tech
1 !Goolle Research, Brain Team University of Alherta

---

Published as a conference paper at ICLR 2023
However, by restricting attention local proposal, these methods tend not to make large jumps and exhibit highly correlated samples. Sun et al. (2021) made the first provably efficient attempt to extend local proposals from 1-Hamming ball to L-Hamming ball, after which subsequent works (Zhang et al., 2022; Sun et al., 2022a; Rhodes & Gutmann, 2022) have shown that using a non-local proposal for the heat kernel K (=) exp(- 2o ||=l3]) can further improve sampling efficiency.
Even though extending locally balanced samplers to non-local proposals has delivered some progress, there remain opportunities for improvement by closing gaps in the current methods. One gap is exemplified by the choice of weight function. To illustrate, consider g() la. For 100 dimensional Bernoulli distribution, we used an informed proposal with the heat kernel (z) exp(-2 20 [|=|²²) and plotted the effective sample size as function of a for different a. Figure shows clearly that performance of varies for different a. In particular, the optimal choice of a monotonically increases with a. When
0, the optimal choice g(t) = vt recovers the locally balanced function. This result indicates that the locally balanced function is no longer optimal for non-local proposals. We will show that good choice of depends on the variance ratio between the target distribution and the kernel. We also give an adaptive algorithm that tunes (a, a.) automatically.
Another gap arises from the gradient approxi- mation. For the local proposal, first order gra- dient approximation is usually sufficient to esti-
e00 of Bernouli Modill
1000
1200
mate the probability ratio. However, for non- local proposal. higher order approximations are generally required to capture correlations be-
1030
800
tween different variables. Extending from re- cent work, consider quadratic approxi- mation of the probability ratio: r(y)/m(<)) exp((y-2:) (x) T V log π (x))+ (y--)) W(y W(y-x)) for non-local proposal, where W is an arbi- trary symmetric real matrix. Unfortunately, the
402
0.4
1.
Figure 1: ESS for different (a, pairs
quadratic heat kernel renders proposal distribution that is pairwise Markov random field in gen- eral, which intractable to directly sample from. However, this difficulty can be addressed by leveraging stochastic factorization via the Gaussian integral trick (Hertz et al., 1991; Zhang et al., 2012), also known as the -aabbarrd-Sttttoovcchh transform (Hubbard, 1959). In particular, we de- compose the quadratic term via (W + & where is a diagonal matrix to make sure W + D is positive semi-definite (PSD) and is standard Gaussian noise. In this paper will show that the quality of the factorization can be characterized by D. While previous work chose to be isotropic, find substantial increase in performance by numerically optimizing over general diagonal ma- trices D.
Closing these two gaps renders our proposal for Any-scale Balanced Sampling (AB Sampling) meth- ods. We extensively demonstrate the advantages of the proposed sampler on both synthetic and real distributions. The results show that, with the proposed numerical optimization of D and the adaptive tuning, the two extensions robustly improve the efficiency of non-local informed proposal.
PRELIMINARIES
Informed Proposal. The informed proposal (Zanella, 2020) is class of Metropolis-Haatinng algo- rithms for discrete spaces, such that the proposal distribution at the current state x has the form:
(1)
(T)
Q3 y)
Ka -/)/Z(x>),, 7(x) = Σ
Ka(
v),
(1)
(xx)
T(x)
ZEX
where π is the target distribution, X' the state space, y : IR+ - R4 is a weight function, Z is the partition function, and K2 is an uninformed kernel with size such that K(z) K(-2) lim K2 (z) = [(==0) and lim (z) = 1. For example, K, can be a Hamming ball kernel K2)z) = [[+[5)), or a heat kernel K2 (x) = exp(- [|=]==//////00 /2a).
Balanced Proposal. Let To denote the reversible distribution associated with the informed proposal Q2: that is, satisfying (z) (z. y) = T (z) (v. æ). We refer to the family {Qa aa as bal- anced proposal if there exists a sequence αı, 22 such that Tay weakly converges to the target
2

---

Published as a conference paper at ICLR 2023
distribution In particular, we say that {Q9}. is locally globally balanced proposal if the sequence of satisfies lim, Jy = 0 or limy oj oc, respectively.
Locally Balanced Sampler. Zanella (2020) showed that the locally balanced function g(t) = tg( () defines the family of locally balanced proposals, which furthermore is asymptotically optimal for locally informed proposals with a 4 0. 'he most commonly used locally balanced function is g(() vl. Grathwohl et al. (2021) introduced a gradient approximation of the probability ratio (()) z exp((y x) T Vog π (2)) to make the locally balanced proposal scalable. Sun et al. (2021); Zhang et al. (2022); Sun et al. (2022a); Rhodes & Gutmann (2022) show locally balanced sampler can be more efficient with large scale and (Sun et al., 2022b) proves that the optimal scaling a for locally balanced proposal is achieved when the average acceptance rate is 0.574.
ANY-SCALE BALANCED SAMPLER
Many recent work (Sun et al., 2021; Zhang et al., 2022; Sun et al., 2022b) have shown using larger kernel Kg can significantly improve the sampling efficiency in locally balanced samplers. Unfor- tunately, in migrating locally balanced samplers to a global proposal regime, these works ignore the fact that locally balanced weight function is no longer optimal and the accuracy of the gradi- ent approximation diminishes. To address such shortcomings, we propose our any-scale balanced samplers. We will consider sampling from the discrete space X' S {1,..., with target distribution x (Ir) (())
3.1 ANY-SCALE BALANCED FUNCTIONS
The first challenge in developing a non-local proposal is that the locally balanced weight func- tion is longer optimal. To determine the proper choice of weight function for kernels K, (00-||=||//20) at different scales, we examine the acceptance rate for the informed proposal in (1) and consider the simple but sufficiently representative weight function class g(() tat for different o. Note that, given a current state x and new state we have the ratio
Ao (v, = T m(x) To (2,
m(y) )a Kr(x - v) Σ.
K (= v)
(2)
(xx)( [(y) )a (y-y-))/_. 3,((2,4
[(() Ko x)
[(a)
n¹ (v)/ Σ, TT (=) v) (1) K (z) म -1 (x)) (z)
(3)
73-0(=)(T***K, (Ta (z)
where (F (x) G(=)G(x - 2) represent the convolution of two functions. Based this formulation, one can easily recover the locally balanced function. In particular, consider the local proposal diminishing scales o, leading to:
lim (T Ko (x) = (<))
lim A, = 1 20 (x)). o-20
(4)
The limit ratio implies that = makes the stationary distribution of To weakly converge to the target distribution π; hence the corresponding weight function g(t) VT, which is one of the most widely used locally balanced functions (Zanella, 2020).
more interesting question how select a for > 0. Since computing the convolution for general target distributions is intractable, we consider continuous relaxation to obtain hint for determining the proper value of a for a given a > 0. In particular, consider a normal target distribution क (.)) ~ (u,,ao)) in the real space IRd, In this case, the ratio has a closed form:
n K₂ ~ (n, (a + oa/a) = π
Ag = π 1-
(v)
(5)
Here, make the proposal balanced, the parameter needs to satisfy
2 vr² + +4
1- aσo 0
ao 5g
2r
One can easily check that this family of balanced functions forms set of interpolants between g(t) = vt and g(t) where the two limiting values are:
r - + 4
r 2
lim a= o=0
lim 0
lim a lim
(7)
2r
2
0
00000
5-000
2r
3

---

Published as a conference paper at ICLR 2023
The first equation recovers the locally balanced function g(t) vt. The second equation shows that, if we consider all states as candidates in the proposal distribution, the optimal choice of weight function is g(t) t, which causes the proposal distribution to become:
T(y)/T(x)
lim (z v) = Exex (2))x() Tty).
(8)
That is, the proposal degenerates to the target distribution. Ignoring computational cost, such a Markov chain draws independent samples from the target distribution in each step and has, in gen- eral, the best efficiency one can expect. Between these two limiting cases, the parameter a E (0, 1) specifies an interpolation that needs be carefully selected based on a to balance the proposal.
To this end, we employ an adaptive algorithm to automatically learn the proper configuration during sampling (Andrieu & Thoms, 2008). Since the hyperparameter pair (o, a)) is highly correlated, can be challenging to directly tune them, so we instead employ a coordinate descent style method that alternatively updates a and a based on the average jump distance. Specifically, we probe the value with (1 + +7)0, a, (1 7) with fixed a and select the new value of a based on which one has the largest average jump distance. And similar method to og see the Appendix A.2 for full details of the adaptation algorithm in Algorithm 3. In our experiments below, we observe that the effective sample size a concave function a for fixed a and vice versa, hence the adaption algorithm is typically able to find good (a, a) configuration efficiently.
3.2 QUADRATIC APPROXIMATION
The second challenge is that, in a non-local proposal, the gradient approximation of the probability ratio becomes less accurate. To capture the correlation between variables, we consider quadratic approximation the log probability change:
f(v) f(x) = (y- x) Vf(x))))))))))))))))(() -)
(9)
When W = fx)(x) is the Hessian matrix, (9) becomes a second order Taylor approximation. In this work, we employ a global W for all states x, hence we choose W as the empirical average Hessian. In particular, the pairs (y - x. Vg(v) Vg(z)) are first collected during burn-in period and W is selected as:
W = arg min - [ - v;) (V)()) - ((=))))
(10)
which can be efficiently solved via gradient descent. Please refer to Appendix A.3 for details. Substituting the quadratic approximation (9), into the informed proposal in (1), with weight function g(f) = t", the quadratic proposal distribution becomes
Q(x,g) x exp [alyy -2) 2) ff(()) + 1/000- ( - =)
(x -))=
1 1 x)-x)T(p-
(11)
20
Although the second order approximation improves proposal quality from the perspective of gra- dient approximation, it also makes (11) become a pairwise Markov random field, which is typi- cally intractable to sample from (Murray, 2007). Therefore, to develop practical sampling algo- rithm, we exploit a stochastic factorization of quadratic proposal distribution known as the Gaus- sian integral trick, which originated in statistical physics (Hubbard, 1959; Hertz et al., 1991) and has been more recently extended in machine learning (Martens & Sutskever, 2010; Zhang et al., 2012). The original Gaussian integral trick is designed for binary random variables, here we show it also works on more general discrete random variables. In particular, for a quadratic distribution (=) x exp(5= W= + TT). and PSD diagonal matrix D that guarantees W + D is PSD, one can introduce a Gaussian auxiliary variable Q(ulz) ~ N((WW + D) I) so that the conditional distribution of given u can be obtained via Bayes' rule:
T W2 + TT (u (W + D)) (u- (W + +D)+2)
Q(2lu) ()) x exp
(12)
2
x exp
2 <T[(W + D) + 2² Dz
(13)
where the square root of a matrix can be obtained by either Cholesky or eigen decomposition. More details of the Gaussian integral trick are provided in Appendix A.5; also see Zhang et al. (2012) for good introduction. To use this trick in sampling from (11), we first sample the auxiliary variable
4

---

Published as a conference paper at ICLR 2023
based the current state Q(u(x)) ~ N(((W + D) xr, I), then propose y according
Quy a u) x exp ay [V)(=) Wx+ (W + D) 2 (aD -00))
(14)
Note that the marginal distribution [ uddu is exactly Q(x, y) in (11), hence we call this a stochastic factorization. By introducing the auxiliary variable u. we avoid calculating the intractable partition function for (11). The M-H acceptance test for this auxiliary sampler is:
A(x, u. y) min { T [ Q {
(15)
Given W, good choice for D should give high acceptance rate and large-variance proposal dis- tribution. However, directly maximizing the acceptance rate and proposal variance at the current sample with respect to D is intractable, therefore, we construct a surrogate.
Consider a continuous relaxation x<x) x exp( f/x0 T Wx), where one can use a 0o, a 1 and the Gaussian integral trick guarantees the acceptance rate is always 1. In this case, the sampling efficiency is only determined by the variance of the proposal distribution. For a current state x and auxiliary variables & and ( (0, Ia), denote u = (W + D)ir + E, and observe
リ D + (W + +)
(16)
for new state in (14), One can compute the variance the change (y - x) in proposal distribution in closed-form:
Eg, ((y - x) - Ely =)) ((y x) Bly =1) " 2D-1
(17)
which is totally determined by the diagonal matrix D. See Appendix A.4 for detailed derivation. Therefore, we would like minimize the diagonal of D for larger variance, thus better sample efficiency, while still keep W + D a PSD matrix.
A common approach to determine diagonal matrix D is X-shift, where D = AT is used with À max (W)), such that e 2 0 is a threshold and Amin is the smallest eigenvalue of W (Martens & Sutskever, 2010). However, such an isotropic choice can suppress movement in dimensions with large variance. For example, consider special case where W is diagonal matrix with Wn 100 and W = for j 2,.... d. Using D 100/ restricts the variance in all dimensions to 0.01, which is inefficient.
Instead, since the quadratic term W is known, one can improve sampling efficiency by a more careful choice of diagonal matrix D. For example, Zhang et al. (2012) claims that the convexity of the proposal distribution depends on the spectrum of W + Following this idea, a straightforward choice is to minimize the largest eigenvalue of W + D. However, instead of only considering one direction, empirically, we find is better to maximize the harmonic mean of D-1 in (17). Intuitively, the harmonic mean provides balanced approach to maximizing the variance of the proposal distribution, as it maximizes variance in all dimensions and puts more weight in directions with smaller variance. Conveniently, recall
arg max D maximizing the harmonic mean of the variance can be reduced to minimizing the trace of D. Under the constraint that W + D is PSD, obtains semi-definite programming (SDP) problem,
arg max D
arg min Σd Dj; D
DDA Dus mingmin
(18)
ii=1 1/D
E=1
min (race(D) s.t. D ž 0, W + D D 0
(19)
D
Empirically, we found there is no need for an exact optimum for (19); a rough solution after early stopping sufficient to characterize the variance scale in each dimension. Using modern solver this estimation step can be typically done in milliseconds for domain with 100 dimensions.
Complete Algorithm. With the any-scale wight function and the quadratic approximation, ready to present our any-scale balanced sampler (AB sampler) in Algorithm 1. The parameters (o,a) are automatically tuned along the way of sampling (also Algorithm 3 and Algorithm for the details on adaptive tuning of these parameters), where W and D are 0 during burn-in, and updated via (10) and (18) right after burn-in.
5

---

Published as a conference paper at ICLR 2023
Algorithm 1: AB sampling algorithm
Algorithm 2: AB M-H step
Input: Initial σ = 0.1, a = 0.5, W = 0, = 0: initial 20
Input: 0, a, W, D: current
Output: MCMC chain To:T and adjusted 0, o, W, D
state Zt
Burn-in period < T1: alternatively update a, a use algorithm 3 while calling M-H step defined in algorithm 2 as subroutine; collect trace Jo:7; along the way
Output: new state 2++1
Sample auxiliary:
u u~ <<W + D))
Estimate W and D using collected TO.T, via (10) and (18) Mixed period Ti < t T: use estimated W and D to continue the alternatively updating 0, a with algorithm 3, while calling M-H step defined in algorithm 2 as subroutine; Return the entire trace To.r and the estimated parameters.
Sample new state: ~ (14)
3 Get acceptance rate A
equation (15)
if rand(0, 1) < A then
r++1 y else Z++1 2/
4 RELATED WORK
The informed proposal, which uses information about the target distribution to guide the proposal for the Metropolis-Hastings (M-H) algorithm has been extensively studied for discrete spaces in recent years. A number of methods have attempted to first map the discrete to a continuous space, using relaxation, apply gradient based methods in the continuous space, then map the new state back to the discrete space, either by using auxiliary variables, uniform dequantization, or VAE flow (Zhang et al., 2012; Pakman & Paninski, 2013; Nishimura al., 2017: Han & Liu, 2018; Zhou, 2020; Jaini et al., 2021). Such methods work in some scenarios, but embedding a discrete into continuous space often destroys its natural topological structure, and can create highly multimodal and irregular target distributions (Zanella, 2020). As shown in previous work, such methods does not scale well to high dimensional discrete settings (Grathwohl et al., 2021).
Another group of methods attempt to directly work within discrete spaces. Titsias & Yau (2017)
and Dai et al. (2020) introduce auxiliary variables to trade off the number of updated variables in block against computational cost, however, by relying on Gibbs sampling, such methods still require significant overhead to make updates. In addition to the related works (Zanella, 2020; Grathwohl et al., 2021; Sun et al., 2021; 2022a; Zhang et al., 2022) already discussed in depth above, a concur- rent work (Rhodes & Gutmann, 2022) has considered preconditioning and also used the Gaussian integral trick to incorporate second order information from the target distribution, but this work does not study how to properly choose the weight function g. the hyperparameter (ar, a). and the diagonal matrix D, making the resulting algorithm less efficient. Another recent work (Sun et al., 2022b) proves that the optimal scale a for locally balanced proposal is achieved when the average accep- tance equals 0.574, and give robust adaptive algorithm for tuning However, its result relies on the property of locally balanced function, and does not apply to more general weight function g(t) = ta with of 0.5.
EXPERIMENTS
We conducted an experimental evaluation on three types of target distributions: 1) quadratic syn- thetic distributions, 2) non-quadratic synthetic distributions, and 3) real distributions. For quadratic synthetic distributions, we focus on demonstrating the benefits of selecting high quality diagonal matrix D. For nuu-quadratic synthetic distributions, we show that the performance of the proposed sampler significantly relies on the choice of weight function g(t) t" For real distributions, compare against baseline samplers on challenging inference problems in deep energy based models trained on MNIST, Omniglot, and Caltech datasets.
5.1 SETTINGS
Samplers. We denote the proposed Any-scale Balanced sampler as AB-trace sampler, which uses the any-scale balanced function g(t) = ta and obtains the diagonal matrix D by minimizing its trace. For comparison, we consider the classical discrete samplers, random walk Metropolis (RWM) and Gibbs sampler. We also compare to a locally balanced sampler (LB), considering representative version DLP in Zhang et al. (2022) that uses a = 0.5 and W = 0 in (11); this is mathematically equivalent to NCG in Rhodes & Gutmann (2022). For RWM and LB, we follow the optimal ac- ceptance rate in (Sun et al., 2022b) and tune the scale of the proposal distribution until the average acceptance is 0.234 and 0.574, respectively.

---

Published as a conference paper at ICLR 2023
Table 1: ESS on selected Quadratic Distributions
Distribution
Grid Ising
BA-4 Ising
Rotation Gaussian
Sparse Gaussian
Sampler
E
ESS.
ESS,,
ESS.
ESS
ESS:
ESS,
ESS,
Gibbs
1.66
0.55
8.53
2.84
2.81
11.26
4.20
16.81
RW
1.27
0.31
4.28
1.07
9.41
2.69
8.67
2.48
DLP
2.96
0.49
120.15
20.89
86.28
15.69
71.23
13.57
DLP-trace
2.88
0.47
143.25
24.90
86.62
15.75
142.86
27.21
AB-Ist
3.99
0.67
242.31
35.89
103.16
16.50
79.88
13.31
AB-shift (PAVG)
7.87
1.12
274.18
39.16
196.88
30.29
154.00
24.64
AB-max
7.83
1.11
287.24
41.03
217.41
33.45
474.75
75.96
AB-trace
7.94
1.13
776.31
110.90
239.09
36.78
1227.94
196.47
To demonstrate the benefit of the proposed methods, consider few variants for ablation: DLP-trace, which uses the same anisotropic diagonal matrix D as AB-trace in kernel (x)z) exp(-= Dz/a) for DLP, AB-Ist, which only uses gradients approximate the probability ratio, and AB-shift and AB- Bax, which obtain the diagonal matrix D via X-shift or minimizing the max- imum eigenvalue of W D, respectively. For all AB- samplers, we tune (a, a) adaptively via the algorithm discussed above (and described in Appendix A.2). Note that the PAVG sampler in (Rhodes & Gutmann, 2022) is equivalent to AB-shift with fixed (a, a) = (oo, 1). Since we find that tnning (ar a)) improves the efficiency, we use AB-shift to represent PAVG. More details about the sampler implementations, such as solving D, are given in Appendix A.I.
Metrics. As in other works (Hoffman et al., 2014; Zanella, 2020), we use effective sample size (ESS) (Lenth, 2001) to characterize the efficiency of the samplers on synthetic distributions. To reduce the effects of implementation, denote the ESS for every 10,000 queries of the log likelihood function, and ESS, denote the ESS for every second of sampling. For each setting and sampler, with 71=20,000 burn-in steps to make sure the chain mixes. For real distributions, we compare the mixing time for different samplers.
report ESS normalized in two different ways: We let ESS,,
100 chains for T=100,000 steps,
5.2 QUADRATIC SYNTHETIC DISTRIBUTIONS
Ising model. The Ising model (Ising, 1924) is mathematical model of ferromagnetism in statistical mechanics. It consists of binary random variables arranged in a graph G = (V, E) and allows cach node to interact with its neighbors. The unnormalized log probability function of the Ising model is:
f(x) = Eiev w,x; + (iis)e JIy
(20)
In this experiment, we consider Ising models on 2D grid graphs and Barabasi-Albert-e graphs (Al- bert & Barabási, 2002). For grid Ising, we set Jú, 0.4407 the critical temperature (Onsager, 1944), so that the model is at its transition phase and hard to sample from; see Appendix B.2 for a more detailed description. We conduct sampling high, medium, and low temperatures. In Table 1, we report results for the medium temperature, where Jig = 0.4407. More results on Ising model are given in Table 2 and Table 3.
The lattice Gaussian model is obtained by restricting the Gaussian distri- bution to Lattice, which is an important distribution in coding and cryptography (Kschischang & Pasupathy, 1993; Micciancio & Regev, 2007). The unnormalized log probability is:
Lattice Gaussian Model.
f(x)
(a - W b).
(21)
In this experiment, we use a finite state space X' {0. 1..... 203 100 and we investigate two settings for the Gaussian model. The first setting is a rotated Gaussian W P AP, where P is orthog- onal matrix and a diagonal matrix. The second setting is a Sparse Gaussian, which a pairwise Markov random field defined on cycle. We constructed the lattice Gaussian models with low, medium, and high conditions. More detailed descriptions of these models are given in Appendix B.3. We report the results for medium condition in Table 1. More results given in Table 4 and 5.
Results Analysis. One can observe that the AB samplers substantially outperform existing samplers on all distributions. Specifically, the first order sampler AB- 1st has ESS consistently larger than LB, which justifies the benefit of selecting the proper weight function. Also, the AB-tra A--trace sampler has comparable efficiency to AB-shift on Grid Ising and Rotation Gaussian, but is significantly better
7

---

Published as a conference paper at ICLR 2023
Demo: gitl 2 on BLR Model
Sampler
ESS..
ESS;
Relative Effocitive Sample
Gibbs
43.81
14.60
RW
25.00
6.25
DLP
222.49
35.60
DLP-trace
247.35
40.69
AB-lst
321.56
51.44
AB-shift (PAVG)
265.89
37.98
AB-max
266.49
38.06
AB-trace
403.78
57.68
0.7
0.8
0.9
1.0
1.1
12
Figure 2: Results on BLR: (1) ESS for different samplers, (r) ESS for AB-trace with different (a,
Demo: g(t) t on Mixture Model
Sampler
ESS.
ESS;
456
Gibbs
13.66
34.14
RW
189.69
34.49
DLP
97.97
12.25
300
DLP-trace
100.65
12.58
250
AB-Ist
256.88
32.11
AB-shift (PAVG)
327.81
36.42
Relative
200
AB-max
355.78
39.53
150
AB-trace
380.50
42.28
0.7
1.1
12
Figure 3: Results QMM: (I) ESS for different samplers, (r) ESS for AB-trace with different (a, o.)
BA-4 Ising and Sparse Ising. The reason is that the variables in Grid Ising and Rotation Gaussian are nearly homogeneous, and an isotropic diagonal matrix D = A/ is not slowed by several hard dimensions. However, in BA-4 Ising and Sparse Gaussian, the variance in different dimensions can be very different, and AB-trace demonstrates significant advantages by employing a general diagonal matrix D that allows different step sizes in different dimensions.
5.3 NON-QUADRATI SYNTHETIC DISTRIBUTIONS
Bayesian Logistic Regression (BLR). Following Zhou (2020), we consider logistic regression
model Y cernoullissigoide (3))), with Y E {0, 11"1 X E IR , E {0, 122 We first generate the sample X, Y. then, using a uniform prior, the target distribution is the posterior of with the unnormalized log probability function:
E V4 log (1 exp(-o,)) + (1-y,) log (1 (xxx))))))), o, - X (22) w.
In this experiment, we considered a d = 100 dimensional regression with 50 samples. More details for generating X, Y are given Appendix B.4. We report the results for ESS Table For AB-trace, the selected configuration is aa. (16, 0.96). To justify the quality of this selection, we also plot the ESS for AB-trace with different (o, a) in Figure
Quartic Mixture Model (QMM). Following Rhodes & Gutmann (2022), consider quartic mixture model, where the unnormalized log likelihood function can be written as:
fxx) log (ΣK=, exp(- exp(-poly{xx))
(23)
such that poly multivariate polynomial with degree 4. In this experiment, we use a finite state 50 components for the mixture model. More details about poly: given in Appendix B.5. We report the ESS results in Figure 3. For AB-trace, the selected configuration (o, a) = (415, 0.92). To justify the quality of this selection, we also plot the ESS for AB-trace with different (a, a) in Figure 3
space X' = {0. 1,... 20]50 and K
Results Analysis. For the non-quadratic synthetic distributions, the AB-trace sampler significantly outperformed the other methods. From the curves for (a, a),, algorithm successfully found optimal configurations. Note that in Figure 2 and Figure 3, the values 1000 can be infinity, since further increasing the does not influence
that the adaptive tuning
64 and
8

---

Published as a conference paper at ICLR 2023
mixing carves on mrist
mixing curves omniglot
mising curves caltach
FW
g100s
gbbs
glbb
1b
ab
Vin sampling stops
sampling stops
sampling stops
Figure 4: Mixing Time on Real Distributions
have some estimation emror for the probability ratio. One interesting phenomenon is that the first order method AB-Ist can be more efficient than second order samplers AB-shift and AB-max in BLR. The reason is that the Gaussian integral trick introduces extra variance in the proposal distribution. If the diagonal matrix D is not properly selected, the benefit of using second order sampler can be reduced.
efficiency. Unlike from (8), the optimal
still not 1
5.4 DEEP EBMS ON REAL DISTRIBUTIONS
Having observed excellent performance of the AB sampler synthetic datasets, considered sampling in more challenging real distributions. In particular, here we trained deep EBMs param- eterized by ResNet (He et al., 2016) the MNIST, Omniglot, and Caltech datasets. In these real image distributions, we are interested in how fast sampling algorithms can find high quality images, so we report the mixing rate in figure 4. Since we comparing behavior during the mixing stage, we do not have samples to estimate W via (10), hence we use AB sampler with bit different from Algorithm 1. In particular, we use the true data (from datasets) estimate the variance var, for each variable 2; and set W as diagonal matrix with W;, = 1/(1 + var;). In this case, we do not need to use the Gaussian integral trick and we do not distinguish the different version of AB samplers. More details are given in Appendix B.6.
In Figure 4, one can see that the AB sampler mixes faster than the LB sampler on all three real distri- butions. The optimal a selected for MNIST, Omniglot, and Caltech are 0.6, 0.55, 0.55, respectively. They are significantly smaller than that in non-quadratic synthetic distributions. We believe the rea- son is that these deep EBMs are much more complicated than the synthetic distributions. Larger estimation errors only allow the sampler to make local movements, and hence have a smaller a.
CONCLUSION
In this work, we proposed an Any-scale Balanced sampler (AB sampler) that substantially improves existing locally balanced samplers for discrete spaces in two respects:
the AB sampler goes beyond considering the locally balanced function as an "optimal" choice for weight function in an informed proposal, and provides an adaptive algorithm for finding the optimal configuration of (a, a));
the AB sampler introduces the Gaussian integral trick, which allows efficient second order approximation to improve proposal quality.
There are still directions for further improvement of the AB sampler. First, current adapting Algo- rithm 3 tunes (a, a:) based the empirical estimation of jump distance, which can vary lot during the mixing process. As a result, our adapting algorithm is not stable until the Markov chain reaches its stationary distribution. This is not a big problem in sampling. But if we want train EBMs via contrastive divergence (Hinton, 2002; Tieleman, 2008), the current adapting algorithm can hardly find the optimal configuration for (a, a) as the model keeps changing. potential solution is to use adaptive algorithms based on acceptance rate (Roberts & Rosenthal, 2001; Sun et al., 2022b). Second, the estimation of W in (10) is rough. In complicated distributions, the Hessian matrix can vary lot at different state ar. More accurate quadratic approximation can be obtained via Rieman- nian (Girolami Calderhead, 201 or quasi-Newtor style algorithms (Zhang & Sutton, 2011) that allow W- W (z) depend on the current state æ. Third, the current quadratic approximation can be inefficient on large models. For example, on a large sparse quadratic model, solving the sparse SDP can be time consuming, and the square root (W + D) is not necessarily sparse. Hence, more sophisticated design are needed to make the quadratic approximation being scalable.

---

Published as a conference paper at ICLR 2023
REFERENCES
Réka Albert and Alsert-László Barabási. Statistical mechanics of complex networks. Reviews of modern physics, 74(1):47, 2002.
Christophe Andrieu and Johannes Thoms. tutorial on adaptive MCMC. Statisties and computing, 18(4):343-373, 2008.
MOSEK ApS. Mosek optimization suite, 2019.
Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiliary-variable local exploration. Advances in Neural Information Processing Systems, 33:10443-10453 2020.
Simon Duane, Anthony D) Kennedy, Brian J Pendleton, and Duncan Roweth. Hybrid monte carlo. Physics letters B, 195(2):216-22,, 1987.
Mark Girolami and Ben Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2):
123-214, 2011.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris J Maddison. Oops I took a gradient: Scalable sampling for discrete distributions. arXiv preprint arXiv:2102.0509, 2021.
Jun Han and Qiang Liu. Stein variational gradient descent without gradient. In International Con- ference on Machine Learning, pp. 1900-1908. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
John Hertz, Anders Krogh, Richard G Palmer, and Heinz Horner. Introduction to the theory of neural computation. Physics Today, 44(12):70, 1991.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800,, 2002.
Matthew D Hoffman, Andrew Gelman, al. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res., 15(1):1593-1623, 2014.
John Hubbard. Calculation of partition functions. Physical Review Letters, 3(2):77, 1959.
Ernst Ising. Beitrag zur theorie des ferro-und paramagnetismus. PhD thesis, Grefe & Tiedemann, 1924.
Priyank Jaini, Didrik Nielsen, and Max Welling. Sampling combinatorial spaces with SurVAE flow augmented MCMC. In International Conference on Artificial Intelligence and Statistics, pp. 3349-3357. PMLR, 2021.
Frank R Kschischang and Subbarayan Pasupathy. Optimal nonuniform signaling for Gaussian chan- nels. IEEE Transactions Information' Theory, 39(3):913-929, 1993.
Russell V Lenth. Some practical guidelines for effective sample size determination. The American Statistician, 55(3): 87-193, 2001.
James Martens and Ilya Sutskever. Parallelizable sampling of Markov random fields. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp- 517-524. JMLR Workshop and Conference Proceedings, 2010.
Daniele Micciancio and Oded Regev. Worst-case to average-case reductions based on Gaussian measures. SIAM Journal on Computing, 37(1):677-302 2007.
lain Murray. Advances in Markov chain Monte Carlo methods. University of London, University College London (United Kingdom), 2007.
10

---

Published as a conference paper at ICLR 2023
Akihiko Nishimura, David Dunson, and Jianfeng Lu. Discontinuous Hamiltonian Monte Carlo for sampling discrete parameters. arXiv preprint arXiv:1705.085/11 853, 2017.
Lars Onsager. Crystal statistics. i. two-dimensional model with an order-disorder transition. Phys- ical Review, 65(3-4):117, 1944.
Ari Pakman and Liam Paninski. Auxiliary-variable exact Hamiltonian Monte Carlo samplers for binary distributions, arXiv preprint arXiv:/311.2166. 2013.
Benjamin Rhodes and Michael Gutmann. Enhanced gradient- based MCMC in discrete spaces. arXiv preprint urXiv:2208.00000 2022.
Christian P Robert, George Casella, and George Casella. Monte Carlo Statistical Methods, volume 2. Springer, 1999.
Gareth Roberts and Jeffrey S Rosenthal. Optimal scaling for various metropolis-hastngg algo- rithms. Statistical science, 16(4):351-36) 2001.
Peter Rossky, JD Doll, and HL Friedman. Brownian dynamics smart Monte Carlo simulation. The Journal of Chemical Physics, 69(10):4628-4633 4633, 1978.
Haoran Sun, Hanjun Dai, Wei Xia, and Arun Ramamurthy. Path auxiliary proposal for MCMC in discrete space. In International Conference on Learning Representations, 2021.
Haoran Sun, Hanjun Dai, Bo Dai, Haomin Zhou, and Dale Schuurmans. Discrete Langevin sampler via Wasserstein gradient flow. arXiv preprint arXiv:2206. 14897, 2022a.
Haoran Sun, Hanjun Dai, and Dale Schuurmans. Optimal scaling for locally balanced proposals in discrete spaces. arXiv preprint urXiv:2209.08133 2022b.
Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference Machine learning, pp. 1064 2071,2008
Michalis K Titsias and Christopher Yau. The Hamming ball sampler. Journal of the American Statistical Association, 12((520):1598 1611, 2017.
Giacomo Zanella. Informed proposals for local MCMC in discrete spaces. Journal of the American Statistical Association, 115(530):852-$66 2020.
Ruqi Zhang, Xingchao Liu, and Qiang Liu. A Langevin like sampler for discrete distributions. In International Conference on Machine Learning, pp. 26375-26396. PMLR, 2022.
Yichuan Zhang and Charles Sutton. Quasi-newton methods for markov chain monte carlo. Advances in Neural Information Processing Systems, 24, 2011.
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.
Guangyao Zhou. Mixed Hamiltonian Monte Carlo for mixed discrete and continuous variables. Advances in Neural Information Processing Systems, 33:17094-17100 2020.

---

Published as a conference paper at ICLR 2023
A SAMPLERS
A.1 SEMI-DEFINITE PROGRAMMING
Solving the diagonal matrix D for AB-trace and AB-max can be formulated as semi-definite pro- gramming (SDP). In particular, for AB-trace, the SDP problem is:
D = arg min (race(D)
(24)
D
s.t. DI O, W WDD20
(25)
For AB-max, the SDP problem is
D arg min min X
(26)
s.t. DL O, A
(27)
Both SDP problems can be efficiently solved by modern SIP solver. In this work, version of Mosek (ApS, 2019). For models have less or equal to 100 variables (Lattice Gaussian, Bayesian Logistic Regression, Quartic Mixture Model), Mosek takes less than 0.05 second solve the SDP problem. For models with 400 variables (ISing), Mosek took less than second to solve the SDP problem. For models with 784 variables, Mosek takes around 10 seconds to solve the SDP problem. For all distributions considered in this work, the time used for solving SDP is negligible comparing to the sampling time. However, for models with several thousands or more variables, the cost for directly solving the SDP could be high and better methods to estimate the diagonal matrix D are needed.
academia
A.2 ADAPTIVE TUNING ALGORITHM
We give the pseudo code for adaptive tuning of (a, a) in Algorithm 3. The basic idea is alternatively updating and a to maximizing the average jump distance. In line 1, 2, and in Algorithm 4, the samples are collected via calling M-H step of AB sampler as in Algorithm 2.
Algorithm 3: Adapting Algorithm
Input: initial 0.1, a 0.5, update rate y 0.2, decay rate B 0.9, initial state zy. buffer size N = 100.
Output: parameters a, a, samples 21, T2,
for i 0, do
(I6IN+1. +3N)+ t Adapting Algorithm Block(f = o, 7, 20)
a' (2GGN .... FBIN+6N Adapting Algorithm Block(I a, 2T)
end
a = ,',a = a' then
7= = B7
else
== σ', = a
end
A.3 QUADRATIC APPROXIMATION
Here, we explain how to efficiently solve the following optimization problem:
W arg min
Elw( - x2)
(28)
W-W
Denote the y1 I; forms a matrix Nxd such that the i-th row M; = yı 2;, Similarly, denote the ff()) (z;) forms matrix Y E Nxx such that the i-th row Y, Vf(m) Vf(x)). Then, the loss function can rewritten as:
W* arg min Σ
|XW., 2
(29)
WT
12

---

Published as a conference paper at ICLR 2023
Algorithm 4: Adapting Algorithm Block
Input: target parameter 0, adapting rate 7, initial state xo
Output: updated parameter 0', samples 31, ... 43N
Using parameter to sample Pi,--.; &N via Algorithm 2
Compute de
Σn Si=l [2 11-1/1
3 Using parameter 0(1 + 7) to sample 2N+1. T2N via Algorithm
Compute d+ Σ 1=1 [2N+s IN+i-1/1
Using parameter { - ?) to sample 22,+1: TaN via Algorithm
Compute d__ = > 12 T2N4i-111
if (axx (eaxx(do,dc d.,44,.0 d_ = d. then
0' 0 77)
else if max do. d+, d_ == d_ then
10
0' 0(1 ?)
else
12
end
where W., represents the j-th column of W. One can easily see that the loss function is a regression.
For the feasible region W = W one can easily check is a d(d+1) 2 dimensional linear subspace. As result, one only efficiently solving this convex optimization problems via projected gradient
descent, where the projection to symmetric matrix space is simply X x+x*
A.4 VARIANCE OF THE PROPOSAL DISTRIBUTION
For & and & ~ N(O Ia), u = (W D) x + ç, and
Dー W DW + D) D-
we have
Ry Re, ((y (1y-x)) - ((y
Ely x1)
(30)
= EE, D " D-
(31)
(32)
-2D-1
(33)
where (32) is because for a normal random variable exp( Wx), the variance of is -W
A.5 GAUSSIAN INTEGRAL TRICK
The Gaussian integral trick, also known Hubbard-SStratonvviin transform (Hubbard, 1959), is first named in (Hertz et al., 1991) and extedned by Martens & Sutskever (2010); Zhang et al. (2012) for efficient Gibbs/HMC sampling inference. The main idea is for discrete-valued pairwise-MRF with within-layer connections, by introducing a real-valued auxiliary variable, the quadratic form in the energy function, x T Wa, will be canceled out. Thus, the inference will be easy earry on. Specifically, we would like to sample from a MRF with pairwise dependency, ie.,
x. Wa
p(z) exp(
+ + b)/Z.
(34)
2
where æ E {0, 1 and 2 = chain Monte Carlo or Gibbs sampling. which although provably converges to the target distribution, but still might stuck in some region.
Σ, exp(x Wa + b). The vanilla sampler for this model is Markov
To accelerate the sampling, one can introduce the auxiliary variables to reformulate the MRF into family of equivalent Boltzmann machines. Concretely, we introduce the auxiliary variable u with conditional distribution:
p(ulx:) (NllA(WW D)a, A(W + D)A
(35)
13

---

Published as a conference paper at ICLR 2023
Figure Visualization: (1) 10 10 Grid Graph, (r) 50-4 BA Graph
where D diag(d). Therefore, we have the joint distribution
" (A-1) (W + D)
A u
p(z, u) exp
A u
6
(36)
2
which cancels the quadratic term w.r.t. x and makes x independent for each dimension. This induces the p(z|u) is multivariate Bernoulli distribution, i.e.,
p(r II (1 p.) 1- p
(37)
with p; I+cxp(-ft.) b; d (A (A-2u).-
One shall notice that in the original Gaussian integral trick (36), the simplification - Dr - dddd using the property that this simplification. As long as D is diagonal, we can factorize the distribution and make efficinet proposal. In this work, we consider using = (W + D)- such that the conditional distribution has following simple forms:
binary random variables. Actually, we don't have
p(u[x) NN(l(WW + D) l);
(38)
p(z(l)) exp (TT(WW D) + b) 2"
1 T Dx)
(39)
B EXPERIMENTS
B.1 HARDWARE
All experiments are running on a virtual machine with CPU: Intel Haswell, GPU: 4x Nvidia V100, System: Debian 10.
B.2
ISING MODEL.
The Ising model (Ising, 1924) is mathematical model of ferromagnetism in statistical mechanics It consists of binary random variables arranged graph G (V, E)) and allows each node to interact with its neighbors. The log probability function of Ising model is:
f(x)
w.T; Σ
Jiyna
(40)
(UJEE
In this experiment, we consider Ising models on 2D grid graphs and Barabasi-Albers graphs (Albert & Barabási, 2002).
grid Ising. We consider 20 x 20 grid graphs. See Figure for visualization of a 10 x 10 grid graph. We zero external force w; and set the interaction Jús 0.3000, 0.4407, 0.7071 for high,
14

---

Published as a conference paper at ICLR 2023
eritical, and low temperatures. We report the results in Table 2. Consistent to the results in statisti- cal physics, phase transition occurs at the critical temperature and makes the sampling much harder (Onsager, 1944). More detailedly, at low temperature, variables in grid Ising model only have strong correlation with variables close to it. At critical temperature, the correlation is global and a variable strongly depends variables far from it. At high temperature, the variables have weak corre- lation to all the other variables. On all these three scenarios, Any-scale (AB) samplers substantially outperforms previous discrete sampling methods, including locally balanced (LB) samplers.
Table ESS on 20 20 Grid Ising
Distribution
Interaction (0.3000)
Interaction (0.4407)
Interaction (0.7071)
Sampler ESS,
ESS,
ESS,,
ESS;
ESS,,
ESS,
Gibbs
4.50
1.50
1.66
0.55
6.11
2.04
RW
4.81
1.20
1.27
0.31
2.02
0.50
LB
43.16
6.91
2.96
0.49
62.06
10.34
AB-Ist
69.91
11.19
3.99
0.67
78.44
13.07
AB-shift 178.70
24.65
7.87
1.12
111.94
15.44
AB-max 178.68
24.65
7.83
1.11
117.06
16.15
AB-trace 182.91
65.23
7.94
1.13
149.06
20.56
BA Ising. We consider 400-4 BA graph, that's to say, a Barabasi-Albert random graph with 400 nodes and 4 attach edges for every node (Albert Barabási, 2002). See Figure B.2 for visualization of 50-4 BA graph. Since we don't know the critical temperature in BA graphs, we keep using the settig grid Ising with zero external force w; = and interaction Jıj 0.3000, 0.4407, 0. 0.7071 for high, critical, and low temperatures. We report the results in Table 3. On all three temperatures, Any-scale (AB) samplers substantially outperforms previous discrete sampling methods, including locally balanced (LBB) samplers. Also, forms other AB samplers using quadratic approximation. low temperature model Ising (0.7071), the first order method BA-1st even beat AB-shift and AB-max using quadratic approximation. The reason is that the variables in BA graph has inhomogeneous topology and a casual selection of D does help. Furthermore, using quadratic approximation has involve randomness in Gaus- sian integral trick and harm the proposal quality. In grid graphs where different variables have very similar topology, thus this drawback is less significant.
notice that AB-trace sampler significantly outper-
Table 3: ESS on 400-4 BA Ising
Distribution
Interaction (0.3000)
Interaction (0.4407)
Interaction (0.7071)
Sampler ESS,,
ESS,
ESS,,
ESS,
ESS,,
ESS,
Gibbs
8.19
3.27
8.53
2.84
8.75
3.50
RW
7.07
2.02
4.28
1.07
3.06
0.88
LB
117.06
20.36
120.15
20.89
153.16
26.64
AB-Ist
207.91
36.16
242.31
35.89
354.50
61.65
AB-shift
367.72
54.48
274.18
39.16
160.03
23.71
AB-max
402.06
59.56
287.24
41.03
172.47
25.55
AB-trace
779.12
115.43
776.31
110.90
848.72
124.85
B.3 LATTICE GAUSSIAN
Rotation Gaussian. For rotation Gaussian, we use bias vector b = 0. Given parameter L, we generate the weight matrix W in the following way: We first generate the diagonal matrix Ig 1100x100 100 such that
VL
A# 1+ 1) 1)/99'
= 1,
100
(41)
15

---

Published as a conference paper at ICLR 2023
6
8
3
0
9
2
8
9
5
5
6
5
2
J
4
1
7
0
5
5
3
1
5 7 7
3
/
0
0
2
4
PT
T
0
1
4
6
3
6
7
JJ
4
/
8
7
1
4
6
V
R
R /
7
1
5
2
8
8
4
5
6 0 7
8
P
Q T F
1
7
5
9
4
6
6
2
4
9
1 5
5
I.
2
8
6
A Q t
F
MNIST
Omniglot
Caltech
Figure 6: Images sampled from the trained EBMs.
Then generate orthogonal matrix P and let W -PAPT. The results for rotation Gaussian with L = 2, 10, 50 are reported Table 4.
Sparse Gaussian For Sparse Gaussian, the weight matrix W in the following way: We first generate the matrix M E IR 10 100 with Wa and W Wi+1, NO, 0.04), for i 1, ..., 100. One shall notice that denote W 1000.101 = 100,1 and W W101,100 = Wi,100 Then we generate the diagonal matrix A, such that
bias vector 0. Given parameter I, generate
99
A; + i ( (L-
i 1,2, 1,2,... 100
(42)
Then, let W AMA. Such log probability function defines graphical model defined a cycle. Since a cycle is sparse, we call it sparse Gaussian. The results for sparse Gaussian with L= 2, 10, 50 are reported in Table 4.
Table 4: ESS 100d Rotation Gaussian
Distribution
L
L 10
L
50
Sampler
ESS,,
ESS;
ESS,,
ESS:
ESS,,
ESS;
Gibbs
4.55
1.82
2.81
11.26
1.64
6.58
RW
14.53
4.15
9.41
2.69
6.20
1.77
LB
222.47
42.38
86.28
15.69
27.28
5.20
AB-Ist
492.12
93.74
103.16
30.29
27.97
5.33
AB-shift
1182.16
189.15
196.88
30.29
35.72
5.71
AB-max
1224.06
195.85
217.41
33.45
36.75
5.88
AB-trace
1244.34
199.09
239.09
36.78
52.84
8.46
Table 5: ESS on 100d Sparse Gaussian
Distribution
L
Sampler
ESS,,
ESS,
ESS,,
ESS,
ESS,,
ESS,
Gibbs
3.80
15.21
4.20
16.81
4.42
17.68
RW
11.80
3.37
8.67
2.48
9.99
2.85
LB
147.91
28.17
71.23
13.57
32.19
6.13
AB-Ist
247.59
47.16
79.88
13.31
36,62
6.98
AB-shift
545.06
87.21
154.00
24.64
45.19
7.23
AB- 8-max
680.15
108.82
474.75
75.96
132.88
21.26
AB-trace
916.76
146.68
1227.94
196.47
1417.57
226.81
16

---

Published as a conference paper at ICLR 2023
B.4 BAYESIAN LOGISTIC REGRESSION
We consider the logistic regression model Y ~ Bernoulii(simmoi( X 8)), with Y E {0.1] 10,10,0, X e
R50x100 3 E {0, 1] 1000_ We first generate X E R Each row X; is realization of the normal distribution (0, 0.), where Σn = 1.25, σz = 0.25, and A₂ = exp( 2xp(--.255 + (i (i- 1)/99). Then, we set the ground truth & that 8; 1 for 2 1.2..... 7 and 3 = for i 8.9. 100. Then, we get the logits v = X3. Then, we sample Y; ~ Bernoulli (ex)) for i = L.... 50, where (t) 1/ exx((-t)) Our target distribution the posterior of and the log probability function is:
m
f(3)
Σ
y: log (1 +1+ cxp(-o,) + (1 log (1 exp(o:)), o, Exno
Xiy
(43)
B.5 QUARTIC MIXTURE MODEL
Following Rhodes& & Gutmann (2022), we consider quartic mixture model, where the log likelihood function can be written as:
K
((x) log Σ exp((- Eexp(-poly((:t)
(44)
where polyi: is multivariate polynomial with degree 4 generated in the following way. For com- ponent = 1, .... 50, the associated bias by k-
49 E R Denote R230 such that
8, = 20 exp(-0.5 + (i 1)/19) for = 1, ..., 20. Then, we have the vector 2k (x b/)/s; s; E 20
(45)
We also generate a rotation matrix P E R0 shared by components and we have
Pck E IR 20x50
(46)
Then, the polynomial is defined as
(x)
Σ d=1
uua + 1)(u + 0.5) 2)
(47)
B.6 DEEP EBM
Model training. In the main text compare the sampling efficiency of different samplers using the trained deep EBMs. Here we provide more details on obtaining these pretrained EBMs.
We follow the existing works parameterize the EBMs using ResNets, where it is trained using persistent contrastive divergence (Tieleman, 2008) framework. Specifically we follow Grathwohl al. (2021); Sun et al. (2021) to maintain buffer of multiple MCMC chains. We use the sampler proposed in Sun et al. (2021) and run 60 steps to obtain samples from the current model per each gradient update. We retain the model after 50,000 steps of training. The models are all reasonable and can produce realistic binary images as the ground truth data.
Estimating W. Since we compare the mixing time on real distributions. we can not use AB sampler as Algorithm which always use W = D = 0 during burn-in stage. Instead, we directly use the true data from the datasets to estimate the variance var; for each variable. Then we set W as diagonal matrix with W. = 1/(1 + var;). In this case, the proposal distribution in equation 11 is naturally factorized and we don't need to use Gaussian integral trick.
Adaptive Tuning. The adaptive Algorithm 3 tuning (a,a) based average jump distance, which could be very unstable during the mixing stage. Hence, we simply apply grid search of the config- urations of (a, ai), and report the best one in Figure 4.
17


Reviews:
[{"review_id": "kCLBBUpsCY", "reviewer": "Reviewer_XxHt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed an any-scale balanced proposal for discrete distribution. Motivated by recent findings that non-local proposals improve sampling efficient, the proposed method closes the gaps of the choice of weight function and the accuracy of first order gradient approximation. Comparison to several baseline method demonstrate the efficiency of the proposed method.", "summary_of_the_review": "The paper provides useful tricks for hyperparameter choice but lacks originality. Also, a very important baseline method that is highly related to the proposed method is missing (PVAG in Rhodes and Gutmann, 2022).", "strength_and_weaknesses": "Strength:\n\n1. The paper is clear and well written.\n2. non-local proposals are helpful for efficient sampling of discrete distributons and the paper introduces several techniques to deal with the associate problems.\n\n\nWeaknesses:\n\n1. The preconditioning (using $D$ instead of $\\sigma^2 I$) seems to be a useful trick. The same trick can be applied to other baseline method (for example, DLP can also use anisotropic diagonal as well), additional comparison to these variants would better illustrate the merits of the second order quadratic approximation.\n\n2. Lack of a comparison to the precondition version PAVG of Rhodes and Gutmann, 2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "P3BK1YxSifL", "reviewer": "Reviewer_XywG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper provides a non-local extension of the locally balanced samplers for improved sampling efficiency. The authors start by providing an analysis of existing approaches to improve the sampling efficiency of locally balanced samplers. They then introduce the any-scale balanced sampler in a step by step manner. Numerical results showing the effectiveness of the proposed algorithm conclude the paper.", "summary_of_the_review": "The paper proposes an any-scale balanced sampler, a non-local extension of locally balanced samplers. The proposed algorithm is clearly presented, numerical results that show the advantages of the proposed sampler are presented and interpreted. Limitations are clearly discussed. Good paper.", "strength_and_weaknesses": "The proposed algorithm is clearly explained, I personally appreciated the step by step manner in which the algorithm is presented. In terms of performances, the algorithm offers notable improvements over its locally balanced counterparts. Results are interpreted and the limitations of the algorithm are clearly stated.\n\nThe authors highlighted rather well themselves the limitations of the proposed algorithm. I would actually count this as a strength, not a weakness. Personally, I don't see any other weaknesses.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "FkZ8ZV4Sw_V", "reviewer": "Reviewer_Ce4Q", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The locally-balanced approach of Zanella has recently proven to be useful for MCMC sampling in discrete state spaces. The article proposes two main innovations when compared to recent articles building on the approach of Zanella:\n1. tuning of the \"balancing function\" $g(t) = t^{\\alpha}$ and length-scale $\\sigma$ of the kernel $K_\\sigma$. The tuning of the parameters $(\\alpha, \\sigma)$ is done quite naturally by attempting the maximise the \"Expected Squared Jumping Distance\" during an adaptation phase.\n2. take second order effect by choosing a global \"preconditioning\" matrix such that $f(y) \\approx f(x) + \\nabla f(x)^\\top (y-x) + \\frac12 \\langle (y-x), W (y-x) \\rangle$. The choice of the global matrix $W$ is done during an adaptation phase by minimising $\\|W(y-x) - (\\nabla f(y) - \\nabla f(x))\\|^2$, which is quite cheap computationally to implement.\n\nUsing a second order approximation of the log target leads to sampling from a Ising-like distribution. As has already been discussed in other article [1], this can be approached by using the standard Gaussian integral trick (i.e. Hubbard-Stratonovich transformation). In order to implement this approach, a diagonal matrix $D$ such that $W+D$ is positive definite needs to be chosen. The authors propose to choose $D$ so that $W+D$ as \"isotropic\" as possible, which is a sensible and standard strategy that appears to work well in this case.\n\n\n### References\n1. Rhodes B, Gutmann M. Enhanced gradient-based MCMC in discrete spaces. arXiv preprint arXiv:2208.00040. 2022 Jul 29.", "summary_of_the_review": "The article proposes an extension of the work of Zanella to take second order effect into account (i.e. preconditioning) when sampling from discrete distributions, which is very important in practical scenarios (i.e. preconditioning can lead to very large improvements in sampling efficiency). Even if the approach is relatively straightforward, it works well and the numerical simulations are convincing.", "strength_and_weaknesses": "## Strengths\nThe paper reads well and the proposed approached is very natural and convincing. Furthermore, the simulations are quite well executed.\n\n## Weakness\nI do not see any particular weakness.\n1. Indeed, the proposed methods is extremely natural after one has read the recent paper [1] but the article seems (?) to have been written at the same time.\n2. I am not really sure we need the discussion at the bottom of page 3 (eg. Equations 5/6/7) since this does not seem to add much to the text. But it certainly does not hurt.\n\n### Very Minor comments\n1. The authors seem to parametrize Gaussian distributions with mean and standard-deviation $N(\\mu, \\sigma)$ instead of mean and variance $N(\\mu, \\sigma^2)$, which is a bit awkward (is it common in some communities?).\n2. Should it be $K_\\sigma$ instead of $H_\\sigma$ in Equation 5. Also, $k_\\sigma$ is used of $K_\\sigma$ on page 2.\n\n### References\n1. Rhodes B, Gutmann M. Enhanced gradient-based MCMC in discrete spaces. arXiv preprint arXiv:2208.00040. 2022 Jul 29.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster