Causal Discovery Foundation Model (CDFM) is a pretrained foundation model for zero-shot causal discovery. Given purely observational data X (N, D), it predicts the causal graph G (D, D) in a single forward pass.
CDFM reframes causal discovery as a unified, general-purpose framework for zero-shot structural inference. By pretraining on a massive, highly diverse space of synthetic structural causal models, CDFM successfully internalizes complex statistical asymmetries.
- State-of-the-art accuracy. Outperforms all baselines across 15 mechanism families and on real-world benchmarks.
- Zero-shot. One pretrained checkpoint works for any N, any D.
- Easy to use. pip-installable, single
model.predict(data)call.
pip install cdfm-baseRequirements: torch>=2.0, numpy>=1.20, safetensors, networkx, huggingface_hub.
The simplest way to use CDFM is to load the model and pass your observational data directly. By default, CDFM automatically calibrates the threshold for edge prediction.
fromcdfmimportCDFMfromcdfm.utilsimportevaluate_graph, edge_aurocimportnumpyasnp# Load from HuggingFace Hubmodel=CDFM.from_pretrained("DMIRLAB/CDFM")
# Load a simple 4-variable nonlinear example (RFF mechanisms)data=np.loadtxt("tests/data/simple/data.csv", delimiter=",")
gt=np.loadtxt("tests/data/simple/adjacency.csv", delimiter=",").astype(np.int32)
# 1. Standard Prediction (Auto-calibrated threshold)result=model.predict(data) # 2. Manual Threshold Controlresult_manual=model.predict(data, threshold=0.5)
print(result.adjacency) # (D, D) binary causal graphmetrics=evaluate_graph(result.adjacency, gt)
auc=edge_auroc(result.logits, gt)
print(f"F1={metrics['f1']:.4f} SHD={metrics['shd']} AUC={auc:.4f}")
# → F1=1.0000 SHD=0 AUC=1.0000CDFM has a built-in imputation head trained with quantile loss. Call model.imputation(data) to fill missing values automatically:
fromcdfmimportCDFMimportnumpyasnpmodel=CDFM.from_pretrained("DMIRLAB/CDFM")
# Load data and create missing values (seed for reproducibility)rng=np.random.default_rng(42)
data=np.loadtxt("tests/data/simple/data.csv", delimiter=",")
data_with_nan=data.copy()
data_with_nan[rng.random(data.shape) <0.2] =np.nan# CDFM imputation — auto-detects NaNimputed=model.imputation(data_with_nan)
# Compare with mean imputationmean_imp=data_with_nan.copy()
forjinrange(data.shape[1]):
col=data_with_nan[~np.isnan(data_with_nan[:, j]), j]
mean_imp[np.isnan(mean_imp[:, j]), j] =col.mean()
missing=np.isnan(data_with_nan)
mae_cdfm=np.abs(imputed[missing] -data[missing]).mean()
mae_mean=np.abs(mean_imp[missing] -data[missing]).mean()
print(f"CDFM MAE: {mae_cdfm:.4f} | Mean MAE: {mae_mean:.4f}")
# → CDFM MAE: 0.3719 | Mean MAE: 0.7817classCDFM:
@classmethoddeffrom_pretrained(
cls,
pretrained_model_name_or_path: str="DMIRLAB/CDFM", # HF Hub or local pathdevice: str="auto", # auto / cpu / cuda:Nthreshold: float|None=None, # None = auto-calibrate
) ->"CDFM"defpredict(
self,
data: np.ndarray, # (N, D) float32threshold: float|None=None, # Probability thresholdstandardize: bool=True, # Apply z-score standardizationmissing_mask: np.ndarray|None=None, ) ->CDFMResult@dataclassclassCDFMResult:
logits: np.ndarray# (D, D) raw edge scoresprobabilities: np.ndarray# (D, D) sigmoid(logits)adjacency: np.ndarray|None# (D, D) binary graphthreshold: float|None# Threshold value usedruntime_sec: float# Wall-clock timeThis project is licensed under Apache 2.0.
If you use CDFM in your research, please cite:
@article{qiao2026cdfm,
title = {{CDFM}: Towards a General-Purpose Causal Discovery Foundation Model},
author = {Jie Qiao and Ruichu Cai and Zijian Li and Weilin Chen and Pengfei Hua and Boyan Xu and Zhengming Chen and Zhifeng Hao and Peng Cui},
journal = {arXiv preprint arXiv:2607.11508},
year = {2026},
}