Causal Effect Estimation with Mixed Latent Confounders and Post-treatment Variables

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 04:48:12

Abstract: Causal inference from observational data has attracted considerable attention among researchers. One main obstacle to drawing valid causal conclusions is handling of confounders. As the direct measurement of confounders may not always be feasible, recent methods seek to address the confounding bias via proxy variables, i.e., variables postulated to be causally related to unobserved confounders. However, the selected proxies may scramble both latent confounders and latent post-treatment variables in practice, where existing methods risk biasing the estimation by unintentionally controlling for variables affected by the treatment. In this paper, we systematically investigate the bias of latent post-treatment variables, i.e., latent post-treatment bias, in causal effect estimation. We first derive the bias of existing covariate adjustment-based methods when selected proxies scramble both latent confounders and latent post-treatment variables, which we demonstrate can be arbitrarily bad. We then propose a novel Confounder-identifiable VAE (CiVAE) to address the bias. CiVAE is built upon a mild assumption that the prior of latent variables that generate the proxy belongs to a general exponential family with at least one invertible sufficient statistic in the factorized part. Based on this assumption, we show that latent confounders and latent post-treatment variables can be individually identified up to simple bijective transformations. We then prove that with individual identification, the intractable disentanglement problem of latent confounders and post-treatment variables can be transformed to a tractable conditional independence test problem. Finally, we prove that the true causal effects can be unbiasedly estimated with transformed confounders inferred by CiVAE. Experiments on both simulated and real-world datasets demonstrate that CiVAE is significantly more robust to latent post-treatment bias than existing methods. The codes are provided in the following anonymous URL: https://anonymous.4open.science/r/CiVAE-demo-E701/readme.md

Main paper: 
 introduction: 
 Causal inference, which seeks to draw conclusions about cause-and-effect relationships among variables of interest, has gained increasing prominence in various fields, such as social science, economics, and public health (Glass et al., 2013; Johansson et al., 2016; Prosperi et al., 2020) . Traditional methods rely on randomized control trials (RCT) to draw valid causal conclusions from experimentation (Cook et al., 2002) . Recently, more attention has been dedicated toward causal inference from observational datasets, which contain samples with passively observed past treatment, the associated outcome, and possibly features, and in which researchers have no control over the treatment assignment mechanism (Shalit et al., 2017; Shi et al., 2019; Wager & Athey, 2018) .One main obstacle to inferring causal relations from observational data is confounding bias, which occurs when past treatments were determined by variables that causally influence the outcome, i.e., confounders. In such cases, the difference in the average outcome between the treatment group and the non-treatment group cannot be attributed solely to the treatment, but may also be due to the systematic difference of samples in the two groups (Mickey & Greenland, 1989) . If the confounders can be observed, a simple strategy to address such a bias is to control them via methods such as covariate adjustment (Pocock et al., 2002) or propensity score re-weighting (Li et al., 2018) . However, confounders are not always measurable (Kuroki & Pearl, 2014) . Therefore, recent methods seek to adjust for the influence of confounders based on their noisy proxies, which are generally covariates researchers postulated to be conducive for the inference of confounders (Miao et al., 2018; Yao et al., 2018; Madras et al., 2019) . One exemplar work from this strain is the causal effect variational auto-encoder (CEVAE) (Louizos et al., 2017) (Fig. 1-(a) ), which has demonstrated that confounding bias can be mitigated by controlling latent variables inferred from proxies of confounders. Although proxy-of-confounder-based methods have achieved substantial progress, we argue that these algorithms may risk controlling latent post-treatment variables (i.e., variables causally affected by the treatment) scrambled in the proxy variables, where post-treatment bias may be unintentionally introduced in the estimated treatment effect. Here, we note that the negative effects of controlling post-treatment variables have been investigated in prior research (Acharya et al., 2016; Elwert & Winship, 2014; King & Zeng, 2006) . For example, Montgomery et al. (2018) found that more than 50% of the papers published in top journals of politics inadvertently control post-treatment variables in the experimental setting, although researchers have complete control over the treatment assignment mechanism and the covariates to control for. On this basis, we postulate that post-treatment bias could be even worse for proxy-based methods in the setting of observational study, as when treatments are passively recorded, it is difficult to determine which variables causally influence the treatment and which variables are influenced by it (as both confounders and post-treatment variables are correlated with the treatment and the outcome). In addition, the post-treatment variables can be latent, which may be scrambled into the observed covariates together with the latent confounders.Consider the following real-world example that researchers from the Company1 have encountered when estimating the average causal effects of switching a job from onsite to online mode to the statistics of the applicants (e.g., average age, gender/geographical diversity, etc.). In this case, the Company collected a dataset of two groups of online (i.e., the treatment group) and onsite jobs (i.e., the non-treatment group), where for each job, the statistics of the applicants (i.e., the average age) are calculated as the outcome. Clearly, the seniority of the job is a confounder between the treatment and the outcome, as less senior jobs (e.g., internships) are more likely to permit online work, and applicants for these jobs tend to be younger on average. The seniority of a job can be difficult to measure. Therefore, the required skills of the job, which the recruiter must provide when publishing a job Ad in the Company, can be used as the proxy of the confounder "seniority". However, a caveat is that, switching to an online working mode may also alter the required skills of a job, thereby affecting the qualification of the applicants (where these altered skills are mediators). Consequently, directly using the required skills as the proxy of the confounder "seniority" could unintentionally control latent mediators, which introduces post-treatment bias in the causal effect estimation results.Addressing the latent post-treatment bias faces multi-faceted challenges. First, there lacks a theoretical formulation of the bias when the selected proxies scramble latent post-treatment variables for proxy-of-confounder-based methods; the trade-off between deconfounding and introducing new post-treatment bias is not clear. In addition, it is difficult to distinguish confounders and post-treatment variables in the latent space. Existing covariate disentanglement-based methods, e.g., TEDVAE (Zhang et al., 2021) , mainly focus on an easier task of disentangling latent confounders with latent adjusters and instrumental variables. This can be achieved by using their different predictive abilities w.r.t. the treatment and outcome (see Fig. 1-(b) ). However, since latent confounders and posttreatment variables correlate with both the treatment and outcome, the two cannot be disentangled by these methods. One solution is to assume the proxy of latent post-treatment variables can be observed, from which post-treatment variables can be inferred and disentangled from the latent confounders. However, this assumption is too strong, as in the previous online/onsite job case, we can never know which skills are causally influenced by the work mode. Finally, even if latent confounders can be distinguished, since general latent variable models have no identifiability guarantee (Khemakhem et al., 2020) , it is unclear whether controlling the inferred latent variables, which may be arbitrary transformations of the true confounders, can provide unbiased estimations of the causal effects.To address the aforementioned challenges, we provide a systematic investigation of the latent posttreatment bias in causal inference. We first analyze the behavior of existing proxy-based causal inference methods when the selected proxies scramble both latent confounders and post-treatment variables, where we demonstrate that the estimated average causal effects can be arbitrarily biased.We then propose the Confounder-identifiable VAE (CiVAE) to address such biases. Specifically, we show that based on a mild assumption that the prior distribution of latent variables (i.e., the latent confounders and post-treatment variables) belongs to a general exponential family with at least one invertible sufficient statistic in the factorized part, latent confounders and latent post-treatment variables can be individually identified up to simple bijective transformations. In addition, based on the causal relations among confounders, mediators, and treatment, we further demonstrate that the inferred confounders (which are actually transformed proxies of the true confounders) could be properly distinguished from the inferred latent post-treatment variables with pair-wise conditional independence tests. Finally, we prove that the true causal effects can be unbiasedly estimated based on transformed confounders inferred by CiVAE. Experiments on both simulated and real-world datasets demonstrate that CiVAE shows more robustness to latent post-treatment bias than existing methods. 
