GEASS: N EURAL CAUSAL FEATURE SELECTION FOR HIGH - DIMENSIONAL BIOLOGICAL DATA
Anonymous authors
Paper under double-blind review
A BSTRACT
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 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 multi-dimensional 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.
I NTRODUCTION
Advances in single-cell omics research enable full characterizations of high-dimensional gene dynamics in biological systems on a either temporal or spatial scale.
An example for the former is single-cell
RNA sequencing (scRNA-seq) trajectories, where cells are sampled from a 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 latter is single-cell level spatial transcriptomics (e.g. SeqFISH+ (Eng et al., 2019), Merfish (Fang et al., 2022)), in which cells from a tissue slice are sequenced with their spatial coordinates preserved (Moses and Pachter, 2022;
Rao et al., 2021; Palla et al., 2022).
Spatial profiling allows investigations of the cellular interplay, corresponding to conditional gene expression change caused by neighborhood 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 or 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´nski 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, a feature in the future cannot have effect on past values of other features.
For data with spatial structure, a similar definition of causal dependency can be established based on the irreversibility of the spatial local diffusion operator.
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 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 a spatial context corresponds to predicting significant relationships between neighboring observations of other variables and the specified variable (Mielke et al., 2020), which is a key insight used in recent works aimed to discover cellular interaction patterns in spatial omics data (Fischer et al., 2021;
Valdés-Sosa et al., 18).
In the nonlinear regime, information-theoretic measures such as directed information, transfer entropy (Schreiber, 2000), and partial transfer entropy (Staniek and Lehnertz, 2008), are used as a counterpart of linear Granger causality.
Moreover, some works consider modeling conditional independence in time-series data to identify the underlying causal graph (Entner and Hoyer, 2010; Malinsky and
Spirtes, 2018; Moneta et al., 2011; Runge et al., 2019a; Pfister et al., 2019; 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-ODE based approaches, resulting in improved prediction power for nonlinear time-series dynamics (Li et al., 2017; Tank et al., 2021; Nauta et al., 2019; Yin and Barucca, 2022; Bellot et al., 2021).
Despite the success of these methods in various systems of interest, multiple challenges limit their use in previously discussed high-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 independence generally scales exponentially with feature size, applying CI-test based methods in data with high feature dimensionality is not realistic.
Meanwhile, Granger-causality based methods build prediction models with the number equal to feature size.
Therefore the time complexity of solving the stacked model is of polynomial level with respect to the feature size.
• Edge sparsity is usually assumed by previous methods to maximize interpretability of the identified causal graph.
However, In biological data, there exists a large proportion of nuisance features (e.g. constitutively expressed genes), therefore the feature sparsity and edge sparsity need to be both included by an ideal model.
This is of particular interest as false discovery rate (FDR) control is favorable in biological data analysis.
• While a 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 applied 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 a single neural network.
GEASS considers sparsity of underlying causal mechanisms instead of link sparsity, thus selects most significant corresponding features for downstream causal discovery.
Our contributions are three-folds.
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 multi-dimensional transfer entropy maximization with theoretical guarantees.
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.
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. R ELATED WORKS Neural Granger causality
Despite the large body of work based on linear Granger causal discovery, neural Granger causality still remains an active area of research.
Various neural network architectures, such as MLP, sequential model, and attention-based architecture (Tank et al., 2021; Nauta et al., 2019; Khanna and Tan, 2019;
Sun et al., 2021), have been proposed for nonlinear Granger causality 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., 2021).
Meanwhile, a neural-ODE 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-based 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 et al., 2011; Runge et al., 2019a).
Meanwhile, seqICP focus on identifying the direct or indirect cause for each feature assuming sufficient interventions in the dataset (Pfister et al., 2019).
SyPI tackles the causal feature selection problem without the assumption of causal sufficiency and avoids issues in multihypothesis testing by construction of the correct conditional set (Mastakouri et al., 2021).
Finally, Guo et al.
(2022) considers dual correction of causal feature selection to control both false positive rates and false negative rates.
M ULTI - DIMENSIONAL T RANSFER ENTROPY
In order to tackle the issue that a neural network may overfit each model therefore overestimates causal interaction numbers, we need a prediction-free loss function that directly indicates causal signficance.
In this work, we propose a novel function, multi-dimensional transfer entropy, based on transfer entropy as a metric of causal interaction significance.
Transfer entropy is a information-theoretic measure of cross dependence.
(Schreiber, 2000) Consider two time series X t and Y t for t ∈ 1 , ..., T .
Transfer entropy aims to model nonlinear dependency across X t and Y t while accounting for possible autocorrelations.
It has been theoretically shown that transfer entropy value is equivalent to Granger causality statistical significance for Gaussian data (Barnett et al., 2009).
The transfer entropy can be formulated as: TE h ( X, Y ) =
I ( X t ; Y t + h | Y t ) .
Denoting the time/space forward operator by W , the transfer entropy can be generalized as:
[Equation]
In this work, we propose a novel function, multi-dimensional transfer entropy to extend the bivariate transfer entropy to higher-dimensional settings.
The key insight is to consider two feature subsets
S 1 , S 2 ⊂ { 1 , ..., p } , such that the time-lagged mutual information over two multivariate random variables X S 1 , X S 2 is maximized.
Definition 2.1.
Assume X ∈ R n × p to be the data of interest with a forward operator W .
S 1 , S 2 ⊂ { 1 , ..., p } each represents a set of feature subset indices.
The multi-dimensional transfer entropy (mTE) is defined as:
[Equation]
And the estimator ( ˆ S 1 , ˆ S 2 ) can be defined as the element with smallest cardinal in the solution set argmax S 1 ,
S 2 mTE W ( X S 1 , X S 2 ) .
Note the constructed mTE function requires strictly stronger dependence and includes no false positives, as shown by the proposition below (Proof can be seen at Appendix A.1):
Proposition 2.2 .
∀ S 1 , S 2 ⊂ { 1 , ..., p
[Equation]
, mTE
[Equation] [Equation]
2 ) >
[Equation]
TE W (
[Equation]
1 , X S 2 ) > 0 .
The purpose of mTE construction is to enable neural optimization as discussed in the next section.
Under some mild assumptions, we are able to provide theoretical justification for mTE maximization for time-series setting.
The complete proof can be seen in Appendix A.3.
Assumptions:
A1 Causal Markov assumption for the causal graph.
A2 Faithfulness for the causal graph.
A3 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).
A5 DAG causal graph: We assume X =
[ s 1 , ..., s l , t 1 , ..., t m ] , where s k are nuisance features that may correlate with t i , and t i are causally interacting signals and form a connected directed acyclic graph (DAG).
An illustration corresponding to the assumption can be seen in Figure 1.
[Equation]
A6 Interaction regularity: Given two interacting features A, B ∈ X, i ∈ { 0 , 1 } with A ⊂ Pa ( W i B ) − B and any other feature set X c , the mutual information difference is larger than zero:
[Equation]
Remark 2.3 .
Here our only additional assumption from prevalent literatures is A6, which aims to filter out features with spurious causations and regularize the algorithmic complexity of causal interactions, thus enabling information-theoretic analysis.
A6 has direct connections with the concept of conditional transfer entropy, which aims to filter out indirect causal relationships via conditioning (Faes et al., 2016; Shahsavari Baboukani et al., 2020).
Further discussions can be seen at Appendix A.2.
Theorem 2.4 .
Given assumption A1-A6, X ∗ ˆ S 1 ∪ ˆ S 2 ⊂ ( t 1 , ..., t m ) .
Moreover, each feature in ˆ S 1 ∪ ˆ Sis connected to other features in the set.
N EURAL OPTIMIZATION OF MULTI - DIMENSIONAL TRANSFER ENTROPY
With theorem 3.1 stated below, we are able to give a theoretical gaurantee of the neural optimization of mTE.
The proof can be seen at Appendix A.4.
Theorem 3.1.
Assume assumptions A1-A6 holds.
Then for small enough λ > 0 : [Equation] S
[Equation]
We have 1.
the optimal h is a one-to-one mapping; 2.
Assume f, g define one-to-one mappings on
X ⊙ 1 S 1 , X ⊙ 1 S 2 respectively.
Then, X ∗ S 1 ∪ S 2 ⊂ ( t 1 , ..., t m ) .
Moreover, each feature in S 1 ∪ S 2 is connected to other features in the set.
Remark 3.2.
The estimation of mutual information by neural networks is an active field itself (Belghazi et al., 2018; Hjelm et al., 2018; McAllester and Stratos, 2020;
Zhang et al., 2019).
In contrast, by this theorem, we show that an accurate estimation of the transfer entropy (such as 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.
Compared with mTE, the original transfer entropy is less robust to selections of mappings f, g, h .
Intuitively, we can construct non-one-to-one g s such that I ( f, Wg | g ) =
I ( f, Wg, g ) − I ( f, g ) > 0 strictly for any features.
This is further illustrated by the empirical performance comparison in the next section.
Given Theorem 3.1, we are able to construct a 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 as:
[Equation]
X ⊙ 1 S 1 ); g ( X ⊙ 1 S 2 )))
+ λ | S 1 ∪ S 2 |
Here f, g, h ∼ Gaussian indicate flows transforming the original feature distributions into Gaussian distributions.
We are able to formulate the target for neural network optimization by the explicit formula for mutual information between Gaussians: I ( X, Y )
= 12 log detΣ X detΣ Y detΣ
[ X,Y ] .
The Gaussian regularization can be applied either by regularizing over the discrepancy between embedding distributions and Gaussian distributions or by applying a adversarial training procedure.
In this work, we use the former approach, constructing mean and variances for each embedding component as learnable parameters and minimize the cross entropy between target distributions and the parametrized
Gaussian distributions.
3. C OMBINATORIAL STOCHASTIC GATES
In order to solve the optimization problem, we need to learn two sparse sets S 1 , S 2 .
Here we use a stochastic gate based approach (Yamada et al., 2020; Lindenbaum et al., 2021).
In order to explicitly construct S 1 and S 2 by stochastic gates, we define two random vectors T 1 and T 2 ranging in [0 , 1] with lengths equal to the feature number, with each element independently sampled from STG distribution defined as: T id = max(0 , min(1 , µ id + ϵ id )) , where ϵ id ∼ N (0 , σ 2 ) with fixed variance and µ id is a parameter trainable by the reparametrization trick (Miller et al., 2017; Figurnov et al., 2018).
In this case, denote 1 A as the indicator function, the loss function can be reformulated as:
[Equation]
In our construction, T 1 controls the sparsity of feature selection, while T 2 corresponds to overlapping level between S 1 and S 2 .
Defining the Gaussian error function as erf() , the regularization term for the first layer is of form:
[Equation]
The regularization term for the second layer can be expressed as:
[Equation]
We are able to show strong consistency for our stochastic-gate based feature selection as the theorem below (Proof can be seen at Appendix A.6):
Theorem 3.4.
Given assumptions A1-A6 and assumption on f, g described in theorem 3.1, denote for the optimal solution of (*), a sample of stochastic gate is given by T 1 , T 2 and ground truth active gates as T S ∗ 1 , T S ∗ 2 , then we have for small enough λ 1 > 0 and λ 2 > 0 ,
[Equation]
In practice, we also have observed the method’s solution highly depends on the stochastic gate initialization.
Here we provide an heuristic initialization scheme that shows superior empirical performance.
Details of the initialization scheme can be seen in Appendix B.
3. P ROPOSED NETWORK ARCHITECTURE
Our proposed network architecture is summarized in Figure 2. For an input dataset X ∈ R n × p and its corresponding graph adjacency matrix W ∈ R n × n , we first pass each feature through two sequential stochastic gate layers (STG layer 1: w , STG layer 2: z ). L0 penalty is conducted on the first STG layer, while the second STG layer is regularized such that large µ d in magnitude is favored (0-1 loss). After passing each feature, denote ˆ z i = 1 − z
i , we have two intermediate embeddings defined by x i ⊙ w
i ⊙ z
i and x i ⊙ w
i ⊙ ˆ
z i respectively.
Then these two embeddings are passed through two
MLPs ( f, g in theorem 3.1) to generate two output embeddings A,
B. Finally we apply the graph adjacency matrix W and another MLP layer (MLP3, h ) on B to generate embedding B*.
We estimate the mTE via the mutual information formula of joint Gaussian distributions and add the cross-entropy penalty between the concatenated (A,B,B*) and a Gaussian distribution with learnable parameters.
𝑥 ! 𝑥 " 𝑥 #$! 𝑥 # ⋮ 𝑤 !
𝑤 " 𝑤 #$!
𝑤 # 𝑧 !
𝑧 " 𝑧 #$!
𝑧 #
Input features 1 - 𝑧 !
1 − 𝑧 " 1 − 𝑧 #$! 1 − 𝑧 #
STG layer 1 STG layer⊗ ⊗ ⊗ ⊗ ⊗ ⊗ ⊗ ⊗ 𝑥 !
𝑤 !
𝑧 !
𝑥 " 𝑤 " 𝑧 " 𝑥 #$!
𝑤 #$!
𝑧 #$!
𝑥 # 𝑤 # 𝑧 # ⋮ 𝑥 !
𝑤 !
̂𝑧 !
𝑥 " 𝑤 " ̂𝑧 " 𝑥 #$!
𝑤 #$!
̂𝑧 #$!
𝑥 # 𝑤 # ̂𝑧 # ⋮
Feature set A
Feature set B
MLP
M L P𝐴 !
𝐴 " 𝐴 # … 𝐵 !
𝐵 " 𝐵 # …
Signal A
Signal B 𝐵 !
∗ 𝐵 "∗ 𝐵 #∗ …
Signal B* Apply W, MLP
MI(A,B) - MI(A,B*) + penalties L0 penalty 0/1penalty
3. O UTPUT INTERPRETATION
Upon the algorithm convergence, 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.
By the construction of GEASS, we are able to get two separate sparse feature subsets as source features and sink features.
These features may be used as inputs to further proper causal analysis, such as LPCMCI (Gerhardus and Runge, 2020), which despite its statistical power in depicting possible lags, identifying latent confounders, and allowing nonlinear tests, can only work on data with intermediate feature size.
Also, these feature themselves may be used in downstream unsupervised analysis to improve generalizations of prediction models.
E XPERIMENTS 4. G AUSSIAN TIME - SERIES WITH POSSIBLE NONLINEARITY
In order to benchmark the method in time-series data, we consider two settings: 1.
Minor effect of latent processes, with autocorrelation present; 2.
Significant effect of latent processes, with autocorrelation present.
Both settings are modelled by Gaussian structural processes with an underlying causal graph.
Further details can be seen in Appendix C.1.
We test the false discovery rate (FDR) and F1 score between ground truth interacting features and recovered features as two metrics for high-dimensional data causal discovery.
We compare GEASS 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ˇcs and Vogt, 2021), and Grid-net which generalizes the definition of Granger causality to Directed Acyclic Graph (DAG)
(Wu et al., 2021); moreover we include two state-of-the-art approaches, DCM and NGM implemented in (Bellot et al., 2021) that use neural ODE to model nonlinear dependence graph.
Table 1 shows our benchmarking result.
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 correlation.
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 are more irregular.
Finally, GEASS consistently gives accurate causal feature identifications (high F1) and low false discovery rate (low FDR) in all settings considered.
[Table]
Furthermore, we evaluate different methods’ scalability with respect to the feature size.
(Experimental details can be seen at Appendix C.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 of these methods for high-dimensional biological data analysis, where the feature number is at the scale of 10 3 − 10 4 .
Meanwhile,
GEASS constructs a single neural network with parameters approximately proportional to p , thus largely reducing the complexity in the high-dimensional regime.
We benchmark PCMCI,
GVAR, GrID-net, NGM, GEASS, and an additional combination 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).
4. S IMULATED 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 autocorrelation is modeled by higher likelihood of same type of cells in the neighborhood, and confounder is modelled by coherent shift of global gene expression for each cell type.
We first simulate scRNA-seq data, then the synthetic scRNA-seq data are assigned to a 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 the dataset is normalized and log1p transformed as the standard pipeline in Scanpy (Wolf et al., 2018).
The majority of methods are not available as their focus is on time-series data.
Therefore in order to perform our benchmarking study, we compare GEASS with Lasso Granger, as well as our implemented L1-regularized version of NCEM, an approach proposed to detect interactions in spatial omics data (Fischer et al., 2021).
Finally, we also implemented a method that maximizes over original transfer entropy to select causal features (TE).
As shown in table 2, original LASSO cannot identify causal features because of strong correlation between features.
L1-NCEM ameliorates the issue by conditioning on the cell type label in regression.
TE outperforms linear methods yet generate a number of false positives, as it may overfit to learn spurious causations as discussed in remark 3.3.
Finally, GEASS consistently outperforms over other methods in identifying causal features of data as shown by both high F1 score and low FDR.
[Table]
4. SC RNA- SEQ PANCREATIC ENDOCRINOGENESIS TRAJECTORY
We test GEASS on the pancreatic endocrinogenesis trajectory data, which is a standard dataset for scRNA-seq trajectory inference task (Bergen et al., 2020; Bastidas-Ponce et al., 2019).
The pancreas trajectory data contains 3696 cells and 27998 genes.
After preprocessing, lowly-expressed genes are filtered out as the standard pipeline in scVelo (Bergen et al., 2020), with remaining 2000 genes for further analysis.
We aim to use GEASS to identify causally-related genes along the developmental trajectory to reveal underlying biology.
(See Appendix C.3 for experimental details).
scRNA-seq data provides a 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 of intron RNAs to infer the underlying dynamics of gene expression change.
Thus, we are able to define a velocity kernel matrix W velo , which provides weighted adjacency relationships of cells based on velocity direction and cell phenotypal proximity.
GEASS identifies 50 causally-related features with high biological relevance.
For example, the gene list includes the key transcriptional regulator Neurog3 , which is required for the specification of a common precursor of the 4 pancreatic terminal states (uni, 2021).
For further quantitative validation, we take advantage of the additional information given by RNA velocity, as significant RNA velocity is a necessary condition for causal dynamics of successor genes along the trajectory.
Therefore, we benchmark different gene selection methods’ performance on selecting dynamical genes.
Specifically, we evaluate the mean velocity likelihood over GEASS causal genes with high-expressed genes (HEG), highly-variable genes (HVG), and genes having high correlation with inferred latent time (HCG).
Our result suggests here GEASS achieves the best performance in selecting dynamical genes.
Mean RNA velocity likelihood
HEG 0.
HVG 0.
HCG 0.
GEASS 0.
4. MERFISH HUMAN CORTEX SINGLE - CELL LEVEL SPATIAL TRANSCRIPTOMICS
Spatial transcriptomics represent a 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 et al., 2022; Jerby-Arnon and Regev, 2022; Fischer et al., 2021).
However, current computational methods revealing interaction modules (Jerby-Arnon and Regev, 2022) or niche effects (Fischer et al., 2021) 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.
Here we use GEASS on a recent published MERFISH data measuring spatially-resolved gene expression of human cortex (Fang et al., 2022).
MERFISH is a multiplexed single-molecule imagingbased technology that enables measurements of spatially-resolved single-cell gene expression in tissues (Chen et al., 2015; Fang et al., 2022).
The data comprises of 3044 cells and 4000 genes; each cell is annotated as 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 (MURAL) as shown by the first panel of figure 6 in Appendix D. We selected 1000 highly variable genes by Scanpy preprocessing pipeline.
After preprocessing, our GEASS analysis selects 9 genes, namely FILIP1, SLC17A7, MYH11, RP1110j21.2, PIRT, C3orf67, TRDMT1, RGS8, SPTLC2 .
An overview of the gene expression is available in the rest panels of figure 6.
Among these genes, Myh11 , RP11-10j21.2 , and TRDMT1 are enriched at the spatial regions with neighboring endothelial and mural cells, corresponding to underlying vascular structures (marked by ellipses in the first panel of figure 6).
The finding is consistent with the MERFISH images, which reveals the spatial selective proximity of oligodendrocyte and microglial cell bodies to vascular structures, while astrocytes contact blood vessels with their end feet (Fang et al., 2022).
Therefore, these identified marker genes of vascular structure encode underlying meaningful cellular interaction, supporting the validity of GEASS.
In order to show GEASS’s ability to facilitate biological findings, we focus on two GEASS identified genes, C3orf67 and PIRT .
These two genes are enriched at the same spatial locations, both expressing prominently in astrocytes.
Further statistical analysis reveals that the model of PIRT affecting neighboring C3orf67 is favored.
This is consistent with the predicted role of PIRT in enabling transmembrane transporter binding activity and being involved in phosphatidylinositol-mediated signaling (Safran et al., 2021), and previous works highlighting the gene’s role in pain sensing (Kim et al., 2008; Wang et al., 2021).
As the role of C3orf67 in human cortex remains unclear, this revealed causal link may lead to further discoveries.
Pearson r
P value (C3orf67, PIRT) 0.
0.014* (C3orf67, PIRT nb ) 0.
0.049* (PIRT, C3orf67 nb ) 0.
C ONCLUSIONS
In this work, we present GEASS, a causal feature selection method based on information-theoretic tools and neural networks.
GEASS is able to scale to high dimensions and identify sparse interacting features.
We provide both theoretical gaurantee and empirical validation of GEASS on synthetic and real biological data.
Our results show GEASS can be integrated into high-dimensional spatiotemporal data analysis pipeline 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 a latent embedding corresponding to different causal mechanisms.
Therefore, in applications where a causal graph output is favored, constraint-based methods may need to be applied after GEASS.
Moreover, when underlying causal graph has a 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 wide use of GEASS in high-dimensional graph-structured data, especially for high-dimensional biological data such as 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 discovery.
A PPENDIX A P ROOFS A. P ROOF OF P ROPOSITION 2.2.
Proposition 2.2 .
∀ S 1 , S 2 ⊂ { 1 , ..., p } , mTE W
( X S 1 , X S 2 ) > 0 ⇒ TE W
( X S 1 , X S 2 ) > 0 .
Proof.
We have
TE W ( X S 1 , X S 2 )
= I ( X S 1 ; WX S 2 | X S 2 ) = I ( X S 1 ; WX S 2 , X S 2 ) − I ( X S 1 ; X S 2 ) = I ( X S 1 ; WX S 2 ) − I ( X S 1 ; X S 2 ) + I ( X S 1 ; X S 2 | WX S 2 ) .
Therefore we have TE W ( X S 1 , X S 2 ) ≥ mTE W ( X S 1 , X S 2 ) holds, therefore mTE W
( X S 1 , X S 2 ) > 0 ⇒ TE
W ( X S 1 , X S 2 ) > 0 .
A
D ISCUSSION 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 is a simple example to see why TE/mTE can have problems with indirect causal interactions: consider the relationship: s t → w t → v t +1 , then we would have: I ( s t , v t +1 ) > 0 although there are no direct causal relationship between s and v . The issue can be resolved by considering the conditional mutual information I ( s t , v t +1 | w t ) , which equals 0. This inspires the concept of conditional transfer entropy:
Definition (Conditional transfer entropy) (Shahsavari Baboukani et al., 2020).
Assume X and Y are the features of interest and the conditioning features are Z . Denote − as [1 , 2 , ..., t ] , then we have cTE Z t ( X, Y ) = [Equation]
( Y t +1 , X − | Y − , Z − ) . The classical formulation of conditional transfer entropy is widely used in high-dimensional observational data to learn direct causal dependencies (Faes et al., 2016; Shahsavari Baboukani et al., 2020).
It implicitly assumes that, there is direct causal relationship between X and Y if ∀ Z, cTE Z t ( X, Y ) > 0 .
Here, we are further assuming that this also holds for cmTE, which we define as an analogy to cTE in the Markovian model setting:
Definition (Conditional mTE) .
Assume X and Y are the feature sets of interest and the conditioning features are Z .
Then we have
[Equation]
We can see not only our definition is similar to cTE in form, the key idea of the conditional transfer entropy is still preserved: If X j causally affects Y j +1 through Z t , or dependency between X j and Y j +1 is because of spurious correlation to Z t , then we have cmTE Z t ( X, Y, j ) ≤ 0 . Meanwhile, the assumption A6 also implies strong interaction as discussed by proposition 2.2.
In summary, our A can be reformulated as ∀ Z t , j, cmTE Z t ( X, Y, j ) > 0 for ground truth interacting X, Y and can be similarly defined for nuisance features.
Theorem 2.4 .
Given assumption A1-A6, X ∗ ˆ S 1 ∪ ˆ S 2 ⊂ ( t 1 , ..., t m ) . Moreover, each feature in ˆ S 1 ∪ ˆ Sis connected to other features in the set.
Proof.
Step 1.
First we prove S 1 ∩ S 2 = ∅ .
If not, assume the overlapping signal to be p , then for arbitrary S 1 , S 2 , assume the corresponding sources are A and B respectively.
Then we have
[Equation]
Therefore removing p would increase the value of mTE, leading to a contradiction.
Step 2.
Now we prove nuisance signals cannot be in either ˆ S 1 or ˆ S 2 .
Otherwise, first we assume S is in ˆ S 2 , as S only interacts with variables at the same time point, S can only interact with WA via indirect links through a subset of interacting features. Assume this feature set as Pa ( WA ) S , then with A6, we have [Equation]
Moreover, we have I
[Equation]
= − I ( WA ; S | B, Pa ( WA ) S − B ) + I ( A ; S | B, Pa ( WA ) S − B )
{+= 
+}
+ I ( A ; S | B, Pa ( WA ) S − B ) ≥ 0 . Therefore we have mTE ( A ; B, Pa ( WA ) S − B ) ≥ mTE ( A ; S, B, Pa ( WA ) S − B ) >
mTE
( A ; S, B ) , thus leading to a contradiction. Then we are able to assume S is in ˆ S 1 , and B does not include any nuisance features.
Then as S only interacts with variables at the same time point, WS can only interact with B via indirect links through a subset of interacting features.
Assume the whole intermediate feature set for B as Pa ( WS ) B , then we first note that we have Pa ( WS ) B ∩ WS = ∅ .
then with A6, we have
I ( WA, WS, Pa ( WS ) B − WS ; B ) − I ( A, S, Pa ( S ) B − S ; B ) − I ( WA, WS ; B ) + I
( A, S ; B )
=
I ( Pa ( WS ) B − WS ; B | WA, WS ) − I ( Pa ( WS ) B − WS ; WB | WA, WS ) > 0 .
Moreover, we have I ( WA, Pa ( WS ) B − WS ; B ) − I ( A, Pa ( S ) B − S ; B ) − I ( WA, WS, Pa ( WS ) B − WS ; B ) + I
( A, S, Pa ( S ) B − S ; B ) = − I ( WS ; B | WA, Pa ( WS ) B − WS ) + I ( S ; B | A, Pa ( S ) B − S )
{+= 
+}
+ I ( S ; B | A, Pa ( S ) B − S ) ≥ 0 . Therefore we have mTE ( A, Pa ( S ) B − S ; B ) ≥ mTE
( A, Pa ( S ) B − S, S ; B ) > mTE ( A, S ; B ) , thus leading to a contradiction.
Step 3.
Moreover, if there exists a component in ˆ S 1 ∪ ˆ S 2 not connected to any other feature components in the set, assume the feature to be q . Then, in this case with A1-4, the feature q is independent of any other features in ˆ S 1 ∪ ˆ S 2 . From step 1 it can be deduced that q cannot be in both ˆ
S 1 , ˆ S 2 . Therefore in this case, we have mTE( ˆ S 1 − q, ˆ S 2 − q ) = mTE( ˆ S 1 , ˆ S 2 ) thus leading to the contradiction of finding an ( ˆ S 1 , ˆ S 2 ) with smaller | ˆ S 1 ∪ ˆ S 2 | .
Theorem 3.1.
Assume assumptions A1-A6 holds.
Then for small enough λ > 0 :
[Equation]
We have 1.
the optimal h is a one-to-one mapping; 2.
Assume f, g define one-to-one mappings on
X ⊙ 1 S 1 , X ⊙ 1 S 2 respectively.
Then, X ∗ S 1 ∪ S 2 ⊂ ( t 1 , ..., t m ) .
Moreover, each feature in S 1 ∪ S 2 is connected to other features in the set. [Equation]
Proof.
Step 1.
function h must be injective in the optimal solution as non-one-to-one mapping reduces mutual information (Standard property of mutual information).
Step 2.
Given the assumption that f, g define injective mappings on 1 S 1 , 1 S 2 respectively, and G and W are commutable as G ( WX t ) = G ( X t +1 ) = WG ( X t ) , and one-to-one transformation does not change mutual information, we have
[Equation]
Therefore using theorem 2.4, in the case of DAG causal graph, we have optimal f ( X ) and g ( X ) define one-to-one mappings to non-overlapping subsets of { t 1 , ..., t m } . Step 3. With our definition of optimal S 1 , S 2 , the minimal gap between the optimal solution and any other solution with smaller | S 1 ∪ S 2 | size is larger than zero.
Denote the minimal gap as δ , then for these other solutions with smaller | S 1 ∪ S 2 | , we have [Equation] Meanwhile, for solutions with larger | S 1 ∪ S 2 | size, we have min [Equation]
[Equation]
Meanwhile, for the desired optimal S 1 , S 2 , we have min [Equation] g, [Equation]
Therefore, when taking λ ∈ (0 , δ | S 1 ∪ S 2 | ) , the desired optimal S 1 , S 2 by mTE is the optimal output of the constructed optimization problem.
Theorem 3.4. Given assumptions A1-A6 and assumption on f, g described in theorem 3.1, denote for the optimal solution of (*), a sample of stochastic gate is given by T 1 , T 2 and ground truth active gates as T S ∗ 1 , T S ∗ 2 , then we have for small enough λ 1 > 0 and λ 2 > 0 ,
[Equation]
Proof.
Step 1.
Given f, g, h projects input distributions into Gaussian distributions, by convergence of Gaussian covariance matrices, we have:
[Equation]
Step 2. Upon step 1, applying Glivenko–Cantelli theorem on the stochastic gate terms, we have (*) equals
[Equation]
Therefore, as N → ∞ , applying step 1 of theorem 2.4, for any λ 2 > 0 , we have ∀ d, P ( T 2 d = 0) = 0 / 1 , P ( T 2 d = 0) = 1 / 0 holds, in which case the penalty for T 2 is zero.
Applying above, noting the stochastics are i.i.d sampled, independent with X , we have
[Equation]
With the non-overlapness of T 1 ⊙ T 2 and T 1 ⊙ (1 − T 2 ) , the second and third term above equals zero.
Therefore we have
I ( [Equation] X
[Equation] Use the same steps, we eventually have [Equation] Also for the first term of mTE we similar have
I ( [Equation] [Equation] Then we eventually have the optimization problem as
[Equation] T 2 ⊙ T 1 > 0) , ((1 − T 2 ) ⊙ T 1 > 0)) + I ( X ⊙ ( T 1 ⊙ T 2 > 0); X ⊙ ( T 1 ⊙ (1 − T 2 ) > 0) | ( T 2 ⊙ T 1 > 0) , ((1 − T 2 ) ⊙ T 1 > 0))+ λ 1 || T 1 ||
From theorem 3.1, with λ 1 small enough, we have above ≥ min T 1 ,T 2 − I ( X ⊙ ( T 1 ⊙ T 2 > 0); WX ⊙ ( T 1 ⊙ (1 − T 2 ) > 0) | ( T 2 ⊙ T 1 > 0) , ((1 − T 2 ) ⊙ T 1 > 0)) + I ( X ⊙ ( T 1 ⊙ T 2 > 0);
X ⊙ ( T 1 ⊙ (1 − T 2 ) > 0) | ( T 2 ⊙ T 1 > 0) , ((1 − T 2 ) ⊙ T 1 > 0))+ λ 1 || T 1 ||
[Equation]
, which can be achieved when P (( T 1 d ⊙ T 2 d > 0) = 1 / 0 when d ∈ / / ∈ S 1 , and P (( T 1 d ⊙ (1 − T 2 d ) > 0) = 1 / 0 when d ∈ / / ∈ S 2 . Note above holds when N → ∞ with probability 1, where we have
[Equation]
Therefore we have proven strong consistency of the stochastic-gate based estimator.
B G ATE INITIALIZATION
Our proposed initialization scheme is based on analysis of the linear case.
Assume f
( X S 1 ) =
Xa, g ( X S 2 ) =
Xb, where a, b ∈ R p represents two feature loadings.
Then:
a, b should be non-overlapping, therefore we expect | a T b | to be small.
We should have f ( X ) ≈
Wg ( X ) to maximize the mTE.
Based on the two criteria above, we formulate an optimization problem as
[Equation] [Equation]
The constraint can be formulated into a regression problem WXb =
Xa , therefore a natural solution is given by a = X † WXb =
( X T X ) − 1 X T WXb.
In this case, || a T b || =
|| b T ( X T X ) − 1 X T WXb ||
= ||
b
|| 2( X T X ) − 1 X T WX .
Given b is normalized, it can be shown that the optimal b corresponds to the eigenvector with least absolute eigenvalue of matrix ( X T X ) − 1 X T WX .
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.
C E XPERIMENTAL DETAILS C. T IME - 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 , 2 , 3 , 4 , 5 , 6 } .
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 a 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 are nuisance features with white noise dynamics.
The forward operator is defined by 5-neighbor lower triangular matrix.
C.1. A LGORITHM IMPLEMENTATION
If not particularly mentioned, default settings of the algorithms are used throughout.
• VAR-LINGAM.
The VAR-LINGAM algorithm is implemented in the Python package LINGAM , available at https://github.com/cdt15/lingam.
VAR-LINGAM gives an weighted matrix as output.
Therefore in our benchmarking study, we choose the most significant edge corresponding features with the number matching the sparsity level.
• PCMCI.
The PCMCI algorithm is implemented in the Python package Tigramite , which gives an weighted matrix as output.
We choose the most significant edge corresponding features with the number matching the sparsity level.
• GVAR.
The GVAR algorithm is implemented at https://github.com/i6092467/GVAR.
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 on the whole time series, giving a binary output and no threshold selection is needed.
• Grid-net.
The Grid-net algorithm is implemented at https://github.com/alexw16/gridnet.
The parameter set: order=5, hidden_layer_size = 10, end_epoch=50, batch_size = 50, lmbd=1 is used throughout our study.
After the training finishes, we choose the most significant edge corresponding features with the number matching the sparsity level.
• DCM, NGM.
The two algorithms are both implemented at https://github.com/alexisbellot/ Graphical-modelling-continuous-time.
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 λ 1 set with 0.04/0.05 based on a validation set, and the rest parameter settings are consistent with default.
C.1. S CALABILITY ANALYSIS
We test PCMCI, GVAR, GrID-net, NGM, GEASS, GEASS+LPCMCI’s running time with consistent settings described in the above section.
(LPCMCI’s setting is consistent with PCMCI’s setting).
We use the same data generation pipeline and select the set of the total feature numbers as [100 , 200 , 400 , 800 , 1600] .
C. S IMULATED 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-specificly expressed.
The restgenes 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.
C.2. A LGORITHM IMPLEMENTATION
If not 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 a linear graph neural network, which in the grid case corresponds to a 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 α = 0 .to match the sparsity level.
• GEASS.
We use the same training parameters in all settings, with the key sparsity regularization parameter λ 1 set 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 parameters in all time-series settings, with the key sparsity regularization parameter λ 1 set with 0.05 based on a validation set, and the latent dimension number is set to be 64 consistent with GEASS.
C. SC RNA- SEQ PANCREAS TRAJECTORY
The data preprocessing is consistent with the scVelo tutorial: https://scvelo.readthedocs.io/ VelocityBasics/. (Bergen et al., 2020)
The parameter set: λ 1 = 0 .
06 , λ 2 = 0 .
Here because the gene regulatory network is fully connected and activated in cascade along the developmental trajectory, we consider the opposite initialization with b be the largest eigenvalues corresponding eigenvectors of the matrix ( X T X ) − 1 X T WX .
C. M ERFISH SPATIAL TRANSCRIPTOMICS DATA
The data is downloaded from Dryad and preprocessed with the standard Scanpy pipeline.
(Wolf et al., 2018)
The forward operator is defined by 5-neighbor adjacency matrix.
The GEASS parameter set is consistent with those used in the spatial omics benchmarking.
D A DDITIONAL EXPERIMENTAL RESULTS