CCPFN: Causal Foundation Models with Continuous Treatments

This repository contains the weights for CCPFN (Continuous Causal Prior-Fitted Network), the first causal foundation model for continuous treatment settings, as presented in the paper Causal Foundation Models with Continuous Treatments.

By leveraging in-context learning, CCPFN estimates the conditional expected potential outcome (CEPO), defined as $𝔼[Y(t) \mid X = x]$, predicting causal effects across a wide variety of unseen tasks without any additional training or fine-tuning.

• Repository (Inference): layer6ai-labs/CCPFN-inference
• Paper: Causal Foundation Models with Continuous Treatments

Installation

You can install the inference package via pip:

pip install ccpfn

Quick Start

Here is a simple example demonstrating how to run CCPFN for CEPO estimation:

import numpy as np
import torch
from ccpfn import CEPOEstimator

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Define true individual treatment-response function
def treatment_response(x, t): 
      return np.cos(x[..., 0]) + 2 * x[..., 1] * t

# Define treatment assignment function
def treatment(x):
      return 1 + np.sin(x[..., 2])

# Create synthetic data - covariates, treatment, outcome
rng = np.random.default_rng(seed=42)
n_samples, n_features = 2048, 3
X = rng.standard_normal((n_samples, n_features))
T = treatment(X)
Y = treatment_response(X, T) + 0.1 * rng.standard_normal((n_samples,))

# Context/query (train/test) split
test_ratio = 0.3
ctx_idx = rng.choice(n_samples, int((1 - test_ratio) * n_samples), replace=False)
qry_idx = np.setdiff1d(np.arange(n_samples), ctx_idx)
X_ctx, X_qry = X[ctx_idx], X[qry_idx]
T_ctx, Y_ctx = T[ctx_idx], Y[ctx_idx]
T_qry = rng.random((X_qry.shape[0],))  # Counterfactual treatments

# CEPO Estimation
estimator = CEPOEstimator(device=device)
estimator.fit(X_ctx, T_ctx, Y_ctx)
cepo_pred = estimator.estimate_cepo(X_qry, T_qry)

# Evaluation and results
cepo_true = treatment_response(X_qry, T_qry)
rmse = np.sqrt(np.mean((cepo_true - cepo_pred) ** 2))
print("Results:")
print(f"RMSE: {rmse:.4f}")

Citation

@misc{stith2026causalfoundationmodelscontinuous,
      title={Causal Foundation Models with Continuous Treatments}, 
      author={Christopher Stith and Medha Barath and Vahid Balazadeh and Jesse C. Cresswell and Rahul G. Krishnan},
      year={2026},
      eprint={2605.15133},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.15133}, 
}