problem formulation and analysis: 
 2.1 PROBLEM FORMULATION Throughout this paper, we assume the causal model in Fig. 1-(c ), where the dashed lines denote indeterminate causal mechanisms that might vary in different cases. We use a binary random variable T to denote the treatment, a random vector X ∈ R K X to denote the observed covariates, and a random scalar Y ∈ R to denote the outcome. Furthermore, observed covariates X are assumed to be generated fromK C independent latent confounders C ≜ [C 1 , C 2 ..., C K C ] and K M latent post-treatment variables M ≜ [M 1 , M 2 ..., M K M ]under the causal influence of treatment T . We use the random vector Z ≜ [C||M ] ∈ R K Z =K C +K M to denote all latent factors. Our aim is to estimate the average causal effects of treatment T on outcome Y with auxiliary confounder information in X, where the estimation should be devoid of both confounding bias and post-treatment bias.2.2 ANALYSIS OF LATENT POST-TREATMENT BIASTo achieve such a purpose, we first formally define the (conditional) average treatment effects (C/ATE) when covariates X scramble both latent confounders C and post-treatment variables M . We then define the post-treatment bias when covariates X are used directly as the proxy of confounders. To facilitate the analysis, we make the following assumption regarding the causal generative process. Assumption 1. (Noisy-Injectivity). We assume X = f (C, M ) + ϵ, where f is a deterministic function that combines latent confounders C and latent post-treatment variables M into observations X and ϵ is random noise. In addition, we assume that the function f is injective; beyond injectivity, f can be arbitrarily nonlinear. We use f † : X → [C||M ] to denote its left inverse. We use f † C : X → C and f † M : X → M to denote the mapping from X to C, M , respectively. Noisy-Injectivity is a common assumption made either explicitly or implicitly in most existing proxyof-confounder-based causal inference algorithms. For example, if both X and C are categorical, Pearl (2012) assumes that X has at least the same number of categories as C, whereas the effect restoration algorithm (Rothman et al., 2008) assumes that the matrix of p(C, X) to be full-rank. Although CEVAE (Louizos et al., 2017) makes no explicit injectivity assumption between C and X, it requires that the joint distribution p(C, X, T, Y ) can be fully recovered from the observations (X, T, Y ). The literature shows that some of the possible identification criteria are 1) multiple independent views of C in X (Edwards et al., 2015) , and 2) C is categorical and X is a mixture of Gaussian components determined by C (that is, X is generated by bijective mapping of C to the mean of the corresponding component with added Gaussian noise) (Anandkumar et al., 2014) .In the following part of this section, we omit the noise ϵ to gain better intuition of latent post-treatment bias (but all the exact conclusions will still hold in the posterior sense). In Section 3, we assume that noise exists and demonstrate that our method can still adequately identify latent confounders.Based on Assumption 1, we are ready to define the estimated average treatment effect (ATE) by controlling the covariates X, as well as the true (conditional) average treatment effects.Definition 1. We define the Difference in Conditional Expected Values (DCEV) asDCEV (x) = E[Y |T = 1, X = x] -E[Y |T = 0, X = x],(1) which is the difference of the expected value of the outcome Y for units with variable X = x in the treatment group and the non-treatment group. Based on DCEV (x), we define the Difference in Expected Value (DEV), i.e., DEV (X) = E p(X) [DCV E(X)] as the expected value DCEV.DEV (X) denotes the ATE estimand by controlling covariates X. If X = ∅, DEV (∅) represents the naive estimator that directly calculates the expected difference of Y between the treatment group and the non-treatment group. With the causal estimand DEV (X) introduced, we then define the true causal effects (i.e., C/ATE) when covariates X scramble both latent confounders and post-treatment variables according to the generative process described in Assumption 1. The main issue that hinders a direct definition of C/ATE with DCEV (x) and DEV (X) is that, since X contains latent posttreatment variables M , conditional on X, the strong ignorability assumption (Imbens & Rubin, 2015) widely used for the identification of causal effects does not hold2 . Accordingly, we have: Definition 2. Under Assumption 1, we define the Conditional Average Treatment Effect (CATE) for individuals with observed covariates X = x as follows:EQUATIONwith the Average Treatment Effect (ATE) of treatment T defined asEQUATIONIn Definition 2, we only consider the latent confounder component of X for CATE in Eq. ( 2), as the causal relationship between the post-treatment variables M and the outcome Y is indeterminate (see Fig. 1-(c )). However, if the specific relationship between M and Y can be further established by the researcher (e.g., all elements of M are latent mediators), more precise forms of CATE can be derived with path-specific counterfactual analysis (Imai et al., 2010; Cheng et al., 2022) .With DEV (X) (the ATE estimator that controls the covariates X), CATE, and ATE defined in Section 2.2.2, in this section, we analyze the latent post-treatment bias of existing proxy-of-confounderbased causal inference methods, such as CEVAE (Louizos et al., 2017) , that control latent variables inferred from the covariates X to estimate the ATE of T on Y , when X scrambles both latent confounders and post-treatment variables. In our analysis, Lemma 2.1 will be frequently used. Lemma 2.1. For an injective function g,E[Y |X = x] = E[Y |g(X) = g(x)] holds.The proof when g is differentiable a.e. can be referred to in Appendix A.1. Since the latent variable models used in existing methods (such as VAE with factorized Gaussian prior in CEVAE) lack identifiability guarantee (i.e., the recovery of the exact latent variables), we assume that these models can recover the true latent space Z = [C, M ] up to invertible transformations f , where the inference process can be represented asẐ = f (X) = f • f † (X).With such an assumption, we have the following theorem regarding the latent post-treatment bias when X mixes post-treatment variables. Theorem 2.2. If the observed covariates X are generated from latent confounders C and latent post-treatment variables M according to Assumption 1, the latent post-treatment bias of a proxy-ofconfounder-based causal inference algorithm that controls latent variables Ẑ inferred fromX via f = f • f † : R K X → R K C +K Mto estimate the ATE can be formulated as follows:EQUATION) which can be arbitrarily bad. Therefore, the estimator of existing proxy-of-confounder-based methods, i.e., DEV ( f (X)), is an arbitrarily biased estimator of the ATE, when the selected proxy of confounders X accidentally mixes in latent post-treatment variables M .The final step of Eq. ( 4) can be proved since f is injective and f bijective, the composite f• f † • f : [C, M ] → Ẑ is bijective, so we can use Lemma 2.1 to remove f • f † • f in the condition.Generally, the latent post-treatment bias defined in Eq. ( 4) cannot be simplified because 1) the causal relationship between M and Y is indeterminate, and 2) the causal influence of C, M , and T on Y can be arbitrary. However, for linear structural causal models with causal relationships determined between M and Y (e.g., M are mediators, which are post-treatment variables that have causal influences on the outcomes), stronger conclusions can be drawn as follows: Corollary 2.3. (MixedMediator) For the linear Structural Causal Model (SCM) defined as:EQUATIONwhere the mixture functionf = A ∈ R K X ×(K C +K M) is a full column-rank matrix, the CATE, ATE, and the bias of proxy-of-confounder-based causal inference model that controls the latent variables Ẑ inferred via Ẑ = f (X) = B T X can be formulated as follows:EQUATIONwhereB ∈ R K X ×(K C +K M ) is another full column-rank matrix. Since γ j • θ j is arbitrary, the estimator DEV ( Ẑ) = E[DCEV (B T X)] is arbitrarily biased for ATE estimation.The proof of Eq. ( 6) is provided in Appendix A.2. In addition, we show that the post-treatment variables M DO NOT necessarily need to have direct causal effects on the outcome Y to incur arbitrary bias in ATE estimation. In Appendix A.3, we provide another example (i.e., MixedCorrelator) in the linear case where M is correlated with Y through unobserved confounders U in Corollary A.1. 
