Learning Identifiable Factorized Causal Representations of Cellular Responses
Authors: Haiyi Mao
Romain Lopez
Kai Liu
Jan-Christian Huetter
David Richmond
Panayiotis V. Benos
Lin Qiu
Venue: NeurIPS 2024 poster (2024)
Keywords: Single Cell
Computational Biology
Causality
Primary area: machine_learning_for_healthcare

Abstract: The study of cells and their responses to genetic or chemical perturbations promises to accelerate the discovery of therapeutics targets. However, designing adequate and insightful models for such data is difficult because the response of a cell to perturbations essentially depends on contextual covariates (e.g., genetic background or type of the cell). There is therefore a need for models that can identify interactions between drugs and contextual covariates. This is crucial for discovering therapeutics targets, as such interactions may reveal drugs that affect certain cell types but not others.
We tackle this problem with a novel Factorized Causal Representation (FCR) learning method, an identifiable deep generative model that reveals causal structure in single-cell perturbation data from several cell lines. FCR learns multiple cellular representations that are disentangled, comprised of covariate-specific (Z_x), treatment-specific (Z_t) and interaction-specific (Z_tx) representations. Based on recent advances of non-linear ICA theory, we prove the component-wise identifiability of Z_tx and block-wise identifiability of Z_t and Z_x. Then, we present our implementation of FCR, and empirically demonstrate that FCR outperforms state-of-the-art baselines in various tasks across four single-cell datasets.

Learning Identifiable Factorized Causal Representations of Cellular Responses
Haiyi Mad 1,2
Romain Lopez" 1,3
Kai Liu
Jan- Can-Christian Huetter
David Richmond
Panayiotis V. B
Lin Qiu 4 University of Florida
'Genentech
University of Pittsburgh
Stanford University
Abstract
The study cells and their responses to genetic or chemical perturbations promises to accelerate the discovery of therapeutic targets. However, designing adequate and insightful models for such data difficult because the response of cell to perturbations essentially depends its biological context (e.g., genetic back- ground or cell type). For example, while discovering therapeutic targets, one may want to enrich for drugs that specifically target certain cell type. This challenge emphasizes the need for methods that explicitly take into account po- tential interactions between drugs and contexts. Towards this goal, propose a novel Factorized Causal Representation (FCR) learning method that reveals causal structure in single-cell perturbation data from several cell lines. Based on the framework of identifiable deep generative models, FCR learns multiple cellular representations that disentangled, comprised of covariate-specific (z.,), treatment-speciie (z.), and interaction-specifi (Ztz) blocks. Based on recent ad- vances in non-linear ICA theory, we prove the component- wiie identifiability of 2tr and block-wise identifiability of and Zy. Then, we present our implementation of FCR, and empirically demonstrate that it outperforms otate-of-the-art baselines in various tasks across four single-cell datasets. The code is available on GitHub (https://github. www com//eenettechfarr)
http://gttuub.cco//eeeeetttci
Introduction
The recent experimental capabilities reached by single-cell perturbation technologies open up new opportunities for characterizing cellular behaviors (Dixit al. (016). For example, high-throughput screening of chemical and genetic perturbations highlighted vulnerabilities in numerous cancer cell lines (McFarland al. [2020). Identifying these vulnerabilities is crucial for pinpointing therapeutic targets, facilitating drug discovery, and furthering our understanding of gene functions (Srivatsan
al. 2020).
The analysis of perturbation data involves modeling how cells respond to diverse treatments across biological contexts. This task is challenging for two main reasons. First, outcomes of perturbation experiments are quantified using single-cell RNA sequencing (scRNA-seq) technologies, whose measurements are noisy as well as high-dimensional (Grün al. 2014). Second, cellular contexts are difficult to comprehensively model because they are extremely variable, encompassing cell types, tissue of origin, and genetic background (Wagner et al.) 2016). This highlights the need consider interaction effects between treatments and contextual covariates while modeling gene expression outcomes (Zapatero et al. 2023).
'This work was conducted during an internship at Genentech. Corresponding author. Email: lin.qiu.stats @ gmail.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

To address these challenges, several computational methods have been developed. Novel approaches based on causal representation learning provide better mechanistic interpretation of single-cell
perturbation data (Lopez et al. 2023 Bereket and Karaletsos 2023: Zhang et al.) 2023). These
methods belong to the family of identifialle deep generative models (Khemakhem et al. 2020
Lachapelle et al. 2022 Zheng et al. 2022) and therefore offer, to some extent, theoretical guarantees but remain limited to the analysis of data from a single cellular context. Another set of studies model
cells across multiple contexts using latent linear additive models (Hetzel et al. 2022 Lotfollahi et al., 2023). However, due to their additive assumption, these models fail to characterize interactions between treatments and biological contexts.
To address these limitations, we introduce the Factorized Causal Representation (FCR) learning framework. This identifiable deep generative model learns representations of cells that take the form of three disentangled blocks, specific to treatments, biological contexts, and their interactions, respectively. We first present the proposed generative model and then provide set of sufficient condi- tions for its identifiability, extending the work of Khemakhem et al. (2020) to the case of interacting covariates. We then present implementation of our FCR method that builds upon the variational auto-encoder framework (Kingma and Welling 2014) as well as adversarial regularization (Ganin al. 2016). We demonstrate that FCR not only effectively identifies interactions but also surpasses -tate-of-the-arr methods in various tasks across four single-cell datasets.
2 Related Work
Learning Representations of Cellular Responses Learning representations of single-cell data is powerful framework, with demonstrated impact in tasks such as imputation (Lopez al. (0118).
clustering (Trapnell et al. 20144 Zhu et al., 2019 Alquicira-Hernannee
et al. 2019). and integration
across modalities (Gayoso et al. 2022). Recent advancements have largely improved our capacity to
predict cellular responses to drug treatments (Lotfollahi et | 2019 Rampasek et al. 2019 Lotfollahi
itt al. [2021] Lopez et al. 2023 Bunne et al. 2023 Zapatero al. 2023). Lotfollahi et al. (2023) and
Hetzel et al. (2022) proposed generative models that additively combine treatment embeddings and biological context embeddings within a latent space. [Wu et al. (202)) cast the prediction problem as a counterfactual inference problem. Despite these advancements, existing methods fail to address how treatments may preferentially impact specific cell types, critical point for understanding the effects of drugs on biological systems.
Identifiable non-linear Independent Component Analysis models A field where disentanglement most important is non- non-linear Independent Component Analysis (non-lineat ICA) (Hyvärinen and Pajunen 1999). where information from set of latent variables is mixed through a non-linear encoding function. It has long been known that such models (i.e.. either the encoding function, or the sources) are non-identifiable without further assumption. However, some recent works (Hyvärinen) ct al, 2019, Lachapelle et al., 2022]] Zheng et al. 2022) proved identifiability was possible in a non-stationary regime. Often, this assumption takes the form of conditional independence of the latent variables given some auxiliary (observed) variables. Notably, the iVAE framework from Khemakhem et al. (2020) offers disentanglement guarantees within this conditional VAE setup. Our work extends the theory of Khemakhem et al. (2020) to prove identifiability of the interaction terms between multiple such auxiliary variables.
This approach differs significantly from specific models the Variational Autoencoder (VAE) literature (Kingma and Welling, 2014) such as the bet (Higgins et al., 2016) and factorVAE (Kim and Mnih (018)) that both address disentanglement learning. Indeed, these latter models lack theoretical foundation regarding the identifiability of their parameters or latent variables (Esmacili ctt al. 2019 Chen and Grosse 2018).
3 Preliminaries
Single-cell perturbation experiments characterize causes and effects at the cellular and molecular levels. Our objective is to disentangle the contributions of treatments, cellular covariates, and their interactions to model the effects of perturbations. Let x E C IRd be a vector of covariates representing intrinsic attributes of single cells, such cell type, tissue of origin, or patient information.

---

