Multitask Gaussian Process Synthetic Control (MTGP)

Multitask Gaussian Process Synthetic Control (MTGP)#

When to use MTGP – and when not to#

Multitask Gaussian Process synthetic control (Ben-Michael et al. 2023) models the no-intervention outcome of every unit as a Gaussian process whose kernel is separable over time and units, then reads the treated unit’s counterfactual off the posterior. Reach for it when the untreated series are smooth functions of time – trends and slow drifts rather than sharp jumps – and you want a counterfactual whose credible band grows the further you extrapolate past the intervention. It is a good choice when:

  • You want a full posterior credible band on the counterfactual and the ATT, and you want that band to widen honestly the further the post-period runs from the last observed treated data – the signature Gaussian-process behavior.

  • The untreated outcomes are smooth in time. A GP with a squared-exponential kernel encodes nearby years look alike; it borrows strength across time, not only across donors.

  • Units differ in size and you want the noise to reflect it. Given a population (or exposure) column, MTGP scales the observation noise as \(\sqrt{1/\text{pop}}\), so a small state’s noisy rate is downweighted relative to a large one, exactly as in the paper’s rate data.

Do not use MTGP when:

Relation to BFSC#

MTGP and Bayesian Factor Synthetic Control (BFSC) are close cousins: both mask the treated post-period cells, impute them from a Bayesian latent-factor model fit on the donors and the treated pre-period, and read the counterfactual off the posterior. They differ in what the factors are allowed to look like over time.

  • Bayesian Factor Synthetic Control (BFSC) (Pinkney) places no smoothness prior on the factors – each period is a free parameter. This is MTGP with an identity time kernel.

  • MTGP (Ben-Michael et al.) puts a Gaussian-process prior on the factors, so the factor paths are smooth functions of time. It adds a shared global time-GP trend on top of the low-rank factor term, and scales the noise by population.

The practical consequence is the post-period band. Because BFSC’s factors are unconstrained, its post-period imputation is pinned only by the donors; because MTGP’s factors are smooth, its band grows with distance from the last treated observation – the further you extrapolate, the less the GP is willing to commit. Reach for BFSC when the shared structure is best left unconstrained; reach for MTGP when the untreated outcomes are genuinely smooth in time and you want the band to reflect extrapolation distance.

Notation#

We use the synthetic-control canon. Let \(j = 1\) denote the treated unit and \(\mathcal{N}_0\) the donor pool, over units \(\mathcal{N} \coloneqq \{1, \ldots, N\}\); time is \(t \in \mathcal{T} \coloneqq \{1, \ldots, T\}\) with the intervention after \(T_0\), pre-period \(\mathcal{T}_1\) and post-period \(\mathcal{T}_2\). MTGP models the no-intervention outcome as

\[y_{jt}^N = \alpha + \kappa_j + g_t + \sum_{r=1}^{R} \ell_{tr}\,k_{rj} + u_{jt}, \qquad u_{jt} \sim \mathcal{N}\!\left(0,\ \sigma^2 / \text{pop}_{jt}\right),\]

with a global intercept \(\alpha\), unit intercepts \(\kappa_j\), a shared global time trend \(g_t\), and a rank-\(R\) factor term (the paper’s \(R\) is n_factors) built from latent time-GP factors \(\ell_{tr}\) and unit coregionalization loadings \(k_{rj}\).

Gaussian-process priors. The global trend \(\mathbf{g} = (g_1,\ldots,g_T)\) and each latent factor \(\boldsymbol{\ell}_r = (\ell_{1r},\ldots,\ell_{Tr})\) are drawn from mean-zero Gaussian processes over time with squared-exponential kernels,

\[\operatorname{Cov}(g_s, g_t) = \sigma_g^2 \exp\!\left(-\frac{(s-t)^2}{2\rho_g^2}\right), \qquad \operatorname{Cov}(\ell_{sr}, \ell_{tr}) = \sigma_f^2 \exp\!\left(-\frac{(s-t)^2}{2\rho_f^2}\right),\]

so nearby periods are correlated and the factor paths are smooth. The length-scales \(\rho_f, \rho_g\) carry InverseGamma priors; the marginal scales \(\sigma_f, \sigma_g\) and the noise scale \(\sigma\) carry half-normal priors. The rank-\(R\) term is an intrinsic-coregionalization model: a single set of \(R\) shared time factors, mixed per unit by the loadings \(k_{rj}\).

Heteroskedastic noise. When a population column is supplied, the per-cell noise variance is \(\sigma^2 / \text{pop}_{jt}\) (internally the estimator carries inv_pop \(= \overline{\text{pop}} / \text{pop}\) so \(\sigma\) remains on the outcome scale), matching the sampling variance of a rate computed from a finite population. With no population column the noise is homoskedastic (\(\text{pop}_{jt} \equiv 1\)).

