MTGP – Multitask Gaussian Process Synthetic Control (Ben-Michael et al. 2023)

MTGP – Multitask Gaussian Process Synthetic Control (Ben-Michael et al. 2023)#

Estimator:

Multitask Gaussian Process Synthetic Control (MTGP) – mlsynth.MTGP

Source:

Ben-Michael, E., Arbour, D., Feller, A., Franks, A., & Raphael, S. (2023), “Estimating the effects of a California gun control program with multitask Gaussian processes,” Annals of Applied Statistics 17(2), 985–1016 [MTGP2023].

Replication type:

cross-validation against the authors’ own reference (the replication package’s Stan program) on the paper’s California panel.

Status:

verified – the NumPyro posterior matches the Stan posterior cell-for-cell (to Monte-Carlo error).

Validation strategy#

Ben-Michael et al. (2023) ship their model as a Stan program in the paper’s replication package. MTGP ports that program to NumPyro (both are NUTS samplers), so the Stan program is the ground truth. We transcribe the Gaussian (normal.stan) model verbatim, run it via rstan on the study’s California panel, and compare the posterior counterfactual of California to the NumPyro estimator on the identical data.

The panel is the paper’s own: gun-homicide rates per 100,000 across the 50 states, 1997–2018 (22 years), with California treated in 2007 – the year the Armed and Prohibited Persons System (APPS) began removing firearms from prohibited owners. Population enters the noise model, so a small state’s noisy rate is downweighted relative to a large one.

Cross-validation – cell for cell#

Both samplers use the same rank-5 coregionalization, squared-exponential time kernels for the global trend and the factors, population-scaled Gaussian noise, and matched NUTS budgets. The posterior counterfactual of California agrees:

Quantity

NumPyro vs Stan

counterfactual, posterior mean (all 22 years)

max \(|\Delta|\) = 0.2% of level; corr 0.99993

mean post-2007 ATT (obs – counterfactual)

\(-1.029\) vs \(-1.031\) per 100,000 (agree to 0.002)

length-scales / noise scale

\(\rho_f, \rho_g, \sigma\) agree to ~2%

The residual is the Monte-Carlo difference between two independent NUTS runs, not a specification gap. A coregionalization model is rotation/sign non-identified, so the raw loadings do not mix (Stan flags this too); the identified quantities – the counterfactual, the ATT, and the length-scales – mix cleanly and match.

The one build note that matters#

The GP Cholesky needs double precision. JAX defaults to float32, whose machine epsilon (\(\approx 10^{-7}\)) is larger than the \(10^{-9}\) jitter added to the kernel diagonal, so in single precision the jitter is lost and the Cholesky factorizes a numerically indefinite matrix – the sampler then produces garbage (negative rates, near-zero correlation with Stan). The fix is a single line, numpyro.enable_x64(), called in the estimator’s backend shim before the model is built. With float64 the port matches Stan to correlation \(0.99993\) and R-hat \(1.001\)–\(1.005\) on the identified quantities. This is the standard trap for GP models on JAX and is the reason the estimator hard-enables x64 rather than leaving it to the caller.

Why NumPyro, not pure numpy#

The model marginalizes a Gaussian process whose kernel hyperparameters (length-scales, marginal scales) are themselves sampled, so the geometry is a correlated, heavy-tailed hierarchy – exactly what a hand-written Gibbs sampler struggles with (the kernel conditionals go ill-conditioned as the length-scales move). NUTS handles it, so MTGP uses NumPyro behind the [bayes] optional dependency, following Bayesian Factor Synthetic Control (BFSC) and Spillover-Detecting Synthetic Control (SPOTSYNTH). The durable case is benchmarks/cases/mtgp_california.py; it feeds the identical California panel to both mlsynth.MTGP and the compiled Stan reference and reports the cell-for-cell agreement, requiring the [bayes] extra for mlsynth and rstan for the reference, and skipping gracefully when either is absent.