Causal Inference II(c) — Hierarchical Models for Clustered Assignment¶
Project STAR: what a cluster-robust standard error patches, and what a model explains¶
The noncompliance and clusters notebook established the problem on Project STAR. Students in the same school share a teacher, a building and a peer group, so their outcomes are correlated and the row count badly overstates how much independent information there is. Ignoring that gives a standard error of 1.042; clustering by school gives 1.850, nearly 1.8× wider. A known-truth simulation in that notebook showed why it matters: naive 95% intervals covered only 54%.
The cluster-robust standard error is the right fix and it is a patch. It corrects the uncertainty on one number and tells you nothing else. It does not say how much schools differ, it cannot produce an estimate for any individual school, and it cannot address whether the small-class effect is the same everywhere — which for a policy that would be rolled out school by school is arguably the more useful question.
A hierarchical model treats the clustering as structure rather than nuisance, and produces four things the published analysis could not:
- The intraclass correlation as a parameter with uncertainty, rather than the 0.225 point proxy.
- School-level intercepts — where a child starts, before any treatment.
- A varying treatment effect, testing whether small classes help more in some schools than others.
- An uncertainty interval that comes from the model rather than from a sandwich correction bolted onto it.
Python/PyMC lead. Third in the Bayesian causal group, and the point where hierarchy stops being a convenience and becomes the model of the design itself.
1. The data, and the frequentist benchmark to beat¶
import numpy as np, pandas as pd, matplotlib.pyplot as plt, warnings
warnings.filterwarnings("ignore")
import pymc as pm, arviz as az
import statsmodels.api as sm
import statsmodels.formula.api as smf
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"
d = pd.read_csv("star_k.csv").dropna(subset=["small","read","school"])
d["school"] = d["school"].astype(int)
sch = np.sort(d.school.unique()); S = len(sch)
sidx = pd.Categorical(d.school, categories=sch).codes
print(f"{len(d):,} students in {S} schools; {int(d.small.sum()):,} in small classes")
print(f"school sizes: min {d.groupby('school').size().min()}, median {int(d.groupby('school').size().median())}, max {d.groupby('school').size().max()}")
ols = smf.ols("read ~ small", data=d).fit()
hc1 = smf.ols("read ~ small", data=d).fit(cov_type="HC1")
cl = smf.ols("read ~ small", data=d).fit(cov_type="cluster", cov_kwds={"groups": d.school})
print(f"\nsmall-class effect on kindergarten reading = {ols.params['small']:.2f} points")
print(f" naive (HC1) SE = {hc1.bse['small']:.3f} 95% CI [{hc1.conf_int().loc['small',0]:.2f}, {hc1.conf_int().loc['small',1]:.2f}]")
print(f" cluster-robust SE = {cl.bse['small']:.3f} 95% CI [{cl.conf_int().loc['small',0]:.2f}, {cl.conf_int().loc['small',1]:.2f}] ({cl.bse['small']/hc1.bse['small']:.2f}x wider)")
# the ICC proxy the frequentist page reports
grand = d.read.mean(); gm = d.groupby("school").read.agg(["mean","size"])
between = ((gm["mean"]-grand)**2 * gm["size"]).sum()/len(d)
icc_proxy = between/d.read.var(ddof=0)
print(f"\nbetween-school variance share (the ICC proxy) = {icc_proxy:.3f} -- a point estimate with no interval")
g++ not available, if using conda: `conda install gxx`
3,743 students in 79 schools; 1,738 in small classes school sizes: min 13, median 43, max 94 small-class effect on kindergarten reading = 5.82 points naive (HC1) SE = 1.042 95% CI [3.78, 7.86] cluster-robust SE = 1.850 95% CI [2.19, 9.44] (1.78x wider) between-school variance share (the ICC proxy) = 0.215 -- a point estimate with no interval
2. The hierarchical model — clustering as structure¶
Instead of correcting a standard error after the fact, put the school in the model:
$$\text{read}_i = \alpha_{s[i]} + \tau\,\text{small}_i + \varepsilon_i, \qquad \alpha_s \sim \mathcal{N}(\mu_\alpha,\ \sigma_\alpha^2), \qquad \varepsilon_i \sim \mathcal{N}(0,\ \sigma_y^2).$$
Now $\sigma_\alpha$ — how much schools differ — is a parameter, and the intraclass correlation follows directly as $\sigma_\alpha^2/(\sigma_\alpha^2+\sigma_y^2)$, with a full posterior rather than a single number. The treatment effect $\tau$ is estimated within schools, which is the comparison the randomization actually supports.
small = d.small.values.astype(float); y = d.read.values.astype(float)
with pm.Model() as m_vi:
mu_a = pm.Normal("mu_a", 450, 50)
sigma_a = pm.HalfNormal("sigma_a", 50)
za = pm.Normal("za", 0, 1, shape=S)
alpha = pm.Deterministic("alpha", mu_a + sigma_a*za)
tau = pm.Normal("tau", 0, 25)
sigma_y = pm.HalfNormal("sigma_y", 50)
pm.Normal("obs", alpha[sidx] + tau*small, sigma_y, observed=y)
icc = pm.Deterministic("icc", sigma_a**2/(sigma_a**2 + sigma_y**2))
id_vi = pm.sample(1500, tune=1500, chains=4, cores=1, target_accept=0.9,
progressbar=False, random_seed=0)
s1 = az.summary(id_vi, var_names=["tau","mu_a","sigma_a","sigma_y","icc"])
cols = [c for c in ["mean","sd"] + [c for c in s1.columns if c.startswith("hdi")] + ["r_hat"] if c in s1.columns]
print(s1[cols].to_string())
t_vi = id_vi.posterior["tau"].values.ravel(); ic = id_vi.posterior["icc"].values.ravel()
print(f"\ntau = {t_vi.mean():.2f}, 95% [{np.percentile(t_vi,2.5):.2f}, {np.percentile(t_vi,97.5):.2f}]"
f" (cluster-robust gave [{cl.conf_int().loc['small',0]:.2f}, {cl.conf_int().loc['small',1]:.2f}])")
print(f"ICC = {ic.mean():.3f}, 95% [{np.percentile(ic,2.5):.3f}, {np.percentile(ic,97.5):.3f}]"
f" (the frequentist page reported {icc_proxy:.3f} with no interval)")
print(f"\ndivergences = {int(id_vi.sample_stats.diverging.values.sum())}")
Initializing NUTS using jitter+adapt_diag...
Sequential sampling (4 chains in 1 job)
NUTS: [mu_a, sigma_a, za, tau, sigma_y]
Sampling 4 chains for 1_500 tune and 1_500 draw iterations (6_000 + 6_000 draws total) took 6 seconds.
mean sd r_hat tau 6.54 0.94 1.00 mu_a 434.1 1.8 1.01 sigma_a 14.75 1.31 1.00 sigma_y 28.311 0.332 1.00 icc 0.214 0.03 1.00 tau = 6.54, 95% [4.63, 8.38] (cluster-robust gave [2.19, 9.44]) ICC = 0.214, 95% [0.160, 0.279] (the frequentist page reported 0.215 with no interval) divergences = 0
3. Does the effect itself vary by school?¶
The model so far lets schools differ in their level — some schools read better than others — but forces the small-class benefit to be identical everywhere. That is a strong assumption, and for a policy question it is the interesting one: rolling small classes out across a state assumes the effect travels.
Adding a varying slope tests it:
$$\text{read}_i = \alpha_{s[i]} + \tau_{s[i]}\,\text{small}_i + \varepsilon_i, \qquad \tau_s \sim \mathcal{N}(\mu_\tau,\ \sigma_\tau^2).$$
$\sigma_\tau$ now measures how much the effect differs across schools. A posterior concentrated near zero says one number describes every school and the policy travels; a posterior clearly away from zero says the average conceals real variation and some schools benefit far more than others.
with pm.Model() as m_vs:
mu_a = pm.Normal("mu_a", 450, 50)
sigma_a = pm.HalfNormal("sigma_a", 50)
za = pm.Normal("za", 0, 1, shape=S)
alpha = pm.Deterministic("alpha", mu_a + sigma_a*za)
mu_t = pm.Normal("mu_t", 0, 25)
sigma_t = pm.HalfNormal("sigma_t", 15)
zt = pm.Normal("zt", 0, 1, shape=S)
tau_s = pm.Deterministic("tau_s", mu_t + sigma_t*zt)
sigma_y = pm.HalfNormal("sigma_y", 50)
pm.Normal("obs", alpha[sidx] + tau_s[sidx]*small, sigma_y, observed=y)
id_vs = pm.sample(1500, tune=2000, chains=4, cores=1, target_accept=0.95,
progressbar=False, random_seed=1)
s2 = az.summary(id_vs, var_names=["mu_t","sigma_t","sigma_a","sigma_y"])
cols2 = [c for c in ["mean","sd"] + [c for c in s2.columns if c.startswith("hdi")] + ["r_hat"] if c in s2.columns]
print(s2[cols2].to_string())
mt = id_vs.posterior["mu_t"].values.ravel(); st = id_vs.posterior["sigma_t"].values.ravel()
print(f"\nmu_t (average effect) = {mt.mean():.2f}, 95% [{np.percentile(mt,2.5):.2f}, {np.percentile(mt,97.5):.2f}]")
print(f"sigma_t (spread of effects) = {st.mean():.2f}, 95% [{np.percentile(st,2.5):.2f}, {np.percentile(st,97.5):.2f}]")
print(f" P(sigma_t < 2 points) = {(st<2).mean():.3f}")
print(f" P(sigma_t < 5 points) = {(st<5).mean():.3f}")
print(f"divergences = {int(id_vs.sample_stats.diverging.values.sum())}")
ts = id_vs.posterior["tau_s"].values.reshape(-1, S)
sch_mean = ts.mean(0)
print(f"\nschool-level effects: min {sch_mean.min():.2f}, median {np.median(sch_mean):.2f}, max {sch_mean.max():.2f}")
nz = int(((np.percentile(ts,2.5,axis=0)>0)|(np.percentile(ts,97.5,axis=0)<0)).sum())
print(f" schools whose 95% interval excludes zero: {nz} of {S}")
print(f" by chance alone you would expect about {0.05*S:.0f}, so {nz} is well above a fluke")
print()
print("One caveat this deserves. The median school has only", int(d.groupby('school').size().median()), "students, so any single school's")
print("effect is poorly determined on its own -- which is exactly why they are partially pooled. sigma_t is")
print("identified from the SPREAD across schools rather than from any one of them, and a spread that large")
print("would be hard to manufacture from noise alone at this many schools. The right reading is that the")
print("small-class effect genuinely differs by school; the wrong one is to trust any individual school's number.")
Initializing NUTS using jitter+adapt_diag...
Sequential sampling (4 chains in 1 job)
NUTS: [mu_a, sigma_a, za, mu_t, sigma_t, zt, sigma_y]
Sampling 4 chains for 2_000 tune and 1_500 draw iterations (8_000 + 6_000 draws total) took 16 seconds.
mean sd r_hat mu_t 6.57 1.64 1.00 sigma_t 11.89 1.44 1.00 sigma_a 15.62 1.44 1.00 sigma_y 27.659 0.319 1.00 mu_t (average effect) = 6.57, 95% [3.33, 9.92] sigma_t (spread of effects) = 11.89, 95% [9.24, 14.85] P(sigma_t < 2 points) = 0.000 P(sigma_t < 5 points) = 0.000 divergences = 0 school-level effects: min -16.00, median 5.56, max 37.52 schools whose 95% interval excludes zero: 18 of 79 by chance alone you would expect about 4, so 18 is well above a fluke One caveat this deserves. The median school has only 43 students, so any single school's effect is poorly determined on its own -- which is exactly why they are partially pooled. sigma_t is identified from the SPREAD across schools rather than from any one of them, and a spread that large would be hard to manufacture from noise alone at this many schools. The right reading is that the small-class effect genuinely differs by school; the wrong one is to trust any individual school's number.
4. The three analyses side by side¶
rows = [
("OLS, naive SE", ols.params["small"], hc1.bse["small"],
hc1.conf_int().loc["small",0], hc1.conf_int().loc["small",1]),
("OLS, cluster-robust SE", ols.params["small"], cl.bse["small"],
cl.conf_int().loc["small",0], cl.conf_int().loc["small",1]),
("hierarchical, varying intercept", t_vi.mean(), t_vi.std(),
np.percentile(t_vi,2.5), np.percentile(t_vi,97.5)),
("hierarchical, + varying slope", mt.mean(), mt.std(),
np.percentile(mt,2.5), np.percentile(mt,97.5)),
]
print(f" {'analysis':34s} {'effect':>8} {'SE/sd':>8} {'95% interval':>22} {'width':>7}")
for nm,e,s_,l,h in rows:
print(f" {nm:34s} {e:>8.2f} {s_:>8.3f} {'[%6.2f, %6.2f]'%(l,h):>22} {h-l:>7.2f}")
print()
print("Read the widths, not just the intervals. The varying-INTERCEPT model is the TIGHTEST of the four --")
print(f"width {rows[2][4]-rows[2][3]:.2f} against cluster-robust's {rows[1][4]-rows[1][3]:.2f} -- and that is not a free improvement.")
print("It is tighter because it assumes one common treatment effect, and section 3 has just shown that")
print(f"assumption is false: sigma_t = {st.mean():.1f} points, with P(sigma_t < 5) = {(st<5).mean():.3f}.")
print("Allow the effect to vary and the interval widens back out to", f"{rows[3][4]-rows[3][3]:.2f}, close to cluster-robust.")
print()
print("So the cluster-robust standard error was RIGHT to be wide. It makes no assumption about how schools")
print("differ, and therefore quietly absorbed heterogeneity the simpler hierarchical model wished away. A")
print("hierarchical model is only better than a robust standard error when its structure is correct; here")
print("the first version was tighter AND wrong, and it took the second version to see it.")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
for i,(nm,e,s_,l,h) in enumerate(rows):
ax[0].errorbar(e, i, xerr=[[e-l],[h-e]], fmt="o", capsize=4,
color=[GREY,ORANGE,BLUE,PURP][i])
ax[0].set_yticks(range(len(rows))); ax[0].set_yticklabels([r[0] for r in rows], fontsize=8)
ax[0].axvline(0, color="k", lw=.6); ax[0].set_xlabel("small-class effect (reading points)")
ax[0].set_title("Same effect, four accounts of the uncertainty"); ax[0].invert_yaxis()
order = np.argsort(sch_mean)
ax[1].errorbar(range(S), sch_mean[order],
yerr=[sch_mean[order]-np.percentile(ts,2.5,axis=0)[order],
np.percentile(ts,97.5,axis=0)[order]-sch_mean[order]],
fmt=".", color=BLUE, alpha=.55, capsize=0, elinewidth=.8)
ax[1].axhline(mt.mean(), color=RED, lw=1.5, label=f"average {mt.mean():.2f}")
ax[1].axhline(0, color="k", lw=.6)
ax[1].set_xlabel("school (sorted by posterior effect)"); ax[1].set_ylabel("small-class effect")
ax[1].set_title(f"{S} school-level effects the frequentist analysis cannot report"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()
analysis effect SE/sd 95% interval width OLS, naive SE 5.82 1.042 [ 3.78, 7.86] 4.08 OLS, cluster-robust SE 5.82 1.850 [ 2.19, 9.44] 7.25 hierarchical, varying intercept 6.54 0.937 [ 4.63, 8.38] 3.74 hierarchical, + varying slope 6.57 1.640 [ 3.33, 9.92] 6.59 Read the widths, not just the intervals. The varying-INTERCEPT model is the TIGHTEST of the four -- width 3.74 against cluster-robust's 7.25 -- and that is not a free improvement. It is tighter because it assumes one common treatment effect, and section 3 has just shown that assumption is false: sigma_t = 11.9 points, with P(sigma_t < 5) = 0.000. Allow the effect to vary and the interval widens back out to 6.59, close to cluster-robust. So the cluster-robust standard error was RIGHT to be wide. It makes no assumption about how schools differ, and therefore quietly absorbed heterogeneity the simpler hierarchical model wished away. A hierarchical model is only better than a robust standard error when its structure is correct; here the first version was tighter AND wrong, and it took the second version to see it.
5. Summary¶
The cluster-robust standard error and the hierarchical model are not competitors on the same question — they answer different numbers of questions.
- On the headline effect they broadly agree, but the widths tell a story. The varying-intercept model produces the tightest interval of the four — and it earns that tightness by assuming one common effect, which the varying-slope model then refutes. Allow the effect to vary and the interval widens back to roughly cluster-robust's. So the cluster-robust standard error was right to be wide: making no assumption about how schools differ, it absorbed heterogeneity the simpler model wished away. A hierarchical model beats a robust standard error only when its structure is correct.
- The ICC becomes a parameter. The frequentist page reports a variance-share proxy as a single number. The model returns a posterior, so "strong clustering" acquires an interval instead of resting on a point estimate nobody could put error bars on.
- School-level effects exist at all. Seventy-nine of them, each partially pooled toward the average by how little that school's own data support it. No sandwich correction produces these, and for a policy delivered school by school they are the quantity a decision-maker actually wants.
- The constant-effect assumption becomes testable, and it fails. $\sigma_\tau$ comes out at 11.9 points with $\Pr(\sigma_\tau<5)=0.000$, against an average effect of 6.6 — the variation across schools is larger than the average benefit. Eighteen of seventy-nine schools have intervals excluding zero where chance would give about four. Small classes do not help equally everywhere, and that question is invisible to an analysis whose entire treatment of clustering is a variance correction.
The honest framing is that cluster-robust inference is the right tool when the effect is all you want and you distrust your model, since it needs no assumption about how schools vary. The hierarchical model asks for more — a distributional assumption on the school effects — and returns much more in exchange. Which is the better trade depends on whether you want a defensible number or a description of the mechanism.
Cross-links. This is the same partial-pooling machinery as the Electric Company example, scaled from four grades to seventy-nine schools and applied to grouping that was a nuisance rather than a design feature. The frequentist treatment supplies the simulation showing what naive standard errors cost, which no amount of hierarchical modelling replaces.