SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Simon Buchholz $^{1*}$
Goutham Rajendran\*
Elan Rosenfeld2
Bryon Aragam<sup>3</sup>
Bernhard Scholkopf<sup>1</sup>
Pradeep Ravikumar2
<sup>1</sup>Max Planck Institute for Intelligent Systems, Tübingen, Germany
$^{2}$ Carnegie Mellon University, Pittsburgh, USA
<sup>3</sup>University of Chicago, Chicago, USA
We study the problem of learning causal representations from unknown, latent interventions in a general setting, where the latent distribution is Gaussian but the mixing function is completely general. We prove strong identifiability results given unknown single-node interventions, i.e., without having access to the intervention targets. This generalizes prior works which have focused on weaker classes, such as linear maps or paired counterfactual data. This is also the first instance of identifiability from non-paired interventions for deep neural network embeddings and general causal structures. Our proof relies on carefully uncovering the high-dimensional geometric structure present in the data distribution after a non-linear density transformation, which we capture by analyzing quadratic forms of precision matrices of the latent distributions. Finally, we propose a contrastive algorithm to identify the latent variables in practice and evaluate its performance on various tasks.
Modern generative models such as GPT-4 [59] or DDPMs [26] achieve tremendous performance for a wide variety of tasks [7]. They do this by effectively learning high-level representations which map to raw data through complicated non-linear maps, such as transformers [89] or diffusion processes [79]. However, we are unable to reason about the specific representations they learn, in particular they are not necessarily related to the true underlying data generating process. Besides their susceptibility to bias [20], they often fail to generalize to out-of-distribution settings [5], rendering them problematic in safety-critical domains. In order to gain a deeper understanding of what representations deep generative models learn, one line of work has pursued the goal of causal representation learning (CRL) [75]. CRL aims to learn high-level representations of data while simultaneously recovering the rich causal structure embedded in them, allowing us to reason about them and use them for higher-level cognitive tasks such as systematic generalization and planning. This has been used effectively in many application domains including genomics [52] and robotics [43, 94].
A crucial primitive in representation learning is the fundamental notion of identifiability [37, 70], i.e., the question whether a unique (up to tolerable ambiguities) model can explain the observed data. It is well known that because of non-identifiability, CRL is impossible in general settings in the absence of inductive biases or additional information [29, 50]. In this work, we consider additional information
in the form of interventional data [75]. It is common to have access to such data in many application domains such as genomics and robotics stated above (e.g. [17, 58, 57]). Moreover, there is a pressing need in such safety-critical domains to build reliable and trustworthy systems, making identifiable CRL particularly important. Therefore, it's important to study whether we can and also how to perform identifiable CRL from raw observational and interventional data. Here, identifiability opens the possibility to provably recover the true representations with formal guarantees. Meaningfully learning such representations with causal structure enables better interpretability, allows us to reason about fairness, and helps with performing high-level tasks such as planning and reasoning.
In this work, we study precisely this problem of causal representation learning in the presence of interventions. While prior work has studied simpler settings of linear or polynomial mixing [84, 88, 2, 71, 11], we allow for general non-linear mixing, which means our identifiability results apply to complex real-world systems and datasets which are used in practice. With our results, we make progress on fundamental questions on interventional learning raised by [75].
Concretely, we consider a general model with latent variables  $Z$  and observed data  $X$  generated as  $X = f(Z)$  where  $f: \mathbb{R}^d \to \mathbb{R}^{d'}$  is an arbitrary non-linear mixing function. We assume  $Z$  satisfies a Gaussian structural equation model (SEM) which is unknown and unobserved. A Gaussian prior is commonly used in practice (which implies a linear SEM over  $Z$ ) and further, having a simple model for  $Z$  allows us to learn useful representations, generate data efficiently, and explore the latent causal relationships, while the non-linearity of  $f$  ensures universal approximability of the model. We additionally assume access to interventional data  $X^{(i)} = f(Z^{(i)})$  for  $i \in I$  where  $Z^{(i)}$  is the latent under an intervention on a single node  $Z_{t_i}$ . Notably, we allow for various kinds of interventions, we don't require knowledge of the targets  $t_i$ , and we don't require paired data, (i.e., we don't need counterfactual samples from the joint distribution  $(X, X^{(i)})$ , but only their marginals). Having targets or counterfactual data is unrealistic in many practical settings but many prior identifiability results require them, so eliminating this dependence is a crucial step towards CRL in the real world.
Contributions Our main contribution is a general solution to the identifiability problem that captures practical settings with non-linear mixing functions (e.g. deep neural networks) and unknown interventions (since the targets are latent). We study both perfect and imperfect interventions and additionally allow for shift interventions, where perfect interventions remove the dependence of the target variable from its parents, while imperfect interventions (also known as soft interventions) modify the dependencies. Below, we summarize our main contributions:
1. We show identifiability of the full latent variable model, given non-paired interventional data with unknown intervention targets. In particular, we learn the mixing, the targets and the latent causal structure.
2. Compared to prior works that have focused on linear/polynomial  $f$ , we allow non-linear  $f$  which encompasses representations learned by e.g. deep neural networks and captures complex real-world datasets. Moreover, we study both imperfect (also called soft) and perfect interventions, and always allow shift interventions.
3. We construct illustrative counterexamples to probe tightness of our assumptions which suggest directions for future work.
4. We propose a novel algorithm based on contrastive learning to learn causal representations from interventional data, and we run experiments to validate our theory. Our experiments suggest that a contrastive approach, which so far has been unexplored in interventional CRL, is a promising technique going forward.
Causal representation learning [75, 74] has seen much recent progress and applications since it generalizes and connects the fields of Independent Component Analysis [13, 28, 30], causal inference [81, 61, 63, 64, 82] and latent variable modeling [40, 3, 78, 96, 32]. Fundamental to this approach is the notion of identifiability [36, 15, 93]. Due to non-identifiability in general settings without inductive biases [29, 50], prior works have approached this problem from various angles — using additional auxiliary labels [36, 27, 6, 76]; by imposing sparsity [75, 42, 54, 103]; or by restricting the function classes [41, 9, 104, 22]. See also the survey [32]. Moreover, a long line of works have proposed practical methods for CRL (which includes causal disentanglement as a special case),
(a) No interventions
(b) An imperfect intervention
(c) A perfect intervention
Figure 1: Illustration of an example latent variable model under interventions. (a) No interventions. (b) An imperfect intervention on node  $t_i = 3$ . Dashed edges indicate the weights could have potentially been modified. (c) A perfect intervention on node  $t_i = 3$ .
[19, 95, 16, 97, 44, 14] to name a few. It's worth noting that most of these approaches are essentially variants of the Variational Autoencoder framework [39, 68].
Of particular relevance to this work is the setting when interventional data is available. We first remark that the much simpler case of fully observed variables (i.e., no latent variables), has been studied in e.g. [23, 62, 83, 34, 18] (see also the survey [82]). In this work, we consider the more difficult setting of structure learning over latent variables, which have been explored in [46, 42, 6, 104, 1, 84, 88, 2, 71, 11, 72]. Among these, [46] assumes that the intervention targets are known, [42] specifically consider instance-level pre- and post- interventions for time-series data and [6, 104, 1] assumes access to paired counterfactual data. In contrast, we assume unknown targets and work in general settings with non-paired interventional data, which is important in various real-world applications [85, 99]. The work [48] require several graphical restrictions on the causal graph and also require  $2d$  interventions, while we make no graphical restrictions and only require  $d$  interventions (which we also show cannot be improved under our assumptions). [71, 11, 84, 88] consider linear mixing functions  $f$  whereas we study general non-linear  $f$ . Finally, [2] consider polynomial mixing in the more restricted class of deterministic do-interventions. Several concurrent works study related settings [100, 35, 65, 45, 92]. For a more detailed comparison to the related works, we refer to Appendix C.
Our proposed algorithm is based on contrastive learning. Contrastive learning has been used in other contexts in this domain [31, 27, 55, 104, 91] (either in the setting of time-series data or paired counterfactual data), however the application to non-paired interventional settings is new to the best of our knowledge.
Notation In this work, we will almost always work with vectors and matrices in  $d$  dimensions where  $d$  is the latent dimension, and we will disambiguate when necessary. For a vector  $v$ , we denote by  $v_{i}$  its  $i$ -th entry. Let  $\mathrm{Id}$  denote the  $d\times d$  identity matrix with columns as the standard basis vectors  $e_1,\ldots ,e_d$ . We denote by  $N(\mu ,\Sigma)$  the multivariate normal distribution with mean  $\mu$  and covariance  $\Sigma$ . For a set  $C$ , let  $\mathcal{U}(C)$  denote the uniform distribution on  $C$ . For two random variables  $X,X^{\prime}$ , we write  $X\stackrel {D}{=}X^{\prime}$  if their distributions are the same. We denote the set  $\{1,\dots ,d\}$  by  $[d]$ . For permutation matrices we use the convention that  $(P_{\omega})_{ij} = \mathbf{1}_{j = \omega (i)}$  for  $\omega \in S_d$ . We use standard directed graph terminology, e.g. edges, parents. When we use the term "non-linear mixing" in this work, we also include linear mixing as a special case.
3 Setting: Interventional Causal Representation Learning
We now formally introduce our main settings of interest and the required assumptions. We assume that there is a latent variable distribution  $Z$  on  $\mathbb{R}^d$  and an observational distribution  $X$  on  $\mathbb{R}^{d'}$  given by  $X = f(Z)$  where  $d \leq d'$  and  $f$  is a non-linear mixing. This encapsulates the most general definition of a latent variable model. As the terminology suggests, we observe  $X$  in real-life, e.g. images of cats and  $Z$  encodes high-level latent information we wish to learn and model, e.g.  $Z$  could indicate orientation and size.
Assumption 1 (Nonlinear  $f$ ). The non-linear mixing  $f$  is injective, differentiable and embeds into  $\mathbb{R}^{d'}$ .
Such an assumption is standard in the literature. Injectivity is needed in order to identify  $Z$  from  $X$  because otherwise we may have learning ambiguity if multiple  $Zs$  map to the same  $X$ . Differentiability is needed to transfer densities. Note that Assumption 1 allows for a large class of complicated nonlinearities and we discuss in Appendix E why proving results for such a large class of mixing functions is important.
Next, we assume that the latent variables  $Z$  encode causal structure which can be expressed through a Structural Causal Model (SCM) on a Directed Acyclic Graph (DAG). We focus on linear SCMs with an underlying causal graph  $G$  on vertex set  $[d]$ , encoded by a matrix  $A$  through its non-zero entries, i.e., there is an edge  $i \rightarrow j$  in  $G$  iff  $A_{ji} \neq 0$ .
Assumption 2 (Linear Gaussian SCM on Latent Variables). The latent variables  $Z$  follow a linear SCM with Gaussian noise, so
$$
Z = A Z + D ^ {1 / 2} \epsilon \tag {1}
$$
where  $D$  is a diagonal matrix with positive entries,  $A$  encodes a DAG  $G$  and  $\epsilon \sim N(0,\mathrm{Id})$ .
For an illustration of the model, see Fig. 1a. It is convenient to encode the coefficients of linear SCMs through the matrix  $B = D^{-1/2}(\mathrm{Id} - A)$ . Then  $Z = B^{-1}\epsilon$  and both  $D$  and  $A$  can be recovered from  $B$ , indeed the diagonal entries of  $B$  agree with the entries of  $D^{-1/2}$ . Note that we can and will always assume that the entries of  $D$  are positive since  $\epsilon \stackrel{\mathcal{D}}{=} -\epsilon$ .
Assuming that the latent variables follow a linear Gaussian SCM is certainly restrictive but nevertheless a reasonable approximation of the underlying data generating process in many settings. The Gaussian prior assumption is also standard in latent variable modeling and enables efficient inference for downstream tasks, among other advantages. Importantly, under the above assumptions, our model has infinite modeling capacity [53, 86, 33], so they're able to model complex datasets such as images and there's no loss in representational power.
The goal of representation learning is to learn the non-linear mixing  $f$  and the high-level latent distribution  $Z$  from observational data  $X$ . In causal representation learning, we wish to go a step further and model  $Z$  as well by learning the parameters  $B$  which then lets us easily recover  $A$ ,  $D$  and the causal graph  $G$ . For general non-linear mixing, this model is not identifiable and hence we cannot hope to recover  $Z$ , but in this work we use additional interventional data, similar to the setups in recent works [2, 84, 88].
Assumption 3 (Single node interventions). We consider interventional distributions  $Z^{(i)}$  for  $i \in I$  which are single node interventions of the latent distribution  $Z$ , i.e., for each intervention  $i$  there is a target node  $t_i$  and we change the assigning equation of  $Z_{t_i}^{(i)}$  to
$$
Z _ {t _ {i}} ^ {(i)} = \left(A ^ {(i)} Z ^ {(i)}\right) _ {t _ {i}} + \left(D ^ {(i)}\right) _ {t _ {i}, t _ {i}} ^ {1 / 2} \left(\epsilon_ {t _ {i}} + \eta^ {(i)}\right) \tag {2}
$$
while leaving all other equations unchanged. We assume that  $A_{t_i,k}^{(i)} \neq 0$  only if  $k \to t_i$  in  $G$ , i.e., no parents are added and  $\eta^{(i)}$  denotes a shift of the mean.
For an illustration, see Fig. 1b. An intervention on node  $t_i$  has no effect on the non-descendants of  $t_i$ , but will have a downstream effect on the descendants of  $t_i$ . In particular,  $A$  is modified so that the weight of the edge  $k \rightarrow t_i$  could potentially be changed if it exists already but no new incoming edges to  $t_i$  may be added. Also, the noise variable  $\epsilon_{t_i}$  is also allowed to be modified via a scale intervention as well as a shift intervention. Note that [84, 2] do not allow shift interventions, whereas we do.
Again, this can be written concisely as  $Z^{(i)} = (B^{(i)})^{-1}(\epsilon + \eta^{(i)}e_{t_i})$  where  $B^{(i)} = (D^{(i)})^{-1 / 2}(\mathrm{Id} - A^{(i)})$ . Let  $r^{(i)}$  denote the row  $t_i$  of  $B^{(i)}$ , then we can write  $B^{(i)} = B - e_{t_i}(B^\top e_{t_i} - r^{(i)})^\top$ . We assume that interventions are non-trivial, i.e.,  $Z^{(i)} \neq Z$ . In our model, we observe interventional distributions  $X^{(i)} = f(Z^{(i)})$  for various interventions  $i \in I$ .
Definition 1 (Intervention Types). We call an intervention perfect (or stochastic hard) if  $A_{t_i}^{(i)} = 0$ , i.e., we remove all connections to former parents and potentially change variance and mean of the noise variable. Note that for nodes without parents, either  $D_{t_i t_i}^{(i)} \neq D_{t_i t_i}$  or  $\eta^{(i)} \neq 0$  so that the intervention is non-trivial. We call an intervention a pure noise intervention if  $A_{t_i}^{(i)} = A_{t_i}$  and  $A_{t_i} \neq 0$  (i.e., node  $t_i$  has at least one parent). In other words, a pure noise intervention targets a
node with parents by changing only the noise distribution through a change of variance (encoded in  $D$ ) or a mean shift (encoded in  $\eta$ ).
We call a single-node intervention imperfect if it is not perfect (some prior works call it a soft intervention). Note that perfect interventions are never of pure noise type because they necessarily change the relation to the parents. For an illustration of perfect and imperfect interventions, see Fig. 1c, Fig. 1b respectively. Finally, we require an exhaustive set of interventions.
Assumption 4 ((Coverage of Interventions)). All nodes are intervened upon by at least one intervention, i.e.,  $\{t_i : i \in I\} = [d]$
This assumption was also made in the prior works [84, 2, 88]. When the interventions don't cover all the nodes, we have non-identifiability as described in Section 4 and Appendix D. We also extensively discuss that none of our other assumptions can simply be dropped in these sections.
Remark 1. Our theory also readily extends to noisy observations, i.e.,  $X = f(Z) + \nu$  where  $\nu$  is independent noise. In this case, we first denoise via a standard deconvolution argument [36, 42] and then apply our theory.
- Similar to the results in [84] we can assume completely unknown intervention targets, i.e., there might be multiple interventions targeting the same node, and we do not need to know the partition. We only require coverage of all nodes. In contrast to their work we assume that we know which dataset corresponds to the observational distribution, but we expect that this restriction can be removed.
To simplify the notation, it is convenient to use  $B^{(0)} = B$ ,  $\eta^{(0)} = 0$  for the observational distribution. We also define  $\bar{I} = I \cup \{0\}$ . Then, all information about the latent variable distributions  $Z^{(i)}$  and observed distributions  $X^{(i)}$  are contained in  $((B^{(i)}, \eta^{(i)}, t_i)_{i \in \bar{I}}, f)$ .
We can now state our main results for the setting introduced above.
Theorem 1. Suppose we are given distributions  $X^{(i)}$  generated using a model  $((B^{(i)},\eta^{(i)},t_i)_{i\in \bar{I}},f)$  such that Assumptions 1-4 hold and such that all interventions  $i$  are perfect. Then the model is identifiable up to permutation and scaling, i.e., for any model  $((\widetilde{B}^{(i)},\widetilde{\eta}^{(i)},\widetilde{t}_i)_{i\in \bar{I}},\widetilde{f})$  that generates the same data, the latent dimension  $d$  agrees and there is a permutation  $\omega \in S_d$  (and associated permutation matrix  $P_{\omega}$ ) and an invertible pointwise scaling matrix  $\Lambda \in \mathrm{Diag}(d)$  such that
$$
\widetilde {t} _ {i} = \omega (t _ {i}), \quad \widetilde {B} ^ {(i)} = P _ {\omega} ^ {\top} B ^ {(i)} \Lambda^ {- 1} P _ {\omega}, \quad \widetilde {f} = f \circ \Lambda^ {- 1} P _ {\omega}, \quad \widetilde {\eta} ^ {(i)} = \eta^ {(i)}. \tag {3}
$$
This in particular implies that
$$
\widetilde {Z} ^ {(i)} \stackrel {\mathcal {D}} {=} P _ {\omega} ^ {\top} \Lambda Z ^ {(i)} \tag {4}
$$
and we can identify the causal graph  $G$  up to permutation of the labels.
This result says that for the interventional model as described in the previous section, we can identify the non-linear map  $f$ , the intervention targets  $t_i$ , the parameter matrices  $B, D, A$  up to permutations  $P_{\omega}$  and diagonal scaling  $\Lambda$ . Moreover, we can recover the shifts  $\eta^{(i)}$  exactly and also the underlying causal graph  $G$  up to permutations.
Remark 2. Identifiability and recovery up to permutation and scaling is the best possible for our setting. This is because the latent variables  $Z$  are not actually observed, which means we cannot (and in fact don't need to) resolve permutation and scaling ambiguity without further information about  $Z$ . See [84, Proposition 1] for the short proof.
When we drop the assumption that the interventions are perfect, we can still obtain a weaker identifiability result. Define  $\prec_{G}$  to be the minimal partial order on  $[d]$  such that  $i\prec_{G}j$  if  $(i,j)$  is an edge in  $G$ , i.e.,  $i\prec_{G}j$  if and only if  $i$  is an ancestor of  $j$  in  $G$ . Note that any topological ordering of  $G$  is compatible with the partial order  $\prec_{G}$ . Then, our next result shows that under imperfect interventions, we can still recover the partial order  $\prec_{G}$ .
Theorem 2. Suppose we are given the distributions  $X^{(i)}$  generated using a model  $((B^{(i)},\eta^{(i)},t_i)_{i\in \bar{I}},f)$  with causal graph  $G$  such that the Assumptions 1-4 hold and none of the interventions is a pure noise intervention. Then for any other model  $((\widetilde{B}^{(i)},\widetilde{\eta}^{(i)},\widetilde{t}_i)_{i\in \bar{I}},\widetilde{f})$  with causal graph  $\widetilde{G}$  generating the same observations the latent dimension  $d$  agrees and there is a permutation  $\omega \in S_d$  such that  $\widetilde{t}_i = \omega (t_i)$  and  $i\prec_{\widetilde{G}}j$  iff  $\omega (i)\prec_{G}\omega (j)$ , i.e.,  $\prec_G$  can be identified up to a permutation of the labels.
Remark 3. We emphasize that neither the full graph nor the coefficients  $B^{(i)}$  or the latent variables  $Z_{i}$  are identifiable in this setting, even for linear mixing functions.
Theorem 1 and Theorem 2 generalize the main results of [84] which assume linear  $f$  while we allow for non-linear  $f$ . The key new ingredient of our work is the following theorem, which shows identifiability of  $f$  up to linear maps.
Theorem 3. Assume that  $X^{(i)}$  is generated according to a model  $((B^{(i)},\eta^{(i)},t_i)_{i\in \bar{I}},f)$  such that the Assumptions 1-4 hold. Then we have identifiability up to linear transformations, i.e., if  $((\widetilde{B}^{(i)},\widetilde{\eta}^{(i)},\widetilde{t_i}),\widetilde{f})$  generates the same observed distributions  $\widetilde{f} (\widetilde{Z}^{(i)})\stackrel {\mathcal{D}}{=}X^{(i)}$ , then their latent dimensions  $d$  agree and there is an invertible linear map  $T$  such that
$$
\widetilde {f} = f \circ T ^ {- 1}, \quad \widetilde {Z} ^ {(i)} = T Z ^ {(i)}. \tag {5}
$$
The proof of this theorem is deferred to Appendix A. In Appendix B we then review the results of [84] and show how their results can be extended to obtain our main results, Theorem 1 and Theorem 2. Now we provide some intuition for the proofs of the main theorems.