SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiaqi Zhang
LIDS, MIT
Broad Institute of MIT and Harvard
Kristjan Greenewald
Chandler Squires
Karthikeyan Shannugam*
Caroline Uhler
Causal disentanglement aims to uncover a representation of data using latent variables that are interrelated through a causal model. Such a representation is identifiable if the latent model that explains the data is unique. In this paper, we focus on the scenario where unpaired observational and interventional data are available, with each intervention changing the mechanism of a latent variable. When the causal variables are fully observed, statistically consistent algorithms have been developed to identify the causal model under faithfulness assumptions. We here show that identifiability can still be achieved with unobserved causal variables, given a generalized notion of faithfulness. Our results guarantee that we can recover the latent causal model up to an equivalence class and predict the effect of unseen combinations of interventions, in the limit of infinite data. We implement our causal disentanglement framework by developing an autoencoding variational Bayes algorithm and apply it to the problem of predicting combinatorial perturbation effects in genomics.
The discovery of causal structure from observational and interventional data is important in many fields including statistics, biology, sociology, and economics [42, 17]. Directed acyclic graph (DAG) models enable scientists to reason about causal questions, e.g., predicting the effects of interventions or determining counterfactuals [46]. Traditional causal structure learning has considered the setting where the causal variables are observed [21]. While sufficient in many applications, this restriction is limiting in most regimes where the available datasets are either perceptual (e.g., images) or very-high dimensional (e.g., the expression of  $>20k$  human genes). In an imaging dataset, learning a causal graph on the pixels themselves would not only be difficult since there is no common coordinate system across images (pixel  $i$  in one image may have no relationship with pixel  $i$  in another image) but of questionable utility due to the relative meaninglessness of interventions on individual pixels. Similar problems are also present when working with very high-dimensional data. For example,
*Currently at Google Research. Contributions to this work were made when affiliated with IBM research.
in a gene-expression dataset, subsets of genes (e.g. belonging to the same pathway) may function together to induce other variables and should therefore be aggregated into one causal variable.
These issues mean that the causal variables need to be learned, instead of taken for granted. The recent emerging field of causal disentanglement [9, 65, 30] seeks to remedy these issues by recovering a causal representation in latent space, i.e., a small number of variables  $U$  that are mapped to the observed samples in the ambient space via some mixing map  $f$ . This framework holds the potential to learn more semantically meaningful latent factors than current approaches, in particular factors that correspond to interventions of interest to modelers. Returning to the image and the genomic examples, latent factors could, for example, be abstract functions of pixels (corresponding to objects) or groups of genes (corresponding to pathways).
Despite a recent flurry of interest, causal disentanglement remains challenging. First, it inherits the difficulties of causal structure learning where the number of causal DAGs grows super-exponentially in dimension. Moreover, since we only observe the variables after the unknown mixing function but never the latent variables, it is generally impossible to recover the latent causal representations with only observational data. Under the strong assumption that the causal DAG is the empty graph, such unidentifiability from observational data has been discussed in previous disentanglement works [28].
However, recent advances in many applications enable access to interventional data. For example, in genomics, researchers can perturb single or multiple genes through CRISPR experiments [13]. Such interventional data can be used to identify the causal variables and learn their causal relationships. When dealing with such data, it is important to note that single-cell RNA sequencing and other biological assays often destroy cells in the measurement process. Thus, the available interventional data is unpaired: for each cell, one only obtains a measurement under a single intervention.
In this work, we establish identifiability for soft interventions on general structural causal models (SCMs), when the latent causal variables are observed through a class of (potentially non-linear) polynomial mixing functions proposed by [3]. Prior works [57, 20, 67] show that the causal model can be identified under faithfulness assumptions, when all the causal variables are observed. We here demonstrate that identifiability can still be achieved when the causal variables are unobserved under a generalized notion of faithfulness. The identifiability is up to an equivalence class and guarantees that we can predict the effect of unseen combinations of interventions, in the limit of infinite data. It then remains to design an algorithmic approach to estimate the latent causal representation from data. We propose an approach based on autoencoding variational Bayes [29], where the decoder is composed of a deep SCM (DSCM) [45] followed by a deep mixing function. Finally, we apply our approach to a real-world genomics dataset to find genetic programs and predict the effect of unseen combinations of genetic perturbations.
Identifiable Representation Learning. The identifiability of latent representations from observed data has been a subject of ongoing study. Common assumptions are that the latent variables are independent [11], are conditionally independent given some observed variable [23, 28], or follow a known distribution [71]. In contrast, we do not make any independence assumptions on the latent variables or assume we know their distribution. Instead, we assume that the variables are related via a causal DAG model, and we use data from interventions in this model to identify the representation.
Causal Structure Learning. The recovery of a causal DAG from data is well-studied for the setting where the causal representation is directly observed [21]. Methods for this task take a variety of approaches, including exact search [12] and greedy search [10] to maximize a score such as the posterior likelihood of the DAG, or an approximation thereof. These scores can be generalized to incorporate interventional data [62, 67, 33], and methods can often be naturally extended by considering an augmented search space [43]. Indeed, interventional data is generally necessary for identifiability without further assumptions on the functions relating variables [54].
Causal Disentanglement. The task of identifying a causal DAG over latent causal variables is less well-studied, but has been the focus of much recent work [9, 65, 30]. These works largely do not consider interventions, and thus require restrictions on functional forms as well as structural assumptions on the map from latent to observed variables. Among works that do not restrict the map, [2] and [6] assume access to paired counterfactual data. In contrast, we consider only unpaired data, which is more common in applications such as biology [55]. Unpaired interventional data
is considered by [3], [53], and as a special case of [37]. These works do not impose structural restrictions on the map from latent to observed variables but assume functional forms of the map, such as linear or polynomial. Our work builds on and complements these results by providing identifiability for soft interventions and by offering a learning algorithm based on variational Bayes. We remark here that the task of causal disentanglement is sometimes called causal representation learning in literature. We adopted the term causal disentanglement mainly following [26], as causal representation learning also includes methods such as Invariant Risk Minimization (IRM) [4] which do not completely learn latent variables. We discuss contemporaneous related work in Appendix I.
We now formally introduce the causal disentanglement problem of identifying latent causal variables and causal structure between these variables. We consider the observed variables  $X = (X_{1},\dots,X_{n})$  as being generated from latent variables  $U = (U_{1},\dots,U_{p})$  through an unknown deterministic (potentially non-linear) mixing function  $f$ . In the observational setting, the latent variables  $U$  follow a joint distribution  $\mathbb{P}_U$  that factorizes according to an unknown directed acyclic graph (DAG)  $\mathcal{G}$  with nodes  $[p] = \{1,\dots,p\}$ . Concisely, we have the following data-generating process:
$$
X = f (U), \quad U \sim \mathbb {P} _ {U} = \prod_ {i = 1} ^ {p} \mathbb {P} \left(U _ {i} \mid U _ {\mathrm {p a} _ {\mathcal {G}} (i)}\right), \tag {1}
$$
where  $\mathrm{pa}_{\mathcal{G}}(i) = \{j \in [p] : j \to i\}$  denotes the parents of  $i$  in  $\mathcal{G}$ . We also use  $\mathrm{ch}_{\mathcal{G}}(i)$ ,  $\mathrm{de}_{\mathcal{G}}(i)$  and  $\mathrm{an}_{\mathcal{G}}(i)$  to denote the children, descendants and ancestors of  $i$  in  $\mathcal{G}$ . Let  $\mathbb{P}_X$  denote the induced distribution over  $X$ .
We consider atomic (i.e., single-node) interventions on the latent variables. While our main focus is on general types of soft interventions, our proof also applies to hard interventions. In particular, an intervention  $I$  with target  $T_{\mathcal{G}}(I) = i \in [p]$  modifies the joint distribution  $\mathbb{P}_U$  by changing the conditional distribution  $\mathbb{P}(U_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$ . A hard intervention sets the conditional distribution as  $\mathbb{P}^I(U_i)$ , removing the dependency of  $U_i$  on  $U_{\mathrm{pa}_{\mathcal{G}}(i)}$ , whereas a soft intervention is allowed to preserve this dependency but changes the mechanism into  $\mathbb{P}^I(U_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$ . An example of a soft intervention is as a shift intervention [50, 69], which modifies the conditional distribution as  $\mathbb{P}^I(U_i = u + a_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)}) = \mathbb{P}(U_i = u + a_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$ .
Figure 1: Example of the data-generating process, where observed gene expressions  $X = f(U)$  and the distribution of  $U$  factorizes with respect to an unknown DAG  $\mathcal{G}$ .
$u\mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$  for some shift value  $a_i$ . In the following, we will use  $\mathbb{P}_U^I = \prod_{i = 1}^p\mathbb{P}^I (U_i\mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$  to denote the interventional distribution, where  $\mathbb{P}^I (U_j\mid U_{\mathrm{pa}_{\mathcal{G}}(j)}) = \mathbb{P}(U_j\mid U_{\mathrm{pa}_{\mathcal{G}}(j)})$  for  $j\neq T_{\mathcal{G}}(I)$ . We denote the induced distribution over  $X$  by  $\mathbb{P}_X^I$ . In cases where the referred random variable is clear from the context, we abbreviate the subscript and use  $\mathbb{P}^I$  instead.
We consider the setting where we have unpaired data from observational and interventional distributions, i.e.,  $\mathcal{D},\mathcal{D}^{I_1},\dots,\mathcal{D}^{I_K}$ . Here,  $\mathcal{D}$  denotes samples of  $X = f(U)$  where  $U\sim \mathbb{P}_U$ ;  $\mathcal{D}^{I_k}$  denotes samples of  $X$  where  $U\sim \mathbb{P}_U^{I_k}$ . We focus on the scenario where we have at least one intervention per latent node. In the worst case, one intervention per node is necessary for identifiability in linear SCMs [53]. We note that having at least one intervention per latent node is a strict generalization of having exactly one intervention per latent node, since we assume no knowledge of which interventions among  $I_{1},\ldots,I_{K}$  target the same node. Throughout the paper, we assume latent variables  $U$  are unobserved and their dimension  $p$ , the DAG  $\mathcal{G}$ , and the interventional targets of  $I_{1},\ldots,I_{K}$  are unknown. The goal is to identify these given samples of  $X$  in  $\mathcal{D},\mathcal{D}^{I_1},\ldots,\mathcal{D}^{I_K}$ .
3 Equivalence Class for Causal Disentanglement
In this section, we characterize the equivalence class for causal disentanglement, i.e., the class of latent models that can generate the same observed samples of  $X$  in  $\mathcal{D}, \mathcal{D}^{I_1}, \dots, \mathcal{D}^{I_K}$ . Since we only have access to this data, the latent model can only be identified up to this equivalence class.
First, note that the data-generation process is agnostic to the re-indexing of latent variables, provided that we change the mixing function to reflect such re-indexing. More precisely, consider an arbitrary permutation  $\pi$  of  $[p]$ . Denote  $U_{\pi} = (U_{\pi(1)}, \dots, U_{\pi(p)})$  and  $f_{\pi}$  as the mixing function such that  $f_{\pi}(U_{\pi}) = f(U)$ . We define  $\mathcal{G}_{\pi}$  as the DAG with nodes in  $[p]$  and edges  $i \to j$  if and only if  $\pi(i) \to \pi(j) \in \mathcal{G}$ . Then the following data-generating process,
$$
X = f _ {\pi} (U _ {\pi}), \quad U _ {\pi} \sim \mathbb {P} _ {U _ {\pi}} = \prod_ {i = 1} ^ {p} \mathbb {P} \big ((U _ {\pi}) _ {i} \mid (U _ {\pi}) _ {\mathrm {p a} _ {\mathcal {G} _ {\pi}} (i)} \big),
$$
satisfies  $X = f(U)$ . The same argument holds when  $U$  is generated from an interventional distribution  $\mathbb{P}^I$ , where this process generates the same  $X$  when  $U_{\pi}$  is sampled from  $\mathbb{P}_{U_{\pi}}^{I_{\pi}}$ . Here  $I_{\pi}$  is such that  $T_{\mathcal{G}_{\pi}}(I_{\pi}) = \pi^{-1}(T_{\mathcal{G}}(I))$  and the mechanism  $\mathbb{P}^{I_{\pi}}\big((U_{\pi})_i\mid (U_{\pi})_{\mathrm{pa}_{\mathcal{G}_{\pi}}(i)}\big) = \mathbb{P}^I\big(U_{\pi (i)}\mid U_{\mathrm{pa}_{\mathcal{G}}(\pi (i))}\big)$ .
We would also observe the same data if each  $U_{i}$  is affinely transformed into  $\lambda_{i}U_{i} + b_{i}$  for constants  $\lambda_{i} \neq 0$  and  $b_{i}$  and the mixing function is rescaled element-wise to accommodate this transformation. To account for these two types of equivalences, we define the following notion of causal disentanglement (CD) equivalence class.
Definition 1 (CD-Equivalence). Two sets of variables,  $\langle U,\mathcal{G},I_1,\dots,I_K\rangle$  and  $\langle \hat{U},\hat{\mathcal{G}},\hat{I}_1,\dots,\hat{I}_K\rangle$  are CD-equivalent if and only if there exists a permutation  $\pi$  of  $[p]$ , non-zero constants  $\lambda_{1},\ldots ,\lambda_{p}\neq 0$  and  $b_{1},\ldots ,b_{p}$  such that
$$
\hat {U} _ {i} = \lambda_ {\pi (i)} U _ {\pi (i)} + b _ {\pi (i)}, \forall i \in [ p ], \quad \hat {\mathcal {G}} = \mathcal {G} _ {\pi}, \quad \text {a n d} \hat {I} _ {k} = \left(I _ {k}\right) _ {\pi}, \forall k \in [ K ].
$$
The same definition applies to  $\langle \mathcal{G}, I_1, \dots, I_k \rangle$  and  $\langle \hat{\mathcal{G}}, \hat{I}_1, \dots, \hat{I}_k \rangle$ , where we say they are CD-equivalent if and only if  $\hat{\mathcal{G}} = \mathcal{G}_{\pi}$ , and  $\hat{I}_k = (I_k)_{\pi}$  for some permutation  $\pi$ .
For simplicity, we refrain from talking about transformations on the mixing function  $f$  and mechanisms of latent variables. These can be obtained once  $U, \mathcal{G}, T_{\mathcal{G}}(I_1), \dots, T_{\mathcal{G}}(I_K)$  are identified. In particular,  $f$  is the map from  $U$  to the observed  $X$ ; and the joint distribution  $\mathbb{P}_U$  (and  $\mathbb{P}_U^{I_k}$ ) can be decomposed with respect to  $\mathcal{G}$  to obtain the mechanisms  $\mathbb{P}_U(U_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$  (and  $\mathbb{P}_U^{I_k}(U_i \mid U_{\mathrm{pa}_{\mathcal{G}}(i)})$ ).
4 Identifiability Results
In this section, we present our main results, namely the identifiability guarantees for causal disentanglement from soft interventions. For this discussion, we consider the infinite-data regime where enough samples are obtained to exactly determine the observational and interventional distributions  $\mathbb{P}_X,\mathbb{P}_X^{I_1},\dots,\mathbb{P}_X^{I_K}$ . Detailed proofs are deferred to Appendices A and B.
Following [3], we pose assumptions on the support of  $U$  and on the function class of the map  $f$ . Our support assumption is for example satisfied under the common additive Gaussian structural causal model [47], and our assumption on the function class is for example satisfied if  $f$  is linear and injective (Lemma 2 in Appendix A), a setting considered in many identifiability works (e.g., [11, 1, 53]).
Assumption 1. Let  $U$  be a  $p$ -dimensional random vector. Following [3], we assume that the interior of the support of  $\mathbb{P}_U$  is a non-empty subset of  $\mathbb{R}^p$ , and that  $f$  is a full row rank polynomial.
Under this assumption, the authors in [3] showed that if  $p$  is known,  $U$  is identifiable up to a linear transformation. This remains true when  $p$  is unknown, as summarized in the following lemma.
Lemma 1. Under Assumption 1, we can identify the dimension  $p$  of  $U$  as well as its linear transformation  $U\Lambda + b$  for some non-singular matrix  $\Lambda$  and vector  $b$ . In fact, with observational data, we can only identify  $U$  up to such linear transformations.
Denote all pairs of  $\mathbb{P}_U, f$  that satisfy this assumption as  $\mathcal{F}_p$ . The proof of this lemma is provided by solving the following constrained optimization problem:
$$
\min  _ {(\mathbb {P} _ {\hat {U}}, \hat {f}) \in \mathcal {F} _ {\hat {p}}} \hat {p} \quad \text {s u b j e c t} \mathbb {P} _ {\hat {f} (\hat {U})} = \mathbb {P} _ {X}.
$$
There exists some integer  $d$ , a full row rank  $H \in \mathbb{R}^{(p + \dots + p^d) \times n}$  and a vector  $h \in \mathbb{R}^n$  such that  $f(U) = (U, \bar{\otimes} U^2, \dots, \bar{\otimes} U^d) H + h$ , where  $\bar{\otimes} U^k$  denotes the size- $p^k$  vector with degree- $k$  polynomials of  $U$  as its entries.
In other words, let  $\hat{p}$  be the smallest dimension such that there exists a pair of  $\mathbb{P}_{\hat{U}},\hat{f}$  in  $\mathcal{F}_{\hat{p}}$  that generates the observational distribution  $\mathbb{P}_X$ . Then  $p = \hat{p}$  and we recover the latent factors up to linear transformation. The intuition is that (1) the support with non-empty interior guarantees that we can identify  $p$  by checking its geometric dimension, and (2) the full-rank polynomial assumption ensures that we search for  $f$  (and consequently  $U$ ) in a constrained subspace.
On the other hand, to show we cannot identify more than linear transformations, we construct a mixing function  $\hat{f}$  for  $\hat{U} \coloneqq U\Lambda + b$  such that the induced distribution  $\mathbb{P}_X$  is the same under both representations. This also means that we cannot identify the underlying DAG  $\mathcal{G}$  up to any nontrivial equivalence class; we give an example showing that any causal DAG can explain the observational data in Appendix A. Next, we discuss how identifiability can be improved with interventional data.
4.2 Identifying ancestral relations
Lemma 1 guarantees identifiability up to linear transformations from solely observational data. This reduces the problem to the case where an unknown invertible linear mixing of the latent variables  $X = f(U) = U\Lambda + b$  is observed. Without loss of generality, we thus work with this reduction for the remainder of the section.
When the causal variables are fully observed, we can identify causal relationships from the changes made by interventions [57]. In particular, an intervention on a node will not alter the marginals of its non-descendants as compared to the observational distribution, i.e.,  $\mathbb{P}(U_j) = \mathbb{P}^I (U_j)$  for  $T_{\mathcal{G}}(I) = i$  and  $j\notin \deg (i)$ . However, it is possible that  $\mathbb{P}(U_j) = \mathbb{P}^I (U_j)$  for some  $j\in \deg (i)$  in degenerate cases where the change made by  $U_{i}$  is canceled out on the path from  $i$  to  $j$ . Hence, prior works<sup>3</sup> defined influentiality or interventional faithfulness [57, 66], which avoids such degenerate cases by assuming that intervening on a node will always change the marginals of all its descendants, i.e.,  $\mathbb{P}(U_j)\neq \mathbb{P}^I (U_j)$  for  $j\in \mathrm{de}_{\mathcal{G}}(i)$ . Under this assumption, we can identify the descendants of an intervention target in  $\mathcal{G}$ , by testing if a node has a changed marginal interventional distribution.
However, if we only observe a linear mixing of the causal variables, interventional faithfulness is not enough to identify such ancestral relations. Consider the following example.
Example 1. Let  $\mathcal{G} = \{1\to 2\}$  with  $\mathbb{P}(U_1) = \mathcal{N}(0,1)$  and  $\mathbb{P}(U_2|U_1) = \mathcal{N}(U_1,1)$ . Suppose that  $T_{\mathcal{G}}(I_1) = 1$ , with  $\mathbb{P}^{I_1}(U_1) = \mathcal{N}(1,1)$ , and that  $T_{\mathcal{G}}(I_2) = 2$ , with  $\mathbb{P}^{I_2}(U_2|U_1) = \mathcal{N}(U_1 + 1,1)$ . Note that this model satisfies interventional faithfulness.
Let  $f$  be the identity map, i.e.,  $X = U$ . Consider latent variables  $\hat{U} = (U_2, U_2 - U_1)$  and  $\hat{f}(\hat{U}) = (\hat{U}_1 - \hat{U}_2, \hat{U}_1)$ . Then  $X = \hat{f}(\hat{U}) = f(U)$ . However, we have  $\hat{\mathcal{G}} = \{2 \to 1\}$  with  $\mathbb{P}(\hat{U}_2) = \mathcal{N}(0,1)$  and  $\mathbb{P}(\hat{U}_1 \mid \hat{U}_2) = \mathcal{N}(\hat{U}_2, 1)$ ,  $T_{\hat{\mathcal{G}}}(\hat{I}_1) = 1$  with  $\mathbb{P}^{I_1}(\hat{U}_1 \mid \hat{U}_2) = \mathcal{N}(\hat{U}_2 + 1, 1)$ , and  $T_{\hat{\mathcal{G}}}(\hat{I}_2) = 2$  with  $\mathbb{P}^{I_2}(\hat{U}_2) = \mathcal{N}(1,1)$ . We thus may reverse ancestral relations between the intervention targets, as illustrated in Figure 2.
This example shows that the effect on  $U_{2}$  from intervening on  $U_{1}$  can be canceled out by linearly combining  $U_{2}$  with  $U_{1}$ . In other words, intervening on  $U_{1}$  does not change the marginal distribution of  $U_{2} - U_{1}$ , even under interventional faithfulness. Thus, we need a stronger faithfulness assumption to account for the effect of linear mixing. In general, we want to avoid the case that the effect of an intervention on a downstream variable  $U_{j}$  can be canceled out by combining  $U_{j}$  linearly with other variables.
Assumption 2. Intervention  $I$  with target  $i$  satisfies linear interventional faithfulness if for every  $j \in \{i\} \cup \mathrm{ch}_{\mathcal{G}}(i)$  such that  $\mathrm{pa}_{\mathcal{G}}(j) \cap \mathrm{de}_{\mathcal{G}}(i) = \emptyset$ , it holds that  $\mathbb{P}(U_j + U_SC^\top) \neq \mathbb{P}^I (U_j + U_SC^\top)$  for all constant vectors  $C \in \mathbb{R}^{|S|}$ , where  $S = [p] \setminus (\{j\} \cup \mathrm{de}_{\mathcal{G}}(i))$ .
This assumption ensures that an intervention on  $U_{i}$  not only affects its children, but that the effect remains even when we take a linear combination of a child with certain other variables. Note that the condition need only hold for the most upstream children of  $U_{i}$ , which may be arbitrarily smaller than the set of all children of  $U_{i}$ . To illustrate this assumption, we give a simple example on a 2-node
<sup>3</sup>A more detailed discussion of interventional faithfulness can be found in Appendix B.1.
DAG where this assumption is generically satisfied. In general, we show in Appendix B that a large class of non-linear SCMs and soft interventions satisfy this assumption.
Example 2. Consider  $\mathcal{G} = \{1\to 2\}$ . Let  $\mathbb{P}(U_2\mid U_1) = \mathcal{N}(\beta U_1^2,\sigma_2^2)$  and  $\mathbb{P}(U_1) = \mathcal{N}(0,\sigma_1^2)$ . Intervention  $I$  that changes  $\mathbb{P}(U_1)$  into  $\mathcal{N}(0,\sigma_1'^2)$  satisfies Assumption 2 as long as  $\beta \neq 0$ . To see this, note that  $\mathbb{P}(U_2 + U_1C)\neq \mathbb{P}^I (U_2 + U_1C)$  for any  $C$ , since  $\mathbb{E}_{\mathbb{P}}(U_2 + U_1C) = \beta \sigma_1^2\neq \beta \sigma_1'^2 = \mathbb{E}_{\mathbb{P}^I}(U_2 + U_1C)$ .