Interventional Causal Representation Learning

Modern deep learning models like GPT-3 (Brown et al., 2020) and CLIP (Radford et al., 2021) are remarkable representation learners (Bengio et al., 2013). Despite the successes, these models continue to be far from the human ability to adapt to new situations (distribution shifts) or carry out new tasks (Geirhos et al., 2020; Bommasani et al., 2021; Yamada et al., 2022). Humans encapsulate their causal knowledge of the world in a highly reusable and recomposable way (Goyal & Bengio, 2020), enabling them to adapt to new tasks in an ever-distribution-shifting world. How can we empower modern deep learning models with this type of causal understanding? This question is central to the emerging field of causal representation learning (Schölkopf et al., 2021).

A core task in causal representation learning is provable representation identification, i.e., developing representation learning algorithms that can provably identify natural latent factors (e.g., location, shape and color of different objects in a scene). While provable representation identification is known to be impossible for arbitrary data-generating process (DGP) (Hyvärinen & Pajunen, 1999; Locatello et al., 2019), real data often exhibits additional structures. For example, Hyvarinen et al. (2019); Khemakhem et al. (2022) consider the conditional independence between the latents given auxiliary information; Lachapelle et al. (2022) leverage the sparsity of the causal connections among the latents; Locatello et al. (2020); Klindt et al. (2020); Ahuja et al. (2022a) rely on the sparse variation in the latents over time.

Most existing works rely on observational data and make assumptions on the dependency structure of the latents to achieve provable representation identification. However, in many applications, such as robotics and genomics, there is a wealth of interventional data available. For example, interventional data can be obtained from experiments such as genetic perturbations (Dixit et al., 2016) and electrical stimulations (Nejatbakhsh et al., 2021). Can interventional data help identify latent factors in causal representation learning? How can it help? We explore these questions in this work. The key observation is that interventional data often carries geometric signatures of the latent factors’ support (i.e., what values each latent can possibly take). Fig. 1 illustrates these geometric signatures: perfect interventions and many imperfect interventions can make the intervened latents’ support independent of their ancestors’ support. As we will show, these geometric signatures go a long way in facilitating provable representation identification in the absence of strong distributional assumptions.

Contributions. This work establishes representation identification guarantees without strong distributional assumptions on the latents in the following settings.

• •

  $do$ interventions. We first investigate scenarios where the true latent factors are mapped to high-dimensional observations through a finite-degree multivariate polynomial. When some latent dimension undergoes a hard $do$ intervention (Pearl, 2009), we are able to identify it up to shift and scaling. Even when the mapping is not a polynomial, approximate identification of the intervened latent is still achievable provided we have data from multiple $do$ interventional distributions on the same latent dimension.

• •

  Perfect & imperfect interventions. We achieve block affine identification under imperfect interventions (Peters et al., 2017) provided the support of the intervened latent is rendered independent of its ancestors under the intervention as shown in Figure 1c. This result covers all perfect interventions as a special case.

• •

  Observational data and independent support. The independence-of-support condition above can further facilitate representation identification with observational data. We show that, if the support of the latents are already independent in observational data, then these latents can be identified up to permutation, shift, and scaling, without the need of any interventional data. This result extends the classical identifiability results from linear independent component analysis (ICA) (Comon, 1994) to allow for dependent latent variables. They also provide theoretical justifications for recent proposals of performing unsupervised disentanglement through the independent support condition (Wang & Jordan, 2021; Roth et al., 2022).

We summarize our results in Table 1. Finally, we empirically demonstrate the practical utility of our theory. From data generation mechanisms ranging from polynomials to image generation from rendering engine (Shinners et al., 2011), we show that interventional data helps identification.

Also, the code repository can be accessed at: github.com/facebookresearch/CausalRepID.

Existing provable representation identification approaches often utilize structure in time-series data, as seen in initial works by Hyvarinen & Morioka (2016) and Hyvarinen & Morioka (2017). More recent studies have expanded on this approach, such as Hälvä & Hyvarinen (2020); Yao et al. (2021, 2022a, 2022b); Lippe et al. (2022b, a); Lachapelle et al. (2022). Other forms of weak supervision, such as data augmentations, can also be used in representation identification, as seen in works by Zimmermann et al. (2021); Von Kügelgen et al. (2021); Brehmer et al. (2022); Locatello et al. (2020); Ahuja et al. (2022a) that assume access to contrastive pairs of observations $(x,\tilde{x})$. A third approach, used in (Khemakhem et al., 2022, 2020), involves using high-dimensional observations (e.g., an image) and auxiliary information (e.g., label) to identify representations.

