pith. sign in

arxiv: 2605.07204 · v1 · submitted 2026-05-08 · 💻 cs.LG

Arrow: A Foundation Model for Causal Discovery

Ryan Thompson , He Zhao , Daniel M. Steinberg , Edwin V. Bonilla This is my paper

Pith reviewed 2026-05-11 02:49 UTC · model grok-4.3

classification 💻 cs.LG
keywords causal discoveryfoundation modelzero-shot learningtransformer architecturedirected acyclic graphobservational datasynthetic pretraininggraph factorization
0
0 comments X

The pith

A transformer pretrained on synthetic causal graphs can discover directed acyclic structures zero-shot on new tabular datasets by predicting skeletons and topological orders.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces Arrow as a foundation model for causal discovery that factorizes any directed acyclic graph into an undirected skeleton plus a topological order to enforce acyclicity. A transformer architecture processes observations to output edge probabilities for the skeleton and scores for variable ordering, which together recover a full graph. Training occurs end-to-end on large numbers of synthetic datasets spanning varied graph families, functional forms, noise, and sizes using a composite likelihood objective. If correct, this yields a reusable model that matches or beats prior methods on both synthetic and real data while requiring far less computation at inference time.

Core claim

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.

What carries the argument

Skeleton-order factorization of a DAG, which separates undirected edge presence from direction via a consistent ordering and allows a transformer to predict both components separately.

If this is right

  • Arrow achieves matching or superior accuracy to existing causal discovery algorithms on in-distribution, out-of-distribution synthetic, semi-synthetic, and real tabular data.
  • Inference cost is substantially lower than that of competitive alternatives while maintaining accuracy.
  • The model can be applied directly to new datasets without retraining or fine-tuning.
  • Large-scale supervised pretraining on diverse synthetic causal data produces reusable zero-shot models.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • This approach suggests that causal discovery could shift toward reusable pretrained models rather than per-dataset optimization.
  • It opens the possibility of scaling the same architecture to higher-dimensional or streaming data where traditional search-based methods become intractable.
  • Domain adaptation techniques could be applied on top of the zero-shot model to handle specific real-world data regimes not captured in the original synthetic mix.

Load-bearing premise

The distribution of synthetic training graphs, functional forms, noise models, and dataset shapes is representative enough of real-world observational data that the pretrained model generalizes zero-shot without domain-specific fine-tuning.

What would settle it

Performance of Arrow falling below competitive baselines on a collection of real-world observational datasets whose variable distributions and causal mechanisms lie outside the span of the synthetic training families.

Figures

Figures reproduced from arXiv: 2605.07204 by Daniel M. Steinberg, Edwin V. Bonilla, He Zhao, Ryan Thompson.

