GEASS: Neural causal feature selection for high-dimensional biological data
Authors: Mingze Dong
Yuval Kluger
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Granger causality
feature selection
neural networks
single-cell genomics
spatial transcriptomics

Abstract: Identifying nonlinear causal relationships in high-dimensional biological data is an important task. However, current neural network based causality detection approaches for such data suffer from poor interpretability and cannot scale well to the high dimensional regime. Here we present GEASS (Granger fEAture Selection of Spatiotemporal data), which identifies sparse Granger causality mechanisms of high dimensional spatiotemporal data by a single neural network. GEASS maximizes sparsity-regularized modified transfer entropy with a theoretical guarantee of recovering features with spatial/temporal Granger causal relationships. The sparsity regularization is achieved by a novel combinatorial stochastic gate layer to select sparse non-overlapping feature subsets. We demonstrate the efficacy of GEASS in several synthetic datasets and real biological data from single-cell RNA sequencing and spatial transcriptomics.

GEASS: NEURAL CAUSAL FEATURE SELECTION FOR HIGH- DIMMENIONNL BIOLOGICAL DATA
Mingze Dong
Yuval Kluger
Yale University mi ngze. dong@yalee.du
Yale University
yuval. k wll.geerrraleeeeu
ABSTRACT
Identifying nonlinear causal relationships in high-dimensione biological data is important task. However, current neural network based causality detection approaches for such data suffer from poor interpretability and cannot scale well to the high dimensional regime. Here we present GEASS (Granger fEAture Selection of Spatiotemporal data). which identifies sparse Granger causal interacting features of high dimensional spatiotemporal data by single neural network. GEASS maxi- mizes paarsit--regllarinn modified transfer entropy with a theoretical guarantee of recovering features with spatial/temporal Granger causal relationships. The sparsity regularization is achieved by a novel combinatorial stochastic gate layer to select sparse non-overlapping feature subsets. We demonstrate the efficacy GEASS in several synthetic datasets and real biological data from single-cell RNA sequencing and spatial transcriptomics.
INTRODUCTION
Advances in single-cell omics research enable full characterizations of high-dimensional gene dy- namics in biological systems on a either temporal or spatial scale. An example for the temporal case single-cell RNA sequencing (scRNA-seq) trajectories, where cells are sampled from dynamical biological process, sequenced, and ordered based on either real sampled time or inferred pseudo- time (Cannoodt et al.) 2016 Saelens et al., 2019). Gene dynamics along the specified cell order encodes information of causal regulation for the underlying biological process. An example for the spatial case is single-cell level spatial transcriptomics (e.g. SeqFISH+ (Eng et al.) 2019). Merfish (Fang) al.] (2022) ), in which cells from tissue slice are sequenced with their spatial coordinates preserved (Moses and Pachter 2022] Rao et al. [202]] Palla et al., (022). Spatial profiling allows investigations of the cellular interplay, corresponding to conditional gene expression change caused by neighbor- hood phenotypic states. However, despite the potential significance, data-driven causal discovery for such data remains largely unexplored, especially for the spatial omics data.
Identifications of causal regulatory patterns in such data can be reformulated into the general task of causal feature selection in observational data with intrinsic structures, e.g. spatial data temporal data. Identifications of causal interactions in time-series has lead to valuable findings in multiple disciplines, including but not limited to, economy, climate science, and biology (Hoover 2006 Kamiñski et al., 2001 Runge et al. 2019a).
Learning directed causal relationships in temporal/spatial data is feasible as time and space both induce asymmetric dependencies. In the case of time-series data, feature in the future cannot have effect on past values of other features. For spatial data, similar definition of causal dependency can be established (Herrera Gómez et al. 2014).
The concept of Granger causality is proposed in order to uncover the assymetric causal dependency (Granger] 1969; Shojaie and Fox (2022) In time-series data, this would translate to identifying one variable's s causal relationship with other variables based on how well the historical observations of other variables can predict the variable's present value. The application of Granger causality in spatial context corresponds to predicting significant relationships between neighboring observations of other variables and the specified variable (Mielke et al., 2020),, which is key insight used in recent works aimed to discover cellular interaction patterns in spatial omics data (Fischer et al. 202T Valdés-Sosa et al. 18)

---

In the nonlinear regime, information-theoree measures such directed information, transfer entropy (Schreiber (2000), and partial transfer entropy (Staniek and Lehnertz (008). are used as a counterpart
of linear Granger causality. Moreover, some works consider modeling conditional independence (CI)
in time-series data to identify the underlying causal graph (Entner and Hoyer 2010; Malinsky and
Spirtes 2018 Moneta et al.] R et al.] 2019a Phister al. [019 Mastakouri et al. 2021).
Two examples are VarLINGAM (Hyvärinen et al. 2010) and PCMCI (Runge et al. 2019b)), which
are generalizations of LINGAM (Shimizu et al. 2006) and PC (Spirtes et al. 2000) respectively.
Finally, multiple recent works have proposed to use neural network approaches to model the nonlinear
Granger causality, including MLP, LSTM, and neural-OD ODr based approaches, resulting in improved
prediction power for nonlinear time-series dynamics (Li al. 2017 Tank et al.) 2021 Nauta et al. 2019 Yin and Barucca
2022 Bellot et al.) (02T).
Despite the success of these methods in various systems of interest, multiple challenges limit their use in nigh-dimensional biological datasets.
Although linear methods (LINGAM, linear Granger causality) have succeeded in various settings and can potentially scale to high feature numbers, these methods may completely fail when the feature dependency in data is highly complex and nonlinear.
As the number of conditional independencies generally scales exponentially at least polynomially with the feature size, applying causal discovery methods which are based on CI tests high-dimensional data is not realistic. Distinctively, Granger-causality based methods are built with a prediction model for each feature in the data. The time complexity of solving the stacked prediction model for all features is of polynomial level with respect to the feature size.
In previous methods, the number of causal edges between features is assumed to be sparse (edge sparsity) to maximize interpretability of the identified causal graph. However, in biological data, there exists a large proportion of nuisance features. Also, one functional gene may activate a large number of downstream genes in neighboring cells. Sparsifying the number of interacting features (feature sparsity) has the potential to improve causal discovery biological systems, which remains to be explored.
While large number of methods are designed for causal discovery in time-series data, only a limited number of present works aim for causal discovery in general graph-structured data. Time-series based methods cannot be directly adopted on data with multi-branch trajectory dynamics or spatial structures.
Our contribution. In this work, we present GEASS (Granger fEAture Selection of Spatiotemporal data), which identifies causally interacting features of high dimensional temporal / spatial data by single neural network. GEASS considers the aforementioned feature sparsity instead of edge sparsity, thus selects most significant interacting features for downstream causal discovery. Our contributions are three-folds.
1. Instead of direct causal discovery in data, we formulate the task as two steps of causal feature selection and causal graph identification. We provide a novel solution of causal feature selection problem in general graph-structured data by the use of modified transfer entropy maximization with theoretical guarantees.
2. In order to solve our proposed optimization problem, we design a novel combinatorial stochastic gate layer to select non-overlapping sparse feature sets with a newly designed initialization procedure.
3. We demonstrate the power of our method by benchmarking it on both temporal data and spatial data of multiple settings. Our method gives accurate and robust causal feature identification and reveals novel biology in real datasets.
1.1 RELATED WORKS
Neural Granger causality. Despite the large body of work based on linear Granger causal discovery, neural Granger causality still remains active area of research. Various neural network architectures. such as MLP, sequential model, and attention-
architecture (Tank et al.) 202 Nauta et al. Khanna and Tan 2019: Sun et al. 2021), have been proposed for nonlinear Granger causality
2019;
2

---

