FADO: Learning Causal Explanations with Prior-Fitted Networks

Code, trained checkpoints and cached results accompanying the paper FADO: Learning Causal Explanations with Prior-Fitted Networks.

FADO (Feature-wise Amortized Do-Operators) is a prior-fitted network trained on synthetic tabular datasets drawn from structural causal models. Given an observational table it predicts, for every feature $X_i$, the feature intervention effect $\Delta_i = E[Y\mid do(X_i=+1)] - E[Y\mid do(X_i=-1)]$, without a causal graph.

This repository contains:

• Training and evaluation code for FADO (causal_probe/), including the SCM prior families and the model architectures (causal_probe/models/).
• Baselines and evaluation harnesses (DoubleML, CausalPFN, Do-PFN, causal forests, DR-Learner, LiNGAM, NOTEARS, GES, do-Shapley / Causal Shapley, etc.).
• Semi-synthetic and real-data benchmarks (ACIC-2016, IHDP, LaLonde, Criteo Uplift, Hillstrom, STAR, RCT_Bench, Causal Chambers, Sachs).
• The §1 SHAP-vs-causal motivating experiments through TabICL (src/).
• Trained checkpoints and cached evaluation artifacts (causal_probe/results/) sufficient to regenerate every script-generated paper table and figure without retraining.

Install

We use uv for environment management. From a fresh clone:

uv venv
uv pip install -e .                   # core deps + tabicl[shap]
uv pip install -e '.[baselines]'      # optional: baseline libraries

Python ≥ 3.11 is required. The tabicl dependency is pinned to upstream soda-inria/tabicl at the SHA in pyproject.toml.

The model

FADO is the foundation model causal_probe/results/probe_intersample_L_seed0.ckpt (57 MB, weights plus architecture and training arguments):

Architecture	axial transformer with inter-sample attention (--arch intersample), d_model=384, 8 heads, 6 feature-attention + 6 sample-attention blocks, 14.2M parameters
Training width	p ∈ [3, 15] features, padded to 15 columns; n = 512 rows per episode
Training prior	mixture of 11 SCM families (_MIXED_PRIOR_DEFAULT in causal_probe/train.py), per-column noise drawn from seven Gaussian and non-Gaussian families (--noise heavy)
Optimisation	150,000 steps, batch 8, AdamW, lr 1e-4 with 5,000 warm-up steps, seed 0
Wall clock	15.3 h (55,077 s of training) on one NVIDIA H100

Every evaluation feeds the model its training shape: inputs with fewer than 15 columns are padded with independent noise columns (PadToTrainingWidth in causal_probe/eval.py, applied automatically by _load_model; --p-pad 15 in the evaluation scripts). Real datasets with more than 15 covariates are capped at 15 (--p-real-cap).

All results in the paper use this model, except the tables that compare model variants (held-out SCMs per training family, extended held-out DAGs, heterogeneous-p, Sachs checkpoints and the checkpoint list). Those use the single-family and earlier checkpoints below, which are also included.

Checkpoint (causal_probe/results/)	Architecture	Training prior	Role
probe_intersample_L_seed0.ckpt (57 MB)	intersample, d=384, 6+6 layers, 8 heads; p ∈ [3, 15]	11 SCM families, 150k steps, batch 8	FADO: every main result
probe_*_p5_*_intersample.ckpt, probe_gauss_control_p5_intersample.ckpt (7 files, 5.5 MB each)	intersample, d=96, 3+3 layers, 4 heads; p=5	one SCM family each, 15k–25k steps, batch 32	cross-family heatmap rows; type-flip ablation
probe_main_p5_100k_mixture_heavy.ckpt (5 MB)	row/column transformer, d=96, 3+3 layers, 4 heads; p=5	5 SCM families (linear 0.35, nonlinear 0.15, MLP 0.15, hidden 0.15, mixed 0.20), 100k steps	earlier small model; model-variant comparisons only
probe_<family>_p<p>_<steps>k.ckpt, Gaussian controls, hetero-p	row/column transformer, d=96	one SCM family each	checkpoint table and the model-variant comparisons

