cs.LGMay 8, 2026

Arrow: A Foundation Model for Causal Discovery

Authors: Ryan ThompsonHe ZhaoDaniel M. SteinbergEdwin V. Bonilla

Organizations: University of Technology Sydney · CSIRO’s Data61

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.

Explore similar work

CardsList
  1. FoundCause: Causal Discovery with Latent Confounders from Observational Data

    Jun 16, 2026Patrick Blöbaum, Krishnakumar Balasubramanian, Shiva Prasad KasiviswanathanCausal Foundation ModelsStructural Causal Model