Predicting Cellular Responses with Variational Causal Inference and Refined Relational Information
Authors: Yulun Wu
Rob Barton
Zichen Wang
Vassilis N. Ioannidis
Carlo De Donno
Layne C Price
Luis F. Voloch
George Karypis
Venue: ICLR 2023 poster (2023)
Keywords: graph neural network
causal inference
variational bayes
asymptotic statistics
single-cell perturbation

Abstract: 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.

Published as a conference paper at ICLR 2023
PREDICTING CELLULAR RESPONSES WITH VARIATIONAL CAUSAL INFERENCE AND REFINED RELATIONAL INFORMATION
Yulun Wu
Robert A. Barton
University of California, Berkeley yulun wu@berkeley.ee edu
Immunai
robert bebbertbbacee www.bbattntnnuuaaa..
Zichen Wang
Vassilis N. Ioannidis
Carlo De Donno Immunai carlo. d immumiimm.mmmmmmmmmmmmm com
Amazon
Amazon
@.chewan@amazoon com
@amazon. prasieei@amazon. ivasi Massileine
com
Layne C. Price
Luis F. Voloch
George Karypis
Amazon
Immunai
Amazon @karypis@amazonm
pril mrilaanneeen aaaieaamazzo...o com
luise immumii.
ABSTRACT
Predicting the responses of cell under perturbations may bring important ben- efits to drug discovery and personalized therapeutics. In this work, we propose novel graph variational Bayesian causal inference framework to predict cell's s gene expressions under counterfactual perturbations (perturbations that this cell did not factually receive), leveraging information representing biological knowl- edge in the form of gene regulatory networks (GRNs) to aid individualized cellular response predictions. Aiming at data-adaptive GRN, we also developed an adja- cency matrix updating technique for graph convolutional networks and used it refine GRNs during pre-training, which generated more insights on gene relations and enhanced model performance. Additionally, we propose robust estimator within our framework for the asymptotically efficient estimation of marginal per- turbation effect, which is yet be carried out in previous works. With exten- sive experiments, we exhibited the advantage of our approach over state-of-the-art deep learning models for individual response prediction.
INTRODUCTION
Studying a cell's response to genetic, chemical, and physical perturbations is fundamental in un- derstanding various biological processes and can lead to important applications such as drug dis- covery and personalized therapies. Cells respond to exogenous perturbations at different levels, in- cluding 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 the development of sev- eral methods for predicting single-cell transcriptional responses (Ji et al. 202T). These methods fall into two broad categories. The first category (Lotfollahi 2019 2020 Rampášek et al.) 2019 Russkikh et al.] [0200] Lotfollahi et al. 2021a) approaches the problem of predicting sin- gle 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-Strein et al. 2014). Notably among those studies, CPA (Lotfollahi et al. 202 uses adversarial autoencoder framework designed to decompose the cellular gene expression sponse to latent components representing perturbations, covariates and basal cellular states. CPA ex- tends the classic idea of decomposing high-ding sigh-dinnnsionll gene expression response into perturbation vectors (Clark et al. 2014 2013). which can be used for finding connections among perturbations (Subramanian al.
2017 However, while CPA's adversarial approach encourages latent indepen-

---

Published as a conference paper at ICLR 2023
dence, it does not have any supervision the counterfactual outcome construction and thus does not explicitly imply that the counterfactual outcomes would resemble the observed response distri- bution. 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 bring- ing the benefits of deep learning graph data, graph neural networks (GNNs) offer 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 includ- ing molccule/protein property prediction (Guo et al. [222] loannidis et al.] 2201;; |Strokach et al.) 2020 |Wu et al. 2022a. Wang al. 2022), perturbation prediction (Roohani et al. 2022] and RNA- sequence analysis (Wang et al.) (02T). In previous work, CCo & Gao (2022) developed GLUE, multi-omics datasets to improve multimodal data integration and response prediction. GEARS (Roohani et al.)
a framework leveraging a fine- tine-grained GRN with nodes corresponding to features
2022] uses GNNs to model the relationships among observed and perturbed genes to predict cellular response. These studies demonstrated that relation graphs informative for predicting cellular sponses. However, GLUE does not handle perturbation response prediction, and GEARS's 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 cunssosase-accessible chromatin using sequencing (ATAC-seq). However, GRNs from these ap- proaches 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 ac- curate GRNs (Mokry et al. 2012 Jiang & Mortazavi] 2018 Angelini & Costa 2014). GRNs are also highly sontext-specific: different cell types can have very distinctive GRNs mostly due to their different epigenetic landscapes (Emerson 2002; Davidson 2010). Hence, GRN derived from the most relevant biological system is necessary to accurately infer the expression of individual genes within such system.
In this work, we employed a novel variational Bayesian causal inference framework to construct the gene expressions of cell under counterfactual perturbations by explicitly balancing individual features emhedded in its factual outcome and marginal response distributions of its cell population. We integrated gene relation graph into this framework, derived the corresponding variational lower bound and designed an innovative model architecture to rigorously incorporate relational informa- tion 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 re- lations in this context) and enhanced model performance. Besides, we propose an asymptotically efficient estimator for estimating the average effect of perturbations under given cell type within our framework. Such marginal inference is of great biological interest because scRNA-seq experi- mental 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 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 a task commonly used in previous works on perturbation predictions. In addition, we carried out ablation studies demonstrate the advantage of using refined relational information for better understanding of the contributions of framework components.
distr--distribution predictions on differentially-exprsse genes
PROPOSED METHOD
In this section we describe our proposed model - Graph Variational Causal Inference (graphVCI), and relation graph refinement technique. A list of all notations can be found in Appendix A
2

---

