SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Under Assumption 2, we can show that we can identify causal relationships by detecting marginal changes made by interventions. In particular, consider an easier setting where  $K = p$ , i.e., we have exactly one intervention per latent node. For a source node  $i$  of  $\mathcal{G}$ ,  $\mathbb{P}(U_i) \neq \mathbb{P}^I (U_i)$  if and only if  $T_{\mathcal{G}}(I) = i$ . Therefore the source node will have its marginal changed under one intervention amongst  $\{I_1,\dots,I_p\}$ . This is a property of the latent model that we can utilize when solving for it.
Since we have access to  $X = U\Lambda + b$ , we solve for  $U_i$  in the form of  $XC^\top + c$  with  $C \in \mathbb{R}^n$ ,  $c \in \mathbb{R}$ , or equivalently,  $UC^\top + c$  with  $C \in \mathbb{R}^p$ . By enforcing that  $V = UC^\top + c$  only has  $\mathbb{P}(V) \neq \mathbb{P}^I(V)$  for one  $I \in \{I_1, \dots, I_p\}$ , Assumption 2 guarantees that  $V$  can only be an affine transformation of a source node and that this  $I$  corresponds to intervening on this source node. Otherwise: (1) if  $C_j \neq 0$  for a non-source node  $j$ , take  $j$  to be the most downstream node with  $C_j \neq 0$ , then  $\mathbb{P}(V) \neq \mathbb{P}^I(V)$  for at least two  $I$ 's targeting  $j$  and its most downstream parents in  $\mathrm{pa}_{\mathcal{G}}(j)$ ; (2) if  $C_{i_1} \neq 0$  and  $C_{i_2} \neq 0$  for two source nodes  $i_1, i_2$ , then  $\mathbb{P}(V) \neq \mathbb{P}^I(V)$  for two  $I$ 's targeting  $i_1$  and  $i_2$ .
In general, we can apply this argument to identify all interventions in  $I_1, \ldots, I_K$  that target source nodes of  $\mathcal{G}$ . Then using an iterative argument, we can identify all interventions that target source nodes of the subgraph of  $\mathcal{G}$  after removing its source nodes. This procedure results in the ancestral relations between the targets of  $I_1, \ldots, I_K$ . Namely, if  $T_{\mathcal{G}}(I_k) \in \mathrm{an}_{\mathcal{G}}(T_{\mathcal{G}}(I_j))$ , then  $I_j$  is identified in a later step than  $I_k$  in the above procedure. We thus have the following theorem.
Theorem 1. Under Assumption 1 and Assumption 2 for  $I_1, \ldots, I_K$ , we can identify  $\langle \hat{\mathcal{G}}, \hat{I}_1, \dots, \hat{I}_K \rangle$  where  $\hat{\mathcal{G}} = \mathcal{T}S(\mathcal{G}_{\pi})$ , and  $\hat{I}_k = (I_k)_{\pi}$  for some permutation  $\pi$ .
Here  $\mathcal{T}S$  denotes the transitive closure of a DAG [57], where  $i\rightarrow j\in \mathcal{T}S(\mathcal{G})$  if and only if  $i\in \mathrm{an}_{\mathcal{G}}(j)$ . Note that this limitation is not due to the linear mixing of the causal variables. It was shown in [57] that with fully observed causal variables, one can only identify a DAG up to its transitive closure by detecting marginal distribution changes. In the next section, we show how to reduce  $\mathcal{T}S(\mathcal{G}_{\pi})$  to  $\mathcal{G}_{\pi}$ , i.e., identifying the CD-equivalence class of  $\langle \mathcal{G},I_1,\dots,I_k\rangle$ .
4.3 Identifying direct edges
DAGs with the same transitive closure can span a spectrum of sparsities; for example, a complete graph and a line graph with the same topological ordering have the same transitive closure. The following example shows that under Assumption 2, in some cases we cannot identify more than the transitive closure.
Example 3. Let  $\mathcal{G}$  be the 3-node DAG shown on the left in Figure 3. Suppose that  $\mathbb{P}(U_1)$  is  $\mathcal{N}(0,1)$ ,  $\mathbb{P}(U_2|U_1)$  is  $\mathcal{N}(U_1^2,1)$ ,
and  $\mathbb{P}(U_3|\bar{U}_1,U_2)$  is  $\mathcal{N}((\bar{U}_1 + U_2)^2,1)$ . Let  $f$  be the identity map and  $I_{1},I_{2},I_{3}$  target nodes 1,2,3, respectively, each changing their conditional variances to 2.
Now consider a different model with variables  $\hat{U} = (U_1, U_1 + U_2, U_3)$  and mixing function  $\hat{f}(\hat{U}) = (\hat{U}_1, \hat{U}_2 - \hat{U}_1, \hat{U}_3)$ . Then  $\hat{f}(\hat{U}) = U = f(U) = X$ . The distributions  $\mathbb{P}(\hat{U}), \mathbb{P}^{I_1}(\hat{U}), \mathbb{P}^{I_2}(\hat{U})$ , and  $\mathbb{P}^{I_3}(\hat{U})$  each factorizes according to the DAG  $\hat{\mathcal{G}}$  that is missing the edge  $1 \rightarrow 3$  (Figure 3), where we let  $I_1, I_2$  and  $I_3$  change the conditional variances of  $\hat{U}_1$ ,  $\hat{U}_2$ , and  $\hat{U}_3$  to 2, respectively.
This example shows that we cannot identify  $1 \to 3$  since  $U_{1} \perp U_{3}|U_{1} + U_{2}$ . In the case when the causal variables are fully observed,  $1 \to 3$  can be identified by assuming  $U_{1} \not\perp U_{3}|U_{2}$ . However, when allowing for linear mixing, we need to avoid cases such as  $U_{1} \perp U_{3}|U_{1} + U_{2}$  in order to be able to identify  $1 \to 3$ . We will show that the following assumption guarantees identifiability of  $\mathcal{G}$ .
4 A source node is a node without parents.
<sup>5</sup>We show in Appendix B that this model satisfies Assumptions 1 and 2.
When  $\mathcal{G}$  is a polytree (a DAG whose skeleton is a tree), this assumption is implied by Assumption 2 under mild regularity conditions (proven in Appendix B). Thus if  $\mathcal{G}$  is the sparsest DAG within its transitive closure, we can always identify it with just the Assumptions 1 and 2.
Assumption 3. For every edge  $i \to j \in \mathcal{G}$ , there do not exist constants  $c_{j}, c_{k} \in \mathbb{R}$  for  $k \in S$  such that  $U_{i} \perp U_{j} + c_{j} U_{i} \mid \{U_{l}\}_{l \in \mathrm{pa}_{\mathcal{G}}(j) \setminus (S \cup \{i\})}$ ,  $\{U_{k} + c_{k} U_{i}\}_{k \in S}$ , where  $S = \mathrm{pa}_{\mathcal{G}}(j) \cap \mathrm{de}_{\mathcal{G}}(i)$ .
Theorem 2. Under Assumptions 1,2,3,  $\langle \mathcal{G},I_1,\dots,I_K\rangle$  is identifiable up to its CD-equivalence class.
4.4 Further remarks
Next we discuss if it is possible to recover  $U$  along with  $\langle \mathcal{G}, I_1, \ldots, I_k \rangle$  up to their CD-equivalence class. Note a simple contradiction with  $\mathcal{G} = \{1 \to 2\}$ : since we consider general soft interventions, there will always be a valid explanation if we add  $U_1$  to  $U_2$ . Therefore even when we can identify  $\langle \mathcal{G}, I_1, \ldots, I_k \rangle$  up to its CD-equivalence class, we still cannot identify  $U$  in an element-wise fashion. However, our identifiability results still allow us to draw causal explanations and predict the effect of unseen combinations of interventions, as we discuss below.
Application of Theorem 1 and 2. Given unpaired data  $\mathcal{D},\mathcal{D}^{I_1},\ldots \mathcal{D}^{I_K}$ , these two theorems guarantee that we can identify which  $I_{1},\dots,I_{K}$  correspond to intervening on the same latent node. Furthermore, Theorem 1 shows that we are able to identify ancestral relationships between the intervention targets of  $I_{1},\dots,I_{K}$ , while Theorem 2 guarantees identifiability of the exact causal structure.
For example, given high-dimensional single-cell transcriptomic readout from a genome-wide knockdown screen, we can under Assumption 1 identify the number of latent causal variables (which we can interpret as the programs of a cell), under Assumption 2 identify which genes belong to the same program, and under Assumption 3 identify the full regulatory relationships between the programs.
Extrapolation to unseen combinations of interventions. Theorems 1 and 2 also guarantee that we can predict the effect of unseen combinations of interventions. Namely, consider a combinatorial intervention  $\mathcal{I} \subset \{I_1, \dots, I_K\}$ , where  $T_{\mathcal{G}}(I) \neq T_{\mathcal{G}}(I')$  for all  $I \neq I' \in \mathcal{I}$ . In other words,  $\mathcal{I}$  is an intervention with multiple intervention targets that is composed by combining interventions among  $I_1, \dots, I_K$  with different targets.
Denote by  $\langle \hat{U},\hat{\mathcal{G}},\hat{I}_1,\dots ,\hat{I}_K\rangle$  the latent model identified from the interventions  $\{I_1,\ldots ,I_K\}$ . Recall from Section 3 that we can also infer the mixing function  $\hat{f}$  and mechanisms from  $\hat{U},\hat{\mathcal{G}},\hat{I}_1,\dots ,\hat{I}_K$ . From this, we can infer the interventional distribution under the combinatorial intervention  $\mathcal{I}$ :
$$
X = \hat {f} (\hat {U}), \hat {U} \sim \mathbb {P} _ {\hat {U}} ^ {\hat {\mathcal {I}}} = \prod_ {\hat {I} \notin \mathcal {I}} \mathbb {P} _ {\hat {U}} \big (\hat {U} _ {T _ {\hat {\mathcal {G}}} (\hat {I})} \mid \hat {U} _ {\mathrm {p a} _ {\hat {\mathcal {G}}} (T _ {\hat {\mathcal {G}}} (\hat {I}))} \big) \cdot \prod_ {\hat {I} \in \mathcal {I}} \mathbb {P} _ {\hat {U}} ^ {\hat {I}} \big (\hat {U} _ {T _ {\hat {\mathcal {G}}} (\hat {I})} \mid \hat {U} _ {\mathrm {p a} _ {\hat {\mathcal {G}}} (T _ {\hat {\mathcal {G}}} (\hat {I}))} \big). (2)
$$
We state the conditions for this result informally in the following theorem. A formal version of this theorem together with its proof are given in Appendix B.5.
Theorem 3 (Informal). Let  $\mathcal{I}$  be a combinatorial intervention (i.e., with multiple intervention targets) combining several interventions among  $I_{1},\ldots,I_{K}$  with different targets. The above procedure allows sampling  $X$  according to the distribution  $X = f(U),U\sim \mathbb{P}^{\mathcal{I}}$ .
5 Discrepancy-based VAE Formulation
Having shown identifiability guarantees for causal disentanglement, we now focus on developing a practical algorithm for recovering the CD-equivalence class from data. As indicated by our proof of Theorem 2, the latent causal graph can be identified by taking the sparsest model compatible with the data. This characterization suggests maximizing a penalized log-likelihood score, a common method for model selection in causal structure learning [10]. The resulting challenging combinatorial optimization problem has been tackled using a variety of approaches, including exact search using integer linear programming [12], greedy search [10, 48, 52], and more recently, gradient-based approaches where the combinatorial search space is relaxed to a continuous search space [70, 34, 38, 61].
Gradient-based approaches offer several potential benefits, including scalability, ease of implementation in automatic differentiation frameworks, and significant flexibility in the choice of components. In light of these benefits, we opted for a gradient-based approach to optimization. In
Figure 4: Our proposed CausalDiscrepancyVAE architecture. Gray boxes represent inputs, white boxes the generated values, blue boxes the trainable modules, and orange boxes the terms of the loss function. Dashed lines indicate copies of the same module or related modules. Highlighted boxes show the procedure to generate virtual counterfactual samples.
particular, we replace the log-likelihood term of our objective function with a variational lower bound by employing the framework of autoencoding variational Bayes (AVB), widely used in prior works for causal disentanglement [36, 6]. To employ AVB, we re-parameterize each distribution  $\mathbb{P}(U_i\mid U_{\mathrm{pa}_g(i)})$  in Eq. (1) into  $U_{i} = s_{i}(U_{\mathrm{pa}_{g}(i)},Z_{i})$ , where  $Z_{i}$  is an independent exogenous noise variable and  $s_i$  denotes the causal mechanism that generates  $U_{i}$  from  $U_{\mathrm{pa}_{g}(i)}$  and  $Z_{i}$ . We let  $p(Z)$  be a prior distribution over  $Z$  and  $p_{\theta ,\varnothing}(X\mid Z)$  be the conditional distribution of  $X$  given  $Z$  under no intervention, thereby defining the marginal distribution  $p_{\theta ,\varnothing}(X)$ . Given an arbitrary distribution  $q_{\phi}(Z\mid X)$ , we have the following well-known inequality (often called the Evidence Lower Bound or ELBO) for any sample  $x$ :
$$
\log p _ {\theta , \varnothing} (x) \geq \mathcal {L} _ {\theta , \phi} ^ {\text {r e c o n}} (x) + \mathcal {L} _ {\phi} ^ {\text {r e g}} (x), \quad \text {w h e r e} \quad \mathcal {L} _ {\theta , \phi} ^ {\text {r e c o n}} (x) := \mathbb {E} _ {q _ {\phi} (Z | x)} \log p _ {\theta , \varnothing} (x \mid Z),
$$
$$
\mathcal {L} _ {\phi} ^ {\operatorname {r e g}} (x) := - D _ {\mathrm {K L}} \left(q _ {\phi} (Z \mid x) \| p (Z)\right).
$$
Putting this into the framework of an autoencoder, we call the distribution  $q_{\phi}$  the encoder and the distribution  $p_{\theta, \varnothing}$  the decoder. In our case, the decoder is composed of two functions. First, a deep structural causal model  $(A_{\theta}, s_{\theta, \varnothing})$  maps the exogenous noise  $Z$  to the causal variables  $U^{\varnothing}$ . In particular, the adjacency matrix  $A_{\theta}$  defines the parent set for each variable, while  $s_{\theta, \varnothing} = \{(s_{\theta, \varnothing})_i\}_{i=1}^p$  denotes the learned causal mechanisms. Second, a mixing function  $f_{\theta}$  maps the causal variables  $U^{\varnothing}$  to the observed variables  $X^{\varnothing}$ . Because of the permutation symmetry of CD-equivalence, we can fix  $A_{\theta}$  to be upper triangular without loss of generality. We add a loss term  $\mathcal{L}_{\theta}^{\text{sparse}} := -\| A_{\theta}\|_1$  to encourage  $A_{\theta}$  to be sparse.
While the observational samples are generated from the distribution  $p_{\theta, \varnothing}$ , the interventional samples are drawn from a different but related distribution  $p_{\theta, I}$ . The modularity of our decoder allows us to replace  $(A_{\theta}, m_{\theta, \varnothing})$  with an interventional counterpart  $(A_{\theta}, m_{\theta, I})$ , while keeping the mixing function  $f_{\theta}$  constant. This is illustrated by the highlighted boxes in Fig. 4. For each intervention label  $I$ , the corresponding intervention target  $i$  and a shift<sup>6</sup>  $a_i$  is determined by an intervention encoder  $T_{\phi}$ , which uses softmax normalization to approximate a one-hot encoding of the intervention target. Given these intervention targets, we generate "virtual" counterfactual samples for each observational sample. Such samples follow the distribution  $\mathbb{P}_{\theta, \phi}(\widehat{X}^{\hat{I}_k})$ , the pushforward of  $\mathbb{P}_X^{\varnothing}$  under the action of the encoder  $q_{\phi}$  and decoder  $p_{\theta, I}$ . These samples are compared to real samples from the corresponding interventional distribution. A variety of discrepancy measures can be used for this comparison. To avoid the saddle point optimization challenges that come with adversarial training, we do not consider adversarial methods (e.g. the dual form of the Wasserstein distance in [5]).
For simplicity, we parameterize interventions in DSCM as shifts, though the theoretical results hold for general nonparametric interventions.
Figure 5: The distribution of generated samples mirrors the distribution of actual samples. Samples are visualized using UMAP. Left: Samples from the 14 single target-node interventions with more than 800 cells. Middle-Right: Samples for target genes SET, CEBPE, and KLF1.
This leaves non-adversarial discrepancy measures, such as the MMD (Maximum Mean Discrepancy) [18], the entropic Wasserstein distance [15], and the sliced Wasserstein distance [63]. In this work, we focus on the MMD measure, whose empirical estimate we recall in Appendix C.1. We take  $\mathcal{L}_{\theta,\phi}^{\mathrm{discrep}} := -\sum_{k=1}^{K} \mathrm{MMD}(\mathbb{P}_{\theta,\phi}(\widehat{X}^{\hat{I}_k}), \mathbb{P}_X^{\hat{I}_k})$ . Thus, the full loss function used during training is
$$
\mathcal {L} _ {\theta , \phi} ^ {\alpha , \beta , \lambda} := \mathbb {E} _ {X \emptyset} \left[ \mathcal {L} _ {\theta , \phi} ^ {\text {r e g}} (X) + \beta \mathcal {L} _ {\phi} ^ {\text {r e c o n}} (X) \right] + \alpha \mathcal {L} _ {\theta , \phi} ^ {\text {d i s c r e p}} + \lambda \mathcal {L} _ {\theta} ^ {\text {s p a r s e}}. \tag {3}
$$
A diagram of the proposed architecture is shown in Fig. 4. Values of the hyperparameters  $\alpha, \beta, \lambda$  used in our loss function as well as other hyperparameters are described in Appendix F.
Our loss function exhibits several desirable properties. First, as we show in Appendix D, the unpaired data loss function lower bounds the paired data log-likelihood that one would directly optimize in the oracle setting where true counterfactual pairs were available. Second, as we show in Appendix E, this procedure is consistent, in the sense that optimizing the loss function in the limit of infinite data will recover the generative process (under suitable conditions). This consistency result also guarantees that the learned model can consistently predict the effect of multi-node interventions; see Appendix E.2.
We now demonstrate our method on a biological dataset. We use the large-scale Perturb-seq study from [44]. After pre-processing, the data contains 8,907 unperturbed cells (observational dataset  $\mathcal{D}$ ) and 99,590 perturbed cells. The perturbed cells underwent CRISPR activation [16] targeting one or two out of 105 genes (interventional datasets  $\mathcal{D}^1,\dots,\mathcal{D}^K$ ,  $K = 217$ ). CRISPR activation experiments modulate the expression of their target genes, which we model as a shift intervention. Each interventional dataset comprises 50 to 2,000 cells. Each cell is represented as a 5,000-dimensional vector (observed variable  $X$ ) measuring the expressions of 5,000 highly variable genes.
To test our model, we set the latent dimension  $p = 105$ , corresponding to the total number of targeted genes. During training, we include all the unperturbed cells from  $\mathcal{D}$  and the perturbed cells from the single-node interventional datasets  $\mathcal{D}^1, \dots, \mathcal{D}^{105}$  that target one gene. For each single-node interventional dataset with over 800 cells, we randomly extract 96 cells and reserve these for testing. The double-node interventions (112 distributions  $\mathcal{D}^{106}, \dots, \mathcal{D}^{217}$ ) targeting two genes are entirely reserved for testing. The following results summarize the model with the best training performance. Extended evaluations and detailed implementation can be found in Appendix F and G. In additional, we also provide ablation studies on biological data and a simple simulation study in Appendix H.
Single-node Interventional Distributions. To study the generative capacity of our model for interventions on single genes, we produce 96 samples for each single-node intervention with over 800 cells (14 interventions). We compare these against the left-out 96 cells of the corresponding distributions. Figure 5 illustrates this for 3 example genes in 2 dimensions using UMAP [41] with all other cells in the dataset as background (labeled by 'NA'). Our model is able to discover subpopulations of the interventional distributions (e.g., for KLF1, the generated samples are concentrated in the middle left corner). We provide a quantitative evaluation for all 105 single-node interventions in Figure 6. The model is able to obtain close to perfect  $R^2$  (on average 0.99 over all genes and 0.95 over most differentially expressed genes).
Double-node Interventional Distributions. Next, we analyze the generalization capabilities of our model to the 112 double-node interventions. Despite never observing any cells from these interventions during training, we obtain reasonable  $R^2$  values (on average 0.98 over all genes and
Figure 6: Our model accurately predicts the effect of single-node interventions. 'All genes' indicates measurements using the entire 5000-dimensional vectors; 'DE genes' indicates measurements using the 20-dimensional vectors corresponding to the top 20 most differentially expressed genes.
0.88 over most differentially expressed genes). However, when looking at the generated samples for individual pairs of interventions, it is apparent that our model performs well on many pairs, but recovers different subpopulations for some pairs (examples shown in Figure 13 in Appendix G). The wrongly predicted intervention pairs could indicate that the two target genes act non-additively, which needs to be further evaluated and is of independent interest for biological research [22, 44].
Structure Learning. Lastly, we examine the learned DAG between the intervention targets. Specifically, this corresponds to a learned gene regulatory network between the learned programs of the target genes. For this, we reduce  $p$  from 105 until the learned latent targets of  $\mathcal{D}^1, \ldots, \mathcal{D}^{105}$  cover all  $p$  latent nodes. This results in  $p = 7$  groups of genes, where genes are grouped by their learned latent nodes. We then run our algorithm with fixed  $p = 7$  multiple times and take the learned DAG with the least number of edges. This DAG over the groups of
targeted genes is shown with example genes in Figure 7 (left). This learned structure is in accordance with previous findings. For example, we successfully identified the edges DUSP9  $\rightarrow$  MAPK1 and DUSP9  $\rightarrow$  ETS2, which is validated in [44] (see their Fig. 5). We also show the interventional distributions targeting these example genes in Figure 7 (right). Among these, MAPK1 and ETS2 correspond to clusters that are heavily overlapping, which explains why the model maps both distributions to the same latent node.
Figure 7: Structure learning on the biological dataset. Left: learned DAG between target genes (colors indicate edge weights). Right: UMAP visualization of the distributions.
We derived identifiability results for causal disentanglement from single-node interventions, and presented an autoencoding variational Bayes framework to estimate the latent causal representation from interventional samples. Identification of the latent causal structure and generalization to multi-node interventions was demonstrated experimentally on genetic data.
Limitations and Future Work. This paper has various limitations that may be useful to address in future work. We provide an overview here, where an extended discussion can be found in Appendix I. In addition, we also provide a brief summary of concurrent related works in Appendix I.
First, we focused on the setting where single-node interventions on each latent node are available. This is overly optimistic for example in the case of chemical perturbations where the available drugs could all have multiple targets. The VAE framework can still be applied in such settings; however, its theoretical guarantees are subject to further investigations. The key techniques in our proofs can be generalized to the multi-node setting, but further assumptions on the set of interventions needed for identifiability are required. Second, we have focused on the infinite data regime for analyzing identifiability; given that obtaining interventional samples tends to be expensive in practice, there is much room for further investigations in terms of sample complexity.
We thank the Causal Representation Learning Workshop at Bellairs Institute for helpful discussions. All authors acknowledge support by the MIT-IBM Watson AI Lab. In addition, J. Zhang, C. Squires and C. Uhler acknowledge support by the NSF TRIPODS program (DMS-2022448), NCCIH/NIH (1DP2AT012345), ONR (N00014-22-1-2116), the United States Department of Energy (DOE), Office of Advanced Scientific Computing Research (ASCR), via the M2dt MMICC center (DE-SC0023187), the Eric and Wendy Schmidt Center at the Broad Institute, and a Simons Investigator Award.