P REDICTING C ELLULAR R ESPONSES WITH V ARIATIONAL C AUSAL I NFERENCE AND R EFINED R ELATIONAL I NFORMATION
Layne
Price
Amazon prilayne@amazon.com
F
Voloch
Immunai luis
Amazon gkarypis@
A BSTRACT
Predicting the responses of a cell under perturbations may bring important benefits to drug discovery and personalized therapeutics.
In this work, we propose a novel graph variational Bayesian causal inference framework to predict a cell’s gene expressions under counterfactual perturbations (perturbations that this cell did not factually receive), leveraging information representing biological knowledge in the form of gene regulatory networks (GRNs) to aid individualized cellular response predictions.
Aiming at a data-adaptive GRN, we also developed an adjacency matrix updating technique for graph convolutional networks and used it to refine GRNs during pre-training, which generated more insights on gene relations and enhanced model performance.
Additionally, we propose a robust estimator within our framework for the asymptotically efficient estimation of marginal perturbation effect, which is yet to be carried out in previous works.
With extensive experiments, we exhibited the advantage of our approach over state-of-the-art deep learning models for individual response prediction.
I NTRODUCTION
Studying a cell’s response to genetic, chemical, and physical perturbations is fundamental in understanding various biological processes and can lead to important applications such as drug discovery and personalized therapies.
Cells respond to exogenous perturbations at different levels, including epigenetic (DNA methylation and histone modifications), transcriptional (RNA expression), translational (protein expression), and post-translational (chemical modifications on proteins).
The availability of single-cell RNA sequencing (scRNA-seq) datasets has led to the development of several methods for predicting single-cell transcriptional responses (Ji et al., 2021).
These methods fall into two broad categories.
The first category (Lotfollahi et al., 2019; 2020; Ramp´aˇsek et al., 2019; Russkikh et al., 2020; Lotfollahi et al., 2021a) approaches the problem of predicting single cell gene expression response without explicitly modeling the gene regulatory network (GRN), which is widely hypothesized to be the structural causal model governing transcriptional responses of cells (Emmert-Streib et al., 2014).
Notably among those studies, CPA (Lotfollahi et al., 2021a) uses an adversarial autoencoder framework designed to decompose the cellular gene expression response to latent components representing perturbations and covariates.
CPA extends the classic idea of decomposing high-dimensional gene expression response into perturbation vectors (Clark et al., 2014; 2015), which can be used for finding connections among perturbations (Subramanian et al., 2017).
However, while CPA’s adversarial approach encourages latent independence, it does not have any supervision on the counterfactual outcome construction and thus does not explicitly imply that the counterfactual outcomes would resemble the observed response distribution.
Existing self-supervised counterfactual construction frameworks such as GANITE (Yoon et al., 2018) also suffer from this problem.
The second class of methods explicitly models the regulatory structure to leverage the wealth of the regulatory relationships among genes contained in the GRNs (Kamimoto et al., 2020).
By bringing the benefits of deep learning to graph data, graph neural networks (GNNs) offer a versatile and powerful framework to learn from complex graph data (Bronstein et al., 2017).
GNNs are the de facto way of including relational information in many health-science applications including molecule/protein property prediction (Guo et al., 2022; Ioannidis et al., 2019; Strokach et al., 2020; Wu et al., 2022a; Wang et al., 2022), perturbation prediction (Roohani et al., 2022) and RNAsequence analysis (Wang et al., 2021).
In previous work, Cao & Gao (2022) developed GLUE, a framework leveraging a fine-grained GRN with nodes corresponding to features in multi-omics datasets to improve multimodal data integration and response prediction.
GEARS (Roohani et al., 2022) uses GNNs to model the relationships among observed and perturbed genes to predict cellular response.
These studies demonstrated that relation graphs are informative for predicting cellular responses.
However, GLUE does not handle perturbation response prediction, and GEARS’s approach to randomly map subjects from the control group to subjects in the treatment group is not designed for response prediction at an individual level (it cannot account for heterogeneity of cell states).
GRNs can be derived from high-throughput experimental methods mapping chromosome occupancy of transcription factors, such as chromatin immunoprecepitation sequencing (ChIP-seq), and assay for transposase-accessible chromatin using sequencing (ATAC-seq).
However, GRNs from these approaches are prone to false positives due to experimental inaccuracies and the fact that transcription factor occupancy does not necessarily translate to regulatory relationships (Spitz & Furlong, 2012).
Alternatively, GRNs can be inferred from gene expression data such as RNA-seq (Maetschke et al., 2014).
It is well-accepted that integrating both ChIP-seq and RNA-seq data can produce more accurate GRNs (Mokry et al., 2012; Jiang & Mortazavi, 2018; Angelini & Costa, 2014).
In this work, we employed a novel variational Bayesian causal inference framework to construct the gene expressions of a cell under counterfactual perturbations by explicitly balancing individual features embedded in its factual outcome and marginal response distributions of its cell population.
We integrated a gene relation graph into this framework, derived the corresponding variational lower bound and designed an innovative model architecture to rigorously incorporate relational information from GRNs in model optimization.
Additionally, we propose an adjacency matrix updating technique for graph convolutional networks (GCNs) in order to impute and refine the initial relation graph generated by ATAC-seq prior to training the framework.
With this technique, we obtained updated GRNs that discovered more relevant gene relations (and discarded insignificant gene relations in this context) and enhanced model performance.
Besides, we propose an asymptotically efficient estimator for estimating the average effect of perturbations under a given cell type within our framework.
Such marginal inference is of great biological interest because scRNA-seq experimental results are typically averaged over many cells, yet robust estimations have not been carried out in previous works on predicting cellular responses.
We tested our framework on three benchmark datasets from Srivatsan et al.
(2020), Schmidt et al.
(2022) and a novel CROP-seq genetic knockout screen that we release with this paper.
Our model achieved state-of-the-art results on out-of-distribution predictions on differentially-expressed genes — a task commonly used in previous works on perturbation predictions.
In addition, we carried out ablation studies to demonstrate the advantage of using refined relational information and using robust marginal estimators for a better understanding of the contributions of framework components.
P ROPOSED M ETHOD
In this section we describe our proposed method.
A list of all notations can be found in Appendix D.
2. C OUNTERFACTUAL C ONSTRUCTION F RAMEWORK
We define outcome Y : Ω → R n to be a n -dimensional random vector (e.g. gene expressions),
X : Ω → E X to be a m -dimensional mix of categorical and real-valued covariates (e.g. cell types,
[Figure]
donors, etc.)
, T : Ω → E T to be a r -dimensional categorical or real-valued treatment (e.g. drug perturbation) on a probability space (Ω , Σ , P ) .
We seek to construct an individual’s counterfactual outcome under counterfactual treatment a ∈ E T from two major sources of information.
One source is the individual features embedded in high-dimensional outcome Y .
The other source is the response distribution of similar subjects (subjects that have the same covariates as this individual) that indeed received treatment a .
We employ the variational causal inference (Wu et al., 2022b) framework to combine these two sources of information.
In this framework, a covariate-dependent feature vector Z : Ω → R d dictates the outcome distribution along with treatment T ; counterfactuals Y ′ and T ′ are formulated as separate variables apart from Y and T with a conditional outcome distribution p ( Y ′ | Z, T ′ =
a ) identical to its factual counterpart p
( Y | Z, T = a ) on any treatment level a .
The learning objective is described as a combined likelihood of individual-specific treatment effect p
( Y ′ | Y, X, T, T ′ )
(first source) and the traditional covariate-specific treatment effect p
( Y | X, T ) (second source):
[Equation]
where D =
( X, Z, T, T ′ , Y, Y ′ ) .
Additionally, we assume that there is a graph structure G =
( V , E ) that governs the relations between the dimensions of Y through latent Z , where V ∈ R n × v is the node feature matrix and E ∈ { 0 , 1 } n × n is the node adjacency matrix.
For example, in the case of single-cell perturbation dataset where Y is the expression counts of genes, V is the gene feature matrix and E is the GRN that governs gene relations.
See Figure 1 for a visualization of the causal diagram.
The objective under this setting is thus formulated as
[Equation]
where D =
( G , D ) .
The counterfactual outcome Y ′ is always unobserved, but the following theorem provides us a roadmap for the stochastic optimization of this objective.
Theorem 1.
Suppose that D
=
( G , X, Z, T, T ′ , Y, Y ′ ) follows a causal structure defined by the
Bayesian network in Figure 1.
Then J ( D ) has the following evidence lower bound:
[Equation]
Proof of the theorem can be found in Appendix A.
We estimate p ( Z | Y, G , X, T ) and p ( Y | Z, T ) (as well as p ( Y ′
| Z, T ′ ) )
with a neural network encoder q ϕ and decoder p θ , and optimize the following weighted approximation of the variational lower bound:
[Equation]
where ω 1 , ω 2 are scaling coefficients;
˜ Y ′ θ,ϕ ∼ E q ϕ
( Z | Y, G ,X,T ) p
θ ( Y ′ | Z, T ′ )
and ˆ p is the covariatespecific model fit of the outcome distribution (notice that p ( Y ′ | G , X, T ′ = a )
=
p ( Y | G , X, T = a ) for any a ).
In our case where covariates are limited and discrete (cell types and donors), we simply let ˆ p ( Y | G , X, T ) be the smoothened empirical distribution of Y stratified by X and T (notice that G is fixed across subjects).
Generally, one can train a discriminator ˆ p ( ·| Y, G , X, T ) with the adversarial approach (Goodfellow et al., 2014) and use log (cid:104) ˆ p (1 | ˜
Y ′ θ,ϕ , G , X, T ′ )
(cid:105) for log (cid:104)
ˆ p ( ˜ Y ′ θ,ϕ | G , X, T ′ )
(cid:105) if p ( Y | G , X, T ) is hard to fit.
See Figure 2 for an overview of the model structure.
Note that the decoder estimates the conditional outcome distribution of Y ′ , in which case T ′ need not necessarily be sampled according to a certain true distribution p ( T ′ | X ) during optimization.
We refer to the negative of the first term in Equation 4 as reconstruction loss, the negative of the second term as distribution loss, and the positive of the third term as KL-divergence.
As discussed in Wu et al.
(2022b), the negative KL-divergence term in the objective function encourages the preservation of individuality in counterfactual outcome constructions.
2. I NCORPORATING R ELATIONAL I NFORMATION
Since the elements in the outcome Y are not independent, we aim to design a framework that can exploit predefined relational knowledge among them.
In this section, we demonstrate our model structure for encoding and aggregating relation graph G within the graphVCI framework.
Denote deterministic models as f · , probabilistic models (output probability distributions) as q · and p · .
We construct feature vector Z as an aggregation of two latent representations:
[Equation]
where aggr G is a node aggregation operation such as sum , max or mean .
The optimization of q ϕ can then be performed by optimizing the MLP encoder f M,ϕ 1 : R n + m + r → R d , the GNN encoder f G ,ϕ 2 : R n × v → R n × d G and the encoding aggregator q H,ϕ 3 : R d + d G → ([0 , 1] d ) Σ where ϕ
= ( ϕ 1 , ϕ 2 , ϕ 3 ) .
We designed such construction so that Z possesses a node-level graph embedding
Z G that enables more involved decoding techniques than generic MLP vector decoding, while the calculation of term KL [ q ϕ ( Z | Y, G , X, T ) ∥
q ϕ ( Z | Y ′ , G , X, T ′ )] can also be reduced to the KLdivergence between the conditional distributions of two graph-level vector representations Z H and
[Equation]
[Figure]
Z ′ H (since Z G is deterministic and invariant across subjects).
In decoding, we construct
[Equation]
where ∥ represents vector concatenation and optimize p θ by optimizing the MLP decoder p M,θ 1 : R d + r → ([0 , 1] d ) Σ and the decoding aggregator f H,θ 2 : R d + d G → R d where θ = ( θ 1 , θ 2 )
The decoding aggregator maps graph embedding Z G ( i ) of the i -th node along with feature vector Y M to outcome Y ( i ) of the i -th dimension, for which we use an attention mechanism and let Y ( i ) = att θ
Z G ( i ) , Y M ) ⊤ · Y M where att( Q ( i ) , K ) gives the softmax attention score for each of K ’s feature dimensions given a querying node embedding vector Q ( i ) (i.e., a row vector of Q ).
One can simply use a key-independent attention mechanism Y ( i ) =
att θ
Z G ( i ) )
⊤ · Y M if GPU memory is of concern.
See Figure 3 for a visualization of the encoding and decoding models.
A complexity analysis of the model can be found in Appendix E.
2. R ELATION G RAPH R EFINEMENT
Since GRNs are often highly context-specific (Oliver, 2000; Romero et al., 2012), and experimental methods such as ATAC-seq and Chip-seq are prone to false positives, we seek to impute and refine a prior GRN by learning from the expression dataset of interest.
We propose an approach to update the adjacency matrix in GCN training while maintaining sparse graph operations to achieve this goal.
We design g ( · ) to be a GCN with learnable and penalized edge weights between all nodes, and aim to acquire an updated gene graph by thresholding updated edge weights post optimization.
In practice, such GCNs with complete graphs performed poorly on our task and had scalability issues.
Hence we apply dropouts to edges in favor of an initial relation graph — edges present in ˜ E
= | E − I |
+ I
( I is the identity matrix) are accompanied with a low dropout rate r l and edges not present in ˜ E with a high dropout rate r h ( r l ≪
r h ).
A graph convolutional layer of g ( · ) is then given as
[Equation]
) where L ∈ R n × n is a dense matrix containing logits of the edge weights and M ∈ R n × n is a sparse mask matrix where each element M i,j is sampled from Bern ( I (˜ E i,j = 1) r l + I (˜ E i,j = 0) r h ) in every iteration; ⊙ is element-wise matrix multiplication; softmax r is row-wise softmax operation;
H l ∈ R n × d l , H l +1 ∈ R n × d l +1 are the latent representations of V after the l -th, ( l + 1) -th layer;
[Figure]
Θ l ∈ R d
l × d l +1 is the weight matrix of the ( l + 1) -th layer; σ is a non-linear function.
The updated adjacency matrix ˆ E ∈ R n × n is acquired by rescaling and thresholding the unnormalized weight matrix W = exp( L ) after optimizing g ( · ) :
[Equation]
12) where α is a threshold level.
We define ˜ W ∈ R n × n to be the rescaled weight matrix having ˜ W i,j = (1 + W − 1 i,j ) − 1 .
With this structure, the convolution layer operates on sparse graphs which benefit performance and scalability, while each absent edge in the initial graph still has an opportunity to come into existence in the updated graph.
We use this approach to obtain an updated GRN ˆ E prior to main model training presented in the previous sections.
We train g ( · ) : R n × (1+ v + m )
→ R n × 1 on a simple node-level prediction task where the output of the i -th node is the expression Y ( i ) of the i -th dimension; the input of the i -th node is a combination O ′ ( i ) =
( Y ( i ) , V ( i ) , X ) with expression Y ( i ) , gene features V ( i ) of the i -th gene and cell covariates X .
Essentially, we require the model to predict the expression of a gene (node) from its neighbors in the graph.
This task is an effective way to learn potential connections in a gene regulatory network as regulatory genes should be predictive of their targets (Kamimoto et al., 2020).
The objective is a lasso-style combination of reconstruction loss and edge penalty:
[Equation]
(13) where ω is a scaling coefficient.
Note that although V ( i ) is not cell-specific and X is not genespecific, the combination of ( V ( i ) , X ) forms a unique representation for each gene of each cell type.
We employ such dummy task since node-level predictions with strong self-connections grants interpretability, but additional penalization on the diagonal of ˜ W can be applied if one wishes to weaken the self-connections.
Otherwise, although self-connections are the strongest signals in this task, dropouts and data correlations will still prevent them from being the only notable signals.
See
Figure 4 for an example of an updated GRN in practice.
2. M ARGINAL E FFECT E STIMATION
Although perturbation responses at single cell resolution offers microscopic view of the biological landscape, oftentimes it is fundamental to estimate the average population effect of a perturbation in a given cell type.
Hence in this work, we developed a robust estimation for the causal parameter Ψ( p ) =
E p ( Y ′ |
X
=
c, T ′ =
a ) — the marginal effect of treatment a within a covariate group c .
We propose the following estimator that is asymptotically efficient when E p θ ( Y ′ | ˜
Z k,ϕ , T ′ k
= a ) is estimated consistently and some other regularity conditions (Van Der Laan & Rubin, 2006) hold:
[Equation]
(14) where ( Y k , X k , T k ) are the observed variables of the k -th individual; ˜ Z k,ϕ ∼ q ϕ
( Z | Y k , G , X k , T k ) ; (1 c , . . .
, n c ) are the indices of the observations having X = c and (1 a,c , . . .
, n a,c ) are the indices of the observations having both T = a and X = c .
The derivation of this estimator can be found in Appendix B. Intuitively, we use the prediction error of the model on the observations that indeed have covariate level c and received treatment a to adjust the empirical mean estimator.
E XPERIMENTS
We tested our framework on three datasets in experiments.
We employ the publicly available sciPlex dataset from Srivatsan et al.
(2020) ( Sciplex ) and
CRISPRa dataset from Schmidt et al.
(2022) ( Marson ).
Sci-Plex is a method for pooled screening that relies on nuclear hashing, and the Sciplex dataset consists of three cancer cell lines (A549, MCF7, K562) treated with 188 compounds.
Marson contains perturbations of 73 unique genes where the intervention served to increase the expression of those genes.
In addition, we open source in this work a new dataset ( L008 ) designed to showcase the power of our model in conjunction with modern genomics.
L008 dataset We used the CROP-seq platform (Shifrut et al., 2018) to knock out 77 genes related to the interferon gamma signaling pathway in CD4 + T cells.
They include genes at multiple steps of the interferon gamma signaling pathway such as JAK1, JAK2 and STAT1.
We hope that by including multiple such genes, machine learning models will learn the signaling pathway in more detail.
Baseline We compare our framework to three state-of-the-art self-supervised models for individual counterfactual outcome generation — CEVAE (Louizos et al., 2017), GANITE (Yoon et al., 2018) and CPA (Lotfollahi et al., 2021b), along with the non-graph version of our framework VCI (Wu et al., 2022b).
To give an idea how well these models are doing, we also compare them to a generic autoencoder ( AE ) with covariates and treatment as additional inputs, which serves as an ablation study for all other baseline models.
For this generic approach, we simply plug in counterfactual treatments instead of factual treatments during test time.
3. O UT - OF - DISTRIBUTION P REDICTIONS
We evaluate our model and benchmarks on a widely accepted and biologically meaningful metric — the R 2 (coefficient of determination) of the average prediction against the true average from the outof-distribution (OOD) set (see Appendix C.2) on all genes and differentially-expressed (DE) genes (see Appendix C.1).
Same as Lotfollahi et al.
(2021a), we calculate the R 2 for each perturbation of each covariate level (e.g. each cell type of each donor), then take the average and denote it as ¯
R 2 .
Table 1 shows the mean and standard deviation of ¯ R 2 for each model over 5 independent runs.
Training setups can be found in Appendix C.3.
[Table]
[Figure]
As can be seen from these experiments, our variational Bayesian causal inference framework with refined relation graph achieved a significant advantage over other models on all genes of the OOD set, and a remarkable advantage on DE genes on Marson.
Note that losses were evaluated on all genes during training and DE genes were not being specifically optimized in these runs.
We also examined the predicted distribution of gene expression for various genes and compared to experimental results.
Fig 5 shows an analysis of the CRISPRa dataset where MAP4K1 and GATA were overexpressed in CD8 + T cells (Schmidt et al., 2022), but these cells were not included in the model’s training set.
Nevertheless, the model’s prediction for the distribution of gene expression frequently matches the ground truth.
Quantitative agreement can be obtained from Table 1.
For graphVCI, we used the key-dependent attention (see Section 2.2) for the decoding aggregator in all runs and there were a few interesting observations we found in these experiments.
Firstly, the key-independent attention is more prone to the quality of the GRN and exhibited a more significant difference on model performance with the graph refinement technique compared to the key-dependent attention.
Secondly, with graphVCI and the key-dependent attention, we are able to get stable performances across runs while setting ω 1 to be much higher than that of VCI.
[Figure]
E STIMATIONS
To evaluate the marginal estimator ˆΨ θ,ϕ in Equation 14, we compute ˆΨ θ,ϕ for treatment a and covariate level c with samples from the training set and calculate its R 2 against the true average of the samples with treatment a and covariate level c in the validation set.
We record the average R 2 of all treatment-covariate combo similar to the previous section, and compare it (robust) to that of the regular empirical mean estimator (mean).
Table 2 shows the results episodically during training.
[Table]
These runs reflects that the robust estimator was able to produce a more accurate estimation of the covariate-stratified marginal treatment effect
E p ( Y ′ |
X = c, T ′ =
a ) with a tigher confidence bound.
3. G RAPH E VALUATIONS
In this section, we perform some ablation studies and analysis to validate the claim that the adjacency updating procedure improves the quality of the GRN.
We first examine the performance impact of the refined GRN over the original GRN derived from ATAC-seq data.
For this purpose, we used the key-independent attention decoding aggregator and conducted 5 separate runs with the original
GRN and refined GRN on the Marson dataset (Fig 6a).
We found that the refined GRN helps the graphVCI to learn faster and achieve better predictive performance of gene expression, suggesting the updated graph contains more relevant gene connections for the model to predict counterfactual expressions.
Note that there is a difference in length of the curves and bands in Fig 6a because we applied early stopping criteria to runs similar to CPA.
[Figure]
Next, we compared the learned edge weights from our method to known genetic interactions from the ChEA transcription factor targets database (Lachmann et al., 2010).
We treat edge weights following graph refinement as a probability of an interaction and treat the known interactions in the database as ground truth for two key TFs, STAT1 and IRF1.
We found the refined GRN obtained by graph refinement is able to place higher weights on known targets than a more naive version of the GRN based solely on gene-gene correlation in the same dataset (Fig 6b).
The improvement is particularly noticeable in the high-precision regime where we expect the ground truth data is more accurate, since it is expected that a database based on ChIP-seq would contain false positives.
C ONCLUSION
In this paper, we developed a theoretically grounded novel model architecture to combine deep graph representation learning with variational Bayesian causal inference for the prediction of single-cell perturbation effect.
We proposed an adjacency matrix updating technique producing refined relation graphs prior to model training that was able to discover more relevant relations to the data and enhance model performances.
Experimental results showcased the advantage of our framework compared to state-of-the-art methods.
In addition, we included ablation studies and biological analysis to generate sensible insights.
Further studies could be conducted regarding complete outof-distribution prediction — the prediction of treatment effect when the treatment is completely held out, by rigorously incorporating treatment relations into our framework on top of outcome relations.
A CKNOWLEDGMENTS
We thank Balasubramaniam Srinivasan, Drausin Wulsin, Meena Subramaniam, and Maxime Dhainaut for the insightful discussions.
Work by author Robert A. Barton was done prior to joining
Amazon.
A PPENDIX A P ROOF OF T HEOREM
Proof.
By the d-separation (Pearl, 1988) of paths on the causal graph defined in Figure 1, we have log [ p ( Y ′ | Y, G , X, T, T ′ )]
= log E p
( Z | Y, G ,X,T )
[ p ( Y ′ | Z, Y, G , X, T, T ′ )] (15) ≥
[Equation]
p ( Z | Y
[Equation]
G ,X,T )
[Equation]
[ p ( Y
[Equation]
Z
[Equation]
Y
[Equation]
G
[Equation]
, T, T
[Equation]
)] (Jensen’s ineq.)
(16) =
[Equation]
Y
[Equation]
G
[Equation]
X
[Equation]
T
[Equation]
T
[Equation]
) p ( Z | Y, G , X, T ) (17) =
E
[Equation]
( Z | Y
[Equation]
G
[Equation]
X
[Equation]
T ) [Equation] p ( Y ′ , Z, Y, G , T | X, T ′ ) p ( Z | Y, G , X, T ) p ( Y, T | G , X ) p ( G ) (18) = E p ( Z | Y, G ,X,T ) log p ( Y ′ | G , X, T ′ )
p ( Z | Y ′ , G , X, T ′ )
p ( Y, T | Z, X ) p ( Z | Y, G , X, T ) − log [ p ( Y, T | G , X )] (19) = log [ p ( Y ′ | G , X, T ′ )]
− KL [ p ( Z | Y, G , X, T ) ∥ p ( Z | Y ′ , G , X, T ′ )]
+ E p
( Z | Y, G ,X,T ) log [ p ( Y | Z, T ) p ( T | X )] − log [ p ( Y,
[Equation]
[Equation]
Reorganizing the terms yields
[Equation]
B D ERIVATION OF THE M ARGINAL E STIMATOR
By Van der Vaart (2000), we derive the efficient influence function of Ψ( p ) and thus provides a mean for asymptotically efficient estimation:
Theorem 2.
Suppose D : Ω → E follows a causal structure defined by the Bayesian network in
Figure 1, where the counterfactual conditional distribution p ( Y ′ , T ′ | Z, X ) is identical to that of its factual counterpart p
( Y, T | Z, X ) .
Then Ψ( p ) has the following efficient influence function:
[Equation]
c
[Equation]
Proof.
Following Van der Vaart (2000), we define a path p
ϵ ( D ) =
p ( D )(1 + ϵS ( D )) on density p of D as a submodel that passes through p at ϵ = 0 in the direction of the score S ( D ) =
ddϵ log [ p ϵ ( D )] (cid:12)(cid:12)(cid:12) ϵ
=0 .
Let X =
( G , X ) and minuscule of a variable denote the value it takes.
By Levy (2019), we have
[Equation]
( z, x ) { S ( d ) − E [ S ( D ) | y ′ , z, x , t ′ , t ] } + (cid:90) y ′ ,z, x
I ( x = c ) p ( x ) y ′ p ( y ′ | z, T ′ =
a ) p ( z, x ) { E [ S ( D ) | z, x ] − E [ S ( D ) | x ] } (31) = (cid:90) d
S ( d ) · I ( x = c, t = a ) p ( x, t ) yp ( d ) − (cid:90) d
E [ S ( D ) | y ′ , z, x , t, t ′ ] p ( y ′ , z, x , t, t ′ ) · I ( x = c, t = a ) p ( x, t ) yp ( y | z, t ) + (cid:90) y ′ ,z, x
E [ S ( D ) | z, x ] p ( z, x ) · I ( x = c ) p ( x ) y ′ p ( y ′ | z, T ′ =
a ) − (cid:90) y ′ ,z, x
E [ S ( D ) | x ] · I ( x = c ) p ( x ) y ′ p ( y ′ | z, T ′ =
a ) p ( z, x ) (32) = (cid:90) d
S ( d ) (cid:26) I ( x = c, t = a ) p ( x, t ) ( y − E
[ Y | z, t ])
+ I ( x = c ) p ( x ) ( E [ Y ′ | z, T ′ =
a ] − Ψ) (cid:27) p
( d ) (33) by assumptions of Theorem 2 and factorization according to Figure 1.
Hence
[Equation]
and we have ˜ ψ
[Equation]
By Theorem 2, we propose the following estimator that is asymptotically efficient among regular estimators under some regularity conditions (Van Der Laan & Rubin, 2006):
[Equation]
where ( Y k , X k , T k ) are the observed variables of the k -th individual and
˜ Z k,ϕ ∼ q ϕ
( Y k , G , X k , T k ) ; ˆ p ( T | X ) is an estimation of the propensity score and ˆ p ( X ) is an estimation of the density of X .
In the context of this work, X and T are discrete, hence ˆ p ( T | X ) and ˆ p ( X ) can be estimated by the empirical density p n ( T | X ) and p n ( X ) .
The above estimator then reduces to:
[Equation]
where (1 c , . .
. , n c ) are the indices of the observations having X = c and (1 a,c , . .
. , n a,c ) are the indices of the observations having both T = a and X = c .
C E XPERIMENTS D ETAILS C. D IFFERENTIALLY -E XPRESSED G ENES
To evaluate the predictions on the genes that were substantially affected by the perturbations, we select sets of 50 differentially-expressed genes associated with each perturbation and separately report performance on these genes.
The same procedure was carried out by Lotfollahi et al.
(2021a).
C. O UT - OF -D ISTRIBUTION S ELECTIONS
We randomly select a covariate category (e.g. a cell type) and hold out all cells in this category that received one of the twenty perturbations whose effects are the hardest to predict.
We use these held-out data to compose the out-of-distribution (OOD) set.
We computed the euclidean distance between the pseudobulked gene expression of each perturbation against the rest of the dataset, and selected the top twenty most distant ones as the hardest-to-predict perturbations.
This is the same procedure carried out by Lotfollahi et al.
(2021a).
C. T RAINING S ETUP
The data excluding the OOD set are split into training and validation set with a four-to-one ratio.
A few additional cell attributes available in each dataset are selected as cell covariates.
For Sciplex, they are cell type and replicate; for Marson, they are cell type, donor and stimulation; for L008, they are cell type and donor.
All these covariates are categorical and transformed into discrete indicators before passing to the models.
All common hyperparameters of all models (network width, network depth, learning rate, decay rate, etc.) are set to the same as Lotfollahi et al.
(2021a).
More details regarding hyperparameter settings can be found in our code repository † .
D L IST OF N OTATIONS
Y outcome random vector / gene expressions
X covariates random vector (or variable) / cell type, donor, replicate, etc.
Z latent features random vector
T treatment random vector (or variable) / perturbation
V ′ counterfactual random vector of a random vector V
D concatenation of all vectors (and variables) above n outcome dimensions / number of genes m combined dimensions of covariates r treatment dimensions / number of perturbations Σ event space
P probability measure on Σ
E V measurable space generated by a random vector (or variable) V on Ω c a covariate value / a sample in E X a a treatment value / a sample in E T
G relation graph / gene graph
E adjacency matrix / gene relation network
V feature matrix / gene feature matrix
D concatenation of D and G (see G as a fixed random vector with its components flattened) v feature dimensions of G / diagonal length of E / column dimensions of V
A ( i ) the i -th row of a matrix
A
A i,j the element on the i -th row and j -th column of a matrix A
J objective function
KL
Kullback–Leibler divergence ω · scaling coefficient p θ estimating model for p ( Y | Z, T ) (and p ( Y ′ | Z, T ′ ) )
q ϕ estimating model for q ( Z | Y, G , X, T ) ˜
Y ′ θ,ϕ a sample from E q ϕ
( Z | Y, G ,X,T ) p
θ ( Y ′ | Z, T ′ )
˜
Z ′ ϕ
a sample from q ϕ ( Z | Y, G , X, T ) p M true model for the distribution of latent Y M q H true model for the distribution of latent Z H
f H true model for the attention score on Y M f M true model for latent Z M f G true model for latent Z G h a,β estimating model for a true model
h a with a parameterization β
[Table]
a probability value
[Table]
a threshold value for probability values
I identity matrix
M mask matrix / probability matrix
L edge weight logits matrix
W unnormalized edge weight matrix ˜
W rescaled weight matrix / matrix after applying element-wise sigmoid function on L
U matrix after thresholding ˜ W
O input node feature matrix to g in the graph refinement task
H l latent representation matrix after the l -th layer of g Θ l weight matrix of the l -th layer of g Ψ causal parameter
E C OMPLEXITY A NALYSIS
Overall, the time complexity of VCI compared to a generic framework like VAE (or the time complexity of graphVCI compared to a generic GNN framework like GLUE (Cao & Gao, 2022) does not over-scale on any factor of any parameter – to put it simply, the workflow of VCI is just twice the forward passes of an VAE for a batch of inputs, with an additional distribution loss which is implemented on the same scale O ( r ) as other losses.
Comparing graphVCI to VCI, graphVCI has additionally a few graph operations.
We give a thorough analysis of the time complexity of each layer in our experiments below:
• Generic method AE has 2 MLP layers with O ( rd ) number of operations and 4 layers of O ( d 2 ) number of operations.
d is the number of hidden neurons.
• VCI has the same layer sizes: 2 MLP layers with O ( rd ) number of operations and 4 layers of O ( d 2 ) number of operations.
But every layer is forward passed twice.
• graphVCI has 1 GNN layer with O ( rvd 2 g + pd g r 2 ) number of operations.
v is number of gene features, d g is number of hidden neurons for GNN (usually a lot less than d since v ≪ r ) and p is the sparsity of GRN (number of connections divided by r 2 , usually around 1% ), and the following layers which are forward passed twice: 1 MLP layer and 1 dotproduct operation each with O ( rd ) number of operations, 1 MLP layer with O ( rd g ( d g + d )) number of operations and 3 MLP layers with O ( d 2 ) number of operations.
So the terms to be concerned compared to AE and VCI are O ( rvd 2 g ) , O ( pd g r 2 ) and O ( rd g d ) .
Since p is around 1% and r is 2000 in our experiment, pd g r 2 ≈ 20 d g r hence the second term is comparably smaller than the other two terms.
Therefore, as long as d g is set to be reasonably small compared to d , the graph approach is reasonably scaled compared to pure MLP approaches.
We note that this is a limitation of ours and GNN approaches in general: if there is a much high number of genes r than 2000 to be considered, or a high number of gene features v for each gene (which results in that d g has to be higher), GNN methods does not scale favorably compared to MLP methods.