Model card: the checkpoint's actual run, RealCause results; remove step snapshots and optimizer/scheduler states

50e9930 verified 2 days ago

6.3 kB

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

# 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

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