Estimating the Partition Function

Comparing two models needs only a ratio

A model defines a distribution through an unnormalised score and a partition function , the sum that makes the probabilities add to one. Suppose two models, and , are to be compared on a test set of points. The book says that model A is the better model on this test set when Equation 18.38 holds, that is, when A has the higher test log-likelihood.

Testing that condition seems to need each model's partition function, because every log-probability contains a logarithm of its partition function. Rearranging the test gives Equation 18.39, which needs only the ratio of the two partition functions.

The ratio itself can be found from a single known partition function. If the ratio is known together with , then follows (Equation 18.40).

The demo fixes model A and changes the width of model B. The test set is 20 points spread evenly from −2 to 2, so every number is exact. The two readouts compute the total log-likelihood difference (A minus B) in two ways. The first uses the normalised densities directly, and the second uses the unnormalised scores and the ratio of the partition functions, as in Equation 18.39. They agree at every width.

The sign of the difference decides the comparison. For a width of 0.5, model A wins by about 24 units. For a width of 1.0, model B wins by about 6.3 units. Because the partition functions enter only through their ratio, the comparison never needs either one on its own.

Simple importance sampling

To estimate , choose a proposal whose partition function is known and which is easy to sample. Equations 18.41 to 18.43 rewrite the integral as an average under the proposal, using the target with its unnormalised form .

Equation 18.44 replaces the average by a sample mean over samples drawn from the proposal. The ratio of partition functions follows in Equation 18.45.

The proposal is a standard Gaussian, so its partition function is . The target is a Gaussian with width , so its exact ratio is . Each dot in the demo is one sample, plotted at its value, with its weight as the height. The weights are clipped at the top of the panel.

At the two distributions are the same, so every weight is exactly one and the estimate is exact for any number of samples. Moving the width away from one spreads the weights out. The readout gives the effective sample size, which counts how many samples really contribute. At a width of 0.8 with 1000 samples it is about 937. At a width of 1.8 with the same 1000 samples it falls to about 500, and the estimate is about 10 percent low in this run.

Above a width of √2 ≈ 1.41 the second moment of the weights is infinite. The readout reports the variance as infinite there, which is a property of the two Gaussians and not of the sample.

How large is the error of the estimate

Equation 18.46 gives the variance of the estimate from the same samples. The book notes that this quantity is largest when the importance weights differ a lot from one another. The typeset form in the book is reproduced below.

This is our check, not a statement about the book. We repeated the estimate 1000 times with , the demo default, and measured the spread of the results. We then compared that spread with the printed form and with a corrected form. The corrected form is added here and is not printed in the book. It follows from , with the sample mean of the weights, , in place of the true mean:

In this check the printed form does not match the measured spread. On the first sample set, its ratio to the corrected form is about 0.9 at a width of 0.3, about 12 at 0.8, and about 1.8 at 1.4, so it is not a fixed rescaling of the corrected form. Averaged over the 1000 repeats, the corrected form is within 3 to 8 percent of the measured spread at these three widths. A single sample set can differ from that average by more, as the readout shows. This check does not decide whether the printed form is a typo or uses another convention.

The histogram shows 1000 estimates of , each from a new set of samples. The dashed line is the exact value . Three numbers are compared in the readout: the variance predicted by the weights, the spread actually measured, and the two formulas evaluated on the first sample set. When the width reaches √2, the predicted variance becomes infinite and the histogram spreads out with it.

The printed-to-corrected ratio depends on the sample and on the width, so it can be small or large. The challenge asks for a setting where the printed form is at least three times the corrected form on the current sample.

Annealed importance sampling

When is large, the two distributions barely overlap and simple sampling fails. Annealed importance sampling (AIS) bridges the gap with a sequence of intermediate distributions (Jarzynski, 1997; Neal, 2001). A common choice is the weighted geometric average of the target and the proposal, given in Equation 18.50.

A transition must leave its intermediate distribution unchanged, as Equation 18.51 requires. The book allows any valid Markov chain method, such as Metropolis-Hastings or Gibbs sampling.

