Large-Scale Targeted Cause Discovery via Learning from Simulated Data
Authors: Jang-Hyun Kim
Claudia Skok Gibbs
Sangdoo Yun
Hyun Oh Song
Kyunghyun Cho
Venue: Accepted by TMLR ()

Abstract: We propose a novel machine learning approach for inferring causal variables of a target variable from observations. Our focus is on directly inferring a set of causal factors without requiring full causal graph reconstruction, which is computationally challenging in large-scale systems. The identified causal set consists of all potential regulators of the target variable under experimental settings, enabling efficient regulation through intervention. To achieve this, we train a neural network using supervised learning on simulated data to infer causality. By employing a subsampled-ensemble inference strategy, our approach scales with linear complexity in the number of variables, efficiently scaling up to thousands of variables. Empirical results demonstrate superior performance in identifying causal relationships within large-scale gene regulatory networks, outperforming existing methods that emphasize full-graph discovery. We validate our model's generalization capability across out-of-distribution graph structures and generating mechanisms, including gene regulatory networks of E. coli and the human K562 cell line.

Published in Transactions on Machine Learning Research (10/2025)
Large-Scale Targeted Cause Discovery via Learning from Simulated Data
Jang-Hyun Kim"
jonghyun@mmlla..nno....
Deportment of Computer Science, Seoul National University
Claudia Skok Gibbs
csg387 edu
Center for Data Science, New York University
Sangdoo Yun NAVER AI Lab
sangdoo.jun www../yuwnnwveevrr..coom
Hyun Oh Song
hyunoh@mmlaabssna. oc.
Department of Computer Science, Scoul National University
Kyunghyun Cho
kgunghyun.cho&myy.uo
Center for Data Science, New York University Prescient Design, Genetech
Reviewed on OpenReview: https: // openreview. net/ forum? NVgg9gg2
Abstract
We propose a novel machine learning approach for inferring causal variables of a target variable from observations. Our focus is on directly inferring a set of causal factors without requiring full causal graph roconstruction, which is computationally challenging in largo- scale systems. The identified causal set consists of all potential regulators of the target variable under experimental settings, enabling efficient regulation through intervention. To achieve this, we train neural network using supervised learning simulated data to in- fer causality. By employing a subsampled-ensembll inference strategy, our approach scales with linear complexity in the number of variables, efficiently scaling up to thousands of variables. Empirical results demonstrate superior performance in identifying causal rela- lionships within large-scale gene regulatory networks, outperforming existing methods that emphasize full- full-graph discovery. We validate our model's generalization capability across out-of-distribution graph structures and generating mechanisms, including gene regulatory networks of E. coli and the human K562 cell line. Implementation codes are available at https://github com/ snu mllab/Targeted- Cause-
1 Introduction
Identifying causal relationships among variables is fundamental challenge in machine learning, with widespread applications in generative modcling, system explanation, and variable control (Pcarl, 2009). Conventional methods have atlempled to infer causal structures using statistical tests or by fitting proba- bilistic models to observations (Spirtes et al., 2001; Chickering, 2002). However, the exponential complexity of causal graphs renders the problem NP-hard (Chickering, 1996), posing significant challenges for large-scale systems (Zanga et al., 2022).
In large-scale systems, inferring the complete set of causal relationships often unnecessary, as various downstream tasks directly benefit from identifying only the causes that relevant to specific outcomes
* Work done while visiting New York University.

---

Published in Transactions on Machine Learning Research (10/2025)
(Magliacane al., 2016). For instance, in drug discovery, it is sufficient to identify causal transcription factors associated with disease-related genes, rather than mapping the entire gene regulatory network (GRN); to effectively guide therapeutic interventions (Huynh-Thu et al., 2010). Similarly, in epidemiology or policy evaluation, interventions can be effectively designed by pinpointing only the key causal drivers of infection rales or targeted outcomes, without modeling the full complexity of inter-variable dependencies (Glass al., 2013). These examples illustrate that identifying the causes of specific targets has practical utility while reducing the complexity inherent in discovering entire causal structures.
In this study, we propose scalable method for identifying the causal variables of target variable from observations in large-scale systems, by leveraging data simulators. Our ap- proach aims to identify both direct and indirect causes that can regulate the target variable under experimental condi- tions (Figure 1). Rather than reconstructing an entire causal graph; we focus on identifying set of causal factors, simpli- fying the estimation process while retaining practical utility. We introduce this problem setting as targeted cause dis- covery.
Target variable
Causes of target
An alternative to our proposed method involves inferring causal variables by traversing the ancestors of larget within a causal graph derived from existing causal discovery tech- niques. However, estimating complete causal graphs becomes computationally prohibitive for systems with thousands of variables (Zanga et al., 2022). Furthermore, imperfect inference from limited observations leads to error ac- cumulation at each step of ancestral traversal, resulting in exponentially increasing inaccuracies. We explore this issue in greater detail in Scction 3.
Illustration of targeted cause dis- covery in a causal graph. Instead of inferring the full causal graph structure, we identify a of causal variables for a target.
Figure 1:
To overcome these limitations, we propose an end-to-end machine learning approach that directly estimates a set of causal variables from observational and experimental data. Our method trains deep neural network simulated data to learn causal discovery algorithm that generalizes to unscen causal structures (Ke et al., 2023). This data-driven approach allows the model to implicitly capture assumptions embedded in the simulated data. A key advantage of our method is to learn complex causal mechanisms from data, such as differential equations, without requiring explicit formal modeling in score-based causal discovery methods (Chickering, 2002; Zheng et al., 2018). We employ the Transformer architecture for our causal discovery model (Vaswani et al., 2017), building on its demonstrated effectiveness in previous works (Lorch et al., 2022). Nonetheless, the quadratic complexity of Transformer's attention mechanism presents computational challenges. To mitigate this, we introduce subsampled-cnnembble inference strategy that scales lincarly with the number of variables, enabling efficient application to large-scale problems involving thousands of variables.
Empirical evaluations demonstrate that our met hod effectively identifies causal relationships within complex systems involving thousands of variables, under the availability of high-fidelity simulators (Section 5). We demonstrate our model's capability to generalize from random causal structures roal-world GRNs, includ- ing E. coli, yeast, and the K562 human cell line (Dibaeinia & Sinha, 2020). We further assess the robustness and generalization capability of our approach through comprehensive evaluations on both synthetic and real- world datasets, examining range of graph structures, causal mechanisms, and levels of simulator fidelity, thereby underscoring its applicability in practical scenarios.
2 Preliminary and Related Work
2.1 Problem Formulation
We consider set of random variables V (x acyclic graph (DAG), g (V.E). We assume the data-generating process involves no latent variables, and for simplicity, we assume there is selection bias (Pearl, 2009). The joint distribution p(V) is defined
.at....,,y, x that has a causal structure represented by a dirccted
2

---

Published in Transactions on Machine Learning Research (10/2025)
II, p(x;/ Pa(x:)) where Pa(x;) means the set of parent variables of in g. We obtain observations of variables through ancestral sampling.
Definition 1. Causes for target node x; is defined as the set of ancestors of x4 in 9, denoted by An(x;)-
Definition 2. A variable x, is a marginal cause of x, if and only if 3c s.t.
p(x; do(x) c)) # p(x.).
The operator do(-) represents an intervention that fixes specific variables to predefined values during the ancestral sampling process (Pearl, 2009). Specifically, the observation for V under do(x, = c) follows 6e(x;) Heg p(x; Pa(x;)), where & is a Dirac delta function at c. Definition 2 implies that interventions on causal variables make a distributional change in the target variable.
Proposition 1. The set of marginal causes of x; is a subset of An(x)).
We provide the proof for Proposition 1 in Appendix A. This proposition implies that the ancestor set An(x:) forms necessary condition for identifying variables that influence the target variable x, in experimental settings. Specifically, in large, sparse networks such as GRNs, efficiently identifying these influential variables can facilitate effective control of target variables. We refer to the task of identifying the set of causes An(x,) of a larget variable x; as largeled cause discovery.
2.2 Amortized Variational Inference for Causal Structure Discovery
Classical causal discovery aims to infer the causal structure G from observation. Lorch et al. (2022) intro- duced a Bayesian perspective causal discovery by defining joint distribution over causal structures and mechanisms as o), where & represents mechanism and noise parameters. Given this prior distribution, the likelihood of observed variables V can be expressed as p(V (, Marginalizing out the mechanism parameters Φ yields the posterior distribution p(g V), which is typically intractable. Learning- causal discovery approaches tackle this intractability by approximating the posterior using amortized variational inference parameterized as (9 VV) (Lorch et al., 2022; Ke et al., 2023).
Rather than estimating the entire graph structure, we specifically focus on inferring the ancestor set An(x:) of target variable x,, represented by the posterior p(An(x;) | V). This ancestor distribution is inherently determined by the posterior over the entire graph P(G V). One possible method is to first fit go V) and then approximate p(An(x;) V) through graph sampling, resembling approximate Bayesian computation (Sisson et al., 2018). However, in large-scale scenarios, such sampling-based methods based on the graph posterior distribution become computationally expensive and present additional technical challenges, as detailed in Scction In our approach, we directly infer the posterior distribution over the ancestor set by optimizing variational parameters 0 of variational distribution qo (An(x,) V)).
2.3 Related Work
Causal discovery Conventional causal discovery methods aim to infer the exact causal graph structure G from observations (Spirtes et al., 2001; Chickering, 2002). These approaches can be broadly categorized into constraint-based and score-based methods. Constraint-based methods use conditional independence tests and formalized directional decision rules to identify causal relationships (Spirtes et al., 1995; 2001). However, they require independence testing over combinatorial sets of variables, leading to exponential complexity the number of variables (Colombo et al., 2014). Subsequent works have sought to reduce this complexity by leveraging sparsity assumptions in causal structures with Gaussian mechanisms (Kalisch & Bühhlan, 2007) or by improving computational efficiency through parallelization (Le et al., 2016). Score-based methods, on the other hand, optimize the goodncss-of-iit of graph structures while balancing complexity constraints (Chickering, 2002; Hauser & Bühlmann, 2012; Zheng et al., 2018; Brouillard et al., 2020). To navigate the combinatorial search space of causal graphs, these approaches rely on specific assumptions about graph structures and data-generating mechanisms (Lopez et al., 2022). An alternative line of rescarch focuses on estimating the topological ordering of causal variables, which helps reduce optimization complexity (Reisach et al., 2021; Sanchez et al., 2023).
3

---

