SLT Notebook 1 — Bayesian inference as statistical mechanics

Question: what does Bayes’ rule look like in the language of statistical mechanics?

Notation: prior \varphi, training loss L_n, population optimum w_0.

The likelihood of the data is e^{-nL_n(w)}, so Bayes’ rule reads

\pi_n(w)=\frac{\varphi(w)\,e^{-nL_n(w)}}{Z_n}, \qquad Z_n=\int\varphi(w)\,e^{-nL_n(w)}\,dw, \qquad F_n=-\log Z_n .

This is a Boltzmann distribution. The total training loss nL_n is the energy and the prior weights the states. The evidence Z_n is the partition function, and the free energy is low when the model gave the data high probability.

Two variants are useful. A power \beta on the likelihood is an inverse temperature, with Bayes at \beta=1. Factoring out the loss at the population optimum, with K_n=L_n-L_n(w_0), leaves a free energy that measures how much prior mass fits the data nearly as well as w_0:

\pi_{n,\beta}\propto\varphi\,e^{-\beta nL_n}, \qquad F_n=nL_n(w_0)+\bar F_n, \qquad \bar F_n=-\log\int\varphi\,e^{-nK_n}\,dw .

The analogy is formal: the energy is random because it depends on the sample, n is not a system size, and the real analogue of the density of states is the volume of parameters at a given loss (Notebook 2). The algebra, though, is exact.

Exercise 1.1 (Free-energy bookkeeping). For a region U of parameter space let F_n(U)=-\log\int_U\varphi e^{-nL_n}.

1. Show that the log posterior odds of two regions are F_n(V)-F_n(U).
2. If k regions partition parameter space and U_1 has the lowest free energy, show F_n(U_1)-\log k\le F_n\le F_n(U_1).
3. Show that \partial_\beta F_n(\beta) is the internal energy n\,\mathbb E_\beta[L_n], and that F_n(\beta) is concave.

Question: what does the posterior look like when n is large?

Notation: H=\nabla^2L(w_0) and H_n=\nabla^2L_n(\hat w) are Hessians; C=\mathrm{Cov}_q[\nabla\ell(x,w_0)] is the gradient covariance. The model is regular: w_0 is the unique minimizer and H is positive definite (plus the usual smoothness and moment conditions).

As n grows the likelihood sharpens while the prior stays fixed. The posterior is squeezed onto \hat w, where only the quadratic part of the loss matters. The centre \hat w is itself random, since the training loss depends on the sample. Two Gaussian statements follow: the Bernstein–von Mises theorem, and the asymptotic normality of the maximum-likelihood estimate.

\pi_n\approx\mathcal N\big(\hat w,\,(nH_n)^{-1}\big), \qquad \hat w-w_0\approx\mathcal N\big(0,\,\tfrac1nH^{-1}CH^{-1}\big).

The first is what the posterior believes. The second is how the estimate actually varies between data sets. Both have width n^{-1/2}, and they agree when the model is well specified, because then C=H. Otherwise the posterior is over- or under-confident.

Exercise 1.2 (The Gaussian posterior).

1. Expand nL_n to second order about \hat w, freeze the prior at its value there, and show that the posterior is the first Gaussian above. Writing w=\hat w+u/\sqrt n, show that the neglected terms are O(n^{-1/2}).
2. Show that \hat w-w_0\approx-H^{-1}\nabla L_n(w_0), and use the central limit theorem to get the second Gaussian.
3. Show that C=H when the model is well specified.
   Hint

   Differentiate \int p(x\mid w)\,dx=1 twice with respect to w.

Question: what is the free energy for large n, and what does it penalize?

Notation: \hat w minimizes L_n, H_n is the Hessian there (positive definite: the model is regular, §1.2), d is the number of parameters.

Laplace’s method replaces the integrand of Z_n by a Gaussian around \hat w (Exercise 1.3). The result is

F_n=nL_n(\hat w)+\frac d2\log n+\frac12\log\det H_n-\log\varphi(\hat w)-\frac d2\log2\pi+o(1).

The first two terms are the Bayesian information criterion (BIC): the best fit plus half a \log n per parameter. The rest is of order one; it depends on the curvature and the prior.

MacKay (1992) reads this as the best-fit likelihood times an Occam factor. For a uniform prior, the Occam factor is the fraction of prior volume that still fits the data. Each parameter’s width shrinks to about n^{-1/2}, and the evidence pays for the shrinkage. No held-out data are involved.

Exercise 1.3 (Laplace, BIC and the Occam factor).

1. Using \int e^{-\frac12u^\top Au}\,du=(2\pi)^{d/2}(\det A)^{-1/2} for positive definite A, derive the formula for F_n.
2. Take a uniform prior on a box of volume V_\varphi, and let V_{\rm post}=\int_{\rm box}e^{-n(L_n(w)-L_n(\hat w))}dw. Show that Z_n=e^{-nL_n(\hat w)}\,V_{\rm post}/V_\varphi exactly.
3. Model B has more parameters than model A and fits better. Using the BIC, how much must B lower the per-sample training loss to be preferred? Compare with what fitting pure noise gains, about \tfrac1{2n} per extra parameter.

