"""Numerical checks of Bishop (2006), Chapter 10, Sections 10.1 to 10.4, using the functions of ch10_variational_inference.py. Each check prints PASS or FAIL; run `python3 checks_ch10a.py`. Covered: the mean-field Gaussian (10.12) to (10.15) and the reverse-KL marginals (10.17); the Gaussian fits of Fig 10.1; the VB treatment of the Gaussian mean and precision (10.26) to (10.33); the variational mixture of Gaussians (10.46) to (10.82), including Monte Carlo checks of the bound terms (10.71) to (10.77) and of (10.65), (10.66), finite-difference checks that each update is a stationary point (and that the misprinted nu_k = nu_0 + N_k + 1 is not); variational linear regression (10.93) to (10.112). """ import numpy as np from scipy import optimize, stats, integrate from scipy.special import digamma, gammaln, logsumexp import common import ch10_variational_inference as c RESULTS = [] def check(name, ok, detail=""): RESULTS.append(bool(ok)) print("%s %s%s" % ("PASS" if ok else "FAIL", name, (" [" + detail + "]") if detail else "")) # --------------------------------------------------------------------------- 10.1.2 def checks_factorised(): mu, Sigma = c.FIG10_2_MU, c.FIG10_2_SIGMA Lam = np.linalg.inv(Sigma) rng = np.random.default_rng(1) path, var = c.mean_field_gaussian(mu, Lam, rng.uniform(0, 1, 2), n_sweeps=400) check("mean-field Gaussian: (10.13), (10.15) converge to the true mean (Exercise 10.2)", np.abs(path[-1] - mu).max() < 1e-10, "error %.1e" % np.abs(path[-1] - mu).max()) check("mean-field Gaussian: factor variances are 1/Lambda_11, 1/Lambda_22 (10.12), (10.14)", np.allclose(var, 1 / np.diag(Lam))) def kl_qp(th): # factorised Gaussian q: means th[:2], log sds th[2:] return c.kl_gauss(th[:2], np.diag(np.exp(2 * th[2:])), mu, Sigma) res = optimize.minimize(kl_qp, [0.3, 0.7, -2.0, -2.0], method="BFGS", options={"gtol": 1e-10}) check("mean-field Gaussian: direct minimisation of KL(q||p) over factorised Gaussians agrees", np.allclose(res.x[:2], mu, atol=1e-6) and np.allclose(np.exp(2 * res.x[2:]), var, rtol=1e-5)) def kl_pq(th): return c.kl_gauss(mu, Sigma, th[:2], np.diag(np.exp(2 * th[2:]))) res = optimize.minimize(kl_pq, [0.3, 0.7, -2.0, -2.0], method="BFGS", options={"gtol": 1e-10}) m_pq, v_pq = c.factorised_kl_pq(mu, Sigma) check("KL(p||q) over factorised Gaussians is minimised by the marginals (10.17)", np.allclose(res.x[:2], m_pq, atol=1e-6) and np.allclose(np.exp(2 * res.x[2:]), v_pq, rtol=1e-5)) # (10.17) for a general (discrete) p, over all factorised q P = rng.uniform(0.1, 1, (3, 4)) P /= P.sum() def kl_disc(th): q1 = np.exp(th[:3] - logsumexp(th[:3])) q2 = np.exp(th[3:] - logsumexp(th[3:])) return np.sum(P * (np.log(P) - np.log(np.outer(q1, q2)))) res = optimize.minimize(kl_disc, np.zeros(7), method="BFGS", options={"gtol": 1e-10}) q1 = np.exp(res.x[:3] - logsumexp(res.x[:3])) q2 = np.exp(res.x[3:] - logsumexp(res.x[3:])) check("KL(p||q1 q2) for a discrete 3 x 4 table is minimised by the two marginals (10.17)", np.allclose(q1, P.sum(1), atol=1e-6) and np.allclose(q2, P.sum(0), atol=1e-6)) # --------------------------------------------------------------------------- Fig 10.1 def checks_fig10_1(): lp = c.skewed_log_density z0, s_lap = c.laplace_1d() m_v, s_v = c.gauss_kl_qp_1d(lp, z0, s_lap) m_e, s_e, lnZ = c.moments_1d(lp) u, w = c.gauss_grid(2001, 8.0) zz = np.linspace(-12, 12, 240001) pz = np.exp(lp(zz) - lnZ) def klqp(m, s): return -0.5 * np.log(2 * np.pi * np.e * s * s) - np.sum(w * lp(m + s * u)) + lnZ def klpq(m, s): return np.trapz(pz * (lp(zz) - lnZ - c.gauss_logpdf_1d(zz, m, s)), zz) check("Fig 10.1: the Laplace mode solves -z + 20(1 - sigma(20z + 4)) = 0", abs(-z0 + 20 * (1 - common.sigmoid(20 * z0 + 4))) < 1e-12, "z0 = %.6f" % z0) h = 1e-5 g = [(klqp(m_v + h, s_v) - klqp(m_v - h, s_v)) / (2 * h), (klqp(m_v, s_v + h) - klqp(m_v, s_v - h)) / (2 * h)] check("Fig 10.1: the variational Gaussian is a stationary point of KL(q||p)", max(map(abs, g)) < 1e-6, "gradient %s" % np.round(g, 9)) kq = [klqp(z0, s_lap), klqp(m_v, s_v), klqp(m_e, s_e)] kp = [klpq(z0, s_lap), klpq(m_v, s_v), klpq(m_e, s_e)] check("Fig 10.1: the variational fit has the smallest KL(q||p) of the three Gaussians", kq[1] < min(kq[0], kq[2]), "Laplace %.4f, variational %.4f, moments %.4f" % tuple(kq)) check("Fig 10.1: the moment-matched fit has the smallest KL(p||q) of the three Gaussians", kp[2] < min(kp[0], kp[1]), "Laplace %.4f, variational %.4f, moments %.4f" % tuple(kp)) check("normalising constant of p(z): grid integral agrees with adaptive quadrature", abs(np.exp(lnZ) - integrate.quad(lambda z: np.exp(lp(z)), -10, 10, points=[-0.2], limit=200)[0]) < 1e-9) # --------------------------------------------------------------------------- 10.1.3 def checks_vb_gaussian(): x = c.VBG_DATA[:c.VBG_N] pr = c.VBG_PRIOR states = c.vb_gaussian(x, pr, *c.VBG_INIT) L = np.array([c.vb_gaussian_bound(x, pr, qm, qt) for qm, qt in states]) lnp = c.gaussian_log_evidence(x, **pr) check("VB Gaussian: the bound never decreases over single-factor updates", np.all(np.diff(L) > -1e-12), "%d updates" % (len(L) - 1)) check("VB Gaussian: the bound stays below ln p(D) at every update", np.all(L <= lnp + 1e-12), "gap at convergence %.4f" % (lnp - L[-1])) # ln p(D) by two-dimensional quadrature of p(D | mu, tau) p(mu | tau) p(tau) def joint(mu, tau): return np.exp(np.sum(-0.5 * tau * (x - mu) ** 2) + 0.5 * len(x) * np.log(tau / (2 * np.pi)) + c.gauss_logpdf_1d(mu, pr["mu0"], 1 / np.sqrt(pr["lam0"] * tau)) + c.gamma_logpdf(tau, pr["a0"], pr["b0"]) - lnp) val = integrate.dblquad(lambda mu, tau: joint(mu, tau), 1e-8, 6, -3, 3, epsabs=1e-11)[0] check("VB Gaussian: the closed-form ln p(D) agrees with 2-D quadrature", abs(val - 1) < 1e-6, "integral of p(D, mu, tau)/p(D) = %.8f" % val) # Monte Carlo value of the bound at the converged q qm, qt = states[-1] rng = np.random.default_rng(3) S = 400000 mu = rng.normal(qm["m"], 1 / np.sqrt(qm["lam"]), S) tau = rng.gamma(qt["a"], 1 / qt["b"], S) lj = (0.5 * len(x) * np.log(tau / (2 * np.pi)) - 0.5 * tau * (np.sum(x ** 2) - 2 * mu * np.sum(x) + len(x) * mu ** 2) + c.gauss_logpdf_1d(mu, pr["mu0"], 1 / np.sqrt(pr["lam0"] * tau)) + c.gamma_logpdf(tau, pr["a0"], pr["b0"]) - c.gauss_logpdf_1d(mu, qm["m"], 1 / np.sqrt(qm["lam"])) - c.gamma_logpdf(tau, qt["a"], qt["b"])) se = lj.std() / np.sqrt(S) check("VB Gaussian: the closed-form bound agrees with its Monte Carlo estimate", abs(lj.mean() - L[-1]) < 4 * se, "closed %.5f, MC %.5f +- %.5f" % (L[-1], lj.mean(), se)) # each update maximises the bound over its factor: finite differences around the q(mu) update qm2 = c.vb_update_mu(x, pr, qt["a"] / qt["b"]) base = c.vb_gaussian_bound(x, pr, qm2, qt) worse = all(c.vb_gaussian_bound(x, pr, dict(m=qm2["m"] + dm, lam=qm2["lam"] * (1 + dl)), qt) < base for dm in (-1e-3, 0, 1e-3) for dl in (-1e-3, 0, 1e-3) if (dm, dl) != (0, 0)) qt2 = c.vb_update_tau(x, pr, qm) base2 = c.vb_gaussian_bound(x, pr, qm, qt2) worse2 = all(c.vb_gaussian_bound(x, pr, qm, dict(a=qt2["a"] * (1 + da), b=qt2["b"] * (1 + db))) < base2 for da in (-1e-3, 0, 1e-3) for db in (-1e-3, 0, 1e-3) if (da, db) != (0, 0)) check("VB Gaussian: (10.26), (10.27) and (10.29), (10.30) are local maxima of the bound in their factor", worse and worse2) # the fixed point (10.33) for the improper prior mu0 = a0 = b0 = lam0 = 0 imp = dict(mu0=0.0, lam0=0.0, a0=0.0, b0=0.0) qm, qt = dict(m=1.0, lam=1.0), dict(a=1.0, b=1.0) for _ in range(200): qm = c.vb_update_mu(x, imp, qt["a"] / qt["b"]) qt = c.vb_update_tau(x, imp, qm) check("VB Gaussian: with mu0 = a0 = b0 = lambda0 = 0, 1/E[tau] = (1/N) sum (x_n - xbar)^2 (10.33, corrected)", abs(qt["b"] / qt["a"] - x.var()) < 1e-12, "1/E[tau] = %.10f, variance %.10f" % (qt["b"] / qt["a"], x.var())) # --------------------------------------------------------------------------- 10.2 def mc_bound_terms(X, r, q, prior, S, rng): """Monte Carlo estimates (mean, standard error) of the expectations in (10.70).""" K, D = q["m"].shape out = {key: [] for key in ("lik", "pz", "ppi", "pml", "qpi", "qml")} lnC0 = gammaln(K * prior["alpha0"]) - K * gammaln(prior["alpha0"]) for _ in range(S): pi = rng.dirichlet(q["alpha"]) lnpi = np.log(np.maximum(pi, 1e-300)) lik = pml = qml = 0.0 for k in range(K): Lk = stats.wishart.rvs(df=q["nu"][k], scale=q["W"][k], random_state=rng) mk = rng.multivariate_normal(q["m"][k], np.linalg.inv(q["beta"][k] * Lk)) lik += np.sum(r[:, k] * common.gauss_logpdf(X, mk, np.linalg.inv(Lk))) pml += (common.gauss_logpdf(mk[None], prior["m0"], np.linalg.inv(prior["beta0"] * Lk))[0] + stats.wishart.logpdf(Lk, df=prior["nu0"], scale=prior["W0"])) qml += (common.gauss_logpdf(mk[None], q["m"][k], np.linalg.inv(q["beta"][k] * Lk))[0] + stats.wishart.logpdf(Lk, df=q["nu"][k], scale=q["W"][k])) out["lik"].append(lik) out["pz"].append(np.sum(r * lnpi)) out["ppi"].append(lnC0 + (prior["alpha0"] - 1) * lnpi.sum()) out["pml"].append(pml) out["qpi"].append(stats.dirichlet.logpdf(pi / pi.sum(), q["alpha"])) out["qml"].append(qml) return {key: (np.mean(v), np.std(v) / np.sqrt(S)) for key, v in out.items()} def checks_vbgmm(): X = common.old_faithful(standardize=True) prior = c.default_gmm_prior(2, 1e-3) res = c.vb_gmm(X, c.FIG10_6_START, prior, keep=[120]) check("VB-GMM (Fig 10.6): the bound never decreases over the iterations", np.all(np.diff(res["bounds"]) > -1e-9), "from %.3f to %.3f" % (res["bounds"][0], res["bounds"][-1])) w = c.expected_pi(res["q"]) check("VB-GMM (Fig 10.6): two components survive with alpha0 = 1e-3", np.sum(w > 1e-3) == 2, "E[pi] = %s" % np.array2string(np.sort(w)[::-1][:3], precision=4)) Nk = res["r"].sum(0) check("E[pi_k] = (alpha0 + N_k)/(K alpha0 + N) (10.69, corrected) after an M step", np.allclose(c.expected_pi(c.vbgmm_m_step(X, res["r"], prior)), (prior["alpha0"] + Nk) / (6 * prior["alpha0"] + len(X)))) # each half-step increases the bound (several random starts, K = 4) rng = np.random.default_rng(5) ok = True for _ in range(3): idx = rng.choice(len(X), 4, replace=False) r, q = c.vbgmm_init(X, idx, prior) L = c.vbgmm_bound(X, r, q, prior) for _ in range(40): q = c.vbgmm_m_step(X, r, prior) L1 = c.vbgmm_bound(X, r, q, prior) r = c.vbgmm_e_step(X, q) L2 = c.vbgmm_bound(X, r, q, prior) ok &= (L1 >= L - 1e-9) and (L2 >= L1 - 1e-9) L = L2 check("VB-GMM: every M step and every E step increases the bound (3 starts x 40 iterations)", ok) # a state that is not converged, for the finer checks prior1 = c.default_gmm_prior(2, 1.0) r, q = c.vbgmm_init(X, [10, 60, 150], prior1) for _ in range(3): q = c.vbgmm_m_step(X, r, prior1) r = c.vbgmm_e_step(X, q) q = c.vbgmm_m_step(X, r, prior1) # finite-difference stationarity of the M step (Svensen and Bishop 2004) base = c.vbgmm_bound(X, r, q, prior1) grads = {} for name in ("alpha", "beta", "nu"): for k in range(3): for sgn in (1, -1): qq = {kk: v.copy() for kk, v in q.items()} qq[name][k] += sgn * 1e-4 grads.setdefault((name, k), []).append(c.vbgmm_bound(X, r, qq, prior1)) for k in range(3): for i in range(2): for sgn in (1, -1): qq = {kk: v.copy() for kk, v in q.items()} qq["m"][k, i] += sgn * 1e-5 grads.setdefault(("m", k, i), []).append(c.vbgmm_bound(X, r, qq, prior1)) for i, j in ((0, 0), (0, 1), (1, 1)): for sgn in (1, -1): qq = {kk: v.copy() for kk, v in q.items()} qq["W"][k, i, j] += sgn * 1e-6 qq["W"][k, j, i] = qq["W"][k, i, j] grads.setdefault(("W", k, i, j), []).append(c.vbgmm_bound(X, r, qq, prior1)) g = {key: (v[0] - v[1]) for key, v in grads.items()} steps = {"alpha": 2e-4, "beta": 2e-4, "nu": 2e-4, "m": 2e-5, "W": 2e-6} gmax = max(abs(v) / steps[key[0]] for key, v in g.items()) check("VB-GMM: the M step (10.58) to (10.63) is a stationary point of the bound (finite differences)", gmax < 1e-4, "largest |dL/dtheta| = %.1e" % gmax) qbad = {kk: v.copy() for kk, v in q.items()} qbad["nu"] = qbad["nu"] + 1.0 Lbad = c.vbgmm_bound(X, r, qbad, prior1) gnu = (c.vbgmm_bound(X, r, dict(qbad, nu=qbad["nu"] + 1e-4), prior1) - c.vbgmm_bound(X, r, dict(qbad, nu=qbad["nu"] - 1e-4), prior1)) / 2e-4 check("VB-GMM: the misprinted nu_k = nu0 + N_k + 1 is not a maximiser (lower bound, non-zero slope)", Lbad < base - 1e-3 and abs(gnu) > 1e-2, "L drops by %.4f, dL/dnu = %.3f" % (base - Lbad, gnu)) # E step: (10.49) maximises the bound over r given q rE = c.vbgmm_e_step(X, q) LE = c.vbgmm_bound(X, rE, q, prior1) rng = np.random.default_rng(9) worse = True for _ in range(20): rp = rE * np.exp(1e-3 * rng.standard_normal(rE.shape)) rp /= rp.sum(1, keepdims=True) worse &= c.vbgmm_bound(X, rp, q, prior1) < LE check("VB-GMM: the responsibilities (10.49) maximise the bound given q (20 random perturbations)", worse) # (10.65) and (10.66) against Monte Carlo ln_pi, ln_lam = c.vbgmm_expectations(q) rng = np.random.default_rng(11) S = 200000 lp = np.log(rng.dirichlet(q["alpha"], S)) ok = np.all(np.abs(lp.mean(0) - ln_pi) < 4 * lp.std(0) / np.sqrt(S)) check("E[ln pi_k] = psi(alpha_k) - psi(alpha_hat) (10.66) agrees with Monte Carlo", ok, "max diff %.1e" % np.abs(lp.mean(0) - ln_pi).max()) ok = True worst = 0.0 for k in range(3): Ls = stats.wishart.rvs(df=q["nu"][k], scale=q["W"][k], size=S, random_state=rng) ld = np.linalg.slogdet(Ls)[1] ok &= abs(ld.mean() - ln_lam[k]) < 4 * ld.std() / np.sqrt(S) worst = max(worst, abs(ld.mean() - ln_lam[k])) check("E[ln|Lambda_k|] (10.65) agrees with Monte Carlo (Wishart samples)", ok, "max diff %.1e" % worst) # the bound terms (10.71) to (10.77) against Monte Carlo t = c.vbgmm_bound_terms(X, r, q, prior1) mc = mc_bound_terms(X, r, q, prior1, 3000, np.random.default_rng(13)) bad = [key for key in mc if abs(mc[key][0] - t[key]) > 4.5 * mc[key][1] + 1e-9] check("VB-GMM: the bound terms (10.71) to (10.74), (10.76), (10.77) agree with Monte Carlo", not bad, "; ".join("%s %.3f vs %.3f+-%.3f" % (key, t[key], mc[key][0], mc[key][1]) for key in mc)) # predictive density (10.81), (10.82) qf = res["q"] g = np.linspace(-9, 9, 721) G1, G2 = np.meshgrid(g, g) dens = c.vbgmm_predictive(np.c_[G1.ravel(), G2.ravel()], qf).reshape(G1.shape) tot = np.trapz(np.trapz(dens, g, axis=1), g) check("the Student-t mixture predictive density (10.81) integrates to 1", abs(tot - 1) < 1e-4, "integral %.6f" % tot) pts = np.array([[0.0, 0.0], [-1.2, -1.1], [0.8, 0.6]]) rng = np.random.default_rng(17) S = 20000 acc = np.zeros((S, len(pts))) w = c.expected_pi(qf) keep = np.where(w > 1e-3)[0] for s in range(S): val = np.zeros(len(pts)) for k in keep: Lk = stats.wishart.rvs(df=qf["nu"][k], scale=qf["W"][k], random_state=rng) mk = rng.multivariate_normal(qf["m"][k], np.linalg.inv(qf["beta"][k] * Lk)) val += w[k] * np.exp(common.gauss_logpdf(pts, mk, np.linalg.inv(Lk))) acc[s] = val pred = c.vbgmm_predictive(pts, dict(alpha=qf["alpha"][keep], beta=qf["beta"][keep], m=qf["m"][keep], W=qf["W"][keep], nu=qf["nu"][keep])) pred *= w[keep].sum() ok = np.all(np.abs(acc.mean(0) - pred) < 4 * acc.std(0) / np.sqrt(S) + 1e-12) check("the predictive (10.81), (10.82) equals the average of N(x | mu_k, Lambda_k^{-1}) under q (Monte Carlo)", ok, "formula %s, MC %s" % (np.round(pred, 5), np.round(acc.mean(0), 5))) # --------------------------------------------------------------------------- 10.3 def checks_vblr(): x, t = c.fig10_9_data() beta = 1 / c.FIG10_9_NOISE_VAR ok = True for M in range(1, 10): Phi = np.vander(x, M + 1, increasing=True) ok &= np.all(np.diff(c.vb_linear_regression(Phi, t, beta, 0.0, 0.0, E_alpha=1e-3)[-1]) > -1e-9) check("VB linear regression (Fig 10.9): the bound never decreases, every degree 1 to 9", ok) Phi = np.vander(x, 4, increasing=True) a0, b0 = 1.0, 0.5 mN, SN, aN, bN, L = c.vb_linear_regression(Phi, t, beta, a0, b0) # ln p(t) = ln int N(t | 0, I/beta + Phi Phi^T/alpha) Gam(alpha | a0, b0) d alpha, by quadrature over ln alpha def integrand(la): a = np.exp(la) C = np.eye(len(t)) / beta + Phi @ Phi.T / a C = 0.5 * (C + C.T) lnN = -0.5 * (len(t) * np.log(2 * np.pi) + np.linalg.slogdet(C)[1] + t @ np.linalg.solve(C, t)) return np.exp(lnN + stats.gamma.logpdf(a, a0, scale=1 / b0) + la - L[-1]) ev = np.log(integrate.quad(integrand, -25, 15, limit=400)[0]) + L[-1] check("VB linear regression: the bound is below ln p(t) and close to it (proper prior a0 = 1, b0 = 0.5)", L[-1] <= ev + 1e-9 and ev - L[-1] < 0.5, "L = %.5f, ln p(t) = %.5f" % (L[-1], ev)) rng = np.random.default_rng(21) S = 300000 al = rng.gamma(aN, 1 / bN, S) W = rng.multivariate_normal(mN, SN, S) M = Phi.shape[1] lj = (stats.norm.logpdf(t[None, :], W @ Phi.T, 1 / np.sqrt(beta)).sum(1) + 0.5 * M * np.log(al / (2 * np.pi)) - 0.5 * al * np.sum(W ** 2, 1) + stats.gamma.logpdf(al, a0, scale=1 / b0) - stats.multivariate_normal.logpdf(W, mN, SN) - stats.gamma.logpdf(al, aN, scale=1 / bN)) se = lj.std() / np.sqrt(S) check("VB linear regression: the bound (10.107) to (10.112) agrees with Monte Carlo", abs(lj.mean() - L[-1]) < 4 * se, "closed %.4f, MC %.4f +- %.4f" % (L[-1], lj.mean(), se)) mN0, SN0, aN0, bN0, _ = c.vb_linear_regression(Phi, t, beta, 0.0, 0.0) check("VB linear regression with a0 = b0 = 0: E[alpha] = M/(m_N^T m_N + Tr S_N) (10.104), the evidence fixed point (9.63)", abs(aN0 / bN0 - M / (mN0 @ mN0 + np.trace(SN0))) < 1e-8 * aN0 / bN0) if __name__ == "__main__": checks_factorised() checks_fig10_1() checks_vb_gaussian() checks_vbgmm() checks_vblr() print("\n%d checks: %d PASS, %d FAIL" % (len(RESULTS), sum(RESULTS), len(RESULTS) - sum(RESULTS)))