DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms Yikang Chen1 , Zhengkang Guan1 , Haoyuan Qian1 , Peng Cui2 , Yi Yang1 , Kun Kuang1∗ Zhejiang University 2 Tsinghua University 1
arXiv:2607.11510v1 [cs.LG] 13 Jul 2026
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 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.
1
Introduction
Causal discovery aims to recover causal structures from data (Pearl, 2009), playing a pivotal role in diverse fields such as bioinformatics (Zhang et al., 2013), epidemiology (Vandenbroucke et al., 2016), sociology (Huber, 2024), and manufacturing (Vuković & Thalmann, 2022). A primary objective of this task is to identify the underlying Directed Acyclic Graphs (DAGs) from observational data. Existing approaches, spanning constraint-based (Spirtes & Glymour, 1991; Spirtes, 1995), score-based (Cooper & Herskovits, 1992; Chickering, 2002), and optimization-based (Zheng et al., 2018; Bello et al., 2022) methods, have been extensively developed. However, without additional assumptions, these methods can only identify causal structures up to Markov equivalence classes. Functional Causal Models (FCMs) introduce supplementary structural assumptions into Structural Causal Models (SCMs), ensuring theoretical identifiability of the underlying DAGs from observational distributions when specific conditions are met. Notable examples include Linear NonGaussian Acyclic Models (LiNGAMs, Shimizu et al. (2006)), Additive Noise Models (ANMs, Hoyer et al. (2008)), Heteroscedastic Noise Models (HNMs, Tagasovska et al. (2020)), and PostNonlinear models (PNLs, Zhang & Hyvärinen (2009)). Building upon the theoretical guarantees of DAG identifiability, FCM-based causal discovery methods have been developed for empirical DAG recovery. However, these methods typically assume homogeneous causal mechanisms, fundamentally limiting their generalizability. To render statistical hypothesis testing tractable, most existing approaches are narrowly tailored to specific theoretical frameworks, such as LiNGAM or ANM (Shimizu et al., 2006; 2011; Peters et al., 2014; Rolland et al., 2022). Alternatively, some methods derive causal orderings by extracting specific statistical features, yet they remain effective only under highly restrictive conditions (Reisach et al., 2021; 2023). Consequently, the robustness and validity of these traditional methods under model misspecification or within highly complex real-world scenarios remain largely questionable. ∗
Corresponding author: [email protected]
1
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
In parallel, amortized algorithms have emerged for causal discovery, empirically demonstrating consistent superiority over traditional methods in complex nonlinear data regimes (Lorch et al., 2022; Dhir et al., 2025; Thompson et al., 2026; Peng et al., 2026; Li et al., 2026). These approaches frame causal discovery as a table-to-graph prediction task, leveraging a foundation model pre-trained on synthetic data generated from complex prior spaces to predict causal graphs for unseen data. However, the correctness and performance of this learning paradigm are highly sensitive to the design of the prior space. For example, in bivariate causal discovery, Dhir et al. (2025) showed that if the prior space is non-identifiable, the output is ambiguous; conversely, Montagna et al. (2025) demonstrated that models trained on well-specified, identifiable prior spaces can achieve consistent results. To address both the limited generalizability of FCM-based methods under heterogeneous mechanisms and the lack of theoretical guarantees in amortized inference, we propose integrating these two paradigms. To this end, we introduce DAG-FM, a foundation model for causal discovery that intrinsically guarantees the output of strictly identifiable DAGs solely from observational data. Specifically, our contributions bridge theory and practice: Theoretically, we establish a design condition detailing how to heterogeneously incorporate FCM assumptions into the prior space. This formulation guarantees that the amortized model converges to a unique, identifiable DAG. Empirically, we decouple the DAG identification process into two synergistic sub-tasks: leaf-node prediction and parent-node prediction, conducting large-scale pre-training for each. Together, these two learned modules constitute DAG-FM, collaborating sequentially to reconstruct the DAG. Furthermore, DAG-FM exhibits exceptional computational efficiency. Without relying on KV caching or half-precision optimizations, it can consistently process either up to 50,000 samples or 500 variables, or jointly handle complex datasets containing 20,000 samples and 100 variables, all strictly within a 24GB VRAM constraint. Extensive experiments demonstrate that DAG-FM achieves state-of-the-art performance, significantly outperforming traditional approaches on both synthetic and real-world datasets, while also surpassing contemporary causal discovery foundation models across multiple evaluation metrics. In summary, our primary contributions are as follows: • We provide a design condition for the prior space that allows for heterogeneous causal mechanisms, while theoretically ensuring that the predictions of the amortized method converge to a unique DAG from observational data. • We develop DAG-FM, a foundation model for causal discovery consisting of two sub-modules, which guarantees the accurate prediction of a unique DAG directly from observational data. • We provide extensive experiments demonstrating that DAG-FM achieves state-of-the-art performance on both synthetic and real-world benchmark tests.
2
Related Works
FCM-based Causal Discovery FCM-based causal discovery aims to recover the underlying causal structure by restricting the causal mechanisms within a SCM. Theoretically, to achieve the ultimate goal of DAG identifiability, FCM approaches impose specific structural constraints on the underlying causal mechanisms. For instance, LiNGAM (Shimizu et al., 2006) elegantly established that causal directions are uniquely identifiable under the assumptions of linear causal relationships and non-Gaussian additive noise. Subsequently, Hoyer et al. (2008) extended this identifiability to ANM with nonlinear causal mechanisms and general noise distributions. Building upon this, Zhang & Hyvärinen (2009) further generalized the identifiability conditions to post-nonlinear settings, characterizing an even broader space of identifiable models. More recently, these identifiability conditions have been further relaxed to accommodate heteroscedastic noises (Immer et al., 2023; Strobl & Lasko, 2023), while contemporary theoretical efforts (Xi et al., 2025; Chen et al., 2026) strive to generalize FCM framework to even broader conditions. In practice, a corresponding lineage of multivariate causal discovery algorithms has been explicitly developed to translate these theoretical guarantees into empirical DAG recovery. For example, within the LiNGAM framework, algorithms such as Direct-LiNGAM (Shimizu et al., 2011) and ICA-LiNGAM (Shimizu et al., 2006) are capable of identifying the entire DAG. For ANMs, a diverse suite of methods has been deployed, encompassing RESIT (Peters et al., 2014), CAM (Bühlmann et al., 2014), SCORE (Rolland et al., 2022), DAS (Montagna et al., 2023b), and NoGAM 2
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
(Montagna et al., 2023a). Similarly, addressing HNMs, recent algorithmic developments include SKEW (Lin et al., 2025) and HOST (Duong & Nguyen, 2022). However, these classical methods heavily premise on the homogeneity of the imposed mechanisms. Consequently, their robustness in handling heterogeneous mechanisms or severe model misspecifications remain largely uncertain. Amortized Causal Discovery Amortized causal discovery employs a large structured-data model (LDM) to predict causal graphs directly from tabular data by treating structured data as joint distributions (Zhang et al., 2025b). This predictive model is pre-trained by sampling random causal graphs along with their corresponding observations from an artificially constructed prior space. Unlike tabular foundation models, which utilize in-context learning to predict missing entries at the cell level, amortized causal discovery operates at the variable level to infer the causal relationships among variables. For instance, AVICI (Lorch et al., 2022) employs an Axial Transformer to model interactions across tabular features, aggregates feature representations through max-pooling, and utilizes a bilinear head to predict the probability of edge existence in the adjacency matrix. Similarly, CSIvA (Ke et al., 2023) predicts edge probabilities autoregressively. Related approaches, such as CauScale (Peng et al., 2026) and TabCausal (Li et al., 2026), further extend this paradigm by scaling the training and inference processes to larger regimes. However, a significant challenge with these methods is their direct modeling of the adjacency matrix, which fails to formally guarantee that the output is a DAG. BCNP (Dhir et al., 2025) and Arrow (Thompson et al., 2026) address this deficiency by decomposing the prediction task into causal ordering and an upper-triangular matrix. Other strategies focus on alternative forms of causal structures: for instance, SiCL (Zhang et al., 2025a) aims to predict Markov equivalence classes, TCD-DL (Kim et al., 2025) focuses on predicting the ancestor set for a target node, and TabOrder (Xu et al., 2026) is dedicated to predicting the causal ordering of variables.
3
Preliminaries
Structural Causal Model An SCM involving d variables is defined as a triplet ϕ = (f , X, U ), where X and U denote exogenous and endogenous variables, respectively. For each i ∈ I = {1, . . . , d}, the relationship is defined as Xi = fi (Xpa(i) , Ui ), where pa(i) ⊆ I \ {i} denotes the set of parents of Xi , and fi represents the causal mechanism. We assume the SCM is Markovian (Pearl, 2009), implying Xpa(i) ⊥ ⊥ Ui for all Xi . The graph G = (I, E), where (j, i) ∈ E if and only if j ∈ pa(i), is known as the causal graph and is a DAG. A topological ordering τ of this graph defines a causal ordering. Amortized Causal Discovery Let p(ϕ) be a prior distribution over an artificially constructed prior space Φ of SCMs. An SCM ϕ ∼ p(ϕ) with d variables corresponds to a specific DAG G. Observational data D = {xk }nk=1 ∈ Rn×d is sampled from the observational distribution p(X | ϕ). The posterior distribution of the causal graph given the observational data is defined as: Z P (G | D) ∝ P (G | D, ϕ) p(D | ϕ) p(ϕ) dϕ. (1) To predict G given D, an amortized model Qθ (G | D) is constructed to approximate the posterior distribution P (G | D). This is achieved by minimizing the KL(P ∥Qθ ) divergence, which is equivalent to maximizing the log-likelihood EG,D [log Qθ (G | D)]. During training, we perform large-scale sampling of random SCMs ϕ from p(ϕ) to generate corresponding pairs of (G, D) for pre-training the model Qθ . During inference, the model takes D as input and produces G through either sampling or maximum a posteriori (MAP) estimation.
4
DAG-FM: The DAG Foundation Model
4.1
DAG Identifiability via Mechanism Families
The asymptotic identifiability of DAGs in amortized methods hinges on whether the MAP estimator arg maxG Qθ (G | D) converges to a unique ground-truth graph as the model parameters θ are optimized and the sample size |D| increases. By the properties of KL divergence and the Law of Large 3
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Numbers, we have: lim Qθ∗ (G | D) = lim P (G | D) = P (G | p(X)),
|D|→∞
|D|→∞
(2)
where θ∗ = arg minθ KL(P ∥Qθ ) denotes the optimal parameters under which Qθ (G | D) matches P|D| the posterior P (G | D). As |D| → ∞, the empirical distribution p̂D (X) = i=1 δxi (X) converges almost surely to the true observational distribution p(X). Consequently, the convergence of the amortized model to a unique graph is equivalent to the uniqueness of the posterior map arg maxG P (G | p(X)). Definition 4.1 (Almost Sure Identifiability). A DAG G is said to be identifiable almost surely in the posterior if there exists a graph G such that P (G | p(X)) = 1. To guarantee this, we incorporate assumptions rooted in established FCM literature, focusing on four primary classes of causal mechanisms: Definition 4.2 (Mechanism Families F). For a mechanism f (Xpa(i) , Ui ), we define: • FLiNGAM : f (Xpa(i) , Ui ) = w⊺ Xpa(i) + Ui , where Ui is non-Gaussian. • FANM : f (Xpa(i) , Ui ) = gi (Xpa(i) ) + Ui , where gi is an arbitrary non-linear function. • FHNM : f (Xpa(i) , Ui ) = gi (Xpa(i) ) + hi (Xpa(i) ) · Ui , satisfying the conditions in (Immer et al., 2023, Theorem 1). • FPNL : f (Xpa(i) , Ui ) = hi (gi (Xpa(i) ) + Ui ), where hi is a non-linear invertible function, satisfying the conditions in (Zhang & Hyvärinen, 2009, Theorem 1). Assumption 4.3 (Prior Space). (i) Causal sufficiency (no unobserved confounders) and faithfulness (conditional independencies imply d-separation) hold. (ii) For any ϕ = (X, f , p(U )) ∈ Φ, each mechanism fi belongs to the family Fi ∈ {FLiNGAM , FANM , FHNM , FPNL }. In stark contrast to traditional FCM-based methods that assume homogeneous data generation, Assumption 4.3 permits the SCM to exhibit heterogeneous mechanisms. That is, an SCM can be elegantly formulated as a composite mixture of mechanisms drawn from fundamentally different families. The following theoretical results guarantee that, within a prior space satisfying Assumption 4.3, the true DAG is almost surely identifiable from the posterior. Theorem 4.4. Under Assumption 4.3, if Xi = fi (Xpa(i) , U ) entails the observational distribution p(Xpa(i) , Xi ) and fi ∈ Fi , then the posterior probability P (Fi | p(Xpa(i) , Xi )) = 1. In essence, under the heterogeneous setup delineated in Assumption 4.3, Theorem 4.4 demonstrates that the exact nature of the mechanism family can be perfectly discriminated by the posterior, given only the local observational distribution. Theorem 4.5. Under Assumption 4.3, if an SCM ϕ = (X, f , p(U )) entails the observational distribution p(X) and the causal graph G, then the posterior probability P (G | p(X)) = 1. Consequently, Assumption 4.3 establishes a sufficient condition for global DAG identifiability. The profound implication for our foundation model is clear: by structurally designing the prior space so that randomly sampled mechanisms are restricted to the four families defined in Definition 4.2 (while preserving causal sufficiency and faithfulness), the true DAG becomes strictly identifiable in the Bayesian posterior. Crucially, this theoretically ensures that the amortized model will inherently converge, mapping observational data to a uniquely identifiable DAG. 4.2
Prior Space Construction
To construct a prior space Φ that is both expressive and adheres to the theoretical requirements stipulated Assumption 4.3, we design random SCMs using the following components. Random Graphs We employ three generative graph models: Erdős–Rényi, Barabási–Albert, and Watts–Strogatz. The hyper-parameters include the number of nodes d ∈ {2, . . . , 100} (drawn from a uniform distribution) and the graph density ρ ∈ [0, 1] (sampled from a truncated exponential distribution to prioritize sparse graphs). 4
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Random Noise Exogenous variables are sampled from a comprehensive library of distributions: Normal, Laplace, Uniform, Student’s t, Log-Normal, Gumbel, Exponential, χ2 , Beta, and Gamma. Each distribution is parameterized by randomly sampled hyper-parameters. These noise terms are i.i.d.-sampled, standardized, and subsequently injected into the forward causal mechanisms under controlled signal-to-noise ratios (SNRs). Random Mechanisms To implement the underlying causal mechanisms, we consider a diverse repertoire of random functions, including random linear functions, Multi-Layer Perceptrons (MLPs), periodic functions, Gaussian Processes with Random Fourier Features (GP-RFF), and Random Forests. In addition to the four strictly defined mechanism families outlined in Definition 4.2, we introduce a generalized causal mechanism formally defined as f (X, U ). This is implemented by treating the exogenous noise U directly as an intrinsic input dimension during the forward pass. Although this general mechanism currently lacks a rigorous theoretical guarantee for strict DAG identifiability, we empirically observe that the implicit inductive biases of the function approximators still endow it with a non-trivial degree of practical identifiability. Consequently, for each local structure within the randomly sampled DAG, we uniformly sample from these mechanism families to instantiate the respective causal assignment. Subsequently, we inject noise into these mechanisms while rigorously controlling the SNR. The SNR values are sampled systematically from a uniform distribution over [−5dB, 10dB]. Specifically, for a generic causal mechanism Y = f (X, U ), we first compute the empirical variance of the deterministic functional output evaluated with zero noise, i.e., f (X, 0). We then adaptively scale the distribution of the noise U such that the ratio of the noise variance to the functional variance strictly matches the assigned SNR bound. Finally, to prevent the model from exploiting spurious statistical artifacts that could artificially trivialize DAG identification on synthetic benchmarks (Reisach et al., 2021; 2023), we enforce strict distribution normalization. Specifically, immediately following each forward inference step of the structural equation Y = f (X, U ), the generated output Y is dynamically standardized to have zero mean and unit variance (Ormaniec et al., 2025). Post-processing Adopting the data augmentation pipeline from Hollmann et al. (2025), we subject a fraction of the features to non-linear warping and discretization, ranged from a truncated exponential distribution in the interval [0, 0.5]. Warping functions include sinh-arcsinh transformations, residual flows, and asymmetric power functions. Discretization is performed via quantile binning, uniform windowing, and ordinal encoding to simulate real-world data characteristics. Sampling Tabular Data After generating the full feature set, we organize the observations into a tabular dataset D ∈ Rn×d and shuffle the feature order. Unlike Hollmann et al. (2025); Qu et al. (2026), we perform no feature dropout to strictly satisfy the causal sufficiency assumption. While the training sample size is fixed at n = 1024, our amortized model is designed to generalize to varying sample sizes during inference. 4.3
Decomposition of DAG Identification
While Theorem 4.5 guarantees the existence of a unique identifiable DAG under Assumption 4.3, existing amortized methods like Lorch et al. (2022) often parameterize the adjacency matrix directly. This approach faces two primary hurdles: the lack of formal DAG constraints and the quadratic scaling of the output space with respect to dimensionality, which impedes scalability to high-dimensional datasets. Following Dhir et al. (2025); Thompson et al. (2026), we decompose the DAG identification problem into two sub-problems: order identification and full DAG pruning, formulated as: X P (G | p(X)) = P (G | τ, p(X)) P (τ | p(X)). (3) τ
Each sub-problem is further decomposed into recursive prediction tasks. Amortized Order Identification via Leaf Prediction Inspired by order-based causal discovery, we identify the causal ordering τ by recursively uncovering the leaves of the DAG. The task of 5
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
identifying a leaf node i from a set of variables X is framed as arg maxi P (Li = 1 | p(X)), where the event Li indicates whether i is a leaf node in the DAG. This transformation reduces the ordering problem to a sequence of binary classification tasks. Amortized Full DAG Pruning via Parent Prediction Given a causal ordering τ , the DAG skeleton is constrained to an upper triangular structure. To avoid the quadratic growth of predicting the entire adjacency matrix, we recursiveley prune the graph by identifying the parents of each leaf node. Specifically, the identification of a parent j for a confirmed leaf node i is framed as arg max P (Eji | Li = 1, p(X)), where Eji denotes the presence of a directed edge (j → i). This process remains a tractable series of binary classification tasks. Mixture of Experts for Leaf Prediction Since our approach relies on the theoretical conditions of FCM-based causal discovery under Assumption 4.3, we analogously draw inspiration from classical FCM practices. Empirically, discovering varying mechanism families requires developing distinct, tailored algorithms; these identification methods exhibit specific focal preferences and generally operate flawlessly only under strict assumption matching. Motivated by this, to effectively handle heterogeneous causal mechanisms, we design a novel framework denoted as Mixture-of-Leaf-Expert (MoLE). This design encourages different experts to dynamically attend to distinct data features, thereby enhancing the holistic predictive capability of the model across each mechanism family. We refine the leaf node prediction process into two hierarchical steps: (i) identifying the underlying mechanism family; (ii) predicting the leaf node conditioned on that family. This is formalized as follows: X P (Li = 1 | p(X)) = P (Li = 1 | Fi , p(X)) P (Fi | p(X)), (4) Fi
where, according to Theorem 4.4, P (Fi | p(X)) effectively acts as a degenerate distribution. This framework offers three primary advantages: First, it modularizes the amortized inference process by assigning specialized experts to distinct mechanism families. Second, it provides a degree of interpretability by explicitly surfacing the inferred mechanism class and its corresponding theoretical foundation. Third, this framework enables human intervention: one may manually adjust P (Fi | p(X)) to incorporate external expert priors, moving beyond the ”black-box” nature of fully automated amortized inference. 4.4
Foundation Model Architecture
DAG-FM employs two Transformer-based sub-modules, Qθ1 (L | D) and Qθ2 (E:,i | Li , D), to amortize the processes of leaf-node prediction and parent-node prediction, respectively. The architecture is illustrated in Figure 1. Tabular Interaction Block Unlike previous approaches (Lorch et al., 2022; Peng Algorithm 1 DAG-FM Inference Procedure et al., 2026; Li et al., 2026) that utilize Require: Tabular data D, threshold α Axial Transformers to alternate between Ensure: Causal DAG G, causal ordering τ row and column interactions or rely on 1: I ← {1, . . . , d}, τ ← [], E ← ∅ max-pooling for feature-wise representa- 2: while |I| > 0 do tions, we leverage the backbone architec- 3: l ← arg maxi∈I Qθ1 (Li | DI ) ture from TabICLv2 (Qu et al., 2026). 4: S ← {j ∈ I \ {l} | Qθ2 (Ejl | Ll , DI ) > α} The block is composed of four cascaded 5: for s ∈ S do modules: a linear embedding layer, a Set 6: E ← E ∪ {(s, l)} Transformer with four ISAB blocks (Lee 7: end for et al., 2019) for row-wise (sample) inter- 8: I ← I \ {l}, τ ← [τ, l] action, a Transformer with four MABs 9: end while (Vaswani et al., 2017) for column-wise 10: return G, τ (feature) interaction, and a PMA (Lee et al., 2019) to aggregate row-specific information. As shown in Figure 1 (a), this transformation yields a representation of size k × e for each feature, where k denotes the number of seed vectors in the PMA and e the embedding dimension. 6
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
(a)
e
Embedding
1 n
d
Set Transformer
Transformer
Pooling
TFCol
TFRow
PMACol
k×e
n
1 d
d ISAB×4
MAB×4
PMA
k×e
1
1
Tabular Interaction
Leaf Prediction
Tabular Interaction
Parent Prediction
n
d
1 d
d
(b)
(c)
Bilinear Bilinear1 Bilinear2 Bilinearm
d
Top-k Router
k×e k×e
MLP d
d
d
Figure 1: Architecture of DAG-FM. (a) Tabular Interaction Block; (b) Leaf Prediction Head; (c) Parent Prediction Head.
Leaf Prediction Head We integrate the MoLE framework into the leaf prediction head. It consists of m + 1 bilinear heads corresponding to the mechanism families in our prior space, outputting both routing probabilities and family-specific expert predictions. Following Lorch et al. (2022), each PE bilinear head is parameterized as sigmoid( E1 i=1 u⊺i vi + bi ), where u, v are generated by linear and LayerNorm layers. We apply a top-k strategy (setting k = 1) to weight the expert predictions for the final leaf probability. Beyond the four families defined in Definition 4.2, we include an expert for the “general mechanism” and a “dummy expert” to handle source nodes and out-of-distribution cases. Parent Prediction Head We employ an MLP followed by a sigmoid activation to predict parentnode probabilities. To distinguish the query target (the identified leaf) from the potential parents, we position the leaf variable at the first feature input and append a learnable type token (Devlin et al., 2019) during the embedding stage. Training Process
The base models are optimized via the following amortized objectives:
arg min −ED,G [log Qθ1 (L | D)] , θ1
arg min −ED,G,Li [log Qθ2 (E:,i | Li , D)] , θ2
(5)
where D, G are sampled from the prior space, and L is the binary label vector of leaf nodes. Within MoLE, the leaf log-likelihood is factored as: X log Qθ1 (L | D) = log (Qθ1 (L | Fi , D) · Qθ1 (Fi | D)) , (6) Fi
where Qθ1 (Fi | D) denotes the routing probability. All objectives are optimized using the standard Binary Cross-Entropy (BCE) loss function. Inference Process
The DAG identification procedure is detailed in Algorithm 1. 7
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
5
Experiments
Baselines We compare DAG-FM against a diverse set of causal discovery baselines: (i) FCMbased algorithms: Direct-LiNGAM (Shimizu et al., 2011), ICA-LiNGAM (Shimizu et al., 2006), RESIT (Peters et al., 2014), CAM (Bühlmann et al., 2014), RCD (Maeda & Shimizu, 2020), SCORE (Rolland et al., 2022), DAS (Montagna et al., 2023b), and NoGAM (Montagna et al., 2023a) (most are capable of outputting causal orderings); (ii) Score-based algorithms: GES (Chickering, 2002); (iii) Optimization-based algorithms: NOTEARS (Zheng et al., 2018), DAGGNN (Yu et al., 2019), GraNDAG (Lachapelle et al., 2020), GOLEM (Ng et al., 2020), and DAGMA (Bello et al., 2022); (iv) Amortized algorithms: AVICI (Lorch et al., 2022), CauScale (Peng et al., 2026), TabCausal (Li et al., 2026) and FoundCause (Blöbaum et al., 2026). Metrics For DAG identification, we report precision, recall, F1-score, and the normalized Structural Hamming Distance (nSHD). To rigorously evaluate the DAG recovery capability, we penalize SHD by 1 for each reversed edge and an additional 1 for each cycle detected in the output graph. For causal order identification, we report the Order Divergence Rate (ODR) (Rolland et al., 2022). The best results will be highlighted in bold, and the second-best will be underlined. Synthetic Benchmark To robustly evaluate the efficacy of our proposed method under structurally heterogeneous causal mechanisms (as strictly defined in Assumption 4.3), we first construct a comprehensive synthetic benchmark. This is achieved by systematically sampling 100 distinct SCMs from the prior distribution space. The comparative results are summarized in Table 1. Notably, on the DAG identification task, DAG-FM significantly outperforms not only a wide array of traditional classical algorithms but also concurrent causal foundation models, achieving state-of-the-art performance across all evaluated metrics. Table 1: Performance evaluation of DAG identification on the synthetic benchmark featuring heterogeneous causal mechanisms. DAG Performance on Synthetic Benchmark (n = 1000, d = 20) Model
Precision ↑
Recall ↑
F1 ↑
nSHD ↓
Direct-LiNGAM ICA-LiNGAM RCD DAS SCORE NoGAM GES NOTEARS DAGMA
0.220.10 0.280.10 0.080.12 0.430.16 0.380.14 0.390.14 0.450.15 0.540.20 0.430.18
0.250.13 0.360.12 0.020.03 0.450.15 0.530.15 0.550.14 0.530.19 0.170.12 0.190.12
0.220.10 0.310.09 0.030.05 0.420.13 0.430.12 0.450.12 0.480.15 0.240.14 0.250.14
0.390.29 0.360.24 0.300.23 0.260.18 0.290.17 0.280.17 0.270.21 0.240.21 0.260.22
AVICI CauScale TabCausal FoundCause DAG-FM
0.580.19 0.570.18 0.630.18 0.570.22 0.660.14
0.300.19 0.340.19 0.500.19 0.660.18 0.650.19
0.370.20 0.410.19 0.540.17 0.570.16 0.640.14
0.230.21 0.230.21 0.210.21 0.230.17 0.190.17
Furthermore, to demonstrate the precision of DAG-FM under strictly homogeneous assumptions, we also isolate and sample specialized subspaces from the prior distribution where data generation is confined to specific causal mechanism families. This experimental setup rigorously tests the model’s accuracy in causal order identification. As detailed in Table 2, DAG-FM consistently establishes remarkable dominance. Although yielding a marginally close second to the highly specialized Direct-LiNGAM under the pure LiNGAM setting (which is perfectly tailored to that specific assumption) our model substantially surpasses all other FCM-based algorithms across all remaining homogeneous configurations. Real-world Benchmark Finally, we evaluate the generalizability of DAG-FM on real-world datasets: (i) Sachs (Sachs et al., 2005): A protein signaling dataset containing 7,466 samples and 8
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Table 2: Performance evaluation of causal order identification across specific homogeneous subspaces and the heterogeneous space within the synthetic benchmark. Order Performance on Synthetic Benchmark (n = 1000, d = 20) Model
LiNGAM
ANM
HNM
PNL
Hetero.
Direct-LiNGAM ICA-LiNGAM SCORE NoGAM DAG-FM
0.020.03 0.180.15 0.360.14 0.330.13 0.120.09
0.410.13 0.500.13 0.290.14 0.280.14 0.210.12
0.410.15 0.490.16 0.330.16 0.320.15 0.230.13
0.400.14 0.470.12 0.380.14 0.370.14 0.190.11
0.380.15 0.480.12 0.350.15 0.340.14 0.210.12
11 variables; (ii) Causal Chamber (Gamella et al., 2025): A diverse set of physical causal systems containing 10,000 samples across 38 variables. Table 3: Performance evaluation of DAG identification on two real-world benchmarks. Performance on Real-world Benchmarks Model Direct-LiNGAM ICA-LiNGAM RESIT CAM RCD SCORE DAS NoGAM GES NOTEARS DAG-GNN GraNDAG DAGMA AVICI CauScale TabCausal FoundCause DAG-FM
Sachs (n = 7466, d = 11) Precision ↑ Recall ↑ F1 ↑ nSHD ↓ 0.19 0.20 0.22 0.29 0.00 0.20 0.22
0.00 0.20
0.35 0.25 0.45 0.28 0.60 0.32 0.75 0.42 0.00 0.00 0.45 0.28 0.35 0.27 Out of memory 0.45 0.31 0.10 0.15 Out of memory 0.00 0.00 0.10 0.13
0.38 0.50 0.80
Out of memory Out of memory 0.30 0.33 0.25 0.33 0.40 0.53
0.23 0.33
Causal Chamber (n = 10000, d = 20) Precision ↑ Recall ↑ F1 ↑ nSHD ↓
0.76 0.85 0.93 0.75 0.36 0.85 0.69
0.05 0.05
0.76 0.40
0.50 0.15
0.40 0.47
1.00 0.04
0.05 0.05 0.05 0.05 Out of memory Out of memory Out of memory Out of memory Out of memory Out of memory 0.62 0.55 0.05 0.08 Out of memory 0.03 0.05 0.03 0.03
0.71 0.52 0.83
Out of memory Out of memory 0.51 0.60 0.59 0.55 0.38 0.53
0.44 0.36 0.25
0.40 0.41
0.19 0.25 0.20 0.32
0.14 0.19 0.14
As shown in Table 3, DAG-FM demonstrates superior scalability and accuracy. Unlike some baselines that struggle with the dimensionality or sample sizes of these datasets, DAG-FM achieves state-of-the-art results, reinforcing its robustness for practical causal discovery applications.
6
Conclusion
In this work, we presented DAG-FM, a highly adaptable and scalable foundation model for causal discovery from observational tabular data. By moving away from direct, monolithic adjacency matrix prediction, DAG-FM reframes causal DAG identification as a sequential prediction process. It employs two amortized Transformer sub-modules to iteratively identify leaf nodes and their respective parents, ensuring the structural validity of the generated causal graph. A key contribution of our architecture is the integration of the MoLE framework, which effectively bridges the gap between theoretical FCM assumptions and practical data by dynamically routing predictions across mixed, diverse, heterogeneous causal mechanisms. Extensive empirical evaluations confirm the superiority of DAG-FM. On synthetic benchmarks featuring mixed identifiable mechanisms, DAG-FM significantly outperforms existing FCM-based, score-based, optimization-based, and amortized baselines in both DAG recovery and causal order identification. Furthermore, its state-of-the-art performance on two real-world benchmarks under9
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
scores its robust generalizability and capacity to handle high-dimensional, large-sample settings where traditional algorithms often struggle. Future Work While DAG-FM demonstrates formidable capabilities in causal discovery under the assumption of causal sufficiency, extending this framework to handle unobserved latent confounders and missing data presents an exciting avenue for future research. Additionally, scaling the foundation model with broader and more diverse pre-training corpora could further enhance its zero-shot inference capabilities across interdisciplinary scientific domains.
References Kevin Bello, Bryon Aragam, and Pradeep Ravikumar. Dagma: Learning dags via m-matrices and a log-determinant acyclicity characterization. Advances in Neural Information Processing Systems, 35:8226–8239, 2022. Patrick Blöbaum, Krishnakumar Balasubramanian, and Shiva Prasad Kasiviswanathan. Foundcause: Causal discovery with latent confounders from observational data, 2026. Peter Bühlmann, Jonas Peters, and Jan Ernest. Cam: Causal additive models, high-dimensional order search and penalized regression. The Annals of Statistics, 42(6):2526–2556, 2014. ISSN 00905364. Yikang Chen, Xingzhe Sun, and Dehui du. Causal discovery via quantile partial effect. In The Fourteenth International Conference on Learning Representations, 2026. David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 3(Nov):507–554, 2002. Gregory F Cooper and Edward Herskovits. A bayesian method for the induction of probabilistic networks from data. Mach. Learn., 9(4):309–347, October 1992. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings 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. Anish Dhir, Matthew Ashman, James Requeima, and Mark van der Wilk. A meta-learning approach to bayesian causal discovery. In The Thirteenth International Conference on Learning Representations, 2025. Bao Duong and Thin Nguyen. Bivariate causal discovery via conditional divergence. In Bernhard Schölkopf, Caroline Uhler, and Kun Zhang (eds.), Proceedings of the First Conference on Causal Learning and Reasoning, volume 177 of Proceedings of Machine Learning Research, pp. 236– 252. PMLR, 11–13 Apr 2022. Juan L Gamella, Jonas Peters, and Peter Bühlmann. Causal chambers as a real-world physical testbed for AI methodology. Nat. Mach. Intell., 7(1):107–118, January 2025. Noah Hollmann, Samuel Müller, Lennart Purucker, Arjun Krishnakumar, Max Körfer, Shi Bin Hoo, Robin Tibor Schirrmeister, and Frank Hutter. Accurate predictions on small data with a tabular foundation model. Nature, 637(8045):319–326, January 2025. Patrik Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou (eds.), Advances in Neural Information Processing Systems, volume 21. Curran Associates, Inc., 2008. Martin Huber. An introduction to causal discovery. Schweiz. Z. Volkswirtsch. Stat., 160(1), October 2024. 10
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Alexander Immer, Christoph Schultheiss, Julia E Vogt, Bernhard Schölkopf, Peter Bühlmann, and Alexander Marx. On the identifiability and estimation of causal location-scale noise models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 14316–14332. PMLR, 23–29 Jul 2023. Nan Rosemary Ke, Silvia Chiappa, Jane X Wang, Jorg Bornschein, Anirudh Goyal, Melanie Rey, Theophane Weber, Matthew Botvinick, Michael Curtis Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. In International Conference on Learning Representations, 2023. Jang-Hyun Kim, Claudia Skok Gibbs, Sangdoo Yun, Hyun Oh Song, and Kyunghyun Cho. Largescale targeted cause discovery via learning from simulated data. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, and Simon Lacoste-Julien. Gradient-based neural dag learning. In International Conference on Learning Representations, 2020. Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3744–3753. PMLR, 09–15 Jun 2019. Zi-Rong Li, Si-Yang Liu, Tian-Zuo Wang, and Han-Jia Ye. Tabcausal: Pretraining across causal environments for tabular causal discovery, 2026. Yingyu Lin, Yuxing Huang, Wenqin Liu, Haoran Deng, Ignavier Ng, Kun Zhang, Mingming Gong, Yian Ma, and Biwei Huang. A skewness-based criterion for addressing heteroscedastic noise in causal discovery. In The Thirteenth International Conference on Learning Representations, 2025. Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized inference for causal structure learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 13104–13118. Curran Associates, Inc., 2022. Takashi Nicholas Maeda and Shohei Shimizu. Rcd: Repetitive causal discovery of linear nongaussian acyclic models with latent confounders. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 735–745. PMLR, 26–28 Aug 2020. Francesco Montagna, Nicoletta Noceti, Lorenzo Rosasco, Kun Zhang, and Francesco Locatello. Causal discovery with score matching on additive models with arbitrary noise. In Mihaela van der Schaar, Cheng Zhang, and Dominik Janzing (eds.), Proceedings of the Second Conference on Causal Learning and Reasoning, volume 213 of Proceedings of Machine Learning Research, pp. 726–751. PMLR, 11–14 Apr 2023a. Francesco Montagna, Nicoletta Noceti, Lorenzo Rosasco, Kun Zhang, and Francesco Locatello. Scalable causal discovery with score matching. In Mihaela van der Schaar, Cheng Zhang, and Dominik Janzing (eds.), Proceedings of the Second Conference on Causal Learning and Reasoning, volume 213 of Proceedings of Machine Learning Research, pp. 752–771. PMLR, 11–14 Apr 2023b. Francesco Montagna, Max Cairney-Leeming, Dhanya Sridhar, and Francesco Locatello. Demystifying amortized causal discovery with transformers, 2025. Ignavier Ng, AmirEmad Ghassami, and Kun Zhang. On the role of sparsity and dag constraints for learning linear dags. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17943–17954. Curran Associates, Inc., 2020. 11
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Weronika Ormaniec, Scott Sussex, Lars Lorch, Bernhard Schölkopf, and Andreas Krause. Standardizing structural causal models. In The Thirteenth International Conference on Learning Representations, 2025. Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, USA, 2nd edition, 2009. ISBN 052189560X. Bo Peng, Sirui Chen, Jiaguo Tian, Yu Qiao, and Chaochao Lu. Causcale: Neural causal discovery at scale. In Forty-third International Conference on Machine Learning, 2026. Jonas Peters, Joris M. Mooij, Dominik Janzing, and Bernhard Schölkopf. Causal discovery with continuous additive noise models. Journal of Machine Learning Research, 15(58):2009–2053, 2014. Jingang Qu, David Holzmüller, Gaël Varoquaux, and Marine Le Morvan. Tabiclv2: A better, faster, scalable, and open tabular foundation model, 2026. Alexander Reisach, Christof Seiler, and Sebastian Weichwald. Beware of the simulated dag! causal discovery benchmarks may be easy to game. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 27772–27784. Curran Associates, Inc., 2021. Alexander Reisach, Myriam Tami, Christof Seiler, Antoine Chambaz, and Sebastian Weichwald. A scale-invariant sorting criterion to find a causal order in additive noise models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 785–807. Curran Associates, Inc., 2023. Paul Rolland, Volkan Cevher, Matthäus Kleindessner, Chris Russell, Dominik Janzing, Bernhard Schölkopf, and Francesco Locatello. Score matching enables causal discovery of nonlinear additive noise models. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 18741–18753. PMLR, 17–23 Jul 2022. Karen Sachs, Omar Perez, Dana Pe’er, Douglas A. Lauffenburger, and Garry P. Nolan. Causal protein-signaling networks derived from multiparameter single-cell data. Science, 308(5721): 523–529, 2005. doi: 10.1126/science.1105809. Shohei Shimizu, Patrik O. Hoyer, Aapo Hyvärinen, and Antti Kerminen. A linear non-gaussian acyclic model for causal discovery. Journal of Machine Learning Research, 7(72):2003–2030, 2006. Shohei Shimizu, Takanori Inazumi, Yasuhiro Sogawa, Aapo Hyvärinen, Yoshinobu Kawahara, Takashi Washio, Patrik O. Hoyer, and Kenneth Bollen. Directlingam: A direct method for learning a linear non-gaussian structural equation model. Journal of Machine Learning Research, 12 (33):1225–1248, 2011. Peter Spirtes. Directed cyclic graphical representations of feedback models. In Proceedings of the Eleventh Conference on Uncertainty in Artificial Intelligence, UAI’95, pp. 491–498, San Francisco, CA, USA, 1995. Morgan Kaufmann Publishers Inc. ISBN 1558603859. Peter Spirtes and Clark Glymour. An algorithm for fast recovery of sparse causal graphs. Social science computer review, 9(1):62–72, 1991. Eric V. Strobl and Thomas A. Lasko. Identifying patient-specific root causes with the heteroscedastic noise model. Journal of Computational Science, 72:102099, 2023. ISSN 1877-7503. Natasa Tagasovska, Valérie Chavez-Demoulin, and Thibault Vatter. Distinguishing cause from effect using quantiles: Bivariate quantile causal discovery. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9311–9323. PMLR, 13–18 Jul 2020. Ryan Thompson, He Zhao, Daniel M. Steinberg, and Edwin V. Bonilla. Arrow: A foundation model for causal discovery, 2026. 12
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
Jan P Vandenbroucke, Alex Broadbent, and Neil Pearce. Causality and causal inference in epidemiology: the need for a pluralistic approach. International Journal of Epidemiology, 45(6): 1776–1786, 12 2016. ISSN 0300-5771. doi: 10.1093/ije/dyv341. 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. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. Matej Vuković and Stefan Thalmann. Causal discovery in manufacturing: A structured literature review. Journal of Manufacturing and Materials Processing, 6(1), 2022. ISSN 2504-4494. doi: 10.3390/jmmp6010010. Johnny Xi, Hugh Dance, Peter Orbanz, and Benjamin Bloem-Reddy. Distinguishing cause from effect with causal velocity models. In Forty-second International Conference on Machine Learning, 2025. Sascha Xu, Sarah Mameche, and Jilles Vreeken. Learning causal orderings for in-context tabular prediction, 2026. Yue Yu, Jie Chen, Tian Gao, and Mo Yu. DAG-GNN: DAG structure learning with graph neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7154–7163. PMLR, 09–15 Jun 2019. Bin Zhang, Chris Gaiteri, Liviu-Gabriel Bodea, Zhi Wang, Joshua McElwee, Alexei A. Podtelezhnikov, Chunsheng Zhang, Tao Xie, Linh Tran, Radu Dobrin, Eugene Fluder, Bruce Clurman, Stacey Melquist, Manikandan Narayanan, Christine Suver, Hardik Shah, Milind Mahajan, Tammy Gillis, Jayalakshmi Mysore, Marcy E. MacDonald, John R. Lamb, David A. Bennett, Cliona Molony, David J. Stone, Vilmundur Gudnason, Amanda J. Myers, Eric E. Schadt, Harald Neumann, Jun Zhu, and Valur Emilsson. Integrated systems approach identifies genetic nodes and networks in late-onset alzheimer’s disease. Cell, 153(3):707–720, 2013. 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. In The 28th International Conference on Artificial Intelligence and Statistics, 2025a. Kun Zhang and Aapo Hyvärinen. On the identifiability of the post-nonlinear causal model. In Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence, UAI ’09, pp. 647–655, Arlington, Virginia, USA, 2009. AUAI Press. ISBN 9780974903958. Xingxuan Zhang, Gang Ren, Han Yu, Hao Yuan, Hui Wang, Jiansheng Li, Jiayun Wu, Lang Mo, Li Mao, Mingchao Hao, Ningbo Dai, Renzhe Xu, Shuyang Li, Tianyang Zhang, Yue He, Yuanrui Wang, Yunjia Zhang, Zijing Xu, Dongzhe Li, Fang Gao, Hao Zou, Jiandong Liu, Jiashuo Liu, Jiawei Xu, Kaijie Cheng, Kehan Li, Linjun Zhou, Qing Li, Shaohua Fan, Xiaoyu Lin, Xinyan Han, Xuanyue Li, Yan Lu, Yuan Xue, Yuanyuan Jiang, Zimu Wang, Zhenlei Wang, and Peng Cui. Limix: Unleashing structured-data modeling capability for generalist intelligence, 2025b. Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. Dags with no tears: Continuous optimization for structure learning. Advances in neural information processing systems, 31, 2018.
13
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
A
Proofs
Theorem 4.4. Under Assumption 4.3, if Xi = fi (Xpa(i) , U ) entails the observational distribution p(Xpa(i) , Xi ) and fi ∈ Fi , then the posterior probability P (Fi | p(Xpa(i) , Xi )) = 1. Proof. According to the analysis in Chen et al. (2026), for each mechanism family, the gradient of the quantile function of the conditional distribution p(Xi | Xpa(i) ) with respect to its parents Xpa(i) can be represented by a uniquely determined set of basis functions. Consequently, different mechanism families exhibit distinct distributional signatures. This unique mathematical formulation guarantees that the underlying causal mechanism family can be fully distinguished and identified solely from the observational distribution p(X). Theorem 4.5. Under Assumption 4.3, if an SCM ϕ = (X, f , p(U )) entails the observational distribution p(X) and the causal graph G, then the posterior probability P (G | p(X)) = 1. Proof. Let I = {1, . . . , d} denote the index set of the d variables. Since the DAG space is discrete, P (G | p(X)) = 1 is mathematically equivalent to showing that the posterior probability of any alternative graph is zero, i.e., ∀G ′ ̸= G, P (G ′ | p(X)) = 0. Without loss of generality, since reordering the indices does not alter the rigorousness of the proof, we assume that the index sequence 1, . . . , d already constitutes a valid topological ordering of the true DAG G. Consider an alternative DAG G ′ = (I, E ′ ) ̸= G. There must exist at least one pair of nodes (i, j) that falls into one of the following three categories regarding structural deviation: (i) Spurious forward edge: i < j, with (i, j) ∈ E ′ but (i, j) ∈ / E; (ii) Missing forward edge: i < j, with (i, j) ∈ E but (i, j) ∈ / E ′ ; (iii) Reversed edge: i > j, with (i, j) ∈ E ′ (i.e., an edge violating the true topological order). We proceed by recursively checking node pairs in reverse topological order, starting from the leaf node d down to 1. Suppose a structural discrepancy exists involving the leaf node d and some node i < d. Under the assumptions of causal sufficiency and faithfulness, the conditional independencies (CIs) in p(X) correspond strictly to d-separation in the graph. Thus, the Markov blanket of Xd uniquely identifies its true parents in G. We can derive a contradiction for each case by confronting the proposed G ′ with the observational distribution p(X): • Spurious forward edge ((i, d) ∈ E ′ but (i, d) ∈ / E): By causal sufficiency, for any SCM satisfying the true graph G where (i, d) ∈ / E, there exists a conditioning set XI ′ (I ′ ⊂ I) such that the conditional independence Xi ⊥⊥ Xd | XI ′ holds in p(X). However, if an alternative SCM posits G ′ where (i, d) ∈ E ′ , this specific d-separation would not exist in G ′ . Due to the faithfulness assumption, such a model cannot truthfully generate p(X). This contradiction invalidates the hypothesis of G ′ . • Missing forward edge ((i, d) ∈ E but (i, d) ∈ / E ′ ): Since Xi is an actual parent of Xd , causal sufficiency and faithfulness guarantee that there exists no conditioning set XI ′ capable of rendering Xi ⊥ ⊥ Xd | XI ′ in the true distribution p(X). If G ′ omits this edge, it intrinsically implies that such a d-separation (and corresponding CI) should exist. The absence of this CI in p(X) contradicts the assumption of G ′ . • Reversed edge ((d, i) ∈ E ′ where i < d): If G ′ contains a reversed edge (d, i), since we proved missing forward edges from true parents cannot exist, it indicates a scenario where the local joint distribution of Xi and Xd needs to be explained by a backward mechanism. Specifically, the distribution natively generated by Xi = fi (Xpa(i) , Ui ) and Xd = fd (Xpa(d) , Ud ) must be reversely generatable. However, according to the well-established identifiability theories of FCMs (Shimizu et al., 2006; Hoyer et al., 2008; Zhang & Hyvärinen, 2009; Immer et al., 2023), if a distribution is induced by a specific identifiable causal mechanism in the forward direction, there does not exist an alternative mechanism within the same valid function class that can perfectly reconstruct the distribution in the reverse direction while ensuring the independence of the noise variables. This fundamental asymmetry introduces a definitive contradiction, thus ruling out G ′ . By applying this logic recursively from the leaf node d up to the root nodes, we dynamically eliminate any structural deviations. Therefore, any G ′ ̸= G inevitably yields a contradiction 14
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
with the observational distribution p(X), proving that P (G ′ | p(X)) = 0, and consequently, P (G | p(X)) = 1.
B
Additional Results
The following sections present additional experiments to demonstrate the scalability of DAG-FM, its robustness under various settings and perturbations, and its generalization capabilities on out-ofdistribution (OOD) data. We compare DAG-FM with other amortized methods using their publicly available checkpoints. All inference experiments were conducted under a 24GB VRAM constraint, and out-of-memory cases are not reported. Except for the specific factors being investigated, the synthetic benchmark configurations remain identical to those in the main text. Additionally, since the other methods cannot output causal orders, the order identification performance is only reported for DAG-FM. B.1
Scaling Capabilities
AVICI
CauScale
TabCausal
0.7
nSHD ( )
0.4
5K 10K 20K 50K
0.250 0.225 0.200
100 200 500 1K
2K
0.150
5K 10K 20K 50K
Number of Samples
0.30 0.25 0.20
0.175
0.3 2K
0.35
0.275
0.5
Number of Samples
DAG-FM
0.300
0.6
100 200 500 1K
FoundCause
0.325
ODR ( )
0.75 0.70 0.65 0.60 0.55 0.50 0.45 0.40 0.35
F1 ( )
Precision ( )
Scaling to Large Samples Figure 1 illustrates the performance of different amortized methods as the sample size increases. Without additional optimization, both AVICI and CauScale encounter OOM errors under large sample sizes. Meanwhile, DAG-FM outperforms TabCausal and FoundCause as the sample size grows.
100 200 500 1K
2K
5K 10K 20K 50K
100 200 500 1K
Number of Samples
2K
5K 10K 20K 50K
Number of Samples
Figure 2: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under different numbers of samples.
Scaling to High Dimensions Figure 2 shows the performance of different amortized methods as the number of features increases. Notably, AVICI, CauScale and FoundCause encounter out-ofmemory issues in high-dimensional settings. In contrast, DAG-FM outperforms TabCausal in terms of precision and nSHD across all evaluated dimensions, with the exception of the 500-dimension case.
0.6
0.6
0.4
0.2
0.0
0.0 5
10
20
50
100 200 500
Number of Features
FoundCause
DAG-FM 0.35
0.8
0.4
0.2
2
TabCausal 1.0
0.30
ODR ( )
0.8
CauScale
nSHD ( )
0.8
F1 ( )
Precision ( )
AVICI
0.6 0.4
0.25 0.20 0.15 0.10
0.2
0.05 2
5
10
20
50
100 200 500
Number of Features
2
5
10
20
50
100 200 500
Number of Features
2
5
10
20
50
100 200 500
Number of Features
Figure 3: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under different numbers of features.
B.2
Robustness
Robustness to Sparse/Dense Graphs Real-world data may contain both sparse and dense graphs. Figure 4 shows the performance of different amortized methods as the graph density 2|E|/d(d − 1) increases, where |E| denotes the number of edges. DAG-FM significantly outperforms the other methods on sparse graphs, while performing second only to FoundCause on dense graphs. 15
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
0.75
0.6
0.65 0.60
0.5 0.4
0.55
0.3
0.50
0.2
0.45
0.05
0.1
0.2
0.3
0.5
DAG Densities
0.75
TabCausal
nSHD ( )
0.70
CauScale
1.0
0.05
0.1
0.2
0.3
0.5
DAG Densities
0.75
FoundCause
DAG-FM
0.8 0.7 0.6 0.5 0.4 0.3 0.2 0.1
1.0
0.21 0.20
ODR ( )
0.7
F1 ( )
Precision ( )
AVICI 0.80
0.19 0.18 0.17 0.16 0.15
0.05
0.1
0.2
0.3
0.5
DAG Densities
0.75
1.0
0.05
0.1
0.2
0.3
0.5
DAG Densities
0.75
1.0
Figure 4: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under different densities.
Robustness to Data Warping Due to complex uncertainties, empirically collected real-world data may contain unknown warpings that affect the alignment with causal discovery assumptions. Figure 5 presents the performance of different amortized methods as the proportion of warped nodes increases. The proportion of warping refers to the ratio of randomly warped variables to the total number of variables. DAG-FM outperforms the other methods, indicating its robustness to data warping.
0.60
0.625
0.55
0.600 0.575
0.50 0.45
0.550
0.40
0.525 0.500
TabCausal 0.26 0.25 0.24 0.23 0.22 0.21 0.20 0.19 0.18
0.35 0.05
0.1
0.2
0.3
0.5
Proportions of Warping
0.75
1.0
FoundCause
DAG-FM 0.24 0.22
0.05
0.1
0.2
0.3
0.5
Proportions of Warping
0.75
1.0
ODR ( )
0.650
CauScale
nSHD ( )
0.65
F1 ( )
Precision ( )
AVICI 0.675
0.20 0.18 0.16
0.05
0.1
0.2
0.3
0.5
Proportions of Warping
0.75
1.0
0.05
0.1
0.2
0.3
0.5
Proportions of Warping
0.75
1.0
Figure 5: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under different proportions of warping.
Robustness to Discretization Another challenge in real-world scenarios is mixed-type data, where features may include both discrete and continuous variables. Figure 6 displays the performance of different amortized methods as the number of discretized variables increases. The proportion of discretization indicates the ratio of discrete variables to the total number of variables. DAG-FM outperforms the other methods and is only marginally surpassed when nearly all variables are discrete, demonstrating its robustness to discretization. TabCausal
nSHD ( )
0.55 0.50 0.45 0.40 0.35 0.30 0.05
0.1
0.2
0.3
0.5
0.75
Proportions of Discretization
1.0
0.05
0.1
0.2
0.3
0.5
0.75
Proportions of Discretization
1.0
0.26 0.25 0.24 0.23 0.22 0.21 0.20 0.19 0.18
FoundCause
DAG-FM 0.26 0.24
ODR ( )
CauScale
0.60
F1 ( )
Precision ( )
AVICI 0.65
0.66 0.64 0.62 0.60 0.58 0.56 0.54 0.52 0.50
0.22 0.20 0.18
0.05
0.1
0.2
0.3
0.5
0.75
Proportions of Discretization
1.0
0.05
0.1
0.2
0.3
0.5
0.75
Proportions of Discretization
1.0
Figure 6: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under different proportions of discretization.
B.3
Generalization to Out-of-Distribution Data
Although amortized causal discovery algorithms enable inference on unseen data, it remains crucial to examine the variation in model performance when encountering data drawn from a distribution outside the prior training space. Figure 7 compares the performance of various amortized methods across four types of OOD scenarios, which include: 16
DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms
• DAG: The causal graphs of the test data are sampled from the SBM (Stochastic Block Model), a random graph model never used during training. • Mech.: The causal mechanisms in the underlying SCM of the test data follow a post-nonlinear heteroscedastic noise model (PNL-HNM), taking the form Y = h(f (x) + g(x) U ), where h is a non-linear invertible function. This specific mechanism is absent from the training data. Although the theoretical identifiability of such a mechanism is yet to be strictly proven, empirical evidence suggests that the model can successfully identify the corresponding DAGs. • Noise: The noise terms in the underlying SCM of the test data are drawn from half-normal and Pareto distributions, both of which are unseen during training. • Hetero.: The PNL-HNM is mixed with other causal mechanisms in the test data to evaluate the models’ performance under OOD heterogeneous mechanisms.
0.4
0.10
0.15 0.10 0.05 0.00
DAG
Noise
DAG
Mech.
Types of Out-of-Distritbution
Hetero.
Types of Out-of-Distritbution
Hetero.
0.00
Noise
0.0 Mech.
0.1
0.0 DAG
0.1
0.05
Types of Out-of-Distritbution
Hetero.
0.2
0.2
0.15
Noise
0.3
0.3
Types of Out-of-Distritbution
Hetero.
0.4
0.20
Noise
0.20
DAG-FM
Mech.
0.5
FoundCause
DAG
0.5
nSHD ( )
0.25
ODR ( )
TabCausal
0.6
Mech.
CauScale
0.6
F1 ( )
Precision ( )
AVICI
0.7
0.7
Figure 7: Performance of DAG identification (Precision, F1, nSHD) and order identification (ODR) under 4 types of out-of-distribution cases. The results indicate that DAG-FM outperforms the other methods in almost all OOD scenarios, demonstrating its strong generalization capabilities.
17