| --- |
| library_name: pfnstudio |
| tags: |
| - pfnstudio |
| - prior-fitted-network |
| - tabular |
| - causal-inference |
| --- |
| |
| # Do-PFN reproduction (PFN Studio) — in-context interventional outcome prediction |
|
|
| An independent reimplementation of [Robertson, Reuter, Guo, Hollmann, Hutter, |
| Schölkopf — *Do-PFN: In-Context Learning for Causal Effect |
| Estimation*](https://arxiv.org/abs/2506.06039) (NeurIPS 2025), built and trained |
| in PFN Studio. A single transformer, pre-trained on synthetic structural causal |
| models, predicts the interventional outcome `E[Y | do(T=t), X]` from an |
| observational context `(X, T, Y)` in one forward pass; the conditional average |
| treatment effect (CATE) is the difference of the two arms. |
|
|
| This is **not** the authors' released checkpoint. No upstream code is copied |
| (the upstream repository at [github.com/jr2021/Do-PFN](https://github.com/jr2021/Do-PFN) |
| declares no license); the prior, model and training objective are written from |
| the paper. |
|
|
| ## The checkpoint |
|
|
| `checkpoint/model.pt` — the end of run `runs/complete-test-2.yaml`, model |
| `models/dopfn-axial-block-smoketest-2.yaml`, prior `priors/do_pfn_scm/`. |
| `checkpoint/topology.json` lists the blocks. |
|
|
| | | | |
| |---|---| |
| | Architecture | Do-PFN input encoder (X in feature groups of 85, value + missing-indicator channels → Linear(170, 192); Y encoded separately → Linear(2, 192)) → 12 post-norm axial attention layers (d_model 192, 6 heads, feed-forward 768, bias-free attention, parameter-free LayerNorm; query rows attend to context rows only) → Y-token pool → bar-distribution head (Linear(192, 768) → GELU → Linear(768, 100), 100 equal-mass buckets, half-normal tails) | |
| | Objective | full-support bar-distribution negative log-likelihood of the interventional outcome at the query positions | |
| | Prior | random-DAG SCM, at most 10 nodes: 1–6 observed covariates, 1 latent node, a binarized treatment, outcome a descendant of the treatment; 2,200 points per task, 75 % context | |
| | Training | 262,144 steps, batch 4, Adam, peak lr 8.5e-5, 32,768 warm-up steps then cosine decay, gradient clip 1.0, weight decay 0, seed 42; a fresh task per sample (no fixed pool) | |
| |
| ## Results on RealCause |
| |
| [RealCause](https://arxiv.org/abs/2011.15007) (Neal, Huang & Raghupathi): semi-synthetic |
| tables sampled from generative models fitted to real data, with both potential |
| outcomes known. 100 tables per dataset from the authors' release (commit `7411cc2`). |
| Per table, rows are permuted with `numpy.random.default_rng(42 + table index)`; |
| the first 90 % are the context and the last 10 % the queries for the CATE; for the |
| average treatment effect (ATE) every row is both context and query. Means over the |
| 100 tables: |
|
|
| | dataset | rows / covariates | PEHE | CATE MSE ÷ MSE of predicting no effect | ATE predicted | ATE true | |
| |---|---|---|---|---|---| |
| | LaLonde CPS | 16,177 / 8 | 9,352 | 0.53 | −4,789 | −7,028 | |
| | LaLonde PSID | 2,675 / 8 | 17,569 | 0.60 | −5,531 | −13,347 | |
| | Twins | 11,383 / 75 | 0.507 | 1.01 | −0.013 | −0.069 | |
|
|
| PEHE = root mean squared error of the per-row CATE on the query rows; a ratio |
| below 1 means the model's CATE beats predicting no effect. |
|
|
| What this shows, in plain terms: on both LaLonde datasets the model recovers the |
| direction of the effect and roughly halves the per-row error of a no-effect guess, |
| but it under-estimates the size of the average effect (by about a third on CPS and |
| more than half on PSID). On Twins it is no better than predicting no effect. |
|
|
| Every RealCause table is outside the range this checkpoint was trained on: the |
| prior's tasks have at most 1,650 context rows and 6 covariates, while these tables |
| have 2,407–14,559 context rows and 8 or 75 covariates. Whether that range is the |
| cause of the shrinkage is not established by these numbers. |
|
|
| ## Load and run |
|
|
| ```python |
| # pip install "pfnstudio-core[torch]==0.9.10" networkx huggingface_hub |
| from pathlib import Path |
| from huggingface_hub import snapshot_download |
| from pfnstudio_core.training.predict import ModelLoader |
| |
| root = Path(snapshot_download("ProfitOpsAI/pfnstudio-dopfn")) |
| loader = ModelLoader(manifest_path=root / "runs/complete-test-2.yaml", |
| project_root=root, checkpoint_dir=root / "checkpoint") |
| model = loader.model # blocks in order; see evals/cid_recovery.py for the exact inference call |
| ``` |
|
|
| Input grid per task: columns `[T, X_1..X_d, Y]`, context rows first with the |
| observed `Y`, query rows after with `Y = NaN` and `T` set to the arm to predict. |
| Standardize `Y` on the context (mean, sample std + 1e-6) before the call and map |
| predictions back; pass `single_eval_pos = number of context rows` to the blocks |
| that ask for it. `evals/cid_recovery.py` is the reference implementation. |
|
|
| ## What's in here |
|
|
| ``` |
| README.md |
| blocks/ do-pfn-input-encoder.py, do-pfn-axial-attention-block.py, bar_distribution_head.py |
| models/ dopfn-axial-block-smoketest-2.yaml (the checkpoint's model), do_pfn.yaml |
| priors/ do_pfn_scm/ prior.py, prior.yaml, requirements.txt (networkx) |
| runs/ complete-test-2.yaml (the checkpoint's run) and earlier runs |
| evals/ cid_recovery.py / .yaml (CID + CATE recovery, PICP calibration) |
| checkpoint/ model.pt, topology.json |
| ``` |
|
|
| ## Not shown here |
|
|
| - **Paper-scale training and wider tables.** The paper pre-trains on far more |
| tasks; this run uses small graphs (≤ 10 nodes, ≤ 6 covariates). |
| - **The paper's structured case studies** (observed confounder, mediator, |
| back-door, front-door, ...) and its Amazon and Law School benchmarks. |
| - **A comparison with the paper's published RealCause numbers.** |
| - **Total-effect CATE under mediators.** The prior's oracle pins observed |
| covariates before intervening, so where a covariate is a descendant of the |
| treatment the evaluation target is a controlled direct effect. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{robertson2025dopfn, |
| title = {Do-PFN: In-Context Learning for Causal Effect Estimation}, |
| author = {Robertson, Jake and Reuter, Arik and Guo, Siyuan and Hollmann, Noah and Hutter, Frank and Sch\"olkopf, Bernhard}, |
| year = {2025}, |
| eprint = {2506.06039}, |
| archivePrefix = {arXiv}, |
| url = {https://arxiv.org/abs/2506.06039} |
| } |
| ``` |
|
|
| ## License |
|
|
| The study files (YAML specs, prior, blocks, evaluation, this card) are Apache-2.0, |
| like the rest of PFN Studio. |
|
|
