Gradient-Guided Importance Sampling for Learning Discrete Energy-Based Models

1 INTRODUCTION . Energy-Based models ( EBMs ) , also known as unnormalized probabilistic models , model distributions by associating unnormalized probability densities . Such methods have been developed for decades ( Hopfield , 1982 ; Ackley et al. , 1985 ; Cipra , 1987 ; Dayan et al. , 1995 ; Zhu et al. , 1998 ; Hinton , 2012 ) and are unified as energy-based models ( EBMs ) ( LeCun et al. , 2006 ) in the machine learning community . EBMs have great simplicity and flexibility since energy functions are not required to integrate or sum to one , thus enabling the usage of various energy functions . In practice , given different data types , we can parameterize the energy function with different neural networks as needed , such as multi-layer perceptrons ( MLPs ) , convolutional neural networks ( CNNs ) ( LeCun et al. , 1998 ) , and graph neural networks ( GNNs ) ( Gori et al. , 2005 ; Scarselli et al. , 2008 ) . Recently , EBMs have been drawing increasing attention and are demonstrated to be effective in various domains , including images ( Ngiam et al. , 2011 ; Xie et al. , 2016 ; Du & Mordatch , 2019 ) , videos ( Xie et al. , 2017 ) , texts ( Deng et al. , 2020 ) , 3D objects ( Xie et al. , 2018 ) , molecules ( Liu et al. , 2021 ; Hataya et al. , 2021 ) , and proteins ( Du et al. , 2020b ) . Nonetheless , learning ( a.k.a. , training ) EBMs is known to be challenging since we can not compute the exact likelihood due to the intractable normalization constant . As reviewed in Section 4 , many approaches have been proposed to learn EBMs , such as maximum likelihood training with MCMC sampling ( Hinton , 2002 ) and score matching ( Hyvärinen & Dayan , 2005 ) . However , most recent advanced methods can not be applied to discrete data directly since they usually leverage gradients over the continuous data space . For example , for many methods based on maximum likelihood training with MCMC sampling , they use the gradient w.r.t . the data space to update samples in each MCMC step . If we update discrete samples using such gradient , the resulting samples are usually invalid in the discrete space . Notably , discrete data is common in our real world , such as texts , graphs , and genome sequences . Therefore , learning EBMs on discrete data remains to be challenging and in demand . Ratio matching ( Hyvärinen , 2007 ; Lyu , 2009 ) is proposed to learn discrete EBMs by matching ratios of probabilities between the data distribution and the model distribution , as detailed in Section 2.2 . However , as analyzed in Section 3.1 , it requires expensive computations and excessive memory usages , which is infeasible if the data is high-dimensional . In this work , we propose to use the gradient of the energy function w.r.t . the discrete data space to guide the importance sampling for estimating the original ratio matching objective . More specifically , we utilize such gradient to approximately construct the provable optimal proposal distribution for importance sampling . Thus , the proposed approach is termed as ratio matching with gradient-guided importance sampling ( RMwGGIS ) . Our RMwGGIS can significantly overcome the limitations of ratio matching . In addition , it is demonstrated to be more effective than the original ratio matching because it can be optimized better in practice . Experimental results on synthetic discrete data , graph generation , and Ising model training demonstrate that our RMwGGIS significantly alleviates the limitations of ratio matching , achieves better performance with obvious margins , and has the ability of scaling to high-dimensional relevant problems . 2 PRELIMINARIES . 2.1 ENERGY-BASED MODELS . Let x be a data point and Eθ ( x ) ∈ R be the corresponding energy , where θ represents the learnable parameters of the parameterized energy function Eθ ( · ) . The probability density function of the model distribution is given as pθ ( x ) = e−Eθ ( x ) Zθ ∝ e−Eθ ( x ) , ( 1 ) where Zθ ∈ R is the normalization constant ( a.k.a. , partition function ) . To be specific , Zθ =∫ e−Eθ ( x ) dx if x is in the continuous space and Zθ = ∑ e−Eθ ( x ) for discrete data . Hence , computing Zθ is usually infeasible due to the intractable integral or summation . Note that Zθ is a variable depending on θ but a constant w.r.t . x . 2.2 RATIO MATCHING . Ratio matching ( Hyvärinen , 2007 ) is developed for learning EBMs on discrete data by matching ratios of probabilities between the data distribution and the model distribution . Note that we focus on d-dimensional binary discrete data x ∈ { 0 , 1 } d in this work . Specifically , ratio matching considers the ratio of p ( x ) and p ( x−i ) , where x−i = ( x1 , x2 , · · · , x̄i , · · · , xd ) denotes a point in the data space obtained by flipping the i-th dimension of x . The key idea is to force the ratios pθ ( x ) pθ ( x−i ) defined by the model distribution pθ to be as close as possible to the ratios pD ( x ) pD ( x−i ) given by the data distribution pD . The benefit of considering ratios of probabilities is that they do not involve the intractable normalization constant Zθ since pθ ( x ) pθ ( x−i ) = e−Eθ ( x ) e−Eθ ( x−i ) = eEθ ( x−i ) −Eθ ( x ) according to Eq . ( 1 ) . To achieve the match between ratios , Hyvärinen ( 2007 ) proposes to minimize the objective function JRM ( θ ) = Ex∼pD ( x ) d∑ i=1 [ g ( pD ( x ) pD ( x−i ) ) − g ( pθ ( x ) pθ ( x−i ) ) ] 2 + [ g ( pD ( x−i ) pD ( x ) ) − g ( pθ ( x−i ) pθ ( x ) ) ] 2 . ( 2 ) The sum of two square distances with the role of x and x−i switched is specifically designed since it is essential for the following simplification . In addition , the function g ( u ) = 11+u is also carefully chosen in order to obtain the subsequent simplification . To compute the objective defined in Eq . ( 2 ) , it is known that the expectation over data distribution ( i.e. , Ex∼pD ( x ) ) can be unbiasedly estimated by the empirical mean of samples x ∼ pD ( x ) . However , to obtain the ratios between pD ( x ) and pD ( x−i ) in Eq . ( 2 ) , the exact data distribution is required to be known , which is usually impossible . Fortunately , thanks to the above carefully designed objective , Hyvärinen ( 2007 ) demostrates that the objective function in Eq . ( 2 ) is equivalent to the following simplified version JRM ( θ ) = Ex∼pD ( x ) d∑ i=1 [ g ( pθ ( x ) pθ ( x−i ) ) ] 2 = Ex∼pD ( x ) d∑ i=1 [ g ( eEθ ( x−i ) −Eθ ( x ) ) ] 2 , ( 3 ) which does not require the data distribution to be known and can be easily computed by evaluating the energy of x and x−i . It is proved that the estimator given by Eq . ( 3 ) is consistent ( Hyvärinen , 2007 ) . That means if it is minimized perfectly , the obtained model distribution will capture the data distribution exactly . Further , Lyu ( 2009 ) shows that the objective function of ratio matching can be reduced to JRM ( θ ) = Ex∼pD ( x ) d∑ i=1 [ pθ ( x−i ) pθ ( x ) ] 2 = Ex∼pD ( x ) d∑ i=1 [ eEθ ( x ) −Eθ ( x−i ) ] 2 . ( 4 ) It is obvious that Eq . ( 3 ) and Eq . ( 4 ) agree with each other since function g ( · ) decreases monotonically in [ 0 , +∞ ) , which aligns with the value range of probability ratios pθ ( x ) pθ ( x−i ) . Intuitively , the objective function of ratio matching , as formulated in Eq . ( 4 ) , can push down the energy of the training sample x and push up the energies of other data points obtained by flipping one dimension of x . Thus , this objective faithfully expect that each training sample x has higher probability than its local neighboring points that are hamming distance 1 from x . 3 THE PROPOSED METHOD . In this section , we analyze the limitations of the ratio matching method from the perspective of computational time and memory usage . Then , we describe our proposed method , ratio matching with gradient-guided importance sampling ( RMwGGIS ) , which utilizes the gradient of the energy function w.r.t . the discrete input x to guide the importance sampling for estimating the original ratio matching objective . Our approach can alleviate the limitations significantly and is shown to be more effective in practice . 3.1 ANALYSIS OF RATIO MATCHING . Time-intensive computations . According to Eq . ( 4 ) , for a given training sample x , we have to compute the energies for all x−i , where i = 1 , · · · , d. In other words , we have O ( d ) evaluations of the energy function for each training sample . This is computationally intensive , especially when the data dimension d is large . Excessive memory usages . Besides the expensive computation , the memory usage of ratio matching is another limitation that can not be ignored , especially when we learn the energy function using modern GPUs with limited memory . As shown in Eq . ( 4 ) , the objective function consists of d terms for each training sample . When we do backpropagation , computing the gradient of the objective function w.r.t . the learnable parameters of the energy function is required . Therefore , in order to compute such gradient , we have to store the whole computational graph and the intermediate tensors for all of the d terms , thereby leading to excessive memory usages especially if the data dimension d is large . Hence , it is challenging to learn EBMs with ratio matching on modern devices , such as GPUs , for high-dimensional discrete data . 3.2 RATIO MATCHING WITH GRADIENT-GUIDED IMPORTANCE SAMPLING . The key idea of our approach is to use the well-known importance sampling technique to reduce the variance of estimating JRM ( θ ) with Monte Carlo method . The most critical and challenging part of using the importance sampling technique is choosing a good proposal distribution . In this work , we propose to utilize the gradient of the energy function w.r.t . the discrete input x to approximately construct the optimal proposal distribution for importance sampling . We describe the details of our method below . The objective for each sample x , defined by Eq . ( 4 ) , can be reformulated as JRM ( θ , x ) = d d∑ i=1 1 d [ eEθ ( x ) −Eθ ( x−i ) ] 2 = dEx−i∼m ( x−i ) [ eEθ ( x ) −Eθ ( x−i ) ] 2 , ( 5 ) where m ( x−i ) = 1d for i = 1 , · · · , d is a discrete distribution . Thus , the objective of ratio matching for each sample x can be viewed as the expectation of [ eEθ ( x ) −Eθ ( x−i ) ] 2 over the discrete distribu- tion m ( x−i ) . In the original ratio matching method , as described in Section 2.2 , we compute such expectation exactly by considering all possible x−i , leading to expensive computations and excessive memory usages as analyzed in Section 3.1 . Naturally , we can estimate the desired expectation with Monte Carlo method by considering fewer terms sampled based on m ( x−i ) . However , such estimation usually has a high variance , and is empirically verified to be ineffective by our experiments in Section 5 . Further , we can apply the importance sampling method to reduce the variance of Monte Carlo estimation . Intuitively , certain values have more impact on the expectation than others . Hence , the estimator variance can be reduced if such important values are sampled more frequently than others . To be specific , instead of sampling based on the distribution m ( x−i ) , importance sampling aims to sample from another distribution n ( x−i ) , namely , proposal distribution . Formally , JRM ( θ , x ) = dEx−i∼m ( x−i ) [ eEθ ( x ) −Eθ ( x−i ) ] 2 = dEx−i∼n ( x−i ) m ( x−i ) [ eEθ ( x ) −Eθ ( x−i ) ] 2 n ( x−i ) . ( 6 ) The detailed derivation of Eq . ( 6 ) is given in Appendix A . Afterwards , we can apply Monte Carlo estimation based on the proposal distribution n ( x−i ) . Specifically , we sample s terms , denoted as x ( 1 ) −i , · · · , x ( s ) −i , according to the proposal distribution n ( x−i ) . Note that s is usually chosen to be much smaller than d. Then the Monte Carlo estimation for JRM ( θ , x ) is computed based on these s terms . Formally , JRM ( θ , x ) ̂n = d 1 s s∑ t=1 m ( x ( t ) −i ) [ eEθ ( x ) −Eθ ( x ( t ) −i ) ] 2 n ( x ( t ) −i ) , x ( t ) −i ∼ n ( x−i ) . ( 7 ) It is known that the estimator obtained by Monte Carlo estimation with importance sampling is an unbiased estimator , as the conventional Monte Carlo estimator . The key point of importance sampling is to choose an appropriate proposal distribution n ( x−i ) , which determines the variance of the corresponding estimator . The optimal proposal distribution n∗ ( x−i ) , which yields the minimum variance , is given by the following theorem . Theorem 1 . Let n∗ ( x−i ) = [ eEθ ( x ) −Eθ ( x−i ) ] 2∑d k=1 [ e Eθ ( x ) −Eθ ( x−k ) ] 2 be a discrete distribution on x−i , where i = 1 , · · · , d. Then for any discrete distribution n ( x−i ) on x−i , where i = 1 , · · · , d , we have V ar ( JRM ( θ , x ) ̂n∗ ) ≤ V ar ( JRM ( θ , x ) ̂n ) . Proof . The proof of Theorem 1 is included in Appendix B . To construct the exact optimal proposal distribution n∗ ( x−i ) given by Theorem 1 , we still have to evaluate the energies of all x−i , where i = 1 , · · · , d. To avoid such complexity , we propose to leverage the gradient of the energy function w.r.t . the discrete input x to approximately construct the optimal proposal distribution . Our approach only needs O ( 1 ) evaluations of the energy function to construct the proposal distribution . It is observed by Grathwohl et al . ( 2021 ) that many discrete distributions are implemented as continuous and differentiable functions , although they are evaluated only in discrete domains . Grathwohl et al . ( 2021 ) further proposes a scalable sampling method for discrete distributions by utilizing the gradients of the underlying continuous functions w.r.t . the discrete input . In this study , we extend this idea to improve ratio matching . More specifically , in our case , even though our input x is discrete , our parameterized energy function Eθ ( · ) , such as a neural network , is usually continuous and differentiable . Hence , we can use such gradient information to efficiently and approximately construct the optimal proposal distribution given by Theorem 1 . The basic idea is that we can approximate Eθ ( x−i ) based on the Taylor series of Eθ ( · ) at x , given that x−i is close to x in the data space because they only have differences in one dimension1 . Formally , Eθ ( x−i ) ≈ Eθ ( x ) + ( x−i − x ) T ∇xEθ ( x ) . ( 8 ) Thus , we can approximately obtain the desired term Eθ ( x ) −Eθ ( x−i ) in Theorem 1 using Eq . ( 8 ) . Note that ∇xEθ ( x ) ∈ Rd contains the information for approximating all Eθ ( x ) −Eθ ( x−i ) , where 1We have this assumption because data space is usually high-dimensional . If the number of data dimension is small , we can use the original ratio matching method with affordable time and memory budgets . Algorithm 1 Ratio Matching with Gradient-Guided Importance Sampling ( RMwGGIS ) Input : Observed dataset D = { x ( m ) } |D| m=1 , parameterized energy function Eθ ( · ) , number of samples s for Monte Carlo estimation with importance sampling 1 : for x ∼ D do ▷ Batch training is applied in practice 2 : Compute Eθ ( x ) 3 : Compute ∇xEθ ( x ) 4 : Compute the proposal distribution ñ∗ ( x−i ) = [ e2 ( 2x−1 ) ⊙∇xEθ ( x ) ] i∑d k=1 [ e2 ( 2x−1 ) ⊙∇xEθ ( x ) ] k ▷ Eq . ( 11 ) 5 : Sample s terms , denoted as x ( 1 ) −i , · · · , x ( s ) −i , according to ñ ∗ ( x−i ) 6 : Compute JRM ( θ , x ) ̂ñ∗ = d 1 s ∑s t=1 m ( x ( t ) −i ) [ e Eθ ( x ) −Eθ ( x ( t ) −i ) ] 2 ñ∗ ( x ( t ) −i ) ▷ Eq . ( 7 ) ( or Eq . ( 12 ) ) 7 : Update θ based on ∇θJRM ( θ , x ) ̂ñ∗ 8 : end for i = 1 , · · · , d. Hence , we can consider the following d-dimensional vector ( 2x− 1 ) ⊙∇xEθ ( x ) ∈ Rd , ( 9 ) where ⊙ denotes element-wise multiplication . Note that we have xi − x̄i = −1 if xi = 0 and xi − x̄i = 1 if xi = 1 , which can be unified as xi − x̄i = 2xi − 1 . Therefore , we have Eθ ( x ) − Eθ ( x−i ) ≈ [ ( 2x− 1 ) ⊙∇xEθ ( x ) ] i , i = 1 , · · · , d. ( 10 ) Afterwards , we can provide a proposal distribution ñ∗ ( x−i ) as an approximation of the optimal proposal distribution n∗ ( x−i ) given by Theorem 1 . Formally , ñ∗ ( x−i ) = [ e2 ( 2x−1 ) ⊙∇xEθ ( x ) ] i∑d k=1 [ e2 ( 2x−1 ) ⊙∇xEθ ( x ) ] k , i = 1 , · · · , d. ( 11 ) Then ñ∗ ( x−i ) is used as the proposal distribution for Monte Carlo estimation with importance sampling , as described by Eq . ( 7 ) . The overall process of our RMwGGIS method is summarized in Algorithm 1 .

Summary: This paper proposes a more efficient version of ratio matching (RM) for training discrete energy-based models. The proposed method subsamples the dimensions to use in the original RM objective and then uses importance sampling to reduce variance. The importance sampling distribution is based on a Taylor-series approximation to the target energy function which approximates the minimal-variance importance sampling distribution. The authors show that this estimator is considerably more efficient than standard RM and enables RM to be applied to training EBMs in higher dimensions. The authors demonstrate their estimator on some toy datasets and to fit EBMs on graph data. The method is evaluated using MMD.