Published as a conference paper at ICLR 2023
Z
g
Y
Y
T
T
T
X
(a) folded
(b) unfolded
Figure I: The causal relation diagram. Each individual has feature state Z following a conditional distribution p( p[Z]g, X). Treatment T (or counterfactual treatment T') along with Z determines outcome Y (or counterfactual outcome Y"). In the causal diagram, white nodes are observed and dark grey nodes unobserved; dashed relations are optional (case dependant). In the context of this paper, graph G is deterministic variable that is invariant across all individuals.
2.1 COUNTERFACTUAL CONSTRUCTION FRAMEWORK
We define outcome Y Ω R" to be r-dimenniioal random vector (e.g. gene expressions), X Q + Ex to be m-dimenniooal mix of categorical and real-valued covariates (e.g. cell types, donors, etc.), T Ω + Er to be r-dimensional categorical or real-valued treatment (e.g. drug perturbation) on a probability space (2, P). We seek to construct an individual's I's counterfactual outcome under counterfactual treatment E Er from source is the individual features embedded in high-dimensional outcome Y. The other source the response distribution of similar subjects (subjects that have the same covariates as this individual) that indeed received treatment a.
major sources of information. One
We employ the variational causal inference (Wu et al. 2022b) framework to combine these two sources of information. In this framework, covariate Serariate-depeneent feature vector Z : 52 - R 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 combined likelihood of undividual-specife treatment effect p(Y'(Y, T, T") (first source) and the traditional covariate-specific treatment effect p( T)) (second source):
(1) where - Z,T,T Y, Y"). Additionally, we assume that there is graph structure G - E) that governs the relations between the dimensions of Y through latent Z, where V E Raxx is the node feature matrix and E E [0,1)aXB the node adjacency matrix. For example, in the case of single-cell perturbation dataset where Y is the expression counts of genes, 2/ is the gene feature matrix and is the GRN that governs gene relations. See Figure for a visualization of the causal diagram. The objective under this setting is thus formulated
J(D) = log (P(Y'IY, X, T')] log [p(||X,)) 7))
(2) = log [P(Y'[Y, G. G,T,T" + log [p(Y G,X,T))
(2)
where D D D). The counterfactual outcome Y' is always unobserved, but the following theorem provides us a roadmap for the stochastic optimization
this objective.
Theorem 1. Suppose that Ø (g X. Z, T, T', Y,Y") follows a causal structure defined by the Then J(D) has the following evidence lower bound:
Bayesian network Figure 7
J(W) > E Ep x.T) log [p Y Z, T) + log [p(Y') (gy Ng,,G,,XT. X,T")
KT. ( V G. X, T) |1 2( g. X, X.))
Proof of the theorem can be found in Appendix We estimate p(Z G, X, T) and (ZY)Z,T) (as well as Z I')) with a neural network encoder 96 and decoder and optimize the following weighted approximation of the variational lower bound:
J [0, Φ) (2). log [pe Y T"] + log ((66.
We KL qo (Z G, X,T)
g (21%% G,X,T")
3

---

Published as a conference paper at ICLR 2023
Reconsiructon Loss
Distribution Loss
KL-divergence
Decoder Block
Decoder
Block
Encoder Block
Enooder Blockk
Y MBN2
7 P
X
ANT2
BOB
BBC
BB
FTEN
Figure 2: Model workflow variational causal perspective. In a forward pass, the graph encoder takes graph G (e.g. gene relation graph), outcome (e.g. gene expressions), covariates X (e.g. cell types, donors, etc.) and treatment T (e.g. drug perturbation) as inputs and generates latent Z: (Z, 7) and (Z, T") where 7" a randomly sampled counterfactual treatment are separately passed into the graphVCI decoder to attain reconstruction of Y and construction of counterfactual outcome Y'; Y' then passed back into the encoder along with G. X, T" to attain counterfactual latent Z'. The objective consists of the reconstruction loss of Y. the distribution loss of Y' and the Kl.-divergence between the conditional distributions of Z and 7'.
where W1, We are scaling coefficients; Vg.0 G. (Y' 2, T' and P is the covariate- specific model fit of the outcome distribution (notice that ", X, T' =a) =p(YY,, X,T, a) for any a). our case where covariates are limited and discrete (cell types and donors), we simply let p(Yl X, T) be the smoothened empirical distribution of Y stratified by X and T (notice that fixed across subjects). Generally, one can train a discriminator p(- G. X, T) with the adversarial
approach (Goodfellow al. 2014) and use log [(1)) 0.00 G, G.X,T)) for log p X, X,T*) if p(Y|G, X,T) 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 7" need not necessarily be sampled according certain true distribution p(T'lxX 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. discussed in Wu et al.] (2022b), the negative KL-divergence term in the objective function encourages the preservation individuality in counterfactual outcome constructions.
Marginal Effect Estimation Although perturbation responses at single cell resolution offers mi- croscopic view of the biological landscape, oftentimes it is fundamental to estimate the average population effect of perturbation in given cell type. Hence in this work, we developed ro- bust estimation for the causal parameter (p) = EY' = c, T' = =) - the marginal effect of treatment within covariate group c. We propose the following estimator that is asymptotically efficient when Epp (Y'l ZR o. TV a) is estimated consistently and some other regularity condi- tions (Van Der Laan & Rubin 2006) hold:
1
Yz - (Y' (Y'l ZE. T' = =a)=]
Σ k=
Eyo (Y'I T' =a) (5)
na,c k=la.
where (Yx, Tx) are the observed variables of the k-th individual; Ze, ~ go( G. XR. Tx);
(1 (1e--... ........e)) n1.) are the indices of the observations having X = c and (Ia,er-... na,e) the indices
4

---

Published as a conference paper at ICLR 2023
MOME
MIOR FINH ARTI
Cross Attention
FF Block
Decoder Block
FF Block
GA Block
ZM
Aggr
GA Block
FF Block
XNN
Layer Norm
f8
FF Block
Add
fu
Batch Norm
Graph
Attention
Feed
Forward
Encoder Block
1903
174
KOOME ARF
MYC MITOR PIEH 28T3
A03
A0T1
8
Figure 3: Model architecture - graph attentional perspective. Structure of graph VCI encoder and decoder defined by Equations (6) (II). Note that in the case of single-cell perturbation datasets, the graph inputs are fixed across samples and graph attention can essentially reduced to weighted graph convolution.
of the observations having both and c. The derivation of this estimator can found in Appendix C.2 and some experiment results can be found in Appendix C.T Intuitively, we use the prediction error the model on the observations that indeed have covariate level and received treatment to adjust the empirical mean. Prior works in cellular response studies only use the empirical mean estimator over model predictions to estimate marginal effects.
2.2 INCORPORATING RELATIONAL INFORMATION
Since the elements in the outcome Y are not independent, aim to design framework that can exploit predefined relational knowledge among them. In this section, we demonstrate our model structure for encoding and aggregating relation graph within the graphVCI framework. Denote deterministic models as f., probabilistic models (output probability distributions) as and p.. We construct feature vector Z as an aggregation two latent representations:
= (Zg. Zn)
(6)
2 ~ Q11 (7,, aggr & (Zg))
(7)
ZM = fY,
(8)
Zg fG(G)
where aggr, a node aggregation operation such as sum, max or mean. The optimization of Is can then be performed by optimizing the MLP encoder fM... gatmtr - Rd the GNN encoder fç.02 1xa - dd and the encoding aggregator qH.d, gd+d, + ((0, 1 where (Ø1, Φ2, 0.). We designed such construction so that possesses a node-level graph embedding Zg that enables more involved decoding techniques than generic MLP vector decoding, while the calculation of term KL [9, Y, X,T)) || 4s (Z)Y", G,X, X")] can also be reduced to the KL- divergence between the conditional distributions of two graph-level vector representations Zn and
5

---

Published as a conference paper at ICLR 2023
Z/1 (since Zg is deterministic and invariant across subjects). In decoding, we construct
(
i1 fH (Zoti)) Xw))
Y
Ya
(10)
YM Pas (an,
(11)
where || represents vector concatenation and optimize pe by optimizing the MLP decoder PM.0;
getr + ((0, 1 and the decoding aggregator fHH0, IRd+d, - He where The decoding aggregator maps graph embedding Zo(n) of the i-th node along with feature vec- tor Yı to outcome Y(s) of the i-th dimension, for which we use an attention mechanism and let Y(a) = attp, (ttg,(Zg(o) YM) T YM where (t) K) gives the attention score for cach of K's feature dimensions given querying node embedding vector Q()) (i.e., row vector of Q). One can simply use key-independent attention mechanism Y() atte, (Z)())) T YAs if GPU memory is of con- cern See Figure for visualization of the encoding and decoding models. A complexity analysis of the model can be found in Appendix D
(0), 02).
2.3 RELATION GRAPH REFINEMENT
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 provide an option to impute and refine prior GRN by learning from the expression dataset of interest. We propose approach to update the adjacency matrix in GCN training while maintaining sparse graph operations, and use it in an auto-regressive pre-training task to obtain an updated GRN.
Let g(+) be a GCN with learnable edge weights between all nodes. We aim to acquire an updated adjacency matrix 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 the initial graph - edges present in E E I[ + I (7 is the identity matrix) are accompanied with low dropout rate ri and edges not present in E with high dropout rate Tb (rt < rh). A graph convolutional layer of g(-) is then given as
H o (softmaxx, (M @ L)H'O')
(12)
where I. E Raxx is a dense matrix containing logits of the edge weights and M E RRux is a sparse mask matrix where each element Mi, is sampled from Bern(T = 1)r: [(Ex, 0)rx) ) in every iteration; @ is element-wise matrix multiplication; softmax, is row-wise softmax operation; H' Baxd, H++1 E =xxd++ are the latent representations of V after the /-th, (1 1)-th layer, O' E [Rddxd++ is the weight matrix of the (2 + 1)-th layer; a is a non-linear function. The updated adjacency matrix É E R"x" is acquired by rescaling and thresholding the unnormalized weight matrix W eXp(L) after optimizing 9()):
Bi = sgn [11 W
(13)
where a is threshold level. We define W E m to be the rescaled weight matrix having W.y (1 - With this design, 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 1 prior to main model training presented in the previous sections. We train g(-) 12 - ggxl on a simple node-level prediction task where the output of the i-th node is the expression Y(a) of the i-th dimension; the input of the i-th node is a combination 0/1) (Y()): Viy, X) with expression Y(a) gene features Vi) of the i-th gene and cell covariates X. Essentially, we require the model to predict the expression of 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 al. 2020 The objective is a lasso-style combination of reconstruction loss and edge penalty:
J(g) =-119(0) Y w ||W I|
(14)
where is a scaling coefficient. Note that although V(e) not cell-specific and X is not gene- specific, the combination of (N(), forms unique representation for each gene of each cell

---

Published as a conference paper at ICLR 2023
(a) original
(b) updated
Figure 4: An example of an updated gene regulatory network U (where U2
al with
0.2) after refining the original ATAC-seq- network Kamimoto et al. 2020) using the Schmidt al. (202) dataset. Source nodes are shown as rows and targets are shown as columns for key immune-related genes. The learned edge weights in (b) recapitulate known biology such as STATI regulating IRFI. Also note that while many the edges are present in the original ATAC-seq data from (a), we see some novel edges in (b) such as IFNg regulating MYC (Ramana et al. 2000).
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 this task, dropouts and data correlations will still prevent them from being the only notable signals. See Figure for an example an updated GRN in practice.
3 EXPERIMENTS
We tested our framework on three datasets in experiments. We employ the publicly available sci- Plex dataset from Srivatsan al. (2020) (Sciplex) and CRISPRa dataset from Schmidt et al. (2022) (Marson). Sci-Plex is method for pooled screening that relies 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 new dataset (I.008) designed to showcase the power of our model in conjunction with modern genomics.
L008 dataset We used the CROP-seg platform (Shifrut et al. 2018 to knock out 77 genes related 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 three state-of-the-arr self suff-uupervieed models for individ- ual counterfactual outcome generation - CEVAE (Louizos et al.) (017). GANITE (Yoon et al.) 2018] and CPA (Lotfollahi al. (02Tb). along with the non-graph version of our framework VCI (Wu et al. (022b).. 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 counter- factual treatments instead factual treatments during test time.
3.1 OUT-OF-DISTRIBUTION PREDICTIONS
We evaluate our model and benchmarks on a widely accepted and biologically meaningful metric the R² (coefficient of determination) of the average prediction against the true average from the out- of-distribution (OOD) set (see Appendix E.2) on all genes and differentially-expesss (see Appendix R. Same Lotfollahi al. (2021a), calculate the R² for each perturbation
(DE) genes

---

Published as a conference paper at ICLR 2023
of each covariate level (e.g. each cell type of each donor), then take the average and denote it as R Table I shows the mean and standard deviation R for each model over independent runs.
Training setups can be found in Appendix E.3
Table 1: R2 of OOD predictions
Sciplex
Marson
L008
genes
DE genes
all genes
DE genes
all genes
genes
AE
0.740 1 0.043
0.421
0.021
0.804
0.020
0.448
0.009
0.948 0.010
0.729
0.041
CEVAE
0.760 0.019
0.436
0.014
0.795
0.014
0.424
0.015
0.941 0.010
0.632 0.034
GANITE1
0.751 +
0.417
0.014
0.795
0.017
0.443
0.025
6.946+00000
0.730 0.030
CPA
0.836 0.002
0.474
0.014
0.876
0.005
0.549
0.019
0,962 0:005
0.849 0.022
VCI
0.828
0.006
0.492
0.011
0.884
0.010
0.604
0.044
0.962
0.002
0.865 1.038
VCT
0.841 0.002
0..997
0.014
0,892 1.892-00000
0.642 0.000
"6,65 0.002
0.31110016
As can be seen from these experiments, our variational Bayesian causal inference framework with refined relation graph achieved 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 being specifically optimized in these runs.
We also examined the predicted distribution of gene expression for various genes and compared to experimental results. Fig shows an analysis of the CRISPRa dataset where MAP4KI and GATA3 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 's prediction for the distribution of gene expression frequently matches the ground truth. Quantitative agreement can be obtained from Table [
Prec.
Tue
MAPPOST
C22
(a) MAP4K1 Overexpression
GATA3 Overexpression
Figure 5; Model predictions versus true distributions for overexpression of genes in CRISPRa exper- iments (Schmidt al. 2022). For two perturbations in CD8+ T cells, (a) MAP4K1 overexpression and (b) GATA3 overexpression, we plot the distribution of gene expressions for unperturbed cells ("Control"), the model's prediction of perturbed gene expressions using unperturbed cells as factual inputs ("Pred")), and the true gene expressions for perturbed cells ("True")) The predicted distribu- tional shift relative to control often matches the direction of the true shift.
For VCT, we used the key-dependent attention (see Section 2.2h for the decoding aggregator in all runs and there were 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 sig- nificant difference on model performance with the graph refinement technique compared to the key-dependent attention. Secondly, with graph VCI and the key-dependent attention, we are able to get stable performances across runs while setting wi to much higher than that of VCI.
9 "Autnencoder with covariates and treatment as additional inputs. 'GANITE's counterfactual block. GANITE's counterfactual generator does not scale with combination of high-dimensional outcome and multi-level treatment, hence made the adaptation Wa al. (202hh
8

---

Published as a conference paper at ICLR 2023
3.2 GRAPH EVALUATIONS
In this section, perform some ablation studies and analysis validate the claim that the adjacency updating procedure improves the quality the GRN. We first examine the performance impact of the refined GRN over the original GRN derived from ATAC-seq data. For this purpose. used the key-independent attention decoding aggregator and conducted separate runs with the original GRN and refined GRN the Marson dataset (Fig 6a We found that the refined GRN helps the graph VCI 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 difference in length of the curves and bands in Fig 6a because we applied early stopping criteria to runs similar CPA.
1.0
STATT PRAAUC Con DUMA
STATT- PRAAIC GON S
PRRAUC
0.8
0.8
0.4
Initiel. all inttal, ce
0.2
upeated,
100
200
300
400
500
000
700
0.0
0.0
0.2
0.4
0.6
0.8
1.0
Epochs
recall
with different GRNs
Agreement of learned graph with TF targets
Figure 6: (a) Graph refinement improves model training. We learned GRN with graph refinement as described in section 2.3] and edges were retained using threshold refinement, the model is better able to reconstruct all genes and differentially expressed (DE) genes during training as can be seen
0.3. Following the
(b) examine whether the edges learned in refinement are accurate by comparing to database of targets for two important TFs, STATI and IRF1. Refined edges between these TFs and their targets agree well with the interactions in the database, according to precision-recall curve where treat the database as the true label and edge weights as a probability of interaction. We also compare the refined edges ("GCN") to gene-gene correlation benchmark ("Corr") and find that the refined graph can outperform the benchmark.
graph of RP vs. number training epochs.
Next, compared the learned edge weights from method to known genetic interactions from the ChEA transcription factor targets database (Lachmann et al. 2010). We treat edge weights probability of an interaction and treat the known interactions in the
following graph refinement database as ground truth for two key TFs, STATI and IRF1. We found the refined GRN obtained by graph refinement is able to place higher weights the GRN based solcly on gene-gene correlation in the same dataset (Fig 6b). The improvement is particularly noticeable 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.
known targets than
naive version of
CONCLUSION
In this paper, we developed 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 rela- tion 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 frame- work compared to state-of of-the--rr methods. addition, we included ablation studies and biological analysis to generate sensible insights. Further studies could be conducted regarding complete out- of-distribution prediction - the prediction of treatment effect when the treatment is completely held out, by rigorously incorporating treatment relations into framework top of outcome relations.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank Balasubramaniam Srinivasan, Drausin Wulsin, Meena Subramaniam, and Maxime Dhain- aut for the insightful discussions. Work by author Robert A. Barton was done prior to joining Amazon.
REFERENCES
Claudia Angelini and Valerio Costa. Understanding gene regulatory mechanisms by integrating chip-seq and ma-seq data: statistical solutions to biological problems. Frontiers in cell and de- velopmental biology, 2:51, 2014.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geomet- ric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-42. 2017.
Zhi-Jie Cao and Ge Gao. Multi-omics single-cell data integration and regulatory inference with graph-linked embedding. Nature Biotechnology, pp. 1-9, 2022.
Neil R Clark, Kevin S Hu, Axel Feldmann, Yan Kou, Edward Chen, Qiaonan Duan, and Avi Ma' ayan. The characteristic direction: geometrical approach identify differentially expressed genes. BMC bioinformatics, 15(1):1-16, 2014.
Neil R Clark, Maciej Szymkiewicz, Zichen Wang, Caroline D Monteiro, Matthew R Jones, and Avi Ma' ayan. Principle angle enrichment analysis (paea): Dimensionally reduced multivariate gene set enrichment analysis tool. In 2015 IEEE International Conference Bioinformatics and Biomedicine (BIBM), pp. 256-262. IEEE, 2015.
Eric H Davidson. The regulatory genome: gene regulatory networks in development and evolution. Elsevier, 2010.
Beverly M Emerson. Specificity of gene regulation. Cell, 109()):267-277, 2002.
Frank Emmert-Streib, Matthias Dehmer, and Benjamin Haibe-Kaii Kains. Gene regulatory networks and their applications: understanding biological and medical problems in terms of networks. Frontiers in cell and developmental biology, 2:38, 2014.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neurul information processing systems, 27, 2014.
Zhichun Guo, Bozhao Nan, Yijun Tian, Olaf Wiest, Chuxu Zhang, and Nitesh V Chawla. Graph- based molecular representation learning. arXiv preprint arXiv:2207. 224009,0,,.04869, 2022.
Vassilis N loannidis, Antonio G Marques, and Georgios B Giannakis. Graph neural networks for predicting protein functions. In 2019 IEEE 8th International Workshop on Computational Ad- vances in Multi-Sensor Adaptive Processing (CAMSAP), pp. 221-225. IEEE, 2019.
Yuge Ji, Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. Machine learning for per- turbational single-cell omics. Cell Systems, 12(6):522-537, 2021.
Shan Jiang and Ali Mortazavi. Integrating chip-seq with other functional genomics data. Briefings in functional genomics, 1742)104-115, 2018
Kenji Kamimoto, Christy M Hoffmann, and Samantha Morris. Celloracle: Dissecting cell identity via network inference and in silico gene perturbation. BioRxiv, 2020.
Alexander Lachmann, Huilei Xu, Jayanth Krishnan, Seth I Berger, Amin R Mazloom, and Avi Ma'ayan. Chea: transcription factor regulation inferred from integrating genome-wide chip-x experiments. Bioinformatics, 26 (6(19))2438-2 2444, 2010.
Jonathan Levy. Tutorial: Deriving the efficient influence curve for large models. arXiv preprint arXiv:1903.000 2019.
10

---

Published as a conference paper at ICLR 2023
Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. segen predicts single-cell perturba- tion responses. Nature methods, 16(8):715-721, 2019.
Mohammad Lotfollahi, Mohsen Naghipourfar, Fabian J Theis, and F Alexander Wolf. Condi- tional out-o--diitributtion generation for unpaired data using transfer vae. Bioinformatics, 36 (Supplement 2):1610-i66 2020.
Mohammad Lotfollahi, Anna Klimovskaia Susmelj. Carlo De Donno, Yuge Ji, Ignacio L Ibarra, F Alexander Wolf, Nafissa Yakubova, Fabian J Theis, and David Lopez-Paz. Compositional perturbation autoencoder for single-cell response modeling. BioRxiv, 2021a.
Mohammad Lotfollahi, Anna Klimovskaia Susmelj, Carlo De Donno, Yuge Ji, Ignacio L Ibarra, F Alexander Wolf, Nafissa Yakubova, Fabian J Theis, and David Lopez-Paz. Learning inter- pretable cellular responses to complex perturbations in tighthhroughut screens. bioRxiv, 2021b.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling, Causal effect inference with deep snrent-varibble models. Advances in neural information pro- cessing systems, 30, 2017.
Stefan R Maetschke, Piyush Madhamshettiwar, Melissa Davis, and Mark A Ragan. Supervised, semi-supervised and unsupervised inference of gene regulatory networks. Briefings in bioinfor- matics, 15(2)1195-21, 2014.
Michal Mokry, Pantelis Hatzis, Jurian Schuijers, Nico Lansu, Frans-Paul Ruzius, Hans Clevers, and Edwin Cuppen. Integrated genome-wide analysis of transcription factor occupancy, ma polymerase ii binding and steady-state rma levels identify differentially regulated functional gene classes. Nucleic acids research, 40(1):145 158, 2012.
Stephen Oliver. Guilt-by-association goes global. Nature, 403(6770):601 602, 2000.
Judea Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan kaufmann, 1988.
Chilakamarti V Ramana, Nicholas Grammatikakis, Mikhail Chernov, Hannah Nguyen, Kee Chuan Goh, Bryan RG Williams, and George R Stark. Regulation of c-myc expression by ifn-y through statl-dependent and-independent pathways. The EMBO journal, 19(2):263-25) 2000.
Ladislay Rampášek, Daniel Hidru, Petr Smirnov, Benjamin Kains, and Anna Goldenberg. Dr. vae: improving drug response prediction via modeling of drug perturbation effects. Bioinformat- ics, (5(19)):743--3751, 2019.
Irene Gallego Romero, llya Ruvinsky, and Yoav Gilad. Comparative studies of gene expression and the evolution of gene regulation. Nature Reviews Genetics, 13(7):505-516. 2012.
Yusuf Roohani, Kexin Huang, and Jure Leskovec. Gears: Predicting transcriptional outcomes of novel multi-gene perturbations. bioRxiv, 2022.
Nikolai Russkikh, Denis Antonets, Dmitry Shtokalo, Alexander Makarov, Yuri Vyatkin, Alexey Zakharov, and Evgeny Terentyev. Style transfer with variational autoencoders is promising approach to rna-seq data harmonization and analysis. Bioinformatics, 36620):5066-5085, 2020.
Ralf Schmidt, Zachary Steinhart, Madeline Layeghi, Jacob W Freimer, Raymund Bueno, Vinh Q Nguyen, Franziska Blaeschke, Chun Jimmie Ye, and Alexander Marson. Crispr activation and interference screens decode stimulation responses in primary human cells. Science, 375(6580): cabj4008, 2022.
Eric Shifrut, Julia Carnevale, Victoria Tobin, Theodore L. Roth, Jonathan M Woo, Christina T Bui, Jonathan Li. Morgan E Diolaiti, Alan Ashworth, and Alexander Marson. Genome-wide crispr screens in primary human t cells reveal key regulators of immune function. Cell, 175(7):1958- 1971, 2018.
François Spitz and Eileen EM Furlong. Transcription factors: from enhancer binding to develop- mental control. Nature reviews genetics, 13(9):613-626, 2012.

---

Published as a conference paper at ICLR 2023
Sanjay R Srivatsan, José L McFaline- Figueroa, Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah A Pliner, Dana Jackson, Riza M Daza, Lena Christiansen, et al. Mas- sively multiplex chemical transcriptomics at single-cell resolution. Science, 367(6473):45-51, 2020.
Alexey Strokach, David Becerra, Carles Corbi-Verge, Albert Perez-Riba, and Philip M Kim. Fast and flexible protein design using deep graph neural networks. Cell systems, 11(4):402-411 2020.
Aravind Subramanian, Rajiv Narayan, Steven M Corsello, David D Peck, Ted E Natoli, Xiaodong Lu, Joshua Gould, John F Davis, Andrew A Tubelli, Jacob K Asiedu, et al. A next generation connectivity map: L1000 platform and the first 1,000,000 profiles. Cell, 173(6):4437-145, 2017.
Mark J Van Der Laan and Daniel Rubin. Targeted maximum likelihood learning. The international journal of biostatistics, 2(1), 2006.
Aad W Van der Vaart. Asymptotic statistics, volume Cambridge university press, 2000.
Juexin Wang, Anjun Ma, Yuzhou Chang, Jianting Gong, Yuexu Jiang, Ren Qi, Cankun Wang, Hongjun Fu, Qin Ma, and Dong Xu. segnn is a novel graph neural network framework for single- cell ma-seq analyses. Nature communications, 12(1):1-11, 2021.
Zichen Wang. Steven A Combs, Ryan Brand, Miguel Romero Calvo, Panpan Xu, George Price, Nataliya Golovach, Emmanuel O Salawu, Colby J Wise, Sri Priya Ponnapalli, al. Lm-gyp: an extensible sequence and structure informed deep learning framework for protein property predic- tion. Scientific reports, 12(1):1-12, 2022.
Yulun Wu, Nicholas Choma, Andrew Chen, Mikaela Cashman, Fírica T Prates, Manesh Shah, Verónica G Melesse Vergara, Austin Clyde, Thomas S Brettin, Wibe A de Jong, et al. Spatial graph attention and curiosity-driven policy for antiviral drug discovery. ICLR, 2022a.
Yulun Wu, Layne C Price, Zichen Wang, Vassilis N loannidis, and George Karypis. Variational causal inference. arXiv preprint arXiv:2209.05933. 05935, 2022b.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International Conference Learning Rep- resentations, 2018
12

---

Published as a conference paper at ICLR 2023
APPENDIX
A
LIST OF NOTATIONS
Y
outcome random vector gene expressions covariates random vector (or variable) / cell type, donor, replicate, etc. latent features random vector
X
Z
treatment random vector (or variable) perturbation counterfactual random vector of random vector V concatenation of all vectors (and variables) above outcome dimensions number of genes combined dimensions of covariates treatment dimensions number of perturbations event space
V
D
m
P
probability measure on measurable space generated by random vector (or variable) V on Ω a covariate value sample in Ex treatment value / a sample in ET relation graph / gene graph adjacency matrix / gene relation network feature matrix / gene feature matrix
Fy
G
V
concatenation of D and G (see G fixed random vector with its compo- nents flattened)
feature dimensions of G diagonal length of column dimensions y the i-th row of a matrix A the element on the i-th row and j-th column of matrix objective function
A
Ard
KL
Kullback-Leible divergence scaling coefficient
estimating model for p(Y| (and p( T') estimating model for (()) X, sample from P4.. sample from 4o G, true model for the distribution of latent YM
po
qo
Yoo
Z'
PM
Q8
true model for the distribution of latent Zu
fH
true model for the attention score on YM
fM
true model for latent Zas
fg
true model for latent Zg estimating model for a true model ha with parameterization 3 dimensions Zas Zu YM
ha.s
dg
feature dimensions of Zg alias for f'r
att
aggr
graph node aggregation operation matrix aggregation operation row-wise softmax function
softmax,
element-wise matrix multiplication
non-linear function graph convolution network
g
probability value
threshold value for probability values identity matrix
M
mask matrix / probability matrix
edge weight logits matrix unnormalized edge weight matrix
W.
W
rescaled weight matrix / matrix after applying element- wise sigmoid function on T,
U
matrix after thresholding W
O
input node feature matrix to g in the graph refinement task latent representation matrix after the [-th layer of
H
O!
weight matrix of the /-th layer of causal parameter
T
13

---

Published as a conference paper at ICLR 2023
B PROOF OF THEOREM 1
Proof. By the d-separation (Pearl 1988) of paths on the causal graph defined in Figure) 1 we have
log (Y g, G,T,TT)) log (,(.I (,(Y" g, X
(15)
xxIx) log (Y 9, XTT.T
(Jensen's ineq.)
(16)
, GAY.G,X,T,,
BPIZIY,G log
(17)
p(Z G,
p(Y', ,2,Y, 9: G.T,
(.( log
(18)
(Y G, G.X,T")p( G.X,TT')
Ep(ZiY,g.x.IT) log
G,X,T)
- log log
(19)
llog log) -LIII/K pl p( G,X,T"
+ 1on12//ggx.n/ogg////////1///1111 log g.X)] (20)
+ 1oggp1>"(..Ax777--k////XV1.... p(2 ((,I'"
(21)
Reorganizing the terms yields
log GX,,, (1 2,111) log (
+ ( 1 G.X,T")
KL
(22)
C MARGINAL EFFECT ESTIMATION
C.I EXPERIMENT
To evaluate the marginal estimator 9 in Equation 5 compute V for treatment and co- variate level with samples from the training set and calculate its R² against the true average of the samples with treatment a and covariate level e in the validation set. We record the average R² of all treatment-covarriate combo similar to Section 3.1 and compare it (robust) to that of the reg- ular empirical mean estimator (mean). Table shows the results on Marson (Schmidt et al. 2022) episodically during training.
Table 2: Comparison of marginal estimators on Marson
All Genes
DE Genes
Episode
mcan
robust
mcan
robust
40
0.9177 +: 0.0015
0.9329 1 0.0008
0.7211 :: 0.01
0.9048 : 0.0040
80
0.9193 = 0.0019
0.9337 + 0.0008
0.7339 + 0.0149
0.9076 + 0.0943
120
0.9178 + 0.0037
0.9340 + 0.0008
0.7234 + 0.0175
0.9104
+ 0.0038
160
0.9191 0.0009
0.9356 + 0.0010
0.7343 0.0079
0.9175 + 0.0034
These runs reflects that the robust estimator able produce
accurate estimation of the covariate-stratified marginal treatment effect E,, = e, T" = a) with tigher confidence bound.
C.2 DERIVATION
By Van der Vaart 2000), derive the efficient influence function of (p)) and thus provides mean for asymptotically efficient estimation:
Theorem 2. Suppose D Ω - E follows causal structure defined by the Bayesian network in Figure I where the counterfactual conditional distribution p(Y', T Z X) is identical to that of its factual counterpart pY, Z, X). Then (p) has the following efficient influence function:
{pp) 1(X = c, T = P(X,T)
(Y - E, (.))) + I(X = c) (E,, | T' = a] = -)).
(23)
p(X)
14

---

Published as a conference paper at ICLR 2023
Proof. Following Van der Vaart (2000), we define path ( p(2)(1 + cS(D) den- sity p of 2 as a submodel that passes through p at e 0 in the direction of the score S = de log [pe (2)) c-0
Let = (G, X) and minuscule of a variable denote the value it takes. By Levy
2019), have
d
v(pe)
Rp P.
(24)
de
c=0
de
c=
d
= o c)]
(25)
de c=0 y/==
/ V
e=0 [pe(0'22, = ()p = c)
(dominated convergence)
(26)
de
1/12
/ = 6)
Pely/ 12, ",,T==)) a)
(27)
de e==l
v/2
d
d'plu'12,T' = a)
Pc = c)
(28)
1/
de c=
I( c,t' = a) plx,t)
tlu,,ll,,,,0//000 2/2, 112 x)(()xx)) de de Pe('(iztt)
p(a,t)
1(x = c) plx)
p(x) yf p(y 12, -a) _ pe(xxx)
+
(29)
1(I = C,t = a)
y/ (54)
p(E,()
1(x cc)
+
(30)
V/ p(z)
T(x c, 1. a)
v, /
p(x,t)
I(x =c)
+
=((= - (31)
M/ p(x)
1(x c,t c,t= up(d)
(dd)
p(x,1)
E[S(D) x, t' plu' 2, x,,,, 1(x 1x= -x= a)
up(ul2,
p(z,t)
E[S(D) x 1(2(2 1(x c) 1 //lly//a 12, =
+
p(z)
I(x 1(==c) v'p(y'l 1,,' a)p(z, x)
E[S(D)
(32)
(0x)
1. S(d)
1(x ((==cl-=0) (y E (Y,,
p(r,t)
+ T(x = c)
(E 22 T'=
= a] v) p(d)
p(d)
(33)
p(x)
by assumptions of Theorem 2 and factorization according Figure 1 Hence
= (S(D), -(x==, -IX=cTT= (Y - E, (Y) D(X,T)
(pp)
de
+ 1(x= c) (E, IY' T" = al 4)
(34)
p(X)
('(()E)
and we have JZ I(X=c,T=o) ((XTT)
(Y Ep (,]) + 1[x=c) p(X)
EE Y Z, v.)
15

---

Published as a conference paper at ICLR 2023
By Theorem 2 propose the following estimator that is asymptotically efficient among regular estimators under some regularity conditions (Van Der Laan & Rubin 2006):
L -
I = a. Xx c) YY -po (X' Zr T% = a)
+ x Epp (Y' (Y' 22 Th =a=
(35)
p(
where (Y₁, Xg, T3) the observed variables of the k-th individual and Zg,o qo(Yx: G. Xg. Tk); X(] is an estimation of the propensity score and p(X) is an estimation of the density X. In the context of this work, X and T are discrete, hence (XT and (XX) can be estimated by the empirical density pr. X) and (x))). The above estimator then reduces to:
Pa.e
00.
1
Σ
{1e -Ep Ey ZY. TA a)}
Na.e le lala
Epa (X'1 T = =a}}
(36)
The k=1=
where (ley..., ne) are the indices of the observations having X c. and (la,e na,e) are the indices of the observations having both = a and X = c.
D COMPLEXITY ANALYSIS
Overall, the time complexity of VCI compared to generic framework like VAE (or the time com- plexity of graph compared to a generic GNN framework like GLUE (Cao Gao 2022) does not over-scale on any factor of any parameter to put simply, the workflow of VCI is just twice the forward passes 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 graph to VCI, graph 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 MLP layers with O(rd) number of operations and 4 layers of (dd) number of operations. d is the number of hidden neurons.
VCI has the same layer sizes: MLP layers with (rd)) number of operations and 4 layers of Odf) number of operations. But every layer is forward passed twice.
graphVCI has GNN layer with Orrdd + pdyr²) number of operations. v is number of gene features, dy is number of hidden neurons for GNN (usually a lot less than d since & r) and is the sparsity of GRN (number of connections divided by r2, usually around 1%), and the following layers which are forward passed twice: MLP layer and dot- product operation each with (rd) number of operations, MLP layer with ((rdy(dy+or) number of operations and 3 MLP layers with (dd) number of operations.
So the terms to be concerned compared to AE and VCI are O(red)), Oppdrr) and O(dd,d). Since p is around 1% and r is 2000 in our experiment, pd,r² 2 xs 20d,r hence the second term comparably smaller than the other two terms. Therefore, as long as dy is set to be reasonably small compared to d, the graph approach reasonably scaled compared to pure MLP approaches. We note that this is a limitation of ours and GNN approaches in general: if there is 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, has to be higher), GNN methods does not scale favorably compared to MLP methods.
E EXPERIMENTS DETAILS
E.
GENES
To evaluate the predictions the genes that substantially affected by the perturbations, select sets of 50 piterentially-xxprsssete genes associated with each perturbation and separately report performance these genes. The same procedure carried out by Lotfollahi et al. (2021a).
16

---

Published as a conference paper at ICLR 2023
E.2 OUT-OF-DISTRIBUTIO SELECTIONS
We randomly select a covariate category (e.g. cell type) and hold out all cells in this category that received one of the twenty perturbations whose effects the hardest predict. We 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-predie perturbations. This is the same procedure carried out by Lotfollahi al. (2021a).
E.3 TRAINING SETUP
The data excluding the OOD set are split into training and validation set with four-to-one ratio. A few additional cell attributes available each dataset are selected as cell covariates. For Sciplex, they 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 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'.
17


Reviews:
[{"review_id": "gekRNiaVTp", "reviewer": "Reviewer_TW7F", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper proposes a graph variational Bayesian causal inference framework to predict a cell's gene expression under counterfactual perturbations, gene regulatory networks were used to aid the individualized cellular response predictions. A robust estimator for the asymptotically efficient estimation of the marginal perturbation effect was also proposed. Empirical studies showed its superiority over state-of-the-art methods.", "summary_of_the_review": "Overall, I like the idea of this paper, and it seems theoretically sound. My only concern is the performance compared to other methods is not significant, and it lacks model complexity study, so I don't know whether the performance of the model outweighs its computational cost.", "strength_and_weaknesses": "Strengths: 1. Paper is well written and easy to follow.\n                  2. The idea of combining graph representation learning with variational bayesian causal inference is interesting, and the method is sound.\n\nWeaknesses: 1. The performance is not significant compared to other competing methods in Table 1.\n                       2. There is no model complexity study.", "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": "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": "WwktdWz2Hs", "reviewer": "Reviewer_kcXz", "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 propose a novel variational causal inference model for constructing the gene expression counts in cells after counterfactual perturbations that relies on information from two sources: individual features embedded in the outcome Y, and response distributions of similar individuals that did indeed receive the treatment. The method is aided by an additional method for updating the connections of the GRN prior to model training. The authors compare this method to known methods and find that theirs produces state-of-the-art R^2 in out-of-distribution predictions.", "summary_of_the_review": "The authors do a good job of motivating the variational causal inference formulation, but not the particulars of the graph latents. While some of the results show clear superiority over existing methods, there are overlapping error bars in Table 1, and Figure 5 seems not to show what is claimed. Clarifying these, as well as better motivating why the particular method for generating graph latents, would make this a stronger paper.", "strength_and_weaknesses": "Strengths:\n- Novel methods leading to state-of-the-art results compared with other models trained for generating individual counterfactual outcomes\n- Paced introduction to help understand the method\n- Excellent figures and legends\n- Makes use of GRNs in novel way\n- Introduces new CROP-seq dataset\n\nWeaknesses:\n- Contains a lot of complex notation and indices, some of which is not intuitive (e.g., indices of graph latents)\n- It is not clear that the manner in which the GRN is incorporated is necessary, and it is reasonable to expect that a simpler approach could yield equivalent results.\n- In Figure 5, the predicted distributions do not often look like the ground-truth distributions. Moreover, many of the ground-truth distributions look like control distributions, and so would be difficult to distinguish from the background anyway.", "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": "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": "ST3ikjI3QjE", "reviewer": "Reviewer_FrCx", "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 authors aim to tackle the problem of predicting the effect on gene expression of cellular perturbations such as gene knock-out. The authors build upon Variational Causal Inference (Wu et al 2022b). VCI takes as input a set of features X, given treatment T (e,g, drug treatment or knockout of a particular gene) and outcome Y (RNA-seq following treatment). It aims to predict, for some alternative treatment T', the corresponding outcome Y'. The authors extend VCI by additionally taking as input a gene regulatory network (GCN). The authors evaluate the method by holding out several perturbations for a particular cell type and using their method to predict the held-out values. ", "summary_of_the_review": "See above.", "strength_and_weaknesses": "Overall, the paper is strong. The text and figures are understandable. The results have significance both in terms of methodology and application. The experimental results are good.\n\nNovelty is moderate because most of the paper is an extension of (Wu et al 2022b). However, the addition of the GCN to VCI is novel.\n\nThere is no description of how the authors go from experimental RNA-seq data to values of X and Y, including in Appendix C, which explains some other parts of the experiments. In particular, I couldn't figure out what X is for their experiments. I think they followed (Lotfollahi et al 2021a) in this, but it was hard to figure out what the authors of this paper did even when examining Lotfollahi et al 2021a. My best guess is that X is a discrete indicator of cell type.\n\nThe setup is rather specific to the application to RNA-seq because it requires that each individual has an information-rich outcome. It doesn't seem obvious that the same setup would work in a setting such as clinical drug response, where each individual has information-rich features but simple outcomes (e.g. symptoms resolved or not resolved). \n\nIt is unclear whether the revised GCN produced by the method of Sec 2.3 is improved according to the desirable criteria for GCNs. There are many existing methods for producing GCNs, so it would be helpful to argue if this improvement method is a better approach for producing a GCN in general, or if it is a preprocessing step useful only for input to a NN method like this. ", "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": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "BbLEW7P-Qz", "reviewer": "Reviewer_qD32", "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": "In this paper, the authors proposed a Graph Variational Causal Inference (GraphVCI) network to estimate gene response under perturbations.\nTo this end, the component of (hypothetical) counterfactual perturbations is integrated into the GraphVCI network to assist the reconstruction training of the proposed Graph autoencoder.  On three commonly used datasets and a soon-to-be open-sourced dataset, the proposed model demonstrated clear improvement compared to several baselines. ", "summary_of_the_review": "Overall, I think this is an interesting paper that brings some new insights to the analysis of the cell responses to different perturbations, which can potentially achieve a good impact on the ML for the biology community.", "strength_and_weaknesses": "Strength: Overall, the paper delivers clear explanations regarding 1. the motivation for cell response under different perturbations, 2. model architectures 3. experimental results discussions. To validate the proposed GraphVCI method, the authors run experiments on several publicly available datasets, which I consider the main experiments are thorough. Besides, the codes are supplemented in the submission, this can help or guide interested readers to reproduce the results. \n\nWeakness: I think the ablation studies such as excluding the counterfactual-related losses and features should be performed in the paper,\nthis can disentangle the contributions that come from different network components. Second, I don't know why the predicted distributional shift presented in Figure 5 matches the true shift, and I feel that they look quite different from each other. Could the authors explain this to me? In the end, the authors introduced 'the causal relation diagram' in Figure 1. I think it would be important to compare it with the standard causal graph definition such as the structural causal model (SCM). ", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster