Learning Identifiable Factorized Causal Representations of Cellular Responses

Authors: haiyi  mao
romain  lopez
kai  liu
jan-christian  huetter
david  richmond
panayiotis v benos
lin  qiu
  genentech
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 22:20:52

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.

Main paper: 
 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) . 
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 A field 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) .We isolate the terms relevant to z tx by strategically comparing equations for different pairs of (x, t) values and subtracting them. 
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).Figure 1 : Data generating process: shaded nodes denote observed variables. Empty nodes denote latent variables.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 zx|x (z x | x), z t is generated from t following distribution z t ∼ p zt|t (z t | t), and z tx from both x and t following distribution z tx ∼ p ztx|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 0x , 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) .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 ẑ such that g(ẑ) and g(z) are equal in distribution, z and ẑ 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}:EQUATIONwhere z j and ẑπ(j) are the π(j)-th components of z and ẑ respectively.Definition 3.2 (Block-wise Identifiability). For 1 ≤ n 1 < n 2 ≤ n, we denote as z [n1:n2] ∈ R n2-n1+1the block of z from index n 1 to n 2 . We say that latent vector z [n1:n2] is block-identifiable from data y if for any other latent vector ẑ such that g(ẑ) and g(z) are equal in distribution, z [n1:n2] and ẑ[n1:n2] are equal up to an invertible function h:EQUATIONwhere ẑ[n1:n2] is the corresponding block in the estimated vector ẑ.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 ẑ that yields the same data distribution than the ground-truth variables z, for all t ∈ T and x ∈ X . Specifically, we have:EQUATIONGiven our assumption of noiseless observations, it is equivalent to equality in distribution of the mixed variables:EQUATIONwhich, after a change of variable, is equivalent to:EQUATIONwhere g -1 : Y → Z denotes the invertible generating function, and h := g -1 • ĝ 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:EQUATIONAccording to the independence relations in the data generating process, we havep z|t,x (z|t, x) = n i=1 p zi|t,x (z i | t, x); p ẑ|t,x (ẑ|t, x) = n i=1 p ẑi|t,x (ẑ i | t, x).Rewriting the notation q i := log p zi|t,x and qi := log p ẑi|t,x yields:log p z|t,x (z|t, x) = n i=1 q i (z i , t, x); log p ẑ|t,x (ẑ|t, x) = n i=1 qi (ẑ i , t, x).Applying the change of variables formula to Equation 24 yieldsEQUATIONwhere 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,EQUATION)Then, we may differentiate Equation 28 with respect to ẑk and ẑ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:EQUATION)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:EQUATIONThen, 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:nx i=1 (b ′′ i (z i , x) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , x)a ′′ i,(k,q) ) + nx+ntx i=nx+1 (b ′′ i (z i , t 0 , x) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t 0 , x)a ′′ i,(k,q) ) + n i=nx+ntx+1 (b ′′ i (z i , t 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t 0 )a ′′ i,(k,q) ) + ∂ 2 log |J h | ∂ ẑk ẑq = 0.(32) Proceeding similarly for (x 0 , t 0 ), we obtain:EQUATIONThen taking the difference between Equation 32 and Equation 33 yields,EQUATIONWe apply the same principle to the pair {(x 0 , t), (x 0 , t 0 )}. We therefore get:nx i=1 (b ′′ i (z i , x 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , x 0 )a ′′ i,(k,q) ) + nx+ntx i=nx+1 (b ′′ i (z i , t, x 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t, x 0 )a ′′ i,(k,q) ) + n i=nx+ntx+1 (b ′′ i (z i , t) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t)a ′′ i,(k,q) ) + ∂ 2 log |J h | ∂ ẑk ẑq = 0, (35) and, nx i=1 (b ′′ i (z i , x 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , x 0 )a ′′ i,(k,q) ) + nx+ntx i=nx+1 (b ′′ i (z i , t 0 , x 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t 0 , x 0 )a ′′ i,(k,q) ) + n i=nx+ntx+1 (b ′′ i (z i , t 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t 0 )a ′′ i,(k,q) ) + ∂ 2 log |J h | ∂ ẑk ẑq = 0.(36) Taking similarly the difference between Equation 35 and Equation 36 yields:nx+ntx i=nx+1 (b ′′ i (z i , t, x 0 ) -b ′′ i (z i , t 0 , x 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t, x 0 ) -b ′ i (z i , t 0 , x 0 ) a ′′ i,(k,q) ) + n i=nx+ntx+1 b ′′ i (z i , t) -b ′′ i (z i , t 0 ) • a ′ i,(k) a ′ i,(q) + b ′ i (z i , t) -b ′ i (z i , t 0 ) a ′′ i,(k,q) = 0.(37) Finally, we consider the pairs {(x, t), (x 0 , t 0 )}, for which we obtain the following:EQUATION)Once again taking the difference between Equation 38 and Equation 39 yields:EQUATION)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 40and Equation 34:EQUATIONFinally, we subtract Equation 41 and Equation 37:nx+n tx i=nx+1 (b ′′ i (zi, t, x) -b ′′ i (zi, t0, x0) -(b ′′ i (zi, t0, x) -b ′′ i (zi, t0, x0) -b ′′ i (zi, t, x0) -b ′′ i (zi, t0, x0) • a ′ i,(k) a ′ i,(q) + b ′ i (zi, t, x) -b ′ i (zi, t0, x0) -b ′ i (zi, t0, x) -b ′ i (zi, t0, x0) -b ′ i (zi, t, x0) -b ′ i (zi, t0, x0) a ′′ i,(k,q) = 0.This last equation may be rearranged as:EQUATIONStep 3 (Establishing Component-wise Identifiability) Given the assumption of linear independence in the Theorem, the linear system is a 2n tx × 2n tx full-rank system. Therefore, the only solution is:a ′ i,(k) a ′ i,(q) = 0 a ′′ i,(k,q) = 0 for i ∈ {n x + 1, . . . , n x + n tx } and k, q ∈ {1, . . . , n}, k ̸ = q.h(•) is a smooth function over Z and its Jacobian can written as:EQUATIONNote 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 ẑ. And because the other ẑ that are not in the block ẑ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 andx ′ ∈ X .(ii) S cannot be expressed asA zx × Z tx × Z t for any A zx ⊂ Z x or as Z x × Z tx × A zt for any A zt ⊂ 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 thatEQUATIONand there also existsEQUATIONthen 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 ] , ẑ+ t = [ẑ tx , ẑt ] for simplification. This proof is comprised of four major steps.1. 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.2. Topological characterization of invariance. We derive equivalence statements for domain invariance of functions.3. Proof of invariance by contradiction. We prove the invariance statement from Step 2 by contradiction. Specifically, we show that if ẑx depended on z + t , the invariance derived in Step 1 would break. 4. 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 ĝ : Z → Y denotes the learned mixing function.We assume that both g and ĝ are invertible, so that their reciprocal functions are well defined. In particular, we denote by ĝ-1 1:nx : 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 ĝ 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:EQUATION)This condition can be written equivalently using the pre-image set of S:∀(t 1 , t 2 ) ∈ T 2 , P ŷ ∈ ĝ-1 1:nx -1 (S) | {x, t = t 1 } = P ŷ ∈ ĝ-1 1:nx -1 (S) | {x, t = t 2 } ,(47) where ĝ-1 1:nx -1 (S) ⊆ Y is the set of estimated observations ŷ originating from covariate-specific variables ẑx in S.Because of the equality of the observed data y and the generated data distribution from the estimated model ŷ, the relation in Equation 47 also holds for the random variable y∀(t 1 , t 2 ) ∈ T 2 , P y ∈ ĝ-1 1:nx -1 (S) | {x, t = t 1 } = P y ∈ ĝ-1 1:nx -1 (S) | {x, t = t 2 } .(48) It follows, by applying the image of the function ĝ-1 1:nx , that:EQUATION) Since g and ĝ are smooth and injective, we may define the function h = ĝ-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 hx := h1:nx : Z → Z x , we have the equivalent condition∀(t 1 , t 2 ) ∈ T 2 , P hx (z) ∈ S | {x, t = t 1 } = P hx (z) ∈ S | {x, t = t 2 } .(50) Now, using the pre-image formulation again, we may write it asEQUATION) and, using an integral notation:EQUATIONwhere h-1 x (S) = z ∈ Z : hx (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 :EQUATIONwhere 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 hx ([z x , z tx , z t ]) = hx z x , z + t is functionally independent of z + t . To achieve this, we initially formulate a set of equivalent statement:1. Statement 1. hx z ⊤ x , z +⊤ t ⊤ does not depend on z + t . 2. Statement 2. ∀z x ∈ Z x , ∃B zx ⊆ Z x \ ∅ : h-1 x (z x ) = B zx × Z tx × Z t . 3. Statement 3. ∀z x ∈ Z x , ∀r ∈ R + , ∃B + zx ⊆ Z x \ ∅ : h-1 x (B r (z x )) = B + zx × 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:EQUATION)Because we assume Statement 2, we have that for all z ′ x ∈ B r (z x ), there exists a setB z ′ x such that h -1 x (z ′ x ) = B z ′ x × Z tx × Z t . Therefore, Statement 3 stands for B + zx = ∪ 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 zB = [z B x , zB tx , zB t ] ∈ Z such that zB x is in the pre-image of {z *x } by h-1x but hx zB ̸ = zx . 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 hx , there exists r > 0 such that hx zB / ∈ B r (z x ). For such r, we have that zB / ∈ h -1 x (B r (z x )). Additionally, the application of Statement 3 suggests that there exists a non-trivial subsetB + zx such that h -1 x (B r (z x )) = B + zx × Z tx × Z t . By definition of zB , it is clear that zB 1:nx ∈ B + zx . The fact that zB / ∈ h -1x (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 hx is continuous, the pre-image h-1x (B r (z x )) is open. In addition, because h is continuous and we have equality of the generated data distributions:EQUATIONwe have that h is a bijection (Klindt et al., 2021) , which ensures that h-1x (B r (z x )) is non-empty. Hence, h-1x (B r (z x )) is both non-empty and open. We now assume, by contradiction, that z x is not block-identifiable. Therefore, hx 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-1x (S * ) cannot be written of the form A × Z tx × Z t for any non-trivial A ⊂ Z x .We may therefore define the setB * z := z ∈ h-1 x (S * ) | {z 1:nx } × Z t × Z tx ⊈ h-1 x (S * ) . In- tuitively, B *z contains the partition of the pre-image h-1x (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 partitionh-1 x (S * ) = ( h-1 x (S * ) \ B * z ) ∪ B * z .We define the following integrals:EQUATIONwhere we have the expected relationT = T 1 + T 2 .We first look at the value of T 1 . In the case where the set h-1x (S * ) \B * z is empty, then T 1 trivially evaluates to 0. Otherwise, there exists a non-empty subsetC * zx of Z x such that h-1 x (S * ) \B * z = C * zx × Z tx × Z t .With this expression, it follows thatT 1 = [zx,z + t ]∈C * zx ×Ztx×Zt p zx|x (z x | x) p z + t |x,t z + t | x, t 1 -p z + t |x,t z + t | x, t 2 dz x dz + t .(59) Because of the separability of the domains, we may apply Fubini's theorem:T 1 = zx∈C * zx p zx|x (z x | x) dz x z + t ∈Ztx×Zt p z + t |x,t z + t | x, t 1 -p z + t |x,t z + t | x, t 2 dz x dz + t (60) T 1 = zx∈C * zx p zx|x z x|x (1 -1)dz x = 0. (61)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 hx , there exists r 0 ∈ R + such that B r0 (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 asA zx × Z tx × Z t for any A zx ⊂ Z x .Therefore, condition (ii) from the theorem indicates that there exists t * 1 , t * 2 , x * , such thatEQUATION)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, hx 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 ẑx .As h is smooth over Z, its Jacobian can written as:J h =      A := ∂zx ∂ẑx B := ∂zx ∂ ẑ tx C := ∂zx ∂ ẑt D := ∂ztx ∂ẑx E := ∂ztx ∂ ẑ tx F := ∂ztx ∂ ẑt G := ∂zt ∂ẑx H := ∂zt ∂ ẑ tx I := ∂zt ∂ ẑt      (63)where we use notation ẑx = h(z) 1:nx and ẑtx = h(z) nx+1:nx+ntx , ẑt = h(z) nx+ntx+1:n .First, we notice that under the condition of Theorem 4.4, there is an invertible mapping between z tx and ẑtx . Therefore, it must be that D = F = 0, and that E is non-singular. Additionally, we have just shown that ẑ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, Jh is non-singular. Therefore, A must be non-singular due to B = C = 0. Relying on analogous assumptions to prove the invariance of ẑ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 ) := hx (z) : Z x → Z x , which takes only the covariates part z x of the input z into hx . 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 ẑ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.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:EQUATION)We apply Jensen's inequality using the fact that the logarithm is a concave function:EQUATION)Algorithm 1 Training of FCR Input: y, y 0 , shared x, and t, t 0 = 0 Output: p, q, z x , z tx , z t while Not converged do Minimization Stage1. 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 0x , z 0 tx | t 0 , x, y) 3. Calculate Kullback-Leibler divergence terms and predict t 4. Calculate similarities {z 0x , z x }, {z 0 t , z t } 5. Permute z tx to get ẑtx and predict permutation labelsl 6. Minimize L ELBO , L sim , L ct , L distx , L distt Maximization Stage 1. Estimate q(z t , z x , z tx | t, x, y), 2.Permute z tx to get ẑtx and predict permutation labels l 3. Maximize L distx , L distt end while Then, we use the factorization of our generative model:EQUATIONNaming the right hand size of Equation 66 as L ELBO , we notice that L ELBO can be written as the following difference:EQUATIONwhere 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:EQUATIONFinally, 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:EQUATIONWe provide additional training information in this section, including a detailed algorithm for the training process of FCR (Algorithm 1).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 .Parameter Values ω 1 {0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0} ω 2 {0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0} ω 3 {0.1, 0.3, 0.5, 0.7, 0.9, 1.0, 3.0, 5.0, 7.0, 10.0} 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 .No 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.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.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. 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. 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.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.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:EQUATIONHere, t represents the treatment variable, uniformly distributed over three discrete values.x ∼ Unif({100, 1000, 5000})x denotes the covariate, also uniformly distributed but over a wider range of values.EQUATIONz x is the latent variable associated with x, following a normal distribution with mean x/2 and unit variance.EQUATIONz t is the latent variable associated with t, also normally distributed with mean t/2 and unit variance.EQUATIONz 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:EQUATIONHere, 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 ẑ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) .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 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:1. Clustering on x: We applied the Leiden clustering algorithm to the different representations and evaluated the results using the labels of x.2. Clustering on t: Similarly, we used the Leiden algorithm on the representations and assessed the outcomes with the labels of t.3. 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:EQUATIONwhere:• 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 multiPlex-9 datasets are presented in Figure 5 .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 PredictionFor 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 0x , 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 ŷ = g(z 0x , 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.R 2 = 1 - n i=1 (y i -ŷi ) 2 n i=1 (y i -ȳ) 2 , (77)where y i is the actual values, ŷi is the predicted values, ȳ 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. 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.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.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:1. Drug Hold-out Setup: We held out two cell lines, ILAM and SKMEL2, from the Multiplex-Tram dataset, which had been treated with Trametinib for 24 hours. We trained 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 2. 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.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. 4. 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 ẑ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 5. Representation Matching: Then we match the ẑh tx by z p x similarity on ILAM and SKMEL2 across Multiplex-tram and Multiplex-7 in the prior knowledge model space. 6. Prediction: We predicted the unseen 24 hours Trametinib responses for the ILAM and SKMEL2 cell lines in holdout dataset using the formula: ŷ = g(z h x , ẑh tx , z h t ),where ẑh tx is the corresponding matched prior contextual representation, z h t is the tramnib-24 average value in other cell lines. 7. 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. • 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). 
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 ẑtx must be reflective of the ground-truth interactions z tx only, and not of any of the other latent variables.Our proof relies on three technical assumptions. v(z tx , t, x) = ∂q nx+1 (z nx+1 , t, x) ∂z nx+1 , • • • , ∂q nx+ntx (z nx+ntx , t, x) ∂z nx+ntx , ∂ 2 q nx+1 (z nx+1 , t, x) ∂z 2 nx+1 , • • • , ∂ 2 q nx+ntx (z nx+ntx , t, x) ∂z 2 nx+ntx ,where q i denotes the logarithm of probability density p zi|t,x for component z i . If in addition to the assumptions 4.1, 4.2, 4.3, we assume that the 2n tx vectorsEQUATIONare 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.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 asA zx × Z tx × Z t for any A zx ⊂ Z x or as Z x × Z tx × A zt for any A zt ⊂ Z t .We have the following identifiability result. If for all S ∈ S(Z), there existsEQUATIONand there also existsEQUATIONthen 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. 
methodology: 
 We now propose a tangible implementation of our method, termed Factorized Causal Representation (FCR) learning. Our approach has three components:1. A variational inference approach to estimate the representations from our FCR model. Our model and inference architecture is specifically designed to learn disentangled representations z x , z tx , z t (Section 5.1). 2. A regularization method that enforces independence between z x and t, and encourages variability of z t with respect to x (Section 5.2). 3. 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 .Model Specification We parameterize the probability distributions as follows:EQUATIONwhere 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: is the component to estimate q(z x | x, y), q(z t | t, y) and q(z t,x | t, x, y). Note that 0 indicates t = 0 representing the control samples. (c) computational diagrams to estimate the Kullback-Leibler divergences, causal structure regularization and permutation discriminators.EQUATIONwhere 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:EQUATION)Following the graphical model from Figure 1 , the Evidence Lower Bound (ELBO) is derived as:EQUATIONwhere θ 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:EQUATIONq(z tx | t, x, y) := Normal(g t,x µ (t, x, y), g t,x σ (t, x, y)),where all the functions introduced above are parameterized by neural networks.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,EQUATIONwhere 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,EQUATIONThe predicted treatment probability vector t is then used for the computation of a cross-entropy lossEQUATIONWe 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 ) M i=1 and permuted samples (X π(i) , Y i , Z i ) Mi=1 with a permutation function π. The corresponding statisticsρ i of (X i , Y i ) M i=1 and ρ π(i) of (X π(i) , Y i ) M i=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 ) M j=1 to yield (z (j) x , z π(j) tx , x (j) = x i ) Mj=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 asEQUATIONwhere ẑ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 lossEQUATIONWe proceed similarly to make sure that z t and z tx are independent conditionally on t.Finally, we specify the overall loss for our model asEQUATIONwhere ω 1 , ω 2 , ω 3 > 0 are hyperparameters. To concurrently train both the representations and the discriminators, we employ an adversarial training approach as follows,EQUATION)The training procedure and the procedure used for hyperparameters selection are detailed in Appendix C and D, respectively. 
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). 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. Baselines We benchmarked our method against several established representation learning methods:(1) scVI (Lopez et al., 2018) , (2) iVAE (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.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, 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.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 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.z x ⊥ ⊥ z tx | x and z t ⊥ ⊥ z tx | t (FigureOur 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. 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 0x , 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 ŷ = g(z 0x , 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. 
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. A.1 Proof of Theorem 4.4 Theorem 4.4. Let us first define v(z tx , t, x) as the vector:v(z tx , t, x) = ∂q nx+1 (z nx+1 , t, x) ∂z nx+1 , • • • , ∂q nx+ntx (z nx+ntx , t, x) ∂z nx+ntx , ∂ 2 q nx+1 (z nx+1 , t, x) ∂z 2 nx+1 , • • • , ∂ 2 q nx+ntx (z nx+ntx , t, x) ∂z 2 nx+ntx ,where q i denotes the logarithm of probability density p zi|t,x of component z i . If in addition to the assumptions 4.1, 4.2, 4.3, we assume that for the 2n tx vectorsEQUATIONare linearly independent, then z tx is component-wise identifiable.Proof. This proof proceeds in three main steps:1. 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. 
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. 
 
 Appendix:  


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

All five reviewers vote towards acceptance.

Reviews: review: paper_summary: paper_summary: The authors develop a modification of the identifiable nonlinear ICA algorithm of Khemakhem et al. 2020, suited for scenarios with two disjoint groups of auxilliary variables x and t and their interactions, with an application of treatment effect on cell gene expression.
The 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.
strength_weakness: strength_weakness: - 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?

- 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?
questions: questions: - Equation (6) (the constraint for interaction between x and t) - how does that guarantee the interactions? The learned embeddings can have empty dimensions, e.g, k(x) = [k_1(x), …, k_n(x), 1, …, 1] and k(t) = [1, …, 1, k_1(t), …, k_m(t)] so that the Hadamard product becomes [k(x), k(t)]
- 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)?
- 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])
- How were the hyperparameters selected for the cellular response task?
limitations: limitations: - the algorithm has 3 hyperparameters, which might make training difficult. Details on hyperparameter selection were ommited
- 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)
paper_summary: paper_summary: The paper presents a novel method, the Factorized Causal Representation (FCR), which leverages identifiable deep generative models to understand cellular responses to genetic and chemical perturbations across multiple cellular contexts. This method improves upon prior models by learning disentangled representations that delineate treatment-specific, covariate-specific, and interaction-specific factors, which are shown to be theoretically identifiable. The effectiveness of FCR is demonstrated on four single-cell datasets. The performance, measured by the R^2 score of conditional cellular responses, often marginally outperforms state-of-the-art methods, showcasing FCR’s potential in accurately modeling cellular dynamics.
strength_weakness: strength_weakness: Choice of Evaluation Metric: A significant concern arises from the choice of the $R^2$ metric in section 6.3 for evaluating the performance of the Factorized Causal Representation (FCR) model. $R^2$ is suitable to assess whether a model adequately explains the variation in the response variable. However, it is not necessarily an indicator of the model's predictive accuracy.  Moreover, the near-identical $R^2$ reported for FCR, VCI, CPA, and sVAE (as shown in Table 1) suggest that $R^2$ gives limited insights into the comparative effectiveness of these methods.