Let E T C Rº represent the treatment or intervention applied to single cells and let E y S 1 denote the gene expression levels (outcome).
3.1 Generative Model
We introduce a low-dimensional latent vector 2 C Rn that encodes cellular states after treatment t and in biological context x. We assume a block structure for
x
ny + Rez Rh. Here, z, captures the effects of contextual covariates x, represents the direct effects of the treatment t, and Zte encodes the interaction effects between both terms.
[Z, Ztz, Zz] with dimension
Ziz
More precisely, the generative model is specified as fol- lows. Latent representation Ze is generated from x ac- cording to distribution 22 (2 x), is gen- Pz.jt(Zz t),
erated from t following distribution
Data generating process: shaded nodes denote observed variables. Empty nodes denote latent variables.
Figure I:
and Ze from both and following distribution Ztæ
Prcs xt,x (Ztz x). The gene expression outcome vector y is then deterministically generated y =9(Z., Ztz, zt),
where g is smooth mixing function. graphical representation of this generative model appears in Figure For the control group (no treatment), the outcome is noted as y", and the representation as
2/. z(].
Learning each element in this triplet of latent variables is sound approach for unravelling interaction effects. Indeed, zu captures covariate-specife patterns that are invariant with respect the perturba- tions, while also capturing the essential biological attributes tied to the cellular covariates. 2, captures the intrinsic effects of the treatments, irrespective of the covariates. Zι unravels the interactions that a treatment could have with specific covariates. This last representation captures the nuanced manner in which distinct cell types, tissues, or patient groups react treatments, reflecting the diversity and specificity of biological responses.
We note that model does not take into account observation noise, since is deterministic function in our assumption. However, our theory may be readily extended to incorporate Gaussian observation noise (Khemakhem et al. 2020), well counting observation noise (Lopez et al. 2024
3.2 Model Identifiability
We now introduce the definitions for the different classes of disentangled models that will appear later in the manuscript. Analogous definitions appear in previous work from Von Kügelgen al. (2021) and Kong et al. (2022). Throughout this section, z E Z denotes a (random) latent vector and y denotes observed vector. g : Z y is unknown mixing function such that y g(z).
Definition 3.1 (Component-wise Identifiaabllit).. We say that latent vector is identifiable from data if for any other latent vector å such that g(ū) and g(z) are equal in distribution, z and á are equal up permutation of the indices, and deformation of each component by invertible scalar functions. More precisely. there exists a permutation π of [I....,,,, and invertible scalar functions hy such that for all; j nI..
n2,,
8j = h)
where 2y and (T)) the th components of and 2 respectively.
Definition 3.2 (Block-wise Identifiabiliy). For 1 S m1 < n2 S n, we denote as Z[nymal E R the block of z from index nı 122. We say that latent vector [nn is block-identifablle from data if for any other latent vector à such that g(2) and gzz) are equal in distribution, Z[mrona] and [m are equal up to an invertible function h:
[|p = h( (i))))))
(2)
where (/m is the corresponding block in the estimated vector 2.
3

---

Identifiability Results
One strong advantage of the FCR framework is that it comes with strong theoretical guarantees concerning the disentanglement of its factorized latent space. We first prove the component identifiability the interaction variable Ztz under the assumption of sufficient experimental variabil- ity (Khemakhem et al. (220)) (Section 4.1 Then, we demonstrate the block-ideniifabbliiy of and Ze by exploiting their invariance with respect to x and t, respectively (Section These guarantees important, they ensure that the obtained latent variables will have desirable semantics. For example, given our theoretical results, the obtained interaction embedding žte must be reflective of the ground-trui truth interactions Ztz only, and not of any of the other latent variables.
4.1 Component-wise identifiability of 2i.
Our proof relies on three technical assumptions. Two are classical from the nonlinear ICA literature, and the last relates to the observability of complementary set of experiments for identifiability of interactions.
Assumption 4.1. The probability density function the prior distribution for the latent variables is smooth and positive, i.e. Pz)t,x(Z t, > 0 for all (z, t, 2 x T X.
Assumption 4.2. The components of are conditionally independent given and x.
Assumption 4.3. (Experimental Sufficiency) There exist at least 2n1 + 1 distinct values for the vector [t. x in the experimental design. One such setting can be referred to as control condition and is noted as fto, xo). Additionally, for any non-control environment (t,, x,) for E {l...., 2n we assume there always exist corresponding switched experiments under the settings (to, x.) and Xo).
Assumption 4.3 is novel and essentially mandates that we conduct a sufficient number of experiments with cross-referenced covariates and treatments. This ensures that we can observe the specific drug response related to each covariate, and is often how such experiments are designed in practice.
Theorem 4.4. Let us first define v(Ztz: t, x) as the vector:
v(zizi t,x) = 04 12n,+1,
Uq new (Zng t, x)
Ozngt1
Oingtnes
8² gn (1n
Qm t (Zn +hi t, x,
02h.+1
Ooy tho t
where qi denotes the logarithm of probability density Pt. for component Zi in addition the assumplions 4.1 we assume that the 280 vectors
{v(zen ti x.) + v(zez, to, xo) v(Zta, to, x:) v(Zez, Li, 201
are linearly independent then Ztz. is component-wise identifiable.
The proof appears in Appendix A
Theorem 4.4 extends the concept of linear independence found in
nonlinear ICA (Khemakhem et al.) 2020). Unlike the original theory, where auxiliary variables must induce sufficient variations of the latent variables, we are confronted with a case where treatments and contexts must have sufficient variability in combination. For example, when v is linear with respect to both and x the vector of interest becomes the null vector. In this trivial case, the theorem's 's assumption is never satisfied (Zt indeed has no purpose in that particular model). However, under rich class of model with complex interaction patterns, our model will be able to infer informative latent variables.
4.2 Block-wise identifiability of z, and z
To prove the block-wise identifiability of z, and z,, we exploit their invariance properties: 2, remains unchanged across different values of x, while z, stable across variations in t. This invariance allows to distinguish these blocks from the interaction terms Zur- Additionally, because this invariance reflects latent features' stability despite perturbations, its utilization can enable deeper biological insights and interpretability. For instance, 2, might represent stable cellular characteristics

---

that persist across different treatments, while Z, could capture consistent treatment effects across various cell types.
We now state our main identifiability result for Z and 2t-
Theorem 4.5. We follow Assumptions 4.1
4.2 4.3
and the from Theorem 4.4 We note S(Z)
the set of subsets $ - 2 of that satisty the following two conditions:
(i) S has nonzero probability measure, i.e. P(ZE (
for any t' T and
x E
anrany (ii) S cannol be expressed as Aa, Ziz Z1 for any An, C By Zz Zta An, for
any Az, C Zt.
We have the following identifiability result. for all S E S(Z), there exists (t1, (2) T X T and x X such that
[ Petexx(z t1, x)dz + [ Pzttxx(z t2 x)d2,
(4)
JEES and there also exists (x1, E x X and t T such that
[ Palex((z t, + [ Pett,x X2222
(5)
then 2, and are block-wise identifiable.
The proof appears in Appendix A and is adapted from Kong et al. (0022) Our main assumption that the conditional distribution Palt.x(2 t, undergoes substantive changes when spanning different treatments t and covariates When treatments differ markedly from each other in their
mechanisms and effects, the probability distributions of the latent variables conditioned on these treatments unlikely to be identical.
5 Methodology
We now propose a tangible implementation of our method, termed Factorized Causal Representation (FCR) learning. Our approach has three components:
1. A variational inference approach to estimate the representations from FCR model. Our model and inference architecture specifically designed to learn disentangled representa- tions 2g, Zt 2, (Section 5.1
2. A regularization method that enforces independence between 2g and t. and encourages variability of 2, with respect (Section 5.2).
3. A second regularization technique to ensure that the conditional independence properties 2rr 11 Ztr and z, Il. Ztz are satisfied (Section 5.3).
The main computational structure of the model is illustrated as schematic in Figure
5.1 Model Specitication and Variational Inference
Model Specification
We parameterize the probability distributions as follows:
2g x) = Normal (orma
(6)
p(z t) = Normal (f))
(7)
p(zzz t, x) = Normal (f" x), ft (t,
(8)
where all the above functions are parameterized by neural networks.
To prevent 2tz from having trivial dependencies with respect to and x, we explicitly encourage its prior to capture interactions between x and by designing the functions ft. and ftr to of the form:
ft. = f (kx (x) o k (t))
(9)
= fo O (t))),
(10)
5

---

(a)
(c)
pize tt
(x(t)
p(z
NLP
xx(x
p Zz x)
pl Zee tt, x)
KL devergencc
q y) (44, y) mia,a... v?
Zar It,
r..
fdiet
Z
y
qiz y) g(z,,,,,, It, q(z
Z
fditx
sin
Boods
2.
O. yd) 10,x, y/ 8
y
21
fat
Zin
Figure The illustration of FCR models. (a) is the component for p(zz x).p(z, t) and P(Zte x. (b) is the component to estimate g( x), ((z t,y) and g(zt.= x. Note
that 0 indicates representing the control samples. (c) computational diagrams estimate the Kullback-Leibler divergences, causal structure regularization and permutation discriminators.
where ki (x) and kg (t) represent the embeddings for the cellular covariates and treatments, respec- tively, while C denotes the Hadamard product.
Variational Inference Because the posterior distribution of the latent variables are intractable, we use the variational autoencoder framework (Kingma and Welling, 2014) to jointly learn the model's parameters and an approximation the posterior, following the approach used in previous causal representation learning work (Khemakhem et al. (020). We consider the following mean-field variational approximation to the posterior distribution:
go (Zzr Ze: Zar x,,, (z(Zz x. x,y)ga(Z; t, (2 tt,,y).
(11)
Following the graphical model from Figure I the Evidence Lower Bound (ELBO) is derived as:
LELO = 2. log Po Bx 21, 2tx)
Dxillo(zz xxy) (z. x))
(12)
DKLLqo(Zz t,y) t))
DKL(qo(Ziz ylpo t,))),
where and & denote the parameters of the generative model and the inference networks, respectively, DEL denotes the Kullback-Leibler divergence between two probability distributions. For simplicity, we omit & and well as script notations in the following sections, wherever appropriate. The
derivation of the ELBO appears Appendix B
The variational distributions defined in Equation are parameterized as follows:
qz, x, y) = (x
(13)
q(z, t,y) = (f (t,y)
(14)
q(zt.r t, x, == Normal(g xt, (t., (t,,,t,))))
(15)
where all the functions introduced above parameterized by neural networks.
5.2 Causal Structure Regularization
exploit both the variability z; and the invariance 2, when comparing control and treated groups that share the same covariates. Specifically, goal is to enforee the resemblance between Ze and z while reducing the congruence of z, and z,. Towards this end, we first add the following score regularizer,
Lrim = Ey(r. m ( (x)) [sim (z,, 2 sim (zz, 20)
(16)

---

where sim(-) denotes the cosine similarity. Second, we introduce a classifier fa to predict the from Ze, Ztn. and |z]. zzz as follows,
treatments
t= Sa (I&,, Zeal [2] =lx)))
(17)
The predicted treatment probability vector t is then used for the computation of a cross-entropy loss
Let
gg. (x.))))))))) 1119 log(E)]-
(18)
5.3 Permutation Discriminators
We want to ensure the conditions of Assumption 4.2 and Theorem 4.4 specifically that Ztx TL 2, and Ziæ IL Ze x. Towards this goal, between exchangeability and conditional independence.
the following proposition, establishing connection
Proposition 5.1. (Bellot and van der Schaar 2019) Ler X,Y and Z be three random variables. Under the assumption of X II Y Z, have the samples (X,, Y,, Z and permuted samples (Xx()) Y,, t4 with permulation function The corresponding statistics Pt of (Xi, Y:) and Pa(s) of (X (9), Y are exchangeable.
This proposition states that permutation will not change the independence between two conditionally independent random variables. We therefore propose to use permutation discriminators for Zy Ztz
and V2, First, we initially permute within the triplet (z.e) 2it (i) x() x:) to yield
a()) x(() x,) Then, we train binary classifier (the discriminator) to predict whether samples have been permuted not. We denote the permutation label I and predict as
(zg Zir
i= faise (Zz, Zta x).
(19)
where žιz could be permuted or con-permuted Zız samples. If z5 and Ztr are indeed independent given x. the discriminator should be unable to distinguish between the permuted and the original samples. For each discriminator, we add regularization term that consists of the cross-entropy loss
Laisæ -Ex. -x..xxxx.... (2.. (x,t) [log)
(20)
We proceed similarly make sure that 2, and Zez independent conditionally on t.
5.4 Objective function
Finally, we specify the overall loss for our model
Liotll
CELBO + Ci + wrCa - wa( WadSsss Ldis.).
(21)
where w1, W2, > O are hyperparameters. To concurenntyy train both the representations and the discriminators, employ an adversarial training approach as follows,
max min Lixaa-
(22)
fang San 8.0.J=
The training procedure and the procedure used for hyperparameters selection are detailed in Ap- pendix C and respectively.
6 Experiments
In this section, we are pursuing three primary objectives. First, we seek to validate the FCR method's proficiency capturing the designated causal structure within the latent space through both clustering analysis (Section 6. and statistical testing of independence (Section 6.2). respectively. Second, we evaluate the method's efficacy predict single-cell level conditional cellular responses (Section 6.3) We note that the current implementation of FCR does not make use of general embeddings for t or and for that reason we do not perform experiments to predict cellular responses unseen treatments and cell types (covariates).
7

---

(8)
(b)
(6)
1.0
0.40
0.23
+1
0.20
0.30
0.5
0.20
0.10
+
6.10
0.5
2000
factriide
betovids
1966
Che
VC
fedevsing Ivrtoone
(actaritl betavis
aethops
rwthtds
aethods
Figure 3: Clustering results for the sciPlex dataset. Normalized Mutual Information (NMI) values for clustering based on: (a) covariates x; (b) combined covariates and treatments [x; t]: (c) treatments t.
Datasets To evaluate the efficacy and robustness of the FCR method, we conducted our study on four real single-cell perturbation datasets (Appendix [D. The first of these the sciPlex dataset (Srivatsan et al. 2020). which provides insights into the impact of several HDAC (Histone Deacetylase) inhibitors on a total of 11,755 cells from three distinct cell lines: A549, K562, and MCF7. Each of these cell lines was subjected to treatment in two independent replicate experiments, using five different drug dosages: nM (control), 10 nM, 100 nM, μM, and 10 μM. The subsequent three datasets soureed from (McFarland al. 2020 which executed several large-scale experiments in varied settings. The multiPlex- Tram dataset contains 13,713 cells from 24 cell lines, treated with Trametinib and DMSO control over durations of 3, 6. 12, 24, and 48 hours. The multiPlex-7 dataset spans 61, ,,552 cells across 97 cell lines, subjected to seven different treatments. Finally, the multiPlex-9 dataset incorporates 19,524 cells from 24 cell lines, undergoing series of nine treatments.
Baselines We benchmarked our method against several established representation learning methods: (1) VI (Lopez et al. 2018)) (2) i'AE (Khemakhem et al. 2020), (3) ,SVAE (Higgins et al. 2016).
(4) factor (Kim and Mnih 2018), (5) VCI (Wu et al. 2023), (6) CPA (Lotfollahi et al. 2023).
(7) scGEN (Lotfollahi et al. 2019h (8) sVAE (Lopez et al. 2023 (9) CINEMA-OT (Dong et al.
2023). For the clustering analysis, we employed all the inferred latent variables from each baseline method. For the conditional independence test, we selected random subset of latent variables for each baseline matching the number of latent variables of FCR. We then tested each subset and repeated the process total of ten times for each baseline using different random subsets to yield the best results. Specifically, CINEMA-OT and seGEN address only binary treatments/petturaations so they only considered in the cellular response predictions tasks.
Results on additional dataset, simulation studies, ablation studies, data visualization, and biological interpretation of the latent variables appear in Appendix
6.1 Clustering Analysis on Covariates, Treatments, and Combined Features
We evaluated the performance of the obtained latent representations in capturing three key aspects: the cellular covariates the treatments t, and their interaction (x, t). To assess each latent representation, we applied clustering and compared the fidelity of the resulting cluster labels with the corresponding or (x, t) from the original data. This approach allowed us to gauge how well the latent representations preserved the underlying structure of the cellular covariates, treatments, and their interactions. We performed clustering using the Leiden algorithm (Traag et al. 2019). To assess the fidelity between two sets cluster labels, we employed the Normalized Mutual Information (NMI) metric (Kim et al. 2019). Higher NMI values indicate better alignment between the clustering results and the original data structure. We conducted this clustering and fidelity assessment on our model's latent variables Z,, Ziz, and Z₁, well
the variables obtained from baseline methods.
Our results highlight that z, has superior performance clustering on covariates x compared to all other available latent representations (Figure 3) Similarly, 2, performs best for clustering on treatments Finally, Ztx outperforms all other methods when clustering jointly on and x, showing that it faithfully represents the combined features of both x and t. Specifically, for the seiPlex datasets,

---

(8)
(b)
(c)
(d)
0.35
0.08
0.30
0.06
0.3
0.25
Poeluw
0.4
Poaluve
0.20
0.04
0.2
0.15
0.2
0.1
0.10
0.05
0.00
fectorin
svAe
FER
beavvk fectorvie
ce
swke
PDR
2 4,44
4+
Figure 4: Statistical Conditional Independence Testing Results (a) p-values for the conditional independence test of be TT x. The red dashed line indicates the 0.05 level. (b) p-values for the conditional independence test of 2 IL x t. (c) p-values for the conditional independence tests Ze TL Ziz x and z Ztz t. (d) HSIC values for assessment of marginal independence of 2₂ with x, t and random numbers (R); well 2, with x, t and random numbers (R).
clustering on x yields better results than clustering based on both (t, x) or on solely t. This can be attributed to the HDAC inhibitors exhibiting minimal impact distinct cell lines until maximal concentration of 10 μM was reached. We report similar results for the other datasets Appendix F Taken together, these results suggest that FCR effectively learns disentangled representations across different datasets.
6.2 Statistical Conditional Independence Testing
We validated the disentanglement of our latent representations via conditional independence testing, implemented as the Kernel Conditional Independence (KCI) (Zhang et al. 2012) tests. Our investiga- tions focused on the following relationships: (1) 2 IL t x: (2) z, II X t; (3) z; TT Ztr t; (4) 2r Ztz x. In conjunction with these tests, we also employed the Hilbert-Schmine Schmidt Independence Criterion (HSIC) (Gretton et al. 2005) to evaluate the (marginal) independence between: (a) 2; and t; (b) 2, and Our aim was two-fold. First, we wanted to assess whether the factors our latent space complied with the necessary conditional independence statements, corroborating well-captured causal structure. Second, we wanted to determine the dependence between our latent representations and their respective observed variables.
We focus our presentation of the results on the sciPlex dataset in the main text (results for the other datasets imply similar conclusions and appear in Appendix [.44). We first examined the results of testing for conditional independence statements zu 1L t and z, 1L x (Figure Ab). In this experiment, all the baseline methods produced p-values smaller than 0.05. This implies rejection of the null hypothesis (conditional independence) for the baseline methods, and suggests that their representations failed to maintain the desired conditional independence statements. Second, we examined the results of testing for conditional independence statements 2, 1L Ztz and IL t (Figure Interestingly, this assessment also quantifies the efficacy of our permutation discriminators. The observed p-values generally exceed 0.05, suggesting that the null hypothesis cannot be rejected. Finally, we use the HSIC to report estimates of mutual information (assessing for marginal independence). Low HSIC values suggest poor dependence between the pairs of random variables. We report the HSIC values for assessing z IL 2, JL t, Z IL R. well as z, IL 2; IL t. and 2 IL R. where R represents simulated random vectors (Figure 4 Contrasting our representations with randomly simulated vectors, we observe that both z, and x, well z, and have HSIC values far from zero, indicating a high dependence. The contrast in results between our approach and the baseline methods highlights FCR's 's nuanced capability in capturing and preserving causal structures.
6.3 Conditional Cellular Responses Prediction
Our analysis demonstrates that treatments often elicit covariate (cell line) specific responses. Conse- quently, accurately predicting outcomes for novel drugs or cell lines becomes challenging without careful consideration of the similarity in Ziz and how t interacts with x. Unlike previous literature,

---

Methods
Datasets
FCR (ours)
VCI
CPA
scGEN
SVAE
CINEMA-OT
sciPlex
00.02
0.86+0.03
0.86
0,56+0.00
0.84+0.02
0.51
multiPlex-Tram
0090
0.89
0.88
0.52+0.00
0.87+0.02
0.33+0.09
multiPle--7
0.83 #0.03
0.8 00.04
0.80
0.49
0.77+0.02
0.41
multiPlex-9
0.78
0.78
0.79 10.02
0.33 00.08
0.75
0.32
Table 1: The R² score of the conditional cellular responses prediction.
do not conduct experiments predict cellular responses to unseen treatments and cell types (covariates). This decision is based on extensive biological research showing that responses to covariates context-specifie (McFarland et al. 2020; Srivatsan al. (020). Without thoroughly examining the similarity of unseen treatments or cell types in the latent space, we cannot confidently predict cellular responses.
Nonetheless, our approach enables the prediction of cellular responses at the single-cell level. This paper focuses predicting cellular responses (expression of 2000 genes) in control cells subjected to drug treatments. is important to note that our comparative analysis is confined to CPA, VCI, they uniquely tailored for this task. We utilize FCR to extract
sVAE, scGEN and CINEMA-OF control's 's 12., z!.- z and corresponding experiments' Zg, Ztz; 21], then use the decoder to predict the gene expression level as y g Ztz; 2,). We measure the R score. From
results (Table D. we observe that FCR generally outperforms other baselines across the first three datasets. However, CPA performs the best the multiPlex-9 dataset. The primary reason for this is that the multiPlex-9 dataset has fewer covariate-specific responses (McFarland et al. 2020). Additionally, seGEN and CINEMA-OT, which are designed for binary perturbations, tend to underperform in these tasks.
Discussion
This paper aimed resolve a current challenge- how to disentangle single-cell level drug responses using latent variables into representations for covariates, treatments and contextual covariate-treatment interactions. To do so, we established a theoretically grounded framework for identifiability of such components Zta. and Expanding upon these theoretical foundations, we have developed the FCR algorithm to factorize the interactions between treatments and covariates. Looking ahead, our aim is to incorporate interpretable components into this framework. This enhancement will aid in pinpointing genes affected by Ze, Z, or Ztz- Such advancements are expected to significantly contribute the progress of precision medicine.
Acknowledgments
The authors would like extend gratitude to Gregory Barlow for his contribution the graphic design.
References
Alquicira-Hemaander Ji, N., and Powell (2019). scPred: accurate supervised method for cell-type classification from single-cell RNA-seq data. Genome Biology, 20:264.
Bellot, A. and van der Schaar, M. (2019). Conditional independence testing using generative adversarial networks. Advances in Neurul Information Processing Systems, 32.
Bereket, M. and Karaletsos, T. (2023). Modelling cellular perturbations with the sparse additive mechanism shift variational autoencoder. In Advances in Neural Information Processing Systems.
Bunne, C., Stark, S. G., Gut, G., Del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L.,, Krause, A., and Rätsch, (2023). Learning single-cell perturbation responses using neural optimal transport. Nature Methods, pages 1-10.
10

---

Chen, L. and Grosse, D. (2018). Isolating sources of disentanglement in VAEs. In Advances in Neural Information Processing Systems, pages 2615-2625.
Dixit, A., Parnas, O., Li, B., Chen, J., Fulco, C. Jenoy-Arnon, L., Marjanovic, N. D., Dionne, D., Burks, T., Raychowdhury, R., et al. (2016). Perturb-seq: dissecting molecular circuits with scalable single-cell RNA profiling of pooled genetic screens. cell, 167(():1853-186)
Dong, M., Wang, B., Wei, J., de Fonseca, A. H., Perry, C. J., Frey, A., Ouerghi, F., Foxman, F., Ishizuka, J. J.. Dhodapkar, R. M., et al. (2023). Causal identification of single-cell wew experimental perturbation effects with cinema-ot. Nature Methods, 20(1 I): 1769-1799.
Esmaeili, B., Wu, H., Jain, S., Bozkurt, A., Siddharth, N., Paige, B., Brooks, D. H., Dy, and de Meent, J.-W. (2019). Structured disentangled representations. In International Conference on Artificial Intelligence and Statistics, volume 89, pages 2525-2534.
Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H.,, Laviolette, F., Marchand, M., and Lempitsky, V. S. (2016). Domain-adversarial training of neural networks. Journal of Machine Learning Research, 17:1-35.
Gayoso, A., Lopez, R., Xing, G., Boyeau, P., Valiollah Pour Amiri, V., Hong. J., Wu, K., Jayasuriya, M., Mehlman, E., Langevin, M., et al. (2022). A Python library for probabilistic analysis of single-cell omies data. Nature Biotechnology, 40(2):163-166.
Gretton, A., Bousquet, O., Smola, A., and Schölkopf. B. (2005). Measuring statistical dependence with bilbert-schmide norms. In International Conference on Algorithmic Learning Theory, pages 63-77.
Grün, D., Kester, L., and Oudenaarden, A. (2014). Validation of noise models for single-cell transcriptomics. Nature Methods, 1166):637-640.
Hetzel, L., Boehm, S., Kilbertus, N., Günnemann, S., Lotfollahi, M., and Theis, J. (2022).
Predicting cellular responses to novel drug perturbations at single-cell resolution. Advances in Neural Information Processing Systems, 35:26711-2677
Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. (2016). beta- VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations.
Hyvärinen, A. and Pajunen, P. (1999). Nonlinear independent component analysis: Existence and uniqueness results. Neural Networks, 12(3):429
Hyvärinen, A., Sasaki, H., and Turner, R. (2019). Nonlinear ICA using auxiliary variables and gener- alized contrastive learning. International Conference on Artificial Intelligence and Statisties, pages 859-868.
Khemakhem, I., Kingma, D., Monti, R., and Hyvärinen, A. (2020). Variational autoencoders and nonlinear ICA: A unifying framework. In International Conference on Artificial Intelligence and Statistics, pages 2207-2217
Kim, H. and Mnih, A. (2018). Disentangling by factorising. In International Conference on Machine Learning, pages 2649-2658.
Kim, T., Chen, I. R.. Lin, Y., Wang, A. Y.-Y., Yang, J. H., and Yang, P. (2019). Impact of similarity metrics on single-cell RNA-seq data clustering. Briefings in bioinformatics, 2016):336
Kingma, D. P. and Welling, M. (2014). Auto-encoding variational Bayes. International Conference Learning Representations.
Klindt, D. A., Schott, L., Sharma, Y., Ustyuzhaninov, I., Brendel, W., Bethge, M., and Paiton, D. (2021). Towards nonlinear disentanglement in natural data with temporal sparse coding. In International Conference Learning Representations.
Kong, I.,, Xie, S., Yao, W., Zheng, Y., Chen, G., Stojanov, P., Akinwande, V., and Zhang, K. (2022). Partial disentanglement for domain adaptation. International Conference on Machine Learning, pages 1142..

---

Lachapelle, S., Rodriguez, P., Sharma, Y., Everett, K. E., Le Priol, R., Lacoste, A., and Lacoste-Julien, S. (2022). Disentanglement via mechanism sparsity regularization: new principle for nonlinear ICA. In Conference on Causal Learning and Reasoning, pages 428
Lopez, R., Huetter, J.-C., Hajiramezanali, E., Pritchard, J. K., and Regev, A. (2024). Toward the identifiability of comparative deep generative models. In Conference on Causal Learning and Reasoning, volume 236, pages 868-912.
Lopez, R., Regier, J., Cole, M. B., Jordan, M. I., and Yosef, N. (2018), Deep generative modeling for single-cell transcriptomics. Nature Methods, 1551201053-11588
Lopez, R., Tagasovska, N., Ra, S., Cho, K., Pritchard, J., and Regev, A. (2023). Learning causal representations of single cells via sparse mechanism shift modeling. In Conference on Causal Learning and Reasoning, pages 662-691.
Lotfollahi, M., Klimovskaia Susmelj. A., De Donno, C., Hetzel, L., Ji, Y., Ibarra, I. L.., Srivatsan, S. R., Naghipourfar, M., Daza, R. M., Martin, B., et al. (2023). Predicting cellular responses to complex perturbations in high-throughput screens. Molecular Systems Biology, page e11517.
Lotfollahi, M., Naghipourfar, M., Theis, F. J., and Wolf, F. A. (2021). Machine learning for perturbational single-cell omics. Cell Systems, 12(6))522-337.
Lotfollahi, M., Wolf, F. A., and Theis, J. (2019), seGen predicts single-cell perturbation responses. Nature methods, 16(8):755-721.
McFarland, J. M.. Paolella, B. R., Warren, A., Geiger-Schuller, K., Shibue, T., Rothberg, M., Kuksenko, O., Colgan, W. N., Jones, A., Chambers, E., et al. (2020). Multiplexed single-cell transcriptional response profiling to define cancer vulnerabilities and therapeutic mechanism of action. Nature communications, 11(1):4296.
McInnes, L.,
Healy, J., and Melville, (2018). UMAP: Uniform manifold approximation and projection for dimension reduction. arXiv.
Rampasek, L., Hidru, D., Smirnov, P., Haibe-Kains, B., and Goldenberg, A. (2019). Dr. improving drug response prediction via modeling drug perturbation effects. Bioinformatics, ((6):3743-37511
Roohani, Y., Huang, K., and Leskovec, J. (2024). Predicting transcriptional outcomes of novel multigene perturbations with GEARS. Nature Biotechnology, 42(6)9927
Srivatsan, S. R., McFaline-Figueroa L., Ramani, V., Saunders, L., Cao, J., Packer, J., Pliner,
H. A., Jackson, D. L., Daza, R. M., Christiansen, L., et al. (2020). Massively multiplex chemical transcriptomics single-cell resolution. Science, (57755555
Traag, V. A., Waltman, L., and Van Eck, N. J. (2019). From Louvain Leiden: guaranteeing well-connected communities. Scientific reports, 9()):2233,
Trapnell, C., Cacchiarelli, D., Grimsby, J., Pokharel, P., Li, S., Morse, M., Lennon, N. Livak, K. J., Mikkelsen, T. S., and Rinn, J. I.. (2014). The dynamics and regulators of cell fate decisions revealed by pseudotemporal ordering of single cells. Nature Biotechnology, 32(4):381-386
Von Kügelgen, J., Sharma, Y., Gresele, L., Brendel, W., Schölkopf, B., Besserve, M., and Locatello, F. (2021). Self-supervised learning with data augmentations provably isolates content from style. Advances in Neural Information Processing Systems, 34: 16411-16467.
Wagner, A., Regev, A., and Yosef, N. (2016). Revealing the vectors of cellular identity with single-cell genomics. Nature Biotechnology, 34(1 10::111--111000
Wu, Y., Barton, R., Wang. Z., Ioannidis, N., De Donno, C., Price, L. C., Voloch, L. F., and Karypis, G. (2023). Predicting cellular responses with variational causal inference and refined relational information. In International Conference on Learning Representations.
Zapatero, M. R., Tong, A., Opzoomer, J. W., O'Sullivan, R., Rodriguez, F. C., Sufi, J., Vickova, P., Nattress, C., Qin, X., Claus, J., et (2023). Trellis tree-based analysis reveals stromal regulation of deriedt-derived organoid drug responses. Cell, 186(25):5606-611)
12

---

Zhang, J., Greenewald, K., Squires, C., Srivastava, A.. Shanmugam, K. and Uhler, C. (2023). Identifiability guarantees for causal disentanglement from soft interventions. Advances in Neural Information Processing Systems.
Zhang, K., Peters, J. Janzing, D., and Schölkopf, B. (2012). Kernel-based conditional independence test and application in causal discovery. arXiv.
Zheng, Y., Ng. I., and Zhang, K. (2022). On the identifiability of nonlinear ICA: Sparsity and beyond. Advances in Neural Information Processing Systems, 35:16411-16422
Zhu, L., Klei, D., and Roeder (2019). Semisoft clustering of single-cell data, The Proceedings of the National Academy of Sciences, 116:466-471.
13

---

Appendix
Proof of Theorems
15
A.1 Proof of Theorem
15
A.2 2 Proof of Theorem 4.5
19
Derivation of the evidence lower bound
22
Training Details
23
Hyperparameter selection
24
Datasets and Preprocessing
24
SciPlex datasel
MultiPle--Tran datasel
25
E.3
Multiplex-7 dataset
25
Multiplex-9 dataset
26
Experimental Setups and Additional Results
26
Training Details
26
Simulation Study
26
Additional Clustering Details and Results
27
Statistical Tests and More results
28
F55
Conditional Cellular Response Prediction
28
H6
Ablation Study
29
Visualization
29
F88
Pilot Study On The Unseen Drug Responses
29
14

---

A Proof of Theorems
A.1 Proof of Theorem 4.4
Theorem 4.4. Let us first define v v(zez: t, x) as the vector:
((tzz, 1, x) = 0gn, +1 (2n,
nn t (Zn, x)
Ozno
Ozng thio Oe Pnnn,tung +he (2n t,x) O2 t
0 qm +l (2n txx)
02
where denotes the logarithm of probability density Px, of component z, If in addition to the
assumptions 4.1 4.2 we assume that for the 27z vectors
(v(zi, L2 x.) (ieez to, xo) v(Zea, to, x:) v(Zta; L, x11)
2012
(23)
linearly independent, then Zer is component-wist identifiable.
Proof. This proof proceeds in three main steps:
1. Derivation of the Fundamental System of Equations: We derive crucial relationship between the true and estimated latent variables by applying the change of variables formula, and differentiating the equality of observed data distributions.
2. Isolation of Interactive Components: We isolate the terms relevant to Zιz by strategically comparing equations for different pairs of (x, t) values and subtracting them.
3. Establishing Component-wise Identifiability: We analyze the structure of the resulting equations and the Jacobian of the transformation between true and estimated latent variables establish the component wise identifiability of Zix-
Step 1 (Derivation of the Fundamental System of Equations) Let us assume there exists another latent representation 2 that yields the same data distribution than the ground-truth variables z, for all t E T and x E X. Specifically, we have:
tyt. Pylt.x
(24)
Given our assumption of noiseless observations, equivalent to equality in distribution of the mixed variables:
Po(2)txxx = Py(x) t.x:
(25)
which, after change of variable, is equivalent to:
(o) |-
Palt,x [_g+t]:
(26)
where g y Z denotes the invertible generating function, and h := g 1 o the transformation between the true latent variable and estimated one. Jy- denotes the determinant of the Jacobian matrix of g g-
Because is invertible, |- + 0. Using this fact, we obtain the equivalent condition:
Pa(a)]txx = Palt.x
(27)
According to the independence relations in the data generating process, we have
(z)t,(2/(txx) II Px. t, x); P (2)t
II P
Rewriting the notation q: = log Pz, and qi = log yields:
n
log Dog.ett.(221e Σ qi(z,, t, x));
logPa).x(211, Σ (i(2, t, x).
15

---

Applying the change of variables formula to Equation 24 yields
" Σ q (2, t, + log Jal Σ t,
Pxlt.x = Pzlt,x |
(28)
where Jn and Jn are the Jacobian matrix of the transformation associated with h- and h, respectively. We now adopt the following notations,
dz
82
10000
" ((((k.))
0&
05,
(29)
0qi(zi 1,
8 Pqilzz x2x)
by x)
by (2i, x.
0z,
(d2;)2
Then, may differentiate Equation 28 with respect &q and & where k q {l..... n2 and # g.
Doing so, we obtain the following fundamental system of equations. For any x E X, t T. for all (kc q) nI!....n]² such that k # g:
82 log Jul
by (2t, x) ooyoo + t, tx)",,
Vz E 2,
+
= 0.
(30)
02p24
Step 2 (Isolation of Interactive Components) We may decompose the sum present on the left- hand-side of Equation 30 across the different block of latent variables, obtaining the following equality:
Σ 6y (zc, t,) d. 6,(2,, ()
Σ by (zi,
d'aw)) 6,(2,, x)d(,(ka))
(31)
na tutn
Σ
by (2,,,, dlw 6, x)d"
im
Σ
(2,, drandd (1(21, t)a(
++
Then, we may substitute according to Equation 31 in the fundamental system of equation (30), and strategically apply it to several pairs of environments. We will then take the difference of the systems of equations to make disappear the unknown quantity related to the Jacobian of
First, we apply this strategy to the pair {(x, to), (xo. to) for any treatment (we assume the existence of a reference treatment to and context Xg). Substituting according to Equation 31 into Equation 30 and applying it to (x, to), we obtain:
(z,,x) de dea b((2 x))(()))))
Σ 1=
(6% (21, to, x) ak + b/(21, to x)))((kg))
a Σ (z,, to) de (, to)
2 log Ju
0.
Iārāg
i=n +0
(32)
Proceeding similarly for (xo, to). we obtain:
ta
(6% (z. xo) al + (,, (2)x)))))(, xo)d
Σ
to, (6 to, xo, (oxxx)/dduueaae deaardie + to, xo,a"
i=
02 20 [J
Σ
(6,, to) dlow + (,,,,ttt) (ok)
DE4.g 0.
tnett1
(33)
16

---

Then taking the difference between Equation 32 and Equation 33 yields,
(6 (21xx) by (2s, xo)) de do + (di(21; x) (2((2, xo)) d((k)))
Σ
(1), - by (zi, to. aaa + (1((2,) b, to xo) of (k) = 0.
11
(34)
We apply the same principle to the pair {xx, t). (xo, to)}. We therefore get:
Ree thtt
(b (z) Xo) (a + (2,, xo)d(
Σ
(6. (zz, t, xu)
(0 + b,,len (1))
1=++1
02 log Jnl 0, 02p22
Σ (z, of + bb2(2 (k)) +
imng +1
(35)
and,
no
na than
Σ
(6 (21, al + b, xo)d
Σ
(b) (2), to, to, (((0)) + b,((2,, to (o)
0 log JAal
Σ
(6) to) dew (, (o)))
0.
+1
(36)
Taking similarly the difference between Equation 35 and Equation yields:
Σ i=
(z,,, xo) - (2,, to. to,xx))) xo)) dea (bi(2), t, xo) bi(=s to xoi) ("(
Σ
(1((a)))) -b (Ei,, (o) deoers + (b((21,)) ()(= b((2), dk((ka)) 0.
+hett
(37)
Finally, we consider the pairs {(x, t), (Xo, bo)}, for which we obtain the following:
(6 (2)x) deo + b2(2,,))" +
Σ
(b" (2,, xxx)
dl + b,(21, t, xx))((
im
82 log [J O. 2ē
Σ (z,, t) (6 + bx(22
t)a" +
+1
(38)
and
Σ (6 (z. xo) daaa
6,(22, (1 + x0)
Σ
(6" (2i, to, xo, deaerdie + bb((zi, to xo,aat xo)a"
i=
02 log [J
Σ
(6) to) dlow + (2(2ttt0))" (ok)
0.
Dēx.q
(39)
17

---

Once again taking the difference between Equation 38 and Equation 39 yields:
(6 (2), 6% (2s, xoo) decc))aica) (bilz)))) b1(2, xo) ((()))))
(Z), (zi to, daaaaaai + (di(zn) t, - 6i(2is to xo)) dk(cks)))
im
n
E 1= +hat
(16( (6 ())( (zn)t)) ((11)) + (b((2tt,)) (2( to) ((()))))) =
(40)
As final step, we are going to combine Equations 34 37 and 40 in order to correctly isolate the interaction components.
We take the difference between Equation 40 and Equation 34
Σ #m
(1i, - 1ig (ze, to, xo)) - ((0" (zi, to,x) x) by (Ei, to: xo)) (ls
+ +((bi(en, - b2(: to, x0)) (bi(2s, to. bi(2n, to. xo)) a (k.)))
Σ ((5())))) (es.:t))) (3) + (bi(2))t) - (((2stt)) ol.(k.))) =
(41)
Finally, we subtract Equation 41 and Equation
Σ
[100 (2i, t,,) - by (Ee, to, xo)) (165" (es, to, x) - by (2,, to, xo)) (b (et xo) - b¹ (Ets tu, xu))] dl
(b((,,,,,))b
be(2,, to, xo)) (bi(2e) (o,,,tt,, to, (bi(21)t,x0) bi(es, to xo)) d"(ck.)) 0.
This last equation may be rearranged as:
Σ
(" 1, (2), to: x) - 03 (zi; xo) + +b (zis to xo) al
(42)
(x,
62 xi) + 6, tu xo) ("((..)) = O. (2,,, to. x). -b,
Step 3 (Establishing Component-wise Identifiability) Given the assumption of linear indepen- dence in the Theorem, the linear system is a 27zz x 2n full-rank system. Therefore, the only solution is:
( 0 ('(...)) = for E {nz + I...., ne + nix and k (1,....nn..1. 9.
1((+) is smooth function over 2 and its Jacobian can written as:
A
d2g 82.
B
Oz..
C
02.g
Ozis
Ozi
Jn
D
Onez 82.
E
02g
F
(43)
Ozis
G
ilg 02.
H
in:
dz.
Oziz
Note that (i (g())) = 0 implies that for each i {ng l,...,, nee Rtw ) d.,a)) * 0 for at most one element E [n]. As a consequence, there's only single non-zero entry in each row indexed 1
by i E {nz + ... Mr + nzz} in the Jacobian matrix Jn. Further strengthening this argument, the invertibility of requires Jj, to full-rank, suggesting there's precisely one non-zero component in each row of matrices D, E, and F.
18

---

We note as It means that each of 2, for i E f because the other that are not in the block 2., do not have dependencies with both and z, must be that the non-zero coefficient is in the block E. Therefore D 0 and F 0. This indicates the invertibility of h; for every in the range (nz + +1.... nr + +tr)) In conclusion, Zt.: are element-wise identifiable, albeit subject to permutations and component- wise invertible transformations.
+ Rey } is attributed to at most one of the 2. And
L
A.2 Proof of Theorem 4.5
Theorem 4.5: We follow Assumptions 4.1 4.2
and the one from Theorem 4.4
(Z) the set of subsets S - Z of Z that satisty the following two conditions;
(i) S has nonzero probability measure, i.e. P(z E S t= =tt 0 for any I' E and
x X.
(ii) S cannot be expressed as Az, x Btx x 2, for any Az, C Z or as 2x x Bex x Az, for any
Az, C Zz.
We have the following identifiability result. If for all S E S(2), there exists (t1, (t) E T x T and x X such that
/ J (z)e,x(z t1. x)dz * / Pelt,x(z t2,x))dz
(44)
and there also exists (x1, X2) E X' x X and E T such that
/ Pajt.x(z t,xi )dz * / ZES
Palt,x(2 txx2)dz,
(45)
JES
then 2, and 2 are block-wise identifiable.
Proof. The proof of the block-wise identifiability of 2, and 2, are similar, so here we focus the proof about the identifiability of zg. Our proof draws parallels to the approach in Kong et al.] (2022), but context specifically pertains to the multi-conditions distributions scenarios. Throughout, use the notations " [tzz zi] 2 = [2a 2,] for simplification. This proof is comprised of four major steps.
1. Integral characterization of domain invariance. We first leverage properties of the gener- ating process and the marginal distribution matching condition to provide characterization of the invariance of block of latent variables by mixing function using an integral condition.
2. Topological characterization of invariance. We derive equivalence statements for domain invariance functions.
3. Proof of invariance by contradiction. We prove the invariance statement from Step 2 by contradiction. Specifically, we show that if 2, depended on z,"" the invariance derived in Step would break.
4. Block-wise identifiability of and 2n. We use the conclusion in Step 3, the regularity properties of h, and the conclusion in Theorem 4.4 to show the identifiability result.
Step 1 (Integral characterization of domain invariance). As reminder to the reader, g 2 y denotes the ground-trute mixing function, and ğ > y denotes the learned mixing function. We assume that both and ğ are invertible, so that their reciprocal functions particular, denote by 91s,, y Z, the estimated transformation from the observation to the coveriate-specific block of the latent space Z. We seek to find an integral characterization of the the domain Z..
well defined. In
invariance of the learned function
We seek to derive condition for which the distribution of latent variables on 2, will be unchanged when the treatment t is changed (it will however depend on x). Let S C 2, designates a subset of Z, and x E X be any context. We seek to characterize the following condition:
V(ti, t2) T. [(T., ( 5} {x, t tiH = = [(TT (ý) E S} {x, t2t]. (46)
19

---

This condition can be written equivalently using the pre-ine rmeiima set of S:
v(t1,t2) E TP {s E (Dith.) (5)} fx.t [x ti}] P [( E (II.h.)
(S]] fx, t =2]
(47)
where ( variables 22 in S.
(S) V is the set of estimated observations y originating from covariate-specific
Because of the equality of the observed data y and the generated data distribution from the estimated model y, the relation in Equation 47 also holds for the random variable y
v(t1, t2) E TP [(x (drthh)
(5))
{x. t = ti) P [{x E (ormhe)
(S)) fx,t= = t2}
(48)
It follows, by applying the image of the function 9im... that:
v(ti, E T (9 (v) E $} {x, 1x,tttto -t] = H (y) E { (x, = 12)]
(49)
Since g and 9 are smooth and injective, we may define the function à و og: 2 + Z. We note that by definition h = h where is introduced in the proof of Theorem 4.4 We now remind the reader that y = (z). Therefore, using the notation h := him,, 2 - 2x,, we have the equivalent condition
((t),t)) E P ((h.(z)) 5} txxt=tt ti}] = P ((, |s (x t2
(50)
Now, using the pre-image formulation again, we may write it as
V(ti, t2) E T [[z E hh (S)) (xxt=t(])) t.3] =
[[z hr (S)} fx, t = b
(51)
and, using an integral notation:
Y(t,, t2) ET? [
2Eh, 22h) (S) pz|t, (z) x, t1) dz = /
12eh, ((S)
Pztt, (z x, tz2) dz,
(52)
where h- (S) (z
3 hz(z)) E $} is the pre-image of S.
By exploiting the factorization of the likelihood, we obtain our final condition, equivalent to the one in Equation 47 for any S Zn and X::
1 [2a07)ek7'(
/ Px (2x x)
(P Ix,t (zt x, t1) - Pat|xxt (zt 1 x, t2)) dan = O, (53)
where the condition must hold for all (t1, t2) E T
Step 2 (Topological characterization of invariance). To demonstrate the
of Zz our objective is to substantiate that hig (lzz, Ziz, 2;)) )> = hr ( 2zz. 2, is functionally independent
of z To achieve this, we initially formulate a set of equivalent statement:
1. Statement 1. hz [Z. zz 1T) does not depend on z,"
2. Statement 2. V22 E Zz, 3B, C Zz 00 h- (zx) = B2. x Zix X Zz.
3. Statement 3. Vzg E Zz, Vr E R E S Zu 10: h- (B, (z+)) = B x Zio x Zt.
where B, (zz) is defined the ball centered around 2n with radius
B,
2// E Z, : 112//- 2n < r).
We note that Statement a mathematical formulation of Statement 1, and that Statement 3 is generalization of Statement 2 from singletons {zz} in Statement 2 to open, non-empty balls B, (zz). We proceed to demonstrating equivalence between those statements.
Statement 2 + Statement 3. Let Zy E 2g and r E R By definition of the pre- mee-maage of a set, we have that:
h (B, (z.)) = Un. (z/z)
(54)
Because we assume Statement 2, we have that for all z'2 E B, (zz), there exists set Bæg such that h- (z() BN, x Zto Z, Therefore, Statement stands for Bt = U2! Bz.
20

---

Statement 2 < Statement 3. We proceed by contradiction. Suppose that Statement 2 is false, then for a certain 2 E Zn, it is possible construct a point z - - z E Z such that z is in the pre-image of {z,} by h-1 but fgg (z²) * Zz. Indeed, this directly means that changing the other components of Z at the imput of h can alter the component of 2g at its output. By continuity of h2. there exists f > 0 such that h (z (z") & B; (zz). For such f, we have that z & h-- (B, (z,,)) Additionally, the application of Statement 3 suggests that there exists non-trivial subset B such that 1-- (B₁ (z.)) B 2 2 By definition of 2 it is clear that z E B The fact that z & hm (B5 (zz)) contradicts Statement 3. Therefore, Statement is true under the premise of Statement 3.
Step 3 (Proof of invariance by contradiction). We first show that the pre-image of any open balls of 2 are non-empty and open sets. For 2 2g and E R+, we note that because B, (2g) is open and hg is continuous, the pre-image h- (B, (z.)) )) is open. In addition, because h is continuous and we have equality of the generated data distributions:
vt E T. Vx E Ply E 8} [t,x}] = Fly S} 1 ft,
(55)
we have that h is bijection (Klindt et al., 2021). which ensures that hiz (B, (zx)) is non-empty. Hence, /1- (B, (zz)) is both non-empty and open.
We now assume, by contradiction, that Zu is not block-identiffable. Therefore, h2 is not invariant with respect to z," Additionally, because of Step 2. we have the existence of a ball S := B,, (z)) centered on the point n Zn and of radius r- Rt such that h (S*) cannot be written of the form A x 2tx x Z, for any non-trivial C 2,.
We may therefore define the set B == {z E 15") (Zim,) } x 2, x Zt & (5)) In- tuitively, Bº contains the partition of the pre-image h (S*) that the t part 2," cannot take any value in 2ta X Z,. It is therefore non-empty by hypothesis. To show contradiction with Equation 53 we evaluate it on the set S and split the integral two domains of the partition
h- (5*) = (h- (S*) BB) U B.
We define the following integrals:
T=
(er>t))hh="(>) Parlx (2x x) Pat ixit (zt x, ti)
P2t 1xtt (zt x, t2)) da (56)
Ti
[ax=)t Px (zx x)) (Pot (z) t1) Patt tx (zz x, t2)) dz,
(57)
Tg [BB Pa,ix (ax x) (axteex (zt x,, t1)
x, t2)) dz Patlexx (zi
(58)
where we have the expected relation T = Ti T2-
We first look at the value of Ti. In the case where the set 1-1 (S*) AB is empty, then 71 trivially evaluates to O. Otherwise, there exists a non-empty subset Ož. of Zz such that h (S*) BB = C2 Zio Ze. With this expression, it follows that
Ti 1 [Kr.e] x2,x22, Pa. x (2x x) (Pe) |x,t (z. |x,ti) Pe! xx,e
Pal (zz Xx, t2)) dz dz,
(59)
Because of the separability of the domains, we may apply Fubini's theorem:
T = /
Pa. ix (Zz x) dzp
(Pat ix.e (z. |x, tz) - -ax bx,t (z) xx t2)) dz dz,
22 E2 Zs
(60)
T Pz. xx (Zx(x) (1 1)dka 0.
(61)
E2,
Therefore, in both cases Ti evaluates to 0 for S*
21

---

Now, we address T2. Towards this goal, we prove that B satisfies the condition for application of the assumption of the theorem. First, we must show that B has non-zero probability measure for all values of and x. For this, it is enough to show that B contains an open set, given that we assume that (z t, x)
over (z, t. x) 2 x X. Let us take one element 2 Bg which is possible because we proved B is non-empty. As discussed above, h (S*) is open and non-empty, and by continuity of hg, there exists To E R* such that B,, (zz) S Bg Therefore, By contains an open set and has non-zero probability. Second, is by definition that B cannot be expressed as Au, Ztx 2, for any Az, Zx.
Therefore, condition (ii) from the theorem indicates that there exists ti, 122 x", such that
Parlx (Zx x²) (Pat 1t.x (zi" ti) Pa, |t,x (zi xx" tÿ) dZz # 0. T2 |B
(62)
Therefore, for such S", we would have 7i T2 * 0 which leads to contradiction with Equation 53 We have proved by contradiction that Statement from Step holds, that is, h2 does not depend on the treatment variable and interaction variable Z,, Zta.
Step 4 (Block-wis) identifiability of 2, and 22). With the knowledge that h₂ does not depend on z: we now show that there exists an invertible mapping between the true content variable Zy and the estimated version 2g.
As h is smooth over Z, its Jacobian can written as:
Uss Ohe
B
Otg Ozi.
C
0z.
Jn
D
U2
E
U2i
F
0212
(63)
2hg
Onis
G
==
d2.z
Ozz
H
dz Oziz
8z: 87..
where we use notation 2. h(z) and بيةة h(z))
(n -ln 2, h(
First, we notice that under the condition of Theorem 4.4 there is an invertible mapping between Ztw and ătz. Therefore, it must be that D F O, and that E is non-singular. Additionally, have just shown that 2n does not depend on the treatment-related variables z Therefore, it follows B C = 0. On the other hand, as h is invertible over J is non-singular. Therefore, A must be non-singular due to B = C = O. Relying on analogous assumptions to prove the invariance of 2, with respect to 2 it follows that G H 0, and that I must be non-singular.
We note that A is the Jacobian of the function W'z (zz) == h (z) Ze - Zz. which takes only the conclude that
covariates part z, of the input z into hr Together with the invertibility of h. 14 is invertible. Therefore, there exists an invertible function Mr between the estimated and the true able to
2, h' (zz), which concludes the proof that 2 is block-identifiable Similarly, conclude 2 is alock-identifable.
B Derivation of the evidence lower bound
We now introduce the classical derivations of the celebrated evidence lower bound for our generative
model. The evidence the logarithm of the marginal data probability, and we calculate by weighting it against the variational distribution:
[Po Zz, Ztzz Zzz t, x)
logp(y x, t) = log E (R. Mi m. (y.),))
(64)
go (Ze, Zizs Ze y. t,
We apply Jensen's inequality using the fact that the logarithm is a concave function:
[y Zt, Zta, Zor 1t xx) (Ze, Ziz. 2z ly, t, x)
logp(y x,t)) t) 2 Egs (21,2ia Rn xx,)tx)) log
(65)
22

---

Algorithm Training of FCR
Input: y,y", shared and t, t 0
Output: p, q. Zx, Ztx, Z4 while Not converged do
Minimization Stage
1. Estimate p(z; t), plzz x).p(Ziz x),), and p(z) tt), p(zz x).p(ziz to x)
2. Estimate (Z Zzz Zex y),
(z,, zΩ, 2/x t ..
3. Calculate Bullback-Leiblle divergence terms and predict
4. Calculate similarities {z.g 2z), ; {z]. zz)
5. Permute Zi to get Liz and predict permutation labels
6. Minimize LELBO, Lsim- Las, Edisty Edist,
Maximization Stage
1. Estimate g 2z, Ztx x,,,y),
2. Permute Ztx to get Ziz and predict permutation labels
3. Maximize Laist Laise,
end while
Then, use the factorization of our generative model:
poly Zz, Zta, Uz Dolpoola Ztz Ux
t, x)
logp(y x,t) CELBO = Eag (2: (2:
2. ly t.xc log
(66)
Tol go( (2, Zen, Zy y, t, x)
Naming the right hand size of Equation 66 LELBO, notice that LELBO can be written as the following difference:
Eng po 2e, Bea, 2x) Egg log go(Zi, Zee: Zz y, x.
LELBO
(67)
po po( Zear Z t,
where the first term is the reconstruction loss, and the second term is the Kullback- divergence between the approximate posterior Ts qa(2,, Zta, Zu y, t, x) and the prior po Ztz 2₂ t, x). Further decomposing this term using the factorization of the generative model and the inference model, we obtain:
Egg log (o(ze, Ziz Zizi Ze y,t,
(z y,t) t)
E (x log
po (2e, Bea Zx t, x)
(2e,
po(z; t))
Joi Ztr
y,
Eng log
(68)
py Ztz
tt
9á Zrz
y.x)
Ey xx) (xx))oo log
po Zzp
x)
Finally, recognizing three Kullback-Leibler divergence terms in the right hand side of Equation 68 and injecting this expression into the evidence lower bound expression of Equation (66] we obtain the desired expression:
LELBO =B log Po Zz 24, Ztx)
DKI(qo(Zz (po)) x))
(69)
-DK(qo(Zz t,y))lpo(2, t))
-DKzlgoZZiz t, t,x)),
C Training Details
provide additional training information this section, including a detailed algorithm for the training process of FCR (Algorithm I.
23

---

D Hyperparameter selection
We split the data into four datasets: train/ train//alidatiio/ees//rrcccccci
following the setup from previous works (Lottollahi et al. 2023 Wu et al. 2023). First we hold out the 20% of the control cells for the final cellular prediction tasks (prediction). Second, we hold 20% of the rest of data for the task of clustering and statistical test (test). Third, the data excluding the prediction and clustering/test sets are split into training and validation sets with a four-to-on ratio.
For the hyperparameter tuning procedure, conduct the exhaustive hyperparameter grid search with eeeepoc==100 the loss assessed the validation data. The hyperparameter search space is shown
in Table
Parameter
Values
W1
{0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0)
W2
{0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0. 7.0, 8.0, 9.0, 10.0)
W3
{0.1, 0.3, 0.5, 0.7, 0.9, 1.0, 3.0. 5.0, 7.0, 10.0)
Table Hyperparameter space
E Datasets and Preprocessing
In this section, we provide detailed descriptions of the four datasets and their corresponding pre- processing procedures. These datasets contain cells from many cell lines that described in Table 3
No.
Cell Line
Origin
1
IALM
Lung
2
SKMEL2
Skin
SHIOTC
Stomach
SQI
Lung
5
BICR31
Upper Aerodigestive Tract
6
DKMG
Central Nervous System
BT474
Breast
TEN
Endometrium
COLO6SON
Ocsophagus
10
CAOV3
Ovary
11
SKMEL3
Skin
12
NCTH226
Lung
13
LNCAPCLONEFGC
Prostate
14
RCCIORGB
Kidney
15
BICR6
Upper Aerodigestive Tract
16
BT549
Breast
17
CCESTTGI
Central Nervous System
18
RERFLCADI
Lung
19
UMUCI
Urinary Tract
20
RCMI
Large Intestine
21
LS1034
Large Intestine
22
SNU1079
Biliary Tract
23
NCIH2347
Lung
24
COV434
Ovary
Table 3: Cell line information for multiplex experiments
F.1 SciPlex dataset
This dataset includes three cancer cell lines exposed to 188 different compounds. In total, this experi- ment profiled approximately 650,000 single-cell transcriptomes across roughly 5,000 independent samples (Srivatsan et al. 2020). For experiments, selected only the HDAC inhibitors that

---

were shown to be effective in these three cell lines (Srivatsan et al. 2020). The following is the list of HDAC inhibitors used in Table 4
HDAC Inhibitor Drugs
Belinostat
Mocetinostat
Panobinostat
Pracinostat
Dacinostat
Quisinostat
Tucidinostat
Givinostat
AR-42
Tacedinaline
CUDC-907
M344
Resminostat
Entinostat
TSA
CUDC-101
Table 4: The list of HDAC inhibitors
We extracted the cells treated with HDAC inhibitors along with their corresponding control groups. After filtering out low-quality cells, we normalized the raw counts. From these, we retained the top 5,000 highly expressed genes. Ultimately, analyzed 90,462 cells, including both treated and control groups. Cell lines repetition numbers were used as covariates, with treatment dosage designated the treatment variable.
E.2 MultiPlex-Tram Tram dataset
This dataset is referred to experiment-5 in the paper McFarland et al. (2020). It contains in total 20,028 Trametinib treated cells, DSMO control cells, and 24 cell lines (McFarland al. 2020). The cells are treated with 100nM Trametinib for 3, 6, 12, 24, 48 hours respectively. We removed the low quality cells and normalize the raw counts. Next, we kept first 5,000 differentially expressed genes. Finally, we have 13,713 cells in total for the experiments and down streaming evaluation.
E.3 Multiplex-7 dataset
This dataset is labeled as experiment-3 McFarland et al. (2020). includes 72,326 cells treated with seven different compounds across 24 cell lines. Details of the seven treatments provided in Table Following the removal of low-quality cells and normalization raw counts, we retained the top 5,000 differentially expressed genes, yielding 61,552 cells for the experiments and downstream analysis.
Drug
Hours
DMSO
6 hours
DMSO
hours
BRD3379
6 hours
BRD3379
24 hours
Dabrafenib
24 hours
Navitoclax hours
Trametinib
24 hours
Table 5: The multiPlex-7 dataset's treatments
25

---

E.4
Multiplex-9 dataset
referred to as experiment- in McFarland et al. (2020), consists of 37,856 cells across 24 cell lines, treated with 9 drugs (including control). The list of drugs can be found in Table After filtering out low-quality cells and normalizing the raw counts, we retained the top 5,000 differentially expressed genes. This resulted in total of 19,524 cells for the experiments and downstream evaluation.
Drug
Hours
DMSO
24 hours
Everolimus
24 hours
Afatinib
24 hours
Taselisib
24 hours
AZD5591
24 hours
JQ1
24 hours
Gemcitabine
24 hours
Trametinib
24 hours
Prexascrtib
24 hours
Table 6: Drugs list of the multiPlex-9 dataset
F Experimental Setups and Additional Results
F1
Training Details
In this subsection, we layout the training parameters for each datasets as follows.
sciPlex For the sciPlex datasets, the dimensions Additionally, we set the hyperparameters to w1 = 3.0, We 3.0, and Wg 5.0, with batch size of 2046. Note that in sciPlex, treat the dosages of the HDAC inhibitors the treatment variable.
follows: zμ is 32, Zir is 64, and is 32.
MultiPlex-Tram or the multiPlex-Tram dataset, the dimensions are set as follows: 2, is 32, 2tz is 32, and 2, 32. The hyperparameters are w1 = 5.0, w2 5.0, w3 = 1.0, with a batch size of 2046. In this dataset, Trametinib treatment time is considered the treatment variable.
multiPlex-7 For the multiPlex-7 dataset, the dimensions are z, 32, Etr = 64, and z, 32. The hyperparameters are w1 1.0, 422 0.5, and Wy 0.1, with batch size of 2046.
multiPlex-9 For the multiPlex-9 dataset, the dimensions 2, 32, 2ir 64, and 32. The hyperparameters are wi = 1.0, We 0.5, and Wg = 0.1, with batch size of 2046.
Additionally, the autoencoder learning rate is set to 3 x 10-4, the discriminator learning rates are also 3 x 10-4, and the number discriminator training steps 10.
F.2 Simulation Study
We provide some empirical assessment of our identifiability theory using simulations.
Data Generation Following the simulation protocol outlined in Kong al. (002)), Khemakhem et al. 2020) and Lopez et al. (2023), we simplify the simulation setup by setting the dimensions of Zu and 2, 1, while Zut has dimension of 4. Specifically, our variables as follows:
sample size of 5,000 and define
t ~ Unif( ({1, 2, 2,3),
(70)
Here, t represents the treatment variable, uniformly distributed over three discrete values.
x Unif({100, 1000, 5000))
(71)
26

---

x denotes the covariate, also uniformly distributed but over a wider range of values.
Z Normal(x/2,1
(72)
2g is the latent variable associated with x, following normal distribution with mean x/2 and unit variance.
2 Normal (1/2, 1)
(73)
is the latent variable associated with t, also normally distributed with mean t/2 and unit variance.
Z Normal(x t, 14)
(74)
Zte represents the interaction between x and following a multivariate normal distribution with mean x t and covariance matrix I4 (the 4-dimensional identity matrix).
Finally, we define output y function of these latent variables:
= g(zz, Zta. Zt)
(75)
Here, 9 implemented as 2-layer MLP with Leaky-ReL.U activation, following the approach of Kong al. (2022) and Khemakhem et al. (2020). The output y is real-valued vector with dimension of 96.
Evaluation To assess the component-wise identifiability of the interaction components, we com- pute the Mean Correlation Coefficient (MCC) between Ztr and h MCC is a standard metric in Independent Component Analysis (ICA) literature, where higher MCC indicates better identifi- ability. MCC reaches I when latent variables are perfectly identifiable (up to a component-wing transformation). We compute the MCC between the original sources and the corresponding latent variables sampled from the approximate posterior. As for the iVAE evaluation framework, we first calculate the correlation coefficients between all pairs of source and latent components. Then, solve a linear sum assignment problem to map each latent component to the source component that correlates best with it, effectively reversing any latent space permutations. A high MCC indicates successful identification the true parameters and recovery of the true sources, up to point-wise transformations. This is standard performance metric in ICA (Khemakhem et al.) 2020).
Results Our results suggest that our method, FCR largely outperforms existing variational autoencoder-basee approaches in identifying the latent interactive components Ztx- As shown in Table FCR achieves an almost perfect Mean Correlation Coefficient (MCC), compared to other methods that have poor performance. Even the iVAE baseline falls short of FCR's capability in recovering the true latent structure. These results indicate that FCR offers a significant advancement in component-win wise identifiability for complex, interacting latent variables in causal representation learning.
Method
MCC
FCR
0.91 + 0.03
8-VAE
0.38 + 0.12
FactorVAE
0.37 + 0.08
iVAE
0.77 + 0.07
Table Mean Correlation Coefficient (MCC) of Zto for different methods
F33 Additional Clustering Details and Results
Clustering Approach Our clustering analysis utilizes the learned representations Zz, Zta. and ze for our method, while all available representations were used for baseline methods. It's important to note that baseline models were trained using their default settings.
We employed the following clustering approach for different scenarios:
Clustering on x: We applied the Leiden clustering algorithm to the different representations and evaluated the results using the labels of x.
27

---

2. Clustering on t: Similarly, we used the Leiden algorithm on the representations and assessed the outcomes with the labels t.
3. Clustering on combined with t: We ran Leiden clustering the combined representa- tions and evaluated the results using both the labels of x and t.
This approach allows us to assess the efficacy of our learned representations capturing the underly- ing structure of both individual and combined variables.
Evaluation Metric follows:
The evaluation metric, Normalized Mutual Information (NMI), is defined as
2 ② Et PMu log
Pei
NMI(Y,C)
(76)
-E ph logr% + Et= of log
where:
Y = {V1, , 99K denotes the set of class labels,
C = (c....., ) denotes the set of cluster labels,
Pkt is the joint probability data sample belonging to class & and cluster
Pk is the marginal probability of sample belonging to class
pf is the marginal probability of a sample belonging to cluster 1.
Note that these probabilities are computed from the same dataset, but with respect the true class labels and the assigned cluster labels, respectively.
Additional Results Additional clustering results for the multiPlex-tram, multiPlex-7, and multiPlex- 9 datasets are presented in Figure 5
F44 Statistical Tests and More results
Due to the computational complexity of kernel calculations, we adopted sampling approach with 2,000 samples, repeating the process 100 times to report the results. For the baseline methods, we sampled their latent spaces match the dimensions of Zz, Zg, and Zir- This sampling was repeated 20 times, and we reported the best results for comparison with FCR. The test results for the multiPlex-tran multiPlex-7, and multiPlex-9 datasets are shown in the Figure
F,5
Conditional Cellular Response Prediction
For this task, we use FCR's 's latent representations to predict gene expression levels and report the corresponding R² scores. Specifically, our approach enables the prediction of cellular responses at the single-cell level. The focus of this paper is on predicting cellular responses (expression of 2,000 genes) in control cells subjected drug treatments. Our comparative analysis includes CPA. VCI, VVE,, scGEN, and CINEMA-OT, these methods are specifically designed for cellular prediction tasks.
We utilize FCR to extract the control's latent representations [z. 20 2!] and the corresponding experimental representations [z,, Zig, 2t). The decoder g is then used predict the gene expression levels as y = g( z z!).. The R² score is used to evaluate the predictions, and we sampled 20% of each dataset for testing, repeating this process five times.
For iVAE and VCI, we used treatments, covariates, and gene expression data to learn latent variables and predict gene expression. In contrast, for VI, 3VAE, and factorVAE, only used gene expression data to learn latent variables and make predictions.
The R² score is key metric for assessing the accuracy of predictive models. It measures the proportion of variance the dependent variable that is explained by the independent variables. R2 score of I indicates perfect prediction accuracy, meaning all variations in the target variable
28

---

fully explained by the model's inputs, while score of 0 suggests that the model fails to capture any variance in the target variable.
R² 1= (3k 0:
(77)
2 (1/i y)²
where is the actual values, Us is the predicted values, y is the mean of average values, n is the number of observations.
We further utilize the Mean Squared Error (MSE) of the top 20 differentially expressed genes (DEGs) for post-treatment. These 20 genes selected for showing statistically significant differences in expression levels for each cell line with drug treatments compared to control samples. The same procedures are also carried out Roohani al. (2024). Note here, we did not compare with CINEMA-OT and scGEN because they are only for binary treatments.
Additional Metric
Dataset
FCR
VCI
CPA
sVAE
sciPlex
0.12 + 0.03
0.15 + 0.04
0.15 + 0.04
0.16 + 0.05
multiPlex-tram
0.10 - 0.07
0.13 0.08
0.13 1 0.08
0.14 + 0.07
multiPlex-7
0.18 + 0.07
0.21 + 0.08
0.22 + 0.08
0.23 + 0.07
multiPlex-9
0.24 + 0.06
0.27 0.04
0.23 + 0.06
0.26 + 0.05
Table Mean Squared Error (MSE) of top 20 Differentially Expressed Genes (DEGs) for different methods across datasets
F.6 Ablation Study
We present the ablation study results for the hyperparameters Wi, 42, and w3- Initially, set wis w2, and wa to [1,1,]]. Then, we varied each parameter independently to [1, 5, 10, 20j. keeping the other parameters fixed 1. Figure iillustrates the NMI scores for clustering on x, t, both x&t.
Figure shows the Ro2 scores for each parameter.
E.7 Visualization
Visualizing latent representations provides intuitive insights into the distinct characteristies and attributes captured by each representation. Using Uniform Manifold Approximation and Projection (UMAP) (McInnes et al. 2018), we visualized the latent representations Zz, Zta, and z, derived from the sciPlex dataset (Figure
From these visualizations, we observe that zμ effectively captures covariate-specife information, showing clear separation in Figure 9 but it does not reflect treatment information. In contrast, the UMAP visualization of reveals clear pattern, where cells treated with higher dosages are positioned at the bottom, and control or lower-dosage treated cells are found the top. Additionally, various cell lines intermingle within the plot, indicating that 2, successfully captures drug responses across different cell lines.
Most importantly, Figure 9 demonstrates that Ztæ captures cell line-specific treatment responses, representing a balanced integration of both covariate and treatment information. Specifically, in the sciPlex dataset, the MCF7 cell line shows a pronounced cell line-specific response, aligning with findings from biological literature (Srivatsan et al., 2020), while the K562 cell line exhibits a less distinct response. This representation confirms the strong, unique responses in certain cell lines, highlighting the validity and precision of our method capturing nuanced biological behaviors.
E.8 Pilot Study On The Unseen Drug Responses
Predicting responses to novel treatments a pivotal and fast-evolving field in drug discovery. However, the biological literature indicates that cellular responses highly context-depenee -onenet-dependent (McFarland et al. 2020). This complexity poses significant challenges for Al-driven drug discovery, which often struggles to achieve success in clinical trials.
Our motivation for developing FCR stems from the need to understand how cellular systems react to treatments and identify conditions that can deepen our understanding of these responses. FCR enables
29

---

the analysis of drug interactions with covariates and contextual variables. Additionally, predicting cellular responses to new treatments necessitates prior knowledge, such chemical structure and molecular function, and comparisons with known treatments. Without this context, predictions can be unreliable.
In this paper, our primary focus is not on predicting responses to unseen treatments. However, given the relevance of this topie, we conducted pilot experiments to showcase the potential future applications of FCR. The multiPex--Tram and multiPlex-7 datasets share the same cell lines and Trametinib-24 24 hours treatment, along with other different treatments. By utilizing these two datasets, we established the following experimental settings for unseen prediction scenario:
1. Drug Hold-out Setup:
We held out two cell lines, ILAM and SKMEL2, from the Multiplex- Tram dataset, which had been treated with Trametinib for 24 hours. We trained FCR model using the remaining data, and this model referred as Mn. We denote the dataset (Multiplex-Trran dataset without Trametinib- 24h treated LLAM and SKMEL2) as D
2. Prior Knowledge Model:We trained another FCR model, Mp. using the Multiplex- 7 dataset, which includes the ILAM and SKMEL2 cell lines treated with Trametinib for 24 hours denoted as DP. We treat this model, Mp, as prior knowledge model.
3. Transfer MLP: We extract both ziz from Mp and z from M, for DP. Then trained 1-layer MLP (the same dimension as z transfer ziy to z!. by minimize the MSE between them.
4. Contextual Prior Representation: We extracted the 2 representations from model Mp for ILAM and SKMEL22 in holdout set. Then transfer ziz by the previous MLP to 2 as prior contextual embedding. For the hold-out cell lines ILAM and SKMEL22 in the Multiplex- -rrm dataset, we extracted z from model M,, and Trametinib-244 z from other treated cell lines D
5. Representation Matching:
Then match the 2h. by z similarity ILAM and SKMEL2 across Multiplex-tram and Multiplex-7 in the prior knowledge model space.
6. Prediction: We predicted the unseen 24 hours Trametinib responses for the ILAM and SKMEL2 cell lines in holdout dataset using the formula: ý g g(zh, 2h. w 2h. is the corresponding matched prior contextual representation. z; is the tramnib-24 average value in other cell lines.
7. Evaluation: We computed the R2 and MSE for the top 20 differentially expressed genes (DEGs) based on the predicted values and compared these results with those from the VCI model. The paper's OOD prediction setups.
Method
R²2
MSE
FCR VCI
0.74 + 0.03
0.52 + 0.11
0.71 + 0.05
0.55 0.08
Table Out-of-Distribuuti (OOD) Performance Results
30

---

(8)
05-
0.6
0.6
g
044
03-
0.3-
63
C9
lachesite
betotng
betabng M90
Bethods
aethots
nethood
(b)
1.0
40
0.3
0.3
lacorsie
betakin
2000
belatid
fecion
lethods
acthoss
nethods
(c)
0.
19
0.5
+0K
+
HHH
0.6
0.3
011
Indeosid
incoorlee
betatig
C96
etthoss
nethods
Figure 5: The clustering results. (a) For the multiPlex-Tran dataset, the first column shows the NMI value for clustering on x, the second column shows the NMI value for clustering on xt, and the third column shows the NMI value for clustering presents the NMI value for clustering on x, the second column shows the NMI value for clustering on xt, and the third column shows the NMI value for clustering on t. (c) For the multiPlex-4 dataset, the first column displays the NMI value for clustering on the second column shows the NMI value for clustering on and the third column shows the NMI value for clustering on
t.(b) For the multiPlex-5 dataset, the first column
31

---

(8)
0.5
0.08
0.25
0.06
Pusluo
0.3
0.15
0.20
nual
0.20
0.04
0.10
0.02
0.05
0.05
0.0
0.00
0.00
Besed You bood
NAE
tetaval
fectorn
CP9
svase
FOR
WAB
tetavve
fectovvig
cog
sWAE
Tozatt
Z224t
Z41
(b)
0.35
0.04
0.6
0.6
0.25
0.4
Pvolus
0.4
0.20
0.15
0.10
0.01
0.0
0.0
0.00
0.00
NCW
NALD
incteriae
CPR
suse
FOR
Sew
WAe
betavsee
tectoriAg
vial
cass
suae
FOR
ZaZett
Zizat
227
201
(0)
(2]
Swar Pott 2XTR
0.20
0.30
0.20
Poalue
0.15
0.10
0.15
0.02
0.10
0.02
0.05
0.00
0.00
0.00
0.00
NAe
tectorice
CPs
SvAe
FRR
acsi
MAE
feciorvig
swae
FOR
Zazat
Figure Conditional independence test results for the multiPlex-7 (a), multiPlex-tram (b), and multiPlex-9 (c) datasets. The first column presents the p-values for the conditional independence test of 2, and conditioned on x, with the red dashed line indicating significance threshold of 0.05. The second column shows the p-values for the conditional independence test of 2, and conditioned on The third column presents the p-values for the conditional independence tests of 2, and 2tz conditioned on and of 2, and biz conditioned on t. The fourth column shows the HSIC values of z with x, t, and random variables (R), as well as the HSIC values of ze with x, t, and random variables (R).
32

---

0.9
0.175
0.7
0.150
0.6
0.125
0.6
0.100
0.076
0.4
0.050
2,
0.3
Z:
0.025
Zox
0.3
10
15
20
10
15
20
10 omega1
15
20
omega1
omega1
0.9
0.7
2x
0.20
21x
0.7
2,
2:
0.15
0.5
Zzc
0.10
0.3
0.3
0.05
omega2
omega2
10
15
20
5
omega2
15
20
0.9
0.70
0.175
0.8
0.85
0.150
0.60
0.125
2,
0.6
0.100
2;
Z
0.50
Ztx
0.5
0.45
0.050
0.4
0.40
0.3
0.35
10
15
20
10
20
15
20
omega3
omega3
omega3
Figure 7: Clustering Ablation Results
0.80
0.75
0.8
0.8
0.7
0.7
0.65
0.60
0.6
0.6
0.55
5
10
15
20
10
15
20
10
15
20
omega_1
omega2
omega3
Figure R² score with different w1, We and W/ values
33

---

(a)
2, UMAP
Z UMAP
2 UMAP
Zg UMAP
Z UMAP
Z, UMAP
Z UMAP
Z, LMAP
Z, UMAP
Z, UMAP
Z, UMAP
Z, UMAP
MCF7_Rapp Rep!
MCF7_Rap2
K562_Rept
A549_Rep2
AS49. Rep1
K362_Rap2
Dosage
(b)
Z, UMAP
Z, UMAP
Z UMAP
Z, UMAP
Z, UMAP
Z, UMAP
Z UMAP
Z, UMAP
7, UAAP
Z, UMAP
2, UMAP
UMAP
Figure 9: (a) UMAP visualization of the sciPlex dataset. The first row distinguishes the data by cell type, while the second row distinguishes by treatment time, with brighter colors indicating longer treatment durations. (b) UMAP visualization of the multiPlex -rrmm dataset. The first row distinguishes by cell type, and the second row distinguishes by treatment time, with brighter colors representing longer treatment durations.
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
35

---

Justification:
The proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer:[Yes Yes]
Justification: The hyperparameters provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
36

---

Answer:
[Yes]
Justification:
The code is provided on GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
coiiioionoly for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ lttppp.cc/pulli/guuiiee gwtcc//ggge///eee CoodeSumiiii wwweeisssonnnae ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:[Yes
Justification:
Details are in the appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: p-values are calculated in the experimental sections.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
37

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide this detail the appendices.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://wwwrr..cooollllccccccaaaa
Answer: [Yes]
Justification: research conducted the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:[Yes]
Justification:
We discuss how this paper would contribute to the precision medicine.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
38

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
It's not applicable in our paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: Not applicable to paper.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
39

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification:
Documentation is available with the code on GitHub.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer:[NA]
Justification:the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with989 human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
40


Reviews:
[{"review_id": "OI1Nmq4ZrR", "reviewer": "Reviewer_vnnM", "summary": "The authors develop a modification of the identifiable nonlinear ICA algorithm of Khemakhem et al. 2020, suited for scenarios with two disjoint groups of auxilliary variables x and t and their interactions, with an application of treatment effect on cell gene expression.\nThe authors propose a learning objective designed to enforce the conditional independence assumptions of the model, and evaluate it against baselines from the disentanglement, ICA, and cellular response estimation.", "questions": "- Equation (6) (the constraint for interaction between x and t) - how does that guarantee the interactions? The learned embeddings can have empty dimensions, e.g, k(x) = [k_1(x), …, k_n(x), 1, …, 1] and k(t) = [1, …, 1, k_1(t), …, k_m(t)] so that the Hadamard product becomes [k(x), k(t)]\n- Equation (7) - shouldn’t f_dis also take x as an argument since you want to enforce a conditional indpendence? Or do you train a separate f_dis for each possible value of x that you condition on (probably not feasible)?\n- If the model is supposed to be fully identifiable wrt. z_xt, then I would suggest evaluating this with the appropriate metrics on synthetically generated data in addition to the clustering objective (e.g., MCC  as in Khemakhem et al. 2020, or any of the disentanglement metrics [Locatello et al. 2019])\n- How were the hyperparameters selected for the cellular response task?", "limitations": "- the algorithm has 3 hyperparameters, which might make training difficult. Details on hyperparameter selection were ommited\n- the evaluation is limited - results of the downstream task of predicting cellular response yield limited improvements, and the identifiability of  z_{t,x} is not properly evaluated (with e.g., the MCC metric using synthetic data)", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the proposed method is justified theoretically\n- the paper is written in a clear and concise manner\n- extensive selection of baselines", "weaknesses": "- I fail to see the reasoning of comparing the conditional independence results (6.3) with the other baselines, which were not designed to have separate latent subspaces for x,t, and tx, and thus it does not seem surprising that the conditional independence does not hold for random subsets of them?\n\n- Results of the cellular response tasks are barely better than the baselines - did the authors conduct any tests to see if the differences are statistically significant?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Comparing the conditional independence results (6.3) with the other baselines}$:\n\n\nWe thank the reviewer for raising this important question.  I would like to clarify the experimental setup. The goal of this experiment is to demonstrate that FCR is capable of learning conditionally independent spaces for $z_x$, $z_t$, and $z_{tx}$. Our algorithm is the first to specifically address the $x$, $t$, and $tx$ subspaces, which makes it challenging to find direct baselines for comparison. However, the baselines such as **iVAE**, **betaVAE**, **factorVAE**, and **sVAE** are semi-supervised or unsupervised disentangled representation learning methods, they aim to learn representations that should be independent or conditionally independent given $x$ and $t$. Additionally, methods like **CPA** and **VCI** enforce certain invariance relations with respect to $t$ and $x$, potentially leading to some level of conditional independence in the learned representations.\n\nHowever, the challenge in making this comparison lies in not knowing which blocks of latent variables induced by baseline methods satisfy the specific conditional independence relations of interest. To ensure a fair comparison, we exhausted different combinations of the representations for the tests and reported the best results. This comparison is intended to demonstrate that **FCR** can effectively learn more meaningful and nuanced conditional independence relations in $x$, $t$ and $tx$, compared to the baselines. We will ensure that this clarification is included in the manuscript.\n\n$\\textbf{Differences are statistically significant}$:\n\nThank you for raising this concern. We tested for significance of changes in mean between the results from the FCR method and those of the second-best performing method (paired t-test). The p-values, along with the corresponding cell number in the testing set appears in the table below.\n#### Table: p-value for $R^2$ performance \n| Dataset         | p-value (# cells) |\n|-----------------|-------------------|\n| sciPlex         | 0.0006 (295)      |\n| Multiplex-Tram  | 0.0056 (266)      |\n| MultiPlex-7     | 0.0001 (318)      |\n\nThe results indicate that FCR significantly out-performed the baselines across cell lines on sci-Plex, multiPlex-Tram and multiPlex-7. On MultiPlex-9, where FCR is slightly outperformed by CPA, the differences are not significant. We will make sure to include all p-values in the updated manuscript.\n\nTo propose an additional metric beyond $R^2$, we also evaluated the mean square error (MSE), for the top 20 Differentially Expressed genes (DEGs). These 20 genes are  measured on the largest difference for each cell line with drug treatments compared to control samples (following [1]). Note here, we didn’t compare with CINEMA-OT and scGEN because they are only for binary treatments (case-control data). For detailed results, please check the results in **General Response to ALL Reviewers a**.\nGenerally speaking, FCR outperforms other baselines in 3 of 4 datasets. \n\n$\\textbf{Concerns on Equation (6) (the constraint for interaction between $\\mathbf{x}$ and $\\mathbf{t}$)}$: \n\nThe reviewer asked the question \"How does that guarantee the interactions? The learned embeddings can have empty dimensions, e.g, $k(x) = [k_1(x), \\ldots, k_n(x), 1, \\ldots, 1]$ and $k(t) = [1, \\ldots, 1, k_1(t), \\ldots, k_m(t)]$ so that the Hadamard product becomes $[k(x), k(t)]$\".\n\nWe apologize for this misleading language. Essentially, we did not want to claim any theoretical guarantee in this case, but instead motivate the specific architecture choice. We therefore changed the language to \"in order to encourage the prior for $z_{xt}$ to capture interactions between $x$ and $t$, we design the functions $f^p_{\\mu}$ and $f^p_{\\Sigma}$  to be of the form $k(x) \\otimes k(t)$, where $\\otimes$ denotes the Hadamard product. \n\n$\\textbf{x as an argument for $f_{dis}$}$:\nWe agree with the reviewer that it would be much clearer to put $x$ as an argument of  $f_{dist}$.  It is already the case that the function $f_{dist}$ implicitly considers $x$ as an argument because we perform a random permutation of the triplet $(z_x, z_{tx}, x)$ under the condition of the same $x$ (i.e., the permutation or shuffle occurs only within a set of $z_x$ and $z_{tx}$ under the same $x$). As a result, the input to $f_{dist}$ is the set of newly generated ($z_x$, $\\hat z_{tx}$), where we ensure that $z_x$ and $\\hat z_{tx}$ correspond to the same $x$. The outputs of $f_{dist}$ simply are the labels indicating whether a permutation occurred or not. Therefore, we can say that $f_{dist}$ effectively takes $x$ into account. We briefly discuss this process in lines 226-230. We will make sure to update it in our final manuscript.\n\n$\\textbf{Simulation studies on synthetic data}$:\nPlease check the **General Response to ALL Reviewers c**.\n\n\n$\\textbf{Hyperparameter selection}$:\nThank you for raising this important question, we appreciate that it is a critical point for evaluating how well the method works. Please check **General Response to All Reviewers b**.", "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": "", "extra_scores": {}}, {"review_id": "0VIZZpG7yX", "reviewer": "Reviewer_xvG5", "summary": "The paper presents a novel method, the Factorized Causal Representation (FCR), which leverages identifiable deep generative models to understand cellular responses to genetic and chemical perturbations across multiple cellular contexts. This method improves upon prior models by learning disentangled representations that delineate treatment-specific, covariate-specific, and interaction-specific factors, which are shown to be theoretically identifiable. The effectiveness of FCR is demonstrated on four single-cell datasets. The performance, measured by the R^2 score of conditional cellular responses, often marginally outperforms state-of-the-art methods, showcasing FCR’s potential in accurately modeling cellular dynamics.", "questions": "Can you elaborate more about the design choice of the evaluating metric?", "limitations": "- Generalizability: The paper does not show the method’s applicability to apply on new dataset. This limitation raises concerns about how well the FCR can adapt to other biological datasets or experimental setups that deviate from the four testing datasets, which restricts the method’s utility.\n- Interpretability: The model lacks mechanisms for clear interpretation, particularly in linking gene expression profile changes directly to specific treatments. So the decomposition does not provide insights into how gene expression is affected by specific treatments.\n\nI understand that both limitations are significantly challenging and it is unrealistic to tackle both of them in one paper while they highlight critical areas for future research.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper successfully demonstrates the identifiability of treatment-specific, covariate-specific, and interaction-specific representations.\n- The inclusion of the cell type classifier $f_{ct}$ and permutations discriminator $f_{dist}$ are innovative loss components.\n- Sections 6.1 and 6.2 are well-designed to rigorously test the proposed method's ability to maintain the integrity of dimension reduction, decomposition, and conditional independence, validating the model's effectiveness in designing and training.", "weaknesses": "Choice of Evaluation Metric: A significant concern arises from the choice of the $R^2$ metric in section 6.3 for evaluating the performance of the Factorized Causal Representation (FCR) model. $R^2$ is suitable to assess whether a model adequately explains the variation in the response variable. However, it is not necessarily an indicator of the model's predictive accuracy.  Moreover, the near-identical $R^2$ reported for FCR, VCI, CPA, and sVAE (as shown in Table 1) suggest that $R^2$ gives limited insights into the comparative effectiveness of these methods.\n\nThere are some minor misformats, which I noticed while reading.\n- Lines 35, 54: citation format issue\n- Line 294: subscript format issue", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Evalution metrics}$:\nPlease check the **General Response to All Reviewers a**.\n\nFor the performance of $R^2$, we assessed the statistical significance of those results between FCR and second best methods using an paired t-test and showed the p-values with the corresponding sample size of the test set.   \n$\\textbf{Table : p-value for $R^2$ performance}$ \n| Dataset         | p-value (# cells) |\n|-----------------|-------------------|\n| sciPlex         | 0.0006 (295)      |\n| Multiplex-Tram  | 0.0056 (266)      |\n| MultiPlex-7     | 0.0001 (318)      |\n\nThe results indicate that FCR statistically significantly out-performed the baselines across datasets on sci-Plex, multiPlex-Tram and multiPlex-7.\n\nWe will incorporate the discussion of evaluation and t-test results in our final manuscript. \n\n$\\textbf{Generalizability}$: \n\nWe concur with the reviewer that generalizability is a very important topic in the computational biology field. Given the high cost of data collection, we are still limited to test and evaluate our models on the few public large scale single cell drug screening datasets that are available. As availability improves, there will be improved possibilities for benchmarking, especially in the out-of-distribution scenario.\n\n$\\textbf{Interpretability}$:\n\nThere are generally two approaches to interpreting latent representations, which we also plan to explore in future work:\n1. **Latent Traversal**: Since $z_{tx}$ is element-wise identifiable, we can use latent traversal across all dimensions of $z_{tx}$. This involves fixing all latent variables $z_{tx}^i$ to their inferred values for each cell's gene expression and then varying the value of one latent at a time to observe the reconstructed gene expression. This method allows us to see how each dimension of $z_{tx}$ influences gene expression levels.\n2. **Interpretable Decoder System**: We can also employ an interpretable decoder, such as LDVAE [1], which uses a linear decoder. By analyzing the weights of the latent representations assigned to each gene, we can interpret how different representations influence each gene.\n3. **On-going Work**: In the light of the neural additive model [2], we are developing a non-linear interpretable encoder as an extension of LDVAE. Specifically, for each latent dimension, there is a sub-neural network to capture information specific to each latent dimension. A logistic regression layer is added to the end of the neural network. So, by analyzing the logistic regression weights (loading matrix), we can conclude how much one latent variable contributes to the reconstruction of a specific gene. **An illustration of this model can be found in the rebuttal PDF file Figure D**.\n\n[1]Svensson, V., et al.(2020). Interpretable factor models of single-cell RNA-seq via variational autoencoders. Bioinformatics (Oxford, England), 36(11), 3418–3421.   \n[2]Agarwal, R., et al. (2021). Neural additive models: Interpretable machine learning with neural nets. Advances in neural information processing systems, 34, 4699-4711.", "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": "", "extra_scores": {}}, {"review_id": "MS2ZdJEwfc", "reviewer": "Reviewer_yZVk", "summary": "The authors present a novel method for causal representation learning in cellular perturbation settings, called Factorized Causal Representation, in which the authors learn disentangled latents for the cellular covariates, the treatment, and the interactions between them. They provide identifiability results following the methods of recent papers, and demonstrate how an implementation of the method, which uses a number of novel regularizers to enforce various causal constraints, achieves top results amongst comparable methods.", "questions": "- why not evaluate on unseen treatments or cell contexts? Did you try and find the results to be comparatively poor?", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- novel causal representation learning method explicitly learning interactions between covariates and treatments\n- novel regularizers for enforcing conditional independence between sets of variables\n- achieves top cellular response prediction ($R^2$) over comparable conditions", "weaknesses": "- the datasets used have a relatively small number of similar treatments (16 at the most), so it is unclear how well the approach will work when applied to large diverse chemical libraries typical of real-life settings\n- the decision not to evaluate on unseen treatments or cell contexts is surprisingly conservative, since the promise of causal models is in generalizing to unseen domains", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Dataset size}$: \n\nDue to the limited availability of public datasets and the high cost of single-cell sequencing, single-cell datasets with a large number of drugs are not yet available. However, as demonstrated by the theorems in our paper, incorporating more treatments with well-designed controls can enable the latent representations to capture more meaningful and disentangled information, thereby facilitating the algorithm's learning process. As more datasets become available in the future, we plan to conduct extensive benchmarks to evaluate our approach on a broader scale.\n\n$\\textbf{Prediction for unseen perturbations}$ :\n\nPredicting responses to novel treatments is a pivotal and fast-evolving field in drug discovery. However, the biological literature indicates that cellular responses are highly context-dependent [1,2,3]. This complexity poses significant challenges for AI-driven drug discovery, which often struggles to achieve success in clinical trials [4].\n\nOur motivation for developing FCR stems from the need to understand how cellular systems react to treatments and identify conditions that can deepen our understanding of these responses. FCR enables the analysis of drug interactions with covariates and contextual variables. Additionally, predicting cellular responses to new treatments necessitates prior knowledge, such as chemical structure and molecular function, and comparisons with known treatments. Without this context, predictions can be unreliable.\n\nIn this paper, our primary focus is not on predicting responses to unseen treatments. However, given the relevance of this topic, we conducted pilot experiments to showcase the potential future applications of FCR. The MultiTram-Plex and Multiplex-7 datasets share the same cell lines and Trametinib-24 hours treatment, along with other different treatments. By utilizing these two datasets, we established the following experimental settings for unseen prediction scenario:\n\n\n- **Drug Hold-out Setup**: We held out two cell lines, ILAM and SKMEL2, from the Multiplex-Tram dataset, which had been treated with Trametinib for 24 hours. We trained a FCR model using the remaining data, and this model is referred to as $M_h$. We denote the dataset (Multiplex-Tram dataset without Trametinib-24h treated ILAM and SKMEL2) as $D^h$\n- **Prior Knowledge Model**: We trained another FCR model, $M_p$, using the Multiplex-7 dataset, which includes the ILAM and SKMEL2 cell lines treated with Trametinib for 24 hours denoted as $D^p$. We treat this model, $M_p$, as a prior knowledge model.\n- **Transfer MLP**: We extract both $z_{tx}^{p}$ from $M_p$ and $z_{tx}^{h}$ from $M_h$ for $D^p$. Then we trained a 1-layer MLP (the same dimension as $z_{tx}^p$) to transfer $z_{tx}^p$ to $z_{tx}^h$, by minimize the MSE between them.\n- **Contextual Prior Representation**: We extracted the $z_{tx}^p$ representations from model $M_p$ for ILAM and SKMEL2 in holdout set. Then transfer $z_{tx}^p$ by the previous MLP to $\\hat z_{tx}^h$ as a prior contextual embedding.  For the hold-out cell lines ILAM and SKMEL2 in the Multiplex-Tram dataset, we extracted  $z_{x}^h$ from model $M_h$, and Trametinib-24h $z_{t}^h$ from other treated cell lines in $D^h$\n- **Representation Matching**: Then we match the $\\hat z_{tx}^h$ by $z_{x}^p$ similarity on ILAM and SKMEL2 across Multiplex-tram and Multiplex-7 in the prior knowledge model space.\n- **Prediction**: We predicted the unseen 24 hours Trametinib responses for the ILAM and SKMEL2 cell lines in holdout dataset using the formula:\n$\\hat y =g(z_{x}^h,\\hat z_{tx}^h, z_{t}^h)$,where $\\hat z_{tx}^h$ is the corresponding matched prior contextual representation, $ z_t^h$ is the tramnib-24 average value in other cell lines.\n- **Evaluation**: We computed the $R^2$ and MSE for the top 20 differentially expressed genes (DEGs)  based on the predicted values and compared these results with those from the VCI model.The paper’s OOD prediction setups.) \n- $\\textbf{The illustration of datasets and experiments setups can be found in the rebuttal PDF file, Figure A and B}$. \n\n\nWe have the following results:   \n\n$\\textbf{Table: Unseen Prediction Results}$\n| Methods | $R^2$ | MSE  |\n|---------|-----|------|\n| FCR     |   $0.74\\pm 0.03$  | $0.52 \\pm 0.11$      |\n| VCI     |   $0.71 \\pm 0.05$  |   $0.55 \\pm 0.08$   |\n\n[1] Sanjay R. Srivatsan et al. ,Massively multiplex chemical transcriptomics at single-cell resolution.Science.  \n[2] McFarland, J.M., Paolella, B.R., Warren, A. et al. Multiplexed single-cell transcriptional response profiling to define cancer vulnerabilities and therapeutic mechanism of action. Nat Commun.  \n[3] Ramos Zapatero, María et al. “Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses.” Cell vol. 186,25 (2023).  \n[4] Derek Lowe \"AI drugs so far\".", "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": "", "extra_scores": {}}, {"review_id": "EtyQKhJnxf", "reviewer": "Reviewer_q35E", "summary": "The authors propose a method for disentangling into factorized representations dependent on only covariates, only treatment and the interaction between treatment and covariates. Introducing a set of assumptions, the authors prove the identifiability of these disentangled variables based on previous proofs on non-linear ICA. The authors implement the model using variational inference and a set of regularizers and discriminators to ensure the needed independence and disentangling. The authors validate the clustering performance, conditional independencies and conditional cellular response prediction against commonly used methods for perturb-seq prediction and show generally very good performance.", "questions": "- How were the hyperparameters chosen for the different data sets?\n- Maybe I am misunderstanding \"control\", but in the conditional cellular response prediction, why is it necessary to use $z_x^0$ instead of $z_x$ in the function $g$ if $z_x$ is independent of $t$?\n- For prediction validation, what is train/val/test protocol? Is the correlation of y in the train set?", "limitations": "The authors could discuss the limitations (based on assumptions) more directly.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Non-linear interactions between covariates and treatment are very likely and have to be modeled\n- Given a set of additional assumptions, the method has some guarantees of disentangling, which is also shown empirically\n- Understanding (e.g. through clustering) interactions between covariates and treatment is important for precision medicine applications\n- Good performance as shown by authors", "weaknesses": "- Complicated implementation with many loss functions and minmax training, a discussion on how easy to train the method is would be good (dependence of performance on hyperparameters and hyperparameters on data set).\n- The ablation study should include a comparison without the different loss terms to be able to see the importance of the different parts of the model\n- Assumption 4.3 could be quite strong in many settings, if for every covariate every drug response has to be available. A comment on sufficient experiments would be good and how common this is.\n- For some of the assumptions (e.g. substantive changes of the distribution) empirically arguing that this is the case in common data sets would make the real world application of the proof (instead of the purely academic pursuit) more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Hyperparameter selection}$:\nThank you for raising this important question, we appreciate that it is a critical point for evaluating how well the method works. Please check **General Response to All Reviewer b**.\n\n$\\textbf{Ablation studies}$:\n\nIn addition, as the reviewer suggested,  we did the components ablation study (without the different loss terms) on the clustering task (sciPlex dataset) with the NMI score reported.    \n$\\textbf{Table 1: Benchmark on $\\mathbf{z}_x$}$\n\n| Parameters      | $z_x$ on $x$ | $z_{x}$ on ${x, t}$ | $z_x$ on $t$ | Mean |\n|-----------------|---------|--------------|---------|------------|\n| $w_1 = 0$        | 0.583   | 0.606        | 0.061   | 0.417      |\n| $w_2 = 0$         | 0.781   | 0.606        | 0.060   | 0.482      |\n| $w_3 = 0$         | 0.742   | 0.633        | 0.059   | 0.478      |\n| $w_1=w_2=w_3=1$      | 0.786 | 0.612    | 0.061 | **0.486**  |\n\n$\\textbf{Table 2: Benchmark on $\\mathbf{z}_{tx}$}$\n| Parameters      | $z_{tx}$ on $x$ | $z_{tx}$ on ${x, t}$ | $z_{tx}$ on $t$ | Mean |\n|-----------------|----------|---------------|----------|------------|\n| $w_1 = 0$          | 0.731    | 0.595         | 0.068    | **0.465**      |\n| $w_2 = 0$          | 0.725    | 0.556         | 0.037    | 0.439      |\n| $w_3 = 0$          | 0.748    | 0.575         | 0.066    | 0.463      |\n| $w_1=w_2=w_3=1$      | 0.732    | 0.594         | 0.068    | **0.465**      |\n\n$\\textbf{Table 3: Benchmark on $\\mathbf{z}_t$}$\n| Parameters      | $z_t$ on $x$ | $z_{t}$ on ${x, t}$ | $z_t$ on $t$ | Mean |\n|-----------------|---------|--------------|---------|------------|\n| $w_1 = 0$          | 0.333   | 0.389        | 0.069   | 0.264      |\n| $w_2 = 0$          | 0.334   | 0.372        | 0.066   | 0.257      |\n| $w_3 = 0$          | 0.331   | 0.401        | 0.081   | 0.271      |\n| $w_1=w_2=w_3=1$      | 0.334   | 0.393        | 0.088   | **0.272**      |\n\nWhen one of ${w_1, w_2, w_3}$ is set to 0, the remaining weights are set to 1.0. Based on this setup, we draw the following conclusions:\n\nThe results of $z_x$ clustering on $x$, {$t, x$}, and $t$ indicate that\n- $w_1$ (the weight for $L_{sim}$) significantly enhances the performance of $z_x$ clustering on $x$ and $z_t$ clustering on $t$. This improvement is due to $L_{sim}$ increasing the similarity among $z_x$ and $z_t$.\n- $w_2$ (the weight for $L_{ct}$) enhances $z_{tx}$ clustering on {$t$, $x$} and $t$, as well as $z_t$ clustering.\n- $w_3$ (the weight for $L_{dist}$) improves $z_{tx}$ clustering on {$t$, $x$}and also enhances the separation within the spaces of $z_{tx}$, $z_t$, and $z_x$.\n\nIt's important to note that when ${w_1, w_2, w_3} = {1,1,1}$, the performance of $z_{tx}$ clustering on {$t$, $x$}is not as strong as $z_x$ clustering on {$t$, $x$}. However, by increasing $w_1$, $w_2$, and $w_3$, the performance of $z_{tx}$ surpasses that of $z_x$ (see Appendix Figure 7).   \n$\\textbf{Necessary to use $z_{x}^0$ instead of $z_{x}$ in the function g, if $z_{x}$ is independent of $t$}$?:\n\nThank you for raising this question. Yes, directly utilizing $z_x$ might achieve higher $R^2$ score, here we followed CPA, VCI and other methods and replace by $z_x^0$ to perform \"counterfactual\" prediction to further evaluate the learned disentangled representations.\n\n$\\textbf{Assumption 4.3 is strong}$:\n\nAssumption 4.3 can be better understood through the lens of biological experimental design. If we aim to determine whether there is a covariate $x_1$-specific response to treatment $t_1$, we need to conduct two reference experiments: $(x_0, t_1)$ and $(x_0, t_0)$, where $x_0$ represents reference cells and $t_0$ represents control/reference treatments.\n\nBy comparing the treatment effects between the pairs $\\{(x_1, t_1), (x_1, t_0)\\}$ and $\\{(x_1, t_1), (x_0, t_1)\\}$, we can assess whether the observed differences are significant and distinct (as described in Assumption 4.4). Without comparing $(x_1, t_1)$ and $(x_1, t_0)$, we cannot determine if $t_1$ has any effect on $x_1$. Similarly, without comparing $(x_1, t_1)$ and $(x_0, t_1)$, we cannot assess whether $t_1$ has a differential effect on $x_1$ than $x_0$. Only when both comparisons are made can we begin to analyze covariate-specific effects. This comparison allows us to confidently identify any $x_1$-specific response to $t_1$.\n\nThis is a common approach in biological experimental design. For instance, an increasing number of experiments are being conducted for drug screening across various micro-environment organoid systems (covariates), as well as with different drugs, dosages, and time points [3]. The main goal is to discover these covariate-specific cellular responses.\n\n \n$\\textbf{Assumptions empirically arguing that this is the case in common data sets}$:\n\nThe substantive changes assumption (iii) means the different cell lines under the same drug treatment, the gene expression is significantly different (non-trivial). It can be seen that there are large parts of gene expression only influenced by cell types, etc. \n(iv) means under different treatments, the same cell (covariates) have significantly different responses to gene expression. \n\nThese phenomena are all observed in biological literatures [1] [2] [3]. The Figures 3.B,C,D ($\\textbf{Figure D in the rebuttal PDF file}$) in the Trellis paper [3] show evidence that these assumptions are common. \n\n$\\textbf{Limiations}$: Due to the limited space, please check the response to $\\textbf{Reviewer WECd}.$\n\n[1] Sanjay R. Srivatsan et al. ,Massively multiplex chemical transcriptomics at single-cell resolution.   \n[2] McFarland, J.M., Paolella, B.R., Warren, A. et al. Multiplexed single-cell transcriptional response profiling to define cancer vulnerabilities and therapeutic mechanism of action.    \n[3] Ramos Zapatero, María et al. Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses.", "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": "", "extra_scores": {}}, {"review_id": "8ES9uwWuKR", "reviewer": "Reviewer_WECd", "summary": "This paper proposes FCR, a causal VAE-based model that aims to decompose and cluster covariates, treatments, and their interactions in the latent space.", "questions": "Is it possible to enhance other model’s interpretability using the factorized $z_x,z_t,z_{tx}$?", "limitations": "I cannot find a section discussing the limitations of this work, but I believe the noise and hidden variables in the single-cell data will cap the model performance.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper has a solid and rigorous mathematical foundation, and the assumptions are validated after the model is trained.", "weaknesses": "This paper is limited in readability, mainly because of word redundancy, making it hard for the reader to follow the authors’ ideas. Besides, Figure 2 outlines the FCR model, but the encoders and decoders are disconnected, so it is not intuitive to understand how we train components such as $p(z_t|t)$ and $p(z_x|t)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Readability}$:\n\nThank you for your feedback regarding the readability of our manuscript. We appreciate your honest assessment and understand the importance of clear and concise writing. We will carefully review the manuscript to eliminate any word redundancy and improve the overall clarity, ensuring that our ideas are communicated more effectively. \n\n$\\textbf{$p(\\mathbf{z}_t|\\mathbf{t})$ and $p(\\mathbf{z}_x|\\mathbf{x})$ learning}$:\n\nThank you for this question. I’d like to clarify that both $x$ and $t$ are inputs to the neural network, which then encodes them to learn the prior mean and standard deviation for $z_t$ and $z_x$. We use the reparameterization trick to sample $p(z_t|t)$ and $p(z_x|x)$. We will add this explanation under the model figure in the manuscript to provide clearer context.\n\n$\\textbf{Interpretability}$:\n\n\nThere are generally two approaches to interpreting latent representations, which we also plan to explore in future work:\n1. **Latent Traversal**: Since $z_{tx}$ is element-wise identifiable, we can use latent traversal across all dimensions of $z_{tx}$. This involves fixing all latent variables $z_{tx}^i$ to their inferred values for each cell's gene expression and then varying the value of one latent at a time to observe the reconstructed gene expression. This method allows us to see how each dimension of $z_{tx}$ influences gene expression levels.\n2. **Interpretable Decoder System**: We can also employ an interpretable decoder, such as LDVAE [1], which uses a linear decoder. By analyzing the weights of the latent representations assigned to each gene, we can interpret how different representations influence each gene.\n3. **On-going Work**: In light of the neural additive model [2], we are developing a non-linear interpretable encoder as an extension of LDVAE. Specifically, for each latent dimension, there is a sub-neural network to capture information specific to each latent dimension. A logistic regression layer is added to the end of the neural network. So, by analyzing the logistic regression weights (loading matrix), we can conclude how much one latent variable contributes to the reconstruction of a specific gene. **An illustration of this model can be found in the rebuttal PDF file Figure D**.\n \n\n\n$\\textbf{Limitation discussions}$:\n\nThank you for raising this important question. Indeed, this paper has some limitations that we plan to address in future work.\n \n\nFirstly, the generating function 𝑔 used in our model is either deterministic or incorporates Gaussian noise. Future research will investigate more general scenarios with stochastic generating functions, which are more representative of real-world applications. Secondly, while this paper does not include interpretability mechanisms for FCR, we have proposed potential approaches to enhance interpretability, as discussed in the previous section. Lastly, due to the limitations of the available public datasets, we recognize that testing FCR on datasets with a broader range of drugs and more complex covariate scenarios will be an intriguing direction for future work.\n\n\n[1] Svensson, V., et al. (2020). Interpretable factor models of single-cell RNA-seq via variational autoencoders. Bioinformatics (Oxford, England), 36(11), 3418–3421.   \n[2] Agarwal, R., et al. (2021). Neural additive models: Interpretable machine learning with neural nets. Advances in neural information processing systems, 34, 4699-4711.", "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": "", "extra_scores": {}}]

Author rebuttal:
$\textbf{General Response to ALL Reviewers}$:

We sincerely thank all the reviewers for your thorough evaluation of our paper and your recognition of our contributions. We want to address some of the common questions and update new experiment results here, and we will have detailed responses for each reviewer below:

$\textbf{a. Evaluation Metrics}$:

The choice of evaluation metrics for the drug response prediction task is indeed a topic of much discussion. We agree that $R^2$ is well-suited to assess how well a model explains the variation in the response variable.

We chose $R^2$ as our evaluation metric for two main reasons:
- Previous works, including **VCI**, **CPA**, and **sVAE**, have also used $R^2$ as a standard evaluation metric, providing a basis for comparison.
- Single-cell transcriptomics data is often noisy, sparse, may contain batch effects, and highly related to library size making direct evaluation of gene prediction accuracy challenging. Given the noisiness of this data, researchers commonly use relative gene enrichment analysis to compare different cell types.

Other evaluation metrics could be:
- the Mean Squared Error (**MSE**) of the top differentially expressed genes (**DEGs**) for post-treatment [1].
- the Spearman correlation between predicted genes and ground truth across all the cells. 
- the cosine similarities between the predicted genes and ground truths. 

We added MSE results for the top 20 DEGs. These 20 genes are selected for showing statistically significant differences in expression levels for each cell line with drug treatments compared to control samples. The same procedures are also carried out in [1]. Note here, that we didn’t compare with CINEMA-OT and scGEN because they are only for binary treatments. 

$\textbf{Table 1: MSE (Top 20 DEGs)}$

| Datasets        | FCR (ours)         | VCI         | CPA         | sVAE        |
|-----------------|-------------|-------------|-------------|-------------|
| sciPlex         | **0.12 (0.03)** | 0.15 (0.04) | 0.15 (0.04)  | 0.16 (0.05)  |
| multiPlex-tram  | **0.10 (0.07)** | 0.13 (0.08) | 0.13 (0.08)  | 0.14 (0.07)  |
| multiPlex-7     | **0.18 (0.07)** | 0.21 (0.08) | 0.22 (0.08)  | 0.23 (0.07)  |
| multiPlex-9     | **0.24 (0.06)** | 0.27 (0.04) | **0.23 (0.06)** | 0.26 (0.05)  |

FCR still outperforms the competing methods in three datasets while CPA achieves slightly better results on multiPlex-9. 

$\textbf{b. Hyperparameter Selection}$: 

We split the data into four datasets: train/validation/test/prediction, following the setup from previous works [2]. First we hold out the 20% of the control cells for the final cellular prediction tasks (pred). Second, we hold 20% of the rest of data for the task of clustering and statistical test (test). Third, the data excluding the prediction and clustering/test sets are split into training and validation sets with a four-to-one ratio.

For the hyperparameter tuning procedure, conduct the exhaustive hyperparameter grid search with n_epoch=100 on the loss assessed on the validation data. The hyperparameter search space is

$w_1 = [0.5, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10]$
$w_2 = [0.5, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10]$
$w_3 = [0.1, 0.3, 0.5, 0.7, 0.9, 1, 3, 5, 7, 10]$

We will make sure to include this in our final manuscript.

$\textbf{c. Simulation Study}$:

To address **Reviewer vnnM**'s question on synthetic data experiments to demonstrate the identifiability of $z_{tx}$. 

By following the [3,4] protocol of simulation data. For the simplicity of simulation,  we set the dimension of $z_x$ and $z_t$ equal to 1, and the $z_{xt}$ dimension as 4. We have the following setups. For the simplicity of the simulation, we output the $y$ (dimension is 96) as real numbers instead of count data with a sample number of 5000.

$t \sim \textrm{Unif}([1,2,3])$,
$x \sim \textrm{Unif}([100, 1000, 5000])$,
$z_x \sim \textrm{Normal}(x/2, 1)$,
$z_t \sim \textrm{Normal}(t/2,1)$,
$z_{tx} \sim \textrm{Normal}(x*t, I_4)$,
$g$ is a 2-layer MLP with the Leaky-ReLU activation function [3].
To measure the component-wise identifiability of the changing components, we compute the Mean Correlation Coefficient (MCC) between $z_{tx}$ and $\hat z_{tx}$ .  A higher MCC indicates a higher extent of identifiability, and MCC reaches 1 when latent variables are perfectly component-wise identifiable.  We computed the mean correlation coefficient (MCC) between the original sources and the corresponding latents sampled from the learned posterior. We followed iVAE [3], and first calculated all pairs of correlation coefficients between source and latent components. We then solve a linear sum assignment problem to assign each latent component to the source component that best correlates with it, thus reversing any permutations in the latent space. The UMAP projections of the ground truth $z_{tx}$ and estimated $\hat z_{tx}$ which showed perfect separation are shown in Figure C of the PDF file.

$ \textbf{Table 2 : Simulation Results on MCC}$
| Methods    | MCC  |
|------------|------|
| FCR        | **0.91 (0.03)**  |
| beta-VAE   |  0.38 (0.12)    |
| iVAE       |    0.77 (0.07)  |
| factor-VAE |    0.37 (0.08)  |

$\textbf{d. Unseen Prediction}$:
For the unseen prediction experiment setup and results, please check the response to **Reviewer yZVk**.

$\textbf{Note: We have new figures in the attached PDF file}$

$\textbf {References}$:  

[1] Roohani, Y., et al. (2024). Predicting transcriptional outcomes of novel multigene perturbations with GEARS. Nature Biotechnology.
[2] Lotfollahi M, et al. Predicting cellular responses to complex perturbations in high-throughput screens. Mol Syst Biol.        
[3] Khemakhem, Ilyes, et al. "Variational autoencoders and nonlinear ica: A unifying framework." International conference on artificial intelligence and statistics. PMLR, 2020.   
[4] Lopez, Romain, et al. "Toward the Identifiability of Comparative Deep Generative Models." Causal Learning and Reasoning. PMLR, 2024.

Meta-review: The method proposes a method for learning disentangled causal representations of covariates, treatment and their interaction 
The learning objective enforces the conditional independence assumptions and is theoretically well-justified.
Performance is tested in and extensive evaluation showing that the method performs significantly better than a large set of baselines. The paper is well-written and tackles a relevant and hard problem. The reviewers criticise minor issues like complexity of the model and choices in the evaluation scheme.

All five reviewers vote towards acceptance.
Decision: Accept (poster)