English

Arrow: A Foundation Model for Causal Discovery

Machine Learning 2026-05-11 v1

Abstract

We introduce Arrow, a foundation model for zero-shot causal discovery on observational tabular data. Arrow factorizes a directed acyclic graph into an undirected skeleton and a topological order, guaranteeing acyclicity by construction. Given a new dataset, it uses a transformer-based architecture to contextualize variables within and across observations, then predicts skeleton edge probabilities and node order scores that together define a graph. Arrow is trained in a supervised fashion on synthetic datasets with ground-truth graphs, using an end-to-end differentiable directed edge composite likelihood induced by the skeleton-order factorization. The training distribution spans diverse graph families, functional forms, noise models, and dataset shapes. Across in- and out-of-distribution synthetic, semi-synthetic, and real datasets, Arrow matches or outperforms existing causal discovery methods at substantially lower inference cost than competitive alternatives. Our results demonstrate that large-scale pretraining on diverse synthetic data can yield zero-shot causal discovery models that are fast, accurate, and reusable on new datasets.

Keywords

Cite

@article{arxiv.2605.07204,
  title  = {Arrow: A Foundation Model for Causal Discovery},
  author = {Ryan Thompson and He Zhao and Daniel M. Steinberg and Edwin V. Bonilla},
  journal= {arXiv preprint arXiv:2605.07204},
  year   = {2026}
}
R2 v1 2026-07-01T12:56:50.488Z