discovery. A recent work uses the information of proxy variable to learn latent confounder for Granger causality by a dual-decoder neural network (Yin and Barucca 2022). One recent biology- oriented work extends the definition of Granger causality to DAGs, where the use of a linear graph neural network is proposed to model underlying Granger causality (Wu et al. (021). Meanwhile, neural-ODS based approach has been proposed to reformulate the Granger causality problem in terms of local dependence graph identification (Bellot et al. 2021).
Causal feature selection. The task of causal feature selection has been considered by multiple groups. Most works in this category uses constraint-baser methods to identify each feature's ' causal relation with all other features, equivalent of identifying the whole causal graph structure, including VARLINGAM, tsFCI. SVAR-FCI, and PCMCI (Hyvärinen et al.) (2010] Entner and Hoyer 2010)
Malinsky and Spirtes, 2018: Moneta al. 2011 1 Runge et al. 2019a). Meanwhile, SeqICP focus on identifying the direct or indirect cause for each feature assuming sufficient interventions in the dataset (Phster et al., 2019). SyPI tackles the causal feature selection problem without the assumption of causal sufficiency and avoids issues in multi-hypoothssis testing by construction of the correct conditional set (Mastakouri et al. 202T). Finally, Guo et al. (2022) considers dual correction causal feature selection to control both false positive rates and false negative rates.
2 MODIFIED TRANSFER ENTROPY (MTE)
In order to tackle the issue that a neural network may overfit each model therefore overestimates the number of causal interactions, need prediction-free loss function that directly indicates causal signficance. In this work, we propose a novel function, modified transfer entropy (mTE), based on transfer entropy (Schreiber 2000) a metric of causal interaction significance.
Transfer entropy is a information-theoretic measure of cross dependence (Schreiber) 2000 Consider lwo vectorized time series x' and y' for t E I,...,'' In a Markovian model, the transfer entropy from æ to at time t is defined as the mutual information between the present value x and the future value y -+1 conditioning on y' eliminate possible autocorrelation: [E, y) = [xx"; y lgy).
By the use of mutual information, transfer entropy is able to model general nonlinear dependencies beyond linear Granger causality. In this work, we further consider the generalization of transfer entropy on graph structured x' and y', where denotes a vertex on the data graph G = (V,
(ar, y) = T(x"; y M(0 (y'), where N()) = (i().)) E E).
Note here the graph can be either directed (the time-series case) or undirected (the spatial case). In this study, we introduce a novel function, modified transfer entropy, that enables the application of bivariate transfer entropy for causal discovery in high-dimensional data. Our key insight is to consider two feature subsets in the dataset that maximizes the mutual information difference:
Definition 2.1. Let X 1"='x"...0"
E p be matrix containing graph-structured vector series D, with i as vertices of the data graph G = E). Suppose S1 and S2 be two subsets of {1, 2. p). The modified transfer entropy mTTE (1,S5) S2) and its maximun mTE; are defined by
N(0) - (S1, S2) = 1(x's sz
1 [ ss)) mTE, = max ('
(2)
Sq.²
Note the mTE function requires strictly stronger dependence than the analogically defined transfer entropy (S 52), as shown by the proposition below (The proof can be seen at Appendix A. Proposition 2.2. SS C [1,....p) mTE S2) > 0 + TE, S2) > 0.
= Let (ST, S2) be one of the maximizers with the smallest size of [S, U S2l, and denote S= Si U
(note St) may not be unique). Under some mild assumptions listed below, we are able to provide the theoretical justification for mTE maximization in the time-series setting (Theorem 2.4 A proof can be seen in Appendix A.3
Assumptions:
A1-A3 Causal Markov assumption, faithfulness, and causal sufficiency for the causal graph.
A4 Ergodicity and Stationarity of the stochastic process defined by the causal graph, meaning the ensemble average equals time average, and the functional relationships encoded by the causal graph do not change by time (or location). This also leads to m (S,, S2) is constant across i.
3

---

A5 DAG causal graph: We assume x" [ti, fm Mm++1 ..., Upl up permutation, where t; are causally interacting features forming a directed acyclic graph (DAG), and ug are nuisance features that may correlate with &.. An illustration based on the time series setting can be seen in Figure I
correlates
4
t
causes
u1
I+1
+++1
t/2
+++
3
l/
causes
L3
Figure 1: Dependence graph for single forward step (left) and the underlying causal graph (right).
A6 Interaction regularity: Given two disjoint feature sets A, B, such that A is a subset of the parent features of B or B is subset of child features of A. Then conditioning on any other feature set C such that I (A BD( I"), BN BM (C)())) > O, we have:
min,in((AA", IN") I(A', B BNOICNN
I(A',
Remark 2.3. Here our only additional assumption from prevalent literatures (Pearl 2009 Spirtes et al. 2000) A6, which aims to filter out features with spurious causations and regularize the algo- rithmic complexity of causal interactions, thus enabling informatiin--heorrtiee analysis. A6 has direct connections with concept of conditional transfer entropy (Faes et al. 2016)) Shahsavari Baboukani et al. 2020); further discussions can be at Appendix
Theorem 2.4. Given AI-A6, S == (Si (5)) S {1,..., m) (the index set of true interacting features described in A5). Moreover, each feature in S* is connected to other features in the set S*.
3 NEURAL OPTIMIZATION OF MODIFIED TRANSFER ENTROPY
With Theorem 3.1 stated below, we are able to give theoretical guarantee of the lo-penalized optimization of mTE. A proof can be seen Appendix A.4 Here @ stands for the Hardmard product. Theorem 3.1. Assume A1-A6 holds and f. 9. define one-to-one mappings on X X@ls, s, (for f) X@
1s, (for h). Then 3Y > O. such that for Q. any solution (S) U U satisfies S == (S) U St) S
{I, m. Moreover, each feature in S* is connected to other features the set.
min S.g.
-(I[((x' 3 Iss)); la(x
(1) @ I(((x
Is))i O 1s,))) + X|Sk U U (4)
Remark 3.2. The estimation of mutual information by various approaches is active field itself BBlggazi et al. 2018; Hjelm et al. 2018: McAllester and Stratos, 2020: [Zhang et al.) 2019). In contrast, by this theorem, we show that accurate estimation of the transfer entropy (such in (Zhang et al. 2019)) may not be needed as optimizing the upper bound of the modified transfer entropy automatically gives the best feature subset selection.
Remark 3.3. Our theoretical guarantee is derived based on one-to-one embeddings f.g. f h. In neural network, the injectivity may be enforced with various architecture designs yet may perfectly hold. Empirically, we have found that the optimization of mTE is robust to the embedding injectivity, compared with the original transfer entropy. This due to our stricter design of the mTE function (Proposition
2.2 and is further illustrated by our experiments in the next section.
able to construct neural network for optimizing the proposed loss function. However, the estimation of mutual information is not directly tractable. In this case, because mutual information is invariant by one-to-one transforms, we can restrict the function class of f.g. h in the optimization problem (4) as flows transforming the original feature distributions into Gaussian distributions with fixed dimensionality. We are able formulate the target for neural network optimization by the explicit formula for mutual information between Gaussians: I(X, Y) det det det det det E|x,y] The Gaussian regularization can be applied either by regularizing over the discrepancy between embedding distributions [f. g. h] and Gaussian distributions or by applying a adversarial training procedure. In this work, we have implemented the former approach, constructing means and covariance matrices for the concatenated embedding as learnable parameters and minimize the cross entropy between target distributions and the parametrized Gaussian distributions.
Given Theorem 3.1
log
4

---

3.1 COMBINATORIAL STOCHASTIC GATES
In order to solve the optimization problem, we need to learn two sparse sets S,, S2, which involves combinatorial optimization, making the task impractical for high-dimensional data. To overcome this issue, we use a stochastic gate based approach (Yamada et al. 2020; Lindenbaum et al., 202T). which performs probabillistie relaxation of deterministic norms. order to explicitly construct Sj and S2 by stochastic gates, we define two random vectors T and 72 ranging in [0, 1] with lengths equal to the feature number, with each element independently sampled from STG distribution defined as: Ti = max(0, min(1, 1, !' (2)), where ca ~ (0, of) is i.i.d. sampled with fixed variance and μu is a parameter trainable by reparametrization (Miller et al., 2017 Figurnov et al. 2018
The new loss function applying stochastic gates can be formulated as:
By [F([(X ): s,,)
s, +
(1P((T) > 0) + A21((T3 (0, (0,1))),
s.1. Xs. = = O T2,
Xs. = x の (1- TT²)
(5)
Here is defined as the empirical Gaussian mutual information: f(x, Y) =
2 log det Ey det Ev and det Spx.r
W is defined as the graph diffusion operator: Wx = N(), In our construction, 7 controls the sparsity of feature selection, while T2 controls the expectation of overlap between Xs, and Xs, Denoting the Gaussian error function erf(), the regularization term for the first layer is of form:
Σ P(T > 0) Σ
2-
2
1
erf Md
)).
(6)
2
2201
d=1
i=1
The regularization term for the second layer can be expressed as:
Σ IPCT3 (0, - Σ -) P(T2 2 1)
1
Σ
(erf( pg erf( 1
V202
/202
d=1
We are able to show strong consistency for our stochastic-gate based feature selection scheme by the theorem below (A proof can be seen at Appendix (A.3):
Theorem 3.4. Assume AI-A6 and f,g, are one-to-one Gaussian embeddings as described above. For the optimal solution of (). denote a sample of stochastic gate as TT, T and denote the ground
truth interacting feature set as S, then there exists A1; X2 > ffr (5) such that as -> D.,
Vi {0, 13. P(B, S S) 0.3h where B {d[T] > 0, 0. =
(8)
In practice, we also have observed the method's solution highly depends on the stochastic gate initialization. Here we provide a heuristic initialization scheme that shows superior empirical performance. Details of the initialization scheme can be seen in Appendix B
3.2 PROPOSED NETWORK ARCHITECTURE
Our proposed network architecture is summarized in Figure 2
For an input dataset X E Rex" and its corresponding graph adjacency matrix A E R we first pass each feature through two sequential stochastic gate layers T T The lo penalty is conducted on the first STG layer, while the second STG layer is regularized with the 0-1 penalty, consistent with the descriptions in the previous section.
After passing each feature, denote 1 = 1- T we have two intermediate embeddings defined by Xs, = X ④ 7 ④ 2 and Xs, = X o T の f2 respectively. Then these two embeddings are passed through MLP1 (f) and MLP2 (g) to generate Gaussian embeddings f fXxs,) ), g(Xss) corresponding to 5 For the design of function h, we consider two crucial elements: an additional layer to aggregate the information from different nodes N()) 2. the injectivity of mappings f, g, h. Note f. h in (5) are automatically enforced to be injective on interacting features to maximize the first term of mTE, but g is not. Therefore, our final design of h is the composition of first applying (enforcing the injectivity of g). mean aggregation layer without self-loop consistent with the GCN design (Kipf and Welling 2016) by multiplying the adjacency matrix A, and another MLP layer (MLP3). Finally, we compute the minus empirical Gaussian mTE ī f) - ī(f, h) and add the cross- coos-entroppy penalty between the concatenated embedding distribution and learnable Gaussian distribution.
5

---

X
X 07
fixe e
Color bar
High values
MLP 1 (1)
Low values
layer
layer
i F( penalties
X ④ (1-2))
x 6x 27))
h(X & & (1 (1-1")
MLP 2 (g)
Aggregation
373
Figure Illustration of the proposed GEASS network architecture.
3.3 OUTPUT INTERPRETATION
Upon the algorithm convergence, GEASS provides both outputs of active features (Bo U Bi) and embeddings (ʃ. g, I)) produced by causally interacting features. In this paper, we emphasize the use of the identified interacting features Bg U B1. The output of embeddings (f,, 9. h) may be complex and nonlinear, potentially requiring additional architectures to maximize its interpretability.
By the construction of GEASS, features B1 and sink features By. These features may be used as inputs to further proper causal analysis, such LPCMCI (Gerhardus and Runge 2020) for time-series data, which despite its statistical power in depicting possible lags, identifying latent confounders, and allowing nonlinear tests, can only work data with moderate feature sizes. Also, these features may be used in other machine learning models for improved model interpretability.
able to get two separate sparse feature subsets source
4 EXPERIMENTS
4.1 GAUSSIAN TIME-SERIES WITH POSSIBLE NONLINEARITY
In order to benchmark the method in time-series data, we consider two settings: Minor effect of latent processes, with autocorrelation present; 2. Significant effect of latent processes, with autocorre- lation present. Both settings are modeled by Gaussian structural processes with an underlying causal graph. Further details can be seen in Appendix C.I
We test the false discovery rate (FDR) and Fl score between ground truth interacting features and recovered features with two categories of methods, namely conditional independence based (CI-based) methods and Granger causality based (GC-based) methods respectively. The first method category includes VAR-LINGAM (Hyvärinen et al. 2010)), PCMCI (Runge et al., 2019b), and LPCMCI (Gerhardus) and Runge, 2020). Among them, despite the statistical power, LPCMCI is not included in our experiment as it fails to converge in given time in our preliminary experiments. The second method category includes a neural- network based generalized vector autoregression model GVAR Granger (Marcinkevičs and Vogt 2021), and Grid-net which generalizes the definition of Granger causality to Directed Acyelic Graph (DAG) (Wu et al. (021); moreover we include two sfate-of-the-art approaches, DCM and NGM implemented in Beelot et al.) 2021 that use neural ODE to model nonlinear dependence graph.
metrics for digheiimensiona data causal discovery. We compare GEASS
Table I shows our benchmarking results. Among the alternative methods, GVAR and GrID-net fail in all settings as they are not designed for causal feature selection. VAR-LINGAM achieves high accuracy in linear settings while fails in nonlinear settings. In contrast, PCMCI fails when latent processes contribute to both true causally interacting features and nuisance features, creating spurious correlations. Empirically we also observe that DCM and NGM achieves comparable performance

---

when the dynamics are linear but performs worse in the nonlinear setting, where the dynamics more irregular. Finally, GEASS consistently gives accurate causal feature identifications (high F1) and low false discovery rate (low FDR) in all settings considered.
Table Comparison of methods on Gaussian linear nonlinear time-series data with different feature numbers and different nuisance feature settings.
Weak confounding interactions
Strong confounding interactions
Linear
nonlinear
Linear
nonlinear
Methods
FDR
F1
FDR
F1 score
FDR
F1 score
FDR
Fi score
LINGAM (CI)
.00 (.00)
.94 (.04)
.83 (.00)
.17 (.00)
.00 (.00)
.94 (.04)
.83 (.00)
.17 (.00)
PCMCI (CI)
(.01)
881 (.04)
.12 (.08)
.85 (.05)
1.0 (.00)
(.00)
.63 (.23)
.36 (.22)
GVAR (GC)
.44 (.00)
.11 (.00)
.94 (.00)
.11 (.00)
.94 (.00)
.11 (.00)
.94 (.00)
.11 (.00)
GrID-net (GC)
1.0 (.00)
.00 (.00)
1.0 (.00)
.00 (.00)
1.0 (.00)
(.00)
1.0 (.00)
.00 (.00)
DCM (GC)
.12 (.20)
.88 (.20)
.65 (.12)
.35 (.12)
.18 (.09)
.2 (.09)
.93 (.11)
.07 (.11)
NGM (GC)
.07 (.08)
.88 (.04)
.48 (.17)
.50 (.17)
.00 (:00)
.11 (.00)
.62 (.25)
.38 (.25)
GEASS (Ours)
.05 (.15)
.97 (.10)
.03 (.06)
.2 (.05)
.03 (.07)
90 (.04)
.00 (.00)
.91 (.00)
Furthermore, we evaluate different methods' scalability with respect to the feature size. (Experimental details can be seen at Appendix 1.1.2). As described before, we anticipate high computational complexity of both conditional independence based methods and neural network based methods with respect to the feature size, which prohibits further use these methods for nigh-dimensional biological data analysis, where the fea- ture number is typically at the scale of 10 10t. Meanwhile, GEASS constructs single neural network with parameters approximately proportional to p, thus largely reducing the com- plexity in the high-dimensional regime. We benchmark PCMCI, GVAR, GrID-net, NGM, GEASS, and an additional combina- tion of GEASS with a downstream CI-test based causal graph identification method LPCMCI. Our experimental result shows the superior performance of GEASS as well as GEASS+LPCMCI in time complexity, consistent with our qualitative analysis (Figure 3
16'
PCMEI
GIARR
GH13 not
OEASB+LPOMO
Foutr 60be
Figure 3: Running time compari- son for methods on time-series data with different feature sizes.
4.2 SIMULATED SPATIAL OMICS DATA WITH CELL TYPE CONFOUNDER
In order to jointly consider spatial confounders and corresponding autocorrelation patterns that are potentially enriched in specific niches, we consider the case of spatial omics data, where the autocorrelation is modeled by higher likelihood of same type of cells in the neighborhood, and the confounder (nuisance features) is modeled by a coherent shift of global gene expression for each cell type. We first simulate scRNA-seq datasets, then each synthetic scRNA-seg dataset is assigned to fixed size grid with cell type labels simulated by Ising model simulation. We then add artificial genes that are spatially correlated with neighboring cell's given gene set. Finally each dataset normalized and loglp lp transformed as the standard pipeline in Scanpy (Wolf et al., 2018).
The majority of methods are not available their focus on time-series data. Therefore in order to perform our benchmarking study, we compare GEASS with Lasso Granger, well as our implemented L.-rregulrrized version NCEM, an approach proposed to detect interactions in spatial omics data (Fischer et al., 2021). Finally, also implemented method that maximizes over the original transfer entropy to select causal features (TE).
As shown in Table the original LASSO cannot identify causal features because of the strong correlation between features. LI-NCEM alleviates the issue by conditioning on cell type labels in regression. TE outperforms linear methods yet generates number of false positives, as it may learn spurious causations as discussed in Remark [33] Finally, GEASS consistently outperforms over other methods in identifying causal features of data shown by both high F1 score and low FDR.
7

---

Table 2: Comparison of methods on simulated spatial transcriptomics data.
Linear
nonlinear
Methods
FDR
F1 score
FDR
Fl score
Lasso
0.050=000055 0.050=00055
0.030=00.000 0.970=0.040
LI-NCEM
0.380=00138
0.620#0.138
0.535=0.134
0.465=0.134
TE
0.190:10.127
0.767110.070
0.141=00087
0.761=0.060
GEASS (Ours)
0.195
0.787-10.088
0.000:10.000
0.775:10.110
4.3 SCRNA-SEQ PANCREATIC ENDOCRINOGENISS, TRAJECTORY
We test GEASS on the pancreatic endocrinogenesis trajectory data, which is a standard dataset for scRNA-seq trajectory inference task (Bergen al. 2020 Bastidas-Ponee et al. (019).. The pancreas trajectory data contains 3696 cells and 27998 genes. After preprocessing, lowly-expressed genes are filtered out the standard pipeline in s scVelo (Bergen al. 2020), with remaining 2000 genes for further analysis. We aim to use GEASS to identity causally-related genes along the developmental trajectory to reveal underlying biology. (See Appendix for experimental details).
scRNA-seq data provides snapshot of cell population distribution therefore time-series based analysis methods cannot be directly applied. However, due to GEASS's flexible setting in forward operator W. we are able to define the time flow by RNA velocity analysis. RNA velocity analysis uses the additional information intron RNAs to infer the underlying dynamics of gene expression change. Thus, we are able to define a velocity kernel matrix Aneto, which provides weighted adjacency relationships of cells based on velocity direction and cell phenotypic proximity.
GEASS identifies 50 causally-relace features with high biological relevance. For example, the gene list includes the key transcriptional regulator NEUROG3, which is required for the specification of common precursor of the 4 pancreatic terminal states (uni 2021). As the ground truth causal interactions here are unknown, for further quantitative validation, We assume the underlying biological process is driven by causal cascade of gene interactions, meaning target genes activated in earlier phases of the trajectory further cause downstream gene activation at later phases. In this case, the higher gene velocity is, the more likely the gene is associated with causal gene-gene relationships. Our benchmarking result here suggests GEASS achieves the best performance in selecting genes with high mean velocity likelihood, compared with alternative gene selection schemes with fixed gene number (50) including high-expressed genes (HEG), variilev genes (HVG), and genes having high correlation with inferred latent time (HCG) (Figure 4).
Mean RNA velocity likelihood
HEG HVG
0.0528
0.1753
HCG GEASS
0. 081889
0.2366
Figure 4: Visualization of pancreas trajectory dataset and comparisons of gene selection criterions.
4.4 MERFISH HUMAN CORTEX SINGLE-CELL LEVEL SPATIAL TRANSCRIPTOMICS
Spatial transcriptomics represent wide category of method that can achieve spatial profiling of gene expression in tissues (Moses and Pachter 2022 Rao et al.) 2021 Palla et al. 2022). By the
additional information of spatial locations, such measurements enable deeper understandings of cellular interactions (Palla al.) 2022 Jerby-Arnon and Regev 2022 Fischer et al. 2021). However, current computational methods revealing interaction modules (Jerby-Amon and Regev 2022 or niche effects (Fischer et al. 2021 Raredon al. (023) for spatial omics data lacks causal interpretation. Applying GEASS, we aim to reveal underlying causal intercellular patterns to fully utilize the potential of spatial omics data for biological discovery.
8

---

Here we use GEASS on recent published MERFISH dataset measuring peatially-resolvee single-cell gene expression of human cortex (Fang et al. 2022). The dataset we used comprises of 3044 cells and 4000 genes; each cell is annotated one of the eight cell types: excitatory neurons (EXC), inhibitory neurons (INC), astrocytes (ASC), microglial cells (MGC), oligodendrocytes (OGC), oligodendrocyte progenitor cells (OPC), endothelial cells (ENDO), and mural cells (MÚRAL) as shown by the first. panel of Figure in Appendix [D] Our GEASS analysis selects 9 genes, namely FILIPI, SLC17A7, 1YH11, RP11-10j21.2. PIRT, C30RF67, TRDMTI, RGS8, SPTLC2 (Appendix O) with further experimental details available in Appendix C.4 Among these genes, MYH11, RP 1T-1102....2 and TRDMTI are enriched at the endothelial cells adjacent with mural cells, corresponding to underlying vascular structures (marked by ellipses in the first panel of Appendix Figure 6 We next aim to verify their expression difference with those of non-adjacent endothelial cells is statistically significant. Indeed, by applying the Wilcoxon rank-sum test, we have found significant enrichments for both MYH11 and TRDMTI, with p-values 0.003 and 0.015 respectively, while the p-value for the gene RP11-10j21.2 is not significant (0.5) due to the gene expression sparsity. The finding is consistent with the MERFISH images, which reveals rich cellular interactions between neuronal cells and the blood vessels (Fang et al.) (022). Therefore, these identified marker genes of vaseular structure may encode underlying meaningful cellular interactions.
Next, we focus on two GEASS identified genes, C3ORF67 and PIRT, which are highly expressed at nearby spatial locations. In order to confirm the possible causal relation- ship between the two genes, we consider three models: I, the two genes are expressed in the same cell without spa- tial causal relationships; 2. The expression of C3ORF67 in each cell causes the expression of PIRT in neighboring cells (C30RF67 + PIRT); 3. The expression of PIRT in each cell causes the expression of C30RF67 in neigh- boring cells (PIRT + C3ORF67). To this end, we first compare Pearson and Spearman p-values of intracellular correlation (model 1), C3ORF67 to neighboring PIRT'
C3arl57
suatai
PIRT
1ualaa
Figure 5: Normalized spatial expression levels of genes C3orf67 and PIRT.
(model 2), and PIRT to neighboring C3ORF67 (model 3). Our comparison shows for the p-values of both correlation measures, model 3 is favored (0.004, 0.001) over model (0.014, 0.003) and model 2 (0.049, 0.004). The validity of model (PIRT + C3ORF67) is further supported by a linear model predicting C30RF67 expression by both intracellular and neighbor expression of PIRT, where the neighboring cell effect coefficient is significant at the confidence level of 0.01 by bootstrap, while the alternative model's corresponding coefficient not significant. Our finding is consistent with the predicted role of PIRT in transmembrane transporter binding and phosphatidylinosinn medid signaling (Safran et al. 2021). As the role of C30RF67 in human cortex remains unclear, this revealed causal link may lead to novel biological discoveries with further experimental validations.
5 CONCLUSIONS
In this work, we present GEASS, a causal feature selection method based on information thheorctic tools and neural networks. GEASS is able to scale high dimensions and identify sparse interacting features. We provide both theoretical gaurantees and empirical validations of GEASS on synthetic and real biological data. Our results show GEASS can be integrated into -ighdiimessiona spatiotemporal data analysis pipelines to provide unique insights for further findings.
Limitations. GEASS is a method designed for nonlinear causal feature selection. GEASS does not provide a causal graph itself as it optimizes latent embedding corresponding to different causal mechanisms. Therefore, in applications where causal graph output is favored, constraint- methods may need to be applied after GEASS. Moreover, when underlying causal graph has large number of vertices, the sparsity assumption is violated and GEASS is not gauranteed to work. Also, further efforts may be taken to incorporate lag selections for GEASS.
Broader impact. We anticipate a wide use of GEASS in nigh-dimensional graph structurd data, especially for high-dimensional biological data such single cell trajectories and spatial omics measurements. Applying GEASS along with causal graph identification methods to a wider range of real biological data may greatly facilitate downstream biological discoveries.

---

ACKNOWLEDGEMENTS
The authors thank Ofir Lindenbaum, Boaz Nadler, Yifei Min, and Ronen Basri for helpful discus- sions. Y.K. acknowledges support by NIH grants RO1GM131642, UMIDA051410, U54AG076043, P50CA121974, and U01DA053628.
REFERENCES
Robrecht Cannoodt, Wouter Saelens, and Yvan Saeys. Computational methods for trajectory inference from single-cell transcriptomics. European journal of immunology, 46(1 1):2496-2506, 2016.
Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature biotechnology. 37(5):547-555 2019.
Chee-Huse Huat Linus Eng, Michael Lawson, Qian Zhu, Ruben Dries, Noushin Koulena, Yodai Takei, Jina Yun, Christopher Cronin, Christoph Karp, Guo-Cheng Yuan, et al. Transeriptome-scale super-resolved imaging in tissues by ma seqfish+. Nature, 368(7751):235-222 2019.
Rongxin Fang, Chenglong Xia, Jennie L Close, Meng Zhang, Jiang He, Zhengkai Huang, Aaron Halpern, Brian Long, Jeremy A Miller, Ed S Lein, et al. Conservation and divergence of cortical cell organization in human and mouse revealed by merfish. Science, 37766601):56-62, 2022.
Lambda Moses and Lior Pachter. Museum of spatial transcriptomics. Nature Methods, 19(5):33-5546 2022.
Anjali Rao, Dalia Barkley, Gustavo S França, and Itai Yanai. Exploring tissue architecture using spatial transcriptomics. Nature, 596(7871):211 2021.
Giovanni Palla, David Fischer, Aviv Regev, and Fabian Theis. Spatial components molecular tissue biology. Nature Biotechnology, 40(3):308-315 2022.
Kevin D Hoover. Causality economics and econometrics. 2006.
Maciej Kamiński, Mingzhou Ding, Wilson A Truccolo, and Steven L Bressler. Evaluating causal relations neural systems: Granger causality, directed transfer function and statistical assessment of significance. Biological cybernetics, 15(2):145-157,22
Jakob Runge, Sebastian Bathiany, Erik Bollt, Gustau Camps- Valls, Dim Coumou, Ethan Deyle, Clark Glymour, Marlene Kretschmer, Miguel D Mahecha, Jordi Muñoz-Marí, et al. Inferring causation from time series earth system sciences. Nature communications, 10(1):1-13, 2019a.
Calls-Valls,
Marcos Herrera Gómez, Manuel Ruiz Marín, and Jesús Mur Lacambra. Testing spatial causality in cross-section data. 2014.
Clive WJ Granger. Investigating causal relations by econometric models and cross-spectral methods. Econometrica: journal of the Econometric Society, pages 424-438. 1969.
Ali Shojaic and Emily B Fox. Granger causality: A review and recent advances. Annual Review of Statistics and Its Application, 9:289-319, 2022.
2022.
Konrad Mielke, Tom Claassen, Huijbregts, Aafke M Schipper, and Tom Heskes. Discovering cause-effect relationships in spatial systems with known direction based on observational data. In International Conference on Probabilistic Graphical Models, pages 305-316. PMLR, 2020.
David S Fischer, Anna C Schaar, and Fabian J Theis. Learning cell communication from spatial graphs of cells. BioRxiv, 2021.
PA Valdés-Sosa, JM Bornot-Sánchez, M Vega-Hernández, I. Melie-García, A Lage-Castellanos, and E Ganales-Rodríguns granger causality on spatial manifolds: applications to neuroimaging. handbook of time series analysis: recent theoretical developments and applications, 18.
Thomas Schreiber. Measuring information transfer. Physical review letters, 85(2):461, 2000.
10

---

Matthäus Staniek and Klaus Lehnertz. Symbolic transfer entropy. Physical review letters, 100(15): 158101, 2008.
Doris Entner and Patrik O Hoyer. On causal discovery from time series data using fci. Probabilistic graphical models, pages 121-128, 2010.
Daniel Malinsky and Peter Spirtes. Causal structure learning from multivariate time series in settings with unmeasured confounding. In Proceedings of 2018 ACM SIGKDD workshop on causal discovery, pages 23-47. PMLR, 2018.
Alessio Moneta, Nadine ChlaB, Doris Entner, and Patrik Hoyer. Causal search in structural vector autoregressive models. In NIPS Mini-Symppa Symposium on Causality in Time Series, pages 95-114. PMLR, 2011.
Niklas Pfister, Peter Bühlmann, and Jonas Peters. Invariant causal prediction for sequential data. Journal of the American Statistical Association, 1 114(527):1264-1277 2019.
Atalanti A Mastakouri, Bernhard Schölkopf, and Dominik Janzing. Necessary and sufficient con- ditions for causal feature selection time series with latent common causes. In International Conference on Machine Leurning, pages 7502-7511. PMLR, 2021.
Aapo Hyvärinen, Kun Zhang, Shohei Shimizu, and Patrik O Hoyer. Estimation of a structural vector autoregression model using non-gaussianity. Journal of Machine Learning Research, 11(5), 2010.
Jakob Runge, Peer Nowack, Marlene Kretschmer, Seth Flaxman, and Dino Sejdinovie. Detecting and quantifying causal associations in large nonlinear time series datasets. Science advances, (1 (1))aaa444996, 2019b.
Shohei Shimizu, Patrik O Hoyer, Aapo Hyvärinen, Antti Kerminen, and Michael Jordan. A linear non-gaussian acyclic model for causal discovery. Journal of Machine Learning Research, 7(10), 2006.
Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffussinn convolutional recurrent neural network: Data-driven traffic forecasting, arXiv preprint arXiv:1707.01926 2017.
Alex Tank, Ian Covert, Nicholas Foti, Ali Shojaie, and Emily B Fox. Neural granger causality. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(8):4267-4279, 2021.
Meike Nauta, Doina Bucur, and Christin Seifert. Causal discovery with attention- based convolutional neural networks. Machine Learning and Knowledge Extraction, 1(1):19, 2019.
Zexuan Yin and Paolo Barucca. Deep recurrent modelling of granger causality with latent confound- ing. arXiv preprint XrXiv:2202.111 1286. 2022.
Alexis Bellot, Kim Branson, and Mihaela van der Schaar. Neural graphical modelling in continuous- time: consistency guarantees and algorithms. In International Conference on Learning Represen- tations, 2021
Saurabh Khanna and Vincent YF Tan. Economy statistical recurrent units for inferring nonlinear granger causality. arXiv preprint arXiv:1911.09879 2019.
Xiangyu Sun, Guiliang Liu, Pascal Poupart, and Oliver Schulte. Nts-notears: Learning nonparametric temporal dags with time-series data and prior knowledge. arXiv preprint 10XX::2009.02266, 2021.
Alexander P Wu, Rohit Singh, and Bonnie Berger. Granger causal inference on dags identifies genomic loci regulating transcription. In International Conference on Learning Representations, 2021.
Xianjie Guo, Kui Yu, Lin Liu, Fuyuan Cao, and Jiuyong Li. Causal feature selection with dual correction. IEEE Transactions Neural Networks and Learning Systems, 2022.
Judea Pearl. Causality. Cambridge university press, 2009.

---

Luca Faes, Daniele Marinazzo, Giandomenico Nollo, and Alberto Porta. An information-theoretii framework to map the spatiotemporal dynamics of the sealp electroencephalogram. IEEE Transac- fions Biomedical Engineering, 63(12):2488-244 249.2 2016.
Payam Shahsavari Baboukani, Carina Graversen, Emina Alickovic, and Jan Østergaard. Estimating conditional transfer entropy in time series using mutual information and nonlinear prediction. Entropy, 22(10):1124, 2020.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint AXXiv:8001004000 2018.
R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint I 808.00660, 2018.
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statistics, pages 875-884. PMLR, 2020.
Jingjing Zhang, Osvaldo Simeone, Zoran Cvetkovic, Eugenio Abela, and Mark Richardson. Itene: Intrinsic transfer entropy neural estimator. arXiv preprint arXiv: 1912.0777, 2019.
Yutaro Yamada, Ofir Lindenbaum, Sahand Negahban, and Yuval Kluger. Feature selection using stochastic gates. In International Conference on Machine Learning, pages 10648-10659, PMLR, 2020.
Ofir Lindenbaum, Moshe Salhov, Amir Averbuch, and Yuval Kluger. LO-sparse canonical correlation analysis. In International Conference on Learning Representations, 2021.
Andrew Miller, Nick Foti, Alexander D' Ammur, and Ryan P Adams. Reducing reparameterization gradient variance. Advances in Neural Information Processing Systems, 30, 2017.
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients.
Advances in neural information processing systems, 31, 2018.
Thomas N Kipf and Max Welling. Semi-supervise classification with graph convolutional networks arXiv preprint arXiv: 1609:06000000 (009.00999....... 2016.
Andreas Gerhardus and Jakob Runge. High-recall causal discovery for autocorrelated time series with latent confounders. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 12615-12625. Curran As- sociates. Inc., 2020. URL URLLsttps://prreeeiinng neuripsss.g.........op....i... http://aww............eaeeaaapee http:pppp..ppppppe////////////
Ričards Marcinkevičs and Julia E Vogt. Interpretable models for granger causality using self- explaining neural networks. arXiv preprint arXiv:2101.07600, 2021.
Alexander Wolf, Philipp Angerer, and Fabian J Theis. Scanpy: large-scale single-cell gene expression data analysis. Genome biology, 19(1):1-5,20) 2018.
Volker Bergen, Marius Lange, Stefan Peidli, F Alexander Wolf, and Fabian J Theis. Generalizing ma velocity to transient cell states through dynamical modeling. Nature biotechnology,388 38(12): 1408-1414, 2020.
Aimée Pustidas-Poncee Sophie Tritschler, Leander Dony, Katharina Scheibner, Marta Tarquis-Medina, Ciro Salinno, Silvia Schirge, Ingo Burtscher, Anika Böttcher, Fabian Theis, et al. Comprchensive single cell mrna profiling reveals a detailed roadmap for pancreatic endocrinogenesis. Development, 146(12):dev) (ed)vve33449 2019.
Uniprot: the universal protein knowledgebase in 2021. Nucleic acids research, 49(D 2021.
Livnat Jerby-Arnon and Aviv Regev. Dialogue maps multicellular programs in tissue from single-cell or spatial transcriptomics data. Nature Biotechnology, pages 1-11, 2022.
12

---

Micha Sam Brickman Raredon, Junchen Yang, Neeharika Kothapalli, Wesley Lewis, Naftali Kamin- ski, Laura E Niklason, and Yuval Kluger. Comprehensive visualization cell-cell interactions in single-cell and spatial transcriptomies with niches. Bioinformatics, 39(1 bbbtac755 2023.
Marilyn Safran, Naomi Rosen, Michal Twik, Ruth BarShir, Tsippi Iny Stein, Dvir Dahary, Simon Fishilevich, and Doron Lancet. The genecards suite. In Practical guide to life science databases, pages 27-56. Springer, 2021.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Dylan Kotliar, Adrian Veres, M Aurel Nagy, Shervin Tabrizi, Eran Hodis, Douglas Melton, and Pardis C Sabeti. Identifying gene expression programs of cell-type identity and cellular activity with single-cell ma-seq. Elife, 2019.
APPENDIX
A PROOFS
A.I PROOF OF PROPOSITION 2.2.
Proposition 2.2. VS1, S2 C {1,.... p}, mT, 0 (S, S2) >
Proof. By standard properties of mutual information Cover 1999 we have
(Xs, Xs,) T(X's (...)E
x3))
S2
I(X X X
((G))E
.x I(X X3s
(9)
(XXss X X (.L 1 (3s,) (XX's ;x's Sx E.
Therefore (Si S2) Z ( S2) holds, thus (S S2) > 0 (SI S2) > 0.
A.2 DISCUSSION OF ASSUMPTION A6.
Our assumption A6 is based on the concept of conditional mutual entropy, which aims to filter out possible indirect causal relationships.
Here are two simple examples to see why TE/TT can have problems with indirect causal interactions the time-series setting: consider the relationships: 8t - Wi + U1415 + W+++1 + Ut+1 Then in both cases, may have: I(st, vt+1) - I(st, vt) > 0 and I(st, tt+1 (0y) > 0 although there direct causal relationship between s and v. Note in our setting, we include the possibility of such indirect interaction by allowing correlation between nuisance features and true interacting features.
The issue can be resolved by considering the conditional mutual information I(st, ut+1 1 (w)) or [(st, ut+1 /w)+1, which equals 0. This insight is also addressed the concept of conditional transfer entropy:
Definition (Conditional transfer entropy) Shahssvaaii Baboukani al. 2020). Assume X and Y are the features of interest and the conditioning features are Z. Denote as 1,2 ...,, 1). then we have
CT,
The classical formulation of conditional transfer entropy is widely used in high-dimensional ob- servational data to learn direct causal dependencies (Faes et al.) 2016 Shahsavari Baboukani ct al., 2020). implicitly assumes that, there is direct causal relationship between X and Y if VZ, t, CTE, (X, 2) > O. Here, we extend this assumption in the context of conditional mE covering both examples described above. The conditional mTEs are defined in analogy to cTE for generalized graph-structured data in the Markovian model setting:
Definition (Two forms of conditional mTE). Assume X and Y are the feature sets of interest and the conditioning features are Z. Then we have
TE (X, Y, 2) T(x (z) 7(X', Y(
13

---

x I(x,, y ( - I(X' -(Xx''''
By controlling the two forms of conditional m'TE to be larger than zero, we rule out both possi- bilities of X - Z + ( and X + ZN(1) - YN()), as mTEE is a stricter version of the original transfer entropy discussed in Proposition 2.2 In summary, our A6 be reformu- lated as AZ, i, Tm (X, Y, Z) > 0: cm (X, Y, Z) 0 for ground truth interacting X, Y in non-degenerating cases, where Z does not fully overlap with X/Y the same point.
A.3 PROOF OF THEOREM 2.4.
Theorem 2.4. Given AI-A6, S* == USTUS) () (the index set of true interacting features described in A5). Moreover, each feature in 8 is connected to other features in the set S
Proof. Step 1. First prove S 0. If not, assume p is an overlapping element. For simplicity, we denote (()) == j) E F}, A = Xst Xsg- Then we have
MTE(S,, (5) M'EES \ (,, S5)
T(AA' (N() N() ((N()))) p'' T(A' BN;))()) ( (p p)
+ +[(A'B p')
= - \ N 10 IA
- -
A (p) <0.
(10)
Therefore removing would increase the value of mTE, leading to contradiction.
Step Now we prove nuisance signals cannot be in either Si Si. Otherwise, first we assume set of nuisance signals U is in Sj. Here we denote A = X B = Xs;. As U' only interacts with variables at the same time point, U can only interact with (N) via indirect links through subset of interacting features at i. Denote this feature set as Pay (B)& C (4i:.., (,.. Im J. and the difference set Pay B = Pay (B)' \ B'. Then we first note (B)' cannot be an empty set. Otherwise, denote S ST NU, noting the non-overlapness between A and B we would have
MTE(S1, 52) TE(S,
I((A' \U, NO) I(A' U' \U',
-(AA (*)
= I(U';BB()) BN0 AA \U)) - A' \U*)
(11)
NO \\) + + 1² U+)
I NO A + + PoU (B)² A \. U*) (Conditioning reduces entropy)
M O.
This means (S1, S3y's mTE is not smaller than (Si, St)'s while having a smaller union size, leading contradiction. Then because Pary (B) does not overlap with either U and B, with A6 we have
TTESS S5) mTE(S) (nd)
= =T(A' D', - T(A' \ B+) - T(A* U U", PBT( BN))
(12)
+1(A') (U,," (B);;
= (B)) B¹ A') I(Pa (B)); B3) 1() v 0.
A6
14

---

The equal sign above is taken iff. (B) _ A'. Further we have
TT U (udex(Pai(B))).))))
mTEE(i U (ndex(Pa
= =I(A' U", Paj - [(A' \ U², Pay)
I(A' \ U², (B) (AA') U² Pay B²)
(13)
= =T(U"; B PN (B) (0)) - I(U"; B Pay (B)', , U*)
- Paj (B), A' A'\(' (U") + +(U''BB Paj (B), A' (0)) S - Pay (B)' A' U*) + h(U (B) A' U 0.
Therefore, in all possible cases, m TE(S U Index(Paç (B)', S²) is either strictly larger than mTE(S}, or equal with TE( 55) but with smaller union size, leading to contradiction.
Next, given the result above, we assume nuisance signal set U is in Sz, and Sj does not include any nuisance features. Then as U only interacts with variables at the same time point, (N) can only interact with Sf via indirect links through a subset of interacting features at (i). Denote the whole intermediate feature set for A as Ch)( C [t) N())
..... tm }. and ChT
,N(0)
Chb(AA) NO) \ (()), Then same above, denote S2 S2 U, if Chy (A) is empty set would have
(B( MTE(S S2)
=I('', BNO \ UNO,, (UNOUNNO) ( I(A': (U,,B - I(A; \ UN(O)
+ +T(A'; B¹ UU)
(14)
=I(A'; \ ( - U
-MUUUNO BBO (UNO, A') M(U* \ A') S -A(UNO) NUMO,4Y + r(U (C)((A)) B U*) 0o.
Above derivation holds due to stationarity (as [N(1) = 1 in the time series setting). Therefore Chū (A) cannot be an empty set. Because of the non-overlapness between (4) and either A or U, with A6, have
TTE(S), (2) MTE(S) S2 U (ndex(((((A))
= =I(A'; BN) (UN(, (ON(O,UNG)) (U) - I[A': (U'UU*)
(15)
-I(A'; BNO) (UNO) (., Cl) + B' (hT()))
= -(AA,CCT - -IAA',ChT
(() - 0.
The equal sign above taken iff. (A) C B Further we have m TE(S, Sz U Index(ChT5
TT.(ST S2 Undex (CTT((()C -[AA'; BNO) NUNO, (O) (O) N(O) 1(A'; \U', UK, (A)))(
(16)
-I() BNG) \ (A)())))))) + ((A:B'' (A)())
= -IAAFUNNIIBN (O, (A) I(A':U"IB' AU, (h)))))) S0.
Therefore, in all possible cases, mTE(S S2 Index Index(Chh (A))) is either strictly larger than mTEES St) or equal with mTE(S), S2) but with smaller union size, leading contradiction.
Step 3. Moreover, there exists a component in S U S not connected to any other feature components, denote the feature as q. Then, in this case with A1-4, the feature g is independent of any other features in ST U US From step it can be deduced that cannot be in both S7, Therefore in this case, have mTE(S; q, S q) mTE(S), St)) thus leading the contradiction of finding an (S1, S2) with the same m'EE but smaller [S] U S2)-
15

---

A.4 PROOF OF THEOREM 3.1.
Theorem 3.1. Assume Al-A6 holds and J 9- h define one-to-one mappings on Ools, (for f] or Xo
1s, (for g. h). Then YE such that for 4 any solution (Si U US) satisfies S = (Si U $
{I,... m.. Moreover, feature in S is connected to other features in the set.
min (II(f(z' の 1s:):h(x 1s,)) I(f(x' @ 1s,): 1s,))) + X|Sk U. S1
Proof.
With A4 (ergodicity and stationarity), the optimization problem equivalent to ([; ( ((((ws,)::( ))) + X|S U2/-
min St
(17)
Given the assumption that f. g, define injective mappings on a's, x's respectively, and one-to-one transformation does not change mutual information, we have the optimization problem is equivalent
min (1's's 52 ((0) I ( sss)) + A|S U S2). $1.22
(18)
Using Theorem 2.4 minimizer of the mTE term with the smallest union size satisfies S* (Sf U US) S (1.... m). Moreover, each feature in S U S is connected to other features in the set. Note that with our definition of optimal S1, S2, the minimal gap between m'TESSi, ()) and any other value m mTE((), S2) with smaller |S| U S2 size is larger than zero. Denote the minimal gap as 6, and take À < [8fUSh]
then for these other solutions, we have
mTE(S), (2) + X|S) S2 2 2-mTE(S1 55) + X U
(19)
2 -TTT(( 52 +
>-mTE(Si (5) iSlk Sy
Meanwhile, for the (S1, with larger union size, with the definition of the mTE, have
TT 52) + X/Sk U.S
2 TTTESS) S5) + X|S U S₂1
(20)
-mTEESS, .5:) X([Si S22 - |S] US Ss5k US!
> -mTFF5T + S U USI
Therefore, when taking (0, [TUS the desired optimal $1, S2 by mTE is the optimal output of the constructed optimization problem.
A.5 PROOF OF THEOREM 3.4.
Theorem 3,4. Assume A1-A6 and f.g. g, l are one-to-one Gaussian embeddings as described above.
Denote for the optimal solution of (5), sample of stochastic gate given T T and denote the ground truth interacting feature set as S, then there exists \1, X2 > for (5) such that as - 20,
Vi E {0, 1}. P(B, S S) 0.2. 1, where B; { > 0, 0 =1.
Proof. In the following proof for simplicity we denote ēs = x TT ; 1 o(1-T2)
Step 1. Given f.9, h projects input distributions into joint Gaussian distributions with fixed dimen- sionality, by convergence of Gaussian covariance matrices, we have:
S(f(5))), N()) 1
> If( h( N No [[ h(
n
h N a.e, E s(as, ).M(
n
)(((a')) 9(9
[[ ): ( ) a.ey Ej(es, )
(21)
16

---

As in the Gaussian case, the mutual information between jointly Gaussian E.V.S is function of the covariance matrix, we have
( - I( (.) ( 4 Sz
N(1):
( g Buli (( )s [ #'s)):
(22)
P( lim Empirical mTE = mTE) = 1. N
Step 2. Importantly, in our formulation eq 5 the Tj, T2 are sampled once in one epoch, mean- ing they are fixed across features for computing mTE. Further note that Σ P(T.) > 0)
1 we have
Eis P(T E (0, = E 1T"000,)) This means denoting the value of eq (5 as I,
L 9 E T2 m TE m
( w > re
(23)
min -m - (reu)13))))) + X || ll + re T:.2
Note with step of the proof of theorem 2.4 for any Ti we have
m (r + \ 1b + X2 T
(24)
which is taken when vd, = 1) -
0) = 1/0. this case,
| = || o 1| o (1-T
Applying theorem 3.1 we have for W in theorem 3.1
min T'T
+ X \ 1o + de LT
(25)
=-mTE(S S5) i/lk S L'
Here (ST, S2) satisfies properties described by theorem 3.1
Note the minimizer may not be unique,
denote the containing all minimizers [(S S }. Then the equal sign in eq (23 holds
if and only if o >0 (1-)(()))) E [(S { = 1. Further noting PTT = 1) P(T = 0) 1/0, and our analysis above holds convergence, we finally have
oc with probability by a.s.
PP lim
P(B1 C S)= = 1)= =1
lim P(Bo S 5) = 1) =1 N
N
holds.
B GATE INITIALIZATION
Our proposed initialization scheme based on analysis of the linear case. Assume
s(Xs,) = Xa, (Xx)) Xb,
where a, IP represents two feature loadings. Then:
1. a, b should be non-overlapping. therefore we expect |a T b) be small.
2. We should have f(X) Wg(X) to maximize the mTE.
The constraint can be formulated into regression problem WXb Xo, therefore nat- ural solution is given by a = Xb (X" xx" In this case, ||a"b||| ||b XX= Xb|l
[ Given is normalized, can shown that optimal corresponds to the eigenvector with least absolute eigenvalue of matrix (X"X)-IxTW
After getting a, b, we select a quantile threshold over a/(a + b) to initialize the second stochastic gate layer. The first stochastic gate layer is initialized with uniform weights.
17

---

C EXPERIMENTAL DETAILS
C.I TIME-SERIES BENCHMARKING STUDY
In the study the causal processes is simulated with Python package Tigramite. Among the total 100 features, there are 6 interacting features {1, 5. 5. The causal links are: 1->2 with time lag 2, 2->3 with time lag 1, 5->4 with time lag 1, 1->5 with time lag 1, 3->6 with time lag 3. These features also have autocorrelations with time lags ranging from 1 to 3. There is also latent confounder modeled by Tigramite interacting with feature 0 and feature 2. In the case of strong latent process, the latent confounder also have effects on other 43 features. All other features (93/50) not mentioned above nuisance features with white noise dynamics. The forward operator is defined by 5-neighbor lower triangular matrix.
C.1.1
ALGORITHM IMPLEMENTATION
If not particularly mentioned, default settings of the algorithms used throughout.
VAR-LINGAM.
The VAR-LINGAM algorithm is implemented in the Python package I 5/lingam. VAR-LINGAM gives a weighted
LINGAM, available at matrix output. Therefore in our benchmarking study, we choose the most significant edge corresponding features with the number matching the sparsity level.
PCMCL The PCMCI algorithm is implemented in the Python package Tigramit which gives a weighted matrix as output. We choose the most significant edge corresponding features with the number matching the sparsity level.
The sparsity parameter is set to be 1. We use the stable training option in GVAR, which trains the first and second half of the time series respectively to optimize over edge selection sparsity level then train the whole time series, giving binary output and threshold selection is needed.
GVAR. The GVAR algorithm implemented at http:///hhubb.coo///99964//7//7
Grid-net. The Grid-net algorithm is implemented at https//ggthuub.co/alewwww 16/gridnet. The parameter set: order=5, hidden_ loaden_layerrsiiee Imbd=1 is used throughout our study. After the training finishes, choose the most sigmificant edge corresponding features with the number matching the sparsity level.
10, end end_epoch=551 batch_size 50,
DCM, NGM. The two algorithms are both implemented at ciiiiiaallmmmeleennnoonooiiiiiii
For DCM, the default setting is used, and we use hidden dim = 10 for NGM. After both training finishes, we choose the most significant edge corresponding features with the number matching the sparsity level.
GEASS. We use the same training parameters in all time-series settings, with the key sparsity regularization parameter λ set with 0.04/0.05 based on a validation set, and the rest parameter settings are consistent with default.
C.1.2 SCALABILITY ANALYSIS
We test PCMCI, GVAR, GrID-net, NGM, GEASS, GEASS+LPCMCT' running time with con- sistent settings described in the above section. (LPCMCT's setting is consistent with PCMCI's setting). We use the same data generation pipeline and select the of the total feature numbers as [100, 200, 400, 800, 1600).
C.2 SIMULATED SPATIAL OMICS DATA BENCHMARKING STUDY
In the study the spatial omics data is simulated with Python package Scsim (Kotliar et al., 2019). 1000 genes are simulated in total, while 990 genes are cell-type-speciifely expressed. The rest 10 genes each has a functional relationship (linear/nonlinear) with one cell-type-specific genes plus the noise term in order to model the cell-type-specific interactions. The data is then normalized and log- transformed according to the standard Scanpy pipeline (Wolf et al. 2018). The forward operator is defined by 4-neighbor adjacency matrix.
18

---

C.2.1 ALGORITHM IMPLEMENTATION
If particularly mentioned, default settings of the algorithms are used throughout.
Lasso Granger. The Lasso algorithm is implemented by Scipy with tuned (0.12) to match the sparsity level.
NCEM. NCEM (Linear) is linear graph neural network, which in the grid case corresponds to standard linear regression based on neighbors and the cell type label. Based on the original work, we implemented our equivalent version by Lasso regression with a = 0.019 to match the sparsity level.
GEASS, We use the same training parameters in all settings, with the key sparsity regu- larization parameter Å with 0.02 based on a validation set, and the latent dimension number is set to be 64.
TE. To give a fair comparison, we use the same architecture as GEASS except for the loss function is changed. We use the same training paraneters all time-series settings, with the key sparsity regularization parameter ê set with 0.05 based on a validation set, and the latent dimension number is set be 64 consistent with GEASS.
C.3 SCRNA-SEQ PANCREAS TRAJECTORY
The data preprocessing is consistent with the scVelo tutorial: https://scvelo. https://swelo.reathhedssss io/ Veloo- ityBasics/ (Bergen et al. 2020 The parameter set: A = 0.06, A = 0.1. Here because the gene regulatory network is fully connected and activated in cascade along the developmental trajectory, consider the opposite initialization with be the largest eigenvalues corresponding eigenvectors of the matrix (x XX
C.4 MERFISH SPATIAL TRANSCRIPTOMICS DATA
The data is downloaded from Dryad and preprocessed with the standard Scanpy pipeline (Wolf et al., 2018): first normalize and log-transform the data by default functions in Scampy then select 1000 highly variable genes by default functions in Scanpy (Wolf et al. 2018). The forward operator is defined by 5-neighbor adjacency matrix. The GEASS parameter set consistent with those used in the spatial omics benchmarking.
D ADDITIONAL EXPERIMENTAL RESULTS
TIPP
SUDITAT
SPIH-1891
Figure 6: Spatial profiling of MERFISH human cortex slice, colored by cell type annotation and GEASS identified gene expressions. The ellipses in the first panel represent examples of vascular structures.
19


Reviews:
[{"review_id": "yJymoaoj8gw", "reviewer": "Reviewer_uZbN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors introduce a framework of causal discovery in data using two steps: a causal feature selection in a general graph structured data and causal graph identification. To do the former, the authors define a multi-dimensional transfer entropy maximization metric and device ways to optimise it. They test out the framework in single cell and spatial data.", "summary_of_the_review": "The paper addresses an interesting problem. The main idea is neat and is benchmarked sufficiently on quite a lot of data. Thus I recommend acceptance. ", "strength_and_weaknesses": "The problem of detecting causal features in spatial transcriptomics is an important problem. The authors promote a nice framework and benchmark on a variety of datasets. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Qp2ZFtFJlZ1", "reviewer": "Reviewer_nQe4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper provides a novel methodology for identifying significant causal features in spatiotemporal data, with primary applications in biology (e.g. scRNA-Seq data). Its primary novelty lies in the judicious use of a transfer entropy characterisation of feature significance/relevance and the design of a combinatorial stochastic gate layer.", "summary_of_the_review": "The paper is an advance over the current state of the art and provides interesting results in several biological applications.", "strength_and_weaknesses": "The strengths are as follows:\n\nThe paper is well-written and coherently organised.\n\nIt has clear novelty in the use of maximum transfer entropy.\n\nAll the non-trivial mathematical statements are supported by proofs in the appendix.\n\nThe weaknesses are as follows:\n\nIt is unclear how the forward operator translates to spatial settings; the authors vaguely mention something to do with diffusion, but that would not be applicable unless the diffusion trajectory could be reconstructed in time.\n\nIt is also unclear how the mutual information is computed (via approximation? based on frequencies in the data?).\n\nMany proofs are hand-wavy; this may make them clear/obvious to experts, but for the general reader they should be made more explicit.\n\nSome of the extensions of the method to specific datasets seem somewhat ad hoc, rather than being motivated by the same principles as the main method.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NOxI80mH81", "reviewer": "Reviewer_WJXd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper identifies causally interacting features of high dimensional temporal/spatial data by considering the sparsity of underlying causal mechanisms instead of link sparsity, which can select the critical corresponding features for downstream causal discovery. Theoretical studies are provided, and empirical evaluations on synthetic and real biological data showed its superiority over competing methods. ", "summary_of_the_review": "I think this is a solid work. ", "strength_and_weaknesses": "Strengths: 1. This paper proposes a novel solution of causal feature selection problem in general graph-structured data. \n                  2. Theoretical analysis is provided.\n                  3. Empirical results showed outperformed performance. \n\nWeaknesses: 1. I think the current version takes too much space on the theoretical analysis while ignoring the model itself. Figure 2 is not clear to me. \n                      2. For the Merfish data experiment, what's the reason to downsample the gene number from 4000 to 1000? In real situations, the gene number can be millions. If the proposed method can not deal with high-dimensional features, then the scalability will be an issue. \n                      3. What's the criteria for GEASS to select 9 genes in Merfish data?  Can the authors elaborate more /give insight on interpreting the selection of this subset for biological discovery?\n                     ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IKmhzOC4Jj", "reviewer": "Reviewer_Qga7", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The manuscript proposed a feature selection model for the causal discovery task based on the optimization of multi-dimensional transfer entropy of the input data by combinatorial stochastic gates. The proposed model was evaluated on by synthetic time series datasets with the presence of latent process, synthetic spatial progression of the scRNA-seq data, as well as the real (pancreatic endocrinogenesis trajectory) data. Experiment results indicate superior performance of the model over both classic and recent causal discovery methods. ", "summary_of_the_review": "A causal discovery framework with fundamental and important algorithmic development (Theorem 3.1 and 3.4), with good practical value as well. ", "strength_and_weaknesses": "Strength: The methodology involved in the optimized feature subset selection for causal discovery (Theorem 3.1, 3.4) and the design of the stochastic gate-based approach is important to the field of data mining. Experiment results show superior performance on both synthetic and real data, further validating the effectiveness of the model.\n \nWeakness: First of all, in section 3.2, the author mentioned that “GEASS provides both outputs of active features and embeddings produced by causally interacting features. In this paper, we emphasize the use of the former as the latter embedding output may be complex and nonlinear, potentially requiring additional architectures to maximize its interpretability.” The statement is confusing as it did not specify where the “embeddings produced by causally interacting features” were obtained: whether they were from the stochastic gates, or the MLPs indicated in Fig. 2? Further, there are no explanations for the design and purpose of the MLPs in Fig. 2.  Also, the term “STG layer” in Fig. 2 has never been mentioned throughout the manuscript. The reviewer guesses it is referring to the stochastic gates but is not sure about it. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%