Each checkpoint stores its full training arguments under the args key. Result JSONs written by the foundation model record the checkpoint under its training filename, probe_foundation_intersample_L_p3_15_intersample_L_seed0.ckpt, which holds the same weights as probe_intersample_L_seed0.ckpt. The ensemble row of the uncertainty-quality table also used foundation-model seeds 1–4 (listed under ensemble_ckpts in causal_probe/results/realdata_2026-09-29/uncertainty_estimator_compare.json); those four checkpoints are not shipped, so that one evaluation cannot be rerun, but the cached JSON is sufficient to rebuild the table. The p=5 checkpoint was trained with an earlier revision of the SCM generators and of the training script (its 5-family prior is no longer the default mixture); retraining it with this code follows the same recipe but is not bit-for-bit identical.

Reproduce paper tables and figures

From a fresh clone, using only the cached artifacts (about one minute on CPU):

mkdir -p paper/tables paper/figures
uv run python scripts/build_tables.py             # main + appendix tables (exits non-zero if any builder fails)
uv run python scripts/make_result_tables.py       # uncertainty, negative controls, IHDP decoys, Chambers
uv run python scripts/paired_delta_baselines.py   # tab:all-baselines-paired-ci
uv run python scripts/build_acic_fp.py            # tab:acic-fp
uv run python scripts/make_selective_prediction_fig.py
PYTHONPATH=. uv run python scripts/cross_family_heatmap.py   # fig: cross-family heatmap (56 cached cells)
PYTHONPATH=. uv run python -m causal_probe.figure_pred_vs_causal  # Fig. 1

Outputs land in paper/tables/ and paper/figures/. Two tables are hand-formatted merges of cached JSONs and have no generating command: tab:all-baselines (the JSONs read by scripts/paired_delta_baselines.py) and tab:wallclock (scripts/wallclock_baselines.py timings). The anchor-audit appendix comes from scripts/run_anchor_audit.py + scripts/plot_anchor_audit.py and the τ̂-vs-τ scatter from causal_probe/plots.py.

Reproduce from scratch (training and evaluation)

reproduce.sh lists every training and evaluation command with the exact settings used, grouped by paper section; commands are commented out by default.

bash reproduce.sh         # documents every step; commands are commented by default

Some benchmarks need external data that is not redistributed here: the ACIC-2016 covariates (data/acic2016/x.csv, from the public aciccomp2016 package), Criteo Uplift (downloaded via scikit-uplift), RCT_Bench, and the Causal Chambers light-tunnel data (downloaded on first use).

Repository layout

causal_probe/        Core training, eval, baselines (no tabicl import)
  ├── train.py       Trainer (--arch rowcol | intersample; --scm-type mixture = 11-family prior)
  ├── eval.py        Held-out SCM evaluator; _load_model + PadToTrainingWidth
  ├── model.py       Row/column transformer (small single-family models)
  ├── models/        Architecture registry (intersample = FADO)
  ├── scm*.py        Synthetic SCM generators by family
  ├── run_baseline.py Cross-regime baseline harness
  ├── semisynth_acic_fp.py, ihdp_decoy_eval.py      Semi-synthetic decoys
  ├── *_rct.py, rct_negative_controls.py, rct_bench_*.py   Real RCTs
  ├── bootstrap_*.py, *uncertainty*.py, selective_prediction_curve.py   Uncertainty
  ├── causal_chambers_eval.py                       Causal Chambers
  ├── shapley_causal.py, eval_shapley_comparison.py do-Shapley / Causal Shapley
  └── results/       Checkpoints and cached eval JSONs / NPZs
scripts/             Paper-artifact glue (table builders, figure makers)
src/                 §1 SHAP-vs-causal motivating experiments via TabICL
data/                IHDP and Tübingen benchmark inputs
pyproject.toml       Pinned upstream tabicl + baseline deps

License

This code is released under the MIT License. See LICENSE.