To understand the factual and counterfactual knowledge used by different works in representation identification, we can classify them according to Pearl’s ladder of causation (Bareinboim et al., 2022). In particular, our work operates with interventional data (level-two knowledge), while other studies leverage either observational data (level-one knowledge) or counterfactual data (level-three knowledge). Works such as Khemakhem et al. (2022, 2020); Ahuja et al. (2022b); Hyvarinen & Morioka (2016, 2017); Ahuja et al. (2021) use observational data and either make assumptions on the structure of the underlying causal graph of latents or rely on auxiliary information. In contrast, works like Brehmer et al. (2022) use counterfactual knowledge to achieve identification for general DAG structures; Lippe et al. (2022b, a); Ahuja et al. (2022a); Lachapelle et al. (2022) use pre- and post-intervention observations to achieve provable representation identification. These latter studies use instance-level temporal interventions that carry much more information than interventional distribution alone. To summarize, these works require more information than is available with level two data in Pearlian ladder of causation.

Finally, a concurrent work from Seigal et al. (2022) also studies identification of causal representations using interventional distributions. The authors focus on linear mixing of the latents and consider perfect interventions. In contrast, our results consider nonlinear mixing function and imperfect interventions.

Causal representation learning aims to identify latent variables from high-dimensional observations. Begin with a data-generating process where some high-dimensional observations $x\in\mathbb{R}^{n}$ are generated from some latent variables $z\in\mathbb{R}^{d}$. We consider the task of identifying latent $z$ assuming access to both observational and interventional datasets: the observational data is drawn from

where the latent $z$ is sampled from the distribution $\mathbb{P}_{Z}$ and $x$ is the observed data point rendered from the underlying latent $z$ via an injective decoder $g:\mathbb{R}^{d}\rightarrow\mathbb{R}^{n}$. The interventional data is drawn from a similar distribution except the latent $z$ is drawn from $\mathbb{P}^{(i)}_{Z}$, namely the distribution of $z$ under intervention on $z_{i}$:

We denote $\mathcal{Z}$ and $\mathcal{Z}^{(i)}$ as the support of $\mathbb{P}_{Z}$ and $\mathbb{P}_{Z}^{(i)}$ respectively (support is the set where the probability density is more than zero). The support of $x$ is thus $\mathcal{X}=g(\mathcal{Z})$ in observational data and $\mathcal{X}^{(i)}=g\big{(}\mathcal{Z}^{(i)}\big{)}$ in interventional data. The goal of causal representation learning is provable representation identification, i.e. to learn an encoder function, which takes in the observation $x$ as input and provably output its underlying true latent $z$. In practice, such an encoder is often learned via solving a reconstruction identity,

where $f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{d}$ and $h:\mathbb{R}^{d}\rightarrow\mathbb{R}^{n}$ are a pair of encoder and decoder, which need to jointly satisfy Eq. 3. The pair $(f,h)$ together is referred to as the autoencoder. Given the learned encoder $f$, the resulting representation is $\hat{z}\triangleq f(x)$, which holds the encoder’s estimate of the latents.

The reconstruction identity Eq. 3 is highly underspecified and cannot in general identify the latents. There exist many pairs of $(f,h)$ that jointly solve Eq. 3 but do not provide representations $\hat{z}\triangleq f(x)$ that coincide with the true latents $z$. For instance, applying an invertible map $b$ to any solution $(f,h)$ will result in another valid solution $b\circ f$, $h\circ b^{-1}$. In practical applications, however, the exact identification of the latents is not necessary. For example, we may not be concerned with the recovering the latent dimensions in the order they appear in $z$. Thus, in this work, we examine conditions of under which the true latents can be identified up to certain transformations, such as affine transformations and coordinate permutations.

We first establish an affine identification result, which serves as a stepping stone towards stronger identification guarantees in the next section. We begin with a few assumptions.

