Algorithms for Estimating the Partition Function of Restricted Boltzmann Machines (Extended Abstract)

1 Introduction

Markov random fields (MRFs, [Kindermann and Snell, 1980]) are undirected probabilistic graphical models that find many applications in AI, for example in computer vision [Blake et al., 2011] and neural computation [Smolensky, 1986]. The distribution modelled by an MRF can be written as

∗This paper is an extended abstract of an article in the Artificial Intelligence Journal (AIJ) [Krause et al., 2020].

†Contact Author

where E(x) is called the energy function and the normalization constant Z is referred to as the partition function. Computing and even estimating Z of an energy-based probabilistic model is typically challenging, because analytical integration is not possible and numerical integration unfeasible. But for many tasks, we would like to know Z with a high accuracy. We may need Z to assess the performance of models, to monitor maximum likelihood learning when adapting model parameters, and to perform likelihood ratio tests. In many cases, the estimation of the normalization constant is complicated by the inability to even generate samples from the distributions in question. Therefore, we would like to be able to obtain a good estimate without the requirements of exact sampling of the distribution.

This explains the popularity of algorithms such as Annealed Importance Sampling (AIS, [Neal, 2001]) which compute the ratio of normalization constants between a reference distribution pref and a target distribution ptarget. If pref is chosen such that its normalization constant is known, the AISestimate of the ratio of the normalization constants gives an estimate of the partition function of ptarget. AIS introduces so called bridging distributions, which interpolate between pref and ptarget as well as a sampling scheme that allows to estimate the normalization constant only requiring exact samples from pref. The performance and limitations of AIS are well known in the AI community [Salakhutdinov and Murray, 2008; Schulz et al., 2010]. It is very general, still, the weak assumptions come at the expense of a possibly large variance if the number of bridging distributions is too small.

There are alternative algorithms such as variants of Bennett's Acceptance Ratio (BAR, [Bennett, 1976]), which was rediscovered several times as Bridge Sampling [Meng and Wong, 1996], Reverse Logistic Regression [Geyer, 1994], and Discriminance Sampling [Liu et al., 2015], which also compute estimates of the ratio of normalization constants. However, they require independent samples from all bridging distributions and consequently there are only few studies applying BAR in the context of energy-based models in the AI community [Desjardins et al., 2011; Liu et al., 2015]. A possible source of samples with small bias from pref, ptarget, and all bridging distributions is Parallel Tempering (PT, [Desjardins et al., 2010]), which is often used for sampling from graphical models during training. Parallel Tempering introduces replica Markov chains to foster faster mixing, which leads to increasing sample quality in all chains. If the distributions of the replica chains are chosen to be the same as the bridging distributions of the BAR-estimator, it is possible to re-use the samples acquired during training for estimation the normalization constant at no additional cost (see, e.g., [Desjardins et al., 2011]).

Still, the theoretical and practical properties of BAR compared to AIS have been little studied in the AI community. Our work contributes to closing this gap using a theoretical framework that can be used to derive AIS as well as different variants of BAR and other algorithms, allowing a concise proof of the fact that BAR is a maximum-likelihood estimator of the normalization constant [Shirts et al., 2003]. In this extended abstract summarizing the work by [Krause et al., 2020], we will put an emphasis on experimental results, showcasing the advantages of the different algorithms. We focus on Restricted Boltzmann Machines (RBMs, [Smolensky, 1986; Hinton, 2002; Fischer and Igel, 2014]) as a particular class of Markov random fields, but our considerations are not limited to these stochastic neural networks. [Krause et al., 2020] also show results for the 2D-Ising model with external magnetic fields and we would like to point out that the results are also relevant for other types of generative models (e.g., see [Wu et al., 2017] for an application scenario in the context of Generative Adversarial Networks).

2 Main Result

Let us briefly sketch our main theoretical result, a generalization of Crooks' equation [Crooks, 2000] to arbitrary sampling distributions. We refer to [Krause et al., 2020] for details.

Let pref = p0, p1, . . . , pN = ptarget be a set of Gibbs distributions over some state space Ω with pi : Ω → R and

where p∗i (x) denotes the unnormalized probability distribution. Our goal is to estimate ZN/Z0.

We now consider a random variable X = (X0,XN, Y ) taking values x = (x0, xN, y) in an extended state space Ω† = Ω2 × Θ. Here y, taking values in the state space Θ, is a placeholder for any set of additional variables an actual estimation method may require. Assume that we can use the set of Gibbs distributions p0, p1, . . . , pN to construct a pair of distributions pF and pR on Ω† with

We call pF the forward distribution, because it creates samples from pN given a sample x0 from p0, and we refer to pR as the reverse distribution.

Consider now any function F on the extended state space Ω† . We are interested in relating expectations of F under pF to expectations of F under pR. [Krause et al., 2020] prove that

where 〈f(x)〉p(x) = ∫ p(x)f(x) dx and

By choosing pF , pR and F, we can now derive several estimation algorithms. By setting F(x) = 1, we obtain

which for particular choices of pF becomes equivalent to AIS or Linked Importance Sampling [Neal, 2005]. Instead of choosing an arbitrary F, BAR uses the variant which minimizes the variance

