Sample, estimate, aggregate: A recipe for causal discovery foundation models
Authors: Menghua Wu
Yujia Bao
Regina Barzilay
Tommi Jaakkola
Venue: Accepted by TMLR ()

Abstract: Causal discovery, the task of inferring causal structure from data, has the potential to uncover mechanistic insights from biological experiments, especially those involving perturbations. However, causal discovery algorithms over larger sets of variables tend to be brittle against misspecification or when data are limited. For example, single-cell transcriptomics measures thousands of genes, but the nature of their relationships is not known, and there may be as few as tens of cells per intervention setting. To mitigate these challenges, we propose a foundation model-inspired approach: a supervised model trained on large-scale, synthetic data to predict causal graphs from summary statistics — like the outputs of classical causal discovery algorithms run over subsets of variables and other statistical hints like inverse covariance. Our approach is enabled by the observation that typical errors in the outputs of a discovery algorithm remain comparable across datasets. Theoretically, we show that the model architecture is well-specified, in the sense that it can recover a causal graph consistent with graphs over subsets. Empirically, we train the model to be robust to misspecification and distribution shift using diverse datasets. Experiments on biological and synthetic data confirm that this model generalizes well beyond its training set, runs on graphs with hundreds of variables in seconds, and can be easily adapted to different underlying data assumptions.

Published in Transactions on Machine Learning Research (03/2025)
Sample, estimate, aggregate:
A recipe for causal discovery foundation models
Menghua Wu
rmuuus at mit edu
Deportment of Computer Science, Massachusetts Institute of Technology
Yujia Bao
guffa. boo ( Jaccnntureccomm
Center for Advanced AI, Accenture
Regina Barzilay
regina reginafat fastffaallccaaaooooo mät edu
Deportment of Computer Science, Massachusetts Institute of Technology
Tommi S. Jaakkola
tommifat mit edu
Department of Computer Science, Massachuselts Institute of Technology
Reviewed on OpenReview: https: // opwrrevee.....eee. net/ forum? id- h4342055X0
Abstract
Causal discovery, the task of inferring causal structure from data, has the potential uncover mechanistic insights from biological experiments, especially those involving perturbations. However, causal discovery algorithms over larger sets of variables tend to be brittle against misspecification when data are limited. For example, single-cell transcriptomics measures thousands of genes, but the nature of their relationships is not known, and there may be as few as tens cells per intervention setting. To mitigate these challenges, we propose a foundation model-ins inseiiedd approach: a supervised model trained on large-scale, synthetic data to prodict causal graphs from summary statistics like the outputs of classical causal discovery algorithms run over subsets of variables and other statistical hints like inverse covariance. Our approach is enabled by the observation that typical errors in the outputs of discovery algorithm remain comparable across datasets. Theoretically, we show that the model architecture is well-specifid,, in the sense that it can recover causal graph consistent with graphs over subsets. Empirically, we train the model to be robust. to misspecification and distribution shift using diverse datasets. Experiments on biological and synthetic data confirm that this model generalizes well beyond its training set, runs on graphs with hundreds of variables in seconds, and can be easily adapted different underlying data assumptions)
1 Introduction
A fundamental aspect of scientific research is to discover and validate causal hypotheses involving variables of interest. Given observations of these variables, the goal of causal discovery algorithms is to extract such hypotheses in the form of directed graphs, in which edges denote causal relationships (Spirtes et al.) 2001). There are several challenges to their widespread adoption in basic science. The core issuc is that the correctness of these algorithıs is tied to their assumptions on the data-generating processes, which are unknown in real applications. In principle, one could circumvent this issue by exhaustively running discovery algorithms with different assumptions and comparing their outputs with surrogate measures that reflect graph quality (Faller et al. 2023). However, this scarch would be costly: current algorithms must be optimized from scratch each lime, and they scale poorly to the graph and dataset sizes present in modern scientific big
data (Replogle et al. 2022
'Qur code available https: //
https://githhb. /gthgb.t.hb......g com/rmvu/sen

---

Published in Transactions on Machine Learning Research (03/2025)
Causal discovery algorithms follow two primary approaches that differ in their treatment of the causal graph. Discrete search algorithms explore the super-exponential space of graphs by proposing and evaluating changes to working graph (Glymour et al. 2019). While these methods are fast on small graphs, the combinatorial space renders them intractable for exploring larger structures. An alternative is to frame causal discovery as a continuous optimization over weighled adjacency matrices. These algorithms either fit a generative model to the data and extract the causal graph as a parameter (Zheng et al., 2018), or train a supervised learning model on simulated data (Petersen et al. 2023). However, these methods may be less robust beyond simple settings, and their optimization can be nontrivial (gg et al. 2024).
In this work, we present SEA: Sample, Estimate, Aggregate, supervised causal discovery framework that aims to perform well even when data-generating processes are unknown, and to easily incorporate prior knowledge when it is available. We train a deep learning model to predict causal graphs from two types of statistical descriptors: the estimates of classical discovery algorithms over small subsets, and graph-level statistics. Each classical discovery algorithm outputs a representation of a graph's equivalency class, and the types of errors that it makes may be comparable across datasets. On the other hand, statistics like correlation or inverse covariance are fast to compute, and strong indicators for a graph's overall connectivity, Thcorctically, we illustrate a simple algorithm for recovering larger causal graphs that are consistent with estimates over subsets, and we show that there exists a set of model parameters that can map sets of subgraph estimates to the correct global graph. Empirically, our training procedure forces the model to predict causal graphs across diverse synthetic data, including on datasets that are misaligned with the discovery algorithms' assumptions, or when insufficient subsets are provided.
Our experiments probe three qualities that view foundation model should fulfill, with thorough comparison to three classical baselines and five deep learning approaches. Specifically, we assess the framework's 's ability to generalize to unseen and to perform well in low-data regimes. SEA attains the state-of-the-art results on synthetic and real causal discovery tasks, while providing 10-1000x faster inference. To incorporate domain knowledge, we show that it is possible to swap classic discovery algorithuns at inference time, for significant improvements on datasets that match the assumptions of the new algorithm. Our models can also be finetuned at fraction of the training cost to accommodate new graph-level statistics that capture different (c.g. nonlincar) relationships. We extensively analyze SEA in terms of low-data performance, scalability, causal identifiability, and other design choices. To conclude, while our experiments focus on specific algorithms and architectures, this work presents& a blueprint for designing causal discovery foundation models, in which sampling heuristics, classical causal discovery algorithms, and summary statistics are the fundamental building blocks.
data; to steer predictions based on prior knowledge; and
2 Background and related work
2.1 Causal structure learning
A causal graphical model is directed graph G -V EV, where each node i E V corresponds to random variable X, E X and each edge (i.5) E E represents a causal relationship from X, + X,. There are number of assumptions that relate data distribution Px to G (Spirtes et al., 2001 Hauser & Bühlmann 2012), which determine whether G is uniquely identifiable, or identifiable up to an equivalence class. In this work, we assume causal sufficiency - that is, V contains all the parents Ty of every node i. Causal graphical models can be used, along with other information, to compute the downstream consequences of interventions. An intervention on node refers to setting conditional P(X; Xx,) to a different distribution P(X, Xx,). Our experiments cover the observational case (no interventions) and the case with perfect interventions on each node, i.e. for all nodes we have access to data where we set P(X, Xx,) + P(X.).
Given a datasct D ~ Px, the goal of causal structure lcarning (causal discovery) is to recover G. There are two main challenges. First, the number of possible graphs is super-exponential in the number of nodes N, so algorithms must navigate this combinatorial search space efficiently. Second, depending on data availability and the underlying data generation process, causal discovery algorithms may or may not be able to recover G in practice. In fact, many algorithms are only analyzed in the infinito-data regime and require at least thousands of data samples for reasonable empirical performance (Spirtes et al. 2001: Brouillard et al. 2020
2

---

Published in Transactions on Machine Learning Research (03/2025)
(A) Design objectives causal discovery foundation model
(B) Our approach: Sample, Estimate, Aggregate (SEA)
Inference framework:
Featurize datasets via summary statistics and outputs classical discovery
1) Outperform dataset specific algorithms, even on real, out-of disttribiit,,,,,,n or misspecified data
algoriihms small subsets, whose errors compurable across datasets Pretrained, supervised model prodicts causal graph.
batches D.
D2 marginal estimates E,
aligned estimates E dig
dataset with unknown assumptions
nodes
FM
ocamples
TXK
NXN
high quality
graph prediction
b examples
batch
global statisties
Pretmined
2) Incorporate domain knowledge at
inference time, or low-cost finetuning
datasct D
NXN
Aggregator Model
graph prediction É
Input
avamppioss
Output
datesct
"non-linear"
"prompe"
FM
Sample subsets of data and variables
Estimate local siructures and global summary statistics
features into causal graphs
Aggregate statistical
"non Goussian"
prediction incorporates assumptions
(C) Training procedure
Update Aggregator
Sub-sample nodes and batches with varying T
3) Scaling up robust low-data performance
Compute statistics, marginal estimates
Aggregator
small daleset
high quality prediction
Large-scale, synthetic datasets
FM
graph predictions
ground truth (synthetic)
Figure Overview of our goals and approach. (A) Criteria we aim to fulfill. (B-C) Inference and training procedure. Green: raw data. Blue: graph / features. Yellow: Learned. Gray: Stochastic, but not learned.
Discrete search algorithms encompass diverse strategies for traversing the combinatorial space of possible graphs. Constraint-based algorithms are based on condit ional independence tests, whose discrete results inform of the presence or absence of edges, and whose statistical power depends directly on dataset size (Glymour et al, 2019 These include the observational FCI and PC algorithms (Spirtes al. 1995), and the interventional JCI algorithm MMooij et al., 2020). Score-based methods define continuous score that guides the search through the discrete space of valid graphs, where the true graph lies at the optimum. Examples include GES (Chickering, 2002)), GIES (Hauser & Bühlmann 2012), CAM (Bühlmann et al. 2014). and IGSP (Wang 2017p. Finally, cemi-parametric methods such as LiNGAM (Shimizu et al. 2006) or additive noise and causal ordering.
et al. models [Hoyer et al. 2008) exploit asymmetries implied by the model class to identify graph connectivity
Continuous optimization approaches recast the combinatorial space of graphs into continuous space of weighted adjacency matrices. Many of these works train generative model to learn the empirical data distribution, which is parameterized through the adjacency matrix (Zheng et al. 2018 Lachapelle et al. 2020 Brouillard et al. 2020). Others focus on properties related to the empirical data distribution, such as a relationship between the underlying graph and the Jacobian of the learned model (Reizinger et al. 2023), or between the Hcssian of the data log-likelihood and the topological order (Sanchez ct al. 2023
Finally, amortized inference approaches (Ke et al. 2023 Lorch et al. 2022 Petersen et al.] 2023 Trame causal discovery as supervised learning problem, where a neural net work is trained predict (synthetic) graphs from (synthetic) datasets. To incorporate new information, current supervised methods must simulate datasets and re-train. In addition, these models that operate datasets. Since this direction is most similar to our own, we include further comparisons in B.2
observations scale poorly to larger
2.2 Foundation models
The concept of foundation models has revolutionized the machine learning workflow in variety of disci- plines: instead of training domain-specific models from scratch, start from pretrained, general-purpose model (Bommasani et al. 2021). This work describes blueprint for designing "foundation models" in the context of causal discovery. The precise definition of a foundation model varies by application, but we aim to fulfill the following properties (Figure [A), enjoyed by modern text and image foundation models (Radford et al. 2021 Brown et al. 2020
foundation model should enable us to outperform domain-specific models trained from scratch, even if the former has never seen similar tasks during training (Radford et. al.] 2019). In the context.
3

---

Published in Transactions on Machine Learning Research (03/2025)
of causal discovery, would like to train model that outperforms any individual algorithm real, misspecified, and/or out-of-distribution datasets.
2. It should be possible to explicitly steer a foundation model's behavior towards better performance on new tasks, either directly at inference time, c.g. "prompting" (Reynolds McDonell 2021 or at low cost compared to pretraining (Ouyang et al. 2022). Here, we would like to easily change our causal discovery algorithm's 's "assumptions" regarding the data, e.g- by incorporating the knowledge of non-linearity, non-Gaussianity.
3. Scaling up foundation model should lead to improved performance in few-shot or data-poor regimes (Brwnn et al. 2020). This aspect we analyze empirically.
In the following sections, we will revisit these desiderata from both the design and experimental perspectives.
3 Methods
Sample, Estimate, Aggregate (SEA) supervised causal discovery framework built upon the intuition that summary statistics and marginal estimates (outputs of classical causal discovery algorithm subsets of nodes) provide useful hints towards global causal structure. In the following sections, we describe how these statistics are efficiently estimated from data (Section 3.1), and how train neural network to predict causal graphs from these inputs (Section 3.2p. We expand upon the model architecture in Section 3.3 and conclude with thcoretical motivation for this architecture in Section 3.4
3.1 Inference procedure
Given new dataset D E RRN,, we sample small batches of nodes and observations; estimate global summary statistics and local subgraphs; and aggregate these information with trained neural network (Figure [
Sample: takes as input dataset D; and outputs data batches {Do, Di...., Dr} and node subsets {SI...., Sr}.
1. Sample T + batches < M obscrvations uniformly at random from D.
Initialize selection scores a NNN (e.g. correlation or inverse covariance, computed over Do or D).
Sample T node subsets of size k. Each subset S, is constructed iteratively as follows.
(a) The initial node sampled with probability proportional to ΣjEV Oig
(b) Each subsequent node is added with probability proportional to Ejes, aij (prioritizing connec- tions to nodes that have already been sampled), until || k.
(c) We update a, down-weighting an proportional to the number of times i,j have been selected.
We include further details and analyze alternative strategies for sampling nodes in B.6
Estimate: takes as inputs data batches, node subsets, and (optionally) intervention targets; and outputs global statisties p and marginal estimates [E E.
1. Compute global statistics p E RNNN over Do (e.g. correlation or inverse covariance).
2. Run discovery algorithm f to obtain marginal estimates f(D)(S.)) = E, for t t= I....
We use D.[S,] to denote the observations in D, that correspond only to the variables in S.. Each estimate E is a adjacency matrix, corresponding to the nodes in S..
A Aggregate: takes as inputs global statisties, marginal estimates, and node subsots. trained aggregator model outputs the predicted global causal graph E (0, N)N (Section 3.3
3.2 Training procedure
The training procedure mirrors the inference procedure (Figure IC). We assumee access to pairs of simulated datasets and graphs, (D, G), where each dataset D is generated by a parametric model, whose dependencies are given by graph G (Section 4.1. Datascts D are summarized into global statistics and marginal estimates via the sampling and estimation steps. The resultant features are input to the aggregator (neural network), which is supervised by graphs G.
4

---

Published in Transactions on Machine Learning Research (03/2025)
graph prediction È (0.1) XXN
FFN
sum
message
FPN
Dropout
FFN
p
Cal Artn
FFN
Axial attention black
Embed bye
Embod usee
Embeed puce
Row Atrn
LN
E
(i.
Drogout
Axial atteation block
Linear W
Col Alin
Marginal Estimate Embedding
concat,
Emnbedding
FFN Col Atta Row Atta
message
h
Dropout Row Attn LN
K
L
Embeed mste
Lincar W
marginal estimate
marginal estimate Ealien gTX global features 1g/D0N
embeddings h global feature
2D input
fi.13
Pgj
embeddings he.f-1
Model Overview
Global Feature Embedding
Axial Attention Block
Axial Attention Layer
Figure Aggregator architecture. Marginal graph estimates and global statistics are embedded into the model dimension. 1D positional embeddings are added along both rows and columns. Embedded features pass through series of axial attention blocks, which attend to the marginal and global features. Final layer global features pass through feedforward network to predict the causal graph.
We trained two aggregator models, which employed the Fci algorithm with the Fisherz test and GIES algorithm with the Bayesian information criterion (Schwarz) 1978). Both estimation algorithms were chosen for speed, but they differ in their assumptions, discovery strategies, and output formats. Though we only trained two models, alternate estimation algorithms that produce the same output type may be used at inference time (experiments in Section 5.2p. The training dataset contains both data that are correctly and incorrectly specified (Section 4.1 so the aggregator forced to predict the correct graph regardless. addition, each training instance samples a random number of marginal estimates, which might not cover every edge. As result, the aggregator must extrapolate to unseen edges using the available estimates and the global statistics. For example, two variables have low correlation, and they are located in different neighborhoods of the already-identified graph, it may be reasonable assume that they unconnected.
3.3 Model architecture
The aggregator is neural network that takes input: global statistics p NxN, marginal estimates EY... E xxxxk and node subsets St..r E [N]T*k (Figure 2h, where & is the set of output edge types for the causal discovery algorithm /
We project global statistics into the model dimension via a learned lincar projection matrix W, : + Rd and we embed edge types via a learned embedding ebdg : - Rdd To collect estimates of the same edge over all subsets, we align entries of E into E alaignn ETTx
[Eiss
Eisss
if i St otherwise
E aligu (e=((h))
(1)
where indexes into the subsets, e 1...A K indexes into the set of unique edges, i.e. the union of pairs (i,j) over all E/. We add learned 1D positional embeddings along both dimensions of each input,
pos-ebd(pi.j) ebdnode (i') ebdnode (i')
pos-ebd(E, (Eaaiggn) ebd (t) FFN((edd (E"), ebd
2PC (Spirtes et 2001 also prodicts CPDAGs Andersson et 1997 its outputs may be input to the CIES model, while FCI's PAG outpuls F.g. "no relationship." "X causes Y "X descendent Y

---

Published in Transactions on Machine Learning Research (03/2025)
where .' index into random permutation on V for invariance to node permutation and graph size D Due the (a)symmetries of their inputs, pos-ebd(pt.)) symmetric, while pos-chd (F,, Eig considers the node ordering. In summary, the inputs to our axial attention blocks are
h (W pP) pos-ebd(pt.))
(2)
h ebdg (F ( + pos-
(3)
for i,j [MP, t [T], E [K]. Note that attention is permutation invariant, so positional embeddings are required for the model to know which edges belong to the same subset, what each edge's endpoints endpoints are.
Axial attention An axial attention block contains two axial attention layers (marginal estimates, global statistics) and a ford-forward network (Figure 2 right). Given a 2D input, an axial attention layer attends first along the rows, then along the columns. For example, on a matrix of size (R, CR, one pass of the axial attention layer is equivalent running standard self-altention along C with batch size R, followed by the reverse. For marginal estimates, R is the number of subsets T, and C is the number of unique edges K. For global statistics, and C are both the total number of vertices N.
Following Rao et al. 22021) each self-antention mechanism is preceded by layer normalization and followed by dropout, with residual connections to the input,
x = x + Dropout (AttataayeaaNaanaaaaaa
(4)
We pass messages between the marginal and global layers to propagate information. Let ΦE, be marginal layer let &p.t be global layer L, and let h denote the hidden representations out of layer l. The marginal to global message m E-op NNxxx contains representations of each edge averaged over subsets,
1 Σ, (10))1))) E-p8_ m,
i,J. St
T. E
(5)
otherwise.
where Te is the number of St containing e, and missing entries are padded to learned constant The global to marginal message m R.xdd is simply the hidden representation itself,
B
oel K
(6)
We update representations based on these messages as follows.
E (b,-()
(marginal feature)
(7)
hoe--1 + 8º [hrie- .m B-p.le
(marginal to global)
(8)
hpt = @pp(th-t-1)
(global feature)
(9)
F. F.t. F.F
(global marginal)
(10)
we [aaddd is a learned linear projection, and [:] denotes concatenation.
Graph prediction For each pair of vertices i $ E V, we predict 0, 1, or 2 for no edge, i i j, and
joi. i. We do not additionally enforce that our predicted graphs are acyclic, similar in spirit to Lippe et al. Given the output of the final axial attention block h², we compute logits
2022
[(.)} = FFN (GAl) E R
(11)
which correspond to probabilities after softmax normalization. The overall output É [0,1) is supervisod by the ground truth E. Our model is trained with cross entropy loss and L2 regularization.
"The random permutation / === aVV); allows to avoid updating positional embeddings of lower order positions more than higher order ones, due the mixing of graph sizes during training.

---

Published in Transactions on Machine Learning Research (03/2025)
Implementation details Unless otherwise noted, inverse covariance is used for the global statistic and selection score, due to its relationship to partial correlation. We sample batches of size b = 500 over k = 5 nodes each (analysis in 5.4). Our model was implemented with 4 layers with 8 attention heads and hidden dimension 64. Our model was trained using the AdamW optimizer with a learning rate of 1c-4 (Loshchilov) See B.4 for additional details about hyperparameters.
et al. 2017
Complexity The aggregator should be be invariant node labeling, while maintaining the order of sampled subsets, so attention- eetention-based architectures were a natural choice (Vaswani et al., 2017). If we concatenated p and E/... into length 22T input, quadratic-scaling attention would cost O(N*T*. Instead, opted for axial atlention blocks, which attend along each of the three axes separately in O(N + N2T2). Both are parallelizable on GPU, but the latter is more efficient, especially on larger N.
3.4 Theoretical interpretation
Marginal graph resolution It is well-established that estimates of causal graphs over subsets of variables can be "merged" into consistent graphs over their union (Faller et al. 2023 Tillman et al.] 2008 Huang et al.) 2020p. In A.. and A.2 we describe a simple algorithm towards this task, based on the intuition that edges absent from the global graph should be absent from at least one marginal estimate, and that v-structures present in the global graph are present in the marginal estimates. We then prove Theorem 3.1 which states that the axial attention architecture is well-spocified as model class. That is, there exists setting of its weights that can map marginal estimates inlo global graphs. Our construction in A.3 follows the same reasoning steps as the simple algorithm in A.2, and it provides realistic bounds on the model size.
Theorem 3.1. Let G (V, be a directed acyclic graph with marimum degree d. For S CV let E's denote the marginal estimate over S. Let Sa denote the superset that contains all subsets S SVV of size at most d. Given { {F'g)seSuz a stack of L arial attention blocks has the capacity to recover G's skeleton and u-structures in ONN) width, and propogate orientations on paths of O(L) length.
There are two practical considerations that motivate a framework like SEA, instead of directly running classic reconciliation algorithms. First, many of these algorithms rely on specific characterizations of the data-generating process, e.g. linear non-Gaussian (Huang et al. 2020). While our proof does not constrain the causal mechanisms or exogenous noise, it assumes that the marginal estimates are correct. These assumptions may not hold on real data. However, the failure modes of any particular causal discovery algorithmn may be similar across datasets and can be corrected using statistics that capture richer information. For example, an algorithm that assumes linearity will make (predictably) asymmetric mistakes on non-linear data and underestimate the presence of edges. However, we may be able to recover nonlincar relationships with statistics like distance correlation (Sx'ekely et al., (207)). By training a deep learning model to reconcile marginal estimates and interpret global statistics, we are less sensitive artifacts of sampling and discretization (e.g. p-value thresholds, statistics & 0). The second consideration is that checking combinatorial number of subsets is wasteful on smaller graphs and infeasible on larger graphs. In fact, if we only leverage marginal estimates, we must check at least O(22) subsets to cover each edge at least once. To this end, the classical Independence Graph algorithm (Spirtes et al. 2001)) motivates statistics such as inverse covariance to initialize the undirected skeleton and reduce the number of independence tests required. This allows us to use marginal estimates more efficiently, towards answering orientation questions. We verify this latter consideration in Scction 5.4 where we empirically quantify the number of estimates a global statistic is "worth."
On identifiability The primary goal of this paper is to develop a practical framework for causal discovery, especially when data assumptions are unknown. Instead of focusing on the identifiability of any particular setting, we provide these interpretations of our model's outputs, and show empirically that our model respects classic identifiability theory (Section 5.3) The model will always output an orientation for all edges, but the graph can be interpreted as one member of an equivalence class. Metrics can be computed with respect to either the ground truth graph (if identifiable) or the inferred equivalence class, e.g. the implied CPDAG. When data do not match the estimation algorithm's 's assumptions, performance is inherently an empirical question, and we show empirically that our model still does well (Section 5.1
7

---

Published in Transactions on Machine Learning Research (03/2025)
4 Experimental setup
Our experiments aim to address the three desiderata proposed in Section 2.2 adaptability, and emergent few-shot behavior. These experiments span both real and synthetic data. Real experiments quantify the practical utility of this framework, while synthotic experiments allow us to probe and compare each design choice in a controlled setting.
namely, generalization,
4.1 Datasets
We pretrained SEA models on 6,480 synthetic datascts, which constitute approximately 280 million individual observations, each of 10-100 variables | To assess generalization and robustness, we evaluate on unseen in-distribution and out-of-distribution synthetic datasets, as well as two real biological datasets (Sachs et al., 2005; Replogle et al., 2022), using the versions from Wang et al. (2017); Chevalley et al. (2025). To probe for emergent few-shot behavior, we down-sample both the training and testing sets. We also include experiments simulated mRNA datasets with unseen datasets in Appendix C.3 Diibaeinia & Sinha 2020
The training datasets were constructed by 1) sampling Rédós-RRnyi and scale free graphs with N = 10, 20, 100 nodes and E = N, 2N, 3N, 4N expected edges; 2) sampling random instantiations of causal mechanisms (Linear, NN with additive aaditivvenno--addite nonaaditivee Gaussian noise); and 3) iteratively sampling observations in topological order (details in Appendix [B.]). From every causal graphical model (steps 1-2), we generated two datasets, each with 1000N points: either all obsorvational, or split equally among regimes (observational and perfect single-node interventions on all nodes). All models that can accommodate interventions were run on the interventional datasets, with complete knowledge of the intervention target identities. The remaining were the observational datasets. We generated 90 training, 5 validation, and 5 testing datasets for each combination. For testing, we also sampled out-of of-iisiiibttii datasets with 1) Sigmoid and Polynomial mechanisms with Gaussian noise; and 2) Linear with additive non-Ganssian noise.
4.2 Metrics
We report standard causal discovery metrics. These include both discrete and continuous metrics, neural networks can be notoriously uncalibrated (Guo el al. 2017), and arbitrary discretization thresholds may impact the robustness of findings (Ng et al. 2024 Schaeffer et al. 2023). For all continuous metrics, we exclude the diagonal since several baselines manually set it to zero (Brouilladd et al. 2020 Lopez et al. 2022).
SHD: Structural Hamming distance is the minimum number of edge edits required to match two graphs (Tsamardinos et al., 2006) (predicted and true DAGs, or the implied CPDAGs). Discretization thresholds are as published default to 0.5.
mAP: Mean average precision computes the area under precision-recal curve per edge and averages over the graph (predicted and true DAGs, or undirected skeletons). The random guessing bascline depends on the positive rate.
Area under the ROC curve (Bradley 1997) computed per edge (binary prediction) and averaged over the graph. For cach edge, 0.5 indicates random guessing, while indicates perfect performance.
AUC:
Orientation accuracy: We compute the accuracy of edge orientations
OA Eamer 1(P(i,)) P(j,))
(12)
|
Since OA normalized by |[E|], it is invariant to the positive rate. contrast to orientation F1 Gelfner et al. 2024 is also invariant the assignment of forward/revers edges as 1/0.
33 mechanisms, 3 graph sizes, 4 sparsities, topologies, 1000 examples, datasets 280.800.000 examples. For sense of scale, single cell foundation models trained on 300K Rosen al. 2024 30M cells CCui al. 2024

---

Published in Transactions on Machine Learning Research (03/2025)
4.3 Baselines
We compare against several deep learning and classical baselines. All baselines were trained and/or from scratch on each testing dataset using their published code and hyperparameters, except AVICI (their rocommended chockpoint was trained on their synthetic train and test sets after publication, Appendix B.2
DCDI (Brouillad et al. 2020) extracts the causal graph as a parameter of a generative model. The and DsF variants use Gaussian or deep sigmoidal flow likelihoods, respectively. DCD-FG (Lopez al.)
2022 follows DCDI-G, but factorizes the graph into product of two low-rank matrices for scalability. DiffAn (Sanchez et al. 2023 uses the trained model's Hessian to obtain a topological ordering, followed by classical pruning am..
AVICI (Lorch et al. 2022) uses amortized inference approach to estimate P[G D) over a class of gentaggnnrgtingg mechanisms via variational inference. VarSort (a.k.a. "sort and regress") (Reisach et al. 2021 sorts nodes by marginal variance and sparsely regresses nodes based on their predecessors. This naive baseline is intended to reveal artifacts of synthetic data generation. FCI, GIES quantify the predictive power of FCI and GIES estimates, when run over all nodes. VARSORT, FcI, and GIES were run using non-parametric bootstrapping (Friedman al.) 1999), with 100 estimates of the full graph (based on 1000 examples each), where the final prediction for each edge is its frequency of appearing as directed. Since these methods are treated as oraclos, the bootstrapping strategy was selected to maximize test performance. Visnalizations of the bootstrapped graph can be found in Figures and S
5 Results
We highlight representative results in each section, with additional experiments and analyses Appendix C
examines the case where we have no prior knowledge about the data. Our models achieve high performance out-of-the-box, even when the data are misspecified or out-of-domain.
Section 5.1
2. Section 5.2 focuses the case where we do know (or can estimate) the class of causal mechanisms
or exogenous noise. We show that adapting our pretrained models with this information at zero/low cost leads to substantial improvement and exceeds the best baseline trained from scratch.
3. Section 5.3 analyzes SEA predictions in context of classic identifiability theory. In particular, we locus on the linear Gaussian case, and show that SEA approaches "oracle" performance (with respect. to the MEC), while simply running a classic discovery algorithm cannot, on our finite datasets.
4. Section 5.4 contains a variety of ablation studies. In particular, SEA exhibits impressive low-data performance, requiring only 400 samples to perform well on N = 100 datasets. We also ablate estimation hyperparameters and the contribution of marginal/globe global features.
5.1 SEA generalizes to out titto---iitribution misspecified, and real datasets
Table summarizes our controlled experiments on synthetic data. SEA exceeds all baselines in the Linear case, which matches the models' assumptions exactly (causal discovery algorithms and inverse covariance). In the misspecified (NN) or misspecified and ant-of-distribution settings (Sigmoid, Polynomial), SEA also attains the best performance in the vast majority of cases, even though DCDI and AVICI both have access to the raw data. Furthermore, our models outperform VARSORT in every single sctting, while most basclines are unable to do so consistently. This indicates that our models do not simply overfit spurious features of the synthetic data generation process.
Table 2 illustrates that we exceed basclines on single cell gene expression data from CausalBench (Chevalley
et al. 2025 Replogle et al. 2022). Furthermore, when we increase the subset size to = 2000, we achieve very high precision (0.838) over 2834 predicted edges. SEA runs within 5s on this dataset of 162k cells and N = 622 genes, while the fastest naive baseline takes minutes and the slowest deep learning baseline takes hours (run in parallel on subsets of genes).
5.2 SEA adapts to new data assumptions with zero minimal finetuning
We illustrate two strategies that allow us to use pretrained SEA models with different implicit assumptions. First, if two causal discovery algorithms share the same output format, they can be used interchangeably
9

---

Published in Transactions on Machine Learning Research (03/2025)
Table 1: Synthetic experiments. Mean/std over 5 distinct Erdós-Rényi graphs, with metrics relative to ground truth DAG. DIFFAN, VARSORT, FCI, SEA (Fi(F run on observational data only. Evaluation w.r.t.
CPDAG and undirected skeleton in Tables 14 and 15
with additional baselines and ablations in Appendix indicates o.o.d. setting. indicates non-parametric bootstrapping. Runtimes with 1 CPU and 1 V100 GPU.
N
Model
Linear
NN add.
Sigmoid
Polynomial
Overall
mAP t
SHD 4
mAP
SHD 4
mAP
SHD 4
mAP 1
SHD 4
Time(s)
DCDI-G
0.59 .2
6.4
0.78
3.0
0.36= 06
42.7..3
0.42
10.4+.4
4735.7
DCDI-DsF
0.66
5.22.3
0.694.18
4.2+.5
0.37=
43.2+.4
0.262.08
15.7+.2
3569.1
DIFFAN
0.19#.09
40.2=4.4
0.164.10 10
38.6=3.1
0.29=
19.2#.6
0.09
19.7+4.6
434.3
AVIcI 0.48
17.2
0.594.09
10.8
0.42..13
17.21.8
0.24
18.4
2.0
20
20
VARSonr*
0.81#.08
10.0#.4
0.81
6.6+.7
0.50=.18
16.1
0.33
17.1
0.4
Fcr*
0.66+..0
19.0
0.42+.19
17.4+.2
0.56=.os
18.5
0.41
18.9+.3
22.2
GIES*
0.84+.u8
7.4+.u
0.79+.07
9.0+.1
0.71=.10
12.5
0.62
13.7
482.1
SEA (Fcr)
0.962.03
3.22.6
0.91
5.0+.8
0.85=.09
6.72.1
0.69
9.82.2
4.2
SEA (Gıes)
0.97#.02
3.0
0.94
3.4+.4
0.84=
8.14.8
0.69
10.1+.9
3.0
DCD-FG
0.05+.00 .0
30681131
0.07+.00
3428
0.13=
36011272
0.12:.03
33161698
1838.2
AVIcr
0.122.02
391
0.17+.01
407
0.10=.02
398+11
0.03
402+19
9.3
100 400
VARSORT *
0.80+.02
124+10
0.18+.08
1139=269
0.51=
350+15
0.27
380+17
5.1
SEA (FCT)
0.84
162
0.04+.00
403
0.63=
247:17
0.34
325
19.2
SEA (GTES)
0.91
116=7
0.27:.10
364+34
0.69
218:21
0.38
328
3.1
Table 2: Results on K562 single cell data, with STRING database (physical) as ground truth. Base-
Table 3: Performance on Sachs C.55
varies depend-
ing on implicit (AviCI training set) and explicit (SEA variants) assumptions.
lines taken from Chevalley al. (2025).
Model
P1
R 1
F1 1 Time(s)
Model
mAP 1 AUC 1 SHD +
GRNBOOST
0.070
0.710
0.127
316
DCDI-DSP
0.20
0.59
20.0
GIFS
0.190
0.020
0.036
2350
AvIcI-L
0.35
0.78
20.0
NOTEERS
0.080
0.620
0.142
32883
AvICI-R
0.29
0.65
18.0
DCDI-G
0.180
0.030
0.051
16561
AVICI-L,+R
0.59
0.83
14.0
DCDI-DSF
0.140
0.040
0.062
5709
SEA (F)
0.23
0.54
24.0
DCD-FG
0.110
0.070
0.086
6368
+ K
0.33
0.63
14.0
SEA (G)+ CoRn
0.491
0.109
0.179
+ CoRR
0.41
0.70
15.0
with = 2000
0.838
0.093
0.167
+ KcI+ CORR
0.49
0.71
13.0
for marginal estimation. On observational, linear non-Gaussian data, replacing the GES algorithm with LINGAM (Shimizu et al., 2006) is beneficial without any other change (Table 4. The same improvement can be observed on Polynomial and Sigmoid non-additive data, when running FcI with a polynomial kernel conditional independence test (KcI, Zhang et al. (2011h) instead of the Fisherz test, which assumes linearity (Table 5) In principle, different algorithms might make different mistakes, so this strategy could lead out-of-distribution inputs for the pretrained aggregator. While comparing across algorithms not a primary focus of this work and requires further theoretical study, we notice similar performance for alternate estimation algorithms with linear Gaussian assumptions (Table 7 regardless of discovery strategy (GRaSP (Lam et al.
2022 The gap is larger for LINGAM, which assumes non- non-Gaussianity (Table 25
Another strategy is to "finetune" the aggregator, either fully or using low-cost methods like LoRA Hu let al.) 2022). Specifically, we keep the same training set and classification objective, while changing the input's fcaturization, c.g. a different global statistic. Here, we show that finctuning our models for distance correlation (DCOR) is beneficial in both Tables 4 and 5 and the combination of strategies results in the highest performance overall, surpassing the best baseline trained from scratch (DeDI-DsF).
10

---

Published in Transactions on Machine Learning Research (03/2025)
Table 4: Adapting SEA to linear non-Gaussian (Uniform) noise. LINGAM without finetuning; SEA(G) finetuned for distance correlation.
Table 5: Adapting SEA to polynomial, sigmoid non- additive (N=10, E=10). Foi with Kei test; SEA(F) finetuned for distance correlation.
Model
N-
E=10
N=20,
E-20
Model
Polynomial
Sigmoid
mAP f SHD
mAP
SHD 1
mAP f SHD + mAP 1 SHD
DCDI-DSF
0.34
22.3
0.32
63.0
DCDI-DSF
0.39
9.8
0.81
13.6
LINGAM
0.34
7.2
0.30
18.8
Fcr*
0.12
10.6
0.53
8.1
SEA (G) 0.26
12.7
0.12
16.6
SEA (F)
0.22
10.6
0.59
4.8
+LINGAM
0.52
10.1
0.22
39.7
+KCI
0.30
10.6
0.59
5.5
+DCOR
0.44
8.0
0.21
33.1
+DCOR
0.45
9.6
0.90
2.1
+LING+DCOM
0.76
4.6
0.67
14.2
DCCI
0.52
8.2
0.86
3.4
Table 6: SEA respects identifiability theory. Observational setting, standardized (-std) N 10, 10 linear Gaussian test datasets with > graph Markov equivalence class (MEC). Top: oracle performance based true MEC (see left).. Bottom: trained SEA approaches oracle performance, while FCI is very noisy.
mean over adjacency matrices
Model
mAP(t)
AUC(1)
SHD(L)
OA(t)
metric
MEC
0.88 1.
0.98 1.03
2.0 11.
0.74 :.22
mean metric over MEC 0.74
1.07
1.2 1.09
0.84
SEA(Fci)-std
0.83 16
0.97 1.04
3.3 12.3
0.69 2:21
SEA CoRR-std CORR-std FcI-std
0.84 =.14
0.96 1.03
5.0 5.0
0.85 #.14
"true"
"metric over mean MEC
0.49 1.28
0.75 1.16
9.3 12.8
0.49 1.29
graph
"mean metric over MEC"
On real data from unknown distributions, these two strategics enable the ability to run causal discovery with different assumptions, which may be coupled with unsupervised methods for model selection (Faller et al. 2023). Table 3 illustrates this idea using the Sachs proteomics dataset. SEA can be run directly with a different estimation algorithm (FCI with polynomial kernel "Kcr"), or finetuned for around 4-6 hours on A6000 and 4 GB of GPU memory (correlation "CORR"). In contrast, methods like AVICI must simulate new datasets based on each new assumption and re-train/finetung on these data (reportedly around 4 days).
5.3 SEA respects identifiability theory
While the identifiability specific causal models is not a primary focus this work, we show that SEA still respects classic identifiability theory. Specifically, while linear Gaussian models are known be unidentifialle, Table might suggest that both SEA and DCDI perform quite well on these data better than would be expected if graphs were only identifiable up their Markov equivalence classes. This empirical "identifiability" may be the consequence of two findings. Common synthetic data generation schemes tend to result in marginal variances that reflect topological order (Reisach et al. 2021), and in additive noise models, it has been shown that marginal variances that are the "same or weakly monotone increasing in the [topological] ordering" result in uniquely identifiable graphs (Park 2020). Data standardization can climinate these artifacts of synthetic data generation. In Table 5.3 we see that after standardizing linear Gaussian data, our model performs no better than randomly selecting graph from the Markov equivalence class (emumerated via pcalg (Kalisch et al., 2012 The classic FCI algorithm is unable to reach this upper bound, suggesting that the amortized inference Tramework allows us to perform better in finite datasets.
5.4 Ablation studies
In addition to high performance and flexibility, one of the hallmarks of foundation models is their ability to act as few-shot learners when scaling up (Brown et al. 2020). We first confirm that SEA indeod data-cfficient, requiring only around 300-400 examples for performance to converge on datasets of N 100 variables, and outperforms inverse covariance (computed with 500 examples) at only 200 examples (Figure 4A). To probe
11

---

Published in Transactions on Machine Learning Research (03/2025)
Table SEA is generally insensitive to swapping estima- tion other algorithms with linear Gaussian assumptions,
SEA
Tiny SEA
1.00
T=50
inference time. Results on N 10 observational
0.75
T=10
mAP
setting, FCI cannot be used with SEA(G) since FCI outputs a PAG, not a CPDAG.
0.50
0.00
Inference
SEA (FCI)
SEA (GIES)
Lin NN
Sig
Poly Lin
NN
Sig Poly
estimator
Lin.
NN Sig.
Poly. Lin.
NN Sig. Poly.
Figure 3: Few-shot learning behavior emerges as train- ing set increases. "Tiny" SEA trained on 1/4 of the
FCI
0.98
0.88 0.83
0.62
PC
0.93
0.85
0.86
0.64
0.96
0.89
0.82
0.58
data is comparable to the full model on N 10 datasets when given T 50 batches, but is less robust.
GES
0.94
0.85
0.80
0.60
0.95
0.88
0.81
0.57
GRaSP
0.93
0.85 0.80
0.61
0.95
0.88
0.81 0.57
with only T 10.
(A) mAP dataset size
(B) mAP batch size
MAP # batches
mAP subset size
1.0
1.0
1.0
1.0
mAp
0.5
0.5
0.5
0.5
Linear
Linear
Linear,b=500 Sigmood, b5000
Lincar
0.0
Sigmeid 0.0
Sigmoid
0.0
Linear,b=200- Sigmoid,b=300
0.0
Sigmoid
100
200
300
400 1000 2000
50 100
150
200 300
500
2
10 20 30 40 50 100
3
4
5
6
7
10
dataset size (M)
batch (b)
number of batches (T)
# variobles subset
Figure 4: Ablations with SEA (GIES) for estimation parameters on N 100, E 100. Error bars indicate 95% confidence interval across the 5 i..d. datasets of each setting. All parameters are set the defaults (Section 3.3 unless otherwise noted. (A) Dashed: inverse covariance at M 500. (C) Variance is unusually high for Sigmoid b = 300 until T = 100, indicating that larger batches result in more stable results.
for how this behavior emerges, trained a "tiny" version of SEA (GIES) on approximately a quarter of the training data (N = 10, 20 datasets, 64.8 million examples). The tiny model performs nearly as well as the original on N = 10 datasots when provided T= 50 batches, but exhibits much poorer fow-shot behavior with only T = 10 batches (Figure 3 This demonstrates that SEA is able to ingest large amounts of data, leading to promising few-shot behavior
We also ablate each parameter of the estimation step to inform best practices. The trado-off between the number and size of batches may be relevant to estimation algorithms that scale poorly with the number of examples, e.g. kernel-based methods (Zhang et al. 2011). When given T = 100 batches, SEA reaches reasonable performance at around 250-300 examples per batch (Figure A3). Figure 4 further illustrates that on the harder Sigmoid datasets, 5 batches of size b = 500 are roughly equivalent to 100 batches of size 300. Finally, increasing the number of variables in each subset has minimal impact (Figure 4D), which is encouraging, as there is no need to incur the exponentially-scalin runtimes associated with larger subsots.
Finally, we analyze the impact of removing marginal estimates or global statistics (Table 8. First, we take fully pretrained SEA (GIES) and set the corresponding hidden representations to O. Performance drops more when h² is set to 0, indicating that our pretrained aggregator relies more on global statistics, though a sizable gap emerges in both uations.. Then, we re-Irain SEA (GIES) on the N = 10 datasets, with and without global statistics, so that lack of p is in-distribution for the latter model, and the training sets are comparable. Here, we see that the "no p" version with T = 50 estimates is on par with the original architecture with T 10 estimates, so the global statistic equivalent to ~ 40 estimates. This roughly aligns with the theory that global statisties can expedite the skeleton discovery process (Section 3.4), as the number of estimates required to discover the skeleton of a N = 10 graph is approximately (10) 45 (Prop, A.9h.
Due computational limitations, we were unable to train larger models, as our existing training set requires several hundred GB in memory, and our file system does not support fast. dynamic loading.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Table 8: Ablating marginal and global features on SEA (GIES), Top: We set marginal and global represen- tations to 0 (lack of / -utrof-distribution)) and observe that the pretrained model is more robust removing E", perhaps since we sample varying T during training. Bottom: Re-train SEA (GIEs) on N 10 datasets, with and without global features (lack of p is in-distributiinn. We observe that global features are "worth" T ~ 40 estimates of k = 5 variables each.
Model
Linear
NN add.
NN.
Sigmoid
Polynomial
mAP
SHD J
mAP f
SHD 4
mAP t
SHD +
mAP 1
SHD
mAP t
SHD J
SEA (GıEs)
0.99
1.2
0.94
2.6
0.91
3.2+.3
0.85+.12
1.0#.5
0.70+.11
5.8t.6
h² <-0 0
0.30
29.2+.4
0.27
29.4
0.19=
29.0=.0
0.35+..1
27.4+.4
0.
27.1+.9
h +0
0.852.09
6.4
0.822.11
10.2:.9
0.78=.07
13.22
0.63:.21
10.2
0.55
13.42.6
T= 2
0.20#.04
31.24.5
0.252.06
29.22.4
0.33
27.8
0.194.07
30.22.9
28.1 0.24
T 7= 10
0.622.16
8.04.8
0.692.11
9.6
0.66=.13
11.22.9
0.62+.20
9.22.6
0.50
8.92.2
T = 50, no
0.63
6.84.1
0.53
6.2+.9
0.68
6.0#.1
0.58
7.14.1
0.50+.14
7.14.5
6 Discussion
Interventional experiments have formed the basis of scientific discovery throughout history, and in recent years, advances in the life sciences have led to datasets of unprecedented scale and resolution (Replogle et al. 2022: Nadig et al. 2024). The goal of these perturbation experiments is to extract causal relationships between biological entitics, such as genes or proteins. However, the sheer size, sparsity, and noise level of these data pose significant challenges to existing causal discovery algorithms. Morcover, these real datasets do not fit cleanly into causal frameworks that are designed around fixed sets of data assumptions, either explicit (Spirtes) et al. 1995) or implicit (Lorch et al., 2022). In this work, we approached these challenges through causal discovery "foundation model." Central to this concept were three goals. First, this model should generalize to unseen datasets whose data-generating mechanisms are unknown, and potentially out-of-distribution. Second, it should be easy to steer the model's predictions with inductive biases about the data. Finally, scaling up the model should lead to data-efficiency. We proposed SEA, a framework that yields causal discovery foundation models. SEA was motivated by the idea that classical statistics and discovery algorithms provide powerful descriptors of data that are fast to compute and robust across datasets. Given these statisties, we trained a deep learning model to reproduce faithful causal graphs. Thcoretically, we demonstrated that it is possible to infer causal graphs consistent with correct marginal estimates, and that our model is well-specified with respect to this task. Empirically, we implemented two proofs of concept of SEA that perform well across a varicty of causal discovery tasks, casily incorporate inductive biases when they are available, and exhibit excellent few-shot behavior when scaled up.
While SEA provides a high-level framework for supervised causal discovery, there are several empirical limitations of the two implementations describe in this paper. These include: 1) an arbitrary, hard-coded maximum of 1000 variables, 2) poor generalization to synthetic cyclic data, 3) erring the side of sparsity real data, 4) numeric instability of inverse covariance on larger graphs, and as a result, 5) training requires full precision. The first three aspects may be addressed by modifying the architecture and/or synthetic training datasets, while the latter two can be addressed with more numerically stable statistics, like correlation.
More broadly, the success of supervised causal discovery algorithms derives from the fidelity of the data simulation procedure. In this paper, we present proofs of concept for the modeling framework, but we do not solve the problem of simulating realistic data. Of the data that may exist in the real world, the training data used here represent only a small, perhaps unrealistic, fraction. To achieve any semblance of trustworthiness on real applications, it is crucial to study the characteristics of cach domain in detail - including common graph topologies and functional forms (Aguirre el al.) 2024); the degree and nal ure of missing data (Hicks) lt al. 2018); and sources of measurement error or other covariates (Tran et al. 2020). When possible, we recommend that any insights be triangulated with other sources of knowledge. For example, while this work does not directly provide an inference framework, the predicted structure could be used to parametrize a generative model, whose likelihood could be evaluated on held-out interventions (Hägele et al.] 2023). It is
13

---

Published in Transactions on Machine Learning Research (03/2025)
also important to check whether the inferred relationships compound upon any existing biases in the data. This is particularly important for sensitive domains like healthcare legal applications.
This work also opens several directions for further investigation. The framework we describe utilizes a single causal discovery algorithm and a single global statistic. Classic causal discovery algorithms leverage diverse insights for identifying causal structure, e-g- the non-Gaussianity of noise (Shimizu et al. 2006) vs.
conditional independence (Spirtes et al. 1995). Thus, different discovery algorithms or summary statistics may reveal different aspects of the causal structure. Learning to resolve these potentially conflicting views remains unexplored, both from experimental and theoretical perspectives. Furthermore, this work only shows that the axial attention architecture is well-specified as model class and probes generalization empirically. This motivates theoretical studies into supervised causal discovery with regards to what information can be provably learned, e.g. in the style of of PAC (probably approximately correct) learning (Allen-2h) Zhu et al., 2019
In summary, hope that this work will inspire a new avenue of research into causal discovery algorithins that are applicable to and informed by real applications.
Acknowledgements
We thank Bowen Jing, Felix Faltings, Sean Murphy, and Wenxian Shi for helpful discussions regarding the writing; well as Jiaqi Zhang, Romain Lopez, Caroline Uhler, and Stephen Bates for helpful feedback regarding the framing of this project. Finally, we thank our action editor Bryon Aragam and our anonymous reviewers for their invaluable suggestions towards improving this paper during the review process.
This material is based upon work supported by the National Science Foundation Graduate Research Fellowship under Grant, No. 1745302. We would like acknowledge support from the NSF Expeditions grant (award 1918839: Collaborative Research: Understanding the World Through Code), Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, and the Abdul Latif Jameel Clinic for Machine Learning in Health.
14

---

Published in Transactions on Machine Learning Research (03/2025)
References
Matthew Aguirre, Jeffrey P Spence, Guy Sella, and Jonathan K Pritchard. Gene regulatory network structure informs the distribution of perturbation effects. bioRriu, 2024. doi: 00.1101/2024.07.44600222
Hirotugu Akaike. Information theory and an extension of the maximum likelihood principle. In 2nd International Symposium on Information Theory, pp. 267 281. Akadémiai Kiadó Location Budapest, Hungary, 1973.
Zeyuan Allen-Zhi Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in Neural Information Processing Systems, 32, 2019.
Steen A. Andersson, David Madigan, and Michael D. Perlman. A characterization of Markov equivalence classes for acyelic digraphs. The Annals of Statistics, 25(2):505 541, 1997. doi:
Bryon Aragam, Arash A Amini, and Qing Zhou. Learning directed acyclic graphs with penalized neighbourhood regression. Xiv preprint MKi:15511000665 2015.
Rishi Bommassni, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bolig, Antoine Bosselut, Emna Brunskill, et al. On the opportunities and risks of foundation models. Xr preprint orXic:2108.079 07255, 2021.
Andrew P. Bradley. The use of the area under the roc curve in the evaluation of machine learning algorithms. Pattern Recognition, (0(7):1155-119) 1997. ISSN 0031-3203. doi: https://doi.orr. https///dii.rg/1001166/// 00142-2.
Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865 21:77765-21877, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranay Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Newral Information Processing Systems, 33:1877-1901, 2020.
Peter Bühhmann, Jonas Peters, and Jan Ernest. CAM: Causal additive models, high-dimensionall order search and penalized regression. The Annals of Stalistics, 42(6):2526 2556, 2014. doi: 10.1214/14AA0S1000
Mathieu Chevalley, Yusuf H Roohani, Arash Mehrjou, Jure Leskovec, and Patrick Schwab. A large-scale benchmark for network inference from single-cell perturbation data. Communications Biology, 8(1):412, 2025.
David Maxwell Chickering. Optimal structure identification with greedy search. Journal of Machine Learning Research, 3:507 554, November 2002.
Haotian Cui, Chloe Wang, Hassaan Maan, Kuan Pang, Fengning Luo, Nan Duan, and Bo Wang, scGPT: toward building foundation model for single-cell multi-omics using generative AL Nature Methods, pp. 1-11, 2024.
George Cybenko. Approximation by superpositions of sigmoidal function. Mathemalics of control, signals and systems, 2(4):303 314, 1989.
Russell Davidson and James G MacKinnon. Several tests for model specification in the presence of alternative hypotheses. Econometrica: Journal of the Econometric Society, pp. 781-793, 1981.
Payam Dibaeinia and Saurabh Sinha. SERGIO: A single-cell expression simulator guided by gene regulatory networks. Cell Systems, 11(3):252 271.e11, 2020. ISSN 2405-4712. doi: http://dooooog//00000////..ee.e/ 08.003.
Frederick Eberhardt, Clark Glymour, and Richard Scheines. N-1 experiments suffice to determine the causal relations among variables. Innovations in machine learning: Theory and applications, pp. 97-112, 2006.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Philipp M. Faller, Leena Chenmuru Vankadara, Atalanti A. Mastakouri, Francesco Locatello, Dominik Janzing Karlsruhe Institute of Technology, and Amazon Research Tuebingen. Self-compatiblliy: Evaluating causal discovery without ground truth. International Conference on Artificial Intelligence and Statistics, 2023.
Nir Friedman, Moisés Goldszmidt, and Abraham Wyner. Data analysis with Bayesian networks: A bootstrap approach. Proceedings of the Fifleenth Conference on Uncertainly in Artificial Intelligence, UA1'99, pp. 196 206, 1999.
Tomas Gef[ner, Javier Antoran, Adam Foster, Wenbo Gong, Chao Ma, Emre Kiciman, Amit Sharma, Angus Lamb, Martin Kukla, Nick Pawlowski, Agrin Hilmkil, Joel Jennings, Meyer Scetbon, Miltiadis Allamanis, and Cheng Zhang. Deep end-to-end causal inference. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Clark Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in Genetics, 10, 2019. ISSN 1664-8021. doi: 10.3389 1..3899/fgene.211.. 2019.0054.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, 2017.
Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional Markov equivalence classes of directed acyelic graphs. Journal of Machine Learning Research, 13(79):2409 2464, 2012.
Stephanie C Hicks, F William Townes, Mingxiang Teng, and Rafael A Irizarry. Missing data and technical
variability in single-cell rna-sequencing experiments. Biostatistics, 19(4):562 -78, 2018.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Newral Networks, 2(5):359-30 1989. ISSN 0893-6080. doi: https://dooi.or/11..0101 (883-6000088))90000-8
Patrik Hoyer, Dominík Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlincar causal discovery with additive noise models. Advances in Neural Information Processing Systems, 21, 2008.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
Biwei Huang, Kun Zhang, Mingming Gong, and Clark Glymour. Causal discovery from multiple data sets with non-identical variable sets. Proceedings of the AAAI Conference on Artificial Intelligence, 34(06): 10153-10161, Apr. 2020. doi: 20.1609/aaai.v34100
Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Learning, pp. 2078 2087. PMLR, 2018.
Alexander Hägele, Jonas Rothfuss, Lars Lorch, Vignesh Ram Somnath, Bernhard Schölkopf, and Andreas Krause. BaCaDl: Bayesian causal discovery with unknown interventions. In International Conference on Artificial Intelligence and Statistics, pp. 1411 1436. PMLR, 2023.
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What is local optimality in nonconvex-nomooonann minimax optimization? In International Conference on Machine Learning, pp. 4880-4888 PMLR, 2020.
Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta. Causal discovery toolbox: Uncovering causal relationships in Python. Journal of Machine Learning Research, 21(3))11-5, 2020.
Markus Kalisch, Martin Mächler, Diego Colombo, Marloes H Maathuis, and Peter Bühlmann. Causal inference using graphical models with the R. package pealg. Journal of statistical software, 47:1-26, 2012.
Nan Rosemary Ke, Silvia Chiappa, Jane X Wang, Jorg Bornschein, Anirudh Goyal, Mclanic Rey, Theophanc Weber, Matthew Botvinick, Michael Curtis Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. In International Conference Learning Representations, 2023.
16

---

Published in Transactions on Machine Learning Research (03/2025)
Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, and Simon Lacoste-Julien. Gradient-based neural DAG learning. In International Conference. on Learning Representations, 2020.
Wai-Yin Lam, Bryan Andrews, and Joseph Ramsey. Greedy relaxations of the spansest permutation algorithm. In James Cussens and Kun Zhang (eds.), Proceedings of the Thirty Eighth Conference Uncertainty in Artificial Intelligence, volume 180 of Proceedings of Machine Learning Research, pp. 1052-1062. PMLR, 01-05 Aug 2022.
Olivier Ledoit and Michael Wolf. A well-conditioned estimator for large-dimensional covariance matrices. Journal of Mullivariale Analysis, 88(2):365 411, 2004. ISSN 0047-259X. doi: htlpss//doi.orr//10.0011 50044722500(03300006---44
Phillip Lippe, Taco Cohen, and Efstratios Gavves. Efficient neural causal discovery without acyclicity constraints. In International Conference on Learning Representations, 2022.
Po-Ling Loh and Peter Bühhmann. digh-dimensional learning of linear causal networks via inverse covariance ostimation. The Journal of Machine Learning Research, 15(1):3065-3105 2014.
Romain Lopez, Can-Christian Hütter, Jonathan K. Pritchard, and Aviv Regev. Large-scale differentiable causal discovery of factor graphs. In Advances in Neural Information Processing Systems, 2022.
Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized inference for causal structure learning. Advances in Newral Information Processing Systems, 35:13104 13118, 2022.
llya Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in Adam. arXiv preprint arXiv:1111.05101 5, 2017.
Vitaly Maiorov and Allan Pinkus. Lower bounds for approximation by MLP neural networks. Nearocomputing, 25(1-3):81-91, 1999.
Peter McCullagh. What is statistical model? The Anmals of Statistics, 1065):1225-1310, 2002.
Joris M Mooij, Sara Magliacane, and Tom Claassen. Joint causal inference from multiple contexts. Journal of Machine Learning Research, 21 (99):1 108, 2020.
Ajay Nadig, Joseph M. Replogle, Angela N. Pogson, Steven McCarroll, Jonathan S. Weissman, Elise B. Robinson, and Luke O'Connor. Transcriptome-wide characterization of genetic perturbations. bioRwie, 2024. doi: 10. /01.//02220777// 0....000...609933.
Yatin Nandwani, Abhishck Pathak, and Parag Singla. A primal dual formulation for deep learning with constraints. Advances in Neurul Information Processing Systems, 32, 2019.
Ignavier Ng, Biwei Huang, and Kun Zhang. Structure learning with continuous optimization: A sober look and beyond. In Causal Learning and Reasoning, pp- 71-105. PMLR, 2024.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feodback. Advances in Neural Information Processing Systems, 01:27730-27744 2022.
Gunwoong Park. Identifiability of additive noise models using conditional variances. Journal of Machine Learning Research, 21(75):1-34, 2020.
Anne Helby Petersen, Joseph Ramsey, Claus Thorn Ekstrøm, and Peter Spirtes. Causal discovery for observational sciences using supervised machine learning. Journal of Data Science, 21(2):255 280, 2023. ISSN 1680-743X. doi: 10.6339/283 JDS1088.
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the Turing completeness of modern neural network architectures. In International Conference Learning Representations, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
17

---

Published in Transactions on Machine Learning Research (03/2025)
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pp. 8748 8763, PMLR, 2021.
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA transformer. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8844-8856. PMLR, 18-24 Jul 2021.
Alexander G. Reisach, Christof Seiler, and Sebastian Weichwald. Beware of the simulated DAG! causal discovery benchmarks may be easy to game. Advances in Newral Information Processing Systems, 34, 2021.
Patrik Reizinger, Yash Sharma, Malthias Bethge, Bernhard Schölkopf, Ferenc Huszár, and Wieland Brendel. Jacobian- basebian-based causal discovery with nonlinear ICA. Transactions on Machine Learning Research, 2023.
ISSN 2835-8856.
J. M. Replogle, R. A. Saunders, A. N. Pogson, J. A. Hussmann, A. Lenail, A. Guna, L. Mascibroda, E. J. Wagner, K. Adelman, G. Yanai,, N. Iremadze, F. Oberstrass, D. Lipson, J. L. Bonnar, M. Jost, T. M. Norman, and J. S. Weissman. Mapping information-rich senotype-phenofyye landscapes with genome-scale Perturb-seq- Cell, 185(14):2559 257-255, Jul 2022.
Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the few-shot paradigm. In Extended abstracts of the 2021 CHT conference on human factors in computing systems, pp-
1 7, 2021.
Yanay Rosen, Maria Brbić, Yusuf Roohani, Kyle Swanson, Ziang Li, and Jure Leskovec. Toward universal cell embeddings: integrating single-cell RNA RNA-seq datasets across species with SATURN. Nature Methods, pp. 1-9, 2024.
Karen Sachs, Omar Perez, Dana Pe'er, Douglas A. Lauffenburger, and Garry P. Nolan. Causal protein- signaling networks derived from multiparamette single-cell data. Science, :523 529, 2005. doi: 10.1126/sciences 1105809.
Pedro Sanchez, Xiao Liu, Alison Q. O'Neil, and Sotirios A. Tsaftaris. Diffusion models for causal discovery via topological ordering. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwandaa May 1-5, 2023. OpenReview.nee 2023.
Clayton Sanford, Danicl J Hsu, and Matus Telgarsky. Representational strengths and limitations of Trans- formers. Advances in Newral Information Processing Systems, 36, 2024.
Rylan Schacffer, Brando Miranda, and Oluwasanmi Koyejo. Are emergent abilities of large language models a mirage? Advances in Newral Information Processing Systems, 0ts/2304.11000 2023.
Gideon Schwarz. Estimating the dimension of a model. The Annals of Statistics, 6(2):461 464, 1978. doi: 10.1214/aus/1173341133
Han Shao, Omar Montasser, and Avrim Blum. A theory of PAC learnability under transformation invariances. Advances Newral Information Processing Systems, 35:13989-14000 14001, 2022.
Zuowei Shen, Haizhao Yang, and Shijun Zhang. Optimal approximation rate of relu networks in terms of width and dopth. Journal de Mathématiques Pures et Appliquées, 157:101-135, 2022.
Shohei Shimizu, Patrik O. Hoyer, Aapo Hyvarinen, and Antti Kerminen. A linear non Gaussian acyclic model for causal discovery. Journal of Machine Learning Research, ((72):2003-203) 2006.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causality from probability. In Conference Proceedings: Advanced Computing for the Social Sciences, 1990.
Peter Spirtes, Christopher Moek, and Thomas Richardson. Causal inference the presence of latent variables and selection bias. In Proceedings of the Eleventh Conference on Uncertainty in Artificial Intelligence, UAT'95, pp. 499 506, San Francisco, CA, USA, 1995. Morgan Kaufmann Publishers Inc. ISBN 1558603859.
18

---

Published in Transactions on Machine Learning Research (03/2025)
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, Prediction, and Search. MIT Press, 2001. doi: https://doi.org/1// 10. 1tt.7551mmmmmms///// 11tp://1111..00000000
G'abor Sz'ekely, Maria L. Rizzo, and Nail K. Bakirov. Measuring and testing dependence by correlation of distances. Annals of Statistics, 35:2769- 27944, 2007.
Robert Tillman, David Danks, and Clark Glymour. Integrating locally learned causal structures with overlapping variables. Advances in Neural Information Processing Systems, 21:1665 1672, 01 2008,
Hoa. Thi Nhu Tran, Kok Siong Ang, Marion Chevrier, Xiaomeng Zhang, Nicole Yee Shin Lee, Michelle Goh, and Jinmiao Chen. A benchmark of batch-effect correction methods for single-cell RNA sequencing data. Genome biology, 21:1-32, 2020.
Ioannis Tsamardinos, Laura E Brown, and Constantin F Aliferis. The max-min hill-climbing Bayesian network structure learning algorithm. Machine learning, 65(1):31-77 17,1178881 2006.
Sara van de Geer and Peter Bühlmann. 10-penalized maximum likelihood for sparse directed acyclic graphs.
The Annals of Statistics, 41(2):536-500 567, 2013.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fo 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, 2017.
Tom S. Verma and Judea Pearl. On the equivalence of causal models. In Proceedings of the Siath Conference on Uncertainty in Artificial Intelligence, 1990.
Yuhao Wang, Liam Solus, Karren Yang, and Caroline Uhler. Permutation- based causal inference algorithms with interventions. In Advances in Neural Information Processing Systems, volume 30, 2017.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are Transformers
universal approximators of sequenco-to-sequence functions? In International Conference Learning Representations, 2020.
Kun Zhang, J. Peters, Dominik Janzing, and Bernhard Scholkopf. Kernel-based conditional independence test and application causal discovery. Conference on Uncertainty in Artificial Intelligence, 2011.
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.
Yujia Zheng, Biwei Huang, Wei Chen, Joseph Ramsey, Mingming Gong, Ruichu Cai, Shohei Shimizu, Peter Spirtes, and Kun Zhang. Causal-learn: Causal discovery in Python. Journal of Machine Learning Research, 2- 2024.
19

---

Published in Transactions on Machine Learning Research (03/2025)
A Theoretical motivations
Our theoretical contributions focus on two primary directions.
We formalize the notion of marginal estimates used in this paper, and prove that given sufficient marginal cstimates, it is possible to recover a pattern faithful to the global causal graph. We provide lower bounds on the number of marginal estimates required for such task, and motivate global statistics as efficient means to reduce this bound.
2. We show that our proposed axial attention has the capacity to recapitulate the reasoning required for marginal estimate resolution. We provide realistic, finite bounds on the width and depth required for this task.
Before these formal discussions, we start with a toy example to provide intuition regarding marginal estimates and constraint-based causal discovery algorithms.
A.1 Toy example: Resolving marginal graphs
Consider the Y-shaped graph with four nodes in Figure 5) Suppose we run the PC algorithm on all subsets of three nodes, and we would like to recover the result of the PC algorithm on the full graph. We illustrate how one might resolve the marginal graph estimates. The PC algorithm consists of the following steps (Spirtes et al. 2001).
1. Start from the fully connected, undirected graph on N nodes.
2. Remove all edges (i where X; I1. Xj.
3. For each edge (i,j) and subsets S C [N] ((.)) of increasing size n = 1, .,2,,...d, d, where d is the
maximum degree in G, and all E S are connected to either or j: if X, X5 S, remove edge
(i.j).
For each triplet ki.j.k), such that only edges ki,k) and (j. remain, if l was not in the set S that eliminated edge (i,j), then orient the "v-structure" as i - k < j-
5. (Orientation propagation) If + j, edge (j, k) remains, and edge (i, k) has been removed, orient
j + k. If there is directed path i ~4 j and an undirected edge (i,j), then orient + j.
Ground truth
(A)
(D)
X
X
X
Nodes
available
7
W
Z
W
W
W
Z
W
Y
Y
X
PC algorithm output
2
W
W
w
Figure 5: Resolving marginal graphs. Subsots of nodes revealed to the PC algorithm (circled in row 1) and outputs (row 2).
In each of the four cases, the PC algorithm estimates the respective graphs follows.
(A) We remove edge (X, Y) via (2) and orient the v-structure.
(B) We remove edge (X, Y) via (2) and orient the v-structure.
(C) We remove edge (X, W) via (3) by conditioning on Z. There are no v-structures, so the edges remain undirected.
(D) We remove edge (Y, W) via (3) by conditioning on Z. There are no v-structures, so the edges remain undirected.
20

---

for all vertices X, Y of G, X and Y are adjacent if and only X and Y are dependent conditional on every set of vertices of G that does not include or Y; and 2. for all vertices X, Y, Z, such that X is adjacent to Y and Y is adjacent to Z and X and Z are not adjacent, X + Y t- Z is a subgraph of G if and only if X, Z are dependent conditional on every set containing Y but not X or Z. Published in Transactions on Machine Learning Research (03/2025)
The outputs (A-D) admit the full PC algorithm output the only consistent graph four nodes.
X and Y are unconditionally independent, so subset will reveal an edge between (X,Y).
There are no edges between (X, W) and (X, W). Otherwise, (C) and (D) would yield the undirected triangle.
X,Y,2 Z must be oriented as X 2 - Y. Paths X + - y and X f- + Y would induce an (X, Y) edge in (B). Reversing orientations X < Z Y would contradict (A).
(Y, must be oriented as Y + Z. Otherwise, (A) would remain unoriented.
A.2 Resolving marginal estimates into global graphs
Classical results have characterized the Markov equivalency class of directed acyclic graphs. Two graphs are observationally equivalent if they have the same skeleton and v-structures (Verma Pearl 1990p. Thus, a pattern P is failhful to graph G if and only if they share the same skeletons and v-structures (Spirtes et al. 1990).
Definition A.1. Let G EV, be directed acyclic graph. A pottern P is set of directed and undirected edges over V.
Definition A.2 (Theorem 3.4 from Spirtes et al. (2001) ) If pattern P is faithful to some directed acyclic graph, then P is faithful to G if and only if
Given data faithful to G, number of classical constraint-basee algorithms produce patterns that are faithful to G. We denote this set of algorithms as F.
Theorem A.3 (Theorem 5.1 from Spirtes et al. (2001)). }). If the inpul to the PC, SGS, PC-1, PC-2, PC", or IG algorithms faithful to directed acyelic graph G, the output is a pattern that represents the faithful indistinguishabiitty class of G.
The algorithms in F are sound and complete if there are no unobserved confounders.
Let Py be a probability distribution that is Markov, minimal, and faithful to G. Let D E MxN Pv be dataset of M obscrvations over all N = [V] nodes.
Consider a subset S C V. Let D[S] denote the subset of D over S,
D[S = (x 5)
(13)
and let G[S) denole the subgraph of G induced by S
G = (S,
E}.
(14)
If we apply any ʃ E F to D[S], the results are not necessarily faithful to G[S), as now there may be latent confounders in V S (by construction). We introduce the term marginal estimate to denote the resultant pattern that, while not faithful to G[S), is still informative.
Definition A.4 (Marginal estimate). A pattern E' is a marginal estimate of G(S) if and only if
1. for all vertices X,Y of S, X and Y are adjacent if and only if X and Y are dependent conditional on every set of vertices of that does not include X or Y; and for all vertices X, X,Y Y, Z, such that X is adjacent to Y and Y is adjacent to and X and Z are not adjacent, X + < Z is a subgraph of if and only if X, Z are dependent conditional on every set containing Y but not X or Z.
21

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm 1 Resolve marginal estimates of E F
1: Input: Data De faithful to G
2: Initialize E' +- Kv as the complete undirected graph on N nodes.
3: for 5 Sat2 do
4: Compute E = ((Dass)
for (i.j) & E's do
Remove (i,j) from E' end for 8: end for
9: for E's E (E'5) do
10: for v-structure j in E's do
11: if {i,j), 1. E and ki & E' then
12: Assign orientation - c k
13: end if
14: end for
15: end for
16: Propagate orientations in E' (optional).
be directed acyclic graph with maximum degree d. For S C V, let E's denote the marginal estimate over S. Let Sa denote the superset that contains all subsets S S V of size at most d. Algorithm maps {F's)seSa12
Proposition A.5. Let G = EV (V,
pallern E' faithful G.
On a high level, lines 3-8 recover the undirected "skeleton" graph E", lines 9-15 recover the v-structures, and line 16 reforences step in Section A.I
Remark A.6. In the PC algorithm (Spirtes et al. 2001 A.1 its derivatives, and Algorithm there is no need to consider separating sets with cardinality greater than maximum degree d, since the maximum number of independence lests required to separate any node from the rest of the graph is equal to number of its parents plus its children (due to the Markov assumption).
Lemma A.7. The undirected skeleton of E is equivalent to the undirected skeleton of E'
C [ (1,5)
((,5)) E or (i,)) E ([i,)) (i.5) E E' or (i,0) E =
(15)
That is, {i,}} E C = + fi.)eC..
Proof. It is equivalent to show that f = (i,s} & C
= If [i,j] g C*, then there must exist a separating set S in G of at most size such that II. y S.S Then (U,j), is a set of at most size + 2, where {ú,j} & C Csurso Thus, {i,j} would have been removed from
C' in line 6 of Algorithm 1
+ If (i,j) & C let S be separating set in Sa42 such that {i,j} & CSsuuss) and i TT LIS. S. S is also a separating set in G, and conditioning on S removes {i,j} from Ch.
Lemma A.8. v-stracture + exists in E* if and only if there exists the same v-structure in E'
Proof. V-structures are oriented + + ls E* if there is an edge between {i,j} and (i,k) k) but not { k); and if j was not in the conditioning set that removed {i, k). Algorithm [ orients v-structures E' if they are oriented as such in any Es; and if {i,j), (j.k) e E", ( E'
t k in
Suppose for contradiction that i + + k is oriented as a v-structure in E", but not in E". There are two cases.
No E's contains the undirected path j-j- k. If either i or j k are missing from any E's then E would not contain (i,j) or (). Otherwise, if all S contain {i. k), then E" would not be missing
(i, k) (Lemma A.7.
22

---

Published in Transactions on Machine Learning Research (03/2025)
2. In every E'g that contains i-j-k, is in the conditioning set that removed {i, i... T k S, 3 j. This would violate the laithfulness property, as j is neither a parent of i or k in E*, and the outputs of the PC algorithm are faithful to the equivalence class of G (Thherem 5.1 Spirtes et al. (2001)).
* Suppose for contradiction that + t- is oriented as a v-structure in E', but not in E". By Lemma A.7 the path - j - must exist in F². There are two cases.
If ii k or j k, then j must be in the conditioning that removes (i, no E's containing ki.j.k} would orient them v-structures.
2. If j is the root of fork i 4- + k, then as the parent of both i and k, j must be in the conditioning set that removes {i, k,, so no E'g containing ji,j,k] would orient them as v-structures.
Therefore, all v-structures in E' are also v-structures in E
Proof of Proposition A.3 Given data that is faithful to G, Algorithm produces a pattern E' with the same connectivity and v-structures as E*. Any additional orientations in both patterns are propagated using identical, deterministic procodures, so =
This proof presents a deterministic but inefficient algorithm for resolving marginal subgraph estimates. In reality, it possible to recover the undirected skeleton and the v-structures of G without checking all subsets S E S242-
Proposition A.9 (Skeleton bounds). Let G EV be a directed acyclic groph with marimum degree d. It takes O(N2) marginal estimates over subsets of size d 2 to recover the undirected skeleton of G.
Proof- Following Lemma A.7 an edge (i,j) is not present in C if it is not present in any of the size d 2 estimates. Therefore, every pair of nodes {i,j} requires only single estimate of size d + 2, so it is possible to recover C' in (N) estimates.
Proposition A.10 (V-structures bounds). Let G EV be directed acyclic graph with mazimum degree d and u-structures. It is possible to identify all v-structures in O(v) estimates over subsets of at most size d 2.
Proof. Each v-structure i + 5- lo falls under two cases.
1. i. II. h unconditionally. Then an estimate over (i.j,k} will identify the v-structure. 2. i IL A | S, where j $ S C V. Then an estimate over SU ji,j,k will identify the v-structure. Note that |S| S + 2 since the degrec of is at least [S] + 1.
Therefore, each v-structure only requires one estimate, and possible to identify all v-structures in ((v) estimates.
There are three takeaways from this section.
1. If exhaustively the estimate of the full graph. However, this is more efficient than running the causal discovery algorithm on the full graph.
constraint- based algorithm all subsets of size d + it is trivial to recover
2. In theory, it is possible to recover the undirected graph in O(N2) estimates, and the v-structures in (v) estimates. However, we may not know the appropriate subsets ahead of time. In practice, if we have a surrogate for connectivity, such as the global statistics used in SEA, then we can vastly reduce the number of estimates used to eliminate edges from consideration, and more effectively focus sampling subsets for orientation determination.
23

---

Published in Transactions on Machine Learning Research (03/2025)
A.3 Model specification
Model specification and universality In classical statistics, a statistical model can be expressed as pair (S, Po) for 0 E O where S is the sample space, P is the family of distributions on S, and θ is the space of parameters (McCullagh, 2002). Let .NN be observations of i..dd random variables in S, and let pa denote their common distribution. A statistical model is well-specified if P = Po for some 0 E O. Identifying the most appropriate class of models has long been an area of interest in statistics (Akaike, 1973). For simple models like linear regression, there are many ways to test for model specification, where the alternative may be that the data follow a non-linear relationship (Davidson & MacKinnon 1981).
In machine learning, however, even the simplest architectures vary immensely. Whether a neural network is "well-specified" depends on many aspects, such as its width (number of hidden units), depth (number of layers), activation functions (sources of non-linearity); and more. Thus, instead of testing whether each neural network architecture is well-specified for each experiment, it is more common to show universality (or lack thereof). That is, given a class of functions F and the space of parameters θ for every f F does there exist a 0 θ that allows the neural network to approximate f to arbitrary accuracy? For example, the most well-cited work for the universality of multi-layer perceptrons (Hornik et al. 1989) showed that:
"standard multilayer feedforward networks with as few as one hidden layer using arbitrary squashing functions are capable of approximating any Borel measurable function from one finite dimensional space to another to any desired degree of accuracy, provided sufficiently many hidden units are available."
Note that this definition excludes neural networks that use rectified linear units (ReLU is unbounded), whose variants are ubiquitous today, as well as neural networks with bounded width. Universality under alternate assumptions has been addressed by many later works, e.g. Maiorov & Pinkus (1999); Shen et al. (2022).
With respect to modern Transformer architectures (Vaswani et al. 2017/, the existing literature are similarly subject to constraints. For example, Pérezz al. (2019) shows that Transformers equipped with positional encodings are Turing complete, given infinile precision. Alternatively, Yun et al. (2020) proves that Trans- formers can approximate arbitrary continuous "sequence-to-ssqqueeee functions, but requires in the worst case, exponential depth. Finally, Sanford et al. (2024) shows that one-dimensional self attention is wnable to detect arbitrary triples along a sequence, unless the depth scales linearly as the input size (sequence length). This last case is particularly relevant to our setting, as v-structures are relations of three nodes, but even modern large language models "only" contain ~ 100 layers (e.g. GPT-4 is reported to have 120 layers). In general, these works present constructive proofs of either (Pérez el al. 2019 Yun el al. 2020 or pathological cases that cannot be handled (Sanford al. 2024).
In this paper, do not consider universality in general sense, but rather limit our scope to the algorithms presented in Section A.2 This will allow us to make more realistic assumptions regarding the model class.
Causal identifiability for continuous discovery algorithms Identifiability is central to causal discovery, as it is important to understand the degree to which a causal model can or cannot be inferred from data. In contrast to classic approaches, causal discovery algorithms that rely on continuous optimization must consider several additional aspects with respect to identifiability - namely, their objective; optimization dynamics; and model specification. We discuss representative works in light of these considerations.
First, the optima of the objectives must correspond to true graphs. NOTEARS (Zheng et al. 2018 focuses on the linear SEM case with least-squares loss, and they cite earlier literature regarding the identifiabiltty of this setup (van de Geer & Bühlmann 2013 Loh & Bühlmann 2014 Aragam et al. 2015). DCDI (Brouillard et all 2020) proves that the graph that maximizes their proposed score is T-Markov equivalent to the ground truth, subject to faithfulness and other regularity conditions. Since amortized causal discovery algorithms tend to be trained on data from a variety of tata-generating processes, the identifiability of an arbitrary graph is less clear. Ke et al. (2023) cites Eberhardt et al.] (2006) and claims that any graph identifiable "in the limit of an infinite amount of [single-node hard] interventional data samples." " In our interventional setting, this also holds. AVICI (Lorch et al. 2022) focuses solely on inferring graphs from data, without addressing whether
24

---

Published in Transactions on Machine Learning Research (03/2025)
they are identifiable. If we suppose that the graphs are identifiable, then in all three cases, the objectives can be written in terms of the KT, divergence between the true and predicted edge distributions, which reaches its minimum when the predicted graph matches the true graph.
Second, the optimization process must not only converge to an optimum, but also (for amortized models) generalize to unseen data. This is, in general, difficult show. Both NOTRARS and DCDI acknowledge that due to the non-convexity of this optimization problem, the optimizer may converge to stable point, which is not necessarily the global optimum. AVICI includes explicit acyclicity constraint, they additionally cite Nandwani et al. 20019) Jin et al. (2020) for inspiration in constrained optimization of a neural network. Regarding generalization, existing results on what can be "provably" learned by neural networks are limited to simple architect ures and/or algorithms (Allen-Zhu et al.) 2019 Shao et al. 2022). Instead, both AVICI and our work assess generalization empirically, by holding out certain causal mechanisms from training.
Finally; the implementation of cach model must be well-specified. Since NOTEARS assumes lincarity, it directly optimizes the (weighled) adjacency, and the only concern is whether linearity holds. DCDI cites that deep sigmoidal flows are universal density approximators (Huang et al.) 2018), which in turn invokes the classic result that a multilayer perceptron with sigmoid activations is universal approximator (Cybenko) 1989p. AVICI does not discuss this aspect of their architecture, which is also based on sparse attention, but is otherwise quite different from ours. The following section proves by construction that our model is well-specified.
Axial attention architecture is well-specifted In the context of Section we show that three axial attention blocks (model depth) are sufficient to recover the skeleton and v-structures in O(N) width, and we require OL)) to propagate orientations along paths of length L. In the following section, we first formalize the notion of a neural network architecture's capacity to "implement" an algorithm. Then we prove Theorem by construction.
Definition A.11.
Let ʃ be a map from finite sets Q to F, and let & be a map from finite sets Qo to Fø. We say o implements f if there exists injection gin : Q + Q& and surjection Jout Fs F such that Vy Q, your Go (0(gi (q))) = [(g)-
(16)
Definition A.12. Let Q, Qo, Fo be finite sets. Let f be map from Q to F, and let + be finite set of maps {Ø: : Qs + F.). We say & has the capacily to implement f if and only if there exists at least one element Φ E & that implements f.
That is, a single model implements an algorithm f if for every input to f. the model outputs the corresponding output. A class of models parametrized by Φ E & has the capacity to implement an algorithm if there exists at least one & that implements f.
Theorem 3.1. Let G (V, be a directed acyelic graph with marimum degree d. For S C V, let E's denote the marginal estimate over S. Let Sa denote the superset that contains all subsets S SV V of size at most d. Given { {Fgg)ESSu2) a slack of L. ariol atlention blocks has the capacity to recover G G's skeleton and o-structures in ONN) width, and propagate orientations on paths of O(L) length.
Proof. We consider axial attention blocks with dot-product attention and omit layer normalization from our analysis, as is common in the Transformer universality literature Yun et al. (2020). Our inputs X E Rxx consist of d-dimension embeddings over R rows and C columns. Since our axial attention only operates over one dimension at a time, use X.. to denote a 1D sequence of length R, given a fixed column c, and X,, to denote 1D sequence of length C. given a fixed row r. A single axial attention layer (with one head) consists of two attention layers and feedforward network,
= X.,e + WoWv - [WWXX. AttDloow (X),)oo((X,e)
(17)
X + Att (X)
Attnrcal (x,.) X,, + WoWv a [(KXX,,) "W
(18)
X <- Attmeaa (X)
FFN(X)= X + W2 ReLU(W x x+b117) + +b2
(19)
25

---

Published in Transactions on Machine Learning Research (03/2025)
where Wo Rdxd Wv, WK, Wo BAxd W2 R W1 E R by R bj E R", and is the hidden layer size of the leedforward network. For concision, we have omitted the and e subscripts on the Ws, but the row and column attentions use different parameters. Any row or column attention can take on the identity mapping by setting Wo, Wv, WK, Wq to x d matrices of zcros.
A single axial attention block consists of two axial attention layers dE and 6p. connected messages (Section
3.3
F (F,
here- woe [hoit-1 m p
he. = (pp((1)t11))
B. Eit -.
where h denote the hidden representations of E and at layer l, and the outputs of the axial attention block are hel B.E
We construct a stack of L 2 3 axial attention blocks that implement Algorithm
Model inputs Consider edge estimate E E in a graph of size N. Let Ci, ey denote the endpoints of
(i,j). Outputs of the PC algorithm can be expressed by three endpoints: {Ø, }. }. A directed edge from
has endpoints ( ), the reversed edge i 4-] j has endpoints o), undirected edge has endpoints
) and the lack of any edge between i, has endpoints (0,
Let one-holx(i) denote the N-dimensional one-hot column vector where element is 1. We define the embedding (i,j) as a 2N + dimensional vector,
one-hotz (e:) one-hoty
9tn (In(E),()))) 11500
(20)
one-hot
one-hotx
To recover graph structures from E we simply read off the indices of non-zero entries (gout). We can haD to any NxxNNN matrix, as do not consider its values in this analysis and discard it during the first step.
Claim A.13. (Consistency) The oulputs of each step
1. are consistent with (20). and are equivariant to the ordering of nodes in edges.
For example, if (i,j) is oriented
.), then we expect (i,i) be oriented
Step 1: Undirected skeleton We use the first axial attention block to recover the undirected skeleton C'. We set all attentions to the identity, set Wa,1 E 2ddd to a d x d zeros matrix, stacked on top of a d x d identity matrix (discard p), and set FFNE to the identity (inputs are positive). This yields
P (z)
P
P
E+tp1 m.j
(21)
one-hot
one-hot ())
where (.. is the frequency that endpoint e; universal approximators of continuous functions Hornik et al.
within the subsets sampled. FFNs with 1 hidden layer are
1989 so we use FFN, to map
i<6
FFN, , (Xiu.))
=
0
i> 6, X =
(22)
Xin
otherwise,
26

---

Published in Transactions on Machine Learning Research (03/2025)
where i [2N + 6] indexes into the feature dimension, and index into the rows and columns. to remove edges not, present in C' from consideration:
This allows
m E.1 h
(i,)) $ C'
E.1 E.1 Wij Wi.
+ m
(23)
F.0 hj
otherwise.
This yields (i.j) E C" and only if heg # O. We satisfy A.13 since our inputs are valid PC algorithm outputs for which (2) Pe, (0).
The second and third axial attention blocks recover v-structures. We run the same procedure twice, once to capture v-structures that point towards the first node in an ordered pair, and one to capture v-structures that point towards the latter node.
Step 2: V-structures
We start with the first row attention over edge estimates, given a fixed subset t. We set the key and query attention matrices
001
0
1000
O
WK = k.
WQ k.
(24)
IN
IN
IN
IN
where & is a large constant, IN denotes the size N identity matrix, and all unmarked entries are Os.
Recall that a v-structure is pair of directed edges that share target node. We claim that two edges (i,j), (u,w) u) form v-structure in E', pointing towards i = u, if this inner product takes on the maximum value
(WKg (WohE.)a..)) 3.
(25)
Suppose both edges (1,5) and (u, v) still remain in C². There are two components to consider.
i u, then their shared node contributes +1 to the inner product (prior to scaling by 1:). If = v, then the inner product accrues -1.
Nodes that do not share the same endpoint contribute the inner product. Of edges that share one node, only endpoints that match at the starting node, or at the ending node contribute +1 to the inner product each. We provide some examples below.
(e,, ez)
(Ca, en)
contribution
note
0
no shared node wrong endpoints one correct endpoint v-structure
0
All edges with endpoints Ø were "removed" in step 1, resulting in an inner product of zero, since their node embeddings were set to zero. We set k lo some large constant (empirically, 122 1000 more than enough) to ensure that after softmax scaling, Je,e" > 0 only if e, form a v-structure.
Given ordered pair e = (i,j), let V; V denote the set of nodes that form a v-structure with with shared node Note that Vú excludes j itself, since setting of W K, Wq exclude edges that share both nodes. We set. Wy to the identity, and multiply by attention weights to obtain
(Wvh o)==(i))
one-hot (i)
(26)
aj binary, n (V,)
where binary (S) denotes the N-dimensional binary vector with ones elements in S, and the scaling factor
aj = (1/1V3l) 1[1V 0) E [0,
(27)
27

---

Published in Transactions on Machine Learning Research (03/2025)
results from softmax normalization.
We set
ON+6
Wo =
(28)
0.5 IN
to preserve the original endpoint values, and to distinguish between the edge's own node identity and newly rocognized v-structures. To summarize, the output of this row attention layer is
x ' yow (x = Xe WoWy e.'
which is equal to its imput h Bi plus additional positive values E (0, 0.5) in the last positions that indicate the presence of rustructures that exist the overall E'.
Our final step is "copy" newly assigned edge directions into all the edges. We the Φø column attention, FFNg and the &p attentions to the identity mapping. We also set w, to d x d zeros matrix, stacked on top of d x d identity matrix. This passes the output of the de row attention, aggregated over subsets, directly to F.N
For endpoint dimensions e [6]. we let FFN implement
10,0,1,0,1,07 - Xe,u 0 E<ENN Xi < 0.5
FFN (Xe, P.Np.2(Xe..)
(29)
0
otherwise.
Subtracting Xe,u, "erases" the original endpoints and replaces them with (>) after the update
E.1 1
E.1
+ mf7
The overall operation translates to checking whether any v-structure points towards i, and if so, assigning edge directions accordingly. For dimensions 1>0 6,
{ Xim,
Ni,u, 0.5 otherwise,
FFN (FN,,2(X,,..)
(30)
effectively erasing the stored v-structures from the representation and remaining consistent to 20
At this point, we have copied all v-structures once. However, our orientations are not necossarily symmetric. For example, given v-structure + j + k, our model orients edges (j.i) and (j, k), but not (i.j) or (k,j).
The simplest way to symmetrize these edges (for the writer and the reader) is to run another axial attention block, which we focus on v-structures that point towards the second node of a pair. The only changes are as follows.
For Wk and Wo, we swap columns 1-3 with 4-6, and columns 7 to N 6 with the last N columns.
()) sees the third and fourth blocks swapped.
Wo swaps the N N blocks that correspond to i and j's node embeddings.
FFNp33 sets the endpoint embedding to [0, 1,0,0,0,,111 Xe,a, if 7,.., N + 6 sum to a value between and 0.5.
The result is E3 with all v-structures oriented symmetrically, satisfying A.13
Step 3: Orientation propagation To propagate orientations, we would like identify cases (i.j), (i, e E', kj. & E' with shared node and corresponding endpoints (. We use Φg identify triangles, and op to identify edges (i.j), (i, k) E' with the desired endpoints, while ignoring triangles.
28

---

Published in Transactions on Machine Learning Research (03/2025)
Marginal layer The row attention in og fixes subset and varies the edge (i.j).
Given edge (i,j), we want to extract all (i, k) that share node i. We set the key and query attention matrices
to
0 0011011
WK. WQ=
(31)
IN
I1
We sot Wv to the identity obtain
(Wvh ()m)(
(32)
o (t)
On binary N (Vk)
where V& is the set of nodes l that share any edge with i. To distinguish between & and VE, again set W. to the same as in (28). Finally, we set FFNE to the identity and pass Λ directly to op. To summarize, we have h equal to its input, with values E (0, 0.5) in the last locations indicating 1-hop neighbors cach edge
Global layer Now would like to identify cases (i, k), (j, with corresponding endpoints We set the key and query atlention matrices
1001
-1
WK = k
We k
(33)
IN
IN
IN
IN
The key allows to check that endpoint is directed, and the query allows check that (i, k) exists in C', and does not already point elsewhere. After softmax normalization, for sufficiently large k, we obtain (is)).i > O if and only if (i, k) should be oriented (e, ), and the inner product attains the maximum possible value
((Wkh")
= 2.
(34)
We consider two components.
If the endpoints match our desired endpoints, we gain a +1 contribution to the inner product. match between the first nodes contributes +1. If the second node shares any overlap (either same edge, triangle), then negative value would be added to the overall inner product.
Therefore, we can only attain the maximal inner product if only one edge is directed, and if there exists no triangle.
We set W. to the same as in (28.. and we add he to the input of the next dg. To summarize, we have h" equal to its input, with values E (0, 0.5) in the last. N locations indicating incoming edges.
Orientation assignment Our final step is to assign new edge orientations. Let the column attention take on the identity mapping. For endpoint dimensions = 64, we let FFN, implement
[0. 0, 11T Xeaur
0 < E X < 0.5
FXN,(Xe...)
(35)
otherwise.
This translates checking whether any incoming edge points towards v, and if so, assigning the new edge direction accordingly. For dimensions 6,
X S 0.5
FFN FFN,(iiu.))
(36)
otherwise,
29

---

Published in Transactions on Machine Learning Research (03/2025)
effectively erasing the stored assignments from the representation. Thus, are left with h that conforms to the same format as the initial embedding in (20).
To symmetrize these edges, we run another axial attention block, in which we focus on paths that point towards the second node of a pair. The only changes are as follows.
For ΦE layer WK and We
31p
we swap IN and +/x.
For Up layer WK and We 33 we swap IN and +Ix-
Wo swaps the N x N blocks that correspond to i and j's node embeddings.
For FFN, 35 we let 21, instead.
The result is h with symmetric 1-hop orientation propagation, satisfying A.13 We may repeat this procedure k times to capture k-hop paths.
To summarize, we used axial attention block 1 to recover the undirected skeleton C", blocks 2-3 to identify and copy v-structures in E', and all subsequent - 3 layers to propagate orientations on paths up [(L -17-33) 3) length. Overall, this particular construction requires O(N) width for OL)) paths.
Final remarks Information theoretically, it should be possible to encode the same information in log N space, and achieve O(log N) width. For case of construction, we have allowed for wider networks than optimal. On the other hand, if we increase the width and encode each edge symmetrically, c.g. (e,, Cj, (, i), we can reduce the number of blocks by half, since we no longer need to run each operation twice. However, attention weights scale quadratically, so we opted for an asymmetric construction.
Finally, a strict limitation of our model is that it only considers 1D pairwise interactions. In the graph layer, cannot compare different edges' estimates at different times in a single step. In the feature layer, we camnot. compare (i,j) to (j in single step either. However, the graph layer does enable to compare all edges at once (sparsely), and the feature layer looks at a time-collapsed version of the whole graph. Therefore, though we opted for this design for computational efficiency, we have shown that it is able to capture significant graph reasoning.
A.4 Robustness and stability
We discuss the notion of stability informally, in the context of Spirtes et al. (2001). There are two cases in which our framework may receive erroncous inputs: low/nois data settings, and functionally misspocified situations. We consider our framework's 's empirical robustness to these cases, terms of recovering the skeleton and orienting edges.
In the case of noisy data, edges may be erroneously added, removed, or misdirected from marginal estimates E'. Our framework provides two avenues to mitigating such noise.
1. We observe that global statistics can be estimated reliably in low data scenarios. For example, Figure D suggests that 300 examples suffice to provide a robust estimate over 100 variables in our synthetic settings. Therefore, even if the marginal estimates are erroneous, the neural network can learn the skeleton from the global statistics.
2. Most classical causal discovery algorithms are not stable with respect to edge orientation assignment. That is, an error in a single edge may propagate throughout the graph. Empirically, we observe that the majority vote of GIES achieves reasonable accuracy even without any training, while FcI suffers in this assessment (Table 10). However both SEA (GIES) and SEA (Fc1) achieve high edge accuracy. Therefore, while the underlying algorithms may not be stable with respect to edge ation, our pretrained aggregator seems to be robust.
It is also possible that our global statistics and marginal estimates make misspecified assumptions regarding the data generating mechanisms. The degree of misspecification vary case by case, it is hard to
30

---

Published in Transactions on Machine Learning Research (03/2025)
provide any broad guarantees about the performance of our algorithm, in general. However, can make the following observation.
If two variables are independent, X; TT X,, they are independent, e.g. under linear Gaussian assumptions. If X, exhibit more complex functional dependencies, they may be erroncously deemed independent. Therefore, any systematic errors are necessarily one-sided, and the model can learn to recover the connectivity based global statistics.
B Experimental details
B.1 Synthetic data generation
Synthetic datasets were generated using code from DCDI (Brouillar el al. 2020 which exlended the Causal Discovery Toolkit data generators to interventional data (Kalainathan et al. 2020).
We considered the following causal mechanisms. Let y be the node in question, let X be its parents, let E be independent noise variable (details below), and let W be randomly initialized weight matrices.
Linear: = XW + E.
Polynomial: y Wo + XW1 X2W, + X
Sigmoid: y= - Wi sigmoid(X,) + xE
Randomly initialized neural network (NN): y = Tanh((X, E) Wont
Randomly initialized neural network, additive (NN additive): y Tanh(XW... + E
Root causal mechanisms, noise variables, and interventional distributions maintained the DCDI defaults.
Root causal mechanisms were set to Uniform(-2,77
Noise was set to E' ~ 0.4 NNO,,²)) where o² ~ Uniform(1,2).
Interventions were applied to all nodes (one at time) by setting their causal mechanisms to N( 1).
Ablation datasets with N 100 nodes contained 100,000 points each (same as N = 100). We set random
seeds for each dataset using the hash of the output filename.
B.2 Related work and baselines
We considered the following basclines. All basclines were rum using official implementations published by the authors.
AVICI Lorch et al. 2022] is the most similar method to this work, though there are significant differences in both the causal discovery strategy and the implementation. Both works simulate diverse datasets for training, which differ in graph topology, causal mechanism, and type of exogenous noise; and both are attention-based architectures (Vaswani et al. 2017). The primary difference that AVICI operates over raw data, while we operate over summary statisties. While access to raw data may allow for richer modeling of relationships, it also increases the computational cost large datasets. In fact, AvIcr's model complexity scales quadratically as the number of samples and quadratically as the number of nodes. Our model does not explicitly depend on the number of data samples, but scales cubically as the number of nodes, as we atten over richer, pairwise features.
Both AVICI and SEA are equivariant to the ordering of nodes. AVICI additionally invariant the ordering of samples. Due to the two-track design our aggregator, we also include message passing operations between the two input types (global statistios, marginal features), while AVIcI directly stacks Transformer layers.
31

---

Published in Transactions on Machine Learning Research (03/2025)
Finally, AVICI explicitly regularizes for acyelicity. We follow the ENCO (Lippe et al. 2022 formulation (edges i - j and j + cannot co-exist), since it is easier to optimize. Empirically, our predictions are still 99% acyelic (Table 18).
AVICI was run on all test datasets using the authors' pretrained scm-vO model, recommended for "arbitrary real-valued data." Note that this model is different from the models described in their paper (denoted AvICI-L and AvicI-R), it was trained on all of their synthetic data, including test sets. We sampled 1000 observations per dataset uniformly at random, with their respective interventions (the maximum number of synthetic samples used in their original paper), except for Sachs, which used the entire dataset (as in their paper). Though the authors provided separate weights for synthetic mRNA data, we were unable to use it since we did not preserve the raw gene counts in our simulated mRNA datasets.
Dedi (Brouilladd et al.) 2020 was trained on each of the N = 10, 20 datasets using their published hyperparamcters. Wc denote the Gaussian and Deep Sigmoidal Flow versions as DCDI-G and DCDI-DSF respectively. DCDI could not scale to graphs with N 100 due to memory constraints (did not fit on a 32GB V100 GPU).
DCD-FG [Lopez ct al. 2022 was trained on all of the test datasets using their published hyperparameters. We set the number of factors to 5, 10, 20 for each of N = 10, 20, 100, based on their ablation studies. Due to numerical instability on N= 100, we clamped augmented Lagrangian multipliers and 7y to 10 and stopped training if elements of the learned adjacency matrix reached NaN values. After discussion with the authors, we also tried adjusting the y multiplier from to 1.1, but the model did not converge within 48 hours.
DECI [Geffner et al. 2024 was trained all of the test datasets using their published hyperparameters. However, on all N 100 cases, the model failed to produce any meaningful results (adjacency matrices nearly all remained Os with AUCs of 0.5). Thus, we only report results on N = 10, 20.
DiffAN (Sanchez et al., 2023 was trained on the each of the N = 10, 20 datasets using their published hyperparameters. The authors write that "most hyperparameters are hard-coded into [the] constructor of the DIPFAN class and we verified they work across a wide set of datasets." " We used the original, non- approximation version of their algorithm by maintaining residue=True in their codebase. We were unable to consistently run DIFFAN with both R and GPU support within Docker container, and the authors did not respond to questions regarding reproducibility, so all models were trained on the CPU only. We observed approximately a 10x speedup in the < 5 cases that were able to complete running on the GPU.
InvCov computes inverse covariance over 1000 examples. This does not orient edges, but it is a strong connectivity baseline. We discretize based on ground truth positive rale.
Corr and D-Corr are computed similarly, using global correlation and distance correlation, respectively (See C.1 for details).
B.3 Metrics
Throughout our evaluation, we compute metrics with respect to the ground truth graph. This means that in the observational setting, the "oracle" value of each metric will vary depending on the size of the equivalence class (o.g. if multiple graphs are observationally equivalent, the expected SHD is > O; see Scction 5.3 for more analysis). While it is more correct lo evaluate (observational) models with respect to the CPDAG that is implied by the predicted DAG, there were several reasons we also chose to evaluate with respect to the ground truth DAG.
This practice is quite common among continuous causal discovery algorithms, regardless of whether interventional data are available (Lorch et. al. 2022 Brouillard et al. 2020 Hägele et al.] 2023) or not (Sanchez et al., 2023). This may be primarily because their predicted graphs contain continuous probabilities, and metrics that reflect uncertainty (mAP, AUROC) are difficult to translate to the equivalence class.
It has been reported that common simulators produce "identifiable" linear Gaussian datasots, if the data unstandardized (hence the VarSort baseline, Reisach et al. (2021)) This would mean
32

---

Published in Transactions on Machine Learning Research (03/2025)
that the MEC is actually smaller than expected. While it is possible to normalize all the data, it is unknown whether there are additional artifacts of the data generating process that may influence the (empirical) identifiability of the graphs, and thus, the "true" size of the equivalence class.
Our models and baselines evaluated on the same graphs, but not all can incorporate interventional information. For comparability, evaluated all predictions against the same ground truth, while noting that the equivalence class is larger in the observational case.
FcI, GIES, and VARSORT don't necessarily produce acyelic graphs, due to bootsrapping, In Ap- pendix C we also include the symmetric summary statistics as baselines, to quantify the "cost" of not directing edges, and of not using marginal estimates. In these cases, we cannot compute a corresponding CPDAG.
Practically, was somewhat expensive to compute the CPDAG, especially for larger dense graphs (N = 100, E 400 and beyond).
B.4 Training and hardware details
Hyperparameters and architectural choices were selected by training the model on 20% of the the training and validation data for approximately 50k steps (several hours). We considered the following parameters in soquence.
learned positional embedding vs. sinusoidal positional embedding
number layers x number heads: {4,8) x {4,8}
learning rate 1j (1e 4,5e 5, 5.1e-5
For our final model, we selected learned positional embeddings, layers, 8 heads, and learning rate 1/ le 4.
The models were trained across 2 NVIDIA RTX A6000 GPUs and 60 CPU cores. We used the GPU exclusively for rumning the aggregator, and retained all classical algorithm exccution on the CPUs (during data loading). The lotal pretraining time took approximately 14 hours for the final FCI model and 16 hours for the final GIES model.
For finetuning, we used rank r adapters on the axial attention model's key, query, and foodforward weights (Hu et al. 2022). We trained until convergence on the validation set (no improvement for 100 epochs), which took 4-6 hours with 40 CPUs and around 10 hours with 20 CPUs. We used single NVIDIA RTX A6000 GPU, but the bottleneck was CPU availability.
For the scope of this paper, our models and datasets are fairly small. We did not scale further due to hardware constraints. Our primary bottlenccks to scaling up lay in availability of CPU cores and networking speed across nodes, rather than GPU memory or utilization. The optimal CPU:GPU ratio for SEA ranges from 20:1 to 40:1.
We are able to inference comfortably over N 500 graphs with T 500 subsets of k = 5 nodes cach, single 32GB V100 GPU. For runtime analysis, used batch size of 1, with 1 data worker per datasel. Our runtime could be further improved if we amortized the GPU utilization across batches.
B.5 Choice of classical causal discovery algorithm
For training, selected FCI (Spirtes et al. 1995) as the underlying discovery algorithm in the observational setting over GES [Chickering, 2002), GRaSP [Lam et al.) 2022), and LiNGAM (Shimizu et al., 2006) due to its speed and superior downstream performance. We hypothesize this may be due to its richer
output (ancestral graph) providing more signal to the Transformer model. We also tried Causal Additive Models (Bühlmann et al. 2014), but its runtime was too slow for consistent GPU utilization. Observational algorithm implementations were provided by the causal-learn library (Zheng al., 2024). The code for running these alternative classical algorithms is available in our codebase.
33

---

Published in Transactions on Machine Learning Research (03/2025)
We selected GIES as the discovery algorithm in the interventional setting because efficient Python implementation was readily available at
We tried incorporating implementations from the Causal Discovery Toolbox via a Docker image (Kalainathan et al. 2020), but there was excessive overhead associated with calling an subroutine and reading/ the inputs/resul results from disk.
Finally, considered other independence tests for richer characterization, such as kernel-based methods. However, due to speed, chose to remain with the default Fisherz conditional independence test for FCI, and BIC for GTES (Schwarz 1978).
B.6 Sampling procedure
Selection scores: We consider three strategies for computing selection scores We include empirical comparison of these strategies in Table 9
Random selection: a is an N N matrix of ones.
Global-statistic--ased selection: a = Uncertainty-basee selection: = (Ee), where denotes the information entropy
04
E p(e) logr(e).
(37)
eE(0,1,2)
Let ci.j be the number of times edge (i,j) was selected in S S.-1, and let a' = a/ V We consider two strategies for selecting S. based on dt--
Greedy selection: Throughout our experiments, we used greedy algorithm for subset selection. We
normalize probabilities to before constructing each Categorical. Initialize
S (i: Categorical(cf....dd
(38)
where of = Ejiie oij While |S,| < k, update
S S, U (j : Categorical(af,ss ... a(v,s))
(39)
where
{ Eies ais j St
aj.s,
(40)
otherwise.
Subsct sclection: We also considered the following subset-level selection procedure, and observed minor performance gain for significantly longer runtime (linear program takes around 1 second per batch). Therefore, opted for the greedy method instead.
We solve the following integer lincar program to select subset S₁ of size k that maximizes Eies, of.y- Let 4 E (0, }} denote the selection of node i, and let 6i,J E {0, 1} denote the selection of edge (i. Our objective is to
maximize subject to
Σng dij fi.j
Σiv; =
subset size
G, 2 14 + Vv - node-edge consistency
6, $ v
G,j Vj.
E [0,1]
4, {0.
for i,j E VVVV, S, is the set of non-zero indices in v.
The final algorithm used the greedy selection strategy, with the first half of batches sampled according global statistics, and the latter half sampled randomly, with visit counts shared. This strategy selected
34

---

Published in Transactions on Machine Learning Research (03/2025)
Table 9: Comparison between heuristics-based sampler (random and inverse covariance) vs. model confidence- based sampler. The suffix -L indicates the greedy confidence Sonfidence-based sampler. Each setting encompasses 5 distinct Erdös-Rényi graphs. The symbol i indicates that SEA was not pretrained on this setting. Bold indicates best of all models considered (including baselines not pictured).
Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial'
mAPf OAf shdj
mAPf OAt shdd
mAPt OAt shid mAPf OA
shi|
mAPf OAt shi]
SEA-F
0.97 0.92
1.6
0.95
0.92
2.4
0.92
0.94
2.8
0.83
0.76
3.7
0.69
0.71
6.7
SEA-G
0.99 0.94
1.2
0.94
0.88
2.6
0.91
0.93
3.2
0.85
0.84
4.0
0.70
0.79
5.8
10 10
SEA-F-L
0.97
0.93
1.0
0.95
0.87
2.4
0.92
0.98
3.4
0.84
0.77
3.9
0.70
0.79
5.8
SEA-G-L
0.98 0.93
1.4
0.91 0.94
2.8
0.91
0.94
4.0
0.88
0.84
3.6
0.70
0.80
5.8
SEA-F
0.90
0.87
14.4
0.91
0.94
11.2
0.87
0.86
16.0
0.81
0.85
22.7
0.81
0.92
33.4
SEA-G
0.94 0.91
12.8
0.91
0.95
10.4
0.89
0.89
17.2
0.81
0.87
24.5
0.89
0.93
29.5
10 40
SEA-F-L
0.91
0.90
15.6
0.91
0.92
15.8
0.88
0.86
14.2
0.81
0.84
23.2
0.82
0.93
33.8
SEA-G-L
0.93
0.91
13.4
0.91
0.93
10.4
0.88
0.85
16.2
0.79
0.83
25.5
0.90
0.94
28.3
SEA-F
0.97
0.92
3.2
0.94
0.97
3.2
0.84
0.93
7.2
0.84
0.85
7.6
0.71
0.80
10.2
SEA-G
0.97 0.89
3.0
0.94
0.95
3.4
0.83
0.94
7.8
0.84
0.83
8.1
0.69
0.78
10.1
SEA-F-L
0.97
0.92
2.8
0.93
0.95
3.8
0.85
0.94
6.8
0.85
0.85
7.5
0.67
0.78
9.9
SEA-G-L
0.97
0.90
2.6
0.94
0.98
3.4
0.83
0.97
7.0
0.84
0.84
7.9
0.67
0.79
10.6
SEA-F
0.86 0.93
29.6
0.55
0.90
73.6
0.72
0.93
51.8
0.77
0.85
12.8
0.61
0.89
61.8
SEA-G
0.89 0.92
26.8
0.58
0.88
71.4
0.73
0.92
50.6
0.76
0.84
45.0
0.65
0.89
60.1
20 80
SEA-F-L
0.86
0.92
32.0
0.55
0.90
74.0
0.74
0.93
49.2
0.76
0.87
41.8
0.59
0.88
62.3
SEA-G-L
0.89
0.92 28.4
0.58
0.89 71.6
0.75
0.92 49.4
0.75
0.85
45.7
0.65
0.88 60.6
heuristically, and we did not observe significant improvements or drops in performance when switching other strategies (e.g. all greedy Maatisiissbbsse greedy uncertainty-based, linear program uncertainty-based, etc.)
Table 9 compares the heuristics-baset greedy sampler (inverse covariance + random) with the model uncertainty-based greedy sampler. Runtimes are plotted in Figure G The latter was run on CPU only, since it was non-trivial to access the GPU within a PyTorch data loader. We ran a forward pass lo obtain an updated selection score every 10 batches, so this accrued over 10 times the number of forward passes, all on CPU. With proper engincering, this model-basod sampler is expected to be much more efficient than reported. Still, it is faster than nearly all baselines.
Average ruatimes on syathetic datasots
N-10
N-20
10
N-100
Wall time (s)
10P
10-1
Ours (FCI)
Our (FCTLL)
Ours (GIES)
Ours (GIES-L)
Figure 6: Runtime for heuristics- based greedy sampler vs. model uncertainty-basse greedy sampler (suffix -L). For sampling, the model was run on CPU only, due to the difficulty of invoking GPU in the PyTorch data sampler.
35

---

Published in Transactions on Machine Learning Research (03/2025)
C Additional analyses
C.1 Choice of global statistic
We selected inverse covariance as our global feature due to its case of computation and its relationship to partial correlation. For context, we also provide the performance analysis of several alternatives. Tables
[1 and [2 compare the results of different graph-level statistics on our synthetic datasets. Discretization thresholds for SHD were obtained by computing the pth quantile of the computed values, where = 1- -E/NN. This is not entirely fair, no other baseline receives the same calibration, but these ablation studies only seek to compare state-of-the-art causal discovery methods with the "best" possible (oracle) statistical alternatives.
Corr refers to global correlation,
E (X;X,) ( E (X,) E (X?) EE (X.)² (X3) EX2)²
PL
(41)
D-Corr refers to distance correlation, computed between all pairs of variables. Distance correlation captures both linear and non-linear dependencies, and D-CORR(X, X,) = 0 if and only if X; IL Xj. Please refer to [Sz "Szeekely et al. (2007) for the full derivation. Despite its power to capture non-linear dependencies, we opted not to use D-CORR because it is quite slow to compute between all pairs of variables.
InvCov refers to inverse covariance, computed globally,
== E - ( (X)))XX(XX)X)
((X))))
(42)
For graphs < 100, inverse covariance was computed directly using NumPy. For graphs N 2 100, inverse covariance was computed using Ledoit-Wolf shrinkage at inference time Ledoit & Wolf (2004). Unfortunately we only realized this after training our models, so swapping to Ledoit- Wolf leads some distribution shift (and drop in performance) SEA results for large graphs.
C.2 Visualization of denoising statistical features
Figure 7 illustrates the input features and our predictions for a N = 10, the inputs, our method is able to produce a much cleaner graph. GIES may orient edges the wrong direction in some of the bootstrap samples, but SEA can generally identify the right direction. This example also illustrates how SEA can triangulate between GIES and the global statistic, to avoid naive predictions of edges wherever the global statistic has high value. In the example marked in tcal, though the global statistic has a moderate value, the edge is absent. from a substantial number of estimates (if we consider both directions). This reflects the theory in A.2 that edges absent from some marginal estimate should be absent from the final skeleton. SEA was able to (rather confidently) reject the edge in the final graph.
linear graph. Compared
In contrast to GIES, Fci tends to identify much fewer edges (Figure 8) Note that this does not hurt performance metrics, as continuous metrics consider "sliding" thresholds, while discrete metrics are computed with respect to the true edge rate (for the oracle FcI baseline). However, it does highlight that FcI has much lower signal-to-noise ratio when used alone.
C.3 Results on simulated mRNA data
We generated mRNA data using the SERGIO simulator Dibaeinia & Sinha (2020). We sampled datasets with the Hill coefficient set to (0.25, 0.5, 2, 4} for training, and 2 for testing (2 was default). We set the decay rate to the default 0.8, and the noise parameter to the default of 1.0. We sampled 400 graphs for each of N = (10, 20} and = (N,2N).
These data distributions are quite different from typical synthetic datasets, as they simulate stcady-state measurements and the data are lower bounded 0 (gene counts). Thus, we trained a separate model on these data using the SEA (Fei) architecture. Table 13 shows that SEA performs best across the board.
36

---

Published in Transactions on Machine Learning Research (03/2025)
Inv. Cov.
GIES bootstrap
Ours (GIES)
True Graph
Figure 7: Left to right: magnitude value of global features with diagonals zeroed; GIES all variables via non-parametric bootstrapping (frequency of edge); SEA predictions; and the ground truth. SEA is able to denoise the inpul features much belter than naive aggregation schemes.
Inv. Cov.
FCI (mean)
FCI (any)
Ours (FCI)
True Graph
Figure 8: Left to right: magnitude value of global features with diagonals zeroed; FcI run over all variables via non-parametric bootstrapping, aggregation via frequency of edge vs. binarized (appear or not); SEA predictions; and the ground truth. Fei produces sparser outputs than GIES, so Foi alone is less able to distinguish between noisy and genuine edges.
C.4 Results and ablation studies synthetic data
For completeness, include additional results and analysis the synthetic datasets. Table 14 contains evaluations in the observational setting, with metrics computed over the inferred CPDAG, rather than the predicted DAG. We find that the SHD varies slightly when evaluating DAGS converted CPDAGs, but the difference is minimal, compared to the overall gaps in performance between methods.
Tables and 20 compare all baselines across all metrics and graph sizes on Frdös-Rényi graphs. Tables 21 and include the same evaluation on scale-free graphs. Tables 23 and 24 assess N 100 graphs.
Table ablates the contribution of the global and marginal features by setting their hidden representations to zero. Note that model has never seen this type of input during training, drops in performance may be conflated with input distributions ional shift. Overall, removing the joint statistics (h? 0) leads to higher performance drop than removing the marginal estimates (h + 0). However, the gap between these ablation studies and our final performance may be quite large in some cases, so both inputs are important to the prediction.
Table 18 shows that despite omitting the DAG constraint, we find that our predicted graphs (test split) are nearly all acyelic, with naive discretization threshold of 0.5. Unlike Lippe al. (2022). which also omits the acyclicity constraint during training but optionally enforces it at inference time, we do not require any post-proccssing to achieve high performance. Empirically, we found existing DAG constraints to be unstable (Lagrangian) and slow to optimize (Zheng et al. 2018 Brouillard et al., 2020). DAG behavior would not emerge until late in training, when the regularization term of 1e-8 scale or smaller.
Alternatively, we could quantify the raw information content provided by these two features through the INvCov, Fci* and GIES baselines (Tables 19 20 21 22). Overall, INvCOV and Fcr* are comparable to worse-performing baselines. GIes performs very well, sometimes approaching the strongest baselines.
37

---

Published in Transactions on Machine Learning Research (03/2025)
Table 10: Synthetic experiments, edge direction accuracy (higher is better). All standard deviations were
within 0.2. The symbol indicates that SEA was not pretrained this setting.
N
E' Model
Linear
NN add
NN
Sig.
Poly.
DCDI-G
0.74
0.80
0.85
0.41
0.44
DCDI-DsB
0.79
0.62
0.68
0.38
0.39
DCD-FG
0.50
0.47
0.70
0.43
0.54
DIFFAN
0.61
0.55
0.26
0.53
0.47
DECI
0.50
0.43
0.62
0.63
0.75
10
10
AVICI
0.80
0.92
0.83
0.81
0.75
Fcr*
0.52
0.43
0.41
0.55
0.40
GIES*
0.76
0.49
0.69
0.67
0.63
SEA (Fei)
0.92
0.92
0.94
0.76
SEA (GrEs)
0.94
0.88
0.93
0.84
0.79
DCDI-G
0.47
0.43
0.82
0.40
0.24
DODI-DSF
0.50
0.49
0.78
0.41
0.28
DCD-FG FG
0.58
0.65
0.75
0.62
0.48
DIPFAN
0.46
0.28
0.36
0.45
0.21
DECI
0.30
0.47
0.35
0.48
0.57
20
80
AVICI
0.57
0.67
0.74
0.63
0.62
Fcr*
0.19
0.19
0.22
0.33
0.23
GIES
0.56
0.73
0.59
0.62
0.61
SEA (Fet)
0.93
0.90
0.93
0.85
0.89
SEA (Gms)
0.92
0.88
0.92
0.84
0.89
DCD-FG
0.46
0.60
0.70
0.67
0.53
AvICI
0.61
0.68
0.72
0.54
0.42
100 400
SEA (Fc1)
0.93
0.90
0.91
0.82
SEA (GIES)
0.94
0.91
0.92
0.87
0.84
Table 11: Comparison of global statistics (continuous metrics). All standard deviations within 0.1.
Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial
mAP AUC t mAP + AUC
mAP
AUC 1 mAP t AUC mAP t AUC f
CORR
0.45
0.87
0.41
0.86
0.41
0.85
0.46
0.86
0.45
0.85
D-CORR
0.42
0.86
0.41
0.87
0.40
0.87
0.43
0.86
0.45
0.89
INvCov
0.49
0.87
0.45
0.86
0.36
0.81
0.44
0.86
0.45
0.83
CORR
0.47
0.53
0.47
0.52
0.46
0.52
0.48
0.53
0.48
0.54
D-CORR
0.46
0.53
0.46
0.51
0.46
0.54
0.48
0.53
0.47
0.54
INvCov
0.50
0.57
0.48
0.52
0.47
0.53
0.47
0.50
0.48
0.52
CORR
0.42
0.99
0.25
0.94
0.25
0.93
0.42
0.98
0.35
0.91
100 100
D-CORR
0.41
0.99
0.25
0.96
0.26
0.96
0.41
0.98
0.37
0.94
INvCov
0.40
0.99
0.22
0.94
0.16
0.87
0.40
0.97
0.36
0.90
CORR
0.19
0.80
0.10
0.63
0.14
0.72
0.27
0.20
0.72
100 400
D-CORR
0.19
0.80
0.10
0.63
0.14
0.75
0.26
0.84
0.21
0.74
INVCOV
0.25
0.91
0.09
0.62
0.14
0.77
0.27
0.86
0.20
0.67
However, there remains a large gap in performance between these ablations and our method, highlighting the value learning non non-linear transformations of these inputs.
Table 16 and Figure show that the current implementations of SEA can generalize to graphs up 4x larger than those seen during training. During training, did not. initially anticipate testing on much larger
38

---

Published in Transactions on Machine Learning Research (03/2025)
Table 12: Comparison of global statistics (SHD). Discretization thresholds for SHD were obtained by computing the pt quantile of the computed values, where p 1 (E/N).
N E Model
Linear
NN add.
NN non-add.
Sigmoid Polynomial
Conn
10.6=2.8
10.2+4.6
12.0+1.9
11.1=43
9.942.
10
10
D-CORR
10.4=2.6
9.844.7
12.2+2.6
10.8=3.3
10.2:3.2
INVCOV
11.0=2.8
11.415.5
13.6+2.9
11.424.1
10.943.5
CORR 39.2=2.4
3..021.8
38.220.7
38.823.3
38. 22.0
10
40
D-CORR
38.8=2.0
38.821.5
37.020.6
38.9=3.2
38.0=2.0
INVCOV
35.8=2.3
39.2
37.642.77
40, 20.2
38.4+1.
CORR 113.0=4.9
132.2
144.625.2
106.5
110.3=6.1
100 100
D-CORR 113.8=5.3
133.2+17.9
141.2+6.7
108.5+11.9
109.5=5.7
INVCOV 124.4=8.1
130.0
158.816.2
112.3:11.8
106.3:4.6
CORR D-CORR 578.2124.7 665.4
580.4+24.6
666.0#13.5
626.2
516.5418.6
562.5+200.1
100
400
626.6
522.3117.6
557 557.2:120.4
INvCOV 557.0+11.7 165.8
639.0
514.7123.1
539.4+18.
Table 13: Causal discovery results on simulated mRNA data. Each setting encompasses 5 distinct scale-free graphs. Data were generated via SERGIO Dibaeinia & Sinha (2020).
N
E
Model
mAP f
AUC f
SHD 4
OA t
DCDI-G
0.48
0.73
16,1+3.3
0.59
DCDI-DsF
0.63
0.84
18.5+2.7
0.79+0.2
10
10
DCD-FG
0.59
0.82
81.020.0
0.79
AvICI
0.5820.2
0.85
6.444.7
0.72+0.2
SEA (Fct)
0.92
0.98
1.91:2.0
0.92
DCDI-G
0.32
0.57
26.2+1.3
0.47+0..2
DCDI-DsF
0.44
0.64
25.7+1.3
0.63
10
20
DCD-FG
0.43
0.69
73.0=0.0
0.67
AVICI
0.22=a.1
0.44
16.8+1.5
0.27+0.3
SEA (Fct) 0.76
0.90
8.8
0.85
DeDI-G
0.48
0.8620.1
37.3+2.8
0.65
DeDI-DsF
0.4520.1
0.9220.0
51.9215.8
0.81 10.1
DCD-FG
0.34
0.8710.0
36110
0.66
AVICI
0.32
0.78
18.7144.
0.66:1.2
SEA (Fcr) 0.54
0.94
16.6+3.3
0.83+0.1
DCDI-G
0.31
0.65
54.7++2.7
0.49
DCDI-DsF
0.4020.1
0.71
54.624.4
0.63+0.1
20
40
DCD-FG
0.3620.1
0.77
31340
0.67
AVICI
0.17
0.54
37.1
0.4640.1
SEA (Fct)
0.50
0.85
31.414.9
0.7810.1
graphs. As result, there are two minor issues with the current implementation with respect to scaling, First, we set an insufficient maximum subset positional embodding size of 500, so it was impossible to encode more subsets. Second, we did not sample random starting subset indices to ensure that higher-order embeddings are updated equally. Since we never sampled up to 500 subsets during training, these higher-order embeddings were essentially random. We anticipate that incrcasing the limit on the number of subsets and ensuring that all embeddings are sufficiently learned will improve the generalization capacity on larger graphs. Nonetheless, current model already obtains reasonable performance larger graphs, out of the box.
39

---

Published in Transactions on Machine Learning Research (03/2025)
Table 14: SHD on synthetic graphs, observational setting, between predicted vs. true DAG and inferred vs. true CPDAG. Mean//td over distinct Prdös-Réént graphs. f indicates o.o.d. setting. indicates non- con-maraettric bootstrapping. DAG results from Table
N Model
Linear
NN add.
Sigmoid"
Polynomial'
DAG 4 CPDAG
DAG +
CPDAG DAG 4 CPDAG DAG
CPDAG
DIFFAN
14.0
16.4
13.622.9
13.623.4
12.0=.0
12.9
15.0+.1
16.3+.6
VARSORT*
6.0+.4
7.0+.3
4.0+.8
4.4.22
7.6=.5
8.04.1
9.3+.5
9.5+.8
10 10
FCr*
10.0
11.0a.8
8.2
8.41.8
9.1=.s
10.3
10.0
10.1
SEA (Fer)
1.0
1.42.7
3.0=.2
2.64.3
3.92.9
3.54.3
6.1
6.4+.1
DIFFAN
40.2
41.8
38.643.1
39.
19.2
21.61.2
49.7.1..6
52.7=5.0
VARSORT*
10.0..4
11.0=.0
6.62.7
7.21.5
16.1=.7
16.81.2
17.11.1
17.71.0
20 20
Fcr
19.0
21.4+33
17.4#.2
18.0+.6
18.5=.5
20.6+.4
18.9+.3
19.1
SEA (Fei)
3.2#.6
2.6
5.04.8
4.8
6.7=a1
6.2+.1
9.8+.2
11.0+.2
Table 15: mAP on synthetic graphs, between predicted vs. true DAG and undiroctod skeletons E Mcan/std over 5 distinet Frdös-Rénny graphs. indicates o.o.d. setting. indicates non-parametrie boolstrapping. DAG results from Table
N
Model
Linear
NN add.
Sigmoid
Polynomial
DAG
DAG 1
E" f
DAG f
E
DAG
E*
DCDI-G
0.74+
0.88+.10
0.79+.12
0.85+.12
0.46+.24
0.64
0.41+.13
0.58+
DeDI-DsP
0.82
0.93
0.574.24
0.88+.11
0.38
0.63
0.61 0.29
DIFFAN
0.252.06
0.54
0.32
0.624.19
1.
0.63
0.20=.os .8
0.50
AVICI 0.45
0.93 .6
0.81
0.981.03
0.52..16
0.861.13
0.31 .08
0.75..13
10 10
VARSORT*
0.70
0.88
0.76...1
0.90
0.52..24
0.88
0.40
0.80+.09
Fer*
0.52
0.70 .20
0.38
0.69
0.56+.16
0.75
0.61 0.41
GIES
0.81
1.0#..0 o0
0.61
0.97+.06
0.70+.14
0.98
0.61
0.89
SEA (Fer)
0.98
1.0
0.881.09
0.961.00
0.83
0.99
0.62 .09
0.87:.06
SEA (GrEs)
0.99
1.0
0.941.00
0.97#.05
0.85
0.98
0.70=.11
0.862.05
DCDI-G
0.59
0.89..04 .04
0.78..07
0.931.07
0.364..0 00
0.84
0.42
0.58
DCDI-DsP
0.66:.16
0.911.04
0.694.18
0.931..10
0.371.04
0.861.05
0.26..06 .08
0.35:1.15
DIFFAN
0.19#.09
0.47#.a6
0.164.10
0.44+.15
0.29
0.67
0.09
0.32+
AvIcI
0.48
0.84
0.59+.09
0.85+.07
0.42+.13
0.78+.10
0.24# .08
0.62+t.08
20
20
VARSORT*
0.81
0.93 04
0.81
0.88
0.50
0.88 09
0.33
0.762.11
Far*
0.66#.07
0.79
0.42
0.54+.22
0.56+.08 08
0.67
0.41
0.61
GIES*
0.84+.0s
0.99 .00
0.79+.07
0.944.04
0.71+.10
0.95+.03
0.62 .0
0.85+.09
SEA (Fet)
0.96+.0d
1.0+
0.914.04
0.954.u3
0.85+.09
0.97+.02
0.69
0.86++.9
SRA (GreS)
0.97
1.02.00 oo
0.944.03
0.96..03
0.84
0.95 os
0.69
0.82
Finally, we note that AVICI scales very poorly to graphs significantly beyond the scope of their training set. For example, N = 100 is only 2x their largest training graphs, but the performance already drops dramatically.
Figure g depicts the model runtimes. SEA continues to run quickly on much larger graphs, while AVICI runtimes increase significantly with graph size.
DCDI learns generative model over each dataset, and its more powerful, deep sigmoidal flow variant seems to perform well in some (but not all) of these harder cases.
C.5 Results on real datasets
The Sachs flow cytometry dataset (Sachs al.) 2005) phospholipids at the single cell level. We use the subset proposed by Wang et al. (2077) The ground truth
measured the expression of phosphoproteins and
40

---

Published in Transactions on Machine Learning Research (03/2025)
Table 16: Scaling to synthetic graphs, larger than those seen in training. Each setting encompasses 5 distinct Erdös-Rényi graphs. All SEA runs in this table used T 500 subsets of nodes, with = 500 examples per batch. For AVICI, we took M 2000 samples per dataset (higher than maximum analyzed in their paper), since it performed better than M 1000. Here, the mean AUC values artificially high due to the high negative rates, actual edges scale linearly as N, while the number of possible edges scales quadratically.
N Model
Linear, E = N
Linear, E = 4N
mAP
AUC f
SHD 4
OA f
mAP t
AUC 1 SHD +
OA f
INvCov
0.43+0.0
0.99
117#7
0.30
0.93
512+11
Conn
0.42+0.0
0.99
113#s
0.19
0.80=a.0
579+25
100
AVIcI
0.03
0.43
109
0.49
0.11=0.0
0.55
894
0.5820.0
SEA (Fe1)
0.97
1.00+t.a
11.6+4.3
0.93+aau
0.88#a.
0.98
129
0.94+v.0
SEA (GIES)
0.97
1.00
12.8+4.7
0.91+a.0
0.91+a.
0.99
105+s
0.95
INvCov
0.45+0.0
1.00+o.a
218+11
0.33=a.0
0.96=a.0
1000
Conn 0.42+0.0
0.99
223+*
0.18=a.0
0.862a.0
1184
200
AVICI
0.00
0.3610.1
107::10
0.41
0.05=0.0
0.53
827
0.54
SEA (Fe1)
0.91
1.00
49.9
0.87+a.0
0.82#a.0
0.97
327+a2
0.92
SEA (GIES)
0.95
1.00
35.4
0.91
0.86
0.98
272
0.92+0.0
INVCOV
0.46.0.0
1.00
308
0.35
0.98
1445
CORR.
0.42+o.0
1.00
326#21
0.20=a.0
0.89
1710#82
300
AVICI
0.01
0.70+0.a
198
0.64+0.0
0.02=a.0
0.50=a.o
1214
0.51
SEA (Fet)
0.8010.0
1.0040.0
121:14
0.7810.0
0.70.0.0
0.95.0.0
693167
0.8610.0
SEA (GIES)
0.88+v.0
1.00+t.a
88.9+11.3
0.84+a.0
0.78
0.96
556
0.87
INvCOv
0.47
1.00+0.0
418
0.36=0.0
0.9820.0
1883
CORR
0.4210.0
1.0010.0
445a14
0.20=0.0
0.91
2269
400
AVIcI
0.01
0.68+0.0
411
0.62+0.0
0.01
0.46=a.0
1614+21
0.47+0.0
SEA (Fc1)
0.49
0.931
3141107
0.61
0.56.0.1
0.90
1103=190
0.7510.1
SEA (GIES)
0.70
0.9910.0
226
0.71
0.70=0.0
0.94
872:44
0.8010.0
INvCov
0.47
1.00
194
0.38
0.99
2300
CORR
0.42+0.0
1.00
543:18
0.21
0.93
2790
500
AVICI
0.00+0.0
0.70+o.a
197+19
0.63
0.01
0.482.00
2001
0.48+0.0
SEA (FcI)
0.27:0.1
0.90
7581297
0.51:0.0
0.29
0.86
1824=273
0.5610.1
SEA (GIES)
0.41
0.9810.0
4851170
0.57
0.48
0.92
1654:505
0.6710.0
"consensus graph" consists of 11 nodes and 17 edges over 5,845 samples, of which 1,755 are observational and 4,091 are interventional. The observational data were generated by a. "general perturbation" which activated signaling pathways, and the interventional data were generated by perturbations intended to target individual proteins. Despite the popularity this datasel in causal discovery literature (due to lack of better alternatives), biological networks are known to be time-resolved and cyclic, so the validity of the ground truth "consensus" graph has been questioned by experts Mooij et al. (202 Nonetheless, benchmark all methods on this dataset in Table 26
41

---

Published in Transactions on Machine Learning Research (03/2025)
Table 17: Causal discovery ablations by setting hidden representations zero. Each setting encompasses 5 distinct Erdös Rényi graphs. The symbol f indicates that SEA was not pretrained on this setting. We set T = 100.
Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial
mAP t SHD
mAP 1
SHD +
mAP 1 SHD + mAP 1 SHD
mAP SHD +
SEA (Fcr)
0.97
1.6
0.95
2.4
0.92
2.8
0.83
3.7
0.69
6.7
10
h" +- 0 A
0.20
29.8
0.27
22.4
0.34
24.2
0.26
22.6
0.24
26.0
0.61
24.0
0.71
28.8
0.76
27.8
0.49
26.7
0.51
26.7
SEA (Gres)
0.99
1.2
0.94
2.6
0.91
3.2
0.85
4.0
0.70
5.8
hP +-0 0 A
0.30
29.2
0.27
29.4
0.19
29.0
0.35
27.4
0.31
27.1
0.85
6.4
0.82
10.2
0.78
13.2
0.63
10.2
0.55
13.4
SEA (FCI)
0.86
29.6
0.55
73.6
0.72
51.8
0.77
42.8
0.61
61.8
20
80
A + 0 A <
0.23
128.4
0.27
110.6
0.22
119.6
0.23
110.7
0.23
111.1
0.79
50.4
0.52
99.6
0.71
76.4
0.58
93.6
0.59
87.5
SEA (GIES)
0.89
26.8
0.58
71.4
0.73
50.6
0.76
45.0
0.65
60.1
he + 0 A < 0
0.25
125.2
0.21
123.0
0.24
113.6
0.27
118.0
0.24
129.8
0.86
35.2
0.55
76.8
0.70
53.4
0.69
47.1
0.59
63.5
SEA (Fct)
0.90
122.0
0.28
361.2
0.60
273.2
0.69
226.9
0.38
327.0
100 400
Ao 0
0.05
726.4
0.04
639.4
0.05
637.0
0.05
760.2
0.04
658.3
R + 0
0.82
167.4
0.04
403.4
0.51
352.4
0.64
263.8
0.33
366.1
SEA (GIES)
0.91
116.6
0.27
364.4
0.61
266.8
0.69
218.3
0.38
328.0
h" <- R -
0.05
780.0
846.0
0.05
715.4
0.04
744.4
769.8
0.86
134.8
0.03
403.4
0.52
357.8
0.67
224.1
0.32
359.6
Table 18: Our predicted graphs are highly acyclic, on synthetic ER test sets.
Acyclic
Total
Proportion
10
434
440
0.99
20
431
440
0.98
100
433
440
0.98
Average runtimes on larger graphs
150
N-100
125
N=200
N=300
100
N=400
Wall time
N-500
25
AVICI
Ours (FCD
Ours (GIES)
Figure SEA scales very well terms of runtime on much larger graphs, while AvIcI runtimes suffer as graph sizes incrcase.
42

---

Published in Transactions on Machine Learning Research (03/2025)
Table 19: Full results on syuthetic datasets (contimuous metrics).
Mean/std over distinct Srdöss-Rény graphs. it indicates o.o.d. sctting. indicates non-parametric bootstrapping. All standard deviations within 0.03 (most within 0.01).
N E Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial
mAP t
AUC f
mAP 1
AUC 1
mAP 1 AUC f
mAP 1
AUC f mAP
AUC 1
DCDI-G
0.74
0.88
0.79
0.91
0.89
0.95
0.46
0.72
0.41
0.68
DCDI-DsF
0.82
0.92
0.57
0.83
0.50
0.81
0.38
0.69
0.29
0.64
DCD-FG
0.45
0.68
0.41
0.67
0.59
0.79
0.40
0.64
0.50
0.72
DIFFAN
0.25
0.73
0.32
0.70
0.12
0.51
0.24
0.70
0.20
0.65
DECI
0.18
0.63
0.16
0.63
0.23
0.71
0.29
0.72
0.46
0.83
10 AVICI
0.45
0.80
0.81
0.97
0.65
0.89
0.52
0.81
0.31
0.70
VARSORT*
0.70
0.84
0.76
0.90
0.83
0.93
0.52
0.72
0.40
0.71
INvCoV
0.46
0.88
0.43
0.86
0.34
0.81
0.43
0.86
0.83
Fcr*
0.52
0.78
0.38
0.71
0.40
0.70
0.56
0.79
0.41
0.71
GIES*
0.81
0.96
0.61
0.93
0.71
0.92
0.70
0.95
0.61
0.87
SEA (Fc1)
0.98
1.00
0.88
0.98
0.88
0.97
0.83
0.97
0.88
SEA (GrEs)
0.99
1.00
0.94
0.99
0.91
0.98
0.85
0.97
0.70
0.89
DCDI-G
0.65
0.72
0.65
0.74
0.84
0.88
0.54
0.59
0.56
0.61
DCDI-Dsr
0.65
0.72
0.59
0.67
0.84
0.90
0.52
0.58
0.57
0.62
DCD-FG
0.51
0.57
0.57
0.63
0.53
0.59
0.65
0.68
0.65
0.68
DIFFAN
0.40
0.49
0.36
0.37
0.41
0.53
0.40
0.45
0.37
0.40
DECI
0.45
0.49
0.50
0.58
0.44
0.52
0.53
0.61
0.63
0.72
AvIcI
0.46
0.47
0.63
0.65
0.79
0.86
0.49
0.53
0.47
0.55
40
VARSORT*
0.81
0.83
0.87
0.89
0.71
0.73
0.69
0.72
0.56
0.62
INvCov
0.49
0.57
0.46
0.50
0.46
0.53
0.47
0.50
0.48
0.53
Fcr*
0.43
0.50
0.50
0.53
0.46
0.52
0.50
0.51
0.45
0.50
GIES
0.49
0.53
0.56
0.64
0.49
0.53
0.44
0.46
0.61
0.62
SEA (FcI)
0.83
0.85
0.85
0.89
0.86
0.90
0.74
0.75
0.67
SEA (GIES)
0.94
0.95
0.91
0.94
0.89
0.92
0.81
0.85
0.89
0.92
DCDI-G
0.59
0.87
0.78
0.94
0.75
0.91
0.36
0.81
0.42
0.74
DeDI-DsF
0.66
0.89
0.69
0.91
0.41
0.83
0.37
0.82
0.26
0.71
DCD-FG
0.48
0.85
0.58
0.91
0.51
0.87
0.50
0.78
0.44
0.76
DIFFAN
0.19
0.73
0.16
0.69
0.20
0.72
0.29
0.79
0.09
0.65
DECI
0.14
0.70
0.14
0.72
0.16
0.73
0.24
0.79
0.35
0.84
AVICI
0.48
0.87
0.59
0.91
0.67
0.90
0.42
0.84
0.24
0.69
20
VARSORT*
0.81
0.91
0.81
0.92
0.57
0.83
0.50
0.76
0.33
0.69
INvCov
0.40
0.90
0.31
0.90
0.31
0.84
0.42
0.92
0.41
0.87
Fcr*
0.66
0.86
0.42
0.74
0.40
0.77
0.56
0.80
0.41
0.76
GIES*
0.84
0.99
0.79
0.97
0.56
0.93
0.71
0.97
0.91
SEA (Fcr)
0.96
1.00
0.91
0.99
0.82
0.97
0.85
0.98
0.69
0.91
SEA (Gres)
0.97
1.00
0.94
0.99
0.83
0.97
0.84
0.97
0.69
0.92
DCDI-G
0.46
0.73
0.41
0.71
0.82
0.93
0.48
0.71
0.37
0.62
DCDI-DsF
0.48
0.75
0.44
0.74
0.74
0.92
0.48
0.71
0.38
0.63
Den-Fa
0.32
0.61
0.33
0.64
0.41
0.73
0.47
0.74
0.49
0.69
DIFFAN
0.21
0.53
0.19
0.41
0.18
0.46
0.22
0.55
0.18
0.37
DECI
0.25
0.57
0.29
0.61
0.26
0.59
0.31
0.66
0.43
0.73
AVICI
0.34
0.63
0.16
0.73
0.49
0.74
0.34
0.64
0.30
0.59
20 80
VARSORT*
0.76
0.86
0.50
0.81
0.47
0.69
0.59
0.76
0.38
0.63
INvCov
0.36
0.72
0.26
0.54
0.30
0.64
0.35
0.72
0.32
0.61
For*
0.30
0.59
0.31
0.57
0.30
0.59
0.41
0.66
0.34
0.61
GIES*
0.41
0.75
0.44
0.74
0.46
0.73
0.50
0.78
0.69
SEA (Fcr)
0.80
0.92
0.55
0.81
0.70
0.89
0.74
0.85
0.55
0.67
SEA (GrEs)
0.89
0.95
0.58
0.84
0.73
0.90
0.76
0.90
0.65
0.84
43

---

Published in Transactions on Machine Learning Research (03/2025)
Table 20: Full results on synthetic datascts (discrete metrics). Mcan/std over 5 distinct Bedöös-Rény graphs. indicales o.o.d. setting. * indicales not-parametrie bootstrapping.
All OA standard deviations within 0.2.
N E Model
Linear
NN add.
NN non-add.
Sigmoid"
Polynomial"
OA
SHD
OA
SHD 4
OA
SHD
OA t
SHD 4 OA f SHD
DeDI-G
0.73
2.8+2
0.84
2.243
0.88
1.0+1
0.46
5.8=3
0.33
8.94a
DCDI-DSF
0.81
2.0+3
0.73
3.0#3
0.60
4.2+1
0.43
6.3=s
0.24
11.2+a
DeD-FG
0.50
20.413
0.47
21.214
0.70
19.214
0.43
19.824
0.54
18.545
DIFFAN
0.61
14.0ss
0.55
13.6:14
0.26
21.8
0.53
12.0=5
0.47
15.026
DECI
0.50
19.4+s
0.43
13.8#6
0.62
16.2+a
0.63
13.9
0.75
7.8+4
AvIcI
0.58
8.2+4
0.79
4.2+3
0.65
5.6+3
0.54
8,3
0.35
9.6+4
10 10
VARSORT*
0.70
6.012
0.74
4.0
0.90
4.243
0.52
7.6
0.48
9.313
INvCov
10.623
10.2
13.623
11.1=4
10.423
Fc1*
0.52
10.0
0.43
8,244
0.41
9.812
0.55
9.1
0.40
10.0
GIES
0.76
3.612
0.49
6.0.5
0.69
4.8+2
0.67
5.9
0.63
7.113
SEA (Fc1)
0.93
1.0+1
0.82
3.0
0.90
3.8+2
0.73
3.9
0.70
6.1
SEA (GIEs)
0.94
1.2:1
0.88
2.6
0.93
3.211
0.84
4.0=3
0.79
5.843
DeDI-G
0.50
19.8#2
0.64
17.423
0.82
8.0+2
0.25
30.8
0.23
31.0+2
DeDI-DsE
0.48
19.8+3
0.55
21.627
0.87
6.0+3
0.25
31.42s
0.25
30.0#3
DCD-FG
0.35
28.614
0.40
26.0.11
0.36
26.814
0.45
24.4
0.41
25.514
DIFFAN
0.41
27.614
0.29
33.414
0.49
26.414
0.38
29.416
0.32
31.814
DECI
0.43
27.6+s
0.55
22.4+4
0.50
25.4+3
0.54
21.5
0.57
18.443
AvIcI
0.44
33.0
0.68
28.4+7
0.85
20.0++
0.50
35.0
0.50
38.8+2
10 40
VARSORT*
0.76
21.8
0.81
15.4
0.61
24.0.44
0.67
33.625
0.46
37.612
INVCOV
40.212
44.6.2
41.043
44.2
42.3+2
Fc1#
0.16
39.2+2
0.24
38.843
0.22
36.2+3
0.22
38.8#2
0.16
39.942
GIES*
0.44
33.444
0.60
33.045
0.51
32.0+3
0.38
36.7#2
0.63
34.643
SEA (Fer)
0.80
18.6#2
0.89
15.4
0.87
18.8+4
0.78
24.426
0.69
26.924
SRA (Gres)
0.91
12.8+4
0.95
10.420
0.89
17.243
0.87
24.5=s
0.93
29.5+3
DeDI-G
0.75
6.4
0.90
3.0#2
0.84
4.4+2
0.39
42.7=6
0.48
10.443
DeDI-DsP
0.77
5.2+3
0.85
4.2+4
0.64
11.623
0.41
43.2=6
0.45
15.725
Den-Fo
0.76
51.2212
0.88
177457
0.85
193_27
0.65
251
0.60
52.7+19
DIFFAN
0.56
40.2a27
0.47
38.64
0.53
35.0=26
0.63
19.2#s
0.42
49.7+115
DECI
0.54
52.0=17
0.56
41.0+14
0.56
39 39.048
0.65
30.0
0.73
18.946
AvICI
0.56
17.2#6
0.69
10.8
0.79
11.2+3
0.53
17.2=6
0.37
18.4+4
20
20
VARSORT*
0.84
10.013
0.88
6.6as
0.74
14.618
0.50
16.1
0.40
17.114
INvCOv
23.616
24.6as
24.6:16
22.9
20.0
For*
0.70
19.01s
0.45
17.4a2
0.50
19.416
0.57
18.524
0.44
18.945
GIES*
0.80
7.4+2
0.77
9.0+a
0.71
14.0+3
0.75
12.5#
0.68
13.7+4
SEA (Fer)
0.92
3.243
0.93
5.023
0.94
8.8+3
0.79
6.7=s
0.76
9.844
SEA (Gres)
0.89
3.0+1
0.95
3.422
0.94
7.843
0.83
8.1
0.78
10.1
DeDI-G
0.54
44.0#6
0.53
61.6+11
0.89
37
37.4=34
0.46
44.2
0.26
59.7+5
DeDI-DsF
0.57
41.213
0.61
60.0+12
0.85
28.4
0.47
43.626
0.30
57.625
DG
0.58
172=27
0.65
156241
0.75
162=49
0.62
80.1
0.48
79.8
DIFFAN
0.46
127+s
0.28
154+10
0.36
145+7
0.45
117421
0.21
157+7
DECI
0.30
87,243
0.47
104+7
0.35
79.649
0.48
71.0
0.57
58.9411
AVICI
0.51
75.6210
0.69
72.826
0.69
61.2
0.50
70.6
0.50
75.527
20 80
VARSORT*
0.82
44.844
0.84
73.6213
0.61
15
0.67
63.426
0.39
75.625
INvCOv
97.616
12144
10449
95.627
97.8
For*
0.19
75.8
0.19
80.215
0.22
74.418
0.33
72.326
0.23
76.615
G
0.56
70.0
0.73
75.2+4
0.59
67.4+7
0.62
65.6=7
0.61
68.1+a
SEA (Fe1)
0.86
39.8+12
0.87
73.8+12
0.92
52.0+*
0.82
42.9
0.69
57.0+5
SEA (GIEs)
0.92
26.8
0.88
71.4as
0.92
50.6
0.84
45.0a7
0.89
60.116
44

---

Published in Transactions on Machine Learning Research (03/2025)
Table 21: Full results on synthetic dalasets (continuous metrics). Mean over distinct scale-free graphs.
indicates o.o.d setting. indicates non parameeeitiic bootstrapping.
All standard deviations were within 0.02.
N E Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial
mAP t
AUC f
mAP 1
AUC t
mAP f AUC f mAP
AUC f mAP
AUC f
DCDI-G
0.54
0.90
0.59
0.88
0.69
0.89
0.48
0.77
0.50
0.73
OcDI-DsF
0.70
0.92
0.71
0.88
0.36
0.83
0.46
0.75
0.49
0.76
DCD-FG
0.56
0.76
0.47
0.72
0.50
0.73
0.44
0.68
0.57
0.75
DIFFAN
0.25
0.73
0.15
0.66
0.16
0.62
0.31
0.75
0.24
0.63
DRCI
0.17
0.65
0.17
0.67
0.20
0.72
0.27
0.73
0.49
0.82
AVICT
0.51
0.87
0.55
0.85
0.76
0.95
0.44
0.81
0.27
0.71
10
VARSORT*
0.67
0.84
0.69
0.86
0.76
0.88
0.45
0.69
0.46
0.73
INVCOV
0.50
0.92
0.41
0.87
0.38
0.84
0.47
0.90
0.45
0.86
Fcr*
0.56
0.80
0.51
0.80
0.43
0.74
0.60
0.82
0.34
0.68
GIES
0.87
0.98
0.61
0.94
0.69
0.94
0.75
0.96
0.71
0.91
SEA (Fc1)
0.96
0.99
0.88
0.98
0.88
0.97
0.79
0.96
0.73
0.90
SEA (GuES)
0.95
0.99
0.94
0.98
0.92
0.98
0.85
0.98
0.74
0.90
DCDI-G
0.70
0.85
0.74
0.85
0.88
0.91
0.56
0.66
0.53
0.64
DCDI-DsF
0.74
0.87
0.73
0.84
0.71
0.90
0.56
0.69
0.51
0.63
DCD-FG
0.37
0.58
0.45
0.61
0.45
0.58
0.49
0.63
0.63
0.73
DIFFAN
0.29
0.50
0.25
0.38
0.28
0.46
0.31
0.53
0.27
0.44
DECI
0.30
0.51
0.41
0.65
0.33
0.51
0.38
0.60
0.59
0.77
AVICT
0.41
0.57
0.65
0.81
0.55
0.67
0.41
0.59
0.40
0.61
10
40
VARSORI*
0.77
0.83
0.74
0.87
0.59
0.71
0.66
0.76
0.50
0.66
INvCOV
0.44
0.71
0.38
0.59
0.42
0.62
0.44
0.67
0.62
Fcr*
0.47
0.64
0.41
0.60
0.40
0.58
0.48
0.64
0.41
0.59
GIES
0.43
0.68
0.43
0.63
0.14
0.61
0.49
0.69
0.59
0.71
SEA (Fc1)
0.85
0.91
0.80
0.90
0.77
0.88
0.76
0.83
0.66
0.71
SEA (GIES)
0.92
0.96
0.84
0.93
0.83
0.90
0.79
0.88
0.78
0.87
DCDI-G
0.41
0.95
0.50
0.94
0.69
0.96
0.37
0.83
0.37
0.77
DCDI-DsF
0.48
0.95
0.55
0.93
0.33
0.90
0.37
0.79
0.35
0.82
DCD-FG
0.51
0.87
0.39
0.83
0.48
0.84
0.56
0.84
0.50
0.84
DIFFAN
0.27
0.80
0.11
0.65
0.11
0.66
0.26
0.77
0.12
0.69
DECI
0.13
0.69
0.15
0.71
0.15
0.73
0.15
0.71
0.25
0.79
AVICI
0.53
0.88
0.66
0.89
0.74
0.92
0.46
0.87
0.32
0.77
20
20
VARSORT*
0.67
0.85
0.84
0.93
0.59
0.86
0.45
0.72
0.44
0.73
INvCOv
0.44
0.94
0.35
0.91
0.30
0.89
0.43
0.93
0.87
Fcr*
0.63
0.84
0.44
0.78
0.43
0.79
0.60
0.86
0.47
0.78
GIES
0.82
0.99
0.58
0.95
0.57
0.96
0.75
0.98
0.61
0.90
SEA (FC1)
0.94
1.00
0.87
0.98
0.83
0.97
0.82
0.98
0.72
0.92
SEA (GIES)
0.93
1.00
0.91
0.98
0.88
0.98
0.82
0.98
0.70
0.91
DCDI-G
0.62
0.88
0.61
0.89
0.76
0.94
0.44
0.76
0.36
0.60
DCDI-DsF
0.58
0.87
0.55
0.86
0.58
0.92
0.43
0.78
0.35
0.66
DCD-FG
0.38
0.70
0.30
0.69
0.48
0.80
0.48
0.75
0.53
0.73
DIFFAN
0.18
0.55
0.15
0.44
0.16
0.53
0.19
0.56
0.15
0.38
DECI
0.21
0.58
0.24
0.64
0.26
0.66
0.30
0.68
0.41
0.75
AVICI
0.29
0.61
0.54
0.80
0.60
0.85
0.35
0.65
0.28
0.63
20 80
VARSORT*
0.79
0.90
0.57
0.83
0.62
0.81
0.57
0.77
0.38
0.64
INVCOV
0.38
0.81
0.22
0.56
0.32
0.73
0.38
0.78
0.33
0.67
Fcr*
0.31
0.63
0.30
0.62
0.30
0.62
0.41
0.68
0.32
0.62
GrES*
0.51
0.87
0.43
0.78
0.47
0.81
0.52
0.82
0.47
0.73
SEA (Fct)
0.87
0.96
0.59
0.87
0.70
0.90
0.73
0.87
0.53
0.71
SEA (GLES)
0.92
0.98
0.63
0.89
0.73
0.91
0.77
0.92
0.62
0.84
45

---

Published in Transactions on Machine Learning Research (03/2025)
Table 22: Full results on synthetic datasets (discrete metrics). Mean/sid over distinct scale- graphs. i All OA standard deviations within 0.2.
indicates o.o.d. setting. indicates non- non-parametric bootstrapping.
N E Model
Linear
NN add.
NN non-add.
Sigmoid'
Polynomial'
OA t SHD
OA t
SHD
OA t SHD
OA t SHD
OA t SHD
DCDI-G
0.51
16.621
0.61
17.412
0.71
16.212
0.59
16.9
0.63
16.613
DCDI-Des
0.70
16.2=2
0.75
15.4
0.36
16.822
0.48
18.1
0.67
18.0
DCD-FG
0.60
16.4*s
0.57
22.2+4
0.57
20.0+4
0.47
17.9+4
0.59
16,8
DIFFAN
0.55
9.2+4
0.49
14.644
0.36
11.644
0.59
7.744
0.42
14.8+6
DECI
0.51
17.425
0.55
17.425
0.58
12.0#2
0.62
12.6
0.74
8.2+6
AVICI
0.64
6.823
0.59
5.443
0.77
3.6#1
0.48
7.8+3
0.35
9.5+3
10 10
VARSORT*
0.82
4.4
0.78
4.8a2
0.80
3.412
0.47
7.313
0.55
8.1
INVCoV
8.8
9.242
9.811
9.74:
10.2
Fcr*
0.62
8.4=2
0.51
7.8+2
0.45
8.0#2
0.61
9.2+3
0.34
9.6+2
GIES
0.83
2.2#2
0.60
6.0
0.75
4.2+2
0.72
4.9+3
0.73
5.7+2
SEA (Fct)
0.89
1.4=2
0.90
2.6
0.93
2.2
0.71
4.0.a3
0.72
5.212
SEA (GLE))
0.85
1.4+1
0.96
1.8+1
0.94
2.0+1
0.86
3.4+3
0.83
5.0
DCDI-G G
0.82
24.0
0.85
27.8
0.87
19.612
0.62
31.413
0.52
32.614
DCDI-DsF
0.79
22.8
0.79
24.4:4
0.82
20.412
0.64
31.612
0.58
33.315
DCD-Fa
0.36
24.8=4
0.41
25.243
0.38
25.6+8
0.41
23.2+6
0.54
18.5+3
DIFFAN
0.40
29.8**
0.28
37.0#2
0.38
32.642
0.45
28.0+a
0.33
32.7
DECI
0.43
27.8
0.66
22.613
0.48
28.613
0.52
22.214
0.66
13.3
AVICT
0.43
20.224
0.84
17.224
0.61
20.4
0.52
24.8
0.49
26.5+2
10 40
VARSORI*
0.75
13.4#2
0.89
12.6#2
0.60
20.6+3
0.65
20.8++
0.50
26.4+2
INVCOV
31.4+3
37.2+4
36.0#4
32.3
34.8
Fcr*
0.33
23.8=2
0.28
27. 27.2
0.25
27.6#4
0.36
26.112
0.24
27.3+2
GIES
0.16
21.8=3
0.50
24.1#2
0.48
25.2+5
0.52
22.743
0.64
22.5+2
SEA (Fct)
0.80
10.0
0.88
14.0.2
0.87
15.8
0.79
15.643
0.64
18.513
SEA (GIES)
0.88
6.6
0.98
14.0as
0.88
14.413
0.87
14.143
0.93
19.1
DCDI-G
0.54
40.4
0.70
44.8
0.88
39.8#6
0.47
41.144
0.53
38.4
DCDI-Des
0.64
40.4aa
0.65
42.4
0.40
42.2
0.37
41.144
0.45
49.3
DCD-FG
0.68
232:23
0.77
183152
0.78
181127
0.70
251145
0.69
278168
DIFFAN
0.67
23.6+13
0.40
12.2+
0.42
34.0+11
0.59
22.6+12
0.50
46.8+13
DECI
0.50
42.0=a
0.54
43.0+11
0.57
40.0+13
0.51
31.747
0.65
25.3+
AVICI
0.68
11.8=3
0.79
9.2#2
0.80
7.8+2
0.60
15.344
0.45
18.325
20 20
VARSORT*
0.74
10.4#3
0.91
6.2+
0.76
13.0#a
0.49
13.7+4
0.51
16.7
INVCOV
20.2+3
24.4+5
23.8+5
21.2+4
20.8+
Fcr*
0.67
13.8#2
0.52
17.4+1
0.53
17.845
0.65
16.744
0.50
18.9
GIES
0.82
6.423
0.71
12.622
0.68
13.243
0.75
11.425
0.73
13.414
SEA (Fc1)
0.90
2.8
0.87
7.0.2
0.91
8.245
0.73
7.743
0.71
9.5
SEA (GIES)
0.85
4.0a2
0.95
3.642
0.93
6.2
0.80
7.9.4
0.82
9.9
DCDI-G
0.81
93
0.73
104
0.91
67.828
0.61
82.54s
0.53
79.7
DCDI-DsF
0.75
103+7
0.73
94.8+11
0.77
63.8+7
0.65
84.4*
0.52
82,645
DCD-FG
0.63
188424
0.70
187+14
0.78
190422
0.71
217+38
0.71
235+a2
DIFFAN
0.42
111
0.30
145+11
0.40
119413
0.41
100
0.21
149411
DECI
0.33
72.2+10
0.48
81.4410
0.48
67.029
0.50
60.4+12
0.58
47.2
AVICI
0.45
56.0as
0.77
47.6a8
0.76
42.613
0.50
54.016
0.47
62.2.5
20 80
VARSORT*
0.85
32.4as
0.84
54.8
0.75
40.2:1
0.64
54.01a
0.37
60.
INVCOV
79.4=c
10645
87.442
78.849
86.2+4
Fcr*
0.26
58. 68.2#6
0.22
58.4+4
0.26
55.0+7
0.37
59.3+6
0.23
62.9+4
GIES
0.65
48.4
0.71
53.623
0.68
48.0a4
0.64
53.644
0.61
54.945
SEA (Fct)
0.91
25. 626
0.88
51.624
0.89
42.2+4
0.83
36.116
0.71
47.415
SEA (GLES)
0.92
17.6#3
0.93
49.2
49.249
0.89
37.2
0.88
35.147
0.89
48.
46

---

Published in Transactions on Machine Learning Research (03/2025)
Table 23: Causal discovery results on synthetic datasets with 100 nodes, continuous metries. Each setting encompasses distinct Erdös- Rendi graphs. The symbol i indicates that the model was not trained on this setting. All standard deviations were within 0.1.
N
E
Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial!
mAP
AUC
mAP
AUC
mAP
ATC
mAP f
AUC f
mAP t AUC ?
DCD-FG
0.11
0.75
0.12
0.71
0.18
0.73
0.20
0.72
0.06
0.60
100 100
INVCOV
0.40
0.99
0.22
0.94
0.16
0.87
0.40
0.97
0.36
0.90
SEA (FCI)
0.96
1.00
0.83
0.97
0.75
0.97
0.79
0.97
0.56
0.88
SEA (GIES)
0.97
1.00
0.82
0.98
0.74
0.96
0.80
0.97
0.54
0.85
DCD-FG
0.05
0.59
0.07
0.64
0.10
0.72
0.13
0.72
0.12
0.64
100 100
INvCoV
0.25
0.91
0.09
0.62
0.14
0.77
0.27
0.86
0.20
0.67
SEA (Fc1)
0.90
0.99
0.28
0.82
0.60
0.92
0.69
0.92
0.38
0.80
SEA (GIES)
0.91
0.99
0.27
0.82
0.61
0.92
0.69
0.91
0.38
0.78
Table 24: Causal discovery results on synthetic datasets with 100 nodes, discrete metries. Each setting encompasses distinet Erdös-Rény' graphs. The symbol f indicates that the model was not trained on this setting.
N
F. Model
Linear
NN add.
NN non-add.
Sigmoid
Polynomial*
OA f SHD
OA
SHD 4
OA
SHD
OA t SHD
OA
SHD
DCD-FG
0.63
3075.8
0.58
2965.0
0.60
2544.4
0.59
3808.0
0.34
1927.9
100 100
INVCOV
124.4
130.0
158.8
112.3
106.3
SEA (Fci)
0.91
13.4
0.90
34.4
0.91
47.2
0.78
40.3
0.69
59.2
SEA (GEES)
0.91
13.6
0.93
32.8
0.91
45.8
0.78
38.6
0.68
60.3
DCD-Fa
0.46
3068.2
0.60
3428.8
0.70
3510.8
0.67
3601.8
0.53
3316.7
100 400
INVCOV
557.0
667.8
639.0
514.7
539.4
SEA (FeI)
0.93
122.0
0.90
361.2
0.91
273.2
0.87
226.9
0.82
327.0
SEA (Gres)
0.94
116.6
0.91
364.4
0.92
266.8
0.87
218.3
0.84
328.0
1.0
1.0
InCOrr
InOOr
Oer
Ger DEATT
0.6
0.6
1L5
02
02
Groph
Gryph
Figure 10: mAP on graphs larger than seen during training. During training, only sampled a maximum of 100 subsets, so performance drop may be due to extrapolation beyond trained embeddings. We did not have time to finetune these embeddings for more samples. These values correspond to the numbers in Table I6
47

---

Published in Transactions on Machine Learning Research (03/2025)
Table 25: Swapping to estimation algorithms with significantly different assumptions (LiNGAM) leads to a larger performance drop. Interestingly, the GIES model seems depend more marginal estimates, while the FCI depends more on global statistics. Results on N 10, E 10 observational setting.
Linear
NN add.
NN non-add.
Sigmoid
Polynomial
Train Inference
mAP f SHD 4 mAP +
SHD 4 mAP t
SHD 4 mAP
SHD
mAP t SHD
FCT
0.96
3.21.6
0.91
5.04.8
0.822.05
8.84.9 0.852.09
6.7..1
0.69
9.8
PC
0.95#.01
3.6#.
0.91=.05
4.6#..6
0.81+.06
10.0#.3
0.84+.07
7.4=.2
0.65=.14
10.8
FCI
GES
0.94
4.4+.3
0.91
4.2+.6
0.81
9.62.6
0.81
8.42.9
0.61=.16
11.5#.2
GRaSP
0.94#.05
4.0+.0
0.91=.05
4.4+.9
0.81+.06
10.0#.0
0.81#.10
8.5=.0
0.61=.16
11.5#.2
LINGAM
0.881.08
8.01.5
0.841.05
6.01.3
0.741..05
10.81.3 0.71
12.5 0.59
14.41.6
PC
0.961.02
1.81.7
0.91
2.81.4
0.891.10
3.21.2
0.821.14
4.1
0.58
6.7
GES
0.95
2.04.9
0.9
2.64.1 0.88
3.42.5
0.
4.12.5 0.57=.19
6.8
GIES
GRaSP
0.951.03
1.81.7
0.92..05
3.01.8
0.881.11
3.21.2
0.81
4.0..4 0.57
6.91.0
LINGAM
0.60
5.64.8
0.40= 23
9.0+.6
0.40
9.4+.2 0.53+.15
7.2=.0 0.51=.14
7.9+.0
Table 26: Complete results on Sachs flow cytometry dataset [aacss et al. by Wang et al. 2017).
2005 using the subset proposed
Model
mAP t
AUC t SHD +
DCDI-G
0.17
0.55
21
DCDI-DsF
0.20
0.59
20
DCD-FG
0.32
0.59
27
DIFFAN
0.14
0.45
37
DECI
0.21
0.62
28
AVICI-L
0.35
0.78
20
AvICI-R
0.29
0.65
18
AVICI-L+R
0.59
0.83
14
Fcr*
0.27
0.59
18
GIES*
0.21
0.59
17
SEA (Fer)
0.23
0.54
24
+KeI
0.33
0.63
14
+CORR
0.41
0.70
15
KcI+ CORR
0.49
0.71
SEA (GIES)
0.23
0.60
14
48


Reviews:
[{"review_id": "a2zOB3EFpb", "reviewer": "Reviewer_RoBV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "There are several places where I find issues with the proposed method or surrounding discussion:\n\nI'm going to make my comments in the order they appear rather than in order of importance. I'll mark what I take to be the most important comments with an asterisk *.\n\n*Abstract and later: in several places, the authors repeat some version of the claim that errors of causal discovery algorithms \"remain comparable across datasets.\" (This is said in the abstract, the introduction, Sec 3.4, and elsewhere.) I'm not sure what this means precisely or what reason there could be to believe this. It seems the suggestion is that different algorithms, which operate under different assumptions, use different statistical principles, and consider different classes of outputs (DAGs vs CPDAGs vs PAGs vs other), will suffer \"similar\" errors across different datasets generated from different kinds of ground truth... I don't see how this could possibly be true and certainly is not true in my experience with causal discovery algorithms. Different algorithms can behave quite differently on the same data in practice, and the same algorithm across different data sets will often have very different performance. Maybe I misunderstand what the authors intend to convey with these sentences, but if so it should be clarified -- and if it is actually essential then the claim should be backed up by evidence.\n\nPg. 1 \"Discrete optimization algorithms\": this is a potentially misleading way to describe the class of causal discovery methods discussed in the cited work, since many of these are not \"optimization\" algorithms in an obvious sense -- while score-based algs do try to find \"optimal\" scoring graphs, constraint-based algs and LiNGaM-type algs do not optimize anything, but rather use testing or other model selection ideas. I would suggest simply calling these \"discrete search algorithms.\"\n\nPg. 1 \"Driven by hypothesis tests\": only some of these are based on hypothesis tests, others do not use hypothesis tests at all.\n\nPg. 1 \"Necessarily impose functional assumptions\": not true, there are nonparametric tests, though these have challenges e.g. no test has power for all alternatives.\n\nPg. 1 \"the former must be trained from scratch per-dataset\": this is a strange way to (negatively) summarize these continuous optimization-based algs. There's always *something* happening \"per-dataset\" (or else you're not using the data), it is only negative if that doesn't work well. Maybe the correct challenge to highlight for the cited methods is that they do not seem robust, they don't perform well outside of simple settings according to some recent papers.\n\nPg. 2 \"which determine whether G is identifiable\": should be \"uniquely identifiable or identifiable up to an equivalence class.\"\n\nPg. 2 \"allow us to perform interventions\": this should be \"can be used, along with other information, to compute the downstream consequences of interventions.\" Computing the post-intervention distribution requires more than just the graph structure.\n\n*Pg. 2 \"the goal of causal structure learning is to recover G\": there is no mention of latent variables or unmeasured confounding in this set up. The implicit assumption here is that the truth is a DAG over variables only in X, a.k.a. causal sufficiency. Even if the setting with unmeasured confounders is not of interest in this paper, which is fine, this should be made explicit here.\n\nPg. 2 \"Discrete optimization methods make atomic changes to a proposal graph until a stopping criterion is met.\" Again, this is misleading description of many algorithms, see above. It is also misleading to describe LiNGaM as score-based. These types of algorithms have been previously described as \"semi-parametric\" procedures since they exploit asymmetries implied by the semi-parametric model classes that are assumed.\n\nPg. 3 \"amortized inference approaches\": missing a reference to Petersen et al. (2023) which should be cited. (This actually came out on arxiv earlier than or around the same time as the cited work.)\n\nPetersen, A. H., Ramsey, J., Ekstrøm, C. T., & Spirtes, P. (2023). Causal Discovery for Observational Sciences Using Supervised Machine Learning. Journal of Data Science, 21(2), 255-280. \n\n*Pg. 4 some of the description in Sec 3.1 is unclear. What are the \"selection scores \\alpha? If they are correlation or precision matrices then they are not in {0,1}^{N x N}, they are real numbers. Are these actually probabilities? Why are they only over D_0? The \\sum_{j \\in S_t} \\alpha_{i,j} is not clear, what's going on with the index i? S_t are subsets of nodes. Why the index t? Overall this description could be cleaned up.\n\nPg. 4 \"Each input is a whole dataset\": make clear here that you are generating synthetic data with known ground truth based on some parametric model. This whole section leaves it somewhat mysterious what the training data is.\n\n*Pg. 4 use of FCI and GIES: I'm confused about these choices. FCI is for the setting with unmeasured confounding and learns a PAG, not a DAG or CPDAG. GIES is for the setting with no confounding but with interventional data. Why are you using FCI (instead of say PC), and how are you translating between the PAG output and the directed graph DG output (esp when computing statistics such as structural hamming distances, etc)? It's not straightforward, because the adjacencies in a PAG and a CPDAG may not be the same even with no unmeasured confounding. You mention using both observational data and interventional data, but how do you do use interventional data with FCI? When there is no interventional data do you use GES instead of GIES? It seems like the inference step is ultimately predicting a (possibly cyclic) directed graph, so how are you using information like \"X is an ancestor of Y\"? You write: \"as long as their outputs are the same format\" -- but they are not the same format? This part is unclear to me.\n\n*Pg. 6 \"graph prediction\": it seems you are learning a (possibly cyclic) directed graph, not a DAG or CPDAG. If you were learning a DAG, at least you could transform it to the corresponding CPDAG using existing results on Markov equivalence. In your evaluations, you seem to computer performance statistics (SHD etc) wrt the true DAG (since the truth is always a DAG) not the true CPDAG, is that correct? This is misleading, since you could estimate a DG where the true CPDAG (best you can do, theoretically) has a *higher* SHD to the true DAG than the estimated DG. But it isn't because you did \"better\" than getting the true CPDAG, but because you got randomly lucky. These issues ought to be addressed.\n\n(For example: say the truth is A-->B-->C. The true CPDAG is A--B--C which has an SHD of 2 from the true DAG. Perhaps you could learn this with PC or GES. But with the proposed algorithm you learn, out of luck, the DG A-->B-->C-->A, which has a SHD of 1 from the truth. You just got \"lucky,\" the best you can do nonparametrically is the CPDAG A--B--C. To say you've done \"better\" than PC/GES in this case would be misleading, PC/GES was perfectly correct.)\n\n*Theorem 3.1. I must admit I don't understand Theorem 3.1 as a formal statement, and I don't follow the proof. It is not that I doubt that the claim is roughly correct, but am not sure I understand what the claim is because of language that is to me (as a statistician) too informal. \"Has the capacity to recover\" is not precise to me. I'm used to claims about identifiability, what a procedure will do as sample size approaches infinity, or what will happen with high probability for some finite sample size, but this is none of those. Again, I don't really doubt that the reasoning is sound. I hope another reviewer could evaluate this claim more carefully. [Note: Appendices A.1 and A.2 seem ok to me, it is the argument in A.3 that I don't follow.]\n\nPg. 7 \"failure modes of causal discovery algorithms may be similar across datasets and can be corrected using statistics that capture richer information\": Again, it is not clear to me what this means and whether it is true (see above).\n\nPg. 7 \"an algorithm that assumes linearity will make (predictably) asymmetric mistakes on non-linear data and underestimate the presence of edges\" -- this depends on the algorithm. It is true that constraint-based methods that use linear association tests may miss edges that correspond to nonlinear dependencies. But how does LiNGaM or NOTEARS behave in the presence of nonlinearities? Less obvious.\n\n*Pg. 7 \"When all assumptions are upheld, and infinite data are available, the model has the capacity to infer a sound graph\": this statement is unclear (what are the assumptions, and what is a \"sound graph\"?) and there is no real evidence for this claim presented specifically. Theorem 3.1 is does not refer to DGP assumptions or infinite data at all.\n\nPg. 7 \"In practice, the model will output an orientation for all edges, but the graph can be interpreted as one member of an equivalence class\": as discussed above, if the graph was acyclic then you could and probably should transform the result to a CPDAG to compare to algs that return a CPDAG. As it stands the procedure can return a cyclic DG for which equivalence theory is more complicated (but anyway the training data all comes from DAGs).\n\nPg. 7 \"both observational and interventional datasets\": how were the interventional data sets generated, what were the interventions and the ground truth here? How was this info used, if at all, with the purely observational discovery algs?\n\n*Pg. 8 \"VarSort, Fci, and Gies were run using non-parametric bootstrapping (Friedman et al., 1999), with 100 subsets of 500 examples each\": how were the bootstraps aggregated into one final structure? Different ways of combining bootstrap with causal discovery algorithms can lead to very different results. Also I'm a little unclear about the sample sizes here. What was the actual sample size for each of the \"inference\" algorithms, n=500? (The training sample is described for the foundation model, but not the testing sample sizes for Table 1. Also, what \\alpha threshold is used with FCI? This obviously affects the sparsity of the discovered graph, and in real applications this hyperparameter needs to be selected somehow.\n\n*Sec 5.3: I find the discussion of \"respecting identifiability theory\" somewhat unclear. It seems to say that when the truth is only identifiable up to Markov equivalence, the proposed procedure recovers a graph that is similar in accuracy to a random graph from the MEC. Why not evaluate whether the procedure gets close to the true MEC/CPDAG instead? It is misleading to evaluate everything wrt the true DAG when that is not identifiable. Also: how is FCI being evaluated here since it returns a PAG? See above.\n\nReferences: some citations are missing a lot of information, needs editing.\n\n*Meta-point: There are a lot of important questions about using these kinds of tools in practice that are not addressed here. In my opinion, foundation models are only as good as the data they are trained on, and as it stands, this are trained on synthetic data which barely resembles real data. All one variable type, simple (even if non-linear) functional relationships, no missingness, no measurement error, etc. What assurances can there be for use on real data? How does one know the \"failure modes\" of this kind of procedure, which elements to trust? (I recognize these questions are also tough for classic causal discovery algorithms, but there I think we have at least imperfect answers. We know roughly when hypothesis tests perform poorly, when score-based algs get stuck in local minima, etc.)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an approach to building foundation models for causal discovery. I'll admit at the outset of this review that I'm a priori quite skeptical of this endeavor: I think in general it is more often than not dangerous to trust these models to arrive at empirical truths, for various reasons. But the approach outlined is interesting, and perhaps could serve as the basis for further discussion. I think the idea of learning subgraphs reliably along with global statistics and stitching them together is a good one for this approach.\n\n(I should also note that my expertise is in causal discovery, graphical models, and statistics, but I have no expertise in foundation models and minimal understanding of recent deep learning advancements, so my review will mainly stick to what I know.)", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The revision should address all the points raised above.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1Ih8A1K9t7", "reviewer": "Reviewer_fspD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper represents a valuable contribution to the active area of using pretraining for causal discovery (in the manner of amortized causal discovery).\n\nThe paper’s findings are mostly empirical. Overall, I find the experimental setup and presented results satisfactory.\n- The training datasets contain sufficient diversity to assess SEA’s performance in different settings. There is also some effort to evaluate the proposed framework on standard benchmark real data (e.g., Sachs et al. protein signaling dataset) \n- Multiple metrics and relevant baselines are presented to evaluate the results in a fair way. For instance, I appreciated using the naïve baseline VarSort to verify that the positive results are not an artifact of the synthetic data generation procedure, and reporting whether the predicted causal graphs obey the theoretical identifiability limits.\n\nAs expected, there are some limitations, which the authors discussed in various places in the paper. The theoretical aspect of the paper is weak, though I think the experimental aspect and the model design are convincing enough that I don’t have an issue with it.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a foundation model-like framework to predict causal graphs, similar to the idea of amortized causal discovery. In summary, the proposed framework Sample, Estimate, Aggregate (SEA) trains a deep learning model using large-scale data. \n\nSEA framework summary:\n- FM-inspired goals: Train a model that outperforms individual algorithms, easily change the assumptions to incorporate the knowledge, and strong performance in few-shot learning.\n- Training the model and inference on test data are performed similarly. In summary, the input is a whole dataset. *Sample* step samples batches and node subsets from the input dataset. *Estimate* step computes global statistics over all nodes (e.g., inverse covariance), and run a causal discovery algorithm on each batch/node subset to obtain marginal estimates (of the induced subgraph over each subset of nodes generated in the sample step).\n- The pretraining the *aggregator* is done in a supervised way: A deep net takes the global statistics, marginal estimates and node subsets to predict the known global causal graph. Aggregator first creates embedding of the inputs, then use axial attention blocks. Finally, all features are passed to a feedforward net to predict the causal graph. \n- Axial attention blocks are used for computational efficiency.\n- Inference (predicting a causal graph of a given dataset), is performed similar to the training: batches and node subsets are sampled, a causal discovery algorithm (e.g., FCI) is used to create marginal estimates on these subsets, and pretrained aggregator module combines these estimates and global statistics to output a global graph estimate.\n\nThe paper presents empirical proofs of concept of SEA that the framework achieves the FM-inspired goals (up to an extent). There is also *some* effort for theoretical demonstration of SEA’s capabilities, in the sense that the proposed model can technically produce the correct global output from marginal estimates.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Requests**:\n- Section 3.1: Overall Section 3.1 is a bit difficult to parse at this point. For instance, the meaning of the selection scores and the role of the global statistics are not immediately clear. Another example is marginal estimate. This is the first time in the paper ``marginal estimate’’ appears, so an explanation sentence would be really helpful. I know that the details become clear later in the paper (some are even in the appendix), but I think Section 3.1 needs to be more clear to improve the flow of the paper.\n- Discussion of (Lorch et al., 2022): While the approaches are different, it seems to me that the study by Lorch et al. 2022 is the most closely related one to this paper. Therefore, I think it warrants a better, more detailed discussion on the similar and different aspects of the present work compared to that paper.\n- I respect the decision of authors to limit the main content to 12 pages to make it a regular TMLR submission. That being said, I suggest authors to consider move some parts of the appendix to the main paper. Limitations discussion in Appendix B.7 and at the end of Appendix A.3 are helpful to assess the paper better, and can be given in the main paper. As a minor request, a very brief summary of Appendix A.2 in the main paper would help the reader -- just the main idea of how can marginal estimates be used to form the global graph theoretically, and the takeaways.\n\n**Minor requests:**\n- Section 2.1, Page 2: Please define the parent set $\\pi_i$ before using it in conditional.\n- Section 3.1, Page 4: Perhaps make the faithfulness statement more explicit.\n\n**Minor questions:**\n- Section 3.2, Page 4: In the aggregator explanation, an example is given for two variables with low correlation. I think the more interesting case would be to determine whether a high correlation pair is an edge (or just spurious correlation due to a confounder). Can you provide any intuition for how aggregator can handle such cases?\n- Section 3.3, Page 5: ``Unique edges’’, can you specify what does unique refer to here? Is it the union of all edges over the marginal estimates?\n- Appendix A.4: 2nd point. Why SEA(FCI) pretrained aggregator is more robust, this is a bit unclear to me. Can you provide any intuition behind this observation?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gp5RkQRL76", "reviewer": "Reviewer_gjMn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n\n* The paper is written in a clear and concise manner, with easy-to-understand illustrative figures.\n* The proposed SEA is a general framework that can be combined with classical causal discovery algorithms specializing in different assumptions. It is very efficient and easy to scale for large graphs in that it operates on the subsets of data.\n* The empirical validation is very comprehensive and strong. It shows the framework (especially the pretrained aggregator) can be generalized to new data (as well as sim2real, small samples), quick adaptation with finetuning for new assumptions, and consistency with the classic identifiability theory. \n\n**Weakness**\n\nI don’t see critical weaknesses in this work, here are some minor issues and questions.\n\n* The baselines selected here are all deep-learning or optimization-based methods, which are hard to learn and might be unstable. The authors should consider including some typical casual discovery algorithms that are common and stable as baselines (PC, GES).\n* The tolerance of changing the subset algorithm in the inference procedure is interesting. Would this be because the algorithms swapped don’t have critical conflicts with the test data properties (LiNGAM would be an extreme one to test)? Or it is because the aggregator is robust enough that the importance of local estimators are not that high?\n* The current implementation of independent test and score functions in the SEA are FCI and GIES is Fisher Z and BIC (which are relatively linear and simple), would complicated score functions/independence tests help? Or the aggregator is enough to handle the complexity and local estimators can be simple and naive.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces SEA (Sample, Estimate, Aggregate), a general framework for causal discovery foundation model that combines the traditional casual discovery methods (e.g. GES, FCI, LiNGAM…) and modern deep learning approaches. This work contributes a foundation model approach that in three steps: (1). Sample subsets from the observational data, (2). Apply classical causal discovery methods to acquire subsets’ graph estimates and acquire global statistics, (3). Aggregate local results via pretraining a neural-network-based aggregator to predict the whole graph.  The authors provide a theoretical analysis showing that the pretrained aggregator could recover the consistent union from local estimates. They conduct comprehensive experiments to demonstrate the generalizability, efficiency, and identifiability of the proposed framework.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The author should consider benchmarking with the classical causal discovery algorithms (PC, GES), especially in large-scale cases where optimization-based methods are usually unstable and ineffective.", "broader_impact_concerns": "The results generally look good and this work is potentially helpful in casual discovery for scientific discovery. This paper would benefit from a discussion of potential misuse and this framework's limitations in sensitive domains where causal discovery could impact decision-making (e.g., healthcare).", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Some concerns about the experimental setup were raised by reviewers which were satisfactorily addressed by the authors. Various clarity and completeness concerns were also addressed in the revision. After the revision, all reviewers agreed to accept the paper as-is.
Decision: Accept (poster)