REVIEW 4 major objections 5 minor 14 references
Learning Causal Graphs at Scale: A Foundation Model Approach
T0 review · 4 major / 5 minor · reviewed 2026-08-15 · deepseek-v4-flash
Pith's one-line read The paper claims that ADAG, an attention network pre-trained on 70,000 synthetic domains, learns a nonlinear map from observed data to weighted causal graphs and infers new DAGs zero-shot with near-zero structural error.
desk verdict Genuinely new architecture and fast inference, but the d=5 zero-shot numbers look like structure memorization; worth a serious referee, with major revisions. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the attention-based kernel map. Data from one domain are tokenized as $H^{(0)} = X \in \mathbb{R}^{d\times n}$, one token per variable; $L$ linear-attention layers compute $$$H^{{(\ell)}}$ = \operatorname{Attn}[$H^{{(\ell-1)}}$]$H^{{(\ell-1)}}$ + $H^{{(\ell-1)}}$,\quad \operatorname{Attn}[H] = \$\sigma$\bigl(\tfrac{1}{\sqrt{d}} H W_Q W_K^\top H^\top\bigr)\in\mathbb{R}^{d\times d},$$ and the final layer outputs $A = W_{\mathrm{out}}^{P,x}\,\sigma\bigl(\tfrac{1}{\sqrt{d}} H^{(L)} W_Q^{\mathrm{out}}(W_K^{\mathrm{out}})^\top (H^{(L)})^\top\bigr)$. The attention matrix is exactly the weighted adjacency matrix: entry $(i,j)$ is the strength of the proposed edge $X_i\to X_j$. Augmented-Lagrangian training with the trace-exponential acyclicity constraint is what makes the unsupervised output a DAG rather than an arbitrary matrix.
What would settle it
Hold the training protocol fixed (ER1, $d=5$, equal-variance Gaussian noise, edge weights in $[-2,-0.5]\cup[0.5,2]$) but reverse the topological order used to generate test-domain DAGs, or switch test graphs to ER2 with two edges per node, then measure zero-shot SHD of the pre-trained ADAG. If SHD jumps from the reported $\sim0.03$ toward the $\sim1$–$3$ range of per-task baselines, the generalization claim is confined to the training distribution rather than being a general causal prior.
Extended reading notes
Core claim
ADAG claims that a nonlinear attention-based kernel map $A[X(1:n); \Theta]$ can take the $n$ observed samples of $d$ variables and output the $d \times d$ weighted adjacency matrix of the underlying linear SEM, $X = A^\top X + E$. The map is trained by minimizing the reconstruction loss $\sum_\eta \|X^\eta - (A^\eta)^\top X^\eta\|_F^2$ across $M$ domains subject to the acyclicity constraint $h(A^\eta)=\operatorname{tr}(e^{A^\eta\circ A^\eta})-d=0$, using augmented Lagrangian iterations. After pre-training on $M=70{,}000$ ER1 domains for $d\in\{5,10,20\}$ with equal-variance Gaussian noise, the learned kernel performs zero-shot inference on 1,000 held-out domains: SHD 0.027 (heterogeneous, $d=5$) and 0.038 (order-consistent, $d=5$), with relative adjacency error around 0.12 and inference runtime 0.0003 seconds, outperforming both single-task optimizers and multi-task learners in accuracy and speed. The paper further claims that the kernel captures the common low-dimensional structure of the training distribution and that this prior keeps performance relatively stable when each domain has only $n=25$ or $n=50$ observations.
Load-bearing premise
The load-bearing premise is that any new task comes from the same generative recipe used for the 70,000 training domains: the same number of variables, the same sparse graph density, the same edge-weight range, equal-variance Gaussian noise, and either the same DAG (heterogeneous) or the same topological order (order-consistent). Outside that recipe the trained kernel has no stated adaptation mechanism.
Editorial extensions
If this is right
- A model pre-trained under ADAG's protocol can replace per-task DAG optimization: a new domain in the same structural family is solved by one forward pass, cutting inference time from seconds (single-task baselines) or minutes (multi-task baselines) to fractions of a millisecond.
- Small-sample performance shifts: at $n=25$ or $n=50$, ADAG keeps SHD roughly 0.05–0.27 for $d=5$, while single-task baselines degrade to SHD 0.54–1.56, so the shared prior behaves like regularization.
- A model trained on $n=100$ transfers to test domains with only 25 or 50 observations when samples are duplicated to the training length, supporting the claim that high-resource pre-training can serve low-resource downstream tasks.
- The learned kernel's covariance matrix and dominant eigenvectors match the ground-truth adjacency structure, making the shared prior visible rather than a black-box effect.
- The same architecture extends in principle to nonlinear SEMs by feeding the last hidden representation $H^{(L)}$ into the reconstruction loss instead of raw $X$, a direction the paper states as a direct extension.
Reading between the lines
- Editorial: The benchmark only tests same-protocol generalization. A decisive next experiment is to train on ER1 and zero-shot test on ER2 graphs or on unequal-variance Gaussian noise; low SHD there would support the shared-prior story, while collapse to baseline SHD would bound the claim to the training generator.
- Editorial: Because each variable is one token and the output is $d\times d$, the architecture does not obviously scale beyond the tested $d\le 20$ regime; adapting it to hundreds of variables would likely require latent grouping or hierarchical attention.
- Editorial: ADAG's training loss is unsupervised, so labeled or interventional samples, when available, could be added as a fine-tuning signal; the paper does not explore this, but it is a natural route to relax the equal-variance Gaussian identifiability dependence.
- Editorial: The theoretical section relies on prior identifiability results and explicitly leaves the small-sample joint identifiability proof to future work; if the empirical prior effect is real, it suggests the sample-complexity conditions of the order-recovery result can be relaxed.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript proposes Attention-DAG (ADAG), a linear-transformer-based model that learns a nonlinear kernel mapping from observational data matrices to weighted adjacency matrices for linear structural equation models. ADAG is pre-trained across many synthetic domains by minimizing a reconstruction loss subject to an acyclicity constraint, then applied to held-out domains in a zero-shot fashion. The paper targets two multi-domain settings: heterogeneous data, where all domains share the same DAG structure with varying edge weights, and order-consistent data, where domains have different DAGs that share a common topological order. The experiments report strong accuracy and very fast inference on synthetic ER1 graphs with d=5,10,20, additional low-sample robustness experiments, an ablation on the number of training domains, non-Gaussian noise ablations, and a Sachs-dataset transfer experiment. The central claim is that a single pre-trained model can zero-shot infer DAGs for new order-consistent or heterogeneous domains and that the pre-trained kernel captures a shared low-dimensional prior that mitigates small-sample ill-posedness.
Significance. If the zero-shot generalization claim is established, ADAG would be a genuinely useful contribution: it would be among the first practical foundation-model-style approaches to DAG learning, with the attractive property that inference at test time is a single forward pass rather than iterative optimization. The paper also has concrete strengths: the training objective is unsupervised and, in principle, does not require ground-truth graphs; the authors provide extensive implementation details, ablation studies on domain count and noise type, and low-sample experiments; and the reported inference speed is orders of magnitude faster than optimization-based baselines. However, the evidence as presented does not fully support the central generalization claim because the experimental design allows structure-level overlap between training and test domains in the d=5 order-consistent setting, and by construction in the heterogeneous setting. The comparison with baselines also conflates the pre-training information advantage with algorithmic superiority.
major comments (4)
- [§4.2, Table 1 (d=5, order-consistent)] The zero-shot comparison does not test genuinely unseen DAG structures. With a fixed topological order on d=5 nodes, the number of possible DAGs is 2^(d(d-1)/2)=1024, while the training protocol uses M=70,000 domains (Appendix C.2). Since all order-consistent domains share the same order (footnote in §2 and dataset settings in §4), every one of the 1024 possible structures appears many times during pre-training. The held-out test domains therefore introduce new data and new edge weights, but almost never a new DAG structure. The SHD of 0.038 in Table 1 and the low-sample results in Table 2 consequently measure within-distribution structure recognition (or memorization) rather than zero-shot generalization to unseen DAGs. This is load-bearing for the paper's central claim, and the d=10/20 results are less affected but do not repair the flagship d=5 experiments. To support the claim, report evaluations on structures excluded from training, with a train/test split performed at the level of graph structures.
- [§4.1-4.2, heterogeneous setting] In the heterogeneous experiments, the test domains are generated from the same DAG as the training domains, as stated in the footnote in §2 and in §4. Thus the heterogeneous setting never presents a new graph structure to the model; it only varies the edge weights. The SHD values of 0.027-0.271 in Table 1 show that the learned kernel can adapt to new causal mechanisms for a fixed structure, but they do not demonstrate structural generalization. The conclusion's statement that ADAG 'generalize[s] effectively to test domains with unseen DAGs and mechanisms' is therefore not supported by the heterogeneous experiments, and the abstract's 'unseen tasks' wording should be qualified accordingly.
- [§4.2, baseline comparison] The comparison against baselines conflates the pre-training information advantage with algorithmic superiority. ADAG is pre-trained on 70,000 domains generated from the same generative protocol as the test domains, while NOTEARS and DAGMA are applied independently to each test domain with no access to training domains, and MetaDAG and MultiDAG are run only on the 1000 test domains. The large apparent gains in SHD, relative error, and runtime therefore include the benefit of having seen 70,000 domains from the target distribution. This is the intended advantage of a foundation model, but the paper's phrasing 'consistently outperforming all state-of-the-art baselines' should be contextualized, and the claim of a fair comparison in §4.2 does not explain why the baselines were not given an equivalent multi-domain training stage.
- [§3.3 and Appendix D] The paper's claim that pre-training on multiple domains reduces the ill-posedness of small-sample DAG learning is not yet backed by a proof for the regime considered. Appendix D.1 explicitly states that the identifiability conditions from Chen et al. (2021) may be violated when n is comparable to d, and that a relaxation is left to future work. Similarly, Lemma D.1 provides an RKHS identifiability statement for the parameter matrix, but it is not shown how the attention-based architecture implements or approximates this RKHS. The empirical low-sample results in Table 2 are valuable, but without a proof or a precise statement of the conditions under which ADAG improves identifiability, the theoretical claim in the abstract and in §3.3 should either be removed or clearly labeled as a conjecture.
minor comments (5)
- [§4.3, second paragraph] The typo 'ADGA' should be corrected to 'ADAG'.
- [§4.1, Figure 2] The caption says there are 'five dominant eigenvalues consistent with the structure of an ER1 graph with d=5', but the plotted eigenvalues appear to be from the weighted adjacency matrices, not from the data covariance. Please clarify in the caption and the text what object is being decomposed.
- [Table 1, d=5 row] The runtime of MultiDAG differs by two orders of magnitude between the heterogeneous setting (0.0151 s) and the order-consistent setting (1.9911 s), which is a surprising discrepancy given the same problem size; please explain this difference.
- [§4.2, threshold choice] A fixed threshold of 0.3 is applied to the inferred weighted adjacency matrices for all methods and all settings. Since the scale of ADAG's kernel output may vary with architecture and data, a sensitivity analysis for this threshold would strengthen the empirical evaluation.
- [Appendix C.1, Table 3] The Sachs evaluation uses a pre-training set generated from an assumed ordering of the Sachs variables plus synthetic DAGs with known ground truth. This is an oracle-assisted transfer setting, and the main text should state that the ordering assumption is used in the pre-training stage; otherwise readers may over-interpret Table 3 as a fully unsupervised real-data result.
Circularity Check
No circular derivation; the training and held-out evaluation loop is self-contained, with only a mild self-citation burden for the shared-prior mechanism.
full rationale
ADAG's derivation is not self-referential: the model is trained by minimizing the reconstruction loss in Eq. (9) under the acyclicity constraint, and the zero-shot numbers in Tables 1 and 2 are computed on 1000 held-out domains (Section 4.2), not on the training domains. The mapping in Eqs. (4)-(7) is a learned nonlinear function, and the test statistics are not encoded in the loss. The main self-citation is Section 3.3's appeal to Lu and Yu [2025] for the claim that a linear transformer discovers a shared low-dimensional prior, and Appendix D.2's referral to Yu et al. [2024] and Lu and Yu [2025] for identifiability of the learned map; these motivate the architecture, but the empirical generalization results and Appendix D.1's external identifiability citations provide independent content. A separate validity caveat, not a circular step, is that for order-consistent d=5 data, there are only 2^10 = 1024 DAG structures for a fixed order while M = 70,000 training domains are used, so all test structures appear in training; the d=5 SHD values therefore do not establish structural generalization to genuinely novel DAGs, although the d=10/20 settings, with astronomically larger structure spaces, do provide such evidence.
Assumptions & free parameters
free parameters (3)
- Learned attention parameters Θ (W_Q_l, W_K_l, W_P_out, W_Q_out, W_K_out) =
Learned on 70,000 synthetic domains; individual values not reported
- Edge threshold 0.3 =
0.3
- Number of training domains M =
70,000
assumptions (6)
- domain assumption Data are generated by a linear SEM X = A^T X + E with mutually independent additive noise.
- standard math The acyclicity constraint h(A)=tr(e^{A∘A})-d=0 from Zheng et al. 2018 is a valid continuous characterization of DAGs.
- domain assumption Per-domain graph identifiability follows from existing results: non-Gaussian noise (Shimizu et al. 2006) or equal-variance Gaussian noise (Peters and Buhlmann 2014).
- domain assumption Test domains are drawn from the same distribution as training domains (same order or same DAG, same edge-weight and noise distributions).
- ad hoc to paper A linear transformer trained across tasks automatically discovers a low-dimensional shared prior that transfers to new tasks.
- ad hoc to paper The attention network parameterization is expressive enough to represent the data-to-DAG mapping.
Cite this review
Pith. "Pith review of Learning Causal Graphs at Scale: A Foundation Model Approach." pith.science (2026). https://pith.science/paper/2YBPTBKF
@misc{pith2026250618285,
author = {Pith},
title = {Pith review of: Learning Causal Graphs at Scale: A Foundation Model Approach},
year = {2026},
howpublished = {\url{https://pith.science/paper/2YBPTBKF}},
note = {Machine review of arXiv:2506.18285}
}
read the original abstract
Due to its human-interpretability and invariance properties, Directed Acyclic Graph (DAG) has been a foundational tool across various areas of AI research, leading to significant advancements. However, DAG learning remains highly challenging, due to its super-exponential growth in computational cost and identifiability issues, particularly in small-sample regimes. To address these two challenges, in this work we leverage the recent success of linear transformers and develop a foundation model approach for discovering multiple order-consistent DAGs across tasks. In particular, we propose Attention-DAG (ADAG), a novel attention-mechanism-based architecture for learning multiple linear Structural Equation Models (SEMs). ADAG learns the mapping from observed data to both graph structure and parameters via a nonlinear attention-based kernel, enabling efficient multi-task estimation of the underlying linear SEMs. By formulating the learning process across multiple tasks as a continuous optimization problem, the pre-trained ADAG model captures the common structural properties as a shared low-dimensional prior, thereby reducing the ill-posedness of downstream DAG learning tasks in small-sample regimes. We evaluate our proposed approach on benchmark synthetic datasets and find that ADAG achieves substantial improvements in both DAG learning accuracy and zero-shot inference efficiency. To the best of our knowledge, this is the first practical approach for pre-training a foundation model specifically designed for DAG learning, representing a step toward more efficient and generalizable down-stream applications in causal discovery.
Figures
Figures from the paper (2 more)
Reference graph
Works this paper leans on
-
[2]
Gradient-based neural dag learning.arXiv preprint arXiv:1906.02226,
Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, and Simon Lacoste-Julien. Gradient-based neural dag learning.arXiv preprint arXiv:1906.02226,
arXiv 1906
-
[3]
DAG learning accuracy is evaluated using Structural Hamming Distance (SHD), and we also report the number of predicted edges. Hyperparameters, including sparsity constraint coefficients and thresholds, are extensively tuned 15 to optimize SHD performance. We first evaluate NOTEARS and DAGMA on the full set of 853 observations, where both methods perform w...
-
[5]
Transformer learns the cross-task prior and regularization for in-context learning
Fei Lu and Yue Yu. Transformer learns the cross-task prior and regularization for in-context learning. arXiv preprint arXiv:2505.12138,
-
[10]
Causality for large language models.arXiv preprint arXiv:2410.15319,
Anpeng Wu, Kun Kuang, Minqin Zhu, Yingrong Wang, Yujia Zheng, Kairong Han, Baohong Li, Guangyi Chen, Fei Wu, and Kun Zhang. Causality for large language models.arXiv preprint arXiv:2410.15319,
-
[11]
Towards causal foundation model: on duality between optimal balancing and attention
Jiaqi Zhang, Joel Jennings, Agrin Hilmkil, Nick Pawlowski, Cheng Zhang, and Chao Ma. Towards causal foundation model: on duality between optimal balancing and attention. InForty-first International Conference on Machine Learning, 2024a. 12 Kun Zhang, Biwei Huang, Jiji Zhang, Clark Glymour, and Bernhard Schölkopf. Causal discovery from nonstationary/hetero...
work page 2017
-
[13]
We leave such theoretical investigations to a future work
and an improved identifiability property under our foundation model setting. We leave such theoretical investigations to a future work. 17 D.2 Parameter Identifiability ofA In addition to the capability of identifying the common topological ordering across all domains, ADAG is also capable of identifying the weighted adjacency matrix parameters, i.e.,A. U...
work page 2024
-
[14]
and Lu and Yu [2025]. 19
work page 2025
-
[2014]
Identifiability of causal graphs using functional models.arXiv preprint arXiv:1202.3757,
Jonas Peters, Joris Mooij, Dominik Janzing, and Bernhard Schölkopf. Identifiability of causal graphs using functional models.arXiv preprint arXiv:1202.3757,
Show all 14 references
-
[2019]
Supervised whole dag causal discovery.arXiv preprint arXiv:2006.04697,
Hebi Li, Qi Xiao, and Jin Tian. Supervised whole dag causal discovery.arXiv preprint arXiv:2006.04697,
2006 arXiv
-
[2020]
Neural interpretable pdes: Harmonizing fourier insights with attention for scalable and interpretable physics discovery.arXiv preprint arXiv:2505.23106,
Ning Liu and Yue Yu. Neural interpretable pdes: Harmonizing fourier insights with attention for scalable and interpretable physics discovery.arXiv preprint arXiv:2505.23106,
-
[2021]
URLhttps://dx.doi.org/10.1088/1361-6420/ac245d
doi: 10.1088/1361-6420/ac245d. URLhttps://dx.doi.org/10.1088/1361-6420/ac245d. Taiyu Ban, Lyvzhou Chen, Xiangyu Wang, and Huanhuan Chen. From query tools to causal architects: Harnessing large language models for advanced causal discovery from data.arXiv preprint arXiv:2306.16902,
-
[2023]
Asymptotic theory of in-context learning by linear attention.arXiv preprint arXiv:2405.11751,
11 Yue M Lu, Mary I Letey, Jacob A Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention.arXiv preprint arXiv:2405.11751,
-
[2024]
Bridging causal discovery and large language models: A comprehensive survey of integrative approaches and future directions
Guangya Wan, Yuqi Wu, Mengxuan Hu, Zhixuan Chu, and Sheng Li. Bridging causal discovery and large language models: A comprehensive survey of integrative approaches and future directions. arXiv preprint arXiv:2402.11068,
-
[2025]
Meta-dag: Meta causal discovery via bilevel optimization
Songtao Lu and Tian Gao. Meta-dag: Meta causal discovery via bilevel optimization. InICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1–5. IEEE,
2023
Reviewed August 15, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.