Annealed Importance Sampling and Estimating log Z

A trained energy-based model can score any input and cannot say how likely it is. The log-likelihood is \(\log p_\theta(x) = -E_\theta(x) - \log Z_\theta\), and the second term is an integral over the whole input space. Every reported likelihood for a restricted Boltzmann machine, deep Boltzmann machine or Markov random field is therefore an estimate, and the standard estimator has a bias whose direction is both known and inconvenient: it makes models look better than they are.

Bridging between a tractable model and the real one

Pick a starting distribution \(f_0\) whose normaliser \(Z_0\) you know, typically the model with all weights zeroed, and the target \(f_K\) you care about. Interpolate with a sequence of unnormalised bridging distributions, usually geometric:

Sample \(x_0\) from \(p_0\), then apply an MCMC transition operator that leaves each \(p_k\) invariant in turn, and accumulate the importance weight

Neal's result is that the Markov chain transitions along such an annealing sequence define a valid importance sampler, with \(\mathbb{E}[w] = Z_K/Z_0\) exactly (Neal, 1998, Annealed Importance Sampling, arXiv:physics/9803008; published as Statistics and Computing 11, 125-139, 2001). Averaging \(w\) over independent runs gives an unbiased estimate of the ratio of partition functions, and since \(Z_0\) is known, of \(Z_K\).

The unbiasedness applies to \(Z\), not to \(\log Z\). By Jensen's inequality, \(\mathbb{E}[\log \hat{w}] \le \log \mathbb{E}[\hat{w}] = \log(Z_K/Z_0)\), so the log estimate is biased downwards. Because \(\log Z\) enters the likelihood with a minus sign, an underestimate of \(\log Z\) is an overestimate of the log-likelihood. The practical consequence is stated plainly in the work that fixed it: standard AIS tends to overestimate MRF log-likelihoods, and does so without warning (Burda, Grosse and Salakhutdinov, 2015, Accurate and Conservative Estimates of MRF Log-likelihood using Reverse Annealing, AISTATS, PMLR 38).

Bounding the error from the other side

The fix is to run the annealing in reverse. The Reverse AIS Estimator starts from an exact sample of the target and anneals back towards the tractable distribution, using the same MCMC transition operators as forward AIS, and yields a stochastic lower bound on the log-likelihood of an approximation to the original model. In experiments on RBMs, deep Boltzmann machines and deep belief networks, RAISE tracks the AIS numbers closely but errs towards underestimating rather than overestimating.

Running both directions gives a sandwich. Forward AIS or sequential Monte Carlo provides stochastic lower bounds on the log marginal likelihood; the same algorithms run in reverse from an exact posterior sample provide upper bounds, and the true value lies between them with high probability (Grosse, Ghahramani and Adams, 2015, Sandwiching the marginal likelihood using bidirectional Monte Carlo, arXiv:1511.02543). The catch is in the premise: an exact posterior sample is available for data simulated from the model, so the tight sandwich is a tool for validating estimators and models on synthetic data, not for reporting a number on a real test set.

Reading an AIS number

Three quantities decide whether an estimate is worth anything, and all three should be reported.

The number of bridging distributions, \(K\). Too few and consecutive distributions overlap poorly, the weights become dominated by a handful of runs, and the estimate collapses towards a point whose error is invisible from the output. Published RBM estimates typically use thousands to tens of thousands of intermediate distributions.

The spread of the log weights. The weights are heavy-tailed, so the usual diagnostic is the effective sample size implied by the variance of \(\log w\). A large number of chains with a tiny effective sample size looks precise and is not.

Which direction was run. A single forward estimate is a biased upper bound on the likelihood. Two models compared by forward AIS alone are being compared through two biases pointing the same way but not by the same amount.

When it breaks

Convergence failure is silent. Nothing in the output of a forward AIS run signals that the chains failed to mix. The estimate is simply too high, and the only way to detect it is a lower bound from the other direction or a sandwich on simulated data.

The bias grows where the model is interesting. The gap between forward and reverse estimates widens with the dimension and with the roughness of the energy landscape, which means the sharpest, most multimodal models are the ones whose likelihoods are least trustworthy.

It does not rescue model comparison across families. An AIS likelihood for an RBM and an exact likelihood for an autoregressive model are not the same kind of quantity. Treating a point estimate with unreported variance as comparable to an exact number has produced more than one misleading leaderboard.

Cost scales with accuracy, not with the model. Halving the gap between the bounds generally means multiplying \(K\), so a trustworthy likelihood for a mid-sized RBM can cost more compute than training it did.