maxdemarzi/duckdo-causalpfn

CausalPFN, exported to ONNX for DuckDo

This repository redistributes CausalPFN (Balazadeh et al., 2025) in the form the DuckDo DuckDB extension loads, so that do_ate(..., model := 'causalpfn') and do_cate(..., model := 'causalpfn') can run it inside DuckDB without Python or PyTorch.

This is not the authors' release. The original weights are at vdblm/causalpfn. Please cite their paper, below.

What changed

Every file here was changed from the upstream release (see NOTICE). DuckDo's scripts/export/export_causalpfn.py traced causalpfn_v0.pt into two ONNX graphs, with weights stored as external data. The first is an encoder that turns the context rows into a key/value cache. The second is a decoder that scores query rows against that cache. Every layer draws its keys and values from the context alone, so the cache is computed once and reused for every chunk of queries. No retraining or fine-tuning was done.

Parity with the PyTorch model

Checked by the exporter on a 512-row synthetic data-generating process with a true ATE of about 3:

check	value
ATE, ONNX vs PyTorch	differs by 0.00087
CATE correlation, ONNX vs PyTorch	0.999399
largest single CATE difference	0.282
largest raw output difference	0.0656

The raw-output gap comes from PyTorch's fused attention kernel versus ONNX Runtime's decomposed one. The exporter gates on the estimand instead: it refuses any export whose ATE differs by more than 0.01 or whose CATE correlation falls below 0.999.

Files

file	bytes	SHA-256
causalpfn_encode.onnx	587,534	63872a723e825a461c5256cdfa5d35215697ed00fa1eab12b583be93157ffb83
causalpfn_encode.weights.bin	68,640,768	96c815bc6c73bfde6b862a79684d2fecb5677ac53dcbd68edb060ff500dd5b72
causalpfn_decode.onnx	326,083	0c8e23951988b0a21366c7e5d557d682b4040c68c91d0972f1cb516f96b915a4
causalpfn_decode.weights.bin	51,764,224	efb48fd595e942a6cffe9620a32f3178a54e1b6cc1eaa5cb02707406dd6af7ec
causalpfn.manifest.json	2,364	6756dc4890bf5a797acee4c66845f1ed1ff0c8ad45db0836c9312c7d322bf366

Use from DuckDB

LOAD duckdo;
SELECT * FROM do_download('causalpfn',
    source := 'https://huggingface.co/maxdemarzi/duckdo-causalpfn/resolve/main');
SELECT estimate, ci_low, ci_high
FROM do_ate('customers', treatment := 'discount', outcome := 'revenue', model := 'causalpfn');

Versions of DuckDo that name this repository as their default source pin a specific commit and verify every file's SHA-256 before using it.

Licence

CausalPFN is distributed under the CausalPFN License, Version 1.0; a copy is in LICENSE. Its terms are Apache-2.0's. DuckDo's export code is MIT-licensed and separate from these files.

Citation

@misc{balazadeh2025causalpfn,
  title={CausalPFN: Amortized Causal Effect Estimation via In-Context Learning},
  author={Vahid Balazadeh and Hamidreza Kamkari and Valentin Thomas and Benson Li and Junwei Ma and Jesse C. Cresswell and Rahul G. Krishnan},
  year={2025},
  eprint={2506.07918},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2506.07918}
}