Counterfactual. The treated unit’s post-period cells are masked – treated as missing and imputed by the model, which contains no treatment-effect parameter – so their posterior is the no-intervention counterfactual \(\widehat{y}_{1t}^N\). The donors contribute all periods, pinning the shared trend and factors; the treated loadings, fit on the pre-period, project that structure forward under the GP smoothness prior. The per-period effect is \(\tau_t \coloneqq y_{1t} - \widehat{y}_{1t}^N\) and the ATT is \(\widehat{\tau} \coloneqq T_2^{-1}\sum_{t\in\mathcal{T}_2}\tau_t\), each with a full posterior credible band.

Assumptions#

  1. Single treated unit, balanced panel. One unit is treated after \(T_0\); every unit is observed at every period.

    Remark. The setup boundary (mlsynth.utils.mtgp_helpers.setup) enforces both through dataprep(), translating a violation to MlsynthDataError.

  2. The no-intervention outcomes are smooth in time. The global trend and the latent factors are draws from squared-exponential Gaussian processes.

    Remark. This is what distinguishes MTGP from Bayesian Factor Synthetic Control (BFSC): smoothness is a prior belief that borrows strength across adjacent periods. It is well suited to trending rate data and ill suited to series with genuine jumps or high-frequency structure, which the kernel would smooth away.

  3. A low-rank coregionalization structure links the units. Donors and the treated unit share \(R\) latent time factors, mixed by unit-specific loadings.

    Remark. As in a factor SC this is weaker than the convex-hull requirement of simplex SC: the treated unit need not lie inside the donors’ hull, only inside their factor span. If the treated loadings fall outside the donors’, the post-period is an extrapolation and – because the factors are smooth GPs – the credible band widens with distance to say so.

  4. Noise variance is known up to scale, inversely proportional to population. \(u_{jt} \sim \mathcal{N}(0, \sigma^2/\text{pop}_{jt})\) with a single \(\sigma\).

    Remark. This matches the sampling variance of a rate from a finite population and lets large units carry more weight. Omit the population column for homoskedastic noise; the estimator then sets \(\text{pop}_{jt}\equiv 1\).

Inference and diagnostics#

MTGP is inferential by construction: res.inference.ci_lower / res.inference.ci_upper give the ATT credible interval, and the counterfactual band is on res.inference_detail (counterfactual_lower / counterfactual_upper). Because the trend and factors are sampled rather than plugged in, the band includes GP uncertainty and widens post-treatment. NUTS diagnostics are surfaced on res.weights.summary_stats – nuts_accept_prob, nuts_divergences, max_rhat – alongside the posterior-mean length-scales (lengthscale_f, lengthscale_global) and noise scale (sigma_post_mean). Read convergence on the counterfactual and \(\sigma\) (the identified quantities), not on the raw loadings: a coregionalization model is rotation/sign non-identified, so individual loadings need not mix while the reported counterfactual does. As a placebo check, refit with each donor as the pseudo-treated unit and compare the treated band to the placebo distribution.

Example#

import pandas as pd
from mlsynth import MTGP   # requires: pip install 'mlsynth[bayes]'

df = pd.read_csv(
    "https://raw.githubusercontent.com/jgreathouse9/mlsynth/main/"
    "basedata/german_reunification.csv")
df["treat"] = df["Reunification"].astype(int)

res = MTGP({
    "df": df, "outcome": "gdp", "treat": "treat",
    "unitid": "country", "time": "year",
    "n_factors": 5, "seed": 0, "display_graphs": False,
}).fit()

print(f"mean post ATT   : {res.att:+.0f}")
print(f"95% credible    : [{res.inference.ci_lower:+.0f}, {res.inference.ci_upper:+.0f}]")
print(f"NUTS max r-hat  : {res.weights.summary_stats['max_rhat']:.3f}")

The counterfactual tracks West Germany through the pre-period and its band widens through the post-period, the further the extrapolation runs from 1990. Supply a population column to scale the noise by unit size.

Verification#

MTGP is cross-validated against the authors’ own reference: the Stan program from the paper’s replication package (Ben-Michael et al. 2023), run on the study’s California panel (the Armed and Prohibited Persons System, homicide rates per 100,000 across the 50 states, 1997–2018, California treated in 2007). Ported to NumPyro, the posterior counterfactual of California matches the Stan posterior cell-for-cell – correlation \(0.99993\), maximum discrepancy \(0.2\%\) of the rate level – and the mean post-2007 ATT agrees to \(0.002\) per 100,000 (\(-1.029\) vs \(-1.031\)), the residual being pure Monte-Carlo error between two independent NUTS runs. See the replication page MTGP – Multitask Gaussian Process Synthetic Control (Ben-Michael et al. 2023) and the durable case benchmarks/cases/mtgp_california.py. The estimator, config validation, the missing-dependency guard, effect recovery, the post-period band-widening, the heteroskedastic-population branch, and the result contract are unit-tested (mlsynth/tests/test_mtgp.py).