Published in Transactions on Machine Learning Research (10/2025)
Local approaches To address the computational challenges of estimating full causal graphs, some efforts have focused on inferring the local causal structure of target node, such as the Markov blanket (Aliferis et al., 2010; Wu et al., 2019), parent set (Gao & Ji, 2015), or parent-cause set (Yin et al., 2008). However, these approaches still suffer from an exponential search space, requiring efficient algorithms for large-scale settings (Gao Ji, 2015). Magliacane al. (2016) introduced the problem of estimating ancestral structures. The key difference from our work is that their approach considers the ordering of ancestral nodes (or canses), whereas we do not impose any ordering constraints. By removing the order constraint, we reformulate our training objective as binary classification mask-determiningg whether a given node is a cause of the target based on observations of all hther nodes. This simplification enables us to develop a scalable method for both training and inference.
Learning-base based approaches Recent efforts have introduced learning- based approaches that leverage large datasets and computational power for causal discovery (Lopez-Paz al., 2015; Löwe et al., 2022; Lorch et al., 2022; Ke et al., 2023; Wu et al., 2024). These methods generate synthetic data with known ground-truth causal graphs and train neural networks to infer graph structures from observations (Torch et al., 2022). In this work, we analyze the generalization performance of these data-driven approaches in large-scale, complex systems. Specifically, we propose a novel strategy that identifies a set of causal variables for a target with linear complexity, enabling efficient scaling to thousands of variables.
3 Targeted Cause Discovery versus Causal Structure Discovery
Targeted cause discovery offers several technical advantages over causal structure discovery, particularly in addressing the shortcomings of conventional methods when the primary ob- jective is to identify the causes of specific target. In this sec- tion, we compare our approach, which directly estimates the set of causes, with an alternative method that infers causes by traversing the ancestors of the larget within an estimated causal graph.
Structure
Cause
False negative rate
1.0
0.8
0.6
0.4
0.2
0.0
1) Error propagation in estimation: With a limited num- ber of observations, causal structure discovery methods have prediction errors (Lorch et al., 2022). When causes are in- ferred from an inaccurately estimated causal graph, these errors propagate exponentially: If a method has an expected prediction error rate of e per edge, the error rate for estimat- ing causes at a distance d from the target is approximately Our approach mitigates this issue by directly inferring causality among distant variables, avoiding the er- ror propagation inherent in graph-based traversal. Figure 2 presents empirical measurements of cause prediction error rates across varying cause-effect distances, where we com- pare methods estimating the full causal structure with those identifying only a set of causes. To ensure a fair compar- ison, we use the same model architecture and dataset for both mothods, diffcring only in training objectives and in- ference strategies. The results indicate that the error rate of method identifying causes remains stable regardless of distance, whereas the error rate of a method trained to es- timate the causal structure increases as the cause-effect dis- tance grows
1
2
3
4
5
6
8
Cause-effect distance
Figure 2: Targeted cause discovery error rate as a function of shortest path length between nodes. Cause refers to our method, directly estimating the causes of target. Struc- ture denotes method that infers causes from an estimated causal graph. Both methods use the same model architecture and dataset but differ in training objectives and inference strategies. Appendix B.6 provides detailed experimental setting.
1
(1 (1--e) e)-
Table Averaged ratio of causes (or parents) per node, i.e., |An(x)//n (or [Pa(x)]/n). We draw statistics from 10 graphs, each with 1000 nodes and an average in-degree of 2.
Graph type
Parent
Cause
Erdös Erdös-Rényi
0.2%
1.1~1.2%
2) Tochnical challenges training: In learning-basch causal discovery, neural networks are typically trained predict the adjacency matrix of causal structure (Ke et al., 2023;
Scale-free
0.2%
0.4~2.22%
Gene regulatory
0.2%
0.5~1.1%
4

---

Published in Transactions on Machine Learning Research (10/2025)
Lorch et al., 2022). However, due to the sparsity of causal graphs, the resulting binary adjacency matrix is highly imbalanced, posing significant. training challenges. For instance, in the E. coli GRN, there are only about 2.3 edges per node among 1,565 nodes (Dibaeinia & Sinha, 2020). Such an extreme imbalance is known to hinder neural network training by providing sparse learning signals from the loss function (Kaur et al., 2019). In contrast, our approach trains neural networks to predict ancestors rather than the full adjacency matrix. This formulation mitigates the sparsity issue as shown in Table 1, thereby improving training stability and effectiveness.
3) Subsampled-inffer infference guarantee: Targeted cause discovery enables subsampled inference, allowing the identification of causal relationships using only subset of the system's variables (Proposition 2). Specifi- cally, a cause of a target variable x; remains a cause within a subsampled variable system (Bongers et al., 2016). This property is particularly useful in large-scale settings, where processing data from all variables is computationally infeasible. Leveraging this property, we propose an efficient algorithm capable of scaling to thousands of variables, as detailed in Section 4.2. It is worth noting that subsampled inference is non-trivial for causal structure discovery, as the direct causal strutueee can change within a subsampled variable system. We provide proof of this proposition in Appendix A.
Proposition 2. For any variable subset V C s.t. containing x; and all root variables in G, let An(x;; V) x4 in the system consisting of V, following the marginalization rule of Bongers et al. (2016). Similarly, let Pa(x;; V) denole the sel of parents of x; within V. Then, An(x; V) An(x:)nV. However, a counterezample exists for parents: Pa(x;; V) + Pa(x;) nv.
denote the set of ancestors xf
4 Method
In this scction, we present our method for Targeted Cause Discovery with Data-driven Learning, termed TCD-DL. We consider a system with n variables V = (xi,..., { having an underlying causal structure G. We denote the observation data as X E RTx*, where m is the number of observation samples. In the interventional setting, we define a boolean matrix M E 10,1)m where I indicates the occurrence of interventions.
Our objective is as follows: Given target variable x, V, predict label 4 E 10,1)" from the observation X and intervention matrix M, where 4Dj]] 1 indicates that xj is cause of x,, i.e., Xj E An(x:). We adopt a binary labeling approach, disregarding proximity in the causal graph, as quantifying causal contributions based on proximity is non-trivial. To tackle this problem probabilistically, we estimate a continuous cause score vector 8; E R"", where st[]] represents the unnormalized likelihood of x, being a cause of x,. A higher score corresponds to a greater likelihood of causality. This continuous relaxation enables the use of gradient- based optimization techniques, which are well-suited for large-scale settings (Bottou, 2010). To obtain a set of estimated causes, we apply thresholding to 8i, considering values greater than zero causes.
4.1 Data-Driven Learning
A straightforward approach to discovering causality involves conducting interventions on every single variable with a statistically sulficient number of trials confirm the hypothesis. However, it is often impractical to intervene at every single variable due to experimental limitations and the high costs associated with conducting a sufficient number of trials (Addanki et al., 2020).
To address this challenge, we develop a parameterized model fo capable of inferring causality among variables from observations X E IR= of arbitrarily limited size. Given the target variable index i, the model processes the entire dataset X with intervention matrix M and returns a cause score vector 8; E R as
s; = MX
(inference)
By leveraging observations from all nodes, we can infer causal relationships more accurately than relying solely on the observations single pair of variables. We discuss this in more detail in Section 5.4.
We implement fo as deep neural network and train on simulated data D = {(X,, Mc. 9x) k E I)} using supervised learning approach. Here, I denotes an index set, and X represents an observation dataset
5

---

