3. Hierarchical models

The eight-schools dataset (Rubin, 1981) is a common stress test for hierarchical Bayesian inference: eight schools, one observed treatment effect and standard error per school, and the question of how much each school's effect borrows strength from the global mean. The model is small, but its centered parameterization can produce funnel-shaped posterior geometry that is difficult for mean-field VI and Hamiltonian Monte Carlo.

This chapter covers:

  • The plate-draw syntax (v : G <- Normal(0, sigma)) for per-group random effects.
  • Centered vs non-centered parameterizations and how to write each in QVR.
  • Running NUTS with R-hat, ESS, and divergence diagnostics.

Eight schools, centered

\[ \mu \sim \mathrm{Normal}(0, 5), \qquad \tau \sim \mathrm{HalfNormal}(5), \qquad \theta_j \sim \mathrm{Normal}(\mu, \tau), \qquad y_j \sim \mathrm{Normal}(\theta_j, \sigma_j). \]

The eight \(\theta_j\) are a per-group random effect over the group object School.

object School : FinSet 8
program eight_schools_centred : School -> School [effects=[Sample, Score]]
    sample mu  <- Normal(0.0, 5.0)
    sample tau <- HalfNormal(5.0)
    sample theta : School <- Normal(mu, tau)
    observe y : School <- Normal(theta, sigma_j)
    return theta

export eight_schools_centred

def eight_schools(sigma_j, y=None):
    mu  = numpyro.sample("mu",  dist.Normal(0., 5.))
    tau = numpyro.sample("tau", dist.HalfNormal(5.))
    with numpyro.plate("schools", 8):
        theta = numpyro.sample("theta", dist.Normal(mu, tau))
        numpyro.sample("y", dist.Normal(theta, sigma_j), obs=y)

The theta : School <- Normal(...) line is a plate-draw: it samples one value per element of School. The plate index is the object's cardinality (8). The observe y : School <- ... line is a vectorized observe over the same index.

The compiler synthesizes a PlateDraw morphism whose codomain is the product space School ⊗ Real; you can index into it like theta[j] inside subsequent let arithmetic.

Centered parameterization and mean-field VI

The centered parameterization puts theta_j inside the prior for mu and tau, which creates a funnel-shaped posterior (Neal, 2003, §8). Mean-field VI doesn't see the funnel and collapses to a tight Gaussian around tau ≈ 0. To confirm:

import torch
from quivers.dsl import loads
from quivers.inference import AutoNormalGuide, ELBO, SVI

CENTRED_SRC = """
object School : FinSet 8

program eight_schools_centred : School -> School
    sample mu  <- Normal(0.0, 5.0)
    sample tau <- HalfNormal(5.0)
    sample theta : School <- Normal(mu, tau)
    observe y : School <- Normal(theta, sigma_j)
    return theta

export eight_schools_centred
"""

program = loads(CENTRED_SRC)
model   = program.morphism

sigma_j = torch.tensor([[15., 10., 16., 11., 9., 11., 10., 18.]])   # (1, 8)
y_obs   = torch.tensor([[28., 8., -3., 7., -1., 1., 18., 12.]])     # (1, 8)

guide = AutoNormalGuide(model, observed_names={"y", "sigma_j"})
elbo  = ELBO(num_particles=1)
optimizer = torch.optim.Adam(
    list(model.parameters()) + list(guide.parameters()), lr=1e-2,
)
svi = SVI(model, guide, optimizer, elbo)
x_tensor = torch.zeros(1, 1)
observations = {"sigma_j": sigma_j, "y": y_obs}
for _ in range(50):                              # bump to ~3000 for real fits
    svi.step(x_tensor, observations)

This short fit may place too much mass near tau = 0, a known failure mode for a diagonal guide on funnel geometry. Compare parameterizations with a longer optimization and a diagnosed MCMC reference rather than treating a particular printed value as expected output.

Non-centered fixes it

The standard fix is to reparameterise (Papaspiliopoulos, Roberts & Sköld, 2007): draw \(\eta_j \sim \mathrm{Normal}(0, 1)\) and define \(\theta_j = \mu + \tau \cdot \eta_j\) deterministically.

object School : FinSet 8
program eight_schools_noncentred : School -> School [effects=[Sample, Score]]
    sample mu  <- Normal(0.0, 5.0)
    sample tau <- HalfNormal(5.0)
    sample eta : School <- Normal(0.0, 1.0)
    let theta = mu + tau * eta
    observe y : School <- Normal(theta, sigma_j)
    return theta

export eight_schools_noncentred