The ratio of partition functions is a chain of ratios over the stages, shown in Equations 18.47 to 18.49. The book prints Equation 18.47 as an identity: each added factor, such as , is a partition function divided by itself, which equals one, so the line says the same thing as its left side. Equation 18.48 regroups the factors so that each one is a ratio of neighbouring partition functions, and Equation 18.49 writes that product with a product sign.

The AIS weight of one run multiplies one ratio per stage, as in Equation 18.52. Averaging the weights over runs gives the estimate in Equation 18.53.

The demo uses Gaussians, where each intermediate distribution is itself Gaussian. Each run starts from a sample of the proposal, passes through stages, and applies two Metropolis steps at each stage. The book does not fix the number of transitions per stage, so two steps is the author's choice for this demo. At every stage the weight multiplies by the ratio of the new and old unnormalised densities at the current point.

The dots show the weight of each of 50 runs, and the solid line is their average, which is the AIS estimate of . The exact value is the dashed line. The runs are seeded, so the same settings always give the same picture.

The book shows that AIS is simple importance sampling on an extended space, so the average over many independent runs approaches the exact value. A single run is noisy. With the default seed, at a width of 1.2 and 50 runs, five stages give an error of about −4 percent, thirty stages about −0.1 percent, and sixty stages overshoot by about 1.6 percent. In this seeded example more stages did not always help.

Bridge sampling and training

Bridge sampling uses one distribution between the proposal and the target. The ratio is estimated by comparing the average weight of under each end. Equation 18.62 gives the estimator, which uses samples from both ends.

The book says that when the bridge overlaps both ends, the two distributions can be much farther apart than simple importance sampling allows. It states that the optimal bridge is with . That depends on the answer being estimated, so the book suggests starting from a coarse estimate of and refining it iteratively (Neal, 2005).

The bridge is a Gaussian of width and the target has width 1.2, so the exact ratio is 1.2. The curve shows the relative error of the estimate at every bridge width, using the same samples throughout. With the default 100 samples per end, the error is within 0.5 percent only for widths of about 1.01 to 1.05. It changes sign near width 1 and grows on both sides, to about −14 percent at 0.5 and about +12 percent at 1.6. The width with the smallest error moves with the number of samples: about 1.04 with 100, about 1.02 with 500, and about 0.56 with 1000. A single sample set does not give a general ranking of widths.

The book describes linked importance sampling (Neal, 2005), which combines the two methods. The book says bridge sampling can be more effective when the two distributions are close, and that AIS with many intermediate distributions spans the gap when one bridge cannot. Linked importance sampling uses the bridge idea to bridge the intermediate distributions of AIS.

The book calls AIS the most common method for estimating the partition function of undirected models, but it is too expensive to run during training. It describes the scheme of Desjardins et al. (2011), which combines bridge sampling, short-chain AIS and parallel tempering to track the partition function of an RBM throughout training.

• Equation 18.39 compares two models through the difference of their unnormalised scores and the ratio of their partition functions, so the test comparison needs only that ratio.
• Simple importance sampling estimates a ratio from samples of a proposal. Its weights are exact when the two distributions are the same, and its variance grows as the target spreads out.
• In our check, the printed variance in Equation 18.46 did not match the measured spread. The corrected form, which squares the constant and centres on the estimate divided by the constant, came within 3 to 8 percent of the measured spread on average at three widths.
• AIS passes samples through a chain of intermediate distributions. Each transition leaves its distribution unchanged, and the estimate is unbiased, although a single run is noisy.
• Bridge sampling uses both ends and an intermediate distribution, and it works best when the two ends overlap with the bridge.

Check yourself

Questions that mix the whole module. Each one is scored, and the score feeds your mastery of this module.

Source: Ian Goodfellow, Yoshua Bengio and Aaron Courville, Deep Learning, MIT Press, 2016, Section 18.7 (pages 623–630), free to read at deeplearningbook.org. Where the book prints equation numbers, the numbers here follow the book. The order of ideas follows the book. The explanations and all demos were written for this page.