arunkatika/csa-pfn-msm

CSA-PFN (MSM) — Amortized Causal Sensitivity Analysis

A prior-data fitted network (PFN) that performs in-context causal sensitivity analysis under the marginal sensitivity model (MSM). Replication of the primary model from:

Javurek, Frauen, Brockschmidt, Schweisthal, Feuerriegel (2026). Amortizing Causal Sensitivity Analysis via Prior-Data Fitted Networks. arXiv:2605.10590.

Built in PFN Studio, with the corresponding author's permission. Trained at the paper's budget (~1B pair-visits per bound) and validated against exact closed-form ground truth: held-out R² 0.893, bounds monotone in Γ 100% of the time.

What the model does

One forward pass, no retraining:

INPUT  = context dataset rows (x1..x10, treatment a, outcome y)
         + query rows (x, a, Γ, bound_flag)
OUTPUT = θ̂ per query row — the upper (bound=0) or lower (bound=1) bound on
         E[Y(a)|x] under the MSM at confounding level Γ ∈ [1, 5]

Sweeping Γ at both bounds traces how causal conclusions degrade with unobserved confounding; CATE bounds follow by combining arms (upper CATE = θ₊(x,1,Γ) − θ₋(x,0,Γ)). The dataset is the context — no per-dataset fitting. The prior's packing/normalization code is part of the model's API: real-data usage must standardize X per feature, normalize y, and pack the 15-column tokens exactly as priors/…/prior.py does.

Held-out validation (exact closed-form ground truth)

100 unseen data-generating processes, 64,000 predictions, seeds disjoint from training (validate_msm_replication.py; artifacts in val_100k/):

Metric	Value
RMSE (normalized-y units)	0.307 (predict-zero baseline: 0.937)
R²	0.893
MAE	0.223
upper ≥ lower (matched pairs)	99.8%
monotone in Γ	100.0%
throughput	~4,600 bounds/s (batched forward pass, CPU laptop)

Per-Γ R² is flat (0.891–0.896 across [1, 5]) — no degradation at high confounding. Γ-monotonicity is learned, not enforced. Throughput is the paper's amortization claim in practice: per-instance methods compute one bound in seconds; this model computes thousands per second.

Usage

import numpy as np, torch, torch.nn as nn

class FlatPFN(nn.Module):
    def __init__(self, d_in=15, dm=128):
        super().__init__()
        self.embed = nn.Linear(d_in, dm)
        layer = nn.TransformerEncoderLayer(dm, 4, 512, 0.0, "relu",
                                           batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(layer, 10)
        self.head = nn.Linear(dm, 1)
    def forward(self, x):
        return self.head(self.encoder(self.embed(x))).squeeze(-1)

sd = torch.load("checkpoint/model.pt", map_location="cpu", weights_only=True)
mapped = {}
for k, v in sd.items():
    if k.startswith("tabular_embedder_0."):
        mapped["embed." + k.split(".")[-1]] = v
    elif k.startswith("head."):
        mapped[k] = v
    elif k.startswith("transformer_encoder_1."):
        mapped["encoder." + k.split(".", 1)[1]] = v
model = FlatPFN(); model.load_state_dict(mapped); model.eval()

def msm_bounds(X, a_col, y_col, x_query, a_query, gammas):
    """X: (N,10) standardized covariates; y_col: normalized outcome.
    Returns (lower, upper) bound arrays, one entry per Γ in `gammas`."""
    N = len(X)
    ctx = np.zeros((N, 15), np.float32)
    ctx[:, :10], ctx[:, 10], ctx[:, 11], ctx[:, 14] = X, a_col, y_col, 1.0
    rows = []
    for b in (0, 1):                       # 0=upper, 1=lower
        for g in gammas:
            r = np.zeros(15, np.float32)
            r[:10], r[10], r[12], r[13] = x_query, a_query, g, b
            rows.append(r)
    seq = torch.tensor(np.vstack([ctx, np.stack(rows)])).unsqueeze(0)
    with torch.no_grad():
        out = model(seq).squeeze(0)[N:].numpy()
    up, lo = out[:len(gammas)], out[len(gammas):]
    return np.minimum.accumulate(lo), np.maximum.accumulate(up)  # monotonized

Training

• Labels: exact closed-form sharp MSM bounds — the paper's "Analytical MSM" pipeline (Table 2), implemented from spec and verified against a brute-force LP solver on 800 cases (exact match; Γ=1 recovers E[Y]). Γ ~ log-uniform [1, 5] redrawn per query group; single MC bank of 128 latent samples reused (paper-exact).
• Prior/DGP: the authors' Setting_standard generative process, vendored verbatim — DAG-structured covariates, MLP propensity and outcome networks, latent confounder U ~ N(0,1), X standardized per feature, Y normalized over both potential outcomes.
• Scale: 100,000 steps × 32 fresh DGPs × 640 labels ≈ 1.02B pair-visits per bound — the paper's budget (10,000 DGPs × 150 epochs) with 320× the DGP diversity (every DGP freshly sampled). Final train loss 0.109. ~9 h on a single RTX 4090.
• Determinism: all training items come from seeded generation; identical runs on different machines reproduce final loss to 4 decimals.

Data pipeline design

Two implementation choices make paper-scale training feasible on one GPU:

1. Vectorized DGP: the covariate DAG is evaluated for all N samples at once (~60× faster than per-sample evaluation; verified statistically identical — only internal RNG call order differs).
2. Parallel prefetch sampling: the trainer requests items with consecutive seeds, so a pool of fork workers precomputes upcoming seeds on all CPU cores while the GPU trains. Outputs are bitwise-identical to serial generation; automatic fallback to serial on any pool failure. Tunable via PRIOR_WORKERS / PRIOR_PREFETCH.

Diagnostics print to stderr ([prior] N=… D=… shows the exact parameters of every sampled item); stdout stays clean JSON for serving.

Architecture

Linear(15→128) → TransformerEncoder (10 pre-norm layers, 4 heads, FF 512, ReLU) → Linear(128→1); ~2M parameters. Matches the paper's Table 1 dimensions (d=128, 4 heads, FF 512, 10 layers).

Documented deviations from the paper: flat row-token attention instead of per-feature attention; single scalar head with a bound flag instead of two 5-component GMM heads with NLL (point estimates rather than posterior predictive distributions); no monotonicity penalty (post-hoc cumulative max/min at inference; empirically 100% monotone regardless); label emission uses fresh random Γ per query group at fixed sequence length instead of a 50-point per-DGP grid (equivalent coverage).

Repository structure

checkpoint/                                 trained weights (model.pt) + topology.json
priors/causal_sensitivity_msm_analytical/   prior: DGP + closed-form labels + packing (the model's API)
models/causal-sensitivity-flat-pfn.yaml     model definition
evals/causal-sensitivity-msm-bounds.yaml    training-time eval
runs/train.yaml                             the training run configuration
validate_msm_replication.py                 held-out validation vs closed-form truth
val_100k/                                   validation report, metrics, plots

Prior — paper Table 2 fidelity ("Analytical MSM" column)

Item	Paper	This prior
Sensitivity model	MSM (closed form)	closed form (verified vs brute-force LP)
Grid range / spacing	[1.0, 5.0], log-uniform	same
Per-DGP grid randomization	yes	redrawn per query group
MC samples	128 (single bank, reused)	same
Grid size	50	5 fresh random Γ per group equivalent coverage
DGP	Setting_standard	vendored verbatim

Intended use & limitations

Research use. Trained for: 10 continuous covariates, binary treatment, continuous outcome, context sizes near 1024, Γ ∈ [1, 5], normalized inputs. Point estimates of bounds (no uncertainty quantification, unlike the paper's GMM heads). Synthetic-prior model: real-data performance depends on resemblance to the prior. Not a substitute for domain-expert sensitivity analysis in high-stakes decisions.

Future work

A companion model for the KL f-sensitivity setting (the paper's optimization-based label pipeline), and mixture-density output heads for full posterior predictive distributions.

Reference & acknowledgements

Javurek, Frauen, Brockschmidt, Schweisthal, Feuerriegel (2026). Amortizing Causal Sensitivity Analysis via Prior-Data Fitted Networks. arXiv:2605.10590. Replicated with the corresponding author's permission. Built with PFN Studio.