32.2 Posterior Inference Methods

1. Markov Chain Monte Carlo Fundamentals

In drug trials, climate models, and gravitational-wave astronomy, every parameter estimate carries a credible interval that shapes real decisions; if the method that produces those intervals silently fails to explore the full posterior, the reported uncertainty is wrong and downstream choices (approve the drug, evacuate the coast, point the telescope) can be dangerously miscalibrated.

What if you could explore a probability distribution you cannot even write down in closed form, by taking a carefully designed random walk through its parameter space? MCMC does exactly that: it constructs a Markov chain whose stationary distribution (the unique probability distribution that the chain converges to regardless of its starting point) is the target posterior \(p(\theta \mid D)\), so that running the chain long enough and discarding the initial burn-in (or warm-up) period yields draws from the posterior. From those draws you can compute any quantity you need: means, medians, credible intervals, predictive distributions.

MCMC algorithms generate correlated random samples from probability distributions too complex to sample directly. The posterior in Bayesian inference almost always involves an intractable integral (the marginal likelihood \(p(D)\)). MCMC sidesteps this integral by constructing a random walk that visits regions of parameter space in proportion to their posterior probability. At each step, the sampler proposes a candidate parameter value and accepts or rejects it according to a rule that forces the chain's long-run frequencies to match the target posterior, regardless of the starting point. Use MCMC when you need calibrated uncertainty estimates and your model has fewer than roughly a thousand parameters; for larger models or when speed matters more than exactness, variational inference (covered later) is a practical alternative.

Metropolis-Hastings

The Metropolis-Hastings (MH) algorithm is the simplest MCMC method. At each step, it proposes a new parameter value \(\theta'\) from a proposal distribution \(q(\theta' \mid \theta)\) (typically a Gaussian centered on the current value), then accepts or rejects the proposal with probability:

The beauty of this formula is that the intractable normalizing constant \(p(D)\) cancels in the ratio. You only need to evaluate the unnormalized posterior \(p(D \mid \theta) \, p(\theta)\), which is the product of likelihood and prior, both of which you can compute directly. Figure 32.2.1 illustrates Hamiltonian Monte Carlo trajectory geometry.

The problem with MH is that it scales poorly with dimensionality. In high dimensions, a random-walk proposal explores the posterior very slowly because most random directions move away from the high-density region. The acceptance rate drops, the chain gets stuck, and the number of samples needed to explore the posterior grows roughly as \(O(d^2)\) in the number of dimensions \(d\). For models with more than about 10 parameters, plain MH is impractical. In short: the posterior holds every answer you need, but only if your sampler actually visits the full landscape; a chain that gets stuck delivers false confidence.

import numpy as np

def metropolis_hastings(log_posterior, theta_init, n_samples, step_size, rng):
    """Minimal Metropolis-Hastings sampler for pedagogical purposes.

    Args:
        log_posterior: function theta -> log p(theta | D), up to a constant
        theta_init: starting parameter vector
        n_samples: number of samples to draw
        step_size: standard deviation of Gaussian proposal
        rng: numpy random generator
    Returns:
        samples array of shape (n_samples, dim), acceptance rate
    """
    dim = len(theta_init)
    samples = np.zeros((n_samples, dim))
    theta = theta_init.copy()
    log_p = log_posterior(theta)
    accepted = 0

    for i in range(n_samples):
        # Propose: random walk in parameter space
        theta_proposal = theta + rng.normal(0, step_size, size=dim)
        log_p_proposal = log_posterior(theta_proposal)

        # Accept/reject
        if np.log(rng.uniform()) < log_p_proposal - log_p:
            theta = theta_proposal
            log_p = log_p_proposal
            accepted += 1

        samples[i] = theta

    return samples, accepted / n_samples

2. Hamiltonian Monte Carlo

Hamiltonian Monte Carlo (HMC) overcomes the random-walk bottleneck by using gradient information to make proposals that follow the curvature of the posterior. Instead of proposing random perturbations, HMC simulates a physical system where the parameter vector \(\theta\) is the "position" of a particle and an auxiliary "momentum" variable \(\rho\) determines its velocity. The particle rolls along the posterior surface, naturally spending more time in high-density regions. Figure 32.2a contrasts the exploration patterns of MH and HMC on the same posterior.

The Hamiltonian (total energy) is:

The first term is the "potential energy" (the negative log-posterior) and the second is the "kinetic energy" with mass matrix \(M\). Hamilton's equations of motion define a trajectory:

In practice, we integrate these equations numerically using the leapfrog integrator, which preserves the symplectic structure of the Hamiltonian system (volume preservation and time reversibility). These properties guarantee that the Metropolis acceptance probability remains high even for large steps in parameter space.

