Title: Learning Identifiable Factorized Causal Representations of Cellular Responses
Venue: NeurIPS 2024
Accepted: yes
Keywords: Single Cell

Computational Biology

Causality
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.

1

Abstract

The study of 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 is difficult because the response of a cell to perturbations essentially depends on its biological context (e.g., genetic background or cell type). For example, while discovering therapeutic targets, one may want to enrich for drugs that specifically target a certain cell type. This challenge emphasizes the need for methods that explicitly take into account potential interactions between drugs and contexts. Towards this goal, we 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 are disentangled, comprised of covariate-specific ( z x ), treatment-specific ( z t ), and interaction-specific ( z tx ) blocks. Based on recent advances in 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 it outperforms state-of-the-art baselines in various tasks across four single-cell datasets. The code is available on GitHub ( https://github.com/Genentech/fcr ).

1 Introduction

The recent experimental capabilities reached by single-cell perturbation technologies open up new opportunities for characterizing cellular behaviors (Dixit et al., 2016). For example, high-throughput screening of chemical and genetic perturbations highlighted vulnerabilities in numerous cancer cell lines (McFarland et al., 2020). Identifying these vulnerabilities is crucial for pinpointing therapeutic targets, facilitating drug discovery, and furthering our understanding of gene functions (Srivatsan et 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 et 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 to consider interaction effects between treatments and contextual covariates while modeling gene expression outcomes (Zapatero et al., 2023).

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 identifiable 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 a set of sufficient conditions for its identifiability, extending the work of Khemakhem et al. (2020) to the case of interacting covariates. We then present an implementation of our FCR method that builds upon the variational auto-encoder framework (Kingma and Welling, 2014) as well as adversarial regularization (Ganin et al., 2016). We demonstrate that FCR not only effectively identifies interactions but also surpasses state-of-the-art methods in various tasks across four single-cell datasets.

2 Related Work

Learning Representations of Cellular Responses Learning representations of single-cell data is a powerful framework, with demonstrated impact in tasks such as imputation (Lopez et al., 2018), clustering (Trapnell et al., 2014; Zhu et al., 2019; Alquicira-Hernandez 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 al., 2019; Rampasek et al., 2019; Lotfollahi et al., 2021; Lopez et al., 2023; Bunne et al., 2023; Zapatero et 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. (2023) 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, a critical point for understanding the effects of drugs on biological systems.

Identifiable non-linear Independent Component Analysis models Afield where disentanglement is most important is non-linear Independent Component Analysis (non-linear ICA) (Hyvärinen and Pajunen, 1999), where information from a 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 et 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 in the Variational Autoencoder (VAE) literature (Kingma and Welling, 2014) such as the betaVAE (Higgins et al., 2016) and factorVAE (Kim and Mnih, 2018) that both address disentanglement learning. Indeed, these latter models lack theoretical foundation regarding the identifiability of their parameters or latent variables (Esmaeili et 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 ∈ X ⊆ R d be a vector of covariates representing intrinsic attributes of single cells, such as cell type, tissue of origin, or patient information.

Let t ∈ T ⊆ R p represent the treatment or intervention applied to single cells and let y ∈ Y ⊆ R k denote the gene expression levels (outcome).

3.1 Generative Model

We introduce a low-dimensional latent vector z ∈ Z ⊆ R n that encodes cellular states after treatment t and in biological context x . We assume a block structure for z = [ z x , z tx , z t ] with dimension n = n x + n tx + n t . Here, z x captures the effects of contextual covariates x , z t represents the direct effects of the treatment t , and z tx encodes the interaction effects between both terms.

More precisely, the generative model is specified as follows. Latent representation z x is generated from x according to distribution z x ∼ p z x | x ( z x | x ) , z t is generated from t following distribution z t ∼ p z t | t ( z t | t ) , and z tx from both x and t following distribution z tx ∼ p z tx | t , x ( z tx | t , x ) . The gene expression outcome vector y is then deterministically generated y = g ( z x , z tx , z t ) ,

where g is a smooth mixing function. A graphical representation of this generative model appears in Figure 1. For the control group (no treatment), the outcome is noted as y 0 , and the representation as z 0 = [ z 0 x , z 0 tx , z 0 t ] .

Learning each element in this triplet of latent variables is a sound approach for unravelling interaction effects. Indeed, z x captures covariate-specific patterns that are invariant with respect to the perturbations, while also capturing the essential biological attributes tied to the cellular covariates. z t captures the intrinsic effects of the treatments, irrespective of the covariates. z tx 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 to treatments, reflecting the diversity and specificity of biological responses.

We note that our model does not take into account observation noise, since g is a deterministic function in our assumption. However, our theory may be readily extended to incorporate Gaussian observation noise (Khemakhem et al., 2020), as well as 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 et al. (2021) and Kong et al. (2022). Throughout this section, z ∈ Z denotes a (random) latent vector and y ∈ Y denotes an observed vector. g : Z → Y is an unknown mixing function such that y = g ( z ) .

Definition 3.1 ( Component-wise Identifiability ) . We say that latent vector z is identifiable from data y if for any other latent vector ˆ z such that g (ˆ z ) and g ( z ) are equal in distribution, z and ˆ z are equal up to a permutation of the indices, and deformation of each component by invertible scalar functions. More precisely, there exists a permutation π of { 1 , . . . , n } , and invertible scalar functions h j such that for all j ∈ { 1 , . . . , n } :

where z j and ˆ z π ( j ) are the π ( j ) -th components of z and ˆ z respectively.

Definition 3.2 ( Block-wise Identifiability ) . For 1 ≤ n 1 < n 2 ≤ n , we denote as z [ n 1 : n 2 ] ∈ R n 2 -n 1 +1 the block of z from index n 1 to n 2 . We say that latent vector z [ n 1 : n 2 ] is block-identifiable from data y if for any other latent vector ˆ z such that g (ˆ z ) and g ( z ) are equal in distribution, z [ n 1 : n 2 ] and ˆ z [ n 1 : n 2 ] are equal up to an invertible function h :

where ˆ z [ n 1 : n 2 ] is the corresponding block in the estimated vector ˆ z .

4 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-wise identifiability of the interaction variable z tx under the assumption of sufficient experimental variability (Khemakhem et al., 2020) (Section 4.1). Then, we demonstrate the block-identifiability of z t and z x by exploiting their invariance with respect to x and t , respectively (Section 4.2). These guarantees are important, as they ensure that the obtained latent variables will have desirable semantics. For example, given our theoretical results, the obtained interaction embedding ˆ z tx must be reflective of the ground-truth interactions z tx only, and not of any of the other latent variables.

4.1 Component-wise identifiability of z tx

Our proof relies on three technical assumptions. Two are classical from the nonlinear ICA literature, and the last one relates to the observability of a complementary set of experiments for identifiability of interactions.

Assumption 4.1. The probability density function of the prior distribution for the latent variables is smooth and positive, i.e. p z | t , x ( z | t , x ) > 0 for all ( z , t , x ) ∈ Z × T × X .

Assumption 4.2. The components of z are conditionally independent given t and x .

Assumption 4.3. ( Experimental Sufficiency ) There exist at least 2 n tx +1 distinct values for the vector [ t , x ] in the experimental design. One such setting can be referred to as a control condition and is noted as [ t 0 , x 0 ] . Additionally, for any non-control environment ( t i , x i ) for i ∈ { 1 , . . . , 2 n tx } , we assume there always exist corresponding switched experiments under the settings ( t 0 , x i ) and ( t i , x 0 ) .

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 ( z tx , t , x ) as the vector:

where q i denotes the logarithm of probability density p z i | t , x for component z i . If in addition to the assumptions 4.1, 4.2, 4.3, we assume that the 2 n tx vectors

are linearly independent then z tx 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 t and x the vector of interest becomes the null vector. In this trivial case, the theorem's assumption is never satisfied ( z tx indeed has no purpose in that particular model). However, under a 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 x and z t

To prove the block-wise identifiability of z x and z t , we exploit their invariance properties: z t remains unchanged across different values of x , while z x is stable across variations in t . This invariance allows us to distinguish these blocks from the interaction terms z tx . Additionally, because this invariance reflects latent features' stability despite perturbations, its utilization can enable deeper biological insights and interpretability. For instance, z x might represent stable cellular characteristics

that persist across different treatments, while z t could capture consistent treatment effects across various cell types.

We now state our main identifiability result for z x and z t .

Theorem 4.5. We follow Assumptions 4.1, 4.2, 4.3, and the one from Theorem 4.4. We note as S ( Z ) the set of subsets S ⊆ Z of Z that satisfy the following two conditions:

(i) S has nonzero probability measure, i.e. P ( z ∈ S | t = t ′ , x = x ′ ) > 0 for any t ′ ∈ T and x ′ ∈ X .

(ii) S cannot be expressed as A z x ×Z tx ×Z t for any A z x ⊂ Z x or as Z x ×Z tx × A z t for any A z t ⊂ Z t .

We have the following identifiability result. If for all S ∈ S ( Z ) , there exists ( t 1 , t 2 ) ∈ T × T and x ∈ X such that

̸

and there also exists ( x 1 , x 2 ) ∈ X × X and t ∈ T such that

̸

then z t and z x are block-wise identifiable.

The proof appears in Appendix A, and is adapted from Kong et al. (2022). Our main assumption is that the conditional distribution p z | t , x ( z | t , x ) undergoes substantive changes when spanning different treatments t and covariates x . When treatments differ markedly from each other in their mechanisms and effects, the probability distributions of the latent variables conditioned on these treatments are 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:

A variational inference approach to estimate the representations from our FCR model. Our model and inference architecture is specifically designed to learn disentangled representations z x , z tx , z t (Section 5.1).

A regularization method that enforces independence between z x and t , and encourages variability of z t with respect to x (Section 5.2).

A second regularization technique to ensure that the conditional independence properties z x ⊥ ⊥ z tx | x and z t ⊥ ⊥ z tx | t are satisfied (Section 5.3).

The main computational structure of the model is illustrated as a schematic in Figure 2.

5.1 Model Specification and Variational Inference

Model Specification We parameterize the probability distributions as follows:

where all the above functions are parameterized by neural networks.

To prevent z tx from having trivial dependencies with respect to t and x , we explicitly encourage its prior to capture interactions between x and t by designing the functions f t,x µ and f t,x σ to be of the form:

(a)

(c)

I.

p(zt|t)

→

MLP

[kt(t)

p(ztIt)

p(Za|x)

X

MLP

p(Zx|x)

p(Zta|t,x)

KLdivergence

q(z,1t, y) q(zlt,x, y) q(zxx, y)

p(Zta|t,x)

II.

Zt

fdist

Ztu

(b)

y

q(zlt, y) q(ztx1t,x, y) q(z1x, y)

Zz

fdisx

Decoder

Zt

sim

sim

0

III.

Zt

y

q(z0, y)q(zt10,x, y)q(z1x, y

fct

Zta

where k x ( x ) and k t ( t ) represent the embeddings for the cellular covariates and treatments, respectively, while ⊙ 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 to the posterior, following the approach used in previous causal representation learning work (Khemakhem et al., 2020). We consider the following mean-field variational approximation to the posterior distribution:

Following the graphical model from Figure 1, the Evidence Lower Bound (ELBO) is derived as:

where θ and ϕ denote the parameters of the generative model and the inference networks, respectively. D KL denotes the Kullback-Leibler divergence between two probability distributions. For simplicity, we omit ϕ and θ as well as script notations in the following sections, wherever appropriate. The derivation of the ELBO appears in Appendix B.

The variational distributions defined in Equation 11 are parameterized as follows:

where all the functions introduced above are parameterized by neural networks.

5.2 Causal Structure Regularization

We exploit both the variability of z t and the invariance of z x when comparing control and treated groups that share the same covariates. Specifically, our goal is to enforce the resemblance between z x and z 0 x while reducing the congruence of z t and z 0 t . Towards this end, we first add the following score as a regularizer,

where sim ( · ) denotes the cosine similarity. Second, we introduce a classifier f ct to predict the treatments t from [ z t , z tx ] and [ z 0 t , z 0 tx ] as follows,

The predicted treatment probability vector ˜ t is then used for the computation of a cross-entropy loss

5.3 Permutation Discriminators

We want to ensure the conditions of Assumption 4.2 and Theorem 4.4, specifically that z tx ⊥ ⊥ z t | t and z tx ⊥ ⊥ z x | x . Towards this goal, we use the following proposition, establishing a connection between exchangeability and conditional independence.

Proposition 5.1. (Bellot and van der Schaar, 2019) Let X,Y and Z be three random variables. Under the assumption of X ⊥ ⊥ Y | Z , we have the samples ( X i , Y i , Z i ) i M =1 and permuted samples ( X π ( i ) , Y i , Z i ) i M =1 with a permutation function π . The corresponding statistics ρ i of ( X i , Y i ) i M =1 and ρ π ( i ) of ( X π ( i ) , Y i ) i M =1 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 z x , z tx and z t , z tx . First, we initially permute z tx within the triplet ( z ( j ) x , z ( j ) tx , x ( j ) = x i ) j M =1 to yield ( z ( j ) x , z π ( j ) tx , x ( j ) = x i ) j M =1 . Then, we train a binary classifier (the discriminator) to predict whether samples have been permuted or not. We denote the permutation label as l and predict it as

where ˆ z tx could be permuted or non-permuted z tx samples. If z x and z tx are indeed independent given x , the discriminator should be unable to distinguish between the permuted and the original samples. For each discriminator, we add a regularization term that consists of the cross-entropy loss

We proceed similarly to make sure that z t and z tx are independent conditionally on t .

5.4 Objective function

Finally, we specify the overall loss for our model as

where ω 1 , ω 2 , ω 3 > 0 are hyperparameters. To concurrently train both the representations and the discriminators, we employ an adversarial training approach as follows,

The training procedure and the procedure used for hyperparameters selection are detailed in Appendix C and D, respectively.

6 Experiments

In this section, we are pursuing three primary objectives. First, we seek to validate the FCR method's proficiency in capturing the designated causal structure within the latent space through both clustering analysis (Section 6.1) and statistical testing of independence (Section 6.2), respectively. Second, we evaluate the method's efficacy in 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 x , and for that reason we do not perform experiments to predict cellular responses to unseen treatments and cell types (covariates).

(a)

(b)

(c)

1.0

工

0.45

0.9

白

0.25

工

0.40

工

0.8

0.6

0.7

NMI

0.35

0.30

+

T

点

NMI

0.15

0.20

0.25

0.5

0.4

0.20

0.10

0.15

0.3

0.2

0.10

+

0.05

SVAE

VCI

VA

SVAE

CPA

VCI

SVAE

CPA

yCI

4+

factorVAE

beta

factorVAE

beta

factorVAE

betaVAB

methods

methods

methods

Datasets To evaluate the efficacy and robustness of the FCR method, we conducted our study on four real single-cell perturbation datasets (Appendix E). The first of these is 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: 0 nM (control), 10 nM, 100 nM, 1 µM, and 10 µM. The subsequent three datasets are sourced from (McFarland et 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 a 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 a series of nine treatments.

Baselines We benchmarked our method against several established representation learning methods: (1) scVI (Lopez et al., 2018), (2) iV AE (Khemakhem et al., 2020), (3) β VAE (Higgins et al., 2016), (4) factorVAE (Kim and Mnih, 2018), (5) VCI (Wu et al., 2023), (6) CPA (Lotfollahi et al., 2023), (7) scGEN (Lotfollahi et al., 2019) (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 a 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 a total of ten times for each baseline using different random subsets to yield the best results. Specifically, CINEMA-OT and scGEN address only binary treatments/perturbations, so they are 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 F.

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 x , 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 x , t , 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 of 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 x , z tx , and z t , as well as on the variables obtained from baseline methods.

Our results highlight that z x has superior performance in clustering on covariates x compared to all other available latent representations (Figure 3). Similarly, z t performs best for clustering on treatments t . Finally, z tx outperforms all other methods when clustering jointly on t and x , showing that it faithfully represents the combined features of both x and t . Specifically, for the sciPlex datasets,

VAE

(a)

(b)

(c)

(d)

0.35

0.08

0.4

0.6

0.30

0.06

0.3

0.25

Pvalue

Pvalue

alue

00000

0.4

0.2

Pval

0.20

JISH

0.04

0.15

0.2

0.1

0.10

0.02

0.05

0.0

0.0

0.00

0.00

SCVI

VA

SCVI

betay

ZZ

x,tx

+

十

十

betal

actorl

actorl

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 on distinct cell lines until a maximal concentration of 10 µM was reached. We report similar results for the other datasets in 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 investigations focused on the following relationships: (1) z x ⊥ ⊥ t | x ; (2) z t ⊥ ⊥ x | t ; (3) z t ⊥ ⊥ z tx | t ; (4) z x ⊥ ⊥ z tx | x . In conjunction with these tests, we also employed the Hilbert-Schmidt Independence Criterion (HSIC) (Gretton et al., 2005) to evaluate the (marginal) independence between: (a) z t and t ; (b) z x and x . Our aim was two-fold. First, we wanted to assess whether the factors in our latent space complied with the necessary conditional independence statements, corroborating a 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 F.4). We first examined the results of testing for conditional independence statements z x ⊥ ⊥ t | x and z t ⊥ ⊥ x | t (Figure 4ab). In this experiment, all the baseline methods produced p-values smaller than 0.05. This implies a 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 z x ⊥ ⊥ z tx | x and z t ⊥ ⊥ z tx | t (Figure 4c). 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 x ⊥ ⊥ x , z x ⊥ ⊥ t , z x ⊥ ⊥ R , as well as z t ⊥ ⊥ x , z t ⊥ ⊥ t , and z t ⊥ ⊥ R , where R represents simulated random vectors (Figure 4d). Contrasting our representations with randomly simulated vectors, we observe that both z x and x , as well as z t and t , have HSIC values far from zero, indicating a high dependence. The contrast in results between our approach and the baseline methods highlights FCR'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. Consequently, accurately predicting outcomes for novel drugs or cell lines becomes challenging without a careful consideration of the similarity in z tx and how t interacts with x . Unlike previous literature,

we do not conduct experiments to predict cellular responses to unseen treatments and cell types (covariates). This decision is based on extensive biological research showing that responses to covariates are context-specific (McFarland et al., 2020; Srivatsan et al., 2020). 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 on predicting cellular responses (expression of 2000 genes) in control cells subjected to drug treatments. It is important to note that our comparative analysis is confined to CPA, VCI, sVAE, scGEN and CINEMA-OT as they are uniquely tailored for this task. We utilize FCR to extract control's [ z 0 x , z 0 tx , z 0 t ] and corresponding experiments' [ z x , z tx , z t ] , then use the decoder g to predict the gene expression level as ˆ y = g ( z 0 x , z tx , z t ) . We measure the R 2 score. From our results (Table 1), we observe that FCR generally outperforms other baselines across the first three datasets. However, CPA performs the best on 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, scGEN and CINEMA-OT, which are designed for binary perturbations, tend to underperform in these tasks.

7 Discussion

This paper aimed to 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 z tx , z x and z t . 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 z x , z t or z tx . Such advancements are expected to significantly contribute to the progress of precision medicine.

Acknowledgments

The authors would like to extend gratitude to Gregory Barlow for his contribution to the graphic design.

References

Alquicira-Hernandez, S., 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 Neural 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, G. (2023). Learning single-cell perturbation responses using neural optimal transport. Nature Methods , pages 1-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. P., Jerby-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(7):1853-1866.

Dong, M., Wang, B., Wei, J., de O. Fonseca, A. H., Perry, C. J., Frey, A., Ouerghi, F., Foxman, E. F., Ishizuka, J. J., Dhodapkar, R. M., et al. (2023). Causal identification of single-cell experimental perturbation effects with cinema-ot. Nature Methods , 20(11):1769-1779.

Esmaeili, B., Wu, H., Jain, S., Bozkurt, A., Siddharth, N., Paige, B., Brooks, D. H., Dy, J., and van 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 omics data. Nature Biotechnology , 40(2):163-166.

Gretton, A., Bousquet, O., Smola, A., and Schölkopf, B. (2005). Measuring statistical dependence with hilbert-schmidt norms. In International Conference on Algorithmic Learning Theory , pages 63-77.

Grün, D., Kester, L., and van Oudenaarden, A. (2014). Validation of noise models for single-cell transcriptomics. Nature Methods , 11(6):637-640.

Hetzel, L., Boehm, S., Kilbertus, N., Günnemann, S., Lotfollahi, M., and Theis, F. J. (2022). Predicting cellular responses to novel drug perturbations at a single-cell resolution. Advances in Neural Information Processing Systems , 35:26711-26722.

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-439.

Hyvärinen, A., Sasaki, H., and Turner, R. (2019). Nonlinear ICA using auxiliary variables and generalized contrastive learning. In International Conference on Artificial Intelligence and Statistics , 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. Y . H., and Yang, P. (2019). Impact of similarity metrics on single-cell RNA-seq data clustering. Briefings in bioinformatics , 20(6):2316-2326.

Kingma, D. P. and Welling, M. (2014). Auto-encoding variational Bayes. In International Conference on 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 on Learning Representations .

Kong, L., Xie, S., Yao, W., Zheng, Y., Chen, G., Stojanov, P., Akinwande, V., and Zhang, K. (2022). Partial disentanglement for domain adaptation. In International Conference on Machine Learning , pages 11455-11472.

Lachapelle, S., Rodriguez, P., Sharma, Y., Everett, K. E., Le Priol, R., Lacoste, A., and Lacoste-Julien, S. (2022). Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ICA. In Conference on Causal Learning and Reasoning , pages 428-484.

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 , 15(12):1053-1058.

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-537.

Lotfollahi, M., Wolf, F. A., and Theis, F. J. (2019). scGen predicts single-cell perturbation responses. Nature methods , 16(8):715-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, J. (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.VAE: improving drug response prediction via modeling of drug perturbation effects. Bioinformatics , 12(6):3743-3751.

Roohani, Y., Huang, K., and Leskovec, J. (2024). Predicting transcriptional outcomes of novel multigene perturbations with GEARS. Nature Biotechnology , 42(6):927-935.

Srivatsan, S. R., McFaline-Figueroa, J. 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 at single-cell resolution. Science , 367(6473):45-51.

Traag, V. A., Waltman, L., and Van Eck, N. J. (2019). From Louvain to Leiden: guaranteeing well-connected communities. Scientific reports , 9(1):5233.

Trapnell, C., Cacchiarelli, D., Grimsby, J., Pokharel, P., Li, S., Morse, M., Lennon, N. J., Livak, K. J., Mikkelsen, T. S., and Rinn, J. L. (2014). The dynamics and regulators of cell fate decisions are 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:16451-16467.

Wagner, A., Regev, A., and Yosef, N. (2016). Revealing the vectors of cellular identity with single-cell genomics. Nature Biotechnology , 34(11):1145-1160.

Wu, Y., Barton, R., Wang, Z., Ioannidis, V. N., De Donno, C., Price, L. C., V oloch, 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., Vlckova, P., Nattress, C., Qin, X., Claus, J., et al. (2023). Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses. Cell , 186(25):5606-5619.

Zhang, J., Greenewald, K., Squires, C., Srivastava, A., Shanmugam, K., and Uhler, C. (2023). Identifiability guarantees for causal disentanglement from soft interventions. In 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.

Appendix

A Proof of Theorems

A.1 Proof of Theorem 4.4

Theorem 4.4 . Let us first define v ( z tx , t , x ) as the vector:

where q i denotes the logarithm of probability density p z i | t , x of component z i . If in addition to the assumptions 4.1, 4.2, 4.3, we assume that for the 2 n tx vectors

are linearly independent, then z tx is component-wise identifiable.

Proof. This proof proceeds in three main steps:

Derivation of the Fundamental System of Equations: We derive a crucial relationship between the true and estimated latent variables by applying the change of variables formula, and differentiating the equality of observed data distributions.

Isolation of Interactive Components: We isolate the terms relevant to z tx by strategically comparing equations for different pairs of ( x , t ) values and subtracting them.

Establishing Component-wise Identifiability: We analyze the structure of the resulting equations and the Jacobian of the transformation between true and estimated latent variables to establish the component-wise identifiability of z tx .

Step 1 (Derivation of the Fundamental System of Equations) Let us assume there exists another latent representation ˆ z that yields the same data distribution than the ground-truth variables z , for all t ∈ T and x ∈ X . Specifically, we have:

Given our assumption of noiseless observations, it is equivalent to equality in distribution of the mixed variables:

which, after a change of variable, is equivalent to:

where g -1 : Y → Z denotes the invertible generating function, and h := g -1 ◦ ˆ g is the transformation between the true latent variable and estimated one. | J g -1 | denotes the determinant of the Jacobian matrix of g -1 .

̸

Because g is invertible, | J g -1 | = 0 . Using this fact, we obtain the equivalent condition:

According to the independence relations in the data generating process, we have

Rewriting the notation q i := log p z i | t , x and ˆ q i := log p ˆ z i | t , x yields:

Applying the change of variables formula to Equation 24 yields

where J h -1 and J h are the Jacobian matrix of the transformation associated with h -1 and h , respectively. We now adopt the following notations,

̸

Then, we may differentiate Equation 28 with respect to ˆ z k and ˆ z q where k, q ∈ { 1 , . . . , n } and k = q . Doing so, we obtain the following fundamental system of equations. For any x ∈ X , t ∈ T , for all ( k, q ) ∈ { 1 , . . . , n } 2 such that k = q :

̸

Step 2 (Isolation of Interactive Components) We may decompose the sum present on the lefthand-side of Equation 30 across the different block of latent variables, obtaining the following equality:

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 h .

First, we apply this strategy to the pair { ( x , t 0 ) , ( x 0 , t 0 ) } , for any treatment x (we assume the existence of a reference treatment t 0 and context x 0 ). Substituting according to Equation 31 into Equation 30, and applying it to ( x , t 0 ) , we obtain:

Proceeding similarly for ( x 0 , t 0 ) , we obtain:

Then taking the difference between Equation 32 and Equation 33 yields,

We apply the same principle to the pair { ( x 0 , t ) , ( x 0 , t 0 ) } . We therefore get:

and,

Taking similarly the difference between Equation 35 and Equation 36 yields:

Finally, we consider the pairs { ( x , t ) , ( x 0 , t 0 ) } , for which we obtain the following:

and

Once again taking the difference between Equation 38 and Equation 39 yields:

As a 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:

Finally, we subtract Equation 41 and Equation 37:

This last equation may be rearranged as:

Step 3 (Establishing Component-wise Identifiability) Given the assumption of linear independence in the Theorem, the linear system is a 2 n tx × 2 n tx full-rank system. Therefore, the only solution is:

̸

h ( · ) is a smooth function over Z and its Jacobian can written as:

̸

Note that a ′ i, ( k ) a ′ i, ( q ) = 0 implies that for each i ∈ { n x +1 , . . . , n x + n tx } , a ′ i, ( k ) = 0 for at most one element k ∈ [ n ] . As a consequence, there's only a single non-zero entry in each row indexed by i ∈ { n x +1 , . . . , n x + n tx } in the Jacobian matrix J h . Further strengthening this argument, the invertibility of h requires J h to be full-rank, suggesting there's precisely one non-zero component in each row of matrices D , E , and F .

It means that each of z i ∈ z tx for i ∈ { n x +1 , . . . , n x + n tx } is attributed to at most one of the ˆ z . And because the other ˆ z that are not in the block ˆ z tx do not have dependencies with both t and x , it must be that the non-zero coefficient is in the block E . Therefore D = 0 and F = 0 . This indicates the invertibility of h i for every i in the range { n x +1 , . . . , n x + n tx } . In conclusion, z tx are element-wise identifiable, albeit subject to permutations and component-wise invertible transformations.

A.2 Proof of Theorem 4.5

Theorem 4.5: We follow Assumptions 4.1, 4.2, 4.3, and the one from Theorem 4.4. We note as S ( Z ) the set of subsets S ⊆ Z of Z that satisfy the following two conditions:

(i) S has nonzero probability measure, i.e. P ( z ∈ S | t = t ′ , x = x ′ ) > 0 for any t ′ ∈ T and x ′ ∈ X .

(ii) S cannot be expressed as A z x ×Z tx ×Z t for any A z x ⊂ Z x or as Z x ×Z tx × A z t for any A z t ⊂ Z t .

We have the following identifiability result. If for all S ∈ S ( Z ) , there exists ( t 1 , t 2 ) ∈ T × T and x ∈ X such that

̸

and there also exists ( x 1 , x 2 ) ∈ X × X and t ∈ T such that

̸

then z t and z x are block-wise identifiable.

Proof. The proof of the block-wise identifiability of z x and z t are similar, so here we focus the proof about the identifiability of z x . Our proof draws parallels to the approach in Kong et al. (2022), but our context specifically pertains to the multi-conditions distributions scenarios. Throughout, we use the notations z + t = [ z tx , z t ] , ˆ z + t = [ˆ z tx , ˆ z t ] for simplification. This proof is comprised of four major steps.

Integral characterization of domain invariance . We first leverage properties of the generating process and the marginal distribution matching condition to provide a characterization of the invariance of a block of latent variables by a mixing function using an integral condition.

Topological characterization of invariance . We derive equivalence statements for domain invariance of functions.

Proof of invariance by contradiction . We prove the invariance statement from Step 2 by contradiction. Specifically, we show that if ˆ z x depended on z + t , the invariance derived in Step 1 would break.

Block-wise identifiability of z t and z x . 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 a reminder to the reader, g : Z → Y denotes the ground-truth mixing function, and ˆ g : Z → Y denotes the learned mixing function. We assume that both g and ˆ g are invertible, so that their reciprocal functions are well defined. In particular, we denote by ˆ g -1 1: n x : Y → Z x the estimated transformation from the observation to the covariate-specific block of the latent space Z . We seek to find an integral characterization of the invariance of the learned function ˆ g on the domain Z x .

We seek to derive a condition for which the distribution of latent variables on Z x will be unchanged when the treatment t is changed (it will however depend on x ). Let S ⊂ Z x designates a subset of Z x and x ∈ X be any context. We seek to characterize the following condition:

This condition can be written equivalently using the pre-image set of S :

where ( ˆ g -1 1: n x ) -1 ( S ) ⊆ Y is the set of estimated observations ˆ y originating from covariate-specific variables ˆ z x in S .

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

It follows, by applying the image of the function ˆ g -1 1: n x , that:

Since g and ˆ g are smooth and injective, we may define the function ¯ h = ˆ g -1 ◦ g : Z → Z . We note that by definition ¯ h = h -1 where h is introduced in the proof of Theorem 4.4. We now remind the reader that y = g ( z ) . Therefore, using the notation ¯ h x := ¯ h 1: n x : Z → Z x , we have the equivalent condition

Now, using the pre-image formulation again, we may write it as

and, using an integral notation:

where ¯ h -1 x ( S ) = { z ∈ Z : ¯ h x ( z ) ∈ S } 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 ⊆ Z x and x ∈ X :

where the condition must hold for all ( t 1 , t 2 ) ∈ T 2 .

Step 2 (Topological characterization of invariance). To demonstrate the block-identifiability of z x , our objective is to substantiate that ¯ h x ([ z x , z tx , z t ]) = ¯ h x ([ z x , z + t ]) is functionally independent of z + t . To achieve this, we initially formulate a set of equivalent statement:

Statement 1 . ¯ h x ( [ z ⊤ x , z + ⊤ t ] ⊤ ) does not depend on z + t .

Statement 2 . ∀ z x ∈ Z x , ∃ B z x ⊆ Z x \ ∅ : ¯ h -1 x ( z x ) = B z x ×Z tx ×Z t .

Statement 3 . ∀ z x ∈ Z x , ∀ r ∈ R + , ∃ B + z x ⊆ Z x \ ∅ : ¯ h -1 x ( B r ( z x )) = B + z x ×Z tx ×Z t ,

where B r ( z x ) is defined as the ball centered around z x with radius r : B r ( z x ) = { z ′ x ∈ Z x : ∥ z ′ x -z x ∥ 2 < r } .

We note that Statement 2 is a mathematical formulation of Statement 1, and that Statement 3 is a generalization of Statement 2 from singletons { z x } in Statement 2 to open, non-empty balls B r ( z x ) . We proceed to demonstrating equivalence between those statements.

Statement 2 ⇒ Statement 3. Let z x ∈ Z x and r ∈ R + . By definition of the pre-image of a set, we have that:

Because we assume Statement 2, we have that for all z ′ x ∈ B r ( z x ) , there exists a set B z ′ x such that h -1 x ( z ′ x ) = B z ′ x ×Z tx ×Z t . Therefore, Statement 3 stands for B + z x = ∪ z ′ x B z ′ x .

̸

Statement 2 ⇐ Statement 3. We proceed by contradiction. Suppose that Statement 2 is false, then for a certain ¯ z ∗ x ∈ Z x , it is possible to construct a point ¯ z B = [¯ z B x , ¯ z B tx , ¯ z B t ] ∈ Z such that ¯ z B x is in the pre-image of { ¯ z ∗ x } by ¯ h -1 x but ¯ h x ( ¯ z B ) = ¯ z x . Indeed, this directly means that changing the other components of Z at the input of ¯ h can alter the component of Z x at its output. By continuity of ¯ h x , there exists ˆ r > 0 such that ¯ h x ( ¯ z B ) / ∈ B ˆ r (¯ z x ) . For such ˆ r , we have that ¯ z B / ∈ h -1 x ( B ˆ r (¯ z x )) . Additionally, the application of Statement 3 suggests that there exists a non-trivial subset B + ¯ z x such that h -1 x ( B ˆ r (¯ z x )) = B + ¯ z x ×Z tx ×Z t . By definition of ¯ z B , it is clear that ¯ z B 1: n x ∈ B + ¯ z x . The fact that ¯ z B / ∈ h -1 x ( B ˆ r (¯ z x )) contradicts Statement 3. Therefore, Statement 2 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 Z x are non-empty and open sets. For z x ∈ Z x and r ∈ R + , we note that because B r ( z x ) is open and ¯ h x is continuous, the pre-image ¯ h -1 x ( B r ( z x )) is open. In addition, because h is continuous and we have equality of the generated data distributions:

we have that h is a bijection (Klindt et al., 2021), which ensures that ¯ h -1 x ( B r ( z x )) is non-empty. Hence, ¯ h -1 x ( B r ( z x )) is both non-empty and open.

We now assume, by contradiction, that z x is not block-identifiable. Therefore, ¯ h x is not invariant with respect to z + t . Additionally, because of Step 2, we have the existence of a ball S ∗ := B r ∗ ( z ∗ x ) centered on the point z ∗ x ∈ Z x and of radius r ∗ ∈ R + such that ¯ h -1 x ( S ∗ ) cannot be written of the form A ×Z tx ×Z t for any non-trivial A ⊂ Z x .

We may therefore define the set B ∗ z := { z ∈ ¯ h -1 x ( S ∗ ) | { z 1: n x } × Z t ×Z tx ⊈ ¯ h -1 x ( S ∗ ) } . Intuitively, B ∗ z contains the partition of the pre-image ¯ h -1 x ( S ∗ ) that the t part z + t cannot take on any value in Z tx × Z t . It is therefore non-empty by hypothesis. To show contradiction with Equation 53, we evaluate it on the set S ∗ and split the integral on two domains of the partition ¯ h -1 x ( S ∗ ) = ( ¯ h -1 x ( S ∗ ) \ B ∗ z ) ∪ B ∗ z .

We define the following integrals:

where we have the expected relation T = T 1 + T 2 .

We first look at the value of T 1 . In the case where the set ¯ h -1 x ( S ∗ ) \ B ∗ z is empty, then T 1 trivially evaluates to 0. Otherwise, there exists a non-empty subset C ∗ z x of Z x such that ¯ h -1 x ( S ∗ ) \ B ∗ z = C ∗ z x ×Z tx ×Z t . With this expression, it follows that

Because of the separability of the domains, we may apply Fubini's theorem:

Therefore, in both cases T 1 evaluates to 0 for S ∗ .

Now, we address T 2 . Towards this goal, we prove that B ∗ z satisfies the condition for application of the assumption of the theorem. First, we must show that B ∗ z has non-zero probability measure for all values of t and x . For this, it is enough to show that B ∗ z contains an open set, given that we assume that p z | t , x ( z | t , x ) > 0 over ( z , t , x ) ∈ Z × T × X . Let us take one element z B ∈ B ∗ z , which is possible because we proved B ∗ z is non-empty. As discussed above, ¯ h -1 x ( S ∗ ) is open and non-empty, and by continuity of ¯ h x , there exists r 0 ∈ R + such that B r 0 ( z B ) ⊆ B ∗ z . Therefore, B ∗ z contains an open set and has non-zero probability. Second, it is by definition that B ∗ z cannot be expressed as A z x ×Z tx ×Z t for any A z x ⊂ Z x .

Therefore, condition (ii) from the theorem indicates that there exists t ∗ 1 , t ∗ 2 , x ∗ , such that

̸

̸

Therefore, for such S ∗ , we would have T 1 + T 2 = 0 which leads to contradiction with Equation 53. We have proved by contradiction that Statement 1 from Step 2 holds, that is, ¯ h x does not depend on the treatment variable and interaction variable z t , z tx .

Step 4 (Block-wise identifiability of z t and z x ). With the knowledge that h x does not depend on z + t , we now show that there exists an invertible mapping between the true content variable z x and the estimated version ˆ z x .

As ¯ h is smooth over Z , its Jacobian can written as:

where we use notation ˆ z x = ¯ h ( z ) 1: n x and ˆ z tx = ¯ h ( z ) n x +1: n x + n tx , ˆ z t = ¯ h ( z ) n x + n tx +1: n .

First, we notice that under the condition of Theorem 4.4, there is an invertible mapping between z tx and ˆ z tx . Therefore, it must be that D = F = 0 , and that E is non-singular. Additionally, we have just shown that ˆ z x does not depend on the treatment-related variables z + t . Therefore, it follows B = C = 0 . On the other hand, as h is invertible over Z , J ¯ h is non-singular. Therefore, A must be non-singular due to B = C = 0 . Relying on analogous assumptions to prove the invariance of ˆ z t with respect to z + x , it follows that G = H = 0 , and that I must be non-singular.

We note that A is the Jacobian of the function ¯ h ′ x ( z x ) := ¯ h x ( z ) : Z x →Z x , which takes only the covariates part z x of the input z into ¯ h x . Together with the invertibility of ¯ h , we can conclude that ¯ h ′ x is invertible. Therefore, there exists an invertible function ¯ h ′ x between the estimated and the true ˆ z x = ¯ h ′ x ( z x ) , which concludes the proof that z x is block-identifiable. Similarly, we are able to conclude z t is block-identifiable.

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 is the logarithm of the marginal data probability, and we calculate it by weighting it against the variational distribution:

We apply Jensen's inequality using the fact that the logarithm is a concave function:

Algorithm 1 Training of FCR

0 0

Input: y , y , shared x , and t , t = 0 Output: p, q, z x , z tx , z t while Not converged do Minimization Stage 1. Estimate p ( z t | t ) , p ( z x | x ) , p ( z tx | t , x ) , and p ( z 0 t | t 0 ) , p ( z 0 x | x ) , p ( z 0 tx | t 0 , x ) 2. Estimate q ( z t , z x , z tx | t , x , y ) , q ( z 0 t , z 0 x , z 0 tx | t 0 , x , y ) 3. Calculate Kullback-Leibler divergence terms and predict ˆ t 4. Calculate similarities { z 0 x , z x } , { z 0 t , z t } 5. Permute z tx to get ˆ z tx and predict permutation labels ˆ l 6. Minimize L ELBO, L sim , L ct , L dist x , L dist t Maximization Stage 1. Estimate q ( z t , z x , z tx | t , x , y ) , 2. Permute z tx to get ˆ z tx and predict permutation labels ˆ l 3. Maximize L dist x , L dist t

end while

Then, we use the factorization of our generative model:

Naming the right hand size of Equation 66 as L ELBO, we notice that L ELBO can be written as the following difference:

where the first term is the reconstruction loss, and the second term is the Kullback-Leibler divergence between the approximate posterior ¯ q ϕ = q ϕ ( z t , z tx , z x | y , t , x ) and the prior p θ ( z t , z tx , z x | t , x ) . Further decomposing this term using the factorization of the generative model and the inference model, we obtain:

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:

C Training Details

We provide additional training information in this section, including a detailed algorithm for the training process of FCR (Algorithm 1).

D Hyperparameter selection

We split the data into four datasets: train/validation/test/prediction, following the setup from previous works (Lotfollahi 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-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 shown in Table 2.

E Datasets and Preprocessing

In this section, we provide detailed descriptions of the four datasets and their corresponding preprocessing procedures. These datasets contain cells from many cell lines that are described in Table 3.

E.1 SciPlex dataset

This dataset includes three cancer cell lines exposed to 188 different compounds. In total, this experiment profiled approximately 650,000 single-cell transcriptomes across roughly 5,000 independent samples (Srivatsan et al., 2020). For our experiments, we 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.

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, we analyzed 90,462 cells, including both treated and control groups. Cell lines and repetition numbers were used as covariates, with treatment dosage designated as the treatment variable.

E.2 MultiPlex-Tram dataset

This dataset is referred to as 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 et 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 in McFarland et al. (2020), includes 72,326 cells treated with seven different compounds across 24 cell lines. Details of the seven treatments are provided in Table 5. Following the removal of low-quality cells and normalization of raw counts, we retained the top 5,000 differentially expressed genes, yielding 61,552 cells for the experiments and downstream analysis.

E.4 Multiplex-9 dataset

referred to as experiment-10 in McFarland et al. (2020), consists of 37,856 cells across 24 cell lines, treated with 9 drugs (including a control). The list of drugs can be found in Table 6. After filtering out low-quality cells and normalizing the raw counts, we retained the top 5,000 differentially expressed genes. This resulted in a total of 19,524 cells for the experiments and downstream evaluation.

F Experimental Setups and Additional Results

F.1 Training Details

In this subsection, we layout the training parameters for each datasets as follows.

sciPlex For the sciPlex datasets, the dimensions are as follows: z x is 32, z tx is 64, and z t is 32. Additionally, we set the hyperparameters to ω 1 = 3 . 0 , ω 2 = 3 . 0 , and ω 3 = 5 . 0 , with a batch size of 2046. Note that in sciPlex, we treat the dosages of the HDAC inhibitors as the treatment variable.

multiPlex-Tram or the multiPlex-Tram dataset, the dimensions are set as follows: z x is 32, z tx is 32, and z t is 32. The hyperparameters are ω 1 = 5 . 0 , ω 2 = 5 . 0 , ω 3 = 1 . 0 , with a batch size of 2046. In this dataset, Trametinib treatment time is considered as the treatment variable.

multiPlex-7 For the multiPlex-7 dataset, the dimensions are z x = 32 , z tx = 64 , and z t = 32 . The hyperparameters are ω 1 = 1 . 0 , ω 2 = 0 . 5 , and ω 3 = 0 . 1 , with a batch size of 2046.

multiPlex-9 For the multiPlex-9 dataset, the dimensions are z x = 32 , z tx = 64 , and z t = 32 . The hyperparameters are ω 1 = 1 . 0 , ω 2 = 0 . 5 , and ω 3 = 0 . 1 , with a batch size of 2046.

Additionally, the autoencoder learning rate is set to 3 × 10 -4 , the discriminator learning rates are also 3 × 10 -4 , and the number of discriminator training steps is 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 et al. (2022), Khemakhem et al. (2020) and Lopez et al. (2023), we simplify the simulation setup by setting the dimensions of z x and z t to 1, while z xt has a dimension of 4. Specifically, we use a sample size of 5,000 and define our variables as follows:

Here, t represents the treatment variable, uniformly distributed over three discrete values.

x denotes the covariate, also uniformly distributed but over a wider range of values.

z x is the latent variable associated with x , following a normal distribution with mean x / 2 and unit variance.

z t is the latent variable associated with t , also normally distributed with mean t / 2 and unit variance.

z tx represents the interaction between x and t , following a multivariate normal distribution with mean x · t and covariance matrix I 4 (the 4-dimensional identity matrix).

Finally, we define our output y as a function of these latent variables:

Here, g is implemented as a 2-layer MLP with Leaky-ReLU activation, following the approach of Kong et al. (2022) and Khemakhem et al. (2020). The output y is a real-valued vector with a dimension of 96.

Evaluation To assess the component-wise identifiability of the interaction components, we compute the Mean Correlation Coefficient (MCC) between z tx and ˆ z tx . MCC is a standard metric in Independent Component Analysis (ICA) literature, where a higher MCC indicates better identifiability. MCC reaches 1 when latent variables are perfectly identifiable (up to a component-wise 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, we 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 of the true parameters and recovery of the true sources, up to point-wise transformations. This is a standard performance metric in ICA (Khemakhem et al., 2020).

Results Our results suggest that our method, FCR largely outperforms existing variational autoencoder-based approaches in identifying the latent interactive components z tx . As shown in Table 7, 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-wise identifiability for complex, interacting latent variables in causal representation learning.

F.3 Additional Clustering Details and Results

Clustering Approach Our clustering analysis utilizes the learned representations z x , z tx , and z t 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 .

Clustering on t : Similarly, we used the Leiden algorithm on the representations and assessed the outcomes with the labels of t .

Clustering on x combined with t : We ran Leiden clustering on the combined representations and evaluated the results using both the labels of x and t .

This approach allows us to assess the efficacy of our learned representations in capturing the underlying structure of both individual and combined variables.

Evaluation Metric The evaluation metric, Normalized Mutual Information (NMI), is defined as follows:

where:

Y = { y 1 , . . . , y K } denotes the set of class labels,

C = { c 1 , . . . , c L } denotes the set of cluster labels,

p kl is the joint probability of a data sample belonging to class k and cluster l ,

p Y k is the marginal probability of a sample belonging to class k ,

p C l is the marginal probability of a sample belonging to cluster l .

Note that these probabilities are computed from the same dataset, but with respect to the true class labels and the assigned cluster labels, respectively.

Additional Results Additional clustering results for the multiPlex-tram, multiPlex-7, and multiPlex9 datasets are presented in Figure 5.

F.4 Statistical Tests and More results

Due to the computational complexity of kernel calculations, we adopted a sampling approach with 2,000 samples, repeating the process 100 times to report the results. For the baseline methods, we sampled their latent spaces to match the dimensions of z x , z t , and z tx . This sampling was repeated 20 times, and we reported the best results for comparison with FCR. The test results for the multiPlex-tram, multiPlex-7, and multiPlex-9 datasets are shown in the Figure 6

F.5 Conditional Cellular Response Prediction

For this task, we use FCR's latent representations to predict gene expression levels and report the corresponding R 2 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 to drug treatments. Our comparative analysis includes CPA, VCI, sVAE, scGEN, and CINEMA-OT, as these methods are specifically designed for cellular prediction tasks.

We utilize FCR to extract the control's latent representations [ z 0 x , z 0 tx , z 0 t ] and the corresponding experimental representations [ z x , z tx , z t ] . The decoder g is then used to predict the gene expression levels as ˆ y = g ( z 0 x , z 0 tx , z 0 t ) . The R 2 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 scVI, β VAE, and factorVAE, we only used gene expression data to learn latent variables and make predictions.

The R 2 score is a key metric for assessing the accuracy of predictive models. It measures the proportion of variance in the dependent variable that is explained by the independent variables. An R 2 score of 1 indicates perfect prediction accuracy, meaning all variations in the target variable are

fully explained by the model's inputs, while a score of 0 suggests that the model fails to capture any variance in the target variable.

where y i is the actual values, ˆ y i is the predicted values, ¯ y is the mean of average values, n is the number of observations.

Additional Metric We further utilize the Mean Squared Error (MSE) of the top 20 differentially expressed genes (DEGs) for post-treatment. 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 Roohani et al. (2024). Note here, we did not compare with CINEMA-OT and scGEN because they are only for binary treatments.

F.6 Ablation Study

We present the ablation study results for the hyperparameters ω 1 , ω 2 , and ω 3 . Initially, we set ω 1 , ω 2 , and ω 3 to [1,1,1]. Then, we varied each parameter independently to [1 , 3 , 5 , 10 , 20] , keeping the other parameters fixed at 1. Figure 7 illustrates the NMI scores for clustering on x , t , or both x & t . Figure 8 shows the R 2 scores for each parameter.

F.7 Visualization

Visualizing latent representations provides intuitive insights into the distinct characteristics and attributes captured by each representation. Using Uniform Manifold Approximation and Projection (UMAP) (McInnes et al., 2018), we visualized the latent representations z x , z tx , and z t derived from the sciPlex dataset (Figure 9).

From these visualizations, we observe that z x effectively captures covariate-specific information, showing clear separation in Figure 9, but it does not reflect treatment information. In contrast, the UMAP visualization of z t reveals a clear pattern, where cells treated with higher dosages are positioned at the bottom, and control or lower-dosage treated cells are found at the top. Additionally, various cell lines intermingle within the plot, indicating that z t successfully captures drug responses across different cell lines.

Most importantly, Figure 9 demonstrates that z tx 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 in capturing nuanced biological behaviors.

F.8 Pilot Study On The Unseen Drug Responses

Predicting 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 (McFarland et al., 2020). This complexity poses significant challenges for AI-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

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.

In 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 multiPlex-Tram 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:

Drug Hold-out Setup : We held out two cell lines, ILAM and SKMEL2, from the MultiplexTram 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

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.

Transfer MLP : We extract both z p tx from M p and z h tx from M h for D p . Then we trained a 1-layer MLP (the same dimension as z p tx ) to transfer z p tx to z h tx , by minimize the MSE between them.

Contextual Prior Representation : We extracted the z p tx representations from model M p for ILAM and SKMEL2 in holdout set. Then transfer z p tx by the previous MLP to ˆ z h tx as a prior contextual embedding. For the hold-out cell lines ILAM and SKMEL2 in the Multiplex-Tram dataset, we extracted z h x from model M h , and Trametinib-24h z h t from other treated cell lines in D h

Representation Matching : Then we match the ˆ z h tx by z p x similarity on ILAM and SKMEL2 across Multiplex-tram and Multiplex-7 in the prior knowledge model space.

Prediction : We predicted the unseen 24 hours Trametinib responses for the ILAM and SKMEL2 cell lines in holdout dataset using the formula: ˆ y = g ( z h x , ˆ z h tx , z h t ) ,where ˆ z h tx is the corresponding matched prior contextual representation, z h t is the tramnib-24 average value in other cell lines.

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.

(a)

0.8

0.6

0.5

0.7

0.5

0.4

0.6

0.5

E0

0.2

宁

0.2

0.4

百

0.1

0.1 -

NAEVAEAE

SVAE

TVAE

factorl

betav

betav

methods

methods

methods

(b)

0.8

白

1.0

0.7

百

口

0.9

百

0.5

0.6

0.8

0.4

0.7

0.4

0.6

0.3

0.5

0.2

0.2

0.4

0.1

0.1

0.3

SVIAE COP NC

VAEVIAEAE

SVAE CPP NC

DNVAEVAEVAE

VAE

VAEVAE

factorl

methods

methods

methods

(c)

0.8

0.7

宁

+

百

0.9

0.5

0.8

0.6

0.7

0.4

M O.5-

N0.6-

0.4

0.5

FE'0

0.4

0.2

宁

一

0.2

0.3

0.1

SVAECPAC

factorl

methods

methods

methods

(a)

0.35

0.5

0.08

0.30

0.30

0.4

0.25

0.25

0.06

0.3

0.2-

0.15

0.15

0.10

0.10

0.02

0.05

0.05

0.0

0.00

0.00

0.00

h.?.." b.d

S

iVAE

S

iVAE

betaVAE

betaVAE

(b)

0.35

0.04

0.6

0.6

0.30

0.25

0.03

Pvalue

0.4

Pvalue

0.2

0.2

0.10

0.01

0.05

0.0

0.0

0.00-

0.00

ScVI

iVAE

betaVAE

betaVAE

(c)

(0)

(1)

(2)

Swarm Plot for ZXTR

0.35

0.30

0.04

0.20

0.30

0.25

0.25 -

0.03

Pvalue

0.20-

Pvalue

0.15

0.10

0.15-

0.10

0.10

0.05

0.01

0.05

0.05

0.00

0.00

0.00-

0.00

....-....

iVAE

iVAE

betaVAE

betaVAE

0.9

0.175

0.7

0.8

0.150

0.7

0.6

0.125

NMI

0.6

Zx

Zx

0.100

Zt

0.5

Zt

0.5

0.075

0.4

0.4

0.050

Zx

0.3

Zt

0.025

0.2

0.3

5

10

15

20

5

10

15

20

5

10

15

20

omega1

omega1

omega1

0.9

0.7

Zx

0.8

Zt

0.20

0.7

0.6

NMI

0.6

Zx

Zt

0.5

Zx

Zt

0.15

0.5

0.4

0.10

0.4

0.3

0.2

0.3

0.05

5

omega2

10

15

20

5

omega2

10

15

20

5

10

15

20

omega2

0.9

0.70

Zx

0.175

Zt

0.8

0.65

0.150

0.60

0.7

0.125

MO.6

Zx

0.55

Zx

Zt

0.100

Zt

0.50

0.075

0.5

0.45

0.050

0.4

0.40

0.025

0.3

0.35

5

10

15

20

5

10

15

20

5

10

15

20

omega3

omega3

omega3

0.9

0.9

0.80

0.75

0.8

0.8

2

0.70

0.7

0.7

0.65

0.60

0.6

0.6

0.55

5

10

15

20

5

10

15

20

5

10

15

20

omega_1

omega2

omega3

(a)

UMAP

UMAP

Z, UMAP

Z

Z_UMAP

Z, UMAP

Z, UMAP

UMAP

UMAP

Z, UMAP

Z

Z,UMAP

Z UMAP

Z, UMAP

MCF7_Rep1

MCF7_Rep2

K562_Rep1

A549_Rep1

K562_Rep2

Dosage

(b)

UMAP

Z, UMAP

Z, UMAP

ZUMAP

Z, UMAP

Z, UMAP

Z, UMAP

UMAP

Z, UMAP

Z, UMAP

Z UMAP

Z, UMAP

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 in 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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer:[Yes]

Justification: The hyperparameters are 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 a 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 it 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 to 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 a 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 submissions 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 it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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. In the case of closed-source models, it 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 to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide this detail in the appendices.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: It's not applicable in our paper.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses 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 our paper.

Guidelines:

The answer NA means that the paper does not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

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 of the dataset/code/model as part of 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 is 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 as 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 not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 an equivalent approval/review based on 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 we expect authors to adhere to 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 as the institution conducting the review.

Peer reviews:
comments: 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.

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)$.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors could discuss the limitations (based on assumptions) more directly.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

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

STRENGTHS:
- novel causal representation learning method explicitly learning interactions between covariates and treatments
- novel regularizers for enforcing conditional independence between sets of variables
- 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
- 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

QUESTIONS:
- why not evaluate on unseen treatments or cell contexts? Did you try and find the results to be comparatively poor?

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nChoice 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\n\nQUESTIONS:\nCan you elaborate more about the design choice of the evaluating metric?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- the proposed method is justified theoretically\n- the paper is written in a clear and concise manner\n- extensive selection of baselines\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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)", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}