Re-running with the non-centered parameterization can improve the diagonal guide's geometry. Whether it agrees with NUTS depends on optimization and MCMC diagnostics.

NUTS

For the centered parameterization (or when you want to trust the posterior mass exactly), reach for the No-U-Turn Sampler (Hoffman & Gelman, 2014):

from quivers.inference import NUTSKernel, MCMC

kernel = NUTSKernel(
    target_accept=0.95,              # high target -> smaller step -> fewer divergences
    max_tree_depth=10,
)
# Small budget for documentation; bump warmup/samples for real fits.
mcmc = MCMC(
    kernel,
    num_warmup=50,
    num_samples=100,
    num_chains=2,
    init_strategy="prior",
)

result = mcmc.run(model, x_tensor, {"sigma_j": sigma_j, "y": y_obs})
print("posterior mean tau:", result.samples["tau"].mean().item())
print("R-hat tau:", result.r_hat["tau"].item())
print("ESS tau:", result.ess["tau"].item())
print("divergences:", result.total_divergences)

A sufficiently long, well-mixed run should have R-hat near 1 for every site (rank-normalized split-R-hat; Vehtari et al., 2021), adequate effective sample sizes, and few or no divergences. The short documentation budget above cannot establish convergence; inspect the values produced by your run.

If R-hat is above 1.01 anywhere, give NUTS more warmup steps (num_warmup=2000). If divergences are nonzero, raise target_accept to 0.99 or reparameterize to non-centered. If ESS is small (<100 per chain), the chain is mixing slowly: either run longer or reparameterize.

Centered or non-centered?

A useful rule of thumb from Betancourt & Girolami (2013): the centered parameterization is good when the data is informative about \(\theta_j\) relative to the prior, that is when the per-group likelihood scale \(\sigma_j\) is small relative to the population scale \(\tau\). The non-centered parameterization is good in the opposite regime: weak per-group likelihoods, where the posterior pinches toward the prior funnel. Eight schools sits in the weak-likelihood regime (\(\sigma_j\) ranges 9-18 against a prior \(\tau\) in single digits), so non-centered is the right call.

Init strategy and mass matrix

init_strategy="prior" draws the chain's initial state from the model prior. The other shipped options are "uniform" (uniform on the unconstrained space) and "value" (you supply a dict). "prior" is the safe default: it starts the chain in regions the prior places mass on.

NUTSKernel(mass_matrix="diagonal") adapts a diagonal mass matrix during warmup, which works for most hierarchical models. "dense" adapts a full matrix and is worth a try when latents are strongly correlated, at quadratic memory cost. "identity" skips adaptation entirely and only makes sense for already-well-scaled models.

Posterior predictive

Predictive accepts either a Guide or an MCMCResult:

from quivers.inference import Predictive

pred = Predictive(model, posterior=result, num_samples=100)
y_hat = pred(x_tensor, {"sigma_j": sigma_j})["y"]  # (100, ..., 8)
print("predictive school 1:", y_hat[..., 0].mean().item(),
      "+/-", y_hat[..., 0].std().item())

What you've seen

  • Plate-draws. v : G <- F(...) declares one draw per index of object G, producing a vector-valued latent.
  • Non-centered parameterization. A small-but-essential trick for hierarchical models; QVR doesn't automate it (yet), but writing it explicitly is two lines.
  • Diagnostics on MCMCResult. R-hat, ESS, divergences are first-class fields, not strings in a log.

Try this

  • Run AutoMultivariateNormal on the centered parameterization. It can sometimes recover the funnel where mean-field can't.
  • Change target_accept from 0.95 to 0.8 and watch divergences appear. The trade-off is step size vs trajectory length.
  • Add a per-school covariate \(x_j\) and lift the model to a varying-intercepts-and-slopes regression.

Next

Chapter 4 introduces the marginalize block: QVR's typed-scope marginalization surface for discrete latents and mixtures.

References

  • Rubin, 1981.
  • Aki Vehtari, Andrew Gelman, Daniel Simpson, Bob Carpenter, and Paul-Christian Bürkner. 2021. Rank-normalization, folding, and localization: An improved \(\hat R\) for assessing convergence of MCMC. Bayesian Analysis, 16(2):667–718.
  • Michael Betancourt and Mark Girolami. 2013. Hamiltonian Monte Carlo for hierarchical models. arXiv preprint arXiv:1312.0906.
  • Omiros Papaspiliopoulos, Gareth O. Roberts, and Martin Sköld. 2007. A general framework for the parametrization of hierarchical models. Statistical Science, 22(1):59–73.
  • Radford M. Neal. 2003. Slice sampling. The Annals of Statistics, 31(3):705–767.