REVIEW 4 major objections 5 minor 14 references
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
T0 review · 4 major / 5 minor · reviewed 2026-08-04 · deepseek-v4-flash
Pith's one-line read This paper claims that a foundation model pre-trained on a heterogeneous mixture of LiNGAM, ANM, HNM, and PNL mechanisms converges to a unique, identifiable causal DAG from observational data, and demonstrates the claim with a leaf/parent p
desk verdict Solid empirical foundation-model work whose headline identifiability theorem doesn't actually cover the continuous prior it trains on; send it to referees, but expect a major theory fix. 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 proof machinery is the Quantile Partial Effect (QPE) — the covariate-gradient of the conditional quantile function — which for the four families takes a finite linear-span form with known basis functions: (1) for LiNGAM, (1) for ANM, (1,y) for HNM, and h for PNL. Non-overlap of these spans lets the posterior identify the mechanism family; a Wronskian condition then singles out the unique effect variable, and recursion over leaves identifies the whole DAG. On the architecture side, a Tabular Interaction Block and two bilinear heads implement leaf and parent prediction, and a Mixture-of-Leaf-Experts router selects the family-specific expert.
What would settle it
For a PNL mechanism Y=h(g(X)+U) with h from the paper's continuous warping family (e.g., sinh-arcsinh), compute the Wronskian W_X(s_{X_j}, eta_X) defined in Assumption A.2(iv) at randomly sampled parameter values. If it is ever identically zero, the additional identifiability condition fails and Theorem 4.5 does not cover that part of the prior space; a null result across many samples would support the paper's finiteness-as-practice stance.
Extended reading notes
Core claim
The central claim is Theorem 4.5: under causal sufficiency, faithfulness, and the requirement that each local mechanism belongs to one of four identifiable families (LiNGAM, additive noise, heteroscedastic noise, post-nonlinear), the posterior probability of the true DAG is one, P(G|p(X))=1, so the DAG is identifiable almost surely from the observational distribution. Theorem 4.4 establishes that the posterior also identifies which mechanism family generated each effect. The authors argue that if the pre-training prior satisfies these conditions, then the amortized model Q_theta trained to maximize log-likelihood converges to the posterior, and hence to the unique graph. The method then cons
Load-bearing premise
The theorem's main identifiability condition assumes only finitely many post-nonlinear warping functions can occur, but the paper's actual prior samples warps continuously, so the guarantee rests on a finiteness approximation that the paper itself marks as 'highly likely' rather than proven.
Editorial extensions
If this is right
- If Theorem 4.5 holds, the amortized model's MAP estimate converges to the true DAG as the sample size grows, so the method does not stop at a Markov equivalence class.
- Heterogeneous mechanisms can be mixed inside one prior without breaking identifiability, so a single foundation model can serve datasets whose edges follow different FCM families.
- The leaf/parent decomposition guarantees a valid DAG at every inference step and avoids quadratic adjacency prediction, enabling the reported scaling to 1,000 variables.
- The routing probabilities from MoLE surface the inferred mechanism family, giving interpretability and a hook for human priors.
- On the paper's benchmarks, DAG-FM reports higher precision/recall/F1 and lower structural Hamming distance than classical FCM algorithms and four concurrent amortized models, including on real-world data where several baselines run out of memory.
Reading between the lines
- The theoretical guarantee is narrower than the empirical prior: the paper's continuous warping family violates its own finiteness assumption, so the strict identifiability result applies only to a finite subset of the prior; an unconditional proof would require either discretizing the warping family or extending the QPE framework to infinite families.
- Because errors in leaf prediction cascade through the recursive pruning, the method's accuracy at high dimensions will depend on the reliability of early leaf choices; the reported scaling results could be stress-tested by measuring order divergence as a function of dimension.
- The QPE basis-function lens suggests a direct test for adding new mechanism families: check whether their QPE spans lie outside the existing span; if they overlap, identifiability is lost.
- The MoLE router's family predictions could be used as a deployment diagnostic: high-entropy routing on a new dataset would signal that the data-generating mechanism falls outside the four identifiable families, warning the user not to trust the recovered graph.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes DAG-FM, a two-stage amortized causal-discovery foundation model: a leaf-node predictor and a parent-node predictor, both implemented with a tabular interaction backbone and a Mixture-of-Leaf-Experts (MoLE) head. The training prior is designed to include LiNGAM, ANM, HNM, PNL, and a general mechanism family. The main theoretical claim is that, under Assumption 4.3 plus additional identifiability conditions, the posterior over DAGs is almost surely degenerate, so the amortized model asymptotically recovers the true DAG; the main empirical claim is state-of-the-art DAG recovery on synthetic heterogeneous benchmarks, homogeneous benchmarks, and two real-world datasets, with scaling to large n and d.
Significance. If the theoretical guarantee were valid, this would be an important step: it would give a principled design condition for amortized causal discovery under heterogeneous mechanisms and would place a foundation-model approach on solid identifiability ground. The empirical study is broad: it includes 24 baselines, multiple metrics with confidence intervals on synthetic data, OOD settings, scalability curves, and an ablation of the MoLE component. The architecture is well motivated by the order-then-prune decomposition, and the scalability to 24GB VRAM is practically relevant. However, the central theoretical result currently does not cover the prior space actually used, and some of the empirical superiority claims are not supported by the reported uncertainty. These issues are substantial but potentially addressable in revision.
major comments (4)
- [§4.1, Appendix A (Assumption A.4), Appendix B.1] Assumption A.4 states that only finitely many post-nonlinear functions h are possible in the prior space, and Corollary A.5 builds the finite basis b = (1, y, h_1, ..., h_k) on this assumption. Section B.1, however, defines the PNL warping functions as continuous families: sinh-arcsinh with continuous (epsilon, tau), residual flow with continuous weights and base slope, and asymmetric power with continuous mu and exponents. The support of the prior over h is therefore uncountable, so A.4 is false for the prior used for training and testing. The paper's defense that only finitely many h are sampled during a finite training run changes the object of the theorem: if A.4 applies only to the finite training realization, Theorem 4.5 guarantees identifiability only on that realization, not on new test SCMs drawn from the continuous prior. This is a load-bearing gap for the headline claim that D
- [§4.1, Theorem 4.5 proof (Appendix A)] The proof of Theorem 4.5 is too sketchy for a central result. In the 'superfluous forward edge' and 'missing forward edge' cases, the argument relies on statements such as 'Markov blanket' and 'no reverse model can yield ... specific directional properties' without a formal derivation from causal sufficiency and faithfulness to the claimed contradictions. The 'reversed edge' case invokes Corollary A.5 but does not precisely show why the existence of two effect variables in a subset contradicts the corollary after conditioning on the rest of the graph. Since Theorem 4.5 is the main theoretical contribution, a rigorous, step-by-step proof is needed. If the proof is essentially that of Lemma A.3 applied recursively, the authors should state that reduction explicitly and verify all conditions after each leaf removal.
- [§4.1, Table 4 and Theorem 4.4] Theorem 4.4 claims that the QPE spaces of the four mechanism families are mutually non-overlapping and that therefore P(F | p(X,Y)) = 1. The proof rules out degenerate cases by saying, for example, that for ANM '∇g will never degenerate into a constant vector' because g is non-linear. But the prior in Section B.1 samples random functions from continuous distributions; there may be measure-zero but nonempty sets where the QPE of one family lies in the span of another family's basis (e.g., an HNM with h(x) constant on a region, or a PNL whose h(y) is locally affine). The proof should quantify these exclusions or show that they have prior probability zero so that the posterior identifiability claim holds almost surely. Without this, Theorem 4.4 is not fully established even for the finite-family case.
- [§5, Table 1 and Table 3] The abstract and Section 5 claim that DAG-FM 'significantly outperforms' other amortized methods. In Table 1, on the Hetero benchmark, the F1 confidence intervals for DAG-FM (0.67 ± 0.15) and FoundCause (0.57 ± 0.16) overlap substantially; nSHD also overlaps (0.17 ± 0.17 vs 0.23 ± 0.17). Only precision appears non-overlapping (0.71 ± 0.14 vs 0.57 ± 0.22), but that is a weak basis for a general 'significantly outperforms' claim. Table 3 reports real-world results without any error bars, making it impossible to assess the significance of the reported improvement on Sachs and Causal Chamber. The authors should add error bars to Table 3 (e.g., bootstrap over data subsets or multiple restarts) and apply a statistical test or at least report overlap-aware intervals for Table 1. If the claim is meant only as 'best point estimates', it should be phrased that way.
minor comments (5)
- [§4.4, Algorithm 1] The threshold alpha is fixed at 0.5 in the text but appears as a free parameter in Algorithm 1. If alpha is not tuned, state so explicitly; if it is, report sensitivity analysis.
- [§5, Real-World Benchmarks] The caption of Table 3 says 'compared with other amortized algorithms' but the table includes only four methods; the full comparison in Table 12 has many 'Out of Memory' entries. A footnote explaining the OOM conditions (e.g., GPU memory limit) would help.
- [Appendix A, after Eq. (7)] The QPE basis for PNL in Table 4 is written as h = h'(h^{-1}), but the notation is not defined at that point; define h, h', and h^{-1} explicitly.
- [§2, Related Work] The reference 'FoundCause (Blöbaum et al., 2026)' is cited without a publication venue. If it is a preprint, say so; if it is not available, the empirical comparison should be reproducible from a public checkpoint, which the paper does state.
- [Appendix C, OOD] In Figure 7, the 'Mechanism' OOD case uses PNL-HNM, but the theoretical identifiability of PNL-HNM is described as open. The text should clearly separate the empirical OOD result from the theoretical guarantee.
Circularity Check
Theorem 4.5's identifiability guarantee rests on a uniqueness lemma imported from the authors' own prior work and on Assumption A.4, which is made to hold by redefining the prior as the finite training realization; the empirical SOTA claims on external benchmarks remain independently grounded.
-
uniqueness imported from authors
[Appendix A, Lemma A.3 and Corollary A.5; footnote 1 and Theorem 4.5 proof]
"Lemma A.3 (Chen et al. 2026, Corollary 3.8). If Assumption A.2 holds, then the variable Xi satisfying condition (iii) in Assumption A.2 is the unique effect variable. ... The proof is then completed by applying Lemma A.3."
Theorem 4.5's conclusion P(G|p(X)) = 1 is obtained by Corollary A.5, whose proof says 'The proof is then completed by applying Lemma A.3.' Lemma A.3 is a uniqueness theorem cited to Chen et al. 2026, whose first author, Yikang Chen, is also the first author of this paper. The lemma is not proved here, and the 'additional identifiability conditions' in footnote 1 point exactly to the assumptions needed to invoke it. Thus the central theoretical guarantee is a self-citation chain: the authors' prior uniqueness theorem is used to declare the DAG forced, rather than being derived independently in this paper. The empirical benchmarks are external, so the circularity is confined to the theoretical layer.
-
self definitional
[Assumption A.4 versus Section B.1 (Post-processing / PNL warps)]
"Assumption A.4. For the post-nonlinear function h in PNL, assume that there is only a finite number of possible functions h = h'(h^{-1}) within the prior space. ... Since the training only proceeds for a finite number of steps, only a finite number of post-nonlinear functions h are sampled from the PNL mechanisms. ... Sinh-Arcsinh: ... skewness parameter ε, sampled continuously from [−1.0,1.0] ... tail-weight parameter τ, sampled continuously from [0.5,2.0]."
Corollary A.5's finite basis b=(1,y,h1,...,hk) requires Assumption A.4, but the prior space used in Section B.1 samples PNL warps from continuous families, so its support is uncountable. The paper's justification of A.4 counts only the finite number of SCMs realized during training and 'consequently' treats the prior as finite. This changes the object: the theorem's guarantee applies to the finite set of already-sampled h, not to the continuous prior from which unseen test SCMs are drawn. The claim that the actual prior space guarantees identifiability is then equivalent to the training-realization assumption, i.e., the conclusion is made to fit the QPE finite-basis framework by defining the prior to be finite.
full rationale
Most of the empirical content is independent: the synthetic benchmarks are held-out SCMs sampled from the same prior (a standard in-distribution held-out evaluation), and Sachs/Causal Chamber are external real-world datasets; those results are not fitted parameters renamed as predictions. The circularity is confined to the theoretical guarantee. Theorem 4.5 depends on Corollary A.5, whose proof is completed by Lemma A.3, quoted as 'Chen et al. 2026, Corollary 3.8' — a uniqueness lemma from the same author group, invoked as an external mathematical fact. In parallel, Assumption A.4 is needed to construct the finite QPE basis, but the actual PNL prior is continuous; the paper's defense that only finitely many functions are realized during training redefines the quantified 'prior space' and therefore makes the identifiability guarantee apply to a finite training realization rather than the generative prior. These are load-bearing definitional/citation issues in theory, not in the externally benchmarked empirical claims. The paper also honestly flags that the 'General' mechanism lacks a rigorous guarantee, further narrowing the theoretical coverage. Overall: some self-citation and a definitional gap, but the central empirical result has independent content.
Assumptions & free parameters
free parameters (6)
- parent threshold α =
0.5
- MoLE top-k =
k=1
- training table size n =
1024
- mechanism mixture probabilities =
uniform over the six train families
- SNR range for noise injection =
[-5, 10] dB uniform
- expected graph density and warp/discretization fractions =
0.25 (truncated exponential on [0,1])
assumptions (6)
- domain assumption Causal sufficiency and faithfulness (Assumption 4.3(i))
- domain assumption QPE identifiability machinery of Chen et al. (2026), imported wholesale (Table 4; Lemma A.3 = Chen et al. Cor. 3.8)
- ad hoc to paper Assumption A.2(iv): Wronskian determinants W(s,η) ≠ 0 for relevant variable pairs
- ad hoc to paper Assumption A.4: finitely many PNL functions h in the prior space
- domain assumption Mutual non-overlap of the four FCM families' QPE spaces (Theorem 4.4)
- standard math Density regularity: densities exist, strictly positive, at least C^k
invented entities (3)
-
General causal mechanism f(X,U) with U as input dimension
-
PNL-HNM mechanism Y=h(a(X)+b(X)U)
-
Dummy expert in MoLE
Cite this review
Pith. "Pith review of DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms." pith.science (2026). https://pith.science/paper/FLVQH5D6
@misc{pith2026260711510,
author = {Pith},
title = {Pith review of: DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms},
year = {2026},
howpublished = {\url{https://pith.science/paper/FLVQH5D6}},
note = {Machine review of arXiv:2607.11510}
}
read the original abstract
Causal discovery from observational tabular data remains fundamentally challenging, primarily due to the heterogeneity of underlying causal mechanisms and the high-dimensional combinatorial search space of Directed Acyclic Graphs (DAGs). In this paper, we propose \textbf{DAG-FM}, a novel foundation model architecture that amortizes causal discovery. Unlike direct matrix prediction, DAG-FM decomposes the causal discovery process into two auto-regressive stages using two specialized Transformer-based sub-modules: a leaf-node predictor and a parent-node predictor. To effectively model complex row-column interactions, we adopt a robust tabular interaction block to output feature-wise representations. Crucially, to handle diverse and unknown Functional Causal Model (FCM) assumptions in real-world scenarios, we introduce Mixture-of-Leaf-Experts (MoLE), allowing the model to dynamically route and adapt to identifiable mechanism families. Through an iterative inference algorithm, DAG-FM seamlessly extracts causal orderings and constructs valid DAGs. Extensive experiments demonstrate that DAG-FM achieves state-of-the-art performance on both synthetic benchmarks and complex real-world datasets, significantly outperforming traditional classical algorithms and recent foundation models in both accuracy and scalability.
Figures
Figures from the paper (13 more)
Reference graph
Works this paper leans on
-
[1]
Post-processing
We then sample 1D random noise variables i.i.d. from these randomly instantiated distributions. Subsequently, these noise samples are normalized and participate in the forward inference of the random causal mechanisms with controlled signal-to-noise ratios (SNR). Random MechanismsThe building blocks of various causal mechanisms are randomly sampled functi...
2025
-
[7]
with respect to the given basis functionsb, it can be shown thatW Xi (si,j,η i)≡0 for anyX j ∈X −i and anyx(Chen et al., 2026, Theorem 3.6 and Corollary 3.8). Assumption A.2 couples this finite linear span form with the cause-and-effect relationship, and rules out coincidences where this necessary condition also identically holds for other variablesX k ∈X...
2026
-
[20]
Except for the specified mechanisms, all other configurations are identical to the prior space construction during training. When evaluating specific properties, we vary only the target hyperparameter during sampling (e.g., varying only sample sizenwhen evaluating sensitivity to sample size), while keeping all other settings identical toHetero. Real-World...
2005
-
[1992]
Bert: Pre-training of deep bidirectional transformers for language understanding
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. InProceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pp. 4171–4186,
2019
-
[2002]
0.23 0.45 0.31 0.18 0.76 0.500.620.55 0.38 0.19 ICA-LiNGAM (Shimizu et al., 2006)0.20 0.45 0.28 0.16 0.85 0.05 0.05 0.05 0.03 0.41Direct-LiNGAM (Shimizu et al., 2011)0.19 0.35 0.25 0.14 0.76 0.05 0.05 0.05 0.03 0.40 CAM (B¨uhlmann et al.,
2006
-
[2005]
doi: 10.1126/science.1105809. Shohei Shimizu, Patrik O. Hoyer, Aapo Hyv ¨arinen, and Antti Kerminen. A linear non-gaussian acyclic model for causal discovery.Journal of Machine Learning Research, 7(72):2003–2030,
-
[2013]
doi: https://doi.org/10.1016/j.cell.2013.03.030
ISSN 0092-8674. doi: https://doi.org/10.1016/j.cell.2013.03.030. Jiaru Zhang, Rui Ding, Qiang Fu, Huang Bojun, zizhen Deng, Yang Hua, Haibing Guan, Shi Han, and Dongmei Zhang. Learning identifiable structures avoids bias in DNN-based supervised causal learning. InThe 28th International Conference on Artificial Intelligence and Statistics, 2025a. Kun Zhang...
-
[2016]
ISSN 0300-5771. doi: 10.1093/ije/dyv341. 12 DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V on Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwan...
Show all 14 references
-
[2018]
13 DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms A Proofs We provide the proofs for the theories in the main text under the identifiability framework based on Quantile Partial Effects (QPE) (Chen et al., 2026). This framework decouples t...
2026
-
[2019]
Out of Memory Out of MemoryDAGMA-Linear (Bello et al., 2022)0.20 0.10 0.13 0.07 0.47 0.04 0.03 0.03 0.02 0.32DAGMA-Nonlinear (Bello et al., 2022)0.25 0.20 0.22 0.12 0.51 Out of Memory A VICI (Lorch et al.,
2022
-
[2020]
0.10 0.10 0.10 0.05 0.65 0.07 0.05 0.06 0.03 0.34GraNDAG (Lachapelle et al., 2020)0.00 0.00 0.00 0.00 0.40 1.000.03 0.05 0.03 0.20DAG-GNN (Yu et al.,
2020
-
[2022]
doi: 10.3390/jmmp6010010
ISSN 2504-4494. doi: 10.3390/jmmp6010010. Johnny Xi, Hugh Dance, Peter Orbanz, and Benjamin Bloem-Reddy. Distinguishing cause from ef- fect with causal velocity models. InForty-second International Conference on Machine Learning,
-
[2025]
0.23 0.50 0.31 0.19 0.80 0.07 0.13 0.09 0.05 0.53 NOTEARS-Linear (Zheng et al., 2018)0.33 0.10 0.15 0.08 0.40 0.15 0.05 0.08 0.04 0.25NOTEARS-Nonlinear (Zheng et al., 2018)0.07 0.05 0.06 0.03 0.58 0.18 0.10 0.13 0.07 0.28GOLEM (Ng et al.,
2018
-
[2026]
Mooij, Dominik Janzing, and Bernhard Sch ¨olkopf
Jonas Peters, Joris M. Mooij, Dominik Janzing, and Bernhard Sch ¨olkopf. Causal discovery with continuous additive noise models.Journal of Machine Learning Research, 15(58):2009–2053,
2009
Reviewed August 4, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.