ProximalFM

ProximalFM estimates conditional average treatment effects (CATEs) from observational data with unobserved confounding and observed proxies. Trained on synthetic causal datasets, it predicts effects in context, without fitting its weights to the new dataset.

Given covariates X, treatment proxy Z, outcome proxy W, treatment A, and outcome Y, it estimates

$$Q(x)=\mathbb{E}[Y(1)-Y(0)\mid X=x]\mathrm{.}Q(x)\; =\; \backslash mathbb\{E\}[Y(1)\; -\; Y(0)\; \backslash mid\; X=x].$$

Checkpoints

File or folder	Contents
proximalfm.ckpt	Main ProximalFM checkpoint. Start here for inference.
proximalfm-ft.ckpt	ProximalFM fine-tuned for larger contexts.
main-training/	Intermediate checkpoints from the main training run.
fine-tuning-larger-context/	Intermediate checkpoints from larger-context fine-tuning.

Main training uses contexts up to approximately 2,048 rows. Fine-tuning extends this range to 5,000 rows with a limited compute budget.

Code and examples

For installation, usage examples, and training details, see the GitHub repository.