def leapfrog(theta, rho, grad_log_posterior, step_size, n_steps):
    """Leapfrog integrator for Hamiltonian dynamics.

    The leapfrog scheme alternates half-steps in momentum with full steps
    in position, preserving the symplectic structure of Hamilton's equations.
    """
    rho = rho + 0.5 * step_size * grad_log_posterior(theta)  # half-step momentum

    for _ in range(n_steps - 1):
        theta = theta + step_size * rho                       # full-step position
        rho = rho + step_size * grad_log_posterior(theta)      # full-step momentum

    theta = theta + step_size * rho                           # final position step
    rho = rho + 0.5 * step_size * grad_log_posterior(theta)   # final half-step momentum

    return theta, -rho  # negate momentum for reversibility


def hmc_step(theta, log_posterior, grad_log_posterior, step_size, n_steps, rng):
    """One Hamiltonian Monte Carlo transition."""
    dim = len(theta)
    rho = rng.normal(size=dim)                          # sample momentum
    theta_new, rho_new = leapfrog(
        theta.copy(), rho.copy(), grad_log_posterior, step_size, n_steps
    )

    # Metropolis correction for numerical integration error
    current_H = -log_posterior(theta) + 0.5 * np.dot(rho, rho)
    proposed_H = -log_posterior(theta_new) + 0.5 * np.dot(rho_new, rho_new)

    if np.log(rng.uniform()) < current_H - proposed_H:
        return theta_new, True
    return theta, False

HMC has two tuning parameters: the step size \(\epsilon\) and the number of leapfrog steps \(L\). The product \(\epsilon L\) determines the total trajectory length. Too short a trajectory and HMC degenerates to a random walk. Too long and the trajectory doubles back on itself, wasting computation. Both parameters require careful tuning, which motivates the NUTS algorithm.

3. The NUTS Sampler: Geometry and Intuition

The No-U-Turn Sampler (NUTS), introduced by Hoffman and Gelman (2014), eliminates the need to choose the trajectory length \(L\) by automatically detecting when the trajectory starts to double back. The name comes from its stopping criterion: NUTS extends the trajectory in both directions (using a binary tree of leapfrog steps, where each level of the tree doubles the trajectory length by appending new leapfrog segments in both the forward and backward directions) until it makes a "U-turn," meaning the trajectory starts moving back toward its starting point.

The geometric intuition is elegant. In the Hamiltonian system, the particle follows a curved path through parameter space. On a well-conditioned posterior (roughly ellipsoidal), this path traces out an arc that covers a large portion of the posterior before returning to the neighborhood of its starting point. NUTS detects the moment the arc starts to close and stops there, giving you a proposal that has traveled as far as possible without wasting computation on the return trip.

import pymc as pm
import arviz as az
import numpy as np

# A model with a highly correlated posterior
rng = np.random.default_rng(42)
n = 50
x = rng.normal(0, 1, size=n)
y = 2.5 * x + 1.0 * (x ** 2) + rng.normal(0, 0.5, size=n)

with pm.Model() as correlated_model:
    beta_0 = pm.Normal("beta_0", mu=0, sigma=5)
    beta_1 = pm.Normal("beta_1", mu=0, sigma=5)
    beta_2 = pm.Normal("beta_2", mu=0, sigma=5)
    sigma = pm.HalfNormal("sigma", sigma=2)

    mu = beta_0 + beta_1 * x + beta_2 * x**2
    pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y)

    # NUTS is the default; these settings are for illustration
    trace = pm.sample(
        draws=2000,
        tune=1000,
        chains=4,
        target_accept=0.9,    # raise for difficult geometries
        random_seed=42,
    )

# Convergence diagnostics
print(az.summary(trace, var_names=["beta_0", "beta_1", "beta_2", "sigma"]))

# Pair plot reveals the posterior correlation structure
az.plot_pair(trace, var_names=["beta_1", "beta_2"], kind="kde")

4. Diagnosing MCMC: R-hat, ESS, and Divergences

MCMC gives you draws that are approximately from the posterior, and the quality of that approximation depends on whether the chain has converged and mixed well. Three diagnostics are essential.

R-hat (\(\hat{R}\)) compares the variance within chains to the variance between chains. If all chains have converged to the same distribution, \(\hat{R} \approx 1.0\). Values above 1.01 indicate that chains are exploring different regions of parameter space and have not yet converged. The split-\(\hat{R}\) variant (Vehtari et al., 2021) also checks for stationarity within each chain by splitting it in half.

Effective sample size (ESS) accounts for autocorrelation between successive MCMC samples. If each sample were independent, 4000 draws from 4 chains would give you 4000 effective samples. In practice, autocorrelation reduces this. An ESS below 400 is a warning sign; below 100 means the posterior summaries are unreliable. The bulk ESS measures the efficiency of exploring the center of the posterior, while the tail ESS measures the efficiency of exploring the tails (important for credible intervals).