methodology: 
 In this section, we introduce the proposed Confounder-identifiable Variational Auto-Encoder (CiVAE) to address latent post-treatment bias. Specifically, we first prove that if the prior distribution of the true latent variables Z = [C, M ] satisfies certain weak assumptions, identifiability criterion holds, and each dimension of the inferred latent variables Ẑ, i.e., Ẑi , corresponds to the invertible transformation of either a true confounder C j or a true post-treatment variable M k . Then, utilizing the causal relations between C, M , and T , we novelly transform the challenging confounderidentifiability problem into a tractable pair-wise conditional independence test problem, which can be effectively solved with kernel-based methods. Finally, we demonstrate that controlling the transformed confounders inferred by CiVAE can yield an unbiased estimation of the true ATE.The fundamental work of deep variational inference with identifiability guarantee, i.e., the identifiable VAE (iVAE) (Khemakhem et al., 2020) , makes a strict assumption that the prior of true latent variables Z (i.e., [C, M ] in our case) is conditionally factorized given the available covariates (i.e., the treatment T and the outcome Y in our case). However, since both latent confounders C and latent post-treatment variables M form fork structures with the outcome Y (see Fig. 1-(c )) (Koller & Friedman, 2009) , C i , C j , M i , and M j are not independent given Y . Recently, Non-Factorized iVAE (NF-iVAE) (Lu et al., 2021) was proposed that allows arbitrary dependence among the true latent variables Z in the conditional priors, where Z can be identified up to arbitrary non-linear transformations, However, the transformation are not necessarily invertible, which is risky for causal inference, as multiple values of the confounders may collapse, leading to bias when estimating the ATE by averaging the DCEV calculated in each stratum of the inferred confounders.The proposed NF-iVAE guarantees the identifiability of Z by putting a general exponential family distribution with at least one invertible sufficient statistic in the factorized part as its prior when conditioning on treatment T and outcome Y , which can be formulated as follows. Assumption 2. Let Z = [C||M ] be the random vector for latent variables that causally generate the observed covariates X according to Assumption 1. We assume that the conditional prior of Z given the outcome Y and the treatment T belongs to a general exponential family with parameter vector λ(Y, T ) and sufficient statistics S(Z) = [S f (Z) T , S nf (Z) T ] T . Specifically, S(Z) is composed of (i) the sufficient statistics of a factorized exponential family, i.e.,S f (Z) = [S 1 (Z 1 ) T , • • • , S K Z (Z K Z ) T ] T ,where all components S i (Z i ) have dimension larger than or equal to 2 and each S i has at least one invertible dimension, and (ii) S nf (Z), where S nf is a neural network with ReLU activation. The density of the conditional prior can be formulated as:EQUATION) where Q(Z) is the base measure and C(Y, T ) not dependent on Z is the normalizing constant.We justify that assumption 2 is weak and practical as follows. 1) Neural networks with ReLU activation have universal approximation ability of distributions (Lu & Lu, 2020) . Therefore, Eq. ( 7) can model arbitrary dependence between true latent confounders C and true post-treatment variables M conditional on T and Y . 2) Although CiVAE makes an extra assumption that ∀i, at least one dimension of S i is invertible, this can be easily satisfied as most commonly used exponential family distributions, such as Gaussian, Bernoulli, etc., has at least one invertible sufficient statistics3 .The reason why we use ReLU as the activation is that, the identifiability of iVAE relies on the condition that the sufficient statistics S have zero second-order cross-derivative. The factorized part, i.e., S f , satisfies it trivially since all cross-derivatives of S f are zero. In addition, since the ReLU neural networks are linear a.e., all second-order derivatives of S nf are zero. Therefore, identifiability holds after adding S nf in the prior that allows the capturing of arbitrary dependence among Z.Combining Assumptions 1 and 2, the generative process of CiVAE can be formulated as follows:p θ (X, Z | Y, T ) = p f (X | Z)p S,λ (Z | Y, T ), (8) p f (X | Z) = p ϵ (X -f (Z)).(9) where θ = (f, λ, S) ∈ Θ are the parameters of the generative distribution4 . Since the generative process of CiVAE is parameterized by deep neural networks, the posterior distribution of Z, i.e., p θ (Z | X, Y, T ), is intractable. Therefore, we resort to variational inference (Blei et al., 2017) , where we introduce approximate posterior q ϕ (Z | X, Y, T ) parameterized by deep neural network with trainable parameter ϕ, and in q ϕ (Z|•) finds the one closes to p θ (Z|•) measured by KL divergence. Minimization of the KL is equivalent to maximization of the evidence lower bound (ELBO) as:L(θ, ϕ) := E q ϕ (Z|X,Y,T )   log p f (X | Z) + log p S,λ (Z | Y, T ) -log q ϕ (Z | X, Y, T ) KL of posterior with prior    . (10)Since the normalization constant C in Eq. ( 7) is generally intractable, it is infeasible to directly learn S, λ by optimizing Eq. ( 10). Therefore, we substitute the KL term in Eq. ( 10) with the widely-used score matching (Hyvärinen & Dayan, 2005) to learn unnormalized densities instead as follows:L(S, λ, ϕ) := E q ϕ (Z|X,Y,T ) ∥∇ Z log q ϕ (Z | X, Y, T ) -∇ Z log p S,λ (Z | Y, T )∥ 2 = E q ϕ (Z|X,Y,T )   K Z j=1 ∂ 2 p S,λ (Z | Y, T ) ∂Z 2 j + 1 2 ∂p S,λ (Z | Y, T ) ∂Z j 2   + const.(11)With the generative process and optimization objective of CiVAE introduced in the previous subsections, we are ready to introduce the final assumption of CiVAE, which, combined with Assumptions 1 and 2, leads to the main theorem of this paper, which states the identifiability of CiVAE. Assumption 3. Assume the following: (i) The set {X ∈ X |ϕ(X) = 0} has measure zero, where ϕ is the characteristic function of the density p f in Eq. ( 9). (ii) The sufficient statistics, S i in S f are all twice differentiable. (iii) The mixture function f in Eq. ( 9) has all second-order cross derivatives.(iv) There exist k+ 1 distinct points (Y, T ) 0 , • • • , (Y, T ) k such that the matrix L = [λ((Y, T ) 1 ) - λ((Y, T ) 0 ), • • • , λ((Y, T ) k ) -λ((Y, T ) 0 )] of size k × k is invertible,where k = Dim(S).(i) -(iii) are trivial for neural networks. (iv) denotes that independent samples of (Y, T ) are required to identify C and M . The identifiability theorem of CiVAE can be formulated as follows.Theorem 3.1. If Assumptions 1, 2, and 3 hold, and if θ, θ ∈ Θ → p θ (X|Y, T ) = p θ (X|Y, T ), the true latent variables Z are identifiable up to permutation and element-wise bijective transformation. Furthermore, in the case of variational inference, if we denote the true parameter that generates the data as θ * , if (i) the distribution family q ϕ (Z|X, Y, T ) contains the posterior p θ (Z|X, Y, T ), and q ϕ (Z|X, Y, T ) > 0, (ii) we optimize Eq. ( 4) w.r.t. both θ, ϕ, then in the limit of infinite data, true parameters θ * can be learned up to a permutation and bijective transformation of Z.The proof of Theorem 3.1 is based on the NF-iVAE paper (Lu et al., 2021) , with the new assumption introduced in CiVAE that each S i has at least one invertible dimension incorporated to ensure that the transformation of each Z i is bijective. The detailed proof is provided in Appendix A.4.Theorem 3.1 ensures that latent variables Ẑ inferred by CiVAE cannot 1) mix confounders and post-treatment variables in each dimension, or 2) collapse different values of the latent confounders into the same value. To further determine the dimensions of confounder and post-treatment variable in Ẑ, we rely on the causal relations between latent variables Z = [C, M ] and treatment T and the associated marginal/conditional dependence properties. These are discussed as follows.• Case 1. Intra-Confounders. Latent confounders C i , C j and the treatment T form the V-structure C i → T ← C j . Therefore, C i and C j are marginally independent, whereas they become dependent when conditioning on the assigned treatment T . • Case 2. Intra-Post Treatment Variables. Latent post-treatment variables M i , M j and the treatment T form a fork-structure M i ← T → M j , where M i , M j are marginally dependent, but they become independent after conditioning on the assigned treatment T . • Case 3. Cross-Confounder and Post-Treatment Variables. Latent confounder C i , latent posttreatment variable M j , and the treatment T forms a chain structure C i → T → M j , where C i , M j are marginally dependent, and they become independent after conditioning on T .From the above analysis we can find that, the dependence between two latent variables Z i and Z j increases after conditioning on the treatment T ONLY in the case of intra-confounders. Therefore, if more than one latent confounders exist, which is highly probable when covariates X are highdimensional, we can conduct independence test Ind( Ẑi , Ẑj ) and CInd( Ẑi , Ẑj |T ) for all pairs of inferred latent variables, which can be implemented via kernel-based methods as (Zhang et al., 2012) , and select the pairs where p-value of CInd is larger than that of Ind as latent confounders.Here, we note that the kernel-based (conditional) independence test incurs N 2 ×K 2 Z complexity in the training phase. However, once the dimensions of the confounders in Ẑ are determined, CiVAE has the same complexity as CEVAE for the estimation of CATE and ATE in the test phase. Therefore, we argue that the additional complexity of model training is worthy due to the substantially increased robustness toward latent post-treatment bias (which will be demonstrated in Section 4).Finally, we show that controlling transformed confounders Ĉ inferred by CiVAE provides an unbiased estimation of ATE. Although assumptions weaker than Assumption 2, e.g., inferred confounders have the same propensity score as the true confounders (i.e., Ĉ does not have to be bijective transformation of C), could lead to the same unbiasedness results (Imbens & Rubin, 2015) , since our main purpose is to analyze the latent post-treatment bias and propose a viable solution accordingly, this introduces unnecessary complexity, which could be explored as a direction for future study. Theorem 3.2. Controlling bijective of confounders is equivalent to controlling true confounders in ATE estimation, i.e., DEV ( Ĉ) = DEV (g(C)) = AT E, if transformation function g is bijective.The proof of Theorem 3.2 for discrete C is trivial (where Ĉ = g(C) represents a simple relabeling of the stratum that we calculate the DCEV and take the expectation). The proof in the continuous case where g is differentiable is provided in Appendix A.5. With Theorem 3.2, we can control the identified latent confounders as true confounders, providing an unbiased estimate of ATE. 