Code. The first cell checks Laplace against the exact free energy of the conjugate model \mathcal N(w,1) with a Gaussian prior; the second computes F_n by quadrature for \mathcal N(\mu,e^{2s}) and checks the slope d/2=1 in \log n.

# 1-D conjugate model: x ~ N(w, 1), prior w ~ N(0, tau^2). Exact Z_n vs quadrature vs Laplace.
tau2, n = 4.0, 50
x = rng.normal(1.0, 1.0, size=n)
xbar = x.mean()
nL_hat = 0.5 * n * np.log(2 * np.pi) + 0.5 * np.sum((x - xbar) ** 2)   # n L_n(w_hat)

F_exact = nL_hat + 0.5 * np.log(n) + 0.5 * np.log(tau2 + 1 / n) + xbar**2 / (2 * (tau2 + 1 / n))
F_laplace = nL_hat + 0.5 * np.log(n) + 0.5 * np.log(tau2) + xbar**2 / (2 * tau2)
F_bic = nL_hat + 0.5 * np.log(n)

w = np.linspace(-10, 10, 20001)
log_integrand = (-0.5 * w**2 / tau2 - 0.5 * np.log(2 * np.pi * tau2)
                 - 0.5 * n * np.log(2 * np.pi) - 0.5 * ((x[None, :] - w[:, None]) ** 2).sum(1))
F_quad = -log_trapz(log_integrand, w)
print(f"F_n exact     {F_exact:.4f}")
print(f"F_n quadrature{F_quad:.4f}")
print(f"F_n Laplace   {F_laplace:.4f}   (difference {F_laplace - F_exact:+.4f}, predicted O(1/n) ~ {1/(2*n*tau2):.4f})")
print(f"F_n BIC terms {F_bic:.4f}")

F_n exact     69.5111
F_n quadrature69.5111
F_n Laplace   69.5094   (difference -0.0017, predicted O(1/n) ~ 0.0025)
F_n BIC terms 68.6570

Code (42 lines)

# 2-D model: x ~ N(mu, exp(2 s)), priors mu ~ N(0,1), s ~ N(0,1). Data from N(1, 0.5^2).
mu_true, sigma_true = 1.0, 0.5
N_all = 3000
x_all = rng.normal(mu_true, sigma_true, size=N_all)
mu_grid = np.linspace(-2.0, 4.0, 601)
s_grid = np.linspace(-4.0, 2.0, 601)
MU, S = np.meshgrid(mu_grid, s_grid, indexing="ij")
log_prior = -0.5 * (MU**2 + S**2) - np.log(2 * np.pi)
dmu, ds = mu_grid[1] - mu_grid[0], s_grid[1] - s_grid[0]

def free_energy_2d(x):
    n = len(x)
    # n L_n(mu, s) = n s + n/2 log 2pi + exp(-2s) * sum (x_i - mu)^2 / 2
    ss = ((x[None, None, :] - MU[:, :, None]) ** 2).sum(-1) if n <= 300 else \
         n * ((MU - x.mean()) ** 2 + x.var())           # same thing, computed from sufficient statistics
    nL = n * S + 0.5 * n * np.log(2 * np.pi) + 0.5 * np.exp(-2 * S) * ss
    return -(logsumexp(log_prior - nL) + np.log(dmu * ds))

def laplace_2d(x):
    n = len(x)
    mu_hat, s_hat = x.mean(), 0.5 * np.log(x.var())
    nL_hat = n * s_hat + 0.5 * n * np.log(2 * np.pi) + 0.5 * n       # exp(-2 s_hat) * n var = n
    H = np.diag([np.exp(-2 * s_hat), 2.0])                            # Hessian of L_n at the optimum
    F_lap = (nL_hat + 1.0 * np.log(n) + 0.5 * np.log(np.linalg.det(H))
             + 0.5 * (mu_hat**2 + s_hat**2) + np.log(2 * np.pi) - np.log(2 * np.pi))
    return nL_hat, F_lap

ns = np.array([10, 20, 50, 100, 200, 500, 1000, 2000, 3000])
rows = []
for n in ns:
    x = x_all[:n]
    nL_hat, F_lap = laplace_2d(x)
    rows.append((n, free_energy_2d(x) - nL_hat, F_lap - nL_hat, 1.0 * np.log(n)))
rows = np.array(rows)
plt.figure(figsize=(5.5, 3.6))
plt.plot(rows[:, 0], rows[:, 1], "o-", label="exact (quadrature)")
plt.plot(rows[:, 0], rows[:, 2], "s--", label="Laplace")
plt.plot(rows[:, 0], rows[:, 3], ":", label=r"BIC: $\frac{d}{2}\log n$")
plt.xscale("log"); plt.xlabel("n"); plt.ylabel(r"$F_n - nL_n(\hat w)$")
plt.title("2-parameter Gaussian model: free energy above the best fit"); plt.legend(); plt.tight_layout(); plt.show()
print("slope of the exact curve against log n over the last four points:",
      np.polyfit(np.log(rows[-4:, 0]), rows[-4:, 1], 1)[0].round(3), " (d/2 = 1)")

