Mixture-of-normals regression — Bayes factors via SMC (PyMC)¶
Normal vs mixture error: marginal likelihoods, Bayes factor, posterior model probabilities¶
Companion to mixreg_python.ipynb (from-scratch Gibbs, model comparison by WAIC) and mixreg_bayesm.ipynb (R). WAIC ranks models by predictive accuracy (≈ leave-one-out). Here we instead compute the marginal likelihood of each model and form a Bayes factor — the evidence the data give for one model over another, and the basis for posterior model probabilities.
$$\text{model } M:\quad y_t=x_t'\beta+\varepsilon_t,\qquad \varepsilon_t\sim\begin{cases}N(0,\sigma^2)&\text{(Normal)}\\ \sum_{j=1}^2 w_j\,N(\mu_j,\sigma_j^2),\ \ \textstyle\sum_j w_j\mu_j=0&\text{(mixture)}\end{cases}$$
Both models share the same prior on $\beta$ and the base scale, so the comparison is fair; the mixture nests the Normal as $w\to(1,0)$ / $\mu\to0$.
The marginal likelihood and why we use SMC¶
A Bayes factor needs the marginal likelihood (model evidence) $$Z_M=p(y\mid M)=\int p(y\mid\theta,M)\,\pi(\theta\mid M)\,d\theta,$$ the average likelihood under the prior. Crucially this is not what NUTS or the Gibbs sampler give you — they target the posterior and never compute its normalising constant.
Sequential Monte Carlo (SMC) does. It samples a sequence of tempered bridge distributions $$p_\beta(\theta)\;\propto\;\pi(\theta)\,p(y\mid\theta)^{\beta},\qquad \beta:\,0\;(\text{prior})\;\longrightarrow\;1\;(\text{posterior}),$$ advancing through adaptively-chosen stages $0=\beta_0<\beta_1<\dots<\beta_S=1$. At each stage it (1) reweights the particles by the incremental likelihood power $p(y\mid\theta)^{\beta_s-\beta_{s-1}}$, (2) resamples to cull low-weight particles, and (3) moves them with a few MCMC steps to restore diversity. The key by-product: the marginal likelihood is the product of the stagewise average weights, $$Z_M=\prod_{s=1}^{S}\frac1N\sum_{i=1}^N p\big(y\mid\theta_i^{(s-1)}\big)^{\beta_s-\beta_{s-1}},$$ a consistent estimator of the evidence. SMC also traverses multimodal posteriors (e.g. the label-switching modes of a mixture) far better than a single MCMC chain — which is exactly why it suits this problem.
pm.sample_smc returns the running $\log Z$ per chain (we take the final-stage value, average across the independent chains, and report the across-chain spread as a stability check).
From evidence to model probability. With $K$ models and prior probabilities $\pi(M_k)$, $$\text{BF}_{\text{mix},N}=\frac{Z_{\text{mix}}}{Z_N},\qquad P(M_k\mid y)=\frac{Z_k\,\pi(M_k)}{\sum_j Z_j\,\pi(M_j)}.$$ Under equally likely priors $\pi(\text{mix})=\pi(N)=\tfrac12$, this collapses to $P(\text{mix}\mid y)=\dfrac{\text{BF}}{1+\text{BF}}=\sigma(\log\text{BF})$.
Caveat (Bartlett–Lindley). Unlike WAIC, a Bayes factor is sensitive to the prior on the extra component — making that prior vaguer mechanically lowers $Z_{\text{mix}}$ (an automatic Occam penalty). We therefore also run a prior-sensitivity sweep so the BF is judged by its robustness, not a single number.
import os, sys
conda_lib = os.path.join(os.environ.get('CONDA_PREFIX',''), 'Library', 'bin')
if os.path.isdir(conda_lib) and conda_lib not in os.environ.get('PATH',''):
os.environ['PATH'] = conda_lib + ';' + os.environ.get('PATH','') # Windows DLL fix (PyMC needs this)
import numpy as np, matplotlib.pyplot as plt, pymc as pm, pytensor.tensor as pt
from scipy.special import gammaln
from scipy.io import loadmat
from scipy import stats
def smc_logZ(idata): # final-stage log marginal likelihood, per chain
return np.array([c[-1] for c in idata.sample_stats['log_marginal_likelihood'].values])
def normal_logml_analytic(y, X, sd_beta, a0, b0, ng=6000):
# exact marginal likelihood of the Normal model: integrate beta analytically, sigma^2 by 1-D quadrature
n, k = X.shape; XtX = X.T @ X; Xty = X.T @ y; yty = float(y @ y)
lg = np.linspace(np.log(1e-3), np.log(20*np.var(y)), ng); s2 = np.exp(lg); dlog = lg[1]-lg[0]
out = np.empty(ng)
for i in range(ng):
sig2 = s2[i]
C = (1/sd_beta**2)*np.eye(k) + (1/sig2)*XtX
quad = (1/sig2)*yty - (1/sig2**2)*(Xty @ np.linalg.solve(C, Xty))
logdet = n*np.log(sig2) + np.linalg.slogdet(np.eye(k) + (sd_beta**2/sig2)*XtX)[1]
logN = -0.5*(n*np.log(2*np.pi) + logdet + quad)
logIG = a0*np.log(b0) - gammaln(a0) - (a0+1)*np.log(sig2) - b0/sig2
out[i] = logN + logIG + np.log(sig2) + np.log(dlog) # d sigma^2 = sigma^2 d(log sigma^2)
m = out.max(); return m + np.log(np.exp(out-m).sum())
SD_BETA, A0, B0 = 10.0, 3.0, 3.0
def normal_model(y, X):
with pm.Model() as m:
beta = pm.Normal('beta', 0, SD_BETA, shape=X.shape[1])
sig2 = pm.InverseGamma('sig2', alpha=A0, beta=B0)
pm.Normal('y', pt.dot(X, beta), pt.sqrt(sig2), observed=y)
return m
def mixture_model(y, X, tau=2.0):
with pm.Model() as m:
beta = pm.Normal('beta', 0, SD_BETA, shape=X.shape[1])
w = pm.Dirichlet('w', a=np.ones(2))
delta = pm.Normal('delta', 0, tau) # component separation (prior scale tau)
off = pt.stack([delta, -(w[0]/w[1])*delta]) # sum_j w_j mu_j = 0 -> error mean 0
sig2 = pm.InverseGamma('sig2', alpha=A0, beta=B0, shape=2)
mu = pt.dot(X, beta)[:, None] + off[None, :]
pm.NormalMixture('y', w=w, mu=mu, sigma=pt.sqrt(sig2), observed=y)
return m
print('helpers ready')
helpers ready
1. Synthetic data (true error is a skewed 70/30 mixture)¶
Same data-generating process as the from-scratch notebook's Section 1, at $n=400$: $\beta=(1,0.5,-0.7)$, error $0.7\,N(-0.3,0.5^2)+0.3\,N(0.7,1.4^2)$.
rng = np.random.default_rng(0); n = 400
X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
beta_true = np.array([1.0, 0.5, -0.7])
z = rng.random(n) < 0.7
eps = np.where(z, rng.normal(-0.3, 0.5, n), rng.normal(0.7, 1.4, n))
y = X @ beta_true + eps
zN = smc_logZ(pm.sample_smc(model=normal_model(y, X), draws=2000, chains=4, cores=1, progressbar=False, random_seed=1))
zM = smc_logZ(pm.sample_smc(model=mixture_model(y, X), draws=2000, chains=4, cores=1, progressbar=False, random_seed=2))
zN_exact = normal_logml_analytic(y, X, SD_BETA, A0, B0)
print('Normal log Z: SMC %.2f (chains sd %.2f) analytic %.2f <- SMC validated' % (zN.mean(), zN.std(), zN_exact))
print('Mixture log Z: SMC %.2f (chains sd %.2f)' % (zM.mean(), zM.std()))
logBF = zM.mean() - zN.mean()
print('\nlog BF (mixture vs Normal) = %.1f BF = %.2e' % (logBF, np.exp(logBF)))
print('Under equal priors: P(mixture | y) = %.5f P(Normal | y) = %.2e' % (1/(1+np.exp(-logBF)), 1/(1+np.exp(logBF))))
Initializing SMC sampler...
Sampling 4 chains sequentially
Initializing SMC sampler...
Sampling 4 chains sequentially
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
Normal log Z: SMC -579.84 (chains sd 0.05) analytic -579.85 <- SMC validated Mixture log Z: SMC -539.08 (chains sd 0.35) log BF (mixture vs Normal) = 40.8 BF = 5.01e+17 Under equal priors: P(mixture | y) = 1.00000 P(Normal | y) = 2.00e-18
2. Prior-sensitivity of the Bayes factor¶
We vary $\tau$, the prior SD on the component-separation $\delta$ (a vaguer prior on the extra component → larger Occam penalty → smaller $Z_{\text{mix}}$). The magnitude of the BF moves with $\tau$; the question is whether the conclusion is robust.
taus = np.array([0.5, 1.0, 2.0, 4.0, 8.0]); logBFs = []
for i, tau in enumerate(taus):
zMi = smc_logZ(pm.sample_smc(model=mixture_model(y, X, tau=tau), draws=1500, chains=4, cores=1, progressbar=False, random_seed=10+i))
logBFs.append(zMi.mean() - zN.mean())
logBFs = np.array(logBFs)
for tau, lb in zip(taus, logBFs):
print('tau=%4.1f log BF = %6.1f P(mix|y) = %.4f' % (tau, lb, 1/(1+np.exp(-lb))))
fig, ax = plt.subplots(1, 2, figsize=(13, 4.4))
ax[0].bar(['Normal','Mixture'], [zN.mean(), zM.mean()], color=['steelblue','firebrick'])
ax[0].set_ylabel('log marginal likelihood (log Z)')
ax[0].set_title('Model evidence (gap = log BF = %.0f)'%logBF)
ax[0].set_ylim(min(zN.mean(), zM.mean())-15, max(zN.mean(), zM.mean())+5)
ax[1].plot(taus, logBFs, 'o-', color='firebrick'); ax[1].set_xscale('log')
ax[1].axhline(0, color='gray', ls='--', label='equipoise (BF=1)')
ax[1].set_xlabel('prior SD τ on the extra component'); ax[1].set_ylabel('log BF (mixture vs Normal)')
ax[1].set_title('Prior sensitivity — magnitude moves, conclusion robust'); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.savefig('mixreg_bayesfactor.png', dpi=120, bbox_inches='tight'); plt.show()
Initializing SMC sampler...
Sampling 4 chains sequentially
Initializing SMC sampler...
Sampling 4 chains sequentially
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters. A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Initializing SMC sampler...
Sampling 4 chains sequentially
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters. A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Initializing SMC sampler...
Sampling 4 chains sequentially
Initializing SMC sampler...
Sampling 4 chains sequentially
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters. A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
tau= 0.5 log BF = 41.9 P(mix|y) = 1.0000 tau= 1.0 log BF = 41.4 P(mix|y) = 1.0000 tau= 2.0 log BF = 41.1 P(mix|y) = 1.0000 tau= 4.0 log BF = 39.9 P(mix|y) = 1.0000 tau= 8.0 log BF = 37.8 P(mix|y) = 1.0000
Left — the model evidence (log marginal likelihood) for the Normal vs the mixture on the synthetic data; the gap between the bars is the log Bayes factor (≈41 → BF ~5×10¹⁷ → P(mixture) ≈ 1). Right — sweeping the prior SD τ on the extra component: the magnitude of the log BF falls as τ grows (the Bartlett–Lindley Occam penalty), but it stays far above the equipoise line (BF = 1) across two orders of magnitude in τ. The number is prior-dependent; the conclusion is not — which is the honest way to report a mixture Bayes factor.
3. PSID earnings (real data)¶
The same log-earnings regression as the from-scratch notebook's Section 2 (positive earners; age, age², education). The OLS residual has kurtosis ≈ 23 — strongly non-normal — so the evidence should overwhelmingly favour the mixture (consistent with the ΔWAIC ≈ 1600 reported there).
M = loadmat('psid.mat')['psid']; age, educ, earn = M[:,0], M[:,1], M[:,2]; pos = earn > 0
y2 = np.log(earn[pos]); a = (age[pos]-age[pos].mean())/age[pos].std(); e = (educ[pos]-educ[pos].mean())/educ[pos].std()
a2 = (a**2 - (a**2).mean())/(a**2).std()
X2 = np.column_stack([np.ones(pos.sum()), a, a2, e])
print('PSID n=%d OLS residual kurtosis=%.1f' % (pos.sum(), stats.kurtosis(y2 - X2 @ np.linalg.lstsq(X2, y2, rcond=None)[0])))
zN2 = smc_logZ(pm.sample_smc(model=normal_model(y2, X2), draws=1500, chains=2, cores=1, progressbar=False, random_seed=21))
zM2 = smc_logZ(pm.sample_smc(model=mixture_model(y2, X2), draws=1500, chains=2, cores=1, progressbar=False, random_seed=22))
logBF2 = zM2.mean() - zN2.mean()
print('Normal log Z %.1f Mixture log Z %.1f log BF = %.1f P(mixture|y) = %.6f'
% (zN2.mean(), zM2.mean(), logBF2, 1/(1+np.exp(-min(logBF2,700)))))
Initializing SMC sampler...
Sampling 2 chains sequentially
PSID n=2733 OLS residual kurtosis=22.8
We recommend running at least 4 chains for robust computation of convergence diagnostics
Initializing SMC sampler...
Sampling 2 chains sequentially
We recommend running at least 4 chains for robust computation of convergence diagnostics
Normal log Z -3810.1 Mixture log Z -3032.9 log BF = 777.2 P(mixture|y) = 1.000000
Results¶
- SMC is validated: its Normal log-evidence matches the closed-form value to ~0.1, with tiny across-chain spread — so we trust the mixture evidence too.
- Synthetic ($n=400$): $\log\text{BF}\approx41$ → $\text{BF}\sim10^{17}$ → $P(\text{mixture}\mid y)\approx1.0000$ under equal priors. Even this modest sample decisively prefers the mixture.
- Prior sensitivity: as $\tau$ grows the log BF falls (the Bartlett–Lindley Occam penalty) — the number is prior-dependent — but it stays far above 0 across two orders of magnitude in $\tau$: the conclusion is robust even though the magnitude is not. (This is the honest way to report a mixture Bayes factor.)
- PSID earnings: the evidence is overwhelming ($\log\text{BF}$ in the hundreds, $P(\text{mixture}\mid y)\approx1$) — the Bayes-factor verdict agrees with the WAIC verdict from the from-scratch notebook.
- WAIC vs Bayes factor: WAIC asks "which predicts new data better" (≈ LOO), the Bayes factor asks "which model is more probable given the data" (marginal likelihood). Here they agree emphatically; in general they can differ, and the BF carries the extra prior-sensitivity caveat shown above.