There are some minor misformats, which I noticed while reading.
- Lines 35, 54: citation format issue
- Line 294: subscript format issue
questions: questions: Can you elaborate more about the design choice of the evaluating metric?
limitations: limitations: - Generalizability: The paper does not show the method’s applicability to apply on new dataset. This limitation raises concerns about how well the FCR can adapt to other biological datasets or experimental setups that deviate from the four testing datasets, which restricts the method’s utility.
- 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.

I 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: - why not evaluate on unseen treatments or cell contexts? Did you try and find the results to be comparatively poor?
limitations: limitations: NA
paper_summary: paper_summary: The authors propose a method for disentangling into factorized representations dependent on only covariates, only treatment and the interaction between treatment and covariates. Introducing a set of assumptions, the authors prove the identifiability of these disentangled variables based on previous proofs on non-linear ICA. The authors implement the model using variational inference and a set of regularizers and discriminators to ensure the needed independence and disentangling. The authors validate the clustering performance, conditional independencies and conditional cellular response prediction against commonly used methods for perturb-seq prediction and show generally very good performance.
strength_weakness: strength_weakness: - 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).
- 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
- 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.
- 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.
questions: questions: - How were the hyperparameters chosen for the different data sets?
- 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$?
- For prediction validation, what is train/val/test protocol? Is the correlation of y in the train set?
limitations: limitations: The authors could discuss the limitations (based on assumptions) more directly.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Is it possible to enhance other model’s interpretability using the factorized $z_x,z_t,z_{tx}$?
limitations: 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.

