---
license: apache-2.0
tags:
- Causal-Inference
- CausalPFN
library_name: causalpfn
---
# CausalPFN Model

This repository contains the model weights for CausalPFN, a transformer-based in-context learning model for causal effect estimation.

## Model Description

CausalPFN is a pre-trained model for amortized causal effect estimation via in-context learning. It allows for accurate estimation of conditional average treatment effects (CATE) and average treatment effects (ATE) without requiring model retraining for each new dataset.

The model is based on a transformer architecture with uncertainty quantification and calibration.

## Requirements


- Python 3.10+
- PyTorch 2.3+
- NumPy
- scikit-learn
- tqdm
- faiss-cpu
- huggingface_hub

## Installation

To use this model, install the CausalPFN library:

```bash
pip install causalpfn
```

## Usage

You can use this model with the CausalPFN library:

```python
import torch
from causalpfn import CATEEstimator, ATEEstimator


# Create a CATE estimator
causalpfn_cate = CATEEstimator(
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    verbose=True,
)

# Fit the model on your data
# X_train: covariates, T_train: binary treatment, Y_train: observed outcome — from observational data
causalpfn_cate.fit(X_train, T_train, Y_train)  

# Estimate CATE
cate_hat = causalpfn_cate.estimate_cate(X_test)

# Create an ATE estimator
causalpfn_ate = ATEEstimator(
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    verbose=True,
)

# Fit and estimate ATE
causalpfn_ate.fit(X, T, Y)
ate_hat = causalpfn_ate.estimate_ate()
```

## Citations

If you use this model in your research, please cite:

```
@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}, 
}
```

## License

This model is licensed under [Apache-2.0](https://github.com/vdblm/CausalPFN/blob/main/LICENSE).