Causal Sensitivity Paper Replica — Sanity v1 (MSM bounds)
A prior-data-fitted network (PFN) that predicts causal sensitivity bounds
(lower / upper) on treatment effects under unobserved confounding, for a given
sensitivity level Γ. Trained in PFN Studio as a replica of "Amortizing Causal
Sensitivity Analysis via Prior-Data-Fitted Networks".
What it does
Given an observational dataset (covariates X, binary treatment A, outcome Y), a
query (x, a), and a confounding level Γ, the model outputs the bound on the
treatment effect — instantly, replacing a per-query optimization.
Architecture
tabular_embedder (d=128) → transformer_encoder (d=128, 4 heads, 10 layers)
→ two scalar heads (upper_bound, lower_bound). ~2M parameters.
Training
- Prior:
causal_sensitivity_optimized — a faithful (scaled) reimplementation
of the authors' SCM + Lagrangian frontier label pipeline (MSM divergence,
warm-started λ-sweep, cumulative monotonicity repair).
- Run: 50 steps, batch 16, lr 1e-3, on a single RTX 3060 (Vast.ai).
- Result: training loss fell from ~80 → ~0.83.
Scope & limitations (honest)
This is a sanity-scale prototype, not the full paper result:
- The prior/label pipeline faithfully reproduces the authors' method (scaled down).
- The model + loss are an approximation: standard transformer + MSE on a single
bound — not the paper's per-feature transformer + GMM heads + GMM-NLL.
- Outputs are point bounds, not calibrated predictive intervals.
- Trained for 50 steps on small synthetic data — a proof the pipeline works
end-to-end, not a benchmarked model.
Files
checkpoint/model.pt — trained weights (the model).checkpoint/topology.json — architecture map.priors/, models/, evals/, runs/ — the full reproducible project.
Reference
Based on Amortizing Causal Sensitivity Analysis via Prior-Data-Fitted Networks
(Javurek, Frauen, et al.).