The assumption that the matrix $G\in\mathbb{R}^{n\times q}$ has a full column rank of $q$ guarantees that the decoder $g$ is injective; see Lemma A.1 in Appendix A.1 for a proof. This injectivity condition on $g$ is common in identifiable representation learning. Without injectivity, the problem of identification becomes ill-defined; multiple different latent $z$’s can give rise to the same observation $x$. We note that the full-column-rank condition for $G$ in LABEL:assm2:g_poly imposes an implicit constraint on the dimensionality $n$ of the data; it requires that the dimensionality $n$ is greater than the number of terms in the polynomial of degree $p$, namely $n\geq\sum_{r=0}^{p}{r+d-1\choose d-1}$. In the Appendix (Theorem A.5), we show that if our data is generated from sparse polynomials, i.e., $G$ is a sparse matrix, then $n$ is allowed to be much smaller.

Under Assumptions 4.1 and 4.2, we perform causal representation learning with two constraints: polynomial decoder and non-collapsing encoder.

We now show that solving the reconstruction identity with these constraints can provably identify the true latent $z$ up to affine transformations.

Theorem 4.4 drastically reduces the ambiguities in identifying latent $z$ from arbitrary invertible transformations to only invertible affine transformations. Moreover, Theorem 4.4 does not require any structural assumptions about the dependency between the latents. It only requires (i) a geometric assumption that the interior of the support is non-empty and (ii) the map $g$ is a finite-degree polynomial.

The proof of Theorem 4.4 is in Appendix A.1. The idea is to write the representation $\hat{z}=f(x)$ as $\hat{z}=f\circ g(z)=a(z)$ with $a\triangleq f\circ g$, leveraging the relationship $x=g(z)$ in Eq. 1. We then show the $a$ function must be an affine map. To give further intuition, we consider a toy example with one-dimensional latent $z$, three-dimensional observation $x$, and the true decoder $g$ and the learned decoder $h$ each being a degree-two polynomial. We first solve the reconstruction identity on all $x$, which gives $h(\hat{z})=g(z)$, and equivalently $H[1,\hat{z},\hat{z}^{2}]^{\top}=G[1,z,z^{2}]^{\top}$. This equality implies that both $\hat{z}$ and $\hat{z}^{2}$ must be at most degree-two polynomials of $z$. As a consequence, $\hat{z}$ must be a degree-one polynomial of $z$, which we next prove by contradiction. If $\hat{z}$ is a degree-two polynomial of $z$, then $\hat{z}^{2}$ is degree four; it contradicts the fact that $\hat{z}^{2}$ is at most degree two in $z$. Therefore, $\hat{z}$ must be a degree-one polynomial in $z$, i.e. a linear function of $z$.

In the previous section, we derived affine identification guarantees. Next, we strengthen these guarantees by leveraging geometric signals specific to many interventions.

In the previous section, we showed that interventions induce geometric structure (independence of supports) in the support of the latents that helps achieve strong identification guarantees. In this section, we consider a special case where such geometric structure is already present in the support of the latents in the observational data. Since we only work with observational data in this section, we set the interventional supports $\mathcal{Z}^{(i)}=\mathcal{X}^{(i)}=\emptyset$, where $\emptyset$ is the empty set. For each $j\in[d]$, define $\beta_{\sup}^{j}$ to be the supremum of the support of $z_{j}$, i.e., $\mathcal{Z}_{j}$. Similarly, for each $j\in[d]$, define $\beta_{\inf}^{j}$ to be the infimum of the set $\mathcal{Z}_{j}$.

Following previous sections, we state a constraint, where the learner leverages the geometric structure in the support in Assumption 6.1 to search for the autoencoder.

The proof of Theorem 6.3 is in Appendix A.4. Theorem 6.3 says that the independence between the latents’ support is sufficient to achieve identification up to permutation, shift, and scaling in observational data. Theorem 6.3 has important implications for the seminal works on linear ICA (Comon, 1994), considering the simple case of a linear $g$. Comon (1994) shows that, if the latent variables are independent and non-Gaussian, then the latent variables can be identified up to permutation and scaling. However, Theorem 6.3 states that, even if the latent variables are dependent, the latent variables can be identified up to permutation, shift and scaling, as long as they are bounded (hence non-Gaussian) and satisfy pairwise support independence.

Finally, Theorem 6.3 provides a first general theoretical justification for recent proposals of unsupervised disentanglement via the independent support condition (Wang & Jordan, 2021; Roth et al., 2022).

In this section, we describe practical algorithms to solve the constrained representation learning problems in § 5 and 6.

To perform constrained representation learning with $do$-intervention data, we proceed in two steps. In the first step, we carry out minimization of the reconstruction objective $f^{\dagger},h^{\dagger}=\operatorname*{arg\,min}_{f,h}\mathbb{E}\big{[}\|h\circ f(X)-X\|^{2}\big{]}$, where $h$ is the decoder, $f$ is the encoder and expectation is taken over observational data and interventional data. In the experiments, we restrict $h$ to be a polynomial and show that affine identification is achieved by the learned $f^{\dagger}$ as proved in Theorem 4.4.

In the second step, we learn a linear map to transform the learned representations and enforce LABEL:eqn:intv_cons. For each interventional distribution, $\mathbb{P}_{X}^{(i)}$, we learn a different linear map $\gamma_{i}$ that projects the representation such that it takes an arbitrary fixed value $z^{\dagger}_{i}$ on the support of $\mathbb{P}_{X}^{(i)}$. We write this objective as

Construct a matrix $\Gamma$ with different $\gamma_{i}^{\top}$ as the rows. The final output representation is $\Gamma f^{\dagger}(X)$. In the experiments, we show that this representation achieves permutation, shift and scaling identification as predicted by Theorem 5.3. A few remarks in order. i) $z_{i}^{\dagger}$ is arbitrary and learner does not know the true do intervention value, ii) for ease of exposition, Eq. 7 assumes the knowledge of index of intervened and can be easily relaxed by multiplying $\Gamma$ with a permutation matrix.

We next describe an algorithm that learns representations to enforce independence of support (leveraged in Theorem 5.8 and 6.3). To measure the (non)-independence of the latents’ support, we follow Wang & Jordan (2021); Roth et al. (2022) and measure the distance between the sets in terms of Hausdorff distance: the Hausdorff distance $\mathsf{HD}$ between the sets $\mathcal{S}_{1},\mathcal{S}_{2}$ is $\mathsf{HD}(\mathcal{S}_{1},\mathcal{S}_{2})=\sup_{z\in\mathcal{S}_{2}}\bigg{(}\inf_{z^{{}^{\prime}}\in\mathcal{S}_{1}}(\|z-z^{{}^{\prime}}\|)\bigg{)}$, where $\mathcal{S}_{1}\subseteq\mathcal{S}_{2}$.

To further enforce the independent support constraint, we again follow a two-step algorithm. The first step remains the same, i.e., we minimize the reconstruction objective. In the second step, we transform the learned representations ($f^{\dagger}(x)$) with an invertible map $\Gamma\in\mathbb{R}^{d\times d}$. The joint support obtained post transformation is a function of the parameters $\Gamma$ and is denoted as $\hat{\mathcal{Z}}(\Gamma)$. Following the notation introduced earlier, the joint support along dimensions $k,m$ is $\hat{\mathcal{Z}}_{k,m}(\Gamma)$ and the marginal support along $k$ is $\hat{\mathcal{Z}}_{k}(\Gamma)$. We translate the problem in 6.2 as follows. We find a $\Gamma$ to minimize

5.6 can be similarly translated.

In this section, we analyze how the practical implementation of the theory holds up in different settings ranging from data generated from polynomial decoders to images generated from PyGame rendering engine (Shinners et al., 2011). The code to reproduce the experiments can be found at https://github.com/facebookresearch/CausalRepID.

In this work, we lay down the theoretical foundations for learning causal representations in the presence of interventional data. We show that geometric signatures such as support independence that are induced under many interventions are useful for provable representation identification. Looking forward, we believe that exploring representation learning with real interventional data (Lopez et al., 2022; Liu et al., 2023) is a fruitful avenue for future work.

Yixin Wang acknowledges grant support from the National Science Foundation and the Office of Naval Research. Yoshua Bengio acknowledges the support from CIFAR and IBM. We thank Anirban Das for insightful feedback that helped us correctly state the precise $\zeta$-thickness conditions.