Causal Representation Learning and Inference for Generalizable Cross-Domain Predictions

Factorized variational distribution $q$: While we acknowledge that the true distribution $p(Z_c, Z_s|X)$ cannot be factored under our SCM assumption, it often becomes intractable due to the conditions that only $X$ and $Y$ are observed in the SCM. The distribution $p(Z_c, Z_s|X)$ is likely to be complex and non-Gaussian. As an alternative solution, we approximate $p(Z_c, Z_s|X)$ by a variational distribution $q(Z_c, Z_s|X)$ with Gaussian assumptions. Although using non-factorized Gaussians might slightly enhance results, it requires more time to estimate the large covariance matrix. There is always a balance between accuracy and efficiency. We believe our assumption, commonly used in many VAEs, is reasonable and effective. Moreover, our empirical results indicate that the factorized Gaussian approximation effectively leads to better OOD generalization. We appreciate your suggestion and will include an analysis of using non-factorized Gaussians in our revised paper.

Identifiability: We emphasize that our proof of identifiability, leveraging the results of [1], is mathematically correct. However, it is important to note that this proof is not our main contribution. Our focus is on designing conditional distributions in our SCM to ensure that the learned latent representations are identifiable. While Kivva et al. have proven identifiability in a variety of cases, our work necessitates only a specific instance to effectively learn our SCM. We model $p(Z_c)$ as a standard Gaussian instead of a mixture of Gaussian, chosen for its simplicity and because we lack additional prior information about $Z_c$. This assumption aligns with the prior distributions commonly used in many VAEs.

[1] Kivva, Bohdan, et al. "Identifiability of deep generative models without auxiliary information." Advances in Neural Information Processing Systems 35 (2022): 15687-15701.

Invariance of $p(Y|X, do(Z_s))$: To address the reviewer’s concerns, we will first emphasize and rectify the assumptions we make regarding the latent variables $U_x, U_{xy}$. Then we will explain that under these assumptions, $p(Y|X, do(Z_s))$ is invariant and transportable across domains. Finally, we will explain how to obtain $Z_s$ values for calculating $p(Y|X, do(Z_s))$ and why it is a reasonable and good choice we have to approximate an invariant transformation between $Z_s$ and $X$.

• Assumptions: Upon careful review, we acknowledge that our assertion regarding the variation of $p(U_{xy})$ was overstated. It is imperative to make the confounding effects between the source domain and target domain consistent to make the proposed interventional distribution invariant and transportable. Hence we correct our initial assumption that the confounding effects, encompassing $p(U_{xy})$ and the causal mechanisms between $U_{xy}$, $Z_s$, and $Y$, remain consistent across training and test domains. Moreover, we assume the distribution of $p(U_x)$ varies across domains. It results in the variation of $p(Z_s)$, and further the variation of $p(X)$. However, we emphasize that the generative mechanism $p(X|Z_s, Z_c)$ remain invariant across domains ($X$ is independent of $U_x, U_{xy}$ given $Z_c, Z_c$). Otherwise, it is impossible to infer $Z$ from any unseen domain.

• Invariance of $p(Y|X, do(Z_s))$: Our proposed interventional distribution, by setting specific values for $Z_s$ to mitigate the influence of the domain-specific variable $U_x$ on $Z_s$, effectively addresses the invariant confounding effects and prevent the $U_x$ from influence $Y$. Hence, it is invariant and transportable across domains.

• The choice of $Z_s$: Under our corrected assumption, the latent confounding effect remains invariant across domains. To infer the label for a test input $x^t$ using our interventional distribution, it is imperative to provide the true value of $Z_s$ for $x^t$ to accurately account for the confounding effects between $x^t$ and $y^t$. The distribution $p(Z_s|X)$ utilized in Eq. (3) can be construed as a proxy distribution that enables us to derive the true values of $Z_s$ from an input $X$. The challenge lies in determining which learned distribution can be employed to approximate this proxy distribution. According to the SCM, we observe that $X$ is independent of $U$ given $Z=[Z_c, Z_s]$, rendering $p(X|Z)$ invariant and transportable. Therefore, a reasonable and good choice is to obtain the identifiable $z_s^t$ from the learned variational distribution $q(Z_s|X=x^t)$ with a high $p(X=x^t|Z_s=z^t_s)$. However, as $q(Z_s|X=x^t)$ represents a variational estimation of the desired proxy distribution, we average the interventional distribution over multiple samples of $z_s$ from $q(Z_s|X=x^t)$. In practice, randomly obtaining a $z_s$ value is likely to yield a low $p(X|Z)$ and contribute minimally to the calculation of the interventional distribution.

Nevertheless, it's crucial to underscore that with the corrected assumption, our CIIRL remains innovative and has proven its efficacy across various benchmark distribution shift datasets: 1) With the revised assumptions, our proposed interventional distribution maintains invariance and transportability, facilitating cross-domain inference. 2) Our training procedure does not explicitly rely on the variation of $p(U_{xy})$. The distinct classes of $U$ can be associated with the diverse domain variable $U_x$. Our training procedure can still effectively distinguish the two types of representations. 3) Empirical results affirm the existence of invariant latent confounders, as evidenced by the overall superior OOD prediction performance of CIIRL compared to predictions solely based on $Z_c$.

The SCM assumption: We would like to emphasize that the causal mechanisms in the proposed SCM in Figure 1 are all assumptions that are widely adopted in the area of causal representation learning [1, 2], including the three following points:

• 1. The latent high-level factors $Z$ can be separated into causal factors $Z_c$ and spurious factors $Z_s$.
• 2. The input $X$ is generated by the the high-level factors $Z$.
• 3. Causal factor $Z_c$ is either direct cause or effect of target $Y$.

We cannot prove these assumptions always hold in real-world data and we admit the effectiveness and soundness of our derived theorem and algorithm is built upon these assumptions. We believe the learned representation $Z_s$ and $Z_c$ are the factors that satisfy the causal graph by

• 1. parameterizing the joint distribution regarding all the variables of interest into conditional distributions adhering to the causal mechanisms in Figure 1;
• 2. establishing the identifiability of the learned $Z_s, Z_c$.

We appreciate the suggestion of further elaboration of these two types of biases and will revise the paper accordingly.

We acknowledge the recommendation to underscore the importance of assumptions and will incorporate this emphasis into our paper accordingly. It's worth noting that our method is not confined solely to image data. The validity of our theorem and the effectiveness of the algorithm persist as long as the assumptions we posit apply to the given data. As an illustration, our approach can be extended to encompass text data, demonstrating the versatility of our framework.

Invariance of $p(Y|X, do(Z_s))$: Upon careful review, we acknowledge that our assertion regarding the variation of $p(U_{xy})$ was overstated. Our proposed interventional distribution, by setting specific values for $Z_s$ to mitigate the influence of the domain-specific variable $U_x$ on $Z_s$, effectively accounts for the invariant confounding effects and prevents the $U_x$ from influencing $Y$. Hence, it is invariant and transportable across domains. However, it is important to note that this framework may not generalize to new domains with different and unknown confounding effects.

In light of this, we correct our initial assumption that the confounding effects, encompassing $p(U_{xy})$ and the causal mechanisms between $U_{xy}$, $Z_s$, and $Y$, remain consistent across training and test domains. This invariant confounding assumption aligns with standard practices widely adopted in works utilizing interventional distributions [1]. We will incorporate such a revision into the paper.

Nevertheless, it's crucial to underscore that with the corrected assumption, our CIIRL remains innovative and has proven its efficacy across various benchmark distribution shift datasets:

• 1. With the revised assumptions, our proposed interventional distribution maintains invariance and transportability, facilitating cross-domain inference.
• 2. Our training procedure does not explicitly rely on the variation of $p(U_{xy})$. The distinct classes of $U$ can be associated with the diverse domain variable $U_x$. Our training procedure can still effectively distinguish the two types of representations.
• 3. Empirical results affirm the existence of invariant latent confounders, as evidenced by the overall superior OOD prediction performance of CIIRL compared to predictions solely based on $Z_c$.

[1] Mao, Chengzhi, et al. "Causal transportability for visual recognition." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

The choice of $Z_s$: Under our corrected assumption, the latent confounding effect remains invariant across domains. To infer the label for a test input $x^t$ using our interventional distribution, it is imperative to provide the true value of $Z_s$ for $x^t$ to accurately account for the confounding effects between $x^t$ and $y^t$. The distribution $p(z_s|x)$ utilized in Eq. (3) can be construed as a proxy distribution that enables us to derive the true values of $Z_s$ from an input $X$. The challenge lies in determining which learned distribution can be employed to approximate this proxy distribution. According to the SCM, we observe that $X$ is independent of $U$ given $Z=[Z_c, Z_s]$, rendering $p(X|Z)$ invariant and transportable. Therefore, a reasonable and good choice is to obtain the identifiable $z^t_s$ from the learned variational distribution $q(Z_s|X=x^t)$ with a high $p(X=x^t|Z_s=z^t_s)$. However, as $q(Z_s|X=x^t)$ represents a variational estimation of the desired proxy distribution, we average the interventional distribution over multiple samples of $z_s$ from $q(Z_s|X=x^t)$. In practice, randomly obtaining a $z_s$ value is likely to yield a low $p(X|Z)$ and contribute minimally to the calculation of the interventional distribution.

The modeling of $p(U_x)$ and $p(U_{xy})$: Firstly, we revise our assumption by rectifying the notion that the distribution of $p(U_{xy})$ remains invariant across domains. In our context, $U_x$ denotes any information specific to the domain, with a common simplification being to assume that $U_x$ represents the domain index [2]. During the training procedure, we utilize a clustering algorithm to estimate the domain index for each training input. The objective of the training process is to acquire the encoder distributions $q(Z_s|X), q(Z_c|X)$ that produce a disentangled representation, which is identifiable and possesses an invariant generative mechanism $p(X|Z)$. The interventional distribution accommodates confounding effects and mitigates the influence of the domain variable $U_x$. Consequently, we can make inferences from the interventional distribution without knowledge of $U_{xy}$ and $U_x$ for the target domain.

[2] Lu, Chaochao, et al. "Invariant causal representation learning for out-of-distribution generalization." International Conference on Learning Representations. 2021.

Assumptions of $U_{x}, U_{xy}$: The assumptions we make about $U_x$ and $U_{xy}$ posit them as random variables representing domain-specific and confounding information, respectively. We adhere to the standard practice of simplifying $U_x$ into a domain index. Importantly, we do not constrain $U_x$ and $U_{xy}$ to be exclusively discrete or continuous. During the training process, a neural network is employed to take in $U_x$ and $U_{xy}$ and generate parameters, including the mean and variances, for the prior distributions of $Z_s$. The neural network accommodates inputs of any type. However, the training procedure disentangles $Z_s$ and $Z_c$ by employing asymmetric prior distributions. As the dimensionality $J=|U|$ increases, optimization becomes more challenging due to potential inaccuracies in estimating domain indices, a growing number of parameters in the prior distribution, and limited improvements (if any) in disentanglement.

Invariance of $p(Y|X, do(Z_s))$: Upon careful review, we acknowledge that our assertion regarding the variation of $p(U_{xy})$ was overstated. Our proposed interventional distribution, by setting specific values for $Z_s$ to mitigate the influence of the domain-specific variable $U_x$ on $Z_s$, effectively accounts for the invariant confounding effects and prevents the $U_x$ from influencing $Y$. However, it is important to note that this framework may not generalize to new domains with different and unknown confounding effects.

In light of this, we correct our initial assumption that the confounding effects, encompassing $p(U_{xy})$ and the causal mechanisms between $U_{xy}$, $Z_s$, and $Y$, remain consistent across training and test domains. This invariant confounding assumption aligns with standard practices widely adopted in works utilizing interventional distributions [1]. We will incorporate such a revision into the paper.

Nevertheless, it's crucial to underscore that with the corrected assumption, our CIIRL remains innovative and has proven its efficacy across various benchmark distribution shift datasets:

• 1. With the revised assumptions, our proposed interventional distribution maintains invariance and transportability, facilitating cross-domain inference.
• 2. Our training procedure does not explicitly rely on the variation of $p(U_{xy})$. The distinct classes of $U$ can be associated with the diverse domain variable $U_x$. Our training procedure can still effectively distinguish the two types of representations.
• 3. Empirical results affirm the existence of invariant latent confounders, as evidenced by the overall superior OOD prediction performance of CIIRL compared to predictions solely based on $Z_c$.

[1] Mao, Chengzhi, et al. "Causal transportability for visual recognition." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

Identifiability: We emphasize that our proof of identifiability, leveraging the results of [2], is mathematically correct. However, it is important to note that this proof is not our main contribution. However, we appreciate the suggestion of the reviewer and will consider dedicating a formal statement with rigorous proof for identifiability.

[2] Kivva, Bohdan, et al. "Identifiability of deep generative models without auxiliary information." Advances in Neural Information Processing Systems 35 (2022): 15687-15701.

Latent $U_x$ and $U_{xy}$: Our approach is specifically designed to enhance out-of-distribution (OOD) prediction in situations where the domain variable $U_x$ and confounder $U_{xy}$ are not known. While domain indices are provided for specific tasks/datasets like PACS and VLCS, acquiring them for general real-world tasks poses significant challenges.

Invariance of $p(Y|X, do(Z_s))$: Upon careful review, we acknowledge that our assertion regarding the variation of $p(U_{xy})$ was overstated. Our proposed interventional distribution, by setting specific values for $Z_s$ to mitigate the influence of the domain-specific variable $U_x$ on $Z_s$, effectively accounts for the invariant confounding effects and prevents the $U_x$ from influencing $Y$. Hence, it is invariant and transportable across domains. However, it is important to note that this framework may not generalize to new domains with different and unknown confounding effects.

In light of this, we correct our initial assumption that the confounding effects, encompassing $p(U_{xy})$ and the causal mechanisms between $U_{xy}$, $Z_s$, and $Y$, remain consistent across training and test domains. This invariant confounding assumption aligns with standard practices widely adopted in works utilizing interventional distributions [1]. We will incorporate such a revision into the paper.

Nevertheless, it's crucial to underscore that with the corrected assumption, our CIIRL remains innovative and has proven its efficacy across various benchmark distribution shift datasets: 1) With the revised assumptions, our proposed interventional distribution maintains invariance and transportability, facilitating cross-domain inference. 2) Our training procedure does not explicitly rely on the variation of $p(U_{xy})$. The distinct classes of $U$ can be associated with the diverse domain variable $U_x$. Our training procedure can still effectively distinguish the two types of representations. 3) Empirical results affirm the existence of invariant latent confounders, as evidenced by the overall superior OOD prediction performance of CIIRL compared to predictions solely based on $Z_c$.

[1] Mao, Chengzhi, et al. "Causal transportability for visual recognition." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

Comparison to [2]: We appreciate that the reviewer recommends this relevant paper. The casual graph for the data generation process closely resembles our SCM. Specifically, our causal graph shares similarities with the one presented in [2] in the following aspects:

• 1. we both separate the latent representation into (invariant) causal representation $Z_c$ and (variant) spurious representation $Z_s$. The input $X$ is generated by both types of representation.
• 2. The spurious representation varies from domain to domain since it is controlled by a domain-specific variable $U_x$.
• 3. There exists a high-level invariance confounder between $Z_s$ and $Y$ ($\tilde{Z}_s$ in their graph).

However, the graph in [2] assumes that causal features as the parent variables to target while we use the child variables.

Algorithmically, both our method and iMSDA from [2] address the confounding issue. iMSDA utilizes domain index information to estimate the confounder, whereas we focus on constructing the interventional distribution. Furthermore, iMSDA is tailored for domain adaptation tasks and necessitates access to test domain data during training. Given these distinctions, a direct comparison between our method and iMSDA would not be equitable.

[2] Partial disentanglement for domain adaptation Lingjing Kong, Shaoan Xie, Weiran Yao, Yujia Zheng, Guangyi Chen, Petar Stojanov, Victor Akinwande, Kun Zhang Proceedings of the 39th International Conference on Machine Learning, PMLR 162:11455-11472, 2022.

Identifiability of representation $Z_c, Z_s$: We strongly agree that establishing the identifiability of the latent variables $Z$ is crucial for performing do-calculus on them. Therefore, we leverage the theoretical results presented by Kivva et al. (2022) to demonstrate the identifiability of the $Z$ obtained through our learning framework. Please refer to the paragraph below the algorithm in Section 3.2 for more details.

[3] Kivva, Bohdan, et al. "Identifiability of deep generative models without auxiliary information." Advances in Neural Information Processing Systems 35 (2022): 15687-15701.