Comments: comment: Thank you for reading our response. Please feel free to let us know if there is anything needs to be clarified after the rebuttal. We value your feedback and want to address your remaining concerns.
Thank you for addressing my review. The unseen cell line experiment is nice, and probably as good as can be done given the datasets. I've raised my score accordingly.
Thank you for the authors' reply addressing much of my concern. I will keep my score unchanged.
Dear Reviewer **yZVk**:

We sincerely appreciate your constructive suggestions on the unseen prediction. During the rebuttal stage, we worked diligently to design and execute the experiment to address your feedback (Please refer to the response above, as well as Figures A and B in the rebuttal FDF file). We would greatly value your feedback on the experiment and would be eager to discuss it further with you before the discussion period ends. Your insights will be crucial in helping us refine our paper.

Thank you,

All the authors
Thank you very much for raising the score. For the MSE evaluation, we selected the top 20 differentially expressed genes, identified by conducting a comparison between treated and untreated cells, and ranked them based on t-test p-values, we considered the rank information to some extent in this evaluation.  Additionally, we plan to add the Spearman correlation comparison in the final version (The table below shows FCR and VCI results). Again, thank you for your valuable comments which helped us improve our work a lot!    
**Table: Spearman Correlation (top 500 highly variable genes)**:
| Datasets         | FCR         | VCI         |
|------------------|-------------|-------------|
| sciPlex          | **0.84(0.07)** | 0.82(0.08) |
| multiplex-Tram   | **0.87(0.05)**  | 0.85(0.06)  |
Thank you for your answers.

While most of my questions were clarified, I would suggest a different choice for the additional metric instead of MSE, e.g., the rank correlation mentioned by authors in the response. Unless I am missing something MSE does not bring too much new insight beyond R2, as they are both measuring the squared residuals, with the main difference being that R2 is assuming standardized targets.

Nevertheless in light of the other answers I raised my score accordingly
Thank you again for acknowledging our paper strengths and raising the score. We will further polish our manuscript and include the  hyperparameter tuning and ablation study in the final version.
Thank you for the extensive response. With additional information on hyperparameter tuning and ablation study, I have increased the score to weak accept.