slope of the exact curve against log n over the last four points: 1.009  (d/2 = 1)

Question: how well does the posterior predict the next sample, compared with the maximum-likelihood fit?

Notation: L is the population loss, minimized at w_0; \hat w minimizes the training loss; S is the entropy of the data; \pi_n is the posterior and F_n the free energy (§1.1); H and C are the Hessian and gradient covariance at w_0; L_{\min}=L(w_0)-S is the best excess loss the model can reach (zero if well specified).

There are three predictors. Bayes averages the model over the posterior, giving the predictive density p^*. Gibbs uses one posterior sample. Maximum likelihood uses \hat w. Each one’s generalization error G_n is its expected loss on a fresh sample minus the entropy, which no predictor can remove. The results below also average over training sets.

The free energy enters through an exact identity (Exercise 1.4.1):

p^*(x)=\int p(x\mid w)\,\pi_n(w)\,dw, \qquad -\log p^*(x_{n+1})=F_{n+1}-F_n .

In the regular case, up to o(1/n):

When the model is well specified the trace is d. Bayes and maximum likelihood then tie, and one posterior sample pays twice as much. Under misspecification Bayes still pays d/2n, while maximum likelihood can pay more or less. So Bayes does not always generalize better; what posterior averaging removes is the dependence on C.

Exercise 1.4 (Three generalization errors).

1. Prove -\log p^*(x_{n+1})=F_{n+1}-F_n.
2. Derive the table. Start from the fluctuation of \hat w (§1.2) for the maximum-likelihood row, then add the posterior spread for Gibbs. For Bayes, expand \log p(x\mid w) to second order about \hat w and use \log\mathbb E[e^Y]\approx\mathbb E[Y]+\tfrac12\mathrm{Var}[Y].
3. (Misspecified.) Model \mathcal N(x;w,1), flat prior, data \mathcal N(0,\sigma^2). Compute the three expected errors exactly. When does Bayes beat maximum likelihood?

Code. Simulates Exercise 1.4.3 and plots n times each expected error; the dotted lines are the table. Set sigma2 below 1 to see Bayes and maximum likelihood swap.

sigma2 = 2.0          # variance of the data; the model assumes 1
tau2 = 100.0          # prior variance (wide, so the prior barely matters)
ns = np.array([5, 10, 20, 50, 100, 200, 500])
trials = 4000
S = 0.5 * np.log(2 * np.pi * np.e * sigma2)
L_min = 0.5 * (sigma2 - 1 - np.log(sigma2))

def errors(xbar, n):
    # TODO: return (G_bayes, G_gibbs, G_mle) minus L_min, given the sample mean xbar
    # posterior: N(m, v) with v = 1 / (n + 1/tau2), m = v * n * xbar; predictive: N(m, 1 + v)
    raise NotImplementedError

# then: average over trials for each n and plot n * G against n, with the three predicted lines

Solution (code)

sigma2 = 2.0          # variance of the data; the model assumes 1
tau2 = 100.0          # prior variance (wide, so the prior barely matters)
ns = np.array([5, 10, 20, 50, 100, 200, 500])
trials = 4000
S = 0.5 * np.log(2 * np.pi * np.e * sigma2)
L_min = 0.5 * (sigma2 - 1 - np.log(sigma2))

def errors(xbar, n):
    v = 1.0 / (n + 1.0 / tau2)
    m = v * n * xbar
    # expected loss of a Gaussian predictor N(m, s2) on x ~ N(0, sigma2): 1/2 log(2 pi s2) + (sigma2 + m^2) / (2 s2)
    bayes = 0.5 * np.log(2 * np.pi * (1 + v)) + (sigma2 + m**2) / (2 * (1 + v)) - S - L_min
    gibbs = 0.5 * np.log(2 * np.pi) + 0.5 * (sigma2 + m**2 + v) - S - L_min
    mle = 0.5 * np.log(2 * np.pi) + 0.5 * (sigma2 + xbar**2) - S - L_min
    return bayes, gibbs, mle

means = []
for n in ns:
    xbar = rng.normal(0.0, np.sqrt(sigma2 / n), size=trials)
    means.append([np.mean(g) for g in errors(xbar, n)])
means = np.array(means)

plt.figure(figsize=(5.5, 3.6))
for j, (lab, pred) in enumerate([("Bayes predictive", 0.5), ("Gibbs (one sample)", (1 + sigma2) / 2),
                                 ("maximum likelihood", sigma2 / 2)]):
    plt.plot(ns, ns * means[:, j], "o-", color=f"C{j}", label=lab)
    plt.axhline(pred, color=f"C{j}", ls=":")
plt.xscale("log"); plt.xlabel("n"); plt.ylabel(r"$n\,\mathbb{E}[G_n - L_{\min}]$")
plt.title(f"model N(w,1), data N(0,{sigma2}); dotted = table"); plt.legend(loc="center right"); plt.tight_layout(); plt.show()