Divergences are NUTS-specific warnings that occur when the leapfrog integrator fails to accurately follow the Hamiltonian trajectory, typically because the posterior has sharp curvature (funnels, ridges, or multimodal regions). A single divergence is a warning; many divergences mean the posterior approximation is unreliable. The fix is usually to reparameterize the model (centered vs. non-centered parameterizations for hierarchical models) or increase target_accept.

import arviz as az

def diagnose_trace(trace, var_names=None):
    """Comprehensive MCMC diagnostic report."""
    summary = az.summary(trace, var_names=var_names)

    # Check R-hat
    rhat_issues = summary[summary["r_hat"] > 1.01]
    if len(rhat_issues) > 0:
        print("WARNING: R-hat > 1.01 for:")
        print(rhat_issues[["r_hat"]])
    else:
        print("R-hat: all parameters < 1.01 (good)")

    # Check ESS
    low_ess = summary[summary["ess_bulk"] < 400]
    if len(low_ess) > 0:
        print("\nWARNING: Low bulk ESS for:")
        print(low_ess[["ess_bulk", "ess_tail"]])
    else:
        print(f"ESS: minimum bulk ESS = {summary['ess_bulk'].min():.0f} (good)")

    # Check divergences (PyMC stores these in the sample stats)
    if hasattr(trace, "sample_stats"):
        n_div = int(trace.sample_stats["diverging"].sum())
        if n_div > 0:
            print(f"\nWARNING: {n_div} divergent transitions detected!")
            print("Consider reparameterizing or increasing target_accept.")
        else:
            print("Divergences: 0 (good)")

    return summary

# Use it
summary = diagnose_trace(trace, var_names=["beta_0", "beta_1", "beta_2", "sigma"])

5. Variational Inference

Convergence diagnostics tell you whether MCMC has worked, but they cannot solve MCMC's deeper limitation: for large models or tight computational budgets, sampling may be too slow.

Variational inference (VI) takes a fundamentally different approach from MCMC. Instead of sampling from the posterior, VI approximates the posterior with a simpler distribution \(q(\theta; \phi)\) (the variational distribution) and optimizes the parameters \(\phi\) to make \(q\) as close to the true posterior as possible. The closeness measure is the Kullback-Leibler (KL) divergence, where KL divergence quantifies how much one probability distribution differs from another (zero when they match exactly, larger as they diverge):

Minimizing this KL divergence is equivalent to maximizing the evidence lower bound (ELBO):

From KL Divergence to the ELBO

The ELBO is a lower bound on \(\log p(D)\), the log marginal likelihood. Maximizing it simultaneously pushes \(q\) toward the posterior and provides an approximation to the model evidence, which is useful for model comparison (Section 32.3).

Automatic Differentiation Variational Inference (ADVI), implemented in PyMC, uses a mean-field Gaussian approximation (where "mean-field" means each parameter is given its own independent distribution, ignoring correlations between parameters): $q(\theta; \phi) = \prod_i \text{Normal}(\mu_i, \sigma_i^2)$. The optimizer finds the best \(\mu_i\) and \(\sigma_i\) for each. ADVI transforms constrained parameters (positive, bounded) to the unconstrained real line, applies the mean-field approximation, and transforms back. This is fast (seconds to minutes for models that take MCMC hours) but comes with a fundamental limitation: the mean-field approximation cannot capture correlations between parameters.

import pymc as pm
import arviz as az

with correlated_model:
    # Variational inference with ADVI
    approx = pm.fit(
        method="advi",
        n=30000,           # optimization iterations
        random_seed=42,
    )

    # Draw samples from the variational approximation
    vi_trace = approx.sample(2000)

# Compare VI and MCMC posteriors
print("MCMC posterior means:")
print(az.summary(trace, var_names=["beta_1", "beta_2"])[["mean", "sd"]])
print("\nADVI posterior means:")
print(az.summary(vi_trace, var_names=["beta_1", "beta_2"])[["mean", "sd"]])

Diagnosing VI convergence. Unlike MCMC, variational inference has no R-hat or ESS equivalent. Instead, monitor the ELBO trace (the value of the objective at each optimization step): a converged run shows the ELBO plateauing with small fluctuations. If the ELBO is still climbing when optimization stops, increase n. If it oscillates wildly, reduce the learning rate. Most importantly, compare VI posteriors against a short MCMC run on at least one representative model configuration; if the credible intervals differ substantially, the variational approximation is too restrictive for your problem.

6. Normalizing Flows for Posterior Inference

