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