Figure 1
Figure 1. Figure 1: Accuracy and runtime on out-of-distribution synthetic causal discovery tasks generated from [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Architecture of Arrow. The model embeds an input dataset via scalar projection and three transformer stages for within-observation variable interaction, within-variable observation aggrega￾tion, and global contextualization. Skeleton and order heads predict undirected edge probabilities and node scores, yielding directed edge probabilities for training and predicted DAGs for inference. For the product-Bern… view at source ↗
Figure 3
Figure 3. Figure 3: Performance as a function of the number of observations and variables on in-distribution [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Performance as a function of distribution shifts on synthetic datasets. Each shift modifies [PITH_FULL_IMAGE:figures/full_fig_p009_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Negative log-likelihood trajectories on the training and validation data. [PITH_FULL_IMAGE:figures/full_fig_p021_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Performance across graph families, functional forms, and noise models on in-distribution [PITH_FULL_IMAGE:figures/full_fig_p022_6.png] view at source ↗
read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript introduces Arrow, a foundation model for zero-shot causal discovery on observational tabular data. It factorizes a directed acyclic graph into an undirected skeleton and a topological order to guarantee acyclicity by construction. A transformer-based architecture contextualizes variables within and across observations to predict skeleton edge probabilities and node order scores. The model is trained in a supervised fashion on synthetic datasets spanning diverse graph families, functional forms, noise models, and dataset shapes, using an end-to-end differentiable directed edge composite likelihood. The central claim is that 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.

Significance. If the zero-shot generalization results hold, the work is significant because it shows that large-scale pretraining on synthetic data can produce a fast, reusable model for causal discovery, addressing the high per-dataset computational cost of traditional methods. The skeleton-order factorization provides a clean architectural guarantee of acyclicity, and the transformer contextualization is a natural fit for tabular data. Explicit credit is due for the end-to-end differentiable training objective and the breadth of the synthetic training ensemble.

major comments (2)
  1. [Abstract and §4] Abstract and §4 (Experiments): The central performance claim asserts that Arrow 'matches or outperforms' baselines across dataset types at lower inference cost, yet the manuscript supplies no quantitative metrics, baseline details, statistical significance tests, ablation studies, or error bars in the abstract and only limited reporting in §4. This is load-bearing for the empirical contribution and prevents assessment of post-hoc selection or robustness.
  2. [§3.2] §3.2 (Training Distribution): The zero-shot generalization to real datasets is the load-bearing assumption. The description of the synthetic ensemble (graph families, functional forms, noise models, dataset shapes) does not include any coverage analysis, ablation, or diagnostic showing that properties such as hidden confounding, non-additive mechanisms, or distribution shifts typical of real observational data are adequately represented. Without this, the transfer claim rests on an untested coverage assumption rather than demonstrated robustness.
minor comments (2)
  1. [§3.1] §3.1: The composite likelihood objective is described at a high level; an explicit equation defining the factorization into skeleton and order terms would improve clarity and reproducibility.
  2. [Figure 1 and §2] Figure 1 and §2: The diagram of the skeleton-order factorization is helpful but would benefit from an accompanying equation reference showing how the predicted probabilities are combined into a final DAG.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive and detailed feedback on our manuscript. We address each major comment point by point below and outline the revisions we will make to strengthen the empirical presentation and justification of generalization.

read point-by-point responses

  1. Referee: [Abstract and §4] Abstract and §4 (Experiments): The central performance claim asserts that Arrow 'matches or outperforms' baselines across dataset types at lower inference cost, yet the manuscript supplies no quantitative metrics, baseline details, statistical significance tests, ablation studies, or error bars in the abstract and only limited reporting in §4. This is load-bearing for the empirical contribution and prevents assessment of post-hoc selection or robustness.

    Authors: We agree that the abstract is high-level and omits specific metrics to remain concise. In the revised manuscript we will update the abstract to include key quantitative highlights such as average performance gains (e.g., in SHD or F1) and inference-time reductions relative to competitive baselines. Section 4 already contains tables and figures comparing Arrow to multiple methods across in-distribution, out-of-distribution, semi-synthetic, and real datasets; however, we acknowledge that additional statistical significance tests, error bars, and ablation studies would improve transparency. We will add these elements, including ablations on the skeleton-order factorization and training diversity, to allow readers to assess robustness and rule out post-hoc selection. revision: yes

  2. Referee: [§3.2] §3.2 (Training Distribution): The zero-shot generalization to real datasets is the load-bearing assumption. The description of the synthetic ensemble (graph families, functional forms, noise models, dataset shapes) does not include any coverage analysis, ablation, or diagnostic showing that properties such as hidden confounding, non-additive mechanisms, or distribution shifts typical of real observational data are adequately represented. Without this, the transfer claim rests on an untested coverage assumption rather than demonstrated robustness.

    Authors: The synthetic training distribution was constructed to span a wide variety of graph families, functional forms, noise models, and dataset shapes precisely to support generalization. While §3.2 describes these choices, we recognize that it lacks explicit coverage diagnostics or targeted ablations for properties such as hidden confounding and non-additive mechanisms. We will revise §3.2 to include a more detailed discussion of how the ensemble addresses these aspects, supported by citations to the empirical results on semi-synthetic and real datasets that serve as the primary evidence of transfer. We will not introduce entirely new large-scale ablations if they fall outside the scope of a revision. revision: partial

Circularity Check

0 steps flagged

No significant circularity; supervised training and external evaluation are independent of inputs

full rationale

The paper trains a transformer model in supervised fashion on synthetic datasets that include ground-truth graphs, using a composite likelihood derived from the skeleton-order factorization. This factorization guarantees acyclicity by construction as an architectural choice, but the reported performance claims (matching or outperforming baselines on held-out synthetic, semi-synthetic, and real datasets) are measured against external methods and data, not derived tautologically from the training inputs or fitted parameters. No load-bearing self-citations, ansatzes smuggled via prior work, or predictions that reduce to the model's own definitions are present in the abstract or described methodology. The zero-shot generalization to real data is an empirical result, not a self-referential derivation.

Axiom & Free-Parameter Ledger

2 free parameters · 2 axioms · 0 invented entities

The central claim rests on the assumption that a broad but finite family of synthetic graphs and mechanisms is sufficient to train a generalizable model; the transformer contains many learned parameters whose values are determined by the training distribution rather than derived from first principles.

free parameters (2)
  • Transformer architecture hyperparameters
    Number of layers, attention heads, embedding dimension, and learning rate schedule are chosen to fit the synthetic training distribution.
  • Composite likelihood weighting
    Relative weighting between skeleton edge probabilities and order scores is tuned during supervised training.
axioms (2)
  • standard math Any DAG can be exactly represented by an undirected skeleton plus a topological order.
    Invoked to guarantee acyclicity by construction; this is a standard graph-theoretic fact.
  • domain assumption The synthetic data generator spans the relevant characteristics of real observational data.
    Required for zero-shot transfer; stated as the training distribution spanning diverse families but not proven.

pith-pipeline@v0.9.0 · 5463 in / 1537 out tokens · 42180 ms · 2026-05-11T02:49:28.657958+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

56 extracted references · 56 canonical work pages

  1. [1]

    Biometrika , volume =

    Ralph Allan Bradley and Milton E Terry , title =. Biometrika , volume =

    work page

  2. [2]

    Mathematical Programming , volume =

    Xiaoye Jiang and Lek-Heng Lim and Yuan Yao and Yinyu Ye , title =. Mathematical Programming , volume =

    work page

  3. [3]

    Léon Bottou and Frank E Curtis and Jorge Nocedal , journal =

    work page

  4. [4]

    Kyle Cranmer and Johann Brehmer and Gilles Louppe , journal =

    work page

  5. [5]

    Noah Hollmann and Samuel Müller and Katharina Eggensperger and Frank Hutter , booktitle =

    work page

  6. [6]

    Jordan Hoffmann and Sebastian Borgeaud and Arthur Mensch and Elena Buchatskaya and Trevor Cai and Eliza Rutherford and Diego de Las Casas and Lisa Anne Hendricks and Johannes Welbl and Aidan Clark and Tom Hennigan and Eric Noland and Katie Millican and George van den Driessche and Bogdan Damoc and Aurelia Guy and Simon Osindero and Karen Simonyan and Eric...

    work page

  7. [7]

    Jared Kaplan and Sam McCandlish and Tom Henighan and Tom B Brown and Benjamin Chess and Rewon Child and Scott Gray and Alec Radford and Jeffrey Wu and Dario Amodei , note =

    work page

  8. [8]

    Samuel Müller and Noah Hollmann and Sebastian Pineda Arango and Josif Grabocka and Frank Hutter , booktitle =

    work page

  9. [9]

    Ryan L Murphy and Balasubramaniam Srinivasan and Vinayak Rao and Bruno Ribeiro , booktitle =

    work page

  10. [10]

    Xun Zheng and Bryon Aragam and Pradeep Ravikumar and Eric P Xing , booktitle =

    work page

  11. [11]

    Kevin Bello and Bryon Aragam and Pradeep Ravikumar , booktitle =

    work page

  12. [12]

    Menghua Wu and Yujia Bao and Regina Barzilay and Tommi S Jaakkola , journal =

    work page

  13. [13]

    Lars Lorch and Scott Sussex and Jonas Rothfuss and Andreas Krause and Bernhard Schölkopf , booktitle =

    work page

  14. [14]

    Bryan Andrews and Joseph Ramsey and Ruben Sanchez-Romero and Jazmin Camchong and Erich Kummerfeld , booktitle =

    work page

  15. [15]

    Ilya Loshchilov and Frank Hutter , booktitle =

    work page

  16. [16]

    Dan Hendrycks and Kevin Gimpel , title =

    work page

  17. [17]

    Ashish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N Gomez and Łukasz Kaiser and Illia Polosukhin , booktitle =

    work page

  18. [18]

    Science , pages =

    Karen Sachs and Omar Perez and Dana Pe'er and Douglas A Lauffenburger and Garry P Nolan , number =. Science , pages =

    work page

  19. [19]

    R Duncan Luce , title =

    work page

  20. [20]

    Journal of the Royal Statistical Society: Series C (Applied Statistics) , pages =

    R L Plackett , number =. Journal of the Royal Statistical Society: Series C (Applied Statistics) , pages =

    work page

  21. [21]

    Peter Spirtes and Clark Glymour and Richard Scheines , address =

    work page

  22. [22]

    Judea Pearl , address =

    work page

  23. [23]

    Developmental Psychology , pages =

    E Michael Foster , number =. Developmental Psychology , pages =

    work page

  24. [24]

    Journal of Economic Literature , pages =

    Guido W Imbens , number =. Journal of Economic Literature , pages =

    work page

  25. [25]

    International Journal of Epidemiology , pages =

    Peter W G Tennant and Eleanor J Murray and Kellyn F Arnold and Laurie Berrie and Matthew P Fox and Sarah C Gadd and Wendy J Harrison and Claire Keeble and Lynsie R Ranker and Johannes Textor and Georgia D Tomova and Mark S Gilthorpe and George T H Ellison , number =. International Journal of Epidemiology , pages =

    work page

  26. [26]

    Anish Dhir and Matthew Ashman and James Requeima and Mark van der Wilk , booktitle =

    work page

  27. [27]

    Achille Nazaret and Justin Hong and Elham Azizi and David Blei , booktitle =

    work page

  28. [28]

    Graham Brightwell and Peter Winkler , journal =

    work page

  29. [29]

    Cristiano Varin and Nancy Reid and David Firth , journal =

    work page

  30. [30]

    Publicationes Mathematicae Debrecen , pages =

    P Erdős and A Rényi , number =. Publicationes Mathematicae Debrecen , pages =

    work page

  31. [31]

    Science , pages =

    Albert-László Barabási and Réka Albert , number =. Science , pages =

    work page

  32. [32]

    Juho Lee and Yoonho Lee and Jungtaek Kim and Adam R Kosiorek and Seungjin Choi and Yee Whye Teh , booktitle =

    work page

  33. [33]

    Proceedings of the 42nd International Conference on Machine Learning , pages =

    Jingang Qu and David Holzmüller and Gaël Varoquaux and Marine. Proceedings of the 42nd International Conference on Machine Learning , pages =

    work page

  34. [34]

    Nature , pages =

    Noah Hollmann and Samuel Müller and Lennart Purucker and Arjun Krishnakumar and Max Körfer and Shi Bin Hoo and Robin Tibor Schirrmeister and Frank Hutter , number =. Nature , pages =

    work page

  35. [35]

    Jimmy Lei Ba and Jamie Ryan Kiros and Geoffrey E Hinton , title =

    work page

  36. [36]

    Journal of Machine Learning Research , volume =

    David Maxwell Chickering , title =. Journal of Machine Learning Research , volume =

    work page

  37. [37]

    Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics , volume =

    Xun Zheng and Chen Dan and Bryon Aragam and Pradeep Ravikumar and Eric Xing , title =. Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics , volume =

    work page

  38. [38]

    Advances in Neural Information Processing Systems , volume =

    Lars Lorch and Jonas Rothfuss and Bernhard Schölkopf and Andreas Krause , title =. Advances in Neural Information Processing Systems , volume =

    work page

  39. [39]

    International Conference on Learning Representations , year =

    Bertrand Charpentier and Simon Kibler and Stephan Günnemann , title =. International Conference on Learning Representations , year =

    work page

  40. [40]

    Edwin V Bonilla and Pantelis Elinas and He Zhao and Maurizio Filippone and Vassili Kitsios and Terry O'Kane , title =

    work page

  41. [41]

    Advances in Neural Information Processing Systems , volume =

    Ryan Thompson and Edwin V Bonilla and Robert Kohn , title =. Advances in Neural Information Processing Systems , volume =

    work page

  42. [42]

    Bo Peng and Sirui Chen and Jiaguo Tian and Yu Qiao and Chaochao Lu , title =

    work page

  43. [43]

    Proceedings of the 41st International Conference on Machine Learning , volume =

    Alan Nawzad Amin and Andrew Gordon Wilson , title =. Proceedings of the 41st International Conference on Machine Learning , volume =

    work page

  44. [44]

    Proceedings of the 36th International Conference on Machine Learning , volume =

    Yue Yu and Jie Chen and Tian Gao and Mo Yu , title =. Proceedings of the 36th International Conference on Machine Learning , volume =

    work page

  45. [45]

    Advances in Neural Information Processing Systems , volume =

    Ignavier Ng and AmirEmad Ghassami and Kun Zhang , title =. Advances in Neural Information Processing Systems , volume =

    work page

  46. [46]

    Journal of Machine Learning Research , volume =

    Markus Kalisch and Peter Bühlmann , title =. Journal of Machine Learning Research , volume =

    work page

  47. [47]

    International Conference on Learning Representations , year =

    Bao Duong and Sunil Gupta and Thin Nguyen , title =. International Conference on Learning Representations , year =

    work page

  48. [48]

    International Conference on Learning Representations , year =

    Nan Rosemary Ke and Silvia Chiappa and Jane X Wang and Jorg Bornschein and Anirudh Goyal and Melanie Rey and Theophane Weber and Matthew Botvinick and Michael Curtis Mozer and Danilo Jimenez Rezende , title =. International Conference on Learning Representations , year =

    work page

  49. [49]

    Journal of Data Science , volume =

    Anne Helby Petersen and Joseph Ramsey and Claus Thorn Ekstrøm and Peter Spirtes , title =. Journal of Data Science , volume =

    work page

  50. [50]

    Naiyu Yin and Tian Gao and Yue Yu , title =

    work page

  51. [51]

    Wenbo Xu and Yue He and Yunhai Wang and Xingxuan Zhang and Kun Kuang and Yueguo Chen and Peng Cui , title =

    work page

  52. [52]

    Omar Swelam and Lennart Purucker and Jake Robertson and Hanne Raum and Joschka Boedecker and Frank Hutter , title =

    work page

  53. [53]

    Weronika Ormaniec and Scott Sussex and Lars Lorch and Bernhard Schölkopf and Andreas Krause , booktitle =

    work page

  54. [54]

    Yuqin Yang and Mohamed Nafea and AmirEmad Ghassami and Negar Kiyavash , booktitle =

    work page

  55. [55]

    Rebecca J Herman and Jonas Wahl and Urmi Ninad and Jakob Runge , booktitle =

    work page

  56. [56]

    Nature , pages =

    Duncan J Watts and Steven H Strogatz , number =. Nature , pages =

    work page