datasets: 
 We establish two simulated datasets, i.e., MixedMediator and MixedCorrelator, that consider two types of post-treatment variables, i.e., 1) mediators and 2) variables that are correlated with the outcome Y via latent confounders U . The generative process of the two datasets can be referred to in Corollary 2.3 and Corollary A.1, respectively, where the latent confounders C are generated from Gaussian as C ∼ Gaussian(0, I K C ). For MixedMediator, γ is set as [-1, -1, -1], θ is set as [1, 1, 1], and τ is set as 2, which results in AT E = -1. For MixedMediator, we set the same γ and θ as MixedMediator, where parameters ϕ = 1 and τ = 1, which results in AT E = 1.In addition, we build a real-world dataset based on the job Ads data from the Company, aiming to estimate the ATE of switching a job from onsite to online working mode to the statistics of the applicants (here we choose the average age as the outcome). In the dataset, treatment T represents the working mode of the job, where T = 1 represents the job is online, whereas T = 0 represents the job is onsite, Y is the standardized age, and X ∈ {0, 1} K X indicates the required skills of the job. We select 3,228 jobs from Bay Area, where a primary study shows that DEV (∅) ≈ -2 years5 (i.e., online job applicants are two years younger than onsite job applicants). To simulate the latent confounder C and post-treatment variables M , we first learn a generative model as follows:EQUATIONwhere M ulti represents multinomial distribution, N N f is a neural network with softmax activation, Z, w ∈ R K Z , K Z = 6, and ⊙ represents the element-wise product operator, respectively. We then treat the first K C = 3 dimensions of Z as the latent confounders C and the remaining K M = K Z -K C dimensions as the latent mediators M . After learning N N f and w according to Eq. ( 12), we draw latent confounders C ∈ Gaussian(0, I), latent mediators M = T • γ, and set the outcome Y = w ⊙ [C||M ] + τ • T , where the true ATE can be calculated as sum(γ ⊙ w -K M : ) + τ .The baselines we include for comparisons can be categorized into three classes. 1) Unawareness, where no information in X is used for ATE estimation. We implement the naive LR0 estimator, which regresses Y on T and uses the coefficient to estimate the ATE (Imbens & Rubin, 2015) (LR0 equals to DEV (∅), i.e., the difference of average outcome between the treatment and non-treatment group). 2) Control-X, which directly controls the covariates X. In this class, LR1 regresses Y on T and X, whereas TarNet uses a two-branch neural network to estimate the DEV (X) 3) Control-Z, which controls latent variables Z learned from the covariates X. Methods from this class include the CEVAE (Louizos et al., 2017) and covariate disentanglement methods (see Fig. 1-(b )), such as DR-CFR Hassanpour & Greiner (2020) and TEDVAE (Zhang et al., 2021) .The comparisons are summarized in Table 1 . From Table 1 , we can empirically verify the correctness of Theorem 2.2 that post-treatment bias indeed poses a serious issue for proxy-of-confounder-based methods, because for the MixedMediator and MixedConfounder datasets, CEVAE is worse than the naive LR0 estimator that directly calculates the difference of mean outcome between the CiVAE achieves significantly better results compared to CEVAE and TEDVAE, which demonstrates its effectiveness in identifying and distinguishing latent confounders from post-treatment variables in proxies. However, we also notice that a downside of CiVAE is the comparatively large variance across ten dataset splits, as misidentifying latent mediators as confounders may result in severe performance degradation when the mediation effects are strong or the number of latent confounders is small.We show the p-value of the pairwise independence test of the true latent variables before and after conditioning on the assigned treatment T . From Fig. 2 we can find that the difference between the three cases discussed in Subsection 3.4 is significant. Here, we should note that the distinction of the intra-confounder case from other cases relies on the assumption that latent confounders are independent. If the latent confounders are correlated, we can first use causal discovery techniques such as the PC algorithm (Spirtes et al., 2000) to find direct parents of T , and use our algorithm as the refinement to determine the true confounders C from the misidentified post-treatment variables. 
conclusions: 
 In this paper, we systematically investigated the latent post-treatment bias in causal inference from observational data. We first prove that unresolved latent post-treatment variables scrambled in the proxy of confounders can arbitrarily bias the ATE estimation. To address the bias, we proposed the Confounder-identifiable VAE (CiVAE), which, utilizing a mild assumption regarding the prior of latent factors, guarantees the identifiability of latent confounders up to bijective transformations. Finally, we show that controlling the latent confounders inferred by CiVAE can provide an unbiased estimation of the ATE. Experiments on both simulated and real-world datasets demonstrated that CiVAE has superior robustness to latent post-treatment bias compared with state-of-the-art methods. 
 
 Appendix: A.1 PROOF OF LEMMA 2.1.Proof. Let Z = g(X) and z = g(x). If g is injective and differentiable a.e., and g † is the left-inverse, we have:EQUATIONwhere f • and f •|• represent the marginal and conditional density function, respectively, and J g † (z) is the Jacobian matrix of function g † evaluated at z. Based on Eq. ( 13), we have:E[Y |X] = y•f Y |X (y|x)dy = y•f Y |Z (y|z)dy = E[Y |Z = z] = E[Y |g(X) = g(x)]. (14) A.2 PROOF OF COROLLARY 2.3. Proof. For X = x, let [c||m] . = [f † C (x)||f † M (x)] . = f † (x) = A † (x -α X ), where A † is the left inverse of the full column-rank matrix A in Eq. ( 2), we have:EQUATIONwhere the first equality is due to the definition of CATE in Eq. ( 2). In addition, the causal estimand and bias of a proxy-of-confounder-based causal inference model that controls the latent variable Z inferred via Ẑ = f (X) = B T X (where B is also a full column-rank matrix) can be formulated as:EQUATIONwhere step (a) is due to the fact that, since both A and B are full column-rank matrices, B T A is an invertible matrix, and the mapping f = B T α X + B T A is bijective. Therefore, we can invoke Lemma 2.1 and apply the left-inverse of f , i.e., f † = (B T A) -1 -B T α X , to the condition of the expectation. The rest steps are based on the structural causal equations defined in Eq. (2).Corollary A.1. (MixedCorrelator) For another Linear Structural Causal Model defined as follows:EQUATIONwhere the mixture functionf = A ∈ R K X ×(K C +K M) is a full column-rank matrix, the CATE, ATE, and the bias of proxy-of-confounder-based causal inference model that controls the latent variable Ẑ inferred via Ẑ = f (X) = B T X can be formulated as follows:EQUATIONwhere B ∈ R K X ×(K C +K M ) is another full column-rank matrix. Since θj •γj ϕj is arbitrary, the estimator DEV ( Ẑ) = E[DCEV (B T X)] is arbitrarily biased for the estimation of ATE.Proof. The proof of the CATE and ATE is trivial. The causal estimand and the bias of a proxyof-confounder-based causal inference model that controls the latent variables Ẑ inferred via Ẑ = f (X) = B T X (where B is also a full column-rank matrix) can be formulated as follows:EQUATIONwhere step (a) and the rest of the proof follow the same logic as the proof in Section 2.3.A.4 PROOF OF THEOREM 3.1The strict definitions of the exponential family, strong exponential (which is assumed for the factorized part of the conditional prior), and identifiability follow (Khemakhem et al., 2020; Lu et al., 2021) , and can be referred to in Appendix E, F of (Lu et al., 2021) , which we omit to avoid redundancy. The proof of Theorem 3.1 is largely based on the NF-iVAE paper (Lu et al., 2021) , where most of the details can be found, with the new assumption introduced in CiVAE that each S f,i has at least one invertible dimension incorporated to ensure that each dimension of the inferred latent variables is a bijective transformation of the corresponding true latent variable.A.4.1 PART IStep I. In this step, we transform the equality of noisy conditional marginal distribution of X given Y, T of two models with parameter θ, θ ∈ Θ into the equality of noise-free distributions.EQUATIONStep Step II. In this step, we transform the equality of the noise-free distributions into the relationship of the sufficient statistics S and S. By taking logarithm of both sides of Eq. ( 20), we have:EQUATIONLet (Y, T ) 0 , • • • , (Y, T ) k be the k + 1 distinct points defined in Assumption 3 -(iv). We obtain k + 1 equations by evaluating the Eq. ( 21) at these points, where the first equation is subtracted from the remaining ones, which leads to the following equation system:EQUATIONLet L be the invertible matrix defined in Assumption 3 -(iv) and L be the counterpart for λ, if we summarize all terms irrelevant to X into a constant b,we have:EQUATIONwhereA = L -T L ∈ R k×k , and c = L -T b ∈ R k .Step III. Ideally, to prove the element-wise bijective identifiability of the latent variables Z, the transformation of the sufficient statistics S derived in Eq. ( 23) should be bijective. We claim that if the conditional prior p S,λ (Z | Y, T ) is strongly exponential and L is invertible, L and A must also be invertible. The proof is omitted, and can be referred to in Appendix H.1.1 of (Lu et al., 2021) .In this part, we prove that, if Assumptions 1, 2 and 3 hold, we can identify the factorized part of the sufficient statistics S(Z), i.e., S f (Z), up to permutation and element-wise transformation. Specifically, if we use v to denote the composite map f † • f : Z → Z, Eq. ( 23) can be rewritten into:EQUATIONWe aim to prove that A in Eq. ( 24) is a block permutation matrix.Step I. We start by showing that v is a component-wise function. If we differentiate both sides of Eq. ( 24) with respect to Z s and Z t , where s ̸ = t, we have:EQUATIONNote that for the factorized part of the sufficient statistics S, i.e., S f , all cross-derivatives are zero, and for the non-factorized part of S, i.e., S nf , which is a neural network with ReLU activation (i.e., linear a.e.), all second-order derivatives are zero. Therefore, the second order cross-derivatives on the LHS. of Eq. ( 25) are zero, which leads to the following equality:EQUATIONEq. ( 26) can be written into the matrix-vector product form as follows:EQUATIONwhereEQUATIONS′ (Z) = ∂ S(v(Z)) ∂v 1 (Z) , • • • , ∂ S(v(Z)) ∂v K Z (Z) ∈ R k×K Z , v ′′ s,t (Z) = ∂ 2 v 1 (Z) ∂Z s ∂Z t , • • • , ∂ 2 v K Z (Z) ∂Z s ∂Z t T ∈ R K Z .If we denote the concatenation asS′′′ (Z) = S′′ (Z), S′ (Z) ∈ R k×2K Z and v ′′ s,t (Z) = v ′ s,t (Z) T , v ′′ s,t (Z) T T ∈ R 2Kz, we have:EQUATIONFinally, if we denote the rows of S′′′ (Z) that correspond to the factorized part of S by S′′′ f (Z), according to Lemma 5 of the iVAE paper (Khemakhem et al., 2020) and the assumption that k ≥ 2K Z , we have that the rank of S′′′EQUATIONZ) ∈ R 2K Z is a zero vector. Therefore, we have v ′ s,t (Z) = 0, ∀s ̸ = t, and we have demonstrated that v is a component-wise function.Step II. Based on Step I, we demonstrate that A is a block permutation matrix. Without loss of generality, we assume that the permutation in v is Identity, where v(Z) = [v 1 (Z 1 ) , • • • , v K Z (Z K Z )]T and each v i is a nonlinear univariate scalar function. Since f and f are injective, v is bijective and 24) can be reformulated as S(Z) = AS(v(Z)). We then apply v -1 to Z on both sides, which givesv -1 (Z) = v -1 1 (Z 1 ) , • • • , v -1 K Z (Z K Z ) T . If we denote S(v(Z)) = S(v(Z)) + A -1 c, Eq. (EQUATIONLet t be the index of an entry in S that corresponds to the factorized part S f . For all s ̸ = t, we have:EQUATIONSince the entries of S are linearly independent, a tj is zero for any j such that ∂S(Z)j ∂Zs ̸ = 0. This includes the entries S j that correspond to 1) the factorized part that does not depend on Z t ; and 2) the non-factorized part S nf . Therefore, when t is the index of an entry in the sufficient statistics S that corresponds to factor i in the factorized part S f , i.e., S f,i , the only non-zero a tj are the ones that map between S f,i (Z i ) and S f,i (v i (Z i )). Therefore, we can construct an invertible submatrix A ′ i with all non-zero elements a tj for all t that corresponds to factor i, such thatEQUATIONwhere c i denotes the corresponding elements of c. Eq. ( 31) means that for each i = 1, • • • , K Z , the matrix block A ′ i of A affinely transforms the i-specific sufficient statistics vector S f,i (Z i ) into Sf,i (v i (Z i )). In addition, there is also an additional block A ′ that affinely transforms S nf (Z) in into S nf (v(Z)). This completes the proof that A is a block permutation matrix.Let Zi = v i (Z i ) = f † (X) i be the ith inferred latent variable. Assume again that the permutation in v is Identity. In this part, we prove that if Assumption 2 holds, each inferred latent variable Zi is the bijective transformation of the true latent variable. The proof is as follows.Proof. Plugging Zi into Eq. (31), we have:EQUATION)According to Assumption 2, there exists one dimension of S f,i , i.e., j, such that S f,ij is bijective. This implies that S f,i is injective, and therefore it has a left-inverse S † f,i . we apply S † f,i to both sides of Eq. ( 32), which gives:EQUATIONSince A ′ i is a block of an invertible block permutation matrix, A i is also an invertible matrix, and therefore A ′ i is a bijective mapping. In addition, since Sf,i is injective, Sf,i is also injective, and therefore the composite map S † f,i A ′ i Sf,i : R → R that applies on Zi is a bijective. This completes the proof that each inferred latent variable Zi is the bijective transformation of the true latent variable in the case of no noise, where Z = f † (X) are the true latent variables. If noise ε exists, the posterior distribution of the latent variables can be identified up to an analogous bijective indeterminacy.Proof. If the family of the variational posterior q ϕ (Z|X, Y, T ) contains the true posterior p θ (Z|X, Y, T ), then by optimizing the loss of Eq. (10) (with the KL term replaced by the score matching loss defined in Eq. ( 11)) over its parameter ϕ, the score matching term will eventually vanish. Therefore, the ELBO term in Eq. ( 10) will be equal to the log-likelihood. Under this circumstance, CiVAE inherits all the properties of maximum likelihood estimation (MLE). Since the identifiability of CiVAE is guaranteed up to permutation and component-wise bijective transformation of the latent variables, the consistency property of MLE means that the model will converge to the true parameter θ * up to such mild indeterminacy of the latent variables in the limit of infinite data.A.5 PROOF OF THEOREM 3.2Proof. Let C be the true latent confounders and Ĉ be the transformed confounders, where the transformation function g is bijective and differentiable a.e. Let g -1 denote its inverse. The ATE estimator that controls transformed confounders Ĉ can be formulated as:EQUATIONSpecifically, for the continuous case where density functions exist, for each term, we have:EQUATION)For the marginal density f Ĉ (ĉ), the following equality holds:EQUATION)As for the conditional density f Y |T, Ĉ (y|t, ĉ), since g is bijective, according to Eq. ( 13), we have:EQUATION)Combining Eqs. ( 36) and ( 37), and given that dĉ = |J g (c)|dc, we have:EQUATIONwhere the term where the last step is due to Eq. (2) in Definition 2, which completes our proof that controlling bijectively transformed confounders provides an unbiased estimation of ATE.Anonymized due to double-blind review policy.Equivalently, we could say that given covariates X, the backdoor criteria between T and Y does not hold, which requires the conditional set of variables contains no descendants of the treatment T(Glymour et al., 2016).There are a few exponential family with no invertible sufficient statistics, e.g., Weibull distribution when shape parameter k is even.Note that although f is a function, we include it in the parameter set to be consistent with the iVAE paper.which leads to -0.178 after standardization. Code demo see https://anonymous.4open.science/r/CiVAE-demo-54B9. 


Decision: This paper addresses causal effect estimation with unobserved confounders, focusing on recovering confounding information from auxiliary proxy variables. Specifically, it tackles the challenge of proxies that capture information about both unobserved confounders and latent post-treatment variables, which can introduce post-treatment bias. The paper proposes a VAE approach to individually recover latent confounders and post-treatment variables up to bijective transformations. It then aims to disentangle these components and adjust for latent confounders in causal effect estimation.

While the reviews are overall positive, some concerns were raised, including the plausibility of stringent assumptions, handling of general confounders and post-treatment variables, and the presentation of experimental results. The rebuttals partially addressed these concerns, but fully resolving the issues, especially regarding general confounders and post-treatment variables, requires significant work. The paper emphasizes the VAE component and briefly mentions using pairwise independence tests to disentangle confounders and post-treatment variables under a restrictive independence assumption. The appendix offers brief discussions on more general cases without independence restrictions. However, this is a very cursory treatment of the disentanglement procedure, one key component that differentiates the work from existing literature. A deeper study of the disentanglement process is needed, including clear procedures, validity explanations, and validation through well-designed ablation studies. 

During the AC-reviewer discussion, another crucial question arose: Is the problem studied practically relevant? It is unclear when the problem with post-treatment proxies would be encountered in real scenarios. Since confounders are pre-treatment, it's generally more reasonable to use pre-treatment information. The paper lacks a discussion on the practical relevance of the problem, providing only a simplified and somewhat unrealistic example in the introduction that may represent only a narrow use case. Additionally, the example involving job mode switch and applicants' age is confusing, as age is an exogenous attribute likely influencing decisions rather than the other way around, and variables like seniority and skills are also ambiguous.

Given these issues, I recommend rejection at this stage.

Reviews: review: paper_summary: paper_summary: In this paper, the authors investigated the issue of latent post-treatment bias in causal inference from observational data. They showed that estimator of existing proxy-of-confounder-based methd, i.e., DEV (f(X)), is an arbitrarily biased estimator of the Average Treatment Effect (ATE), when the selected proxy of confounders X accidentally mixes in latent post-treatment variables (Theorem 3.2). To address this issue, they proposed the Confounder-identifiable VAE (CiVAE), which identifies latent confounders up to bijective transformations under a mild assumption regarding the prior of latent factors. They showed that controlling for latent confounders inferred by CiVAE can provide an unbiased estimation of the ATE. Experiments on both simulated and real-world datasets demonstrate that CiVAE exhibits superior robustness to latent post-treatment bias compared to state-of-the-art methods.
strength_weakness: strength_weakness: The presentation can be improved.
questions: questions: Is Fig. 1(c) general enough? It assumes that all latent variables are either confounders or post-treatment variables. However, there can be other types of latent variables, such as:
1. Pre-treatment Variables: These latent variables influence the treatment (T) but do not directly affect the outcome (Y) or the additional observation (X). They exist before the treatment is applied and can introduce selection bias.
2. Latent Interaction Variables: These latent variables interact with the treatment (T) to influence the outcome (Y). They are not confounders because they do not influence the treatment directly, nor are they post-treatment variables.
3. Latent Mediator Variables: These latent variables mediate the effect of the treatment (T) on the outcome (Y) and are not directly observed.
4. Latent Variables Influencing Both Pre-treatment and Post-treatment States: These latent variables influence the state of the system both before and after the treatment but do not fit the typical definition of confounders or post-treatment variables. For example, a latent mental state might affect both a person's initial willingness to undergo treatment and their behavior or responses post-treatment.

Can the proposed method handle these types too (with some extension), or some of the types are quite disruptive to the proposed methodology?
limitations: limitations: n.a.
paper_summary: paper_summary: This paper addresses the challenge of causal inference with observational data, particularly when direct measurement of confounders is infeasible. The authors propose a new method, Confounder-identifiable Variational Autoencoder (CiVAE), to mitigate post-treatment bias using observed proxies for both latent confounders and latent post-treatment variables. The paper provides a theoretical analysis under specific assumptions and validates the proposed approach through experiments on both simulated and real-world datasets.
strength_weakness: strength_weakness: * In practical scenarios, interactions among latent factors are often present and can significantly impact the estimation. It would be beneficial if the authors could elaborate on how their method addresses these interactions and whether there are any theoretical guarantees regarding their handling in the proposed approach.

* The theoretical guarantees rely on strong assumptions, and the assumptions are hard to verify in practice. In assumption 1, the paper assumes an injective function of latent confounders and latent post-treatment variables into the observed proxy. This is a strong assumption,  and it will be much harder to meet the assumption in general when the function is nonlinear. The specific setup with strong assumptions limits the practical applicability of the proposed approach. It would be helpful if the authors could provide examples where these assumptions hold and demonstrate how they can be verified.