Published in Transactions on Machine Learning Research (10/2025)
Training
Inference
Dataset (X)
Subsampled dataset (Xy
Partial cause score
Subsampled dataset (Ny,))
Partial cause score
Estimated cause score (s.)
Dataset (X)
Subsample
fo
Variables
Subsample
fo
Average
Variabia
Target (x))
Observations (m)
Compute
Subsampled
Label (L)
label (lay)
Target (x.)
Subsample
1
Observations (m)
Figure 3: Overview of our method. The left figure depicts a single training step, while the right figure illustrates the inference procedure with multiple subsampling and ensembling. Note that the intervention matrix undergoes the same subsampling as X, resulting the stacked input [Xv,o, My.ol of shape n! m which is then fod into the model fa- We omit the intervention matrix the figure for simplicity.
sampled from a synthetic causal graph 9k with an intervention matrix Mg. We generate synthetic datasets using simulator on random graphs, detailed in Section 5. Empirically, demonstrate that the model generalizes to unseen causal structures and causal mechanisms, achieving strong performance datasets (Section 5.3). For causal graph Ga with variables, we obtain label l.,1 E {0, 1)ax for each variable = 1,.... using topological sorting. The label is defined such that [ = 1 if the j-th variable is a cause of the i-th variable in Gk. Our training objective, with prediction (X Ms, i) E R is
real-world
minimize ER-IBin<(I...,, [f(fo(Xx Ms. 4,,)),
(training)
where L the loss function. In this study, we use binary cross- cross-entropy with logits for (Wei al., 2022). ignoring the loss at the target node's position.
Acyclicity. Our problem focuses estimating a set of causes rather than inferring the full causal graph structure, and therefore does not require explicitly enforcing acyclicity constraint. Moreover, super- vised lcarning approach leverages training labels extracted from DAG, which contains no cycles. Training on such data provides implicit guidance for our model lo generale causal predictions that are consistent and non-contradictory.
Model architecture. The model fe comprises two sequential modules, a feature extractor ge and a score calculator h, designed to optimize compute efficiency. The feature extractor go processes the stack [X E m2xmxx produce features F1 E Rnxd where each variable index i corresponds to two d- dimensional features F [1]. Pa[i] E Rdd We employ an axial-Transformer the feature extractor (Ho et al., 2019), matrix-input variant of Transformers that has been a conventional architecture in prior learning- based approaches (Lorch et al., 2022; Ke et al., 2023). Notably, our framework is compatible with general model architectured ures designed for matrix-shaped inputs. A detailed description of the model architecture is provided in Appendix B.1.
The score calculator h computes the cause score vector s; for a target index using dot-product between features 1((F1, i) F2 [2 RR,, where F1 Raxd and F2li] E Rd. To sum up, So(X, M, h(go(X, a) h(F,, i2, The use of separate feature extractor allows for the reuse of features across multiple target indices enhancing training efficiency with batch data processing.
4.2 Subsampled Inference for Scaling Up
The axial Transformer comprises three main operations: two attention layers over each variable and obser- vation dimension, and feed-forward layer (Ho et al., 2019). For input X Raxm, the complexities of

---

Published in Transactions on Machine Learning Research (10/2025)
attention layers are (n?mm) and O(nm²), while the complexity of the forwafo layer is O(nm). The quadratic complexity of the attention mechanism poses challenges for large inputs in terms of computational time and memory usage. In Appendix B.4, we present a detailed analysis of the experimental scales and discuss the computational challenges.
To address this issue, we propose a subsampled inference strategy (Figure 3), supported by Proposition 2, which estimates the causes of a target variable using subsampled variables and observations. Specifically, define subsample sizes n' < n for variables and m' < m for observations, set independently of and m, based on available computing resources. Given target variable X,, randomly sample a subset of variables V C V with [V] = n', ensuring X; E V. We then extract the corresponding observation matrix Xv E Rn'xm and intervention matrix My E {0, 1/11"""' To avoid distortion from interventions on variables outside the subsampled set V, we select only those observations (columns of Xv) where no variables in V V are intervened. From these, we randomly subsample m' observations, yielding Xy.0 E R xm' and Mv.0& E 10,1] / xm', where O denotes the selected observation indices. We denote this subsampling process as V,O ~ Mi).
We then infer the causal effects on x; from variables in V using fo My.o. i) E R Repeating this process across multiple subsamplings of variables and observations, we aggregate and average the results to compute the final cause score vector 8;- We call this process subsampled-onsemblee inference. In practice, perform subsampling to ensure that each variable is considered identically T times (Algorithm 2). For variable Xj E V, we denote the estimated cause score to x, as fo(Xv.o, B)D] (Here, j refers to the index in V, but for clarity, we retain the same index notation within V.) The final estimation is then
s + =
My,o x v
(subsampled-ensemblee inference)
For training, we apply the identical random subsampling strategy on inputs and target labels. For X and a target variable index i, we denote the subsampled data as X E qg''xm" and the corresponding target label as I,.iv E {0, In practice, we generalize the results from Proposition 2, defining An(x;: V) = AV for arbitrary subsets of variables. Specifically, each label entry is the corresponding variable belongs to An(x;) n V, and otherwise. This ensures consistent labeling, i.e., for Xy E An(x,, the variable x, receives label of 1 regardless of V, providing consistent supervision signals throughout training. The subsampled version of our training objective becomes
minimize EL..TEBrtt.... bog.:a/h...nn..ooo.oo...oaaaaaad Mev.o,
(subsampled training)
We optimize fo using stochastic gradient descent with the AdamW optimizer (Loshchilov & Hutter, 2019). Algorithms and 2 describe pseudo codes of our final training and inference algorithms. We leave detailed hyperparameters in Appendix B.2.
Algorithm complexity. Our algorithm reduces the inference complexity from quadratic to linear with respect to the number of variables n. Specifically, we fix the subsample size n' as a constant across all experiments, regardless of the variable size n. By splitting the inputs into fixed-size chunks of n'' and repeating the Transformer computation n/n' times, we achieve a computational complexity of O(n" ((n".n/n')) ) = O(nn'), reducing the original complexity of O(n²). Proposition 3 provides detailed complexity analysis of our inference algorithm, with a proof included Appendix A.
Proposition 3. Let n' and m' denote the subsampled variable and observation sizes, and let T denote the ensemble size. For a dataset X with n variables, the inference complezity of our algorithm is O(nm'T n + m'))
7

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 1 Training (batch version)
Algorithm 2 Inference
inputs: D = {(Xx, Ma, G&) k E
inputs: X E R=xm M E {0. pxm, target in-
parameters: subsample sizes n' and m', batch
dex i
size b
parameters: subsample sizes n' and m', #en- semble 7
initialize 0
repeat
initialize 8; + O,
X.Y + 0,0
for = 1 to T do
for to do
I I permute
n \ {i} U0 j= where b b= [ and |T S n'
k < sample from I
split I =
V,0 < subsample given XR, M&
for j = to b do
X' - XU [U MAV.O. X4 V}
1 < U U [i] and V + fxx k E 1,}
y + y u
0 - subsample given V,X,M M
end for
s; [ + s, s,[T,] + Mo.
g - calculate gradients VoL(so(),,) on X,y 0 update using gradients g
end for
end for + 8:/T return 8;
until convergence
return 0
5
Experiment
We present an experimental analysis of out-of-distribution (OOD) settings for targeted cause discovery. In Section 5.2, we verify that our model, trained on random graphs, effectively identifies causal relationships in biological networks consisting of several thousand variables. In Section 5.3, we examine the model's general- ization capabilities with respect to novel causal mechanisms and varying noise levels, including analyses on real-world human gene dataset. Finally, in Section 5.4, we analyze interventional settings and assess sen- sitivity to hyperparameters. To compare against conventional algorithms such PC and GIES, we conduct additional small-scale experiments detailed in Appendix C.2.
5.1 Setup
Simulated dataset. We generate the training set and test set using the SERGIO GRN simulator (Dibaeinia & Sinha, 2020). This simulator produces single-cell gene expression data, modeling the stochastic nature of transcription based user-provided causal graph structure. We conduct experiments over varying levels of simulator's observational fidelity, where higher fidelity yields expression data closer to the population parameters. Details about the simulator, including the generation mechanisms and the definition of fidelity levels, are provided in Appendix B.3. For the training dataset, we use random graph structures including Erdös-Róny: (ER), Sealo-Free (SF), and Stochastic Block Model (SBM) with 1,000 variables (Drobyshevskiy & Turdakov, 2019). The test. data generated from biological structures of E. coli (1,565 genes) and yeast (4,441 genes) as obtained from Marbach et al. (2009).
Intervention. We adopt the intervention settings from Lorch al. (2022). Specifically, we perform single- variable interventions by knocking out gene expressions, i.e., setting their transcription rates to zero. We gen- erate 10 samples per intervention, along with 500 observational samples (details provided in Appendix B.3). In Figure we analyze model performance under various intervention scenarios, including different ratios of intervened nodes and multi-variable interventions.
Identifiability. Our method shares similar identifiability considerations with existing learning bassedg-based ap- proaches (Lorch et al., 2022; Ke et al., 2023). In particular, we focus experimental settings with single-node interventions, where the true causal structure is theoretically identifiable in the limit of infinite intervent ional data from each variable (Eberhardt et al., 2006). In practice, however, the number available interventional samples is typically limited, precluding full identifiability guarantecs. To address this, we propose a proba- bilistic framework designed to provide calibrated estimates of causal relationships given the available data, rather than enforcing strict identifiability. This design reflects practical trade-off, enabling the model to
8

---

Published in Transactions on Machine Learning Research (10/2025)
[ Fidelity high 0 Fidelity medium
D Fidelity low
GENEES
AVICI
TCD-DL (ous)
100
1.00
False negative rate
90
0.98
AUROC (%)
80
70
0.94
60
0.78
50
0.76
40
Random
Correlation
Sort nregresss
PMF-GRN
GENIE3
DCD-FG
AVICI
TCD-DL (ours)
2
4
Cause-effect distance
(a) Performance across varying levels of simulator fidelity
(b) Error by causal distance
Figure 4: Benchmarking results. (a) Performance on E. coli GRN with 1565 genes over varying levels of simulator's observational fidelity. We provide AUROC, AP, and Fl score values, including standard deviations in Table 10. (b) The cause prediction error rate as a function of the shortest path length between variables a causal graph
generate causal predictions that are consistent with observed data, even low-sample regimes. Empirically, we find that our approach yields meaningful and robust causal inferences under realistic data constraints, as illustrated Figure 8.
Baseline. We comprehensively evaluate causal discovery methods drawn from various methodological frameworks. As baselines, we include a random guessing model, the absolute correlation score, and a regression- based method known as sortnregress (Reisach et al., 2021). Additionally, we compare against the score-based approach DCD-FG (Lopez et al., 2022), an improved method designed to address instability issues by DCDI (Brouillard et al., 2020) and NO-TEARS (Zheng et al., 2018) in large-scale settings (Kaiser & Sipos, 2022). We evaluate the learning-baid based approach AVICI (Lorch et al., 2022) by utilizing their re- leased model, which was trained on the same simulator and comparable amount of data as ours. It is worth noting that AVICI closely aligns with the concurrent method CSIVA (Ke et al., 2023), while CSIVA does not provide source code or trained models. Conventional algorithms such as PC and GIES (Spirtes et al., 2001; Hauser & Bühlmann, 2012) are evaluated separately in small-scale experiments due their limited scala- bility (Appendix C.2). We assess representative GRN inference methods, including the tree-based method GENIE3 (Huynh-Thu et al., 2010) and the linear factor model PMF-GRN (Skok Gibbs et al., 2024). All mentioned methods, except PC, compute likelihood scores indicative of (direct) causal relationships among variables. We utilize these likelihood scores to construct baseline causal score vectors
Evaluation metric. We evaluate the targeted cause discovery performance variables having at least one causal variable. For each target variable, we compare the predicted cause score vector against the ground-truth binary label, where indicates a causal relationship between variables. We employ AUROC, Average Precision (AP), and F1 score for this binary classification task (Rainio et al., 2024). To calculate the Fl score, we threshold the cause scores to match the number of positive predictions with the ground-truth labels. We measure average performance across all valid variables (those with at least one cause) in test datasets. We obtain statistics using expression data with 10 different. random seeds for sampling.
5.2 Generalization on Unseen Causal Structure
Benchmarking. Figure 4-a illustrates the performance of targeted cause discovery on the E. coli GRN, which has a novel causal structure not observed during training. Our approach consistently achieves the best performance by a large margin, demonstrating the effectiveness our data-driven learning approach. The results reveal the shortcomings of existing methods relying on specific assumptions. Lincar models (correlation, sortnregress) fail to identify causality in our settings with complex generation mechanisms. As sanity check, observe that correlation achieves 70.7% AUROC causal graphs with linear generation

---

Published in Transactions on Machine Learning Research (10/2025)
AVICI (best)
TCD-DL (ours)
GENIE3 (best)
TCD-DL (ours)
100
100
AUROC (%)
90
AUROC (%)
90
80
70
60
60
50
50
Graph
Craph Moch.
Craph
Graph
Graph | Mech.
Graph
+ +Nocee
+
(a) E. coli
(b) Yeast
Figure 5: OOD performance. Box plots illustrating targeted cause discovery performance on novel graph structures, mechanism parameters (mech.), and noise configurations unseen during training. For the coli dataset, AVICI is the best-performing baseline method, whereas GENIE3 achieves the best results for yeast.
mechanisms, showing moderate performance under valid assumptions. The score-based approach (DCD-FG) performs nearly at random, likely due to invalid assumptions about the factorizability of graph structures and simplistic generation mechanisms (Lopez et al., 2022). Our method outperforms the learning- basod approach (AVICI), which focuses on direct causality and does not leverage subsampled inference, thereby demonstrating the effectiveness of our scalable algorithm.
Robustness to causal distance. To gain deeper insight into the observed performance improvements, measure the felse-ne rale as a function of the shortest path length in the ground-truth causal graph (Figure 4-b). To ensure fair comparison, we threshold the causal scores of the best-performing methods so that each produces an identical number positive prodictions. The results indicate that bascline methods exhibit increasing error rales as causal distance grows, whereas our method maintains consistent performance. These findings highlight key advantage of our approach that reliably identifies both proximal and distant causes without performance degradation.
Error analysis. While model demonstrates strong perfor- mance, there still remain errors, particularly as simulator fidelity decreases. To investigate the sources of these errors, we com- pare our model's performance on random graphs sampled from the training setting (i.e., validation set) to its performance on E. coli graphs (i.e., test set). Table 2 reveals that both valida- tion and test performance decline as simulator fidelity decreases. This parallel degradation suggests that the primary source of error not a generalization issue, but rather stems from other factors. This finding raises questions about causal identifiabilit in low-fidelity scenarios, indicating that the aset itself may lack sufficient information for accurate causal inference.
Table AUROC (%) on random graphs (validation) and E. coli GRN (test).
Data \ Fidelity High
Medium
Low
Validation Test
92.6
83.3
70.1
94.6
81.7
71.5
Runtime measurement. The runtime of our inference algorithm for processing each target variable is 2.5 seconds for F. coli (1,565 genes) and 7.8 seconds for yeast (4,441 genes), as measured with an NVIDIA RTX 3090 GPU. These results highlight that our model, with subsampled-onsemblee inference, operates within seconds for large-scale systems. We provide a comparison of the runtime with baseline methods in Appendix C.1, where some baselines, such as DCD-FG and PMF-GRN, take several hours for inference.
5.3 Extended Out-of-Distribution Analysis
OOD simulator configuration. We further study the generalization capabilities of our model by testing on causal mochanism and noise configurations that differ from the training. We adopt the configurations from Lorch et al. (2022), as described in Appendix B.3. Figure shows that our method largely outperforms the baselines across all settings, demonstrating robust generalization. Notably, when only the graph structure
10

---

Published in Transactions on Machine Learning Research (10/2025)
differs from training, model shows high prediction capability. However, generation mechanisms diverge from the training setting, the performance variance increases. These findings show both the strengths of our approach and the challenges inherent in generalizing to diverse generation mechanisms.
Real-world human cell analysis. We test our simulator- trained model in a real-world scenario using a Perturb-seq dataset derived from the K362 cell line of patient with chronic myeloge- nous leukemia (Replogle al., 2022). We focus on the gene MYC, key oncogene involved in cell proliferation, growth, and apop- tosis, frequently overexpressed in cancers (Dhanasekaran et al., 2022). We compute cause scores for 1,868 genes and select the top 20 genes as predicted causes of MYC expression (Appendix B.5). One major challenge here that the exact structure of the hu- man GRN has not yet been fully clucidatod. To evaluate our pro- dictions, we compare them with the following three sources: (1) The STRING database, which provides associations between pro- teins but does not indicate the causal direction (Szklarczyk et al., 2023). (2) Existence of literature that supports causality between a specific gene pair (Table 7). (3) Top 20 genes with the highest. expression correlations to MYC. Figure shows that our model demonstrates strong predictive accuracy, achieving 90% preci- sion compared to STRING. Notably, 30% of our predictions are validated by existing literature (5 greens and 1 yellow), demon- strating that our model uncovers meaningful causal relationships An interesting observation is that only the gene EEFIA1 shows high expression correlation with the target, suggesting that our method identifics novel causal factors not captured by correlation ranking. In Appeudix C.3, we provide additional analysis, including the quantitative comparison to corre- lation ranking and the predicted influence of causal genes. These results highlight our model's potential for real-world applications in understanding and manipulating gene regulation, particularly in the context of personalized medicine and targeted cancer therapies.
RPSI8
RPLM
TURN
HISPAR
ACTS
RPIZIA
NCL
RPSISA
MYC
NPSI
REPS14
RPL226
LST1
REFLAI
PTMA
RPL4
TCD-Dt.
TCD-DL. STRING
RSPSS
1 CCDaLL STRING, Litenatree TCD.DL STRING, Likenture, Concelatios
Figure 6: Causal genes of MYC identified by our method. Genes are categorized by validation from STRING, cxisting litera- ture, and top-20 expression correlations. We note that existing literature validates РТMA and RPS6 as effects of MYC.
Ablating training sources. To further assess the generalization capabilities of our approach, we perform ablation studies by excluding specific causal structures and mechanism types from the training data, and evaluating the resulting models. In Figure 7-a, we conduct experiments with various graph structures, including Erdös Rényi (ER), directional Scale-Free (SF-direct), Scale-Free (SF), and the Stochastic Block Model (SBM) (Drobyshevskiy & Turdakov, 2019). In Figure 7-b, we vary the causal generation mechanisms to include different analytic functions (linear, non-linear multi-layer perceptron (MLP), polynomial, and sigmoid), while maintaining scale-free network structure, following Wu et al. (2024).
Figure 7 illustrates the relative performance for each test case, scaled from 0 (random) to 100 (best). Both subfigures exhibit similar patterns, where diagonal entries have relatively lower performance, indicating a generalization gap when the testing data type is omitted from the training sources. Nonetheless, relative per- formance consistently exceeds 90, highlighting the robust generalization capability of our method. Notably, models trained on a diverse combination of all data types consistently achieve near-optimal performance. This data-scaling effect is consistent with observations from large-scale language models (Brown el al., 2020), suggesting that diversifying training data improves overall causal discovery performance.
5.4 Additional Analysis
Ratio of intervened nodes. We evaluate our model by varying the ratio intervened nodes and the number of intervention samples per node on the E. coli test datasct. To maintain a fixed dataset size, we adjust the number of observational samples accordingly each scenario. We select intervened nodes randomly, and use single model for evaluation across all settings. In this experiment, disable the dropout
11

---

Published in Transactions on Machine Learning Research (10/2025)
Test
Test
ER
9p 9F-diract
SE
SBM
Linear
MLP
Polynomial
Sigmoid
100
100
w/o ER
97.8
99.7
99.0
98.5
w/o Linear
91.1
99.2
100
99.4
95
98
w/o SF-direct
99.7
77.6
99.5
99.3
w/o MLP
97.4
92.2
99.4
98.4
90
96
Train
Train
w/o
100
100
94.6
100
w/o Polynomial
100
100
91.9
100
85
94
w/o SBM
99.1
99.8
99.8
98.9
w/o Sigmoid
96.9
95.1
99.2
93.2
80
92
Include all
99.8
99.8
100
99.7
Include all
97.3
98.3
99.2
98.4
75
90
Graph structure
(b) Generation mechanism
Figure 7: Training source ablation analysis. The metric used the relative AUROC, with 100 indicating the best model's performance on each test case and 0 corresponds to random prediction, i.e, 100(p Prandom) )/ (peest Prandom): where p denotes the AUROC score of given model.
100
100
AUROC (%)
90
AUROC (%)
90
AUROC (%)
90
Accuracy (%)
90
80
None
Cause
70
70
70
Target
70
Both
60
60
2
8
10
0.2 0.4 0.6 0.8
0 0.4 0.6 0.8
0
0.2 0.4 0.6
# Intervention
Intervened node ratio
Intervened node ratio
Intervened node ratio
(a) Intervention samples per node
(b) Ratio of intervened nodes
(c) Ratio of intervened nodes (categorized)
(d) Causal direction
prediction accuracy
Figure 8: Intervention analysis. Performance on the E. coli GRN test dataset with varying intervention settings. (a) Controlling the nunber of intervention samples per node in the test dataset. (b) Controlling the ratio of intervened nodes among all nodes. zero ratio indicates that only the observational samples are used. (c) Performance with categorized intervention types. We categorize target nodes into four cases and report average scores for each case: no intervention on the target and ancestors (None), intervention on least ancestor and not the target. (Cause), intervention the target and not ancestors ( (Targel), and intervention both (Both). (d) Causal direction prediction accuracy of our model across varying ratios of intervened nodes.
effect in the simulator (Appendix B.3), as the random removal of information can weaken interventional signals and lead mislcading analysos.
From Figure 8-a, observe that even with only one intervention sample per node (complemented by obser- vational data), the model achieves reasonable performance, reaching AUROC of 90.1%. As decrease the ratio of intervened nodes (Figure 8-b), performance decreases sublincarly. These results demonstrato that our method effectively leverages intervent ional information to infer causality. Furthermore, the model consistently provides meaningful predictions, even with limited intervention data, and surprisingly, performs modcrately well with purely obscrvational data, achieving 65.3% AUROC. In Figure &-c, we present cate- gorized performance based on whether the target node itself one of its causes has been intervened upon, while varying the ratio of intervened nodes. Interestingly, indicated by the None category in the figure,
12

---

Published in Transactions on Machine Learning Research (10/2025)
95
95
AUROC (%)
AUROC (%)
94
93
93
91
92
89
91
5
10
25
50
30
50
100
200
300
71'
(a) Ensemble size
(b) Subsampled input size
Figure 9: Hyperparameter analysis. Performance on E. coli GRN with varying (a) ensemble sizes T and (b) imput variable subsample sizes n'.
observe that even when neither the largel node nor its causes are intervened on, the model achieves moderate performance (~82% AUROC) by leveraging interventional information on other variables.
Multi-node intervention. We conduct experiments involving multi-node interventions by increasing the number of interventions sample from 1 to 3 during both training and testing while keeping all other configurations unchanged. In this scenario, we observe a slight performance drop from 94.6% to 93.7% AUROC on the E. coli test dataset. We suspect this decline results from increased complexity due to the combinatorial explosion possible interventions, making more challenging for the model fully capture relevant causal information. Nonetheless, exploring optimal intervention settings remains an important direction for future research.
Causal direction prediction. To verify that our model indeed infers causality beyond mere correlation learning, measure the causal direction prediction accuracy (Figure 8-d). Specifically, for every pair of nodes i and j with known causal relationship, the causal direction by comparing the cause scores s:[j] and sylر] predicted by our model. If s:Li] > sy[i], we predicted as the cause of i. Note that random guessing achieves an accuracy of 50% for this task. As shown in Figure 8-d, our model accurately predicts causal directions under the full-node intervention setting, while achieving 68% accuracy in the observational setting. This clearly demonstrates that our model indeed infers causality beyond mere correlation learning.
varying ratios of intervened nodes estimated
Algorithm hyperparameter. We analyze the impact design choices in our subsampled-ensemble infer- ence strategy by sweeping the ensemble size T and the subsampled variable size n' Algorithm 2. Figure 9-a demonstrates that performance improves as ensemble size incrcases, platcauing around 25. This result vali- dates the effectiveness of our ensembled inference approach. Figure 9-b illustrates the effect of subsampled variable size per input. Performance increases with input size up to 200, suggesting that processing larger variable sets through single Transformer forward pass allows the model to utilize richer relational infor- mation. However, performance declines for input sizes exceeding 300, indicating conflicting effects between input complexity and information richness. These results validate our approach of processing a subset. of variables, highlighting its effectiveness compared to processing all variables simultaneously.
6 Conclusion and Discussion
In this work, propose effective and scalable approach for targeted cause discovery, aiming to identify set of causal variables for given target. Our method trains neural network to learn causal discovery algorithms from simulated data. To handle large-scale systems, we introduce subsampled-cnsemble infer- ence strategy that achieves linear complexity with respect to the number of variables. Empirical results on gene regulatory networks demonstrate that our approach significantly outperforms existing causal discov- ery basclines, exhibiting strong generalization across diverse graph structures and generation mechanisms. By shifting the focus from explicit causal modeling to data-driven framework, our method aligns with recent. advances in deep learning. We anticipate further performance gains through data scaling within
13

---

Published in Transactions on Machine Learning Research (10/2025)
scalable framework. However, our reliance data-driven black-box models comes at the cost of reduced interpretability. Balancing causal interpretability with the advantages of data-driven methods remains crucial future direction for targeted cause discovery.
Broader Impact Statement
Our work introduces a scalable technical approach to targeted cause discovery in large-scale causal inference problems, specifically designed for computational efficiency and robustness in identifying relevant causal vari- ables. Positive impacts our method include improved computational performance and increased accuracy in practical causal inference tasks, enabling researchers efficiently analyze complex systems such as gene regulatory networks. Although this research primarily advances foundational machine learning techniques without direct societal applications, potential negative impacts might arise if applied irresponsibly- for ex- ample, misinterprotation of causal findings leading to incorrect conclusions in biological or medical contexts.
Acknowledgement
The work was supported by the National Science Foundation (under NSF Award 1922658). Kyunghyun Cho is supported by the Samsung Advanced Institute of Technology (under the project Next Generation Deep Learning: From Pattern Recognition to Al). Jang-Hyun Kim and Hyun Oh Song are supported by SNU-NAVER Hyperscale AI Center, Institute of Information & Communications Technology Planning Evaluation (ITTP) grant funded by the Korea government (MSIT) (No. RS2020-II200888 SW STAR LAB, Development. of deployable learning intelligence via self-sustainable and trustworthy machine learning), and the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS- 2024-00354033). Claudia Skok Gibbs is supported by an NSF-GRFP award (DGE-2234660), Kyunghyun Cho is the corresponding author.
References
Raghavendra Addanki, Shiva Kasiviswanathan, Andrew McGregor, and Cameron Musco. Efficient interven- tion design for causal discovery with latents. ICML, 2020.
Constantin F Aliferis, Alexander Statnikov, Toannis Tsamardinos, Subramani Mani, and Xenofon D Kout- soukos. Local causal and markov blanket induction for causal discovery and feature selection for classifi- cation. Journal of Machine Learning Research, 11, 2010.
Jimmn Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ar preprint arXiv:1607.06456 2016.
Pau Badia-i Mompel, Lorna Wessels, Sophia Müileer-ott, Rémi Trimbour, Ricardo O Ramirez Flores, Ricard Argelaguet, and Julio Saez-Rodriguez. Gene regulatory network inference in the era of single-cell multi- omics. Nature Reviews Genetics, 24(11):739-78) 754, 2023.
Stephan Bongers, Jonas Peters, Bernhard Schölkopf, and Joris M Mooij. Structural causal models: Cycles, marginalizations, exogenous reparametrizations and reductions. arXiv preprint arXiv, 2016.
Léon Bottou. Large-scale machine learning with stochastic gradient descent. COMPSTAT, 2010.
Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. NewrIPS, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, et al. Language models are few-shot learners. NewrIPS, 2020.
Wenan Chen, Yan Li, John Easton, David Finkelstein, Gang Wu, and Xiang Chen. Umi-count modeling and differential expression analysis for single-cell rma sequencing. Genome biology, 19, 2018.
David Maxwell Chickering. Learning bayesian networks is np-complete. Learning from data: Artificial intelligence and statistics 1996.
14

---

Published in Transactions on Machine Learning Research (10/2025)
David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 3, 2002.
Dominique Chu, Nicolae Radu Zabet, and Boris Mitavskiy. Models of transcription factor binding: sensitivity of activation functions to model assumptions. Journal of Theoretical Biology, 257, 2009.
Diego Colombo, Marloes H Maathuis, et al. Order-independent constraint- causal structure learning. Journal of Machine Learning Research, 15, 2014.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. CVPR, 2009.
Renumathy Dhanasckaran, Anja Deutzmann, Wadic D Faabauaa-Forrandez. Aida S Hansen, Arvin M Gouw, and Dean W Felsher. The mye oncogene- the grand orchestrator of cancer growth and immune evasion. Nature revieus Clinical oncology, 19(1), 2022.
Payam Dibaeinia. and Saurabh Sinha. Sergio: a single-cell expression simulator guided by gene regulatory networks. Cell systems, 11, 2020.
Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
Mikhail Drobyshevskiy and Denis Turdakov. Random graph modeling: survey of the concepts. ACM compuling surveys (CSUR), 52, 2019.
Frederick Eberhardt, Clark Glymour, and Richard Scheines. N-1 experiments sulfice to determine the causal relations among variables. Innovations in machine learning: theory and applications, 2006.
Ayako Egoh, Shin Nosuke Kanesashi, Chie Kanei-Ishii Teruaki Nomura, and Shunsuke Ishii. Ribosomal protein 14 positively regulates activity of ac-myb proto-oncogene product. Genes to Cells, 15(8):829 881, 2010.
Brunangelo Falini, Lorenzo Brunetti, Paolo Sportoletti, and Maria Paola Martelli. Npml-mutated acute myeloid leukemia: from bench to bedside. Blood, The Journal of the American Sociely of Hematology, 136 (15), 2020.
Tian Gao and Qiang Ji. Local causal discovery of direct causes and effects. NewrIPS, 2015.
Thomas A Glass, Steven N Goodman, Miguel A Hernán, and Jonathan M Samet. Causal inference in public health. Annual review of public health, 34(1), 2013.
Dalian Gong, Xinxu Rao, Ziqian Min, Xiaowen Lin, Huan Xin, Peijun Zhou, Lifang Yang, and Dan Li. Ube2s targets rpl26 for ubiquitination and degradation to promote non-small cell lung cancer progression via regulating c-myc. American Journal of Cancer Research, 13(8). 2023.
Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equiva- lence classes of directed acyclic graphs. The Journal of Machine Learning Research, 13, 2012.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint XrXi:: 122.12211 12180, 2019.
Zhe Hong, Chengdang Xu, Shengfeng Zheng, Xinan Wang, Yiran Tao, Yao Tan, Guowen Lin, Denglong Wu, and Dingwei Ye. Nucleophosmin cooperates with brd4 to facilitate c-mye transcription to promote prostate cancer progression. Cell Death Discovery, 9(1), 2023.
Vān Anh Huynh-Thin Alexandre Irrthum, Louis Wehenkel, and Pierre Geurts. Inferring regulatory networks from expression data using tree- based methods. PloS one, 5, 2010.
Marcus Kaiser and Maksim Sipos. Unsuitability of notears for causal graph discovery when dealing with dimensional quantities. Neural Processing Letters, 2022.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Markus Kalisch and Peter Bühlman. Estimating high-dimensional directed acyclic graphs with the pe- algorithm. Journal of Machine Learning Research, 8, 2007.
Harsurinder Kaur, Husanbir Singh Pannu, and Avleen Kaur Malhi. A systematic review on imbalanced data challenges in machine learning: Applications and solutions. ACM computing surveys (CSUR), 52, 2019.
Nan Rosemary Ke, Silvia Chiappa, Jane Wang, Anirudh Goyal, Jorg Bornschein, Melanie Rey, Theophane Weber, Matthew Botvinic, Michael Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. RCI.R, 2023.
Thuc Duy Le, Tao Hoang, Jiuyong Li, Lin Liu, Huawen Liu, and Shu Hu. A fast pe algorithm for high dimensional causal discovery with multi-core pcs. IEEE/ACM transactions on computational biology and bioinformatics, 16(5):1483-1495 2016.
Mingli Li, Lu Yang, Anthony KN Chan, Sheela Pangeni Pokharel, Qiao Liu, Nicole Mattson, Xiaobao Xu, Wen-Han Chang, Kazuya Miyashita, Priyanka Singh, al. Epigenetic control of translation checkpoint and tumor progression via ruvbll-eeflal axis. Advanced Science, 10(17), 2023.
Jiayu Liang, Zhihong Liu, Zijun Zou, Xiangxiu Wang, Yongquan Tang, Chuan Zhou, Kan Wu, Fuxun Zhang, and Yiping Lu. Knockdown of ribosomal protein s15a inhibits hurnan kidney cancer cell growth in vitro and in vivo. Molecular Medicine Reports, 19(2), 2019.
Yi-Te Lin, Hsing Pang Lu, and Chuck C-K Chao. Oncogenic c-myc and prothymosin-alpha protect hepato- collular carcinoma cells against sorafenib-induced apoptosis. Biochemical pharmacology. 93(1), 2015.
Romain Lopez, Jan-Christian Hütter, Jonathan Pritchard, and Aviy Regev. Large-scale differentiable causal discovery of factor graphs. NewrIPS, 2022.
David Lopez-Paz, Krikamol Muandet, Bernhard Schölkopf, and lliya Tolstikhin. Towards learning theory of cause-effect inference. ICML, 2015.
Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized inference for causal structure learning. NewrIPS, 2022.
Πya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. ICLR, 2017.
Πya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR, 2019.
Sindy Löwe, David Madras, Richard Zemel, and Max Welling. Amortized causal discovery: Learning to infer causal graphs from time-series data. Conference on Causal Learning and Reasoning, 2022.
Sara Magliacane, Tom Claassen, and Joris M Mooij. Ancestral causal inference. Advances in Neural Infor- mation Processing Systems, 2016.
Daniel Marbach, Thomas Schaffter, Claudio Mattiussi, and Dario Floreano. Generating realistic in silico gene networks for performance assessment of reverse engineering methods. Journal of computational biology, 16, 2009.
Adam A Margolin, Ilya Nemenman, Katia Basso, Chris Wiggins, Gustavo Stolovitzky, Riccardo Dalla Favera, and Andrea Califano. Aracne: an algorithm for the reconstruction of gene regulatory networks a mammalian cellular context. In BMC bioinformatics, volume 7, pp. 1 1-. Springer, 2006.
Judea Pearl. Causality, Cambridge university press, 2009.
Oona Rainio, Jarmo Teuho, and Riku Klén. Evaluation metries and statistical tests for machine learning. Scéentific Reports, 14, 2024.
Michael J Ravitz, Chen, Mary Lynch, and Emmett Schmidt. c-mye repression of tsc2 contributes control of translation initiation and myc-induced transformation. Cancer research, 67(23), 2007.
16

---

Published in Transactions on Machine Learning Research (10/2025)
Alexander Reisach, Christof Seiler, and Sebastian Weichwald. Beware of the simulated dag! causal discovery benchmarks may be easy to game. NearIPS, 34, 2021.
Joseph M Replogle, Reuben A Saunders, Angela N Pogson, Jeffrey A Hussmann, Alexander Lenail, Alina Guna, Lauren Mascibroda, Eric J Wagner, Karen Adclman, Gila Lithwick-Yanal et al. Mapping information-rich venotype-phenolyy landscapes with genome-scale perturb-seq. Cell, 185(14), 2022.
Karen Sachs, Omar Perez, Dana Pe'er, Douglas A Lauffenburger, and Garry Nolan. Causal protein- signaling networks derived from multiparameter single-cell data. Science, 308(5721), 2005.
Pedro Sanchez, Xiao Liu, Alison Q O'Neil, and Sotirios A Tsaftaris. Diffusion models for causal discovery via topologiaal ordering. ICLR, 2023.
Scott A Sisson, Yanan Fan, and Mark A Beaumont. Overview of abc. Handbook of opprorimate Bayesian computation, 2018.
Claudia Skok Gibbs, Omar Mahmood, Richard Bonneau, and Kyunghyun Cho. Pmf-grn: a variational inference approach single-cell gene regulatory network inference using probabilistic matrix factorization. Genome Biology, 25, 2024.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Peter L Spirtes, Christopher Meek, and Thomas S Richardson. Causal inference in the presence of latent variables and selection bias. UAI, 1995.
Damian Szklarezyk, Rebecca Kirsch, Mikaela Koutrouli, Katerina Nastou, Farrokh Mchryary, Radja Hachilif, Annika Gable, Tao Fang, Nadezhda Doncheva, Sampo Pyysalo, et al. The string database in 2023: protein-protedd association networks and functional enrichment analyses for any sequenced genome of interest. Nucleic acids research, 51 (D1), 2023.
Mathias Uhlén, Linn Fagerberg, Björn M Hallström, Cecilia Lindskog, Per Oksvold, Adil Mardinoglu, Asa Sivertsson, Caroline Kampf, Evelina Sjöstedt, Anna Asplund, et al. Tissuc-based map of the human proteome. Science, 347 (6220),, 2015.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NewrIPS, 2017.
Hongxin Wei, Renchunzi Xie, Hao Cheng, Lei Feng, Bo An, and Yixuan Li. Mitigating neural network overconfidence with logit normalization. ICML, 2022.
Rachel B Wilson, Alexandra. M Koulov, Helia Hatain Tehrani, Jessica S Twumasi-Ankrah,, Yun Jin Chen, Matthew J Borrelli, Cynthia G Sawyez, Siddhant Maini, Trevor G Shepherd, Robert C Cumming, et al. Elongation factor lal regulates metabolic substrate preference in mammalian cells. Journal of Biological Chemistry, 33)(3), 2024.
Menghua WI, Yujia Bao, Regina Barzilay, and Tommi Jaakkola. Sample, estimate, aggregate: recipe for causal discovery foundation models. arXiv preprint arXiv:2402001929 2024.
Xingyu Wu, Bingbing Jiang, Kui Yu, Huanhuan Chen, et al. Accurate markov boundary discovery for causal feature selection. IEEE transactions on cybernetics, 50(12):4983 4996, 2019.
Jianxin Yin, You Zhou, Changzhang Wang, Ping He, Cheng Zheng, and Zhi Geng, Partial orientation and local structural learning of causal networks for prediction. In Causation and prediction challenge, pp. 93- 105105. PMLR, 2008
Paul E Young, Rashmi Issa, L Jeffrey Modeiros, and Sergej Konoplev. Chronic myeloid leukemia, ber-abll-positigg carrying npml mutation first case series from a single institution. Leukemia Research, 111, 2021.
Shimin Hu, Guilin Tang, Beenu Thakral, Naval Daver, Ghayas C
17

---

Published in Transactions on Machine Learning Research (10/2025)
Alessio Zanga, Elif Ozkirimli, and Fabio Stella. A survey causal discovery: theory and practice. Interna- tional Journal of Approsimate Reasoning, 151, 2022.
Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric Xing. Dags with no tears: Continuous optimization for structure learning. NewrIPS, 2018.
Xiang Zhou, Qian Hao, Jun-ming Liao, Peng Liao, and Hua Lu. Ribosomal protein s14 negatively regulates c-myc activity, Journal of Biological Chemistry, 288(30), 2013.
18

---

Published in Transactions on Machine Learning Research (10/2025)
A Proofs and Definitions
Proposition 1. A set of marginal causes of x; is the subset of An(x,).
Proof. Consider the ancestor structure of node Xi, defined the subgraph of g containing the node set x; U An(x)) and all corresponding edges. We denote this subgraph by G. For each node x, E x; Un((x)), we define the parent set restricted to this subgraph as Pa(x, Gi)= = Pa(xx))An(xx) Under the assumption of no latent variables in G, the joint distribution over the ancestor structure factorizes as follows: p(x, U (n())) = Is p(x; Pa(x, Ge)). For any node Iy & x, U An(x;), a do-operation on 26K docs not affect the ancestor structure since the node ia and all edges connected to I are absent from Ge. Also, there are no causal effects on the ancestor set 9 originating from nodes outside this set. Cousequently, for all values c, we have p(x: do(xx = c)) = p(x)). According to Definition 2, it follows that Ia & x; An(x;) is not a marginal cause of x4. This completes the proof by contraposition.
Proposition 2. For any variable subset V C. V s.t. containing x; and all root variables in G, let An(x; V) denote the set of ancestors of x, in the system consisting of V, following the marginalization rule of Bongers et al. (2016). Similarly, let Pa(x;; V) denole the set of parents of x; within V. Then, An(x; An. However, a counterezample exists for parents: Pa(x;; V) + Pa(x))VV
Proof. Note that the marginalization rule of Bongers et al. (2016) on node x operates follows: (1) every direct path x; - X1 - xj is replaced with direct edge x; - Xj, and (2) all remaining edges connected to X are removed. This marginalization rule preserves all causal relationships (i.e., the existence of a directed path between nodes) after the removal of a node. Therefore, by definition, we have An(x;; V) An(x:) n V. In contrast, for direct causes (parents), the relationship may not hold. Consider a causal chain graph with three variables: X1 - X2 + X3. By definition, Pa(xs) = {x2]- However, in the subsampled system with V = x3], xi becomes the direct cause of x3, i.e. Pa(x3; V) = {xi]. On the other hand, Pa(xs) n 0. Thus, Pa(x;; V) # nV.
Proposition (Algorithm complexity). Let n' and m' denote the subsampled variable and observation sizes, and let T denote the ensemble size. For dataset X with variables and observations, the inference complerity of our algorithm is O(n''T(n' + m'))).
Proof. The set of variables can be partitioned into [ inputs. For each input, the complexities of attention layers are O(n"mm') and O(n'm²?,, while the complexity of the feed-forward layers is O(r'm'). Thus processing each variable once results in a complexity of ( (n'mm' + n'm(²)) O(nm'(n' + m²)). Considering an ensemble size of T, the overall computational complexity becomes O(nm'T + mm)). ]
B Experimental Settings
B.1 Model Architecture
For the feature extractor 90, we utilize an axial-Transformer without positional encoding to ensure permuta- tion equivariance (Ho et al., 2019), i.e., if the variables are permuted, the model produces a correspondingly permuted causality prediction. Each Transformer layer comprises two attention layers, one along the variable dimension and one along the observation dimension, followed by a feed-forward layer (Figure 10). Both at- tention and food-forward layers include layer normalization and a skip connection (Ba ct al., 2016). Detailed configuration is provided in Table 3.
As described in Section 4.1, the input to the Transformer is stack [X, Raxmx2 and the output is two feature matrices Fi, F2 E Raxd In basic axial Transformer, the output size for an input of size n X m x 2 is n x m x d, where d is the embedding dimension (Figure 10). We denote this output as H. To obtain the feature matrix of size n x d, we first average the output H of size n x m x along the observation dimension m, resulting in a matrix of size n x d. We then apply two feod-forward layers and concatenate the results produce two feature matrices Fi, F1, E Baxd The total number of trainable parameters of our model is 62k. We will release the code for our model open-source.
19

---

Published in Transactions on Machine Learning Research (10/2025)
Player
Rxd
#XOXX
XXXiXd
Variable-axis
Observation-axis Attention
HXXXXX
Average
Rxd
MLP
F
Input matrix
Embedding
Attention
MLP
Pooling
MLP
nXi
F2
Computational flow of our feature extractor go- We indicate the size of the inputs for each module above the arrows, with only the embedding and pooling modules changing the sizes. For simplicity, normalization and skip connections are excluded from the figure. MLP indicates two-layer perceptron with non-linear activation (Vaswani et al., 2017).
Figure 10:
Table 3: Architecture configuration.
Argument
Value
Number of Transformer layers
10
Embedding dimension
16
Number attention heads
layer hidden dimension
96
B.2 Training/Inference Configuration
Training configuration. We train neural network using the AdamW optimizer (Loshchilov Hutter, 2019), with training configurations detailed in Table 4. The batch size chosen to fully utilize our GPU memory (24GB). Training for maximum of 40,000 steps, with early stopping applied if the validation accuracy does not improve a span of 4,000 steps, chocked every 200 steps. The limit of 40,000 training steps is set empirically, we observed no further performance gains beyond this point on the training dataset specified in Table 6.
We tune the learning rate from the set (6c-4, 8c-4, 1c-3), finding that 8c-4 yields the most stable training performance across all experimental seltings. Notably, the optimal learning rale depends primarily on the model configuration rather than on the dataset type. To further enhance training stability, we apply a cosine learning rate scheduler, which reduces scnsitivity to initial learning rates (Loshchilov & Hutter, 2017). Dropout is not used in experiments, it empirically decreases performance. We hypothesize that the subsampled training and inference scheme, in which predictions are made based on partial information, inherently reduces the need for dropout regularization.
Inference configuration. Figure 9 describes the configurations required for our inference procedure (Al- gorithm 2). Hyperparameters are selected based on the analysis in Section 5.4. Note that we use the same sizes, n' and m', during training. The observation size is set to 200, identical to the variable size, to fit within our GPU memory constraints during training.
Computing environment. We conduct all experiments including training and inference, using a NVIDIA RTX 3090 GPU with 24GB memory. The training time for models Section 5.2 is approximately 9h, while inference takes about few seconds per target (Table 8).
Table 4: Training configuration.
Table 5: Inference configuration.
Argument
Value
Argument
Value
Batch size
32
Subsampled variable size n'
200
Training step
40,000
Subsampled observation size m'
200
Learning rate
Se-4
Ensemble size T
50
Learning rate scheduler
cosine
Weight decay
le-5
20

---

Published in Transactions on Machine Learning Research (10/2025)
B.3 GRN Simulator
We use the SERGIO GRN simulator in our experiments (Dibaeinia & Sinha, 2020). Given a user-defined GRN, the simulator generates a gene expression matrix based on a specified cell type configuration.
Generation mechanism. The simulator samples gene expressions from the steady state of a dynamic system modeled as stochastic differential equations (Dibaeinia Sinha, 2020). Master regulators (i.e. root nodes in the causal graph) operate independently without external regulatory inputs, evolving with fixed production and decay rales. The regulatory influence of each gene is represented by Hill function with pre-determined interaction parameters (Chu et al., 2009). This mechanism captures non-linear relationships and time-lagged effects, providing a realistic model of gene behavior.
The simulator produces datasets that reflect the statistical properties of real-world single-cell experimental data, incorporating several types of measurement errors and technical noisc. The simulator applies these technical noises sequentially to the expression data sampled from the stochastic differential equations:
Technical noise.
1. Dropouts: A high proportion of gene expressions (typically 60-95%) are randomly set to zero, simulating the dropout effect common in single-cell technologics.
2. Outlier genes: Some genes are assigned unusually high expression levels, replicating the presence of outliers.
3. Library size: The total expression level for each cell (known as library size) follows a log-normal distribution, reflecting the variability.
Observational fidelity. To generate the unique molecular identifier (UMI) count expression matrix, a quantification scheme single-cell RNA-sequencing, the simulator applies Poisson random sampling to the expression values after incorporating the technical noises (Chen et al., 2018). That is, the final observation value v is derived cause discovery performance and to determine the performance ceiling of our method, control the fidelity of the Poisson sampling and define three levels:
Poisson(A). To evaluate the impact of this Poisson sampling process targeted
High fidelity: Uses expression À directly as the observation.
Medium fidelity: Uses the mean of 100 samples drawn from Poisson(A).
Low fidelity: Uses a single sample drawn from Poisson(A).
In Figure 4, analyze the performance differences across varying levels of the simulator's observational fidelity. Unless otherwise specified, we use the high-fidelity setting for our analysis.
Intervention. We use the interventional setting identical to Lorch et al. (2022). Specifically, we perform gene knockout by sctting the expression level of a specific gene to zero. We sample 10 intervened observations per gene. For example, given a GRN with 1000 genes, this results in an interventional dataset with 10,000 observations. From these observations and variables, we randomly sample subsets of size n' 200 and m' 200, as provided in Table 5.
B.4 Simulated Dataset
We summarize the configurations of datasets used in our experiments in Table adopted from Lorch et al. (2022). We generate training data using random graphs while testing biological GRNs, E. coli (1,565 genes) and yeast (4,441 genes), obtained from Marbach et al. (2009). Figure 11 shows the degree histograms of these graph structures, highlighting the different patterns between biological graphs and random graphs. We generate interventional data by performing gene knockout on each gene, obtaining 10 observations per intervention. For yeast, we obtain 5 observations per intervention due to its larger gene count. We include 500 obscrvational data points and conduct inference using mixture of obscrvational and interventional data. We preprocess each observation matrix using log, counts-per per-million (CPM) normalization following previous works (Lorch et al., 2022).
21

---

Published in Transactions on Machine Learning Research (10/2025)
Erdös-Rényl Rényi
Scale-Tree
102
10a
10a
102
Count
Connt
Count
Count
10
10
10a
10a
100
10
10
20
30
10
30
Degree (in)
Degree (out)
Degree (in)
Degree (out)
Directional scale-free (out)
Directional scale- free (in)
10a
10a
10a
10a
Count
10a
Count
10a
Count
10a
Count
10a
10a
10a
10
10a
109
109
100
2
50
100
100
200
2
Degree (in)
Degree (out)
Degree
(in)
Degree (out)
Stochastic block model
E. coli
103
Count
10a
Connt
102
Count
10a
Count
10a
10a
10
10a
10a
10
6
8
0
8
5
100
200
300
Degree (in)
Degree (out)
Degree (in)
Degree
(out)
Figure 11: Edge degree histograms of graph structures considered in our experiments. All graphs have an average degree of 2. Blue represents the in-degree histogram, and red represents the out-degree histogram.
For the out-of-distributee (OOD) analysis in Figure 5, we adopt configurations from Lorch et al. (2022)-Table 4. These configurations include varying mechanism function parameters, such as Hill function coefficients and decay rates, which differ from the training settings. We also test models on OOD technical noise types, as described in Table 6. These noise types reflect the statistics of different experimental datasets, which have varying dropout percentages, outlier ratios, and library size distributions (Lorch et al., 2022).
Computational challenge. For the yeast dataset, each input comprises approximately 100 million en- tries tries-2,0000 times larger than the number of pixels in image from ImageNet. (Deng et al., 2009) and 500,000 times larger than the number of input tokens for Vision Transformers (Dosovitskiy, 2021). This substantial size presents significant computational challenges for deep neural networks in memory usage and computation time, especially inducing out-of-memory issue in our computing environment described in Appendix B.2. To overcome this, we propose a subsampled-ensemblee inference strategy that efficiently processes datasets across all considered scales.
B.5 Human Cell Dataset
This scction describes a Perturb-scq dataset used in the human cell experiments (Figure 6). The datasct contains gene expression data from the K562 cell line, which is derived from a patient with chronic myeloge- nous leukemia (Replogle et al., 2022). The dataset. includes both interventional and observational data on
22

---

Published in Transactions on Machine Learning Research (10/2025)
Table 6: Dataset configuration. Note for abbreviations used: ER (Erdós Rényi), SF (Scale-Free), SF- direct (directional Crale-Free), and SBM (Stochastic Block Model) (Drobyshevskiy & Turdakov, 2019). For the training data, we randomly select the graph structure and edge degree independently from the candidate sets. We use slash (/) symbol to separately denote the statistics for E. coli and yeast GRNs. For the exact configuration of technical noise, please refer to Table 4 in Lorch et al. (2022).
Argument
Training set
Test set
Graph structure
{ER, SF, SF-direct, SBM}
E. coli/yeast
Average edge degree
{2,4,6]
2.3/ 2.3/2.1 2.
Dataset size [D] per graph structure
150
10
Variable size (n)
1,000
1. 1,565/4 /4,441
Number observations per intervention
10
10/5
Obscrvation size (interventional)
10,000
15, 15,650 15,650/22,200 /2,
Observation size (observational)
500
500
Number of cell types Technical noise type Technical noise type (OOD)
10
10
10x-chromium
10x-chromium fillumina, drop-seq, smart-seq)
Table 7: Predicted causal gene interactions for target gene MYC by our method. It is worth note that the absence of literature does not imply the invalidation of causality but rather indicates that the causality has not yet been clucidated.
Target gene Causal gene
Supported by
EEFIAI
TCD-DL, STRING,
Literature (Li et al., 2023; Wilson et al., 2024), Correlation Literature (Hong al., 2023)
NPM1
TCD-DL, STRING,
РТMА
TCD-DL, STRING,
Literature (Lin et al., 2015)
RPL26
TCD-DL, STRING, G.
Literature (Gong et al., 2023)
RPLA
TCD-DI. STRING,
Literature (Egoh et al., 2010)
RPS14
TCD-DL, STRING,
Literature (Zhou et al., 2013)
RPSI5A
TCD-DL, STRING,
Literature (Liang et al., 2019)
RPS6
TCD-DL, STRING,
Literature (Ravitz et al., 2007)
RPL23A
TCD-DL, STRING
NCL
TCD-DL, STRING
MYC
HSPA8
TCD-DI., STRING
HISTIH2AE
TCD-DL, STRING
RPL13
TCD-DL, STRING
ACTB
TCD-DL, STRING
TUBB
TCD-DL, STRING
TUBAIB
TCD-DL, STRING
RPL.26
TCD-DL, STRING
RPSI8
TCD-DL, STRING
LST
TCD-DL
VPS37C
TCD-DL
gene expression. The intervention is performed through gene knockouts, identical to our simulation sctting. From the Perturb-seq datasel, we obtain 1,868 genes that have undergone inlervention. We conduct the inference among these genes. For each intervention, we randomly subsample 10 observations. We also sam- ple 500 observational data points, consistent with the number used in our training sctting (Table 6). Using this subsampled dataset, we run our TCD-DL inference algorithm to obtain cause scores for the target gene MYC, applying the same CPM normalization scheme used in simulation data (Appendix B.3).
23

---

Published in Transactions on Machine Learning Research (10/2025)
B.6 Setting for Error Propagation Analysis
This section details the experimental setup for Figure 2 (Section 3). We generate scale-free random graphs with 100 nodes and define causal mechanisms using two-layer perceptron with Tanh activation function (Wu et al., 2024). Root variables are sampled from a uniform distribution. Both approaches estimating the full causal structure and identifying a set of causal variables are trained using the same Transformer architecture. However, their final layers are modified to align with their respective objectives. Both methods maintain the same number of model parameters and are trained on datascts of oqual size. To enable clear comparison of how the false negative rate (FNR) evolves with increasing cause-effect distance, we threshold the cause scores to ensure similar FNRs at cause-effect distance of 1.
C Additional Experimental Results
C.1 Runtime Comparison
Table 8 compares the inference times of methods for targeted cause discovery. Some baseline met (sortnregress and GENIE3) can individually compute a cause score vector for a target, while other baselines (AVICI, DCD-FG, PMF-GRN) roquire the calculation of the entire n X n score matrix to obtain cause score vector for larget. From the table, our method conducts inference within seconds even for the yeast gene regulatory network (GRN) comprising 4,441 genes. In contrast, certain baselines encounter memory issues (AVICI) or require substantial computation time (DCD-FG, PMF-GRN).
Table 8: Runtime measurement. We measure inference time for identifying the causes of a target variable with an NVIDIA RIX 3090 GPU. 00M refers to the GPU out-of-memory error.
Species
Sortnregress
GENIE3
AVICI
DCD-FG PMF-GRN TCD-DL (ours)
E. coli (1,565 genes)
0.9s
1.7s
3.1s
1h 55m
3h 2m
2.5s
yeast (4,441 genes)
2.8s
2.7s
OOM
98h 31m
10h 46m
7.8s
C.2 Scall-Scale Settings
In Table 9, we conduct experiments small-scale settings, comparing our approach to conventional methods such as PC and GIES (Spirtes et al., 2001; Hauser Bühhmann, 2012). We train our model exclusively on synthetic datasets generated with randomly parameterized analytic functions (linear, non-linear multi-layer perceptron (MLP), polynomial, and sigmoid), while maintaining scale-free random network structure with 10 nodes. For testing, we sample datasets with new causal mechanism coefficients and graph structures. Additionally, we evaluate the methods on a real-world dataset called Sachs, which contains 11 variables (Sachs et al., 2005). We obtain interventional data from belearn including six intervention types, each applied to a single node. We sample a total of 100 observations for inference.
Given the DAG predicted by baseline methods, we apply topological sorting obtain binary predictions for ancestor relationships. For our approach, we threshold the cause scores at 0 to obtain binary predictions. measure binary classification accuracy for every pair of nodes to
After inferring the causal structure, determine if the method accurately predicts causal (ancestor) relationships. Table 9 demonstrates that our method generalizes effectively to real-world settings with fewer nodes, highlighting its promising capabilities.
C.3 Analysis on Human Cells
Comparison to correlation ranking. As illustrated in Figure 6, our method identifics novel causal factors for the gene MYC that are not captured by correlation ranking. To quantify the disparity between our model's predictions and correlation ranking, we analyze statisties across 1,868 genes from the Perturb-seq dataset. The average rank correlation between our model's cause scores and correlation- based rankings is
24

---

Published in Transactions on Machine Learning Research (10/2025)
Table 9: Performance on small graphs. We measure binary classification accuracy on a pair of nodes to determine whether the method accurately predicts causality. TCD-DL (direct) refers to a variant of our method that trains model to predict direct causal structures. Each column denotes test datasets with different causal structures and mechanisms. The number of nodes is 10, except for Sachs (Sachs et al., 2005),
which has 11. For PC, use only observational samples.
Method
Linear
MLP
Polynomial
Sigmoid
Sachs
PC
57.2
70.1
69.9
67.6
71.8
GIES
94.3
84.8
70.2
82.3
74.6
TCD-DL (direct)
96.5
86.4
82.3
89.7
75.5
TCD-DL
96.7
88.8
83.0
90.1
76.4
0.068, indicating low ranking similarity. When comparing the top 20 predictions from each approach, on average only 1.1 genes appear in both sets. Notably, an average of 6.4 genes from our top 20 predictions are validated by the STRING database, underscoring that our model identifies novel causal factors not captured by correlation.
Identifying causes of leukemia-related genes. To further validate our predictions from the K562 Perturb-seq dataset, we use the Human Protein Atlas to identify a set of 29 target genes associated with leukemia (Uhlén et al., 2015). We identify the top 10 causal predictions for each of these target genes, including those with support from the STRING database, and visualize these interactions in Figure 12. We present the resulting GRNs for cach leukemia target gene and its predicted causal regulators through network diagrams in Figures 13 and 14. These visualizations highlight the polential regulatory roles of our identified causal genes, providing insights into the predicted interactions driving leukemia. Validation against STRING demonstrates that our approach generates well-supported and highly relevant causal predictions.
Predicted influence of causal genes. We investigate the regulatory influence of each predicted causal gene over lenkemia-related target genes in Figure 15. Notably, nucloophosmin (NPM1) is predicted to regulate 25 of the 29 leukemia target genes, receiving support from the STRING database for 21 of these predictions. NPM1 mutations are prevalent in approximately one-third of adult Acute Myeloid Leukemia (AML.) cases, leading to an abnormal cytoplasmic localization of the NPM1 protein (Falini et al., 2020). Although NPM1 mutations are primarily associated with AML, recent studics have identified them in a small subset of Chronic Myeloid Leukemia (CML.) patients (Young et al., 2021). The prediction that NPM1 regulates substantial number of muukemia-associatte target genes is particularly significant as it provides insights into potential key regulatory mechanisms underlying leukemia pathology. Understanding how NPM1 influences these target genes in CML could reveal critical pathways involved in leukemia progression and help identify novel therapeutic targets.
D Extended Related Works
GRN inference Gene regulatory network (GRN) inference aims to uncover directed causal influences among genes rather than merely identifying co-expression or correlation patterns (Badia-i Mompel et al., 2023). Representative approaches include ARCANE, which employs cermmation-thereete algorithms (Mar- golin et al., 2006), and GENTE3, which uses an ensemble of regression trees to predict each gene's expression based on all other genes (Huynh-Thu al., 2010). While conventional causal discovery methods can be applied to GRN inference, they face significant challenges due to the large network sizes, complex causal mechanisms, and technical noise (Dibaeinia & Sinha, 2020). Recent advancements in sequencing technolo- gics have enabled the collection of interventional data, such as gene knockout experiments, which provide more reliable identification of causal regulators (Replogle et. al., 2022). Additionally, recent efforts have focused on developing simulators that generate single-cell gene expression data by modeling the stochastic nature of transcription based on user-provided causal graph structures (Dibacinia Sinha, 2020). Our data- driven learning approach leverages these advancements in data collection and simulation, demonstrating the potential of data-driven methods for GRN inference.
25

---

Published in Transactions on Machine Learning Research (10/2025)
Causal-Target Gene Interactions
ARCN1
BCLAF1
BCR
FIP1L1
GNPNAT1
HNRNPU
HUWE1 1
МAT2A
MCL 1
MRPL15
MYBL2
Target gene
NCL
NFYB
NPM1
NUMA1
NUP214
NUP35
NUP98
RNMT
RPL6
RPS16
SARNP
SDAD1
SET
SRSF3
SRSF7
STIL
TFRC
XRCC5
ACTB
ARL4D
C9or116
CCAR1
CDK1
CHTF18
DNAJA1
DNAJC9
EEF1A1
HHST112AE
HSPAS
HSPA8
KLC2
LST1
MCM4
MTBP
MYC
NCL
NPM1
PTMA
RAD21
RPL13
RPL14
RPL15
2PL23A
RRL22
RPL27A
RPL299
RPL32
RPL37A
RPL4
RPL7A
RPS12
RP$14
RPS15A
RPS18
RPS2
RPS27A
RPS3
RPS6
RPSA
TAF1C
TIMM23B
TUBA1B
TUBA1C
TUBB
VPS37C
WBP1L
WDR24
HIST2H2A4
Causal gene
TCD-DL TCD-DL, STRING
Figure 12: Predicted causes of leukemia. The matrix illustrates the predieted causality of leukemia-
related larget genes by TCD-DL predictions (blue) and TCD-DL predictions supported by STRING (green).
26

---

Published in Transactions on Machine Learning Research (10/2025)
RPLT3
ARIONT
RGLAPT
DGR
APILI
RPS
RP33
PUPOTT
SPCCR
RPGZ
CMPHAAT
HHRAIPU
HUWET
MAFZA
FTMA
NPRI
HPMIT
HCL
FTMA
HPMM
NCL
KOPAA
RPLID
RPLIS
HOP
PINA
REPTAT
ACTI
AP512
SPER
RPER
TURB
TURO
RPOTO
RPL23
RPSISA
RPL
RPS3
RPSS
JePornal
RPSS
1
MRPL11
NYBL2
RPLES
RPLIZA
LOTT
APoia
IAPSIGA
upssan
NFYB
NPM1
NUMA1
PTMA
KPPH
COKT
CCARE
HTMMI
RPLI3
DHAJAK
RPLI3
NFYS
NPMM
MUMMAT
Target gene
SIAJJOO
RPLI
Causal gene (TCD-DL)
GESTAA
MP812
HOPAN
HPOT2
Causal gene (TCD-DL, STRING)
RPSA
APSTE
MCNA
IPS1R
8PS3
NYC
RADZY
RP22
Figure 13: Predicted causes of a leukemia-related gene. GRNs illustrate the causal genes predicted by TCD-DL (blue) and TCD-DL predictions supported by STRING (green) for each lenkemia-related target gene (pink).
27

---

Published in Transactions on Machine Learning Research (10/2025)
RPRR
APLI3
RPLT3
MUP221
HUPSS
MUrre
RWMT
RPS12
RP512
TUDKIB
RPOTS
RPSZ
RP33
RPL8
RP516
SARNP
SOADT
HOPAY
HEPAS
HAPAS
RPLI
HPM
RPLIT
MPMI
RPLESA
RPLLIA
KLCC
MOCAA
APILEN
MCM
MYC
WBPIL
WOPFL
TURS
RPOS
NYC
RPPLZAA
RPSIIA
SPS14
VPE3TC
RPSS
IePersal
RPOTE
SRSF3
SRBF7
RPS4
RPG2
TFRC
XROCS
HEPAI
IISTIDAA
PTMA
NFMII
LET1
RPLI4
TERC
KROOS
Target gene
NCL
RPLES
Causal gene (TCD-DL)
ACTIB
MPM1
Causal gene (TCD-DL, STRING)
TUBB
PTMA
RPL32
SPLA0
RPQ2TAL
RPER
Figure 14: Predicted causes of a leukemia-related gene. GRNs illustrate the causal genes predicted by TCD-DL (blue) and TCD-DL predictions supported by STRING (green) for each lenkemia-related target gene (pink).
28

---

Published in Transactions on Machine Learning Research (10/2025)
Number of Target Genes per Causal
TCD-DL, STRING TCD-DL
24
Number of target genes
10
NPM1
RPL 26
EEF1A1
RPS18
RPL133
RPS12
RPS15A
HSPA8
TUBB
NCL
RPL23A
RP52
RPS14
VPS37C
LST1
RPL15
TUBA1B
DNAJA1
KLC2
WBP1L
HSPAS
MCM4
RPL4
RPSA
RPL27A
RPL7A
DNAJC9
MTBP
CHTF118
WDR24
C9or116
CC4R1
CDK1
ARL4D
HIST2H2AC
TIMM238
TAFIC
RAD21
TUBA1C
HIST1H2AE
RPL14
RPLS2
RPS27A
Causal gene
Figure 15: Predicted influence of causal genes. The histogram shows the number of leukemia-related
target genes predicted to be regulated by each causal gene, with predictions made by TCD-DL (blue) and
TCD-DI. supported by STRING (green).
Table 10: Benchmarking results. Targeted cause discovery performance on E. coli GRN with 1565 genes varying levels of simulator's observational fidelity. All measurements are expressed as percentages (%).
Fidelity high
Fidelity medium
Fidelity low AP
Method
AUROC
AP
F1
AUROC
AP
AUROC
F1
Random
50.0
0.5 0.0
0.5
50.0
0.0
0.5 0.0
0.5 + 0.0
50.0 0.0
0.5 0.0
0.5 + 0.0
Correlation
51.1
0.7
0.8
51.0
0.7
0.8 0.1
49.8
0.7
0.8
Sortnregress
1.4
49.9 0.8
1.4
1.2 1.4004
50.1
0.7
0.7
PMF-GRN
52.6
0.9
52.9 3.2
1.0
51.5
1.
1.1
GENTE3
56.0
2.9
2.1 0.7
54.3 = 39
2.2 1.1
0.9
37
11
+ 02
DCD-FG
50.0
0.6 0.0
50.0 = 0.0
0.6 0.0
0.0
0.0
0.5
1.0 0:00
AVICI
56.5
1.0 + 0.2
0.4 0.3
62.7 + 5.5
3.1 + 1.7
3.7 4 2.6
56.5 8.4
1.2 = 0.9
1.3 + 1.1
TCD-DL, (ours)
94.6 + 1.8
98.6 + 6.3
36.3 6.1
81.7 # 11.3
14.0 11.7
13.4 + 12.0
71.5 + 3.7
2.4
2.0 + 1.8
29


Reviews:
[{"review_id": "55ILxKfKTb", "reviewer": "Reviewer_nsJv", "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 recasts targeted causal discovery as binary classification of ancestors (order-free) and combines it with local subsampling + ensembling, yielding practical scalability on large GRNs.\n- The paper includes comprehensive experiments with real-world datasets and demonstrates significant practical impact.\n\n### Weakness\n\n1. **Heavy reliance on high-fidelity simluators**\n    \n    The primary limitation of this approach is its fundamental reliance on the quality and realism of the simulator used for training. The model's outstanding performance is contingent on the SERGIO simulator's ability to generate synthetic data that accurately captures the complex dynamics of real-world gene regulatory networks.\n    \n    This raises a critical question about the method's broader applicability: its success in new domains (e.g., economics, climate science) would depend entirely on the existence of an equally high-fidelity simulator for that specific domain. Therefore, while the method itself is powerful, its generalizability is ultimately bounded by our ability to accurately simulate the target system, a significant challenge in many scientific fields.\n    \n2. **Strong assumption in Proposition 2** \n    \n    The assumptions in Proposition 2 appear quite strong. To apply this result to the local subsampling, it seems that every local subsample would need to include the root set. Is this actually the case? If so, how does the method guarantee this condition?", "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": "### Summary\n\nThis paper provides a method for predicting an ancestral set of a targeted node from a mixture of observational and interventinoal datasets. The proposed framework (TCD-DL) takes a  simulated training dataset, which is a triplet of an observational sample, an interventional sample and a graph. In the training phase, the model for computing the score vector, which indicates which nodes are ancestors of the target node, is trained from the random subsets of the training dataset with random subset of nodes with marginalized causal graphs.  In the test phase, the model can predict the ancestors of a given target node. \n\n### Contribution\n\nThis paper recasts the causal discovery task into the classification task using so-called “Sim2Real” or “synthetic-data pretraining method”. This approach is interesteing, practically strong and scalable when the data generating process for training and test data are believed to be shared. Under this setting, the superior performance is shown from many experiments in the paper.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Issues: Overclaiming of Title and Contribution\n\nI think the title and the contribution is overly claimed, because the method is confined to those which has the high-fidelity simulators. This is a very special type of causal discovery setting, which should be highlightened in the introduction and the title. This is not the same as a standard causal discovery problem. This kind of confinement should be reflected in Introduction and the title. \n\n### Questions: Performance Trade-off of Local Subsampling\n\nWhat is the performance loss when comparing the use of local subsampling versus not using it? I understand that local subsampling has been implemented to achieve scalability, but to what extent does it trade off with performance?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "6CyEELnFrk", "reviewer": "Reviewer_JL6c", "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 strengths of the paper are:\n\n1) a reasonably successful method for causal identification is presented and evaluated.\n2) the writing is reasonably clear\n3) there are quite a few useful and interesting comparative experimental results presented\n4) due attention is paid to computation issues\n\nThe main weakness is that the rationale for the 'learning based approach' should be fleshed out a little more. Here are some specific points:\n\n(1) Fig 2 compares results of a particular 'Structure' method versus a\nparticular 'Cause' (ie the authors') method. They conclude from this\nthat: \"The results show that our method maintains a stable error rate\nregardless of distance, whereas error rates in causal structure\ndiscovery methods increase as the cause-effect distance grows.\" This\nis an invalid inference from results using a particular causal\nstructure discovery method to \"causal structure discovery methods\" in\ngeneral. This is wrong, of course.\n\n(2) Why is the (reasonable) subsampling approach mentioned in S4.2 called\n'local inference'?\n\n(3) So we take top 20 genes found by the method. Unsurprisingly, most of\nthem have been found to be associated to the target. It is claimed\nthat 30% are validated by the literature (the 7 green genes in Fig 6)\nbut in only 5 cases is the causality the correct way round.\n\nTYPOS:\n\nA set of marginal causes of xi is the subset of An(xi ).\n->\nThe set of marginal causes of xi is a subset of An(xi ).\n\nExtened\n->\nExtended", "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 continues a line of work using what the authors call\n'learning-based approaches' to causal learning. The distinguising\ncharacteristic of this paper is that the goal is to identify causes of\na particular target variable and not to infer an entire causal\nstructure, or even to distinguish between direct and indirect causes.\nGiven this restricted goal the learning method can be tailored to it,\nand so it is no surprise that the authors' method does better at\nidentifying causes than other methods aiming to infer causal structure\nmore generally.\n\nThe rationale for 'learning-based approaches' is better given by Lorch\net al than in the current paper, so I would encourage the authors to\nspend a little time improving this aspect of the paper. Leaving\ninterventions aside, the method of generating simulated data (See B.4)\ndefines a joint distribution: P(G)P(Theta|G)P(x|G,Theta) where G is\nthe graph, Theta the \"mechanism parameters (mech.), and noise\nconfigurations\" and x is an observation sampled from the distribution\ndefined by (G,Theta). The terms in the joint distribution\nfactorisation are structure prior, parameter prior and likelihood\nrespectively.\n\nThis way of describing the problem is essentially given by Lorch et al\nand would encourage the authors to also include such a description\nsince it it's very helpful in allowing a reader to grasp the\nessentials of the presented method.\n\nIf G(l_i[j]) indicates that xj is an ancestor of xi in G, then the\nauthors' goal is to estimate s_i[j](X) = Sum_{G:\nG(l_i[j]=1)}P(G|X). Rather than attempt to do this by, say, sampling\ngraphs from the posterior P(G|X) using MCMC, the authors take a more\ndirect approach: use a transformer to learn (all the) s_i[j](X) from\nsimulated data. So I would argue that what we have here is a\nparticular way of doing Bayesian statistics, so it's success depends\non the choice of prior (how the random graphs are generated, how the\n'mechanisms' are chosen). The method is somewhat similar to\n\"Approximate Bayesian computation\".\n\nThe goal of identifying causes of a given target is sufficiently\nimportant that I think a method for doing just this is justified.\nAlthough the description of the method could be improved (see above) I\ncould work out what they were doing (although much was contained in\nthe appendices). The authors' method has high computational\ncomplexity, but this issue has been addressed properly and a\nreasonable subsampling method used.\n\nThe experimental results support the usefulness of the method.\nComparative results show the authors' method in a good light, although\none might view the comparison as a little unfair, given that the other\nmethods were not specifically designed for the task at hand.\n\nFig 8d shows that we really need interventions to do a good job of\ngetting the causal direction correct, but 68% on purely observational\ndata seems reasonable.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Address the weakness mentioned above by:\n\n1) improving the motivation/rationale for the 'learning-based approach'\n2) Tone down the claims mentioned in (1) and (3) above\n3) Fix small things like typos.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "GlM3uK2qUF", "reviewer": "Reviewer_gFu9", "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": "### Strengths\nI find the paper clearly written and easy to follow. The Method section describes the proposed algorithm in detail, with all design choices properly explained and motivated. The Experiment section is extensive, including relevant ablations and analysis. Overall, together with the hyperparameter details in the appendix, I feel that the work is complete and reproducible.\n\nThe proposed local inference strategy is a strong practical contribution of the paper, as it is simple and theoretically motivated by the Local Inference Guarantee in Proposition 2, while providing scaling benefits in the number of variables in practice as showcased by Section 5.2, which is a long-standing issue of causal structure learning.\n\n### Weaknesses\n\nIt is not immediately clear how useful the task of cause discovery is compared to causal structure discovery, as the example given in the introduction mainly focuses on local causal discovery. While I agree that causal structure discovery is simply infeasible for systems with too many variables, I wish the authors provided more context and examples of use-cases that cause discovery, i.e. recovering only the set of ancestors, can solve. Are there other works in the literature that consider cause discovery, and if so, what is their motivation besides the task being more feasible?", "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": "The paper considers the problem of cause discovery: finding all ancestors/causes of a variable in the causal graph. The authors detail arguments against using standard causal structure discovery to solve this problem, namely that causal structure discovery does not scale well with the number of variables, and that small errors in the predicted causal graph propagate, increasing the overall error rate for predicting ancestry. Instead, the authors propose Targeted Cause Discovery with Data-driven Learning (TCD-DL), which inputs a target variable as well as observational and interventional data with known intervention targets, and directly predicts for each other variable whether it is its ancestor or not, without reconstructing the causal graph.\n\nTCD-DL is implemented by a Transformer architecture, and trained by stochastic gradient descent with a supervised learning objective using simulated data from synthetic graphs and known ancestry information as labels. The goal is to generalise the model to unseen causal structures and causal mechanisms. The main technical novelty of the TCD-DL architecture is its ability to scale linearly with the number of variables by utilising the *local inference strategy* proposed by the authors, which subsamples the input data both on the variable and the sample level.\n\nThe authors evaluate TCD-DL using causal structure discovery baselines. They show that it generalises better to unseen, realistic data after being trained on fully synthetic data even across different noise levels. The authors also validate their method on a real world human cell dataset. Finally, the authors provide extensive analysis and ablation studies for the impact of graphs and causal mechanisms seen during training, and other training and algorithm parameters.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Adjustments critical to securing recommendation for acceptance\n\nIn the example of the introduction, the authors claim that their approach “prioritizes transcription factors that are more accessible and cost-effective to manipulate”. However, as I understand, the proposed algorithm returns a binary label indicating whether a variable is an ancestor of some target variable or not, without outputting any further prioritization for cost-effectiveness. I suggest the authors describe how the algorithm provides such prioritization, or correct this claim otherwise.\n\n### Adjustments that would strengthen the work\n\nAs described in the weaknesses, I believe the paper would benefit from providing more motivation for why “cause discovery” is an interesting and useful task.\n\nThe term “ensemble” and its parameter $T$ should be introduced in the following sentence: “Repeating this process across multiple subsamplings of variables and observations …”. Otherwise it is not introduced anywhere else before it is already used in the name of “ensembled local-inference”, as well as Proposition 3 and Algorithm 2.\n\nIn Section 5, the paper sometimes uses the phrase “targeted causal discovery” instead of “targeted cause discovery”. Since the former phrase was not introduced earlier, I assume this is a typo that should be corrected. Otherwise, please introduce targeted “causal” discovery and describe the difference compared to targeted “cause” discovery.\n\nI find the usage of the term “local” in the paper confusing or overloaded. In Section 2.2, local approaches involve methods that discover the local neighborhood around a node, such as its Markov blanket or parent set. On the other hand, in Section 3 and later, local inference refers to the fact that subsampling the variables preserves ancestral relationships. I suggest the authors clarify how local inference in Section 3 and later relates to local methods in Section 2.2, or how it is “local” in general.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)