cdfm-base 0.1.0

CDFM: Towards a General-Purpose Causal Discovery Foundation Model

CDFM: Towards a General-Purpose Causal Discovery Foundation Model

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.

---

Installation

pip install cdfm-base

Requirements: torch>=2.0, numpy>=1.20, safetensors, networkx, huggingface_hub.

---

Usage

1. Causal Discovery

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.

from cdfm import CDFM
from cdfm.utils import evaluate_graph, edge_auroc
import numpy as np

# Load from HuggingFace Hub
model = 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 Control
result_manual = model.predict(data, threshold=0.5)

print(result.adjacency) # (D, D) binary causal graph
metrics = 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.0000

2. Missing value imputation

CDFM has a built-in imputation head trained with quantile loss. Call model.imputation(data) to fill missing values automatically:

from cdfm import CDFM
import numpy as np

model = 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 NaN
imputed = model.imputation(data_with_nan)

# Compare with mean imputation
mean_imp = data_with_nan.copy()
for j in range(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.7817

---

API Reference

CDFM Class

class CDFM:
    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str = "DMIRLAB/CDFM",   # HF Hub or local path
        device: str = "auto",                                  # auto / cpu / cuda:N
        threshold: float | None = None,                        # None = auto-calibrate
    ) -> "CDFM"

    def predict(
        self,
        data: np.ndarray,                     # (N, D) float32
        threshold: float | None = None,       # Probability threshold
        standardize: bool = True,             # Apply z-score standardization
        missing_mask: np.ndarray | None = None, 
    ) -> CDFMResult

CDFMResult Object

@dataclass
class CDFMResult:
    logits: np.ndarray           # (D, D) raw edge scores
    probabilities: np.ndarray    # (D, D) sigmoid(logits)
    adjacency: np.ndarray | None # (D, D) binary graph
    threshold: float | None      # Threshold value used
    runtime_sec: float           # Wall-clock time

---

Links

• Paper (arXiv)
• HuggingFace Model
• GitHub Repository

License

This project is licensed under Apache 2.0.

Citation

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