Unfortunately, this expression involves the unknown normalization constant itself. However, inserting into (1) and simplifying leads to a fixed-point problem which can be solved for Z0 ZN . [Krause et al., 2020] show that this is indeed the maximum-likelihood estimator. A practical variant can be obtained by setting pF (x) = ∏Ni=0 pi(xi) and pR(x) = pN(xN) ∏Ni=1 pi(xi-1). In this case, we can compute the estimator using independent samples from the N bridging distributions as produced by parallel tempering.

3 Experiments

We empirically compared different algorithms for partition function estimation. Here, we will only show results for two of these algorithms, vanilla AIS and BAR (which [Krause et al., 2020] refer to as BARPT-ind). We asked:

1. How accurate are the methods depending on the number of bridging distributions and the choice of the reference distribution?
2. Do the methods tend to over- or underestimate the normalization constant?
3. How do the methods perform in an online setting, where PT samples are used both for training as well as for partition function estimation?

To address these questions, we added several partition function estimation methods to the open-source machine learning library Shark [Igel et al., 2008]. In this extended abstract, we present only results for MNIST data set [LeCun et al., 1998]. In all experiments, we considered the task of estimating ln ZN Z0 with a fixed budget of sampling steps. We measured the mean relative error of an estimate C defined as

.

All our experiments were based on the same setup. For a given RBM we took two sets of samples, one created by PT as a way to get approximately independent samples and one using the AIS sampling-scheme. Given a reference distribution pref and the target distribution p(x) = 1 ZN e-E(x) specified by the RBM parameters, we constructed the bridging distributions as

with βi = iN , for i = 1, 2 . . . ,N - 1. Thus, we chose a geometric mean of distributions with a uniform spacing of β-values. This is often done in practice and fits the theoretical results by [Fischer and Igel, 2015]. In our experiments, we considered two different choices of reference distribution pref, the uniform distribution, which is used often in AIS, and a distribution which uses the pixel-wise marginal probabilities as computed from the MNIST training set. The second distribution is also called "base-rate" by [Salakhutdinov and Murray, 2008].

To answer the first question raised above, we used a larger RBM with 500 hidden units trained with CD-25 by [Salakhutdinov and Murray, 2008]. We compared standard AIS with BAR using samples generated by PT. We varied the number of bridging distributions while keeping the overall budget of sampling steps fixed. For estimation, we used 1,200,000 samples in total where 50 % were used as burn-in time by PT. The results can be seen in Figure 1. In our experiments, BAR performed better with a small amount of bridging distributions, while AIS required a very large amount. A large amount of distributions can be detrimental for BAR as performing initial burn-in of PT becomes very expensive. Furthermore, all algorithms performed better when using the "base-rate" reference distribution, even outperforming our pre-computed ground truth estimate of lnZn/Z0 (which was computed by AIS using a uniform distribution and a very large budget).

In the next experiment, we investigated the second question. For this, we trained a smaller RBM with 16 hidden neurons where the normalization constant can be computed exactly. We took the RBM obtained at the end of the training, estimated ln(ZN/Z0) 1000 times for the different algorithms (i.e., BAR with 10 and AIS with 10,100, or 1000 bridging distributions), and computed histograms over these estimates. The resulting histograms are given in Figure 2. The results show that the distribution of estimates obtained by AIS when using the uniform reference distribution is skewed towards small values. All algorithms underestimated the true value, but BAR was slightly better, most likely because PTsamples had a longer time to converge to the true distribution. When using the "base-rate" reference distribution, all algorithms performed significantly better. However, AIS with 10 bridging distributions still showed a larger variance than BAR as well as a visible skew. This is in line with the general knowledge that AIS tends to underestimate the normalization constant, even though it is unbiased in expectation.

In our third experiment, we investigated the online performance. Again, we trained an RBM with 16 hidden units and computed the normalization constant exactly. To generate the samples required for training, we used PT with 50 Markovchains and re-used the samples to compute an estimate of the normalization constant. We did not use any burn-in, which is in line with usual training procedures. We compared our results to AIS and BAR using separate samples (and using burnin for PT). The results can be seen in Figure 3. While OnlineBAR performed slightly worse than the offline-variants, it was still able to accurately track the progress of training. To visualize this, we plotted the true log-likelihood curve and added shaded regions depicting the error-ranges as obtained using the BAR-estimator.

4 Conclusions

[Krause et al., 2020] derived a generalized form of Crooks' equality, which can be used to devise generalizations of known estimators for the partition functions of Markov Random Fields (MRFs), including Annealed Importance Sampling (AIS) and Bennett's Acceptance Ratio method (BAR). Various methods for generating samples are covered such as Parallel Tempering (PT), path sampling (used by AIS), and Linked Importance Sampling [Neal, 2005].

When empirically comparing PT-based estimators with vanilla AIS for estimation the normalization constants we found different regimes: PT based estimators (as BAR shown here) worked well with a small number of bridging distributions but required many samples, while AIS required many bridging distributions but only a few samples of those. This makes BAR a particularly good candidate for monitoring training process when PT is used during training.

Another important result is that choosing a reference distribution which is close to the target distribution can have a major impact on the quality of the estimate.

Acknowledgements

This work was supported by the Innovation Fund Denmark through the Danish Center for Big Data Analytics Driven Innovation (DABAI).