Core API#

Multitask Gaussian Process synthetic control (MTGP).

Implements the Gaussian model of:

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.

MTGP models the control potential outcomes with a multitask Gaussian process whose kernel is separable over time and units: a global time-GP trend, a low-rank intrinsic-coregionalization term whose latent factors carry a squared-exponential smoothness prior over time, and unit intercepts. The treated unit’s post-period cells are masked and imputed; the posterior of that imputation is the counterfactual, giving a credible band that widens post-treatment. The posterior is drawn with NUTS (NumPyro, double precision), so MTGP needs the [bayes] optional dependency.

See mlsynth.utils.mtgp_helpers for the algorithmic pieces.

class mlsynth.estimators.mtgp.MTGP(config: MTGPConfig | dict)#

Bases: object

Multitask Gaussian Process synthetic control (Ben-Michael et al. 2023).

Parameters:

config (MTGPConfig or dict) – Configuration object. See mlsynth.config_models.MTGPConfig.

Returns:

MTGPResults – Posterior counterfactual with a credible band that widens post-treatment, the ATT posterior mean + credible interval, and NUTS diagnostics. Requires mlsynth[bayes].

fit() → MTGPResults#

Run the MTGP NUTS sampler and assemble results.

Configuration#

class mlsynth.config_models.MTGPConfig(*, df: ~pandas.DataFrame, outcome: str, treat: str, unitid: str, time: str, display_graphs: bool = False, save: bool | str = False, counterfactual_color: ~typing.List[str] = <factory>, treated_color: str = 'black', plot: ~mlsynth.config_models.PlotConfig = <factory>, n_factors: ~typing.Annotated[int, ~annotated_types.Ge(ge=1)] = 5, population: str | None = None, n_warmup: ~typing.Annotated[int, ~annotated_types.Ge(ge=1)] = 1000, n_samples: ~typing.Annotated[int, ~annotated_types.Ge(ge=1)] = 1000, n_chains: ~typing.Annotated[int, ~annotated_types.Ge(ge=1)] = 4, target_accept: ~typing.Annotated[float, ~annotated_types.Gt(gt=0.0), ~annotated_types.Lt(lt=1.0)] = 0.9, max_tree_depth: ~typing.Annotated[int, ~annotated_types.Ge(ge=1)] = 13, ci_alpha: ~typing.Annotated[float, ~annotated_types.Gt(gt=0.0), ~annotated_types.Lt(lt=1.0)] = 0.05, seed: int = 0, verbose: bool = False)#

Configuration for Multitask Gaussian Process synthetic control (MTGP).

Implements the Gaussian model of:

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).

MTGP models the control potential outcomes with a multitask Gaussian process whose kernel is separable over time and units: a global time-GP trend, a low-rank (rank n_factors) intrinsic-coregionalization term whose latent factors carry a squared-exponential smoothness prior over time, and unit intercepts. The treated unit’s post-intervention cells are masked and imputed, so the posterior of that imputation is the counterfactual, giving the treatment effect a credible band that widens post-treatment. The posterior is drawn with NUTS (NumPyro, in double precision), so this estimator requires the [bayes] optional dependency (pip install mlsynth[bayes]).

Parameters:

  • n_factors (int) – Rank of the unit (intrinsic-coregionalization) kernel – the number of shared latent time-GP factors. Paper default 5.

  • population (str, optional) – Column giving each unit-period’s population (or exposure). When provided, the observation noise scales as sqrt(1 / population) (heteroskedastic by size, as in the paper); otherwise the noise is homoskedastic.

  • n_warmup (int) – NUTS warm-up (adaptation) iterations per chain.

  • n_samples (int) – Retained NUTS samples per chain (total, across chains).

  • n_chains (int) – Number of NUTS chains.

  • target_accept (float) – NUTS target acceptance probability (0.9 by default, as in the paper).

  • max_tree_depth (int) – NUTS maximum tree depth (13 by default, as in the paper).

  • ci_alpha (float) – Two-sided level for the credible interval (0.05 -> 95% band).

  • seed (int) – PRNG seed (makes a fit reproducible).

  • display_graphs (bool) – Plot the observed-vs-counterfactual path with the credible band.

  • verbose (bool) – Show the NUTS progress bar.

ci_alpha: float#
display_graphs: bool#
max_tree_depth: int#
model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid'}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

n_chains: int#
n_factors: int#
n_samples: int#
n_warmup: int#
population: str | None#
seed: int#
target_accept: float#
verbose: bool#

Result Containers#

MTGP.fit() returns a MTGPResults – an EffectResult whose standardized sub-models carry the ATT, the counterfactual and gap paths, and the credible interval, with the posterior draws, GP length-scales, and NUTS diagnostics on res.posterior.