Mean-field VI sacrifices correlation structure for speed, and full-rank VI recovers correlations but cannot represent multimodality; a more expressive variational family is needed to close that gap.

Normalizing flows offer a middle ground between the exactness of MCMC and the speed of VI. A normalizing flow is an invertible neural network $f_\phi: \mathbb{R}^d \to \mathbb{R}^d$ that transforms a simple base distribution (usually a standard Gaussian) into a complex target distribution. The change-of-variables formula (which relates the density of a transformed random variable to the original density scaled by the Jacobian determinant of the transformation) gives the density of the transformed distribution:

Because the transformation is a neural network, normalizing flows can represent posteriors with correlations, multimodality, and other complex structures that mean-field VI cannot capture. Training proceeds by maximizing the ELBO, just as in VI, but with a much more expressive variational family.

The flowMC library (Wong et al., 2023) combines normalizing flows with MCMC: it trains a flow to approximate the posterior, then uses that flow as a global proposal distribution for MCMC sampling. The flow provides efficient long-range proposals; the MCMC correction preserves asymptotic exactness. This combination handles multimodal posteriors in astrophysics, gravitational-wave inference, and cosmological parameter estimation.

import jax
import jax.numpy as jnp
from flowMC.sampler.MALA import MALA
from flowMC.sampler.Sampler import Sampler
from flowMC.nfmodel.rqspline import MaskedCouplingRQSpline

def log_posterior_jax(theta, data):
    """Log-posterior for a simple model, JAX-compatible."""
    x, y = data
    beta_0, beta_1, log_sigma = theta[0], theta[1], theta[2]
    sigma = jnp.exp(log_sigma)
    mu = beta_0 + beta_1 * x
    # Log-likelihood (Normal)
    ll = -0.5 * jnp.sum(((y - mu) / sigma)**2) - n * jnp.log(sigma)
    # Log-prior (weakly informative)
    lp = -0.5 * (beta_0**2 / 100 + beta_1**2 / 100 + log_sigma**2 / 4)
    return ll + lp

# Set up the normalizing flow model
n_dim = 3
n_layers = 4
key = jax.random.PRNGKey(42)

# Rational quadratic spline flow (expressive, stable)
flow_model = MaskedCouplingRQSpline(
    n_features=n_dim,
    n_layers=n_layers,
    hidden_sizes=[64, 64],
    num_bins=8,
    key=key,
)

# Local sampler for the MCMC phase
local_sampler = MALA(log_posterior_jax, True, {"step_size": 0.01})

# Combined flow + MCMC sampler
sampler = Sampler(
    n_dim=n_dim,
    rng_key_set=jax.random.split(key, 3),
    local_sampler=local_sampler,
    nf_model=flow_model,
    n_local_steps=50,       # local MCMC steps per iteration
    n_global_steps=50,      # flow-based global steps per iteration
    n_epochs=30,            # flow training epochs
    n_chains=10,
)

7. Choosing an Inference Method

The choice of inference method depends on the model complexity, the number of parameters, the posterior geometry, and the computational budget. Table 32.2 summarizes the trade-offs.

For most scientific Bayesian models (10 to a few hundred parameters, unimodal posterior, need for calibrated uncertainty), NUTS is the default recommendation. It is the default in PyMC, Stan, and NumPyro, and it works well out of the box for the vast majority of models encountered in practice. Use ADVI as a fast screening tool when you are iterating on model structure and do not yet need precise posteriors. Use normalizing flows or flowMC when your posterior is multimodal or when you need amortized inference (where a single trained model can produce posterior samples for new datasets without rerunning the full inference procedure) across many datasets.

Exercises

1. (Conceptual) Explain why the leapfrog integrator must be both time-reversible and volume-preserving (symplectic) for HMC to satisfy detailed balance (the condition requiring that the probability of transitioning from state A to state B equals the probability of transitioning from B to A under the stationary distribution). What would go wrong if we used a standard Euler integrator instead? Hint: consider what happens to the acceptance probability.
2. (Coding) Implement the Metropolis-Hastings sampler from this section and use it to sample from a 2D Gaussian posterior with correlation 0.9. Measure the effective sample size as a function of the proposal step size. Then repeat with HMC (using the leapfrog code provided). Compare ESS per gradient evaluation for MH and HMC. At what dimensionality does HMC become clearly superior?
3. (Analysis) Take the quadratic regression model from this section and fit it with both NUTS and mean-field ADVI. Compare the posterior credible intervals for \(\beta_1\) and \(\beta_2\) from both methods. Plot the ADVI posterior as a 2D contour overlaid on the NUTS posterior samples. Where does ADVI misrepresent the posterior? How does the mismatch change if you switch to fullrank_advi?