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