P ATH A UXILIARY P ROPOSAL FOR MCMC IN D ISCRETE S PACE
A BSTRACT
Energy-based Models (EBMs) offer a powerful approach for modeling discrete structure, but both inference and learning of EBM are hard as it involves sampling from discrete distributions.
Recent work shows Markov Chain Monte Carlo (MCMC) with the informed proposal is a powerful tool for such sampling.
However, an informed proposal only allows local updates as it requires evaluating all energy changes in the neighborhood.
In this work, we present a path auxiliary algorithm that uses a composition of local moves to efﬁciently explore large neighborhoods.
We also give a fast version of our algorithm that only queries the evaluation of energy function twice for each proposal via linearization of the energy function.
Empirically, we show that our path auxiliary algorithms considerably outperform other generic samplers on various discrete models for sampling, inference, and learning.
Our method can also be used to train deep EBMs for high dimensional discrete data.
I NTRODUCTION
Many real-world problems involve discrete structured data modeling, such as syntax trees for natural language processing(Tai et al., 2015), graphical model for molecules(Gilmer et al., 2017) etc.
A powerful approach for modelling the distribution over structured data is Energy Based Models(LeCun et al., 2006) (EBMs).
EBMs deﬁne the distribution with an unnormalized energy function, which allows great ﬂexibility to ﬁt the target distribution.
However, this ﬂexibility also results in the difﬁculties in inference and learning, as they require sampling from the EBM (Andrieu et al., 2003; Hinton, 2002), where the partition function is intractable in many cases.
Markov Chain Monte Carlo (MCMC) algorithms are one of the most widely used methodologies to sample from intractable distributions (Robert & Casella, 2013).
The efﬁciency for MCMC depends drastically on the proposal distribution.
For example, in continuous space, Metropolis-adjusted
Langevin algorithm (MALA) exploits the gradient of the target in single step walk and biases the proposal distribution toward high probability region (Rossky et al., 1978; Roberts & Rosenthal, 1998; Welling & Teh, 2011); Hamiltonian Monte Carlo (HMC) employs multi-step walk and explores the distribution more efﬁciently(Neal, 2004; Girolami & Calderhead, 2011; Hoffman et al., 2014).
These methods substantially improve the performance of the MCMC algorithm in theory and in practice.
However, their proposal distributions are derived as discretization of continuous diffusion process and it is still not clear how to appropriately extend such methods into discrete space(Zanella, 2020).
Recently, Zanella (2020) proposed a general framework called pointwise informed proposal (PIP) that shows promising results on directly sampling from discrete distributions.
PIP utilizes the energy change in the neighborhood of current state to propose a new state.
Following this work, Grathwohl et al.
(2021) propose a more efﬁcient sampler that uses Taylor series to estimate the energy change in the neighborhood.
However, both methods only focus on proposing from a small neighborhood, for example, a 1-Hamming ball.
This is due to the computational expense of evaluating the energy change or approximation error of Taylor series for a large neighborhood.
As a result, the samples in
Markov chain will have strong correlation or even be trapped at local optimum, which deteriorates the sampling efﬁciency.
In order to efﬁciently explore a larger neighborhood, we propose a path auxiliary sampler, which is an auxiliary sampler (Liang et al., 2011) that uses auxiliary path to propose new states.
In construction of the path, we employ a local proposal distribution to make a sequence of small movements to reach a new state in long distance.
A typical challenge for such multi-step proposal is that the accept ratio could decrease very fast when the number of steps increases.
In this work, we show that the accept ratio of our path auxiliary proposal is independent of the path and only determined by the property of the current state and the proposed state.
As a result, our algorithm is able to maintain a high accept ratio with multiple local steps.
We can provably show that the locally balancing functions are asymptotically optimal for our path auxiliary sampler.
We also introduce a scalable version of the algorithm that uses linearization and can be applied to smooth target distributions.
We empirically evaluate our methods in inference, sampling, and learning on various discrete EBMs.
We demonstrate that our methods signiﬁcantly improve the sampling efﬁciency on parity model, weighted permutation model, Ising model, Restricted Boltzmann Machine, and Factorized Hidden
Markov Model.
Our method can also learn competitive deep EBMs on discrete image data.
The code can be found at https://github.com/ha0ransun/Path-Auxiliary-Sampler.git .
B ACKGROUND
Energy Based Model : Let X be a ﬁnite state space.
An EBM deﬁnes an energy function f :
X → R with target distribution π ( x ) =
e − f ( x ) /Z , where Z = (cid:80)
z ∈X e − f
( z ) (LeCun et al., 2006;
Wainwright & Jordan, 2008; Du & Mordatch, 2019).
The unnormalized energy function provides great ﬂexibility to characterize complex distribution.
However, this can also make the partition function Z intractable to calculate and exacerbate the difﬁculties of learning and inference.
Metropolis-Hastings : MCMC is a commonly used methodology to sample from an intractable distribution.
Metropolis-Hastings (MH) is one of the most commonly used framework for
MCMC(Metropolis et al., 1953; Hastings, 1970).
Given the current state x , a proposal distribution Q ( x, · ) gives a new state y , then MH algorithm accept y with probability of min { 1 , π ( y ) Q ( y,x ) π ( x ) Q ( x,y ) } to guarantee the Markov chain is π -reversible.
The efﬁciency of MH algorithm highly depends on the selection of proposal distribution Q .
Peskun Ordering :
Peskun ordering provides a method to compare the efﬁciency for two MCMC algorithms (Peskun, 1973; Tierney, 1998).
Let P 1 , P 2 be π -reversible Markov transition kernels on
X such that P 1 ( x, y ) > cP 2 ( x, y ) for all x (cid:54) =
y for a ﬁxed c > 0 , then we say P 1 is c -times more efﬁcient than P 2 as the following holds: 1) var π ( h, P 1 ) ≤ 1 c var π
( h, P 2 )
+
− c c var π ( h ) , ∀ h : X → R ; 2) Gap ( P 1 ) ≥
c
Gap ( P 2 ) .
A smaller asymptotic variance var π ( h, P ) means a better estimation for the expectation of h and a larger spectral gap Gap ( P ) means a faster convergence of the Markov chain.
Pointwise Informed Proposals and Locally Balancing Function : A pointwise informed proposal (PIP) is a MH algorithm in discrete space (Zanella, 2020).
PIP uses proposal distribution
Q g ( x, y ) =
g ( π ( y ) π
( x ) )
I ( x, y ) /Z
g ( x ) , where the I ( x, y )
= 1 { y ∈N ( x ) } is the membership indicator w.r.t a symmetric neighborhood N ( · ) , and the normalizer Z g ( x )
= (cid:80) z ∈X g ( π ( z ) π ( x ) )
I ( x, z ) .
The weight function g : R + → R + determines the efﬁciency of PIP.
Zanella (2020) show that the family of locally balancing functions G =
{ g : R + → R + , g ( t ) =
tg ( 1 t ) , ∀ t > 0 } is asymptotically optimal for PIP w.r.t.
Peskun ordering.
Empirically, Zanella (2020) shows g ( t ) = √ t and g ( t ) =
t t + have the best performance.
Gibbs with Gradient GWG (Grathwohl et al., 2021) is a scalable version of PIP, where the target distribution π ( · ) is approximated via Taylor’s series.
Despite being powerful, PIP and GWG requires the calculate the weight and sample from the neighborhood.
Hence, only a small neighborhood, usually a 1-Hamming ball, is used in existing methods.
PATH AUXILIARY P ROPOSAL
Problems occur in point-wise informed proposal (PIP) when merely considering a small neighborhood, especially for distributions with many local optima.
For example, consider a parity distribution where the state space X
= { 0 , 1 } p , and the energy function f ( x ) =
U if the number of 1s in x is odd, otherwise f ( x ) = 0 .
When a 1-Hamming ball neighborhood is used, by symmetry, a PIP will have a uniform probability to propose a new state from neighborhood.
Then, the expected time to leave a low energy state is O ( e U ) , which could be very inefﬁcient when U is large.
When a r -Hamming ball larger neighborhood is used, a PIP can efﬁciently escape a low energy state.
However, the neighborhood will contain O ( n r ) states, which could be computationally prohibitive when n or r is large.
To address such problems with PIP, we propose a path auxiliary sampler.
Instead of directly sampling from a large neighborhood, path auxiliary sampler sequentially samples new state from a local proposal distribution Q 0 .
The computation cost at each step is still manageable as Q 0 still samples locally.
Then the composition of small movements forms a path that can lead to a new state that is distant from the current state.
The complexity of the sampling grows linearly instead of exponentially w.r.t.
r .
In this section, we ﬁrst present the framework of our algorithm.
Then, we discuss the choice of the weight function for path auxiliary sampler.
And ﬁnally, we introduce a fast version of path auxiliary sampler.
3. P ROPOSAL VIA AUXILIARY PATH
We ﬁrst deﬁne an auxiliary path.
Given the neighborhood function N , the set of auxiliary paths on
X is deﬁned as
[Equation]
To obtain the auxiliary path, we employ a PIP Q 0 ( · , · ) to make local movements.
We also sample path length from a prior α ( · ) to assure our chain is aperiodic.
The path auxiliary sampler is deﬁned as follows
Sample a path length
L from a prior α ( L ) .
Denote the current state x t = σ 0 , sample σ l ∼ Q 0 ( σ l − 1 , · ) for l = 1 , ..., L . 3.
Accept x t +1 =
σ L in probability
A ( x, σ, L ) =
min (cid:110) 1 , π ( σ L ) (cid:81)
Ll =1 Q 0 ( σ l ,σ l − 1 ) π ( σ 0 ) (cid:81)
Ll =1 Q 0 ( σ l − 1 ,σ l ) (cid:111) , else x t +1 = x t .
Theorem 1.
The path auxiliary proposal transition rule described above satisﬁes the detailed balance and induces a reversible Markov chain with π as its invariant distribution.
3. B ENEFIT FROM A UXILIARY P ATH
We ﬁrst analyze the performance of path auxiliary sampler on parity distribution mentioned above.
Assume we use a uniform prior α (1) =
α (2) = 12 .
If we sample L = 2 , then we will transit back to a new state with same energy as the current state.
As a result, we can always escape from a state in O (1) steps in expectation.
Having a path auxiliary sampler to sample from L -Hamming ball at each step can be more efﬁcient than performing L steps of MH sampling in a 1-Hamming ball.
To mathematically justify it, we compare their accept ratio. For a path ( σ, L ) :
[Equation]
We can notice that the accept ratio for path auxiliary sampler is the product of the probability for single-step method before the minimum operator.
As a result, an auxiliary path can have a high accept ratio as long as the product is large.
On the contrary, when performing L steps MH sampling, the transition can probably be blocked at some steps in the path with low accept ratios.
3. B ALANCED P ROPOSAL
Although the product in equation 2 allows the Markov chain to escape from a local optimum when some steps in the path have high accept ratio, it can also exponentially decrease the accept ratio w.r.t.
the path length L when every step has a low accept ratio.
Hence, it is important to select a good local proposal Q 0 .
In this section, we show that the locally balancing function Zanella (2020) is asymptotically optimal for path auxiliary sampler.
We ﬁrst deﬁne a sub-class of weight function named as ideal function:
Deﬁnition 1.
F is the set of ideal function.
A function f ∈ F if and only if following conditions holds: 1) f : R + → R + ; 2) f
= 1 ; 3)
f
( t ) is monotonic increasing;4)
f ( t ) f
( 1 t ) t [Equation]
The next theorem shows that locally balancing function G is asymptotically better than ideal function
F in Peskun Ordering.
Theorem 2.
Consider the state space in Cartesian products X = × ni
=1 X i , where each X i is a ﬁnite space with M elements, and the neighborhood is deﬁned as 1-Hamming ball.
Let d n be the maximum degree in conditional independence graph.
If 1) lim n →∞ d n n = 0 ; 2) the target distribution satisﬁes π ( y ) π ( x ) ≤ C < ∞ , ∀ y ∈ N ( x ) ; 3) the path length prior α is bounded by U .
Then for any f ∈ F , we have a function g ∈ G = { g ( · ) : g ( t )
=
tg ( 1 t ) } that is asymptotically more efﬁcient than f , which implies G is asymptotically optimal in F .
The idea to prove this theorem is to use ¯ f ( t ) :
= (cid:113) f
( t ) f
( 1 t ) t .
By deﬁnition, we have:
[Equation]
which means ¯ f ∈ G .
Then, the theorem is proved by showing the π -reversible Markov transition kernel P ¯ f ( x, y ) ≥ P f ( x, y ) asymptotically holds.
The proof is given in Sec A.
Although Deﬁnition 1 has some constraints, it already contains almost all natural choices of weight functions, such as f ( t )
=
2 t 1+ t , f ( t ) =
min { 1 , t } , f ( t ) =
max { 1 , t } , and f ( t ) =
t α , α ≥ 0 .
Hence, the asymptotic optimality for locally balanced function G in ideal function F strongly suggest that locally balanced function is a good choice for path auxiliary sampler.
More discussion for ideal function can be found in Sec A.
Property 1.
When using a locally balancing function g ∈ G as the weight function, the accept ratio for path auxiliary sampler can be written as:
[Equation]
The path independent form in Property 1 shows that if the energy function is smooth enough, the accept ratio in equation 4 will be high when y is close to x .
By selecting g ( t ) = √ t ∈ G , we name our auxiliary sampler (Liang et al., 2011) as path auxiliary sampler (PAS) and give the algorithm in Algorithm 1.
3. S CALABLE PATH AUXILIARY A LGORITHM
Although path auxiliary proposal reduces the complexity from exponential to linear, it still requires O ( nL ) evaluations of the probability π ( x ) in total.
In cases where the computational bottleneck is evaluating π ( x ) , this cost will be expensive.
Fortunately, most distributions of interest have differentiable energy functions, such as deep EBMs , and we can use the linearization as approximation (Grathwohl et al., 2021).
Given the current state σ 0 , the linearization is
[Equation]
we can estimate ˜ π σ 0 ( y ) /
˜ π
σ 0 ( x ) =
[Equation]
f
[Equation]
σ 0 ) ,y − x
[Equation]
, and deﬁne our proposal distribution ˜ Q g,σ 0 ( x, y )
=
g ( ˜ π 0 ( y ) ˜
π 0 ( x ) ) /
˜ Z g,σ 0 ( x ) , with the normalizer Z g,σ 0 ( x ) = (cid:80) z ∈N ( x ) g ( ˜ π 0 ( z ) ˜ π 0 ( x ) ) .
After sampling σ , we use the linearization ˜ f
σ L at σ L to compute the accept ratio w.r.t.
the auxiliary path
[Equation]
[Table]
[Table]
Denote σ 0 = x t , sample σ l ∼ Q g ( σ l − 1 , · ) or σ l ∼ ˜ Q g,σ 0 ( σ l − 1 , · ) for l = 1 , ..., L .
[Table]
in probability A ( x, σ, L ) [Table]
min {
[Table] g
[Table]
σ L ) }
[Table]
[Table]
ﬁnish
[Table]
By using approximation, we only need to evaluate π ( y ) , π ( x ) , ∇ f ( y ) , ∇ f ( x ) once in every proposal, which can signiﬁcantly reduce the computational cost. We name this algorithm as path auxiliary fast sampler (PAFS) (see Algorithm 1). The PAS framework also allows different approximations. For example, GWG (Grathwohl et al., 2021) can be seen as a special case where all indices are sampled from ˜ Q g,σ 0 ( σ 0 , · ) rather than ˜
Q g,σ 0 ( σ l , · ) . Besides, one way to further encourage long range movements is sampling sites without replacement, that’s to say, the auxiliary path σ does not modify a site more than once. However, such a sampler requires a decent choice of the path length L . When L is too large, the acceptance rate drops exponentially fast.
Hence, in this work, we focus on PAFS.
Similar to GWG (Grathwohl et al., 2021), the decrease of the proposal quality from PAFS can be bounded.
Theorem 3.
Assume the energy function f ( x ) is differentiable, ∇ f ( x ) is K -Lipschitz.
Consider we use weight function g ( t ) = √ t , path length prior α bounded by U , and 1-Hamming ball as neighborhood.
Denote P and ˜ P as the transition kernel by path auxiliary sampler and path auxiliary faster sampler, respectively, we have
[Equation]
Using the Peskun ordering, we know PAFS is at least e − K
[Equation]
U + [Equation] times as efﬁcient as PAS.
When
K is small, theorem 3 justify PAFS is a good approximation for PAS.
When K is large, a tighter bound needs more assumptions for the target distribution, e.g. conditional independence.
See the proof and discussion in Sec A.7.
R ELATED W ORKS
Informed proposal for Metropolis-Hastings (MH) algorithm has been extensively studied in the continuous space(Robert & Casella, 2013).
The most famous algorithms are Metropolis-adjusted
Langevein algorithm (MALA) (Roberts & Rosenthal, 1998) and Hamiltonian Monte Carlo (HMC).
MALA, HMC, and their variants (Girolami & Calderhead, 2011; Hoffman et al., 2014; Welling &
Teh, 2011; Titsias & Papaspiliopoulos, 2018) exploit the gradient of the target distribution to bias the proposal distribution towards high probability regions.
Although gradient-based methods having brought substantial improvements in continuous space, it is still unclear how to extend them to discrete space.
A number of methods try to map the discrete space to a continuous space using relaxation, apply informed methods in continuous space, and then map the new state back into discrete space(Zhang et al., 2012; Pakman & Paninski, 2013; Nishimura et al., 2017; Han & Liu, 2018;
Jaini et al., 2021) via Gaussian Integral Trick, uniform dequantization, or SurVAE ﬂow.
Such methods work in some scenarios, but a key challenge is the embedding of discrete space into continuous space can destroy the inherent discrete structure, resulting in highly multi-modal and irregular target distribution in continuous spaces.
Another group of methods directly work on discrete space.
Dai et al.
(2020) introduces the path as latent variable in the variational distribution for initializing PCD, but still relies on slow Gibbs sampling for improvement; Titsias & Yau (2017) augment the discrete space with auxiliary variable and perform Gibbs sampling in the augmented space based on informed proposal.
Zanella (2020) shows that a family of locally balancing function is asymptotically optimal for informed proposal.
Following Zanella (2020), Power & Goldman (2019) extends the framework to Markov jumping process, and Sansone (2021) parameterize locally balanced function to tune it via mutual information to select good weight function from the locally balanced class.
When the target distribution is smooth enough, Grathwohl et al.
(2021) employs a Taylor’s series to approximate the target distribution and further improve the sampling efﬁciency.
Though these informed proposal algorithms successfully show orders of magnitude improvements when using 1-Hamming ball as neighborhood, they are not able to explore a large neighborhood.
An extension of GWG (Grathwohl et al., 2021) can partially address this problem via sampling multiple dimensions to modify in one step.
However, such procedure can easily lead to backtracks thereby reducing its efﬁciency.
E XPERIMENTS 5. S AMPLERS UNDER C ONSIDERATION
In this section, we empirically evaluate the sampling efﬁciency of path auxiliary sampler (PAS) and path auxiliary fast sampler (PAFS).
We choose the path length prior α as a uniform distribution on { 1 , ..., 2 X − 1 } as suggested in Hoffman et al. and denote the corresponding samplers as PAS- X and PAFS- X .
We compare our methods with ﬁve types of baselines: random walk sampler (RW), Gibbs sampler (Gibbs), Hamming ball sampler (HB), locally balanced sampler (LB), Gibbs with gradient sampler (GWG).
RW is an informed proposal with g ( t ) =
1 that uniformly propose new state from neighborhood.
Gibbs partitions the dimension of a state x into two groups x u and x − u , then updates the state from conditional distribution p ( x u | x − u ) .
We denote it as Gibbs-X, where X refers to the dimension of x u .
HB is a two-stage Gibbs sampler on the extended state space ( x, x (cid:48) )
(Titsias & Yau, 2017).
We use HB-10-1, where 10 is the block size, and 1 is the hamming ball size.
LB (Zanella, 2020) is implemented as a single-step PIP with g ( t ) =
√ t .
GWG (Grathwohl et al., 2021) is a scalable version of LB, which draw the index to ﬂip via ﬁrst order Taylor’s series of the target distribution.
Although the original paper focuses on ﬂipping one site per step, it is possible for GWG to draw multiple indices to ﬂip in every step.
We denote the algorithm as GWG-X, where X indicates that the number of indices to modify per step is uniformly sampled from { 1 , ..., 2 X − 1 } .
The difference between PAFS and GWG is that GWG is more likely to sample the same index repeatedly thereby reducing the efﬁciency especially after mixing, e.g in ﬁgure 6.
5. I NFERENCE ON E NERGY B ASED M ODEL
Parity Model:
We ﬁrst demonstrate the beneﬁt of path auxiliary sampler versus single-step sampler in a inference task that estimates the mean µ of a parity distribution.
A parity distribution has state space X
= { 0 , 1 } p and energy function
[Equation]
i.e. a state has energy U if it has an odd number of 1s, otherwise 0 .
The neighborhood is deﬁned as 1-Hamming ball.
We run the simulation with p = 100 and U = { 1, 3, 5 } , and report the estimation
[Figure]
error E n := (cid:107) ˆ µ n − µ (cid:107) 2 , where ˆ µ n = (cid:80) ni
=1 X
i is the sample mean of the Markov chain at step n .
For each setting and method, we run 5 chains for 20,000 steps, and we plot the mean and standard deviation of the estimation error in ﬁgure 1.
We can observe that the efﬁciency of single step sampler [Figure] [Figure] decreases exponentially when U increases.
On the contrary, the path auxiliary sampler can always escape from the local optima in parity distribution and estimate the distribution mean efﬁciently.
Weighted Permutation Model: We consider an optimization task on weighted permutation model.
The state space X = S p is a symmetric group, i.e. for any ρ ∈ X , ρ is a permutation of 1 , 2 , ..., p .
The energy function is deﬁned as:
[Equation]
Following Zanella (2020), the weight { w ij } pi,j =1 are i.i.d.
sampled from Gaussian (0 , σ 2 ) .
The neighborhood is determined by one exchange of the permutation.
We use p = 100 , σ = 3 , 5 , 10 , and run 5 simulations for each conﬁgurations.
From ﬁgure 2, we can see that LB sampler is trapped at local optimum, especially when σ is large and the distribution is sharp.
On the contrary, PAS can always efﬁciently ﬁnd good solutions via auxiliary path.
5. S AMPLING ON E NERGY B ASED M ODEL
Lattice Ising Model: Consider the state space X
= {− 1 , 1 } V p , where ( V p , E p ) is the p × p square lattice graph.
For each x ∈ X , the energy function is deﬁned as:
[Equation]
α i ∈ R are bias terms representing the property of x i and λ is a global interaction term.
We ﬁrst compare PAS and PAFS in different path lengths.
For each length, we run 100 chains with 100,000 steps.
We report the burn-in stage as well as effective sample size (ESS) * in terms of both MCMC steps and energy function evaluation times in ﬁgure 3.
We can see that, 1) PAS has slight improvement in terms of energy function evaluations when we increase the path length; 2) when compared to PAS, PAFS has very similar proposal quality and signiﬁcantly better efﬁciency.
We also compare our sampler with other competitors.
For each sampler, we run 5 Markov chains with 1,000,000 steps and report the ESS to compare the proposal quality in ﬁgure 4.
We can see our sampler leads in both quality and efﬁciency.
More results are given in Sec B.1.
Factorized Hidden Markov Model: FHMM is a statistical model that use latent variables in X = { 0 , 1 } N × K to characterize time series data y ∈ R N .
Denote p ( x ) for hidden variables, and p ( y | x ) for likelihood:
[Equation]
Given data y , we use MCMC to sample x from the posterior p ( x | y ) and compare the mixing for different samplers.
We choose parameters N = 1000 , K = 10 , P ( x n, 1 = 1) = 0 .
[Figure]
P ( X n,k = x n,k − 1 ) = 0 . 85 , w ∈ Gaussian (0 , I K ) , b ∈ Gaussian (0 , 1) , and σ 2 = 0 .
We run each chain 5 times and report the mean and std for the energy (negative log joint density)
E (ˆ x ) = − log p (ˆ x ) p ( y | ˆ x ) and the reconstruction error (cid:107) y − w ˆ x (cid:107) 2 w.r.t.
the number of evaluations of energy function.
From ﬁgure 5 we can see that both GWG and PAFS gain signiﬁcant improvements in mixing by using a larger neighborhood.
More results are given in Sec B.
Restricted Boltzmann Machine: RBM is bipartite latent-variable model, whose energy function is deﬁned as:
[Equation]
(12) where { W, b, c } are parameters and x ∈ { 0 , 1 } D .
We follow Grathwohl et al.
(2021) to train a
RBM with 500 hidden units on the MNIST dataset using contrastive divergence(Hinton, 2002).
We generate samples via various MCMC samplers on the trained RBM.
Besides reporting ESS, we also estimate the maximum mean discrepancy (MMD) (Gretton et al., 2012) w.r.t.
a set of “ground truth” samples generated by structure known Block-Gibbs sampler.
Figure 7 shows PAFS has large ESS and can efﬁciently match the “ground truth” samples.
We also demonstrate an ablation study on different the approximations of the auxiliary path.
Speciﬁcally, we evaluate GWG, PAFS with different lengths.
For each sampler, we run 100 chains and report their ESS and average hop distance.
We can notice GWG suffers from a large path length
L , while our PAFS obtains robust
[Figure]
improvements for all L by employing a soft no-replacement sampling.
This result indicates that our PAFS provides an efﬁcient framework to sample from a large neighborhood.
More results and discussions, including a no-replacement sampler, are given in Sec B.
[Figure]
[Figure] [Table]
5. L EARNING ON E NERGY B ASED M ODEL
Learning an EBM is a challenge task.
Consider the target distribution is π and our energy function f θ is parameterized by θ .
The gradient for the likelihood of π θ ( x ) ∝ e − f θ ( x ) is:
[Equation]
The ﬁrst expectation can be estimated using the data from true distribution.
The second expectation requires samples from the current model, which is typically obtained via MCMC.
Hence, the success of training an EBM relies on efﬁcient MCMC algorithms to get an accurate estimation of the second expectation.
Ising model We ﬁrst learn an Ising model following the setting in Grathwohl et al.
(2021), where the energy function f ( x ) = θx T Jx .
Given a set of samples { x i } , the task is to learn an EBM via recovering the connectivity matrix J .
We generate a 25 × 25 2D cyclic lattice with θ = 0 .
25 and sample training data with a long-run Gibbs chain.
We train the models to maximize the likelihood of samples using persistent contrastive divergence (PCD)(Tieleman & Hinton, 2009) and report the root mean squared error (RMSE) between the inferred connectivity matrix ˆ J and the true matrix J .
For fair comparison, we also include GWG-eq, which runs GWG-1 with more steps such that it has the same time cost as our PAFS-5.
Figure 8 shows larger neighborhood can effectively help learn the
Ising model.
GWG-5 having similar performance as PAFS-5 as, using PCD, the learning efﬁciency is mainly determined by the mixing of the Markov chain.
Deep EBM We evaluate path auxiliary sampler by learning a deep EBM.
The experiment follows the setting in (Grathwohl et al., 2021).
We train deep EBMs paramterized by Residual Networks(He et al., 2016) on small binary image datasets using PCD(Tieleman & Hinton, 2009) with a replay buffer(Du & Mordatch, 2019).
We compare our methods with Variational Autoencoders Kingma &
Welling (2013), GWG (Grathwohl et al., 2021), RBM and a Deep Belief Network(Hinton, 2009).
We estimate the likelihoods using Annealed Importance Sampling(Neal, 2001) and report the results in table.1.
The results for VAE is taken from Tomczak & Welling (2018), RBM and DBN are taken from Burda et al.
(2015), GWG is taken from Grathwohl et al.
We can see that our approach improves the log-likelihoods for deep EBMs on all datasets.
D ISCUSSION AND C ONCLUSION
The problem for the selection of the path length remains open.
In this work, we propose a soft no-replacement method PAFS.
PAFS is not sensitive to the path length larger than the optimal length, hence has robust proposal quality for different choice of path length. We also considered no-backtrack samplers that ﬂip sites without replacement.
Intuitively, such auxiliary path encourages to propose states in larger distance.
However, for no-backtrack samplers, a too large path length is harmful for the proposal quality and signiﬁcant reduce the efﬁciency of the sampler.
As a result, ﬁnding a principle way to decide the path length is very important.
In continuous space, the optimal step size can be characterized via acceptance rate (Gelman et al., 1997; Roberts & Rosenthal, 1998; 2001; Beskos et al., 2013).
PAS is the gradient based sampler in discrete space, we believe it is possible to derive its optimal path length in a similar manner as continuous case.
We will investigate it in our future work.
In summary, informed proposal has shown good results for inference, sampling, and learning EBMs in discrete spaces.
Our path auxiliary sampler provides an approach allowing informed proposal to efﬁciently explore large neighborhoods.
We believe there is considerable room for future works to improve the sampling methods in discrete space.
Acknowledgement : This research was supported in part by the Defense Advanced Research
Projects Agency (DARPA) under Contract FA8750-20-C-0542 (Systemic Generative Engineering).
The views, opinions, and/or ﬁndings expressed are those of the author(s) and should not be interpreted as representing the ofﬁcial views or policies of the Department of Defense or the U.S.
Government.
A P ROOFS A. P ROOF FOR T HEOREM
Our path auxiliary sampler is a proposal distribution augmented auxiliary sampler.
Before diving into the proof, we ﬁrst introduce some background knowledge for auxiliary sampler, based on the results from page 86 in Liang et al.
Given a proposal distribution T ( y | x ) , we can augment it by an auxiliary variable u , such that T ( y | x ) =
(cid:82) T 1 ( u | x ) T 2 ( y | x, u ) du .
Using this proposal distribution, we can deﬁne a marginal sampler and an auxiliary sampler.
For a marginal sampler, the accept rate is
[Equation]
and for y (cid:54) = x , its transition kernel is:
[Equation]
For an auxiliary sampler, given x , we ﬁrst sample auxiliary u ∼ T 1 ( u | x ) , then we sample the new state y ∼ T 2 ( y | x, u ) .
The corresponding accept rate is:
[Equation]
Since the proposal and the accept rate depend on the auxiliary u , we need to integrate the auxiliary variable u to obtain the transition from x to y :
[Equation]
Though a marginal sampler is Peskun better than an auxiliary sampler (Titsias & Papaspiliopoulos, 2018), a marginal sampler is intractable in most of scenarios, as it requires to integrate over auxiliary variables u in both proposal and accept rate calculation.
Hence, our sampler is implemented in an auxiliary fashion.
Going back our theorem
[Equation]
the auxiliary variable u is the auxiliary path ( σ, L ) .
Condition on L , we have:
[Equation]
Given path length
L
[Equation]
with a little abuse of notation
[Equation]
the transition kernel for our path auxiliary sampler K ( x, y | L ) is:
[Equation]
[Equation]
When σ L (cid:54) =
y , the accept rate above is not well-deﬁned.
To avoid making extra deﬁnition
[Equation]
we absorb T 2 , the indicator, as a constraint of the domain in the integration.
Then, we have:
[Equation]
With these knowledge prepared
[Equation]
we begin to proof our theorem
[Equation]
Proof.
Denote K ( x, y ) as the probability that x transit to a different state y , then we have:
[Equation]
The key idea for the proof is that ( σ, L ) is symmetric w.r.t.
its two ends, when σ is a path from x to y , σ is also a path from y to σ .
Hence, we are able to exchange the orientation of the path in equation 27.
A. P ROOF FOR L EMMA
Lemma 1.
X is a ﬁnite state space with distribution π and a neighborhood function N such that every state has equal number of states.
F is a class of weight functions, such that for any f ∈ F , 1) f : R + → R + , 2) f (1) = 1 , 3) f
( t ) is monotonically increasing, 4) f ( t ) f
( 1 t ) t ≤ 1 , ∀ t ≤ 1 .
∀ f ∈ F , deﬁne ¯ f ( t ) =
(cid:113) f
( t ) f
( 1 t ) t , we have
[Equation]
Proof.
Since X is ﬁnite, we can always ﬁnd state x 1 , x 2 having highest and lowest probability, respectively.
Then, we have:
[Equation]
A. D ISCUSSION FOR I DEAL F UNCTION
We deﬁne ideal function as a stepping stone to show the advantage of locally balanced function.
We name it ”ideal” as the its properties are ideal assumption for a weight function in PIP.
Conditions 1) f : R + → R + is a natural requirement for weight function.
Condition 2) f (1) = 1 can be easily realized by substituting f ( t ) by f ( t ) /f
(1) , as PIP only depends on the ratio of weights.
condition 3) f ( t ) is monotonically increasing, and condition 4) f ( t ) f
( 1 t ) t ≤ 1 , ∀ t ≤ 1 , though, are technical requirements, it is reasonable for weight functions.
Monotonic increasing in 3) indicates the proposal distribution match the target distribution, where a point z has a higher probability in target distribution should also has a higher probability in proposal distribution.
Condition 4) is weaker than the following condition:
f ( t ) ≤ t, ∀ t ≥ 1 , hence condition 4) is easier to satisfy.
Currently deﬁnition of ideal function class already contains most of commonly used weight function, such as f ( t ) = 2 t 1+ t , f ( t ) =
max { 1 , t } , f ( t ) =
min { 1 , t } , and f ( t ) =
t α , where α ≥ 0 .
Also, we can following properties
Property 2.
The logarithm of ideal function is closed in convex combination.
That’s to say, ∀ f 1 , ..., f n ∈ F , and ∀ λ 1 , ..., λ n , s.t λ i ≥ 0 and (cid:80) ni
=
λ i = 1 , we have
[Equation]
Proof.
Let f i
and λ i be deﬁned as above.
Obviously, F is positive and F (1)
= 1 , hence satisﬁes condition 1) and 2).
For condition 3), we have:
[Equation]
[Equation]
Hence, ∀ t ≤ 1 , we have:
[Equation]
which implies F ( t ) ∈ F .
Property 3.
Ideal function is a superset for normalized monotonically increasing locally balanced function G I := { g ∈ G : g (1) = 1 , g is monotonically increasing } .
Proof.
∀ g ∈ G
I , condition 1), 2) and 3) are obviously satisﬁed.
For condition 4), since g ( t ) is locally balanced, we have:
[Equation]
implies g ∈ F .
Considering the ideal function cover such a large number of functions, we believe the asymptotic optimality over ideal function strongly suggests locally balanced function is a good choice for path auxiliary sampler.
A. P ROOF F OR L EMMA
The (undirected) conditional independence graph corresponds to Markov random ﬁelds:
Deﬁnition 2.
The conditional independence graph of X is the undirected graph G =
( K, E ) ,
where
K = { 1 , ..., k } and ( i, j ) / ∈ E
if and only if X i ⊥
X
j | X K \{ i,j } .
More details can be found in p.60 Whittaker (2009).
Lemma 2.
Consider the state space in Cartesian products X = × ni
=1 X i , where each X i is a ﬁnite space with M elements, and the neighborhood is deﬁned as 1-Hamming ball.
Let d n be the maximum degree in conditional independence graph.
Deﬁne c ( n ) g
[Equation]
the target distribution satisﬁes π ( y ) π ( x )
[Equation]
C [Equation] ∞ , ∀ y ∈ N ( x ) , then we have:
[Equation]
To simplify the notation, we assume M = 2 in this proof.
It is straightforward to extend the proof to any ﬁnite M .
Proof.
On one side, by deﬁnition, we can easily see that c ( n ) g ≥ 1 .
On the other side, given x , we use index i represent that we select y ∈ N ( x ) by ﬂipping index i for x .
We denote g x ( j ) =
g ( π ( z )) /g
( π ( x )) where z is obtained by ﬂipping index j for x , and we denote g i ( j )
=
g ( π ( z )) /g
( π ( y )) , where z is obtained by ﬂipping index j for y .
We also denote B ( i )
Then
[Equation]
we can write
[Equation]
If B ( i ) is not empty
[Equation]
the ﬁrst term in equation 40 equals to 1 .
If B ( i ) is empty, the ﬁrst term in equation 40 does not exist, hence it can still be bounded by [Equation]
Remark :
For a ﬁxed ideal function g ∈ F , Lemma 2 shows c ( n ) converges to 1 at a rate 1+ O ( d n n )
A. P ROOF FOR T HEOREM
We prove the theorem by using Lemma 1 and Lemma 2.
Proof.
We ﬁrst show that, for all g ∈ F , the transition probability:
[Equation]
where c g is deﬁned in Lemma 2.
We temporarily ignore the superscription and will add it back at the end of the proof.
Consider a path ( σ, L ) [Equation] we denote t j,k = π ( σ k ) π ( σ j ) .
Then the probability that x = σtransits to y =
σ L is:
[Equation]
By deﬁnition of c g
[Equation]
we have:
[Equation]
Use the property of ¯ g
[Equation]
we have
[Equation]
Combining equation 47, equation 48, equation 49, and Lemma 1 we have:
[Equation]
The inequality holds for arbitrary auxiliary path ( σ, L ) and the path length
L ≤ U , hence we prove the ﬁrst step.
Then ∀ g ∈ F , we use the estimation of c ( n ) g in Lemma 2, we have
[Equation] [Equation]
where
[Equation]
This indicates ¯ g is asymptotically better than g and proves the theorem.
Remark :
In the proof, equation 54 shows the convergence rate depends on g , this factor prevents a uniform convergence rate for the ideal function class F .
We can obtain a uniform rate when we constraint in smaller function class, for example, if we restrict to F M := { f ∈ F : f ( C ) /f ( 1 C ) ≤
M } , the convergence rate is 1 − O ( d n n ) .
A. P ROOF FOR P ROPERTY
Proof.
[Equation]
A. P ROOF FOR T HEOREM
The idea to proof theorem 3 is similar to Grathwohl et al.
(2021), which uses the Lipschitz condition to bound the estimation error.
The different is, in PAFS, we use the linearization at each state σ l to propose the next state σ l +1 , hence to need to accumulate the error at each step to obtain our ﬁnal result.
Proof.
For any path ( σ, L ) , we ﬁrst bound the estimation error ˜ f
( z ) − ˜ f 0 ( σ l − 1 ) , for z ∈ N ( σ l − 1 ) .
Since f is M -smooth
[Equation]
we have:
[Equation]
Similarly, we also have:
[Equation]
With this estimation
[Equation]
we have:
[Equation]
Similarly, we can also obtain
[Equation]
Now, consider the transition kernel
[Equation]
Remark : Though the theorem proved above is very loose, it provides a framework to prove approximation bound for path auxiliary sampler via single step approximation in equation 66.
More sophisticated bounds can be obtained by improving the single step estimation in equation 66 via the property of the conditional independence.
Speciﬁcally, a local modifying of the state usually does not inﬂuence most of the indices, which for example can be characterized by the Markov boundary.
Then, with high probability, the auxiliary path can avoid manipulating correlated indices and hence result in the approximation error in equation 66 be controlled by a constant, rather than e − Kl .
For example, in Ising model, an auxiliary path has high probability to avoid manipulating two adjacent nodes, and in this case, the linearization is accurate and we loose nothing in approximation.
In deep
EBMs, the conditional independence is usually not available, but we can still expect low correlation between most of dimensions.
Such analysis relies on more accurate assumption of conditional independence and we will leave it to our future work.
[Figure]
[Figure]
B D ETAILS FOR E XPERIMENTS B. S AMPLING ON I SING
Following Zanella (2020), we set the interaction term λ = 1 , and we set α i = µ + Z i if i in the center of the square lattice, and α i = − µ + Z i otherwise, where µ = 2 and Z i ∼ Unif ( − 3 , 3) .
We generate the lattice Ising model in four sizes: p = 50 , 100 , 150 , 200 .
We ﬁrst run PAS and PAFS with path lengh L = 1 , 2 , 3 , 5 , 10 on Ising model with size p = 50 , 100 , 150 , 200 .
For each path length and each model size, we run 100 chains 100,000 steps and compute the ESS using the last 50,000 steps.
The results for p = 200 is given in ﬁgure 3.
We give the results for p = 50 , 100 ,in the following.
From ﬁgure 9, ﬁgure 10, ﬁgure 11, and ﬁgure 3, we can see that the difference between PAS and PAFS in terms of MCMC steps is decreasing when the model size is increasing.
The reason is that when the model becomes larger, most of the nodes are conditionally independent.
When the auxiliary path does not involve two adjacent nodes, the approximation has zero error.
This observation provides a possibility to further improve the approximation bound in Theorem 3.
We will study this in our future work.
For comparing with other samplers, we report the results for each problem size p = 50 , 100 , 150 ,in ﬁgure 12, ﬁgure 13, ﬁgure 14, ﬁgure 4.
Our path auxiliary samplers substantially outperform other competitors.
B. S AMPLING ON FHMM
We compare PAS and PAFS with different path length on FHMM.
We choose parameters K = 10 , P ( x n, 1 = 1) = 0 .
05 , P ( X n,k = x n,k − 1 ) = 0 .
85 , w ∈ Gaussian (0 , I K ) , b ∈ Gaussian (0 , 1) , and σ 2 = 0 .
We simulate PAS, PAFS with path length
L = 1 , 2 , 3 , 5 , 10 on FHMM with
N = 500 , 1000 , 1500 , 2000 .
For each conﬁguration, we run 100 chains and report the burn in period, as well as the ESS in terms of MCMC steps and energy function evaluations.
[Figure] [Figure]
[Figure]
We can observe that: 1) in terms of MCMC steps, PAFS has very similar burn in time in energy compared to PAS, while the the gap in ESS is larger than the gap in Ising.
The reason is FHMM is not as close to linear model as Ising and PAFS has larger estimation error.
2) in terms of energy function evaluations, PAFS still leads the performance as other models.
PAS only obtains faster mixing
[Figure]
and no improvements in ESS.
The results show that whether PAS can help depends on the property of the target distribution.
B. S AMPLING ON RBM
We compare GWG, PAFS
[Figure]
and NB with different path lengths on RBM trained on MNIST dataset.
GWG (Grathwohl et al., 2021) samples all indices to manipulate based on the linearization at starting state of the path σ 0 , and we classify it as sampling with replacement.
Our PAFS samples the indices based on the linearization at current state σ l and we classify it as sampling with soft no replacement.
Speciﬁcally, consider a binary case, when we ﬂip index i at state σ l , then in the remaining of the path, the energy change for index i will be reversed.
Hence, if we ﬂip an index and reduce the system’s energy, we will have small probability ﬂip it back.
A third type sampler is No Backtrack (NB) sampler.
It is also a path auxiliary sampler and it rules out the indices have been selected to assure
[Figure]
[Figure] [Figure]
it is sampling with hard no replacement.
Speciﬁcally, NB ﬁrst calculate the probability for indices based on ∇ f ( σ 0 ) as GWG.
Then, once selecting an index i at σ l , NB manipulates the probability for selecting index i in the remaining path.
For each sampler with path length L = 1 , ..., 20 , we runchains and report the MMD (Gretton et al., 2012) w.r.t.
a set of “ground truth” samples generated by structure known Block-Gibbs sampler.
We can see that, when increasing the path length, PAFS has steady fast mixing, GWG decreases the efﬁciency in mixing, and NB fails in mixing.
For this reason, we only compare the statistics for GWG and PAFS in our main text in ﬁgure 6.
B. R UNNING T IME IN L EARNING I SING
We follow the experiments in Grathwohl et al.
GWG runs faster 1.4x than our PAFS-5 as the energy function computation is cheap in this experiment.
To obtain a fair comparison, we add
GWG-equivalent-length (GWG-eq), which we allow GWG-1 to run 1.4x more sampling steps such that its running time is the same as our PAFS.
When our PAFS runs 5, 10, 25, 50 steps, GWG-eq runs 6, 13 34, 69 steps.
Hence, when we draw ﬁgure 8, the point for GWG-eq at step 5, 10, 25,using the RMSE obtained at setp 6, 13, 34, 69.
B. D ETAILS IN L EARNING D EEP EBM S
DataSet
Static MNIST
Omniglot
Caltech
PAFS0.
PAFS-
[Table]
[Figure] [Figure]
[Figure]
We use the same setting as Grathwohl et al.
(2021), including the batch size, number of iterations, the PCD hyper parameters, etc.
In principle, both PAFS and GWG requires two energy function evaluations per MCMC step, despite that PAFS explores much larger neighborhood space via path auxiliary.
As the energy function evaluation and gradient calculation usually dominates the computation, we expect these two should have similar runtime.
In practice
[Figure]
we report the average running time per batch of 100 chains in table 2.
We can see for deep EBMs, our PAFS with different path lengths run slightly faster than GWG.
Given that both of the methods implemented using pytorch , the minor speed up might be due to the implementation issue.
So overall we think it is fair to say the two methods run equally fast in most cases.
For this reason, we use the same number of steps when we train the deep EBMs.
As is reported in
Grathwohl et al.
, we the same number of steps (40) to train the binary deep EBMs.
For our method, we also tune the expected path length in { 3 , 5 , 7 } , and report the result with the best validation likelihood.