SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiaoyu Liu<sup>1</sup>, Jiaxin Yuan<sup>2</sup>, Bang An<sup>1</sup>, Yuancheng Xu<sup>2</sup>, Yifan Yang<sup>1</sup>, Furong Huang<sup>1</sup>
$^{1}$  Department of Computer Science,  $^{2}$  Department of Mathematics
{xliu1231, jyuan98, bangan, ycxu, yang7832, furongh}@umd.edu
Representation learning assumes that real-world data is generated by a few semantically meaningful generative factors (i.e., sources of variation) and aims to discover them in the latent space. These factors are expected to be causally disentangled, meaning that distinct factors are encoded into separate latent variables, and changes in one factor will not affect the values of the others. Compared to statistical independence, causal disentanglement allows more controllable data generation, improved robustness, and better generalization. However, most existing works assume unconfoundedness (i.e., there are no common causes to the generative factors) in the discovery process, and thus obtain only statistical independence. In this paper, we recognize the importance of modeling confounders in discovering causal generative factors. Unfortunately, such factors are not identifiable without proper inductive bias. We fill the gap by introducing a framework named Confounded-Disentanglement (C-Disentanglement), the first framework that explicitly introduces the inductive bias of confounder via labels/knowledge from domain expertise. In addition, we accordingly propose an approach to sufficiently identify the causally-disentangled factors under any inductive bias of the confounder. We conduct extensive experiments on both synthetic and real-world datasets. Our method demonstrates competitive results compared to various SOTA baselines in obtaining causally disentangled features and downstream tasks under domain shifts. code available here
Causally disentangled representation learning methods endeavour to identify and manipulate the underlying explanatory causes of variation (i.e., generative factors) within observational data through obtaining causally disentangled representations [24, 8]. Such a representation encodes these factors separately in different random variables in the latent space, where changes in one variable do not causally affect the others. Pursuing causal independence<sup>1</sup> makes it possible to identify the ground truth generative factors that are not statistically independent, such as the color, shape and size in a fruit dataset as shown in Figure 1, which is more realistic and allows more controlled data generation, improved robustness, and better generalization
$p(color, shape, size) \neq p(color)p(shape)p(size)$
Figure 1: Causally independent generative factors (color, shape, size) may appear statistically correlated in observational data.
<sup>1</sup>We will use causally disentangled and causally independent interchangeably.
Despite great success, existing methods suffer from the identifiability issue in discovering semantically meaningful generative factors. The main reason is that most of them equate disentangling in the latent space with enforcing statistical independence [9, 11, 6] in the latent space, or requiring no mutual information[8]. In other words, they explicitly or implicitly assume that the observational dataset is unconfounded in the generative process, which is not realistic in practice. A confounder C is a common cause of multiple variables. It could create a correlation among causally disentangled generative factors in the observational distribution. In the previous example of a fruit dataset, the type of fruit is a confounder, creating a correlation between size and shape (apples are usually red and round objects), thus size and shape can never be captured by a learning objective that requires statistically independent latent representations. Assuming unconfoundedness in such a dataset leads to discrepancies in characterizing the statistical relationship of the generative factors, and makes them non-identifiable.
Considering the limitation of the unconfoundedness assumption, a few of the recent studies take into account confounders [24, 23, 21] in the problem formulation. Specifically, [23, 21] proposed evaluation metrics in the confounded causal generative process. [24] introduced an unsupervised learning scheme by regulating the learning process with the Independence-Of-Support Score (IOSS) under an unobserved confounder, following the theory that if a latent representation is causally disentangled, then it must have independent support. However, it has been shown almost impossible to obtain disentangled representations through purely unsupervised learning without proper inductive biases [10, 16]. While there are a few semi-supervised or weakly supervised methods [16, 17] developed seeking statistical independence, none of the above-mentioned methods considers how to provide inductive bias for confounder in learning causally disentangled representations.
In this paper, we recognize the importance of providing inductive bias to confounder so that the ground truth generative factors can be identified and we fill the gap by introducing the inductive bias via knowledge from domain expertise. Specifically, we propose a framework called ConfoundedDisentanglement (C-Disentanglement). C-Disentanglement is, to the best of our knowledge, the first framework that discusses the identifiability issue of generative factors regarding the inductive bias of confounder, and thus opens up the possibility to discover the ground truth causally disentangled generative factors which are correlated in the observational dataset. Importantly, C-Disentanglement is a general framework that includes existing methods (in which unconfoundedness is assumed) as special cases, allowing customized inductive bias of confounders according to diverse user needs.
Under the framework, We develop an algorithm to discover the causally disentangled generative factors in the latent space with inductive bias  $\mathbf{C}$ , where  $\mathbf{C}$  is a label set. Instead of enforcing global statistical independence among variables on the observational dataset, we partition the dataset into subsets according to realizations of  $\mathbf{C}$ , and regulate these variables within each subset.
Our proposed discovery methodology is general and can be applied to various disentangle mechanisms. But as a prototype for a proof of concept, we implement our method, cdVAE, under the context of VAE [12]. Specifically, we formulate a mixture of Gaussian model in which each Gaussian component, conditioned on a specific value of  $\mathbf{C}$ , represents the distribution of a latent variable that estimates the underlying generative factor. Intuitively, under the previous example of the fruit dataset, if we fix the type of fruit (e.g., the apple) and check the correlation between the color and the shape, we could find that the color change within the apples does not affect the shape distribution, thus can be captured by regulating the statistical independence conditioned on  $\mathbf{C}$  (a fixed fruit type).
Our methodology is a novel learning paradigm that can discover latent factors through weak supervision labels (i.e., C).
We experiment with several tasks from image reconstruction and generation, to classification under the out-of-distribution scenario on both real-world and synthetic datasets. We demonstrate that our method outperforms various baselines in discovering generative factors in the latent space.
Summary of contributions: (1) We recognize the identifiability issue of discovering generative factors in the latent space. We accordingly introduce a framework, named Confounded Disentanglement (C-Disentanglement). It is the first framework that discusses how inductive bias of confounder could be explicitly provided via labels/knowledge from domain expertise. (2) We propose an algorithm, cdVAE, that discovers causally disentangled generative factors in the latent space. The algorithm sheds light on the easy injection of inductive bias into existing methods.(3) We conduct extensive experiments and ablation studies across various datasets and tasks. Empirical results verify that
cdVAE outperforms existing methods in terms of inferring causally disentangled latent representation and also show cdVAE's superiority in downstream tasks under OOD generalization.
In this section, we introduce basic concepts in causal inference and then show how to evaluate the causal relationship among variables in latent space.
Causal graph through DAG. The causal relationship among variables can be reflected by a Directed Acyclic Graph (DAG). Each (potentially high-dimensional) variable is denoted as a node. The directed edges indicate the causal relationships and always point from parents to children.
Intervention and do-operator. Intervention is one of the fundamental concepts in causal inference. When we intervene on a variable, we set the value of a variable and cut all incoming causal arrows since its value is thereby determined only by the intervention [20]. The intervention is mathematically represented by the do-operator  $\text{do}(\cdot)$ . Let  $Z_{1}$  and  $Z_{2}$  be two variables,  $P(Z_{2}| \text{do}(Z_{1} = z_{1}))$  characterizes an interventional distribution and reflects how a change of  $Z_{1}$  affects the distribution of  $Z_{2}$ . The do-operation and the interventional distribution should be estimated on the interventional dataset. However, in practice, the true distribution of the data is unavailable but an observational subset. As a result, we estimate the interventional distribution from the observational set following do-calculus introduced by [19]. A detailed introduction can be found in the Appendix A.
Confounders bring in spurious correlation. Although the detailed definitions vary from literature, a confounder usually refers to a common cause (i.e., a common parent in the causal graph) of multiple variables and it brings in a certain level of correlation among these variables. Consequently, to estimate the causal effect of one variable on the other from the observational dataset, we have to eliminate the correlation introduced by the confounders. For example, when analyzing the causal relationship between the number of heat strokes and the rate of ice cream consumption, we may find that there is a correlation. However, the temperature is a confounder in the situation. If we eliminate this spurious correlation by conditioning on the temperature, we may find out that under the same temperature, heat stroke is not correlated to the rate of ice cream consumption.
Evaluation of interventional distribution. We specifically consider the case where there are variables  $Z_{1}$  and  $Z_{2}$ , and we analyze how the existence of parental nodes affects the estimation of  $P(Z_{2}|do(Z_{1} = z_{1}))$  on the observational distribution.
Proposition 2.1. Let  $Z_{1}$  and  $Z_{2}$  be two random variables,  $\mathbf{C}^{*}$  be the ground truth confounder set. If  $\mathbf{C}$  is a superset of or is equivalent to  $\mathbf{C}^{*}$ , i.e.,  $\mathbf{C}^{*} \subseteq \mathbf{C}$ , with  $c$  being a realization of  $\mathbf{C}$ , we have
$$
P \left(Z _ {2} \mid d o \left(Z _ {1}\right)\right) = \sum_ {c \in \mathbf {C}} P \left(Z _ {2} \mid Z _ {1}, \mathbf {C} = c\right) P (\mathbf {C} = c) \tag {1}
$$
if no  $C\in \mathbf{C}$  is a descendant of  $\mathbf{Z}$
The proof can be found in Appendix B.1.
Intuitively speaking, Proposition 2.1 states that to accurately estimate the causal relationship between variables, we have to eliminate the spurious correlation by conditioning on confounders. In addition, conditioning on additional variables will not affect the estimation if they are not decedents of  $\mathbf{Z}$ .
Causal disentanglement or statistical independence. If we had access to a series of ground-truth features that generated a fruit dataset(e.g., color, shape, texture), we can encode each of the features into a variable and concatenate them for a disentangled representation. However, these features in the observational dataset are usually correlated, depending on the type of fruits. If not considered from the causal perspective, traditional methods aiming for statistically independent representations, such as  $\beta$ -VAE [9], may not be able to consider correlated features that are disentangled. In fact, neither causal disentangled entails statistical independence nor vice versa because of the existence of confounders. We, therefore, formulate the task of inferring causally disentangled features from the observational under a certain level of confoundedness.
3 Confounded causal disentanglement via inductive bias
We formally frame the task of discovering a set of generative factors from the observational dataset from a causal perspective as shown in Figure 2. Let  $\mathbf{X}$  denote the observational data, the confounded causal generative process [23, 21] assumes that  $\mathbf{X}$  are generated from  $K$  ground-truth causes of variations  $G = [G_{1}, G_{2}, \dots, G_{K}]$  (i.e.,  $G \to X$ ) that do not cause each other. These generative factors are generally not available, and they are confounded
by some unobserved confounding variables  $\mathbf{C}^*$ . The task of interest is to discover those generative factors in the latent space (denoted as  $\mathbf{Z} = [Z_1, Z_2, \dots, Z_D]$ ) that best approximates  $\mathbf{G}$  from  $\mathbf{X}$ .
Figure 2: The left figure shows the ground truth generative process while the right figure demonstrates the learning task.
Causal disentanglement among latent generative factors. Generative factors, encoded in the latent representational space, are causally disentangled, meaning that intervening on the value of one factor does not affect the distribution of the others. It is formally defined as:
Definition 3.1 (Causal Disentanglement on Data  $\mathbf{X}([19,23,24])$ ). A representation is disentangled if, for  $i\in \{1,\dots ,D\}$
$$
P \left(Z _ {i} | \operatorname {d o} \left(Z _ {- i} = z _ {- i}\right), \mathbf {X}\right) = P \left(Z _ {i} | \mathbf {X}\right), \quad \forall z _ {i}. \tag {2}
$$
where  $-i = \{1,2,\dots,D\} /i$  indicates the set of all indices except for  $i$
Challenge of unobserved  $\mathbf{C}^*$ . The generative factors  $\mathbf{G}$  are not identifiable without proper inductive bias of  $\mathbf{C}^*$  [16]. Previous works on discovering the generative factors either obtain disentanglement by enforcing statistical independence on latent variables [9, 15, 6], or require that latent variables do not capture information of each other [8]. Such a setting is equivalent to assuming unconfoundedness of the generative factors. It ignores the possibility that correlated latent variables can also be causally disentangled in the observational distribution, and hence is an assumption too restrictive. Fortunately, even though the ground truth generative factors are unobserved, domain expertise may inform a "reasonable" or "likely" inductive bias of the confounder from an accessible label set, denoted as  $\mathbf{C}$ . This  $\mathbf{C}$  is used to account for all correlation among  $\mathbf{Z}$ .
Note here that we do not assume the accessible label set  $\mathbf{C}$  equals to  $\mathbf{C}^*$  (but hope that it is close to  $\mathbf{C}^*$ ). The relationship between  $\mathbf{C}^*$  and label set  $\mathbf{C}$  must fall into one of the following scenarios, despite the immeasurability of their exact relationships.
Case 1 The label set contains no information about the confounders, i.e.,  $\mathbf{C} = \emptyset$
Case 2 The label set contains partial information about the confounders, i.e.,  $\mathbf{C} \subset \mathbf{C}^{*}$ ,
Case 3 The label set contains all information about the confounders, i.e.,  $\mathbf{C} = \mathbf{C}^{*}$
One may argue that  $\mathbf{C}$  may contain information irrelevant to the ground truth confounders. We show in Appendix B that in the confounded generative process described in this paper, irrelevant information in  $\mathbf{C}$  does not affect the evaluation of the interventional distribution, and therefore can be ignored. We only take into consideration how much information in  $\mathbf{C}^*$  is captured here without loss of generality.
According to Proposition 2.1, in case 3, we can estimate  $P(Z_{i}|do(Z_{-i} = z_{-i},\mathbf{X})$  on the observational set with inductive bias from C. However, in the rest of the cases, the equation does not hold. Therefore, we introduce an operator  $do^{c}(\cdot)$  to estimate the interventional distribution on the observational set under inductive C. The framework that applies  $do^{c}(\cdot)$  is named C-Disentanglement, shorten for Confounded Disentanglement.
3.2 C-Disentanglement as the learning objective
We formally introduce C-Disentanglement and  $do^c$  as follows:
Definition 3.2 (C-Disentanglement and  $do^c$ ). Let  $\mathbf{X}$  be the observational data,  $\mathbf{Z} = [Z_1, Z_2, \dots, Z_D]$  be a concatenation of  $D$  random variables,  $\mathbf{C}$  be a label set selected from domain expertise to provide inductive bias for confounders of the observational data, we define  $do^c$  operation as:
$$
P \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i} = z _ {- i}\right), \mathbf {X}\right) = \sum_ {c \in \mathbf {C}} P \left(Z _ {i} \mid \mathbf {X}, Z _ {- i} = z _ {- i}, \mathbf {C} = c\right) P (\mathbf {C} = c) \tag {3}
$$
$\forall i \in 1,2,3,\ldots,D$ , where  $c$  is realizations of  $\mathbf{C}$ .  $\mathbf{Z}$  obtains C-Disentanglement on  $\mathbf{X}$  given  $\mathbf{C}$  if
$$
P \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i} = z _ {- i}\right), \mathbf {X}\right) = P \left(Z _ {i} \mid \mathbf {X}\right). \tag {4}
$$
$\mathbf{C}$  as an approximation of  $\mathbf{C}^*$ . Whether the ground truth generative factors can be discovered depends on whether their statistical relationship on the observational distribution is correctly characterized, and it requires knowledge of  $\mathbf{C}^*$ . However, due to the unobservable nature of the confounders, we use  $\mathbf{C}$  as an approximation of  $\mathbf{C}^*$ , as in Equation (4). Then the identifiability of the ground truth generative factors depends on the relationship between  $\mathbf{C}$  and  $\mathbf{C}^*$ . We thus examine three possibilities as listed in section 3.1 together with a case study as in Figure 3.
(a)  $\mathbf{C} = \emptyset$
(b)  $\mathbf{C}\subset \mathbf{C}^{*}$
(c)  $\mathbf{C} = \mathbf{C}^{*}$
Figure 3: Case study: a fruit dataset. Suppose a fruit dataset is generated by color, shape and size. These factors are correlated depending on the type of fruit and the place of origin. One way to estimate the causal effect within latent variables, proposed as C-Disentanglement, is to condition on realizations of the confounders  $\mathbf{C}$ , which partitions the data into subsets and inspects within each set. Figure (a)-(d) demonstrate partitions under different  $\mathbf{C}$ . In fig. 3a, when  $\mathbf{C} = \emptyset$ , color, shape and size cannot be identified, as the requirement of statistical independence contradicts the fact that they are correlated in the observational set. In fig. 3b, only partial information is given. It still helps as it can identify that size and shape are causally independent although entangled by color. In fig. 3c and fig. 3d, the ground truth generative factors can be discovered.
(d)  $\mathbf{C}\supset \mathbf{C}^*$
In Case 1,  $\mathbf{C}$  is an empty set as shown in fig. 3a. It is equivalent to assuming that the generative process is unconfounded and Equation (3) degrades to statistical independence on the observational set:
$$
P \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i} = z _ {- i}\right), \mathbf {X}\right) = P \left(Z _ {i} \mid Z _ {- i} = z _ {- i}, \mathbf {X}\right) = P \left(Z _ {i} \mid \mathbf {X}\right). \tag {5}
$$
This contradicts the fact that these factors are actually correlated and such an assumption makes the ground truth generative factors unidentifiable when there exists a underlying confounder.
In Case 3, as shown in fig. 3c and fig. 3d,  $\mathbf{C} \supseteq \mathbf{C}^*$ ,  $do$  is equivalent to  $do^c$ . Thus the ground truth generative factors can be identified.
In Case 2, partial information about the confounders is given as shown in fig. 3b. Heuristically speaking, this relaxes the constraint of statistical independence on the observational set and therefore enlarges the solution set, providing an increased possibility for the ground truth factors to be recovered. We empirically show in Section 5 that even partial information of confounders improves accuracy, outperforming existing methods in various tasks, even under distribution shifts.
4 cdVAE: identify causally disentangled factors in the latent space.
In this section, we provide an algorithm, cdVAE, for confounded disentangled VAE, to identify the latent causally disentangled generative factors under confounder  $\mathbf{C}$  in the context of VAE.
Figure 4: Learning paradigm of cdVAE. The input data  $\mathbf{X}$  is partitioned by realizations of the confounder  $\mathbf{C}$ . We infer from each partition a Gaussian distribution and form a mixture of Gaussian model to characterize the distribution of  $\mathbf{Z}$  on the observational distribution. We assume soft assignment of samples and further infer  $\pi^c (\mathbf{X})$  that resembles  $\mathbf{C}$ .
4.1 Learning objective
Given  $\mathbf{X}$  as the dataset, we hope to find a deterministic function  $f$ , parameterized by  $\theta$ , where  $f: \mathcal{Z} \to \mathcal{X}$  such that (1)  $P(X) = \int f(\mathbf{Z}; \theta) P(\mathbf{Z}) dz$  is maximized and (2) each  $Z_i$  encodes causally disentangled generative factors, with  $\mathbf{Z} = [Z_1, Z_2, \dots, Z_D]$  as a concatenation of random variables in the latent space.
More concretely,  $f(z; \theta)$  is characterized by a probability distribution  $P(\mathbf{Z}|\mathbf{X};\theta)$ , and  $p(\mathbf{Z})$  is a prior distribution from where  $\mathbf{Z}$  can be easily sampled. For each  $Z_{i} \in \mathbf{Z}$ , we require
$$
P \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i}\right), \mathbf {X}\right) = P \left(Z _ {i} \mid \mathbf {X}\right), \quad \forall c \in \mathbf {C}. \tag {6}
$$
Applying the variational Bayesian method [12], the learning object is to optimize the evidence lower bound (ELBO) while satisfying the constraints on causal disentanglement:
$$
\max  _ {\theta , \phi} \mathbb {E} _ {\mathbf {Z} \sim Q} [ \log P (\mathbf {X} | \mathbf {Z}; \theta) ] - D [ Q (\mathbf {Z} | \mathbf {X}; \phi) | | P (\mathbf {Z}) ] \tag {7}
$$
$$
s. t. \quad P \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i}\right), \mathbf {X}\right) - P \left(Z _ {i} \mid \mathbf {X}\right) = 0 \quad \forall i \in 1, 2, \dots , D \tag {8}
$$
For simplicity, we omit all model parameters  $\theta ,\phi$  in writing.
4.2 Learning strategy
We start from the estimation of the causal disentanglement constraint as shown in Equation (6). Because the probability distribution is hard to directly calculate, inspired by [23], we resort to its first-order moment as an approximation. Specifically, we estimate the  $L_{1}$  distance between the expectation of probability  $P(Z_{i}|do^{c}(Z_{-i}),\mathbf{X})$  and  $P(Z_{i}|\mathbf{X})$  as follows:
$$
l _ {c} = \sum_ {i = 1} ^ {D} d \left[ \mathbb {E} \left(Z _ {i} \mid d o ^ {c} \left(Z _ {- i}\right), \mathbf {X}\right), \quad \mathbb {E} \left(Z _ {i} \mid \mathbf {X}\right) \right]. \tag {9}
$$
We show in the following theorem that eq. (9) is satisfied if the learned latent variable  $\mathbf{Z}$  subjects to a Gaussian distribution with a diagonal covariance matrix. Proof can be found in appendix B.2.
Theorem 4.1. Suppose that the latent variable  $\mathbf{Z}$  on dataset  $\mathbf{X}$  given  $\mathbf{C} = c$  subjects to the Gaussian Distribution  $\mathcal{N}(\mu^c (\mathbf{X}),\Sigma^c (\mathbf{X}))$ . Specifically,
$$
P (\mathbf {Z} | \mathbf {C} = c, \mathbf {X}) = (2 \pi) ^ {- D / 2} \det (\Sigma^ {c}) ^ {- 1 / 2} \exp \left(- \frac {1}{2} (\mathbf {Z} - \mu^ {c}) ^ {\mathsf {T}} (\Sigma^ {c}) ^ {- 1} (\mathbf {Z} - \mu^ {c})\right),
$$
where  $\mathbf{Z} \in \mathbb{R}^D$ . If  $\Sigma^c(\mathbf{X})$  is diagonal for all  $c$ , we have