* The experiment lacks sufficient details on setup and implementation. Could the authors provide more specific information to enhance understanding of the empirical results?
questions: questions: See Weaknesses.
limitations: limitations: * The proposed method relies on very strong assumptions to ensure identifiability, which can be challenging to verify in practical applications.
paper_summary: paper_summary: The authors deal with latent post-treatment bias for proxy-based methods which are employed for causal effect estimation.
They show that post-treatment variables can be latent and mixed into the observed covariates along with the latent confounders.
The authors transform the confounder-identifiability problem into a tractable pair-wise conditional independence test problem.
They prove that the latent confounders and latent post-treatment variables can be identified up to bijective transformations. Finally, they provide experimental analysis for their approach.
strength_weakness: strength_weakness: Here I provide some weaknesses of the paper:
* Bi-directed edges in Figure 1 are not defined properly.
* Do-operator in equation 3 is not defined in detail.
* Assumptions in Assumption 2 should be described in more detail.
* The proposed method seems to depend on a lot of assumptions. Assumptions 1,2,3 each contain multiple assumptions. The authors should explain how their assumptions hold for the real-world scenarios they considered in their experiment section.
questions: questions: Here I provide some questions:
* Why do the authors assume that the models can recover the true latent space up to invertible transformation (line 151) ? How realistic is that assumption?
* Do the proxy-based methods claim to perform well for the causal graphs in Fig 1c?
* How do these assumptions hold when X is high-dimensional?
* What values of K_C and K_M are considered?
limitations: limitations: The authors discussed a very few limitations of their paper but more discussion should be done.
paper_summary: paper_summary: This paper investigates the problem of latent post-treatment bias in causal models where there exists some proxy variables of the latent confounder and post-treatment variables. The authors first derive a general form of latent post-treatment bias which is intractable in most situations (except in special cases such as linear SCM). The authors state that the latent post-treatment bias can be arbitrarily bad for existing proxy-based causal inference methods. They then propose an identifiable VAE-based causal inference algorithm under the assumption that at least one dimension of each sufficient statistic of the latent prior is invertible. The proposed method is evaluated on both synthetic and real-world datasets to demonstrate its causal effect estimation capability with the presence of both latent confounders and post-treatment variables.
strength_weakness: strength_weakness: •	For the illustrative example in the introduction, it might be better to explicitly specify what the post-treatment variable is.

•	Other existing works [1-3] on identifying latent confounder/mediators based on the iVAE architecture should also be included in the related work.

•	The role of post-treatment variables $\boldsymbol{M}$ seems to be a bit ambiguous. To be specific, is Theorem 4.1 valid for all types of relationships between $\boldsymbol{M}$ and $Y$?

•	The illustration of (iv) in Assumption 3 is a little confusing, as it assumes one extra degree of freedom on the prior parameters of $\boldsymbol{Z}$ and is critical to the identifiability of $\boldsymbol{Z}$ from $\boldsymbol{X}$. More explanation on this point will be appreciated.

•	The empirical evaluation consists of only one real-world dataset, which somehow limits the applicability of the proposed method.

References:

[1]. Zhou, D., & Wei, X. X. (2020). Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE. Advances in Neural Information Processing Systems, 33, 7234-7247.

[2]. Sorrenson, P., Rother, C., & Köthe, U. (2020). Disentanglement by nonlinear ica with general incompressible-flow networks (gin). arXiv preprint arXiv:2001.04872.

[3]. Jiang, Z., Liu, Y., Klein, M. H., Aloui, A., Ren, Y., Li, K., ... & Carlson, D. (2023). Causal Mediation Analysis with Multi-dimensional and Indirectly Observed Mediators. arXiv preprint arXiv:2306.07918.
questions: questions: • Is $DEV(\tilde{f}(\boldsymbol{X})) = \mathbb{E}[\mathbb{E}[Y | T = 1, \tilde{f}(\boldsymbol{X})] - \mathbb{E}[Y | T = 0, \tilde{f}(\boldsymbol{X})]]$ in Eq. 4 also based on Lemma 3.1? If yes, it should be explicitly stated.

• In what cases can the bias in Theorem 3.2 be arbitrarily bad besides the causal model assumed by linear SCM in Corollary 3.3?

• What is the rationale behind the simulation procedure of $\boldsymbol{C}$ and $\boldsymbol{M}$ in Eq. 11 if the latent confounder represents “job seniority” as elaborated in the introduction? How do you anticipate the estimation error to change if we increase the complexity of the neural network $NN_f$ in Eq. 11?
limitations: limitations: The authors do not include a paragraph discussing the limitations and potential societal impact of this work.

Comments: title: Response to author's comment
Response to rebuttal
comment: **Dear Reviewer kHjK,**

Thank you for the acknowledgment. We are glad that our responses partially address your concerns. We will take your remaining comments seriously and further polish the paper, and we are committed to integrating your valuable comments into the paper. Thank you again for your time and efforts.

Best,    
Authors
I thank the authors for providing detailed responses, and these partially address the concerns. I will keep my rating.
**Dear Reviewer jEuj,**

Thank you for the acknowledgment. We are committed to integrating your valuable comments into the paper. Thank you again for your time and efforts.

Best,   
Authors
Thank you for your clarification. I have no further questions at this point.
**Dear reviewer jEuj,**

Thank you for the acknowledgment. We are glad to hear that our responses mostly address your concerns. We will try our best to integrate your valuable comments into the paper. Here, we provide responses to your further questions.

**1.** In the current version, we report the difference between the true ATE and the estimated ATE as the error. We will also provide the absolute difference to make the comparison between different methods more straightforward.

**2.**  We note that $DEV(X^{\prime})$ is the ATE estimator when controlling an arbitrary variable $X^{\prime}$. $DEV(X^{\prime})$ is defined as: $$DEV(X^{\prime}) := \mathbb{E}_{p(X^{\prime})}[DCEV(X^{\prime})] := \mathbb{E}[\mathbb{E}[Y|T=1, X^{\prime}] - \mathbb{E}[Y|T=0, X^{\prime}]],$$ where ":=" denotes the RHS is the definition of the LHS. Here, $X^{\prime}$ is an arbitrary variable. 

For Eq. (4), we are discussing the bias when estimating the ATE when controlling the variable $\tilde{f}(X)$, i.e., the latent variable inferred from $X$ via $\tilde{f}$. The bias is defined as the difference between the true ATE and the estimated ATE, i.e., $ATE - DEV(\tilde{f}(X))$. Therefore, the second step simply uses the definition of $DEV$, i.e.,  $$ATE - DEV(\tilde{f}(X)) := ATE - \mathbb{E}_{p(\tilde{f}(X))}[DCEV(\tilde{f}(X))] := \mathbb{E}[\mathbb{E}[Y|T=1,\tilde{f}(X)] - \mathbb{E}[Y|T=0, \tilde{f}(X)]],$$
where no derivation is involved in this step. The actual derivations occur in the 3rd and 4th step, where the 4th step uses Lemma 3.1 to remove the injective in the condition.  We will change "=" in the second equation of Eq. (4) to ":=" to avoid confusion. Thank you again for raising this important question to make our paper clearer.

Best,   
Authors
Thank you for your response to my comments, though the negative ATE error on the Jobs dataset looks a bit weird. Also, for the second step in Eq. 4 where you replace $\boldsymbol{X}'$ with $\tilde{f}(\boldsymbol{X})$, I mean this substitution probably needs the injectivity statement in Lemma 3.1.

I think the authors' response mostly addresses my concerns, and I'm willing to raise my score.
**Dear reviewer VF9h,**

Thank you for the acknowledgment. We will try our best to integrate your valuable comments into the paper. Thank you again for your time and efforts.

Best,    
Authors
Thank you. I've read your rebuttal, responses, and the other reviews. I will keep an eye on the reviewers' discussion phase, if there is one.