Arrow: A Foundation Model for Causal Discovery

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.

• The paper introduces a pre-trained, zero-shot model that efficiently maps observational data to acyclic causal graphs.
• It employs a skeleton–order factorization with a composite edge likelihood to ensure identifiability and enforce acyclicity.
• Arrow’s transformer-based architecture yields scalable and robust performance across synthetic, semi-synthetic, and real benchmarks.

Arrow: A Foundation Model for Zero-Shot Causal Discovery

Arrow introduces a pre-trained, zero-shot foundation model for causal discovery on observational tabular data. Leveraging a skeleton–order factorization of DAGs and a transformer-based encoder, Arrow provides a tractable and efficient map from raw data to acyclic causal graphs, showing strong performance and transferability across diverse synthetic, semi-synthetic, and real benchmarks.

Skeleton–Order Factorization and Model Formulation

Arrow formalizes a directed acyclic graph (DAG) as the combination of two structures: an undirected skeleton and a node order (permutation). The skeleton $A$ encodes adjacency, while the topological order $\pi$ is used to orient edges. The DAG is then constructed as $G = A \odot M(\pi)$, where $M(\pi)$ is an order-consistency mask ensuring that all directed edges follow the specified order, guaranteeing acyclicity by construction.

This decomposition yields a tractable conditional likelihood for training, parameterized by a product-Bernoulli distribution for the skeleton and a Plackett–Luce distribution for the node order. Although the exact likelihood is intractable due to the latent ordering, Arrow adopts a composite directed edge likelihood, which the authors formally show to be both identifiability-preserving and structure-aligning within this model class. At inference, Arrow predicts the most probable skeleton and ordering, yielding an acyclic graph in a single forward pass.

Transformer-Based Architecture

Arrow’s architecture contextualizes variables within and across samples using three transformer modules:

1. Observation Transformer: Permutation-invariant, applies self-attention over variables within each observation.
2. Variable Transformer: Aggregates across observations for each variable using learned summary tokens and cross-attention.
3. Context Transformer: Aggregates variable embeddings globally for mutual context.

Prediction heads then transform these contextualized variable embeddings to (i) undirected skeleton edge probabilities (via a permutation-invariant MLP), and (ii) order scores (via a linear head), respecting symmetries and equivariances in the data.

Pretraining Distribution and Optimization

Arrow's pretraining spans a broad and heterogeneous synthetic task distribution:

• Graphs are sampled from Erdős–Rényi and scale-free families with variable sizes ($n$ from $100$ to $2000$, $p$ from $2$ to $100$).
• Causal mechanisms alternate between linear and nonlinear (small MLP) SEMs, with multiple activation functions.
• Noise is introduced via normal, uniform, or beta distributions, in both homogeneous and heterogeneous configurations.
• Task-level $\pi$0 and signal/noise ratios are randomized.

Training proceeds via large-scale stochastic optimization on streams of freshly generated tasks, using the composite edge-wise negative log-likelihood with AdamW.

Empirical Evaluation

Arrow is systematically benchmarked against modern task-specific and pretrained competitors. Key findings include:

• Strong out-of-distribution performance: Arrow consistently matches or outperforms alternatives on synthetic causal discovery tasks with shifted graph structures, functional mechanisms, and noise distributions, while incurring a small, sometimes negligible, degradation under distributional shifts.

  Figure 1: Accuracy and runtime on out-of-distribution synthetic causal discovery tasks; Arrow occupies the best point on the accuracy-runtime tradeoff.

• Scalability across data regimes: On in-distribution synthetic datasets, Arrow maintains minimal normalized structural Hamming distance (nSHD) and superior precision/recall (F1, AP) across varying observation counts and variable dimensions, while maintaining consistently sub-second inference runtime.

Figure 2: Effect of sample and variable count on performance on synthetic datasets (average/SE over 100 draws).

• Robustness to distributional shift: Arrow retains strong performance under interventions on the underlying task distribution, such as swapping to small-world graphs, spline functional relationships, or introducing gamma-distributed noise; shifts in functional mechanism degrade accuracy more than others, but Arrow remains among the strongest methods.

  Figure 3: nSHD and predictive metrics under graph family, functional, and noise model shifts (100 datasets per condition).

• Broad pretraining robustness: Direct ablations on graph families, functional forms, and noise models confirm Arrow's stability, with consistently low error rates in each scenario.

Figure 4: Fine-grained breakdown of performance by generator component (graph, functional, and noise families).

• Empirical convergence: Both training and held-out negative log-likelihoods decrease smoothly over training.

  Figure 5: Negative log-likelihood (composite) on streamed training data and fixed validation set during pretraining.

• Transfer to semi-synthetic and real-world data: On Bayesian Network Repository benchmarks and the Sachs et al. flow cytometry dataset, Arrow provides best or near-best accuracy (nSHD, AP, F1) with fastest inference, showing strong out-of-domain generalization beyond synthetic generator details.

Theoretical and Practical Implications

Arrow demonstrates that large-scale supervised pretraining over a synthetically diverse set of graph discovery tasks can yield a foundation model that is:

• Zero-shot and reusable: Arrow predicts DAGs for novel datasets in a single forward pass, with no task-specific search, optimization, or adaptation.
• Acyclicity-guaranteed: The skeleton–order design ensures all predictions are globally valid DAGs.
• Empirically robust: Broad synthetic pretraining translates into nontrivial transfer to disparate real and semi-synthetic settings.
• Efficient: Inference time is orders of magnitude lower than those of task-specific and aggregation-based alternatives.

From a theoretical perspective, Arrow’s composite likelihood—while a significant simplification compared to the full DAG likelihood—is provably sufficient for parameter identifiability within the skeleton–order model class. This factorization implies that the edgewise statistics capture all model-relevant latent structure, as long as the pairwise directional preferences are globally consistent (i.e., the log-odds field is curl-free).

Outlook and Future Directions

Arrow’s success establishes several new research avenues:

• Latent confounding and interventional data: Extension of the skeleton–order framework to partial observability, latent variables, and experimental data is an open challenge.
• Ultra-large-scale graphs: Model and memory scaling to $\pi$1, possibly via sparse attention or distributed encoding.
• Distribution shift and robustness guarantees: Formal quantification of Arrow's failure modes under hard distribution shifts would further elucidate its domain of applicability.
• Integration with prior knowledge: Adaptation to settings with expert or partial graph constraints.
• Application-oriented deployment: Continued assessment in high-stakes real-world domains (biological, economic, medical) leveraging Arrow as an initial discovery step.

Conclusion

Arrow provides a foundation for scalable, accurate, and efficient causal structure discovery by treating the problem as one of cross-domain supervised learning with a neural inductive bias enforcing acyclicity. Its formulation, pretraining diversity, and transformer architecture collectively yield a zero-shot, acyclicity-guaranteed model that generalizes across a remarkable range of synthetic and real data regimes. Arrow is positioned as a reusable scientific tool, accelerating data-driven hypothesis generation and inviting further advances at the intersection of foundation models and causal inference.