MediPFN

MediPFN estimates the population natural indirect effect (NIE) and its uncertainty from observed tabular data. The StatsArena checkpoint takes one 500-row table and combines frozen TabPFN-v3 nuisance predictions with a cross-fitted mediation score, a fine-tuned error head, and independent calibration.

$${\mathrm{NIE}}_1=\mathbb{E}[Y(1,M(1))-Y(1,M(0))]\mathrm{.}\backslash operatorname\{NIE\}\_1\; =\; \backslash mathbb\{E\}[Y(1,M(1))\; -\; Y(1,M(0))].$$

Paper · GitHub · Evaluation results

Checkpoints

File	Contents
Arena method ZIP	Complete MediPFN inference bundle with the files below, code, and a synthetic example
Mediation adapter	Trained adaptation for nuisance prediction
Error head	Fine-tuned model for quantiles of estimation error
Calibration	Independently fitted interval adjustments

The Quick Start also downloads the official classifier and regressor. Prior Labs TabPFN-v3 weights

Quick Start

Use Python 3.10. The first run downloads about 446 MB of official TabPFN-v3 weights from Prior Labs. Both Hugging Face repositories support anonymous downloads.

python -m pip install "huggingface_hub==2.4.0"
hf download shibo769/MediPFN arena_n500/requirements-research.txt --revision 315046cf9838ebcf2b920dad0f87e95bbb5bbc04 --local-dir medipfn_hf
python -m pip install -r medipfn_hf/arena_n500/requirements-research.txt

Save the following as quickstart.py and run python quickstart.py. The example downloads a fixed MediPFN release and the two matching TabPFN-v3 checkpoints, then scores the included synthetic table on CPU.

import json
from pathlib import Path
import sys

from huggingface_hub import hf_hub_download, snapshot_download

model_dir = Path(snapshot_download(
    repo_id="shibo769/MediPFN",
    revision="315046cf9838ebcf2b920dad0f87e95bbb5bbc04",
    allow_patterns=["arena_n500/**"],
    local_dir="medipfn_hf",
    token=False,
)) / "arena_n500"

backbone_repo = "Prior-Labs/tabpfn_3"
backbone_revision = "e923ba9be85784206c9e2f43b0035c84d5fd5747"
classifier, regressor = [
    hf_hub_download(
        repo_id=backbone_repo,
        filename=f"tabpfn-v3-{role}-v3_default.ckpt",
        revision=backbone_revision,
        token=False,
    )
    for role in ("classifier", "regressor")
]

sys.path.insert(0, str(model_dir.resolve()))
from predict import Predictor, read_observed

predictor = Predictor(classifier, regressor, device="cpu", threads=4)
observed = read_observed(model_dir / "example" / "observed_n500.npz")
result = predictor(observed, fold_seed=18031)
print(json.dumps({
    "estimate": result["estimate"],
    "standard_error": result["standard_error"],
    "intervals": result["intervals"],
}, indent=2))

With the example seed, the NIE is approximately 0.24574, the score standard error is 0.06337, and the 95% confidence interval is [0.06425, 0.42723].

For Arena, initialize Predictor once and reuse it across tables. Pass the evaluator-assigned fold_seed on each call. Supply an observed 500-row table with x, a, m, y, x_types, and mediator_types.

See the input specification for the table format, model details for identification assumptions, and license and attribution for usage terms. NOTICE.