Title: Causal Effect Estimation with Mixed Latent Confounders and Post-treatment Variables
Venue: NeurIPS 2024
Accepted: no
Keywords: causal effect estimation

latent post-treatment bias

identifiable VAE
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

Abstract

Causal inference from observational data has attracted considerable attention among researchers. One main obstacle is the handling of confounders. As direct measurement of confounders may not be feasible, recent methods seek to address the confounding bias via proxy variables, i.e., covariates postulated to be conducive to the inference of latent confounders. However, the selected proxies may scramble both confounders and post-treatment variables in practice, which risks biasing the estimation by controlling for variables affected by the treatment. In this paper, we systematically investigate the bias due to latent post-treatment variables, i.e., latent post-treatment bias , in causal effect estimation. Specifically, we first derive the bias when selected proxies scramble both confounders and post-treatment variables, which we demonstrate can be arbitrarily bad. We then propose a novel Confounderidentifiable V AE (CiV AE) to address the bias. Based on 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, CiV AE individually identifies latent confounders and latent post-treatment variables up to bijective transformations. We then prove that with individual identification, the intractable disentanglement problem of latent confounders and post-treatment variables can be transformed into a tractable 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 significantly improved robustness of CiVAE.

1 Introduction

Causal inference, which aims to infer cause-and-effect relations from data, has gained increasing 23 prominence in various fields, such as social science, economics, and public health [10, 17, 34]. 24 Traditional methods rely on the golden standard of randomized control trials (RCT) to draw valid 25 causal conclusions via experimentation [6]. Recently, more attention has been dedicated to causal 26 inference from observational data, where treatments, outcomes, and unit features are passively 27 observed, and researchers have no control over the treatment assignment mechanism [36, 37, 40]. 28

One main obstacle to inferring valid causal relations from observational data is the confounding 29

bias, which occurs when we fail to account for the systematic difference between the treatment and 30

non-treatment group due to variables that causally influence the past treatments and the outcome, i.e., 31

unobserved confounders [16]. If the confounders can be measured, a simple strategy to address the 32

33

bias is to control them via covariate adjustment [33] or propensity score re-weighting [24]. However,

confounders are not always measurable [23]. Therefore, recent methods seek to adjust for the 34

influence of unobserved confounders based on their proxies, which are easily acquirable covariates 35

postulated to be causally related with the unobserved confounders [29, 42, 28]. One exemplar work 36

(a) SCMAssumed by CEVAE

(c) SCMAssumed by the Proposed CiV AE

T

Y

C

X

T

Y

C

1

X

C

K

C

M

1

M

K

(b) SCMAssumed by TEDVAE

T

Y

C

X

I

A

Instrumental

Variables

Adjusters

Confounders

Post-treatment

Variables

Confounders

M

predictive

for only T

predictive

for only Y

correlate with

both T and Y

correlate with

both T and Y

is the causal effect variational auto-encoder (CEV AE) [25], which has demonstrated that confounding 37 bias can be mitigated by controlling latent variables inferred from the proxies of confounders. 38

39

40

41

42

43

44

45

46

47

48

Although proxy-based methods have achieved substantial progress in recent years, they may risk controlling latent post-treatment variables scrambled in the proxies, where latent post-treatment bias can be introduced. Here, we note that the negative effects of controlling observed post-treatment variables have been investigated in prior research [1, 9, 21]. For example, Montgomery et al. [30] found that more than 50% of the papers published in top journals of politics inadvertently control post-treatment variables in the experimental setting, even though researchers have complete control over which covariates to control for. On this basis, we postulate that the post-treatment bias could be even worse for proxy-based methods in the setting of observational study where variables are passively recorded. In addition, the post-treatment variables can be latent and scrambled into the observed covariates together with the latent confounders, which makes them difficult to disentangle.

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

Consider a real-world example from the Company 1 . We found that changing a job from onsite to online mode causes applicants to make different decisions, and we want to estimate the causal effects of switching a job from onsite to online mode to the decisions of the applicants (reflected by statistics of applicants that apply for the job). In this case, the Company collected two groups of online (treated) and onsite (control) jobs, where the statistics of the applicants (e.g., the average age) are calculated as the surrogate outcome. Clearly, job seniority is a confounder, since less senior jobs are more likely to permit online work, and applicants for these jobs tend to be younger. However, the seniority level of a job can be difficult to measure. Therefore, the required skills of the job can be used as the proxy of the confounder "seniority", as senior jobs tend to require more advanced skills. However, a caveat is that switching to an online work mode may also alter the required skills of a job, thereby affecting the qualification and, therefore, the decision of the applicants. Consequently, directly using the skills as the proxy of the confounder "seniority" for adjustment could unintentionally control latent mediators (changed skills), which introduces latent post-treatment bias in the causal effect estimation.

Addressing the latent post-treatment bias faces multi-faceted challenges. First, there lacks a theoretical formulation of the bias when selected proxies scramble latent post-treatment variables for existing proxy-based methods. In addition, it is difficult to distinguish confounders and posttreatment variables in the latent space due to their similar observed behaviors. Existing covariate disentanglement-based methods, e.g., TEDVAE [44], focus on an easier task of disentangling latent confounders with latent adjusters and instrumental variables, which can be achieved by leveraging their different predictive abilities w.r.t. the treatment and outcome. However, since both latent confounders and post-treatment variables correlate with the treatment and the outcome, they cannot be disentangled by these methods. Finally, even if latent confounders can be distinguished from posttreatment variables, since most existing latent variable models have no identifiability guarantee [19], it is unclear whether controlling the inferred latent variables, which may be arbitrary transformations of the true confounders, can provide unbiased estimations of true causal effects.

To address the aforementioned challenges, we first analyze existing proxy-based methods when se74 lected proxies scramble both latent confounders and post-treatment variables and show the estimation 75 can be arbitrarily biased. We then propose a novel Confounder-identifiable V AE (CiV AE) to address 76 the latent post-treatment bias. Specifically, we prove that based on a mild assumption that the prior 77 of latent variables that generate the observed proxy (i.e., the latent confounders and post-treatment 78 variables) belong to a general exponential family with at least one invertible sufficient statistic in the 79 factorized part, latent confounders and latent post-treatment variables can be individually identified up 80 to simple bijective transformations . With such identifiability guarantee, based on the causal relations 81 among confounders, mediators, and treatment, we further demonstrate that the inferred confounders 82

(which are actually transformed proxies of the true confounders) could be properly distinguished 83 from the latent post-treatment variables with pair-wise conditional independence tests. Finally, we 84 prove that the true causal effects can be unbiasedly estimated based on transformed confounders 85 inferred by CiVAE. Experiments on both simulated and real-world datasets demonstrate that CiVAE 86 shows more robustness to latent post-treatment bias than existing methods. 87

88

2 Problem Formulation

In this paper, we assume the causal model in Fig. 1-(c). We use a binary random variable T to 89 denote the treatment, a random vector X ∈ R K X to denote the observed covariates (i.e., the proxy), 90 and a random scalar Y ∈ R to denote the outcome. Furthermore, the observed covariates X are 91 assumed to be generated from K C independent latent confounders C ≜ [ C 1 , C 2 ..., C K C ] causally 92 influencing both T and Y , and K M latent post-treatment variables M ≜ [ M 1 , M 2 ..., M K M ] under 93 the causal influence of the treatment (where the relation between M and Y can be arbitrary). We use 94 the random vector Z ≜ [ C || M ] ∈ R K Z = K C + K M to denote all latent factors. Our aim is to estimate 95 the average causal effects of treatment T on outcome Y with auxiliary confounder information in X , 96 where the estimation should be devoid of both confounding bias and post-treatment bias. 97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

3 Theoretical Analysis of Latent Post-Treatment Bias

3.1 Preliminaries and Assumptions

To achieve such a purpose, we first 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 directly used 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, [31] assumes that X has at least the same number of categories as C , whereas the effect restoration algorithm [35] assumes that the matrix of p ( C , X ) to be full-rank. Although CEVAE [25] 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 ) . [2] show that some of the possible identification criteria for the recovery include 1) having multiple independent views of C in X [8], 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).

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 [19]). In Section 4, we assume noise exists and demonstrate that our method can still properly identify the latent confounders.

3.2 Causal Estimand and the True ATE

Based on Assumption 1, we are ready to define the estimand of average treatment effect (ATE) through controlling the covariates X ′ , as well the as the true (conditional) average treatment effects.

Definition 1. (DCEV & DEV) . We define the Difference in Conditional Expected Values (DCEV) as:

which is the difference of the expected value of 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) as DEV ( X ′ ) = E p ( X ′ ) [ DCEV ( X ′ )] as the expectation of DCEV w.r.t. p ( X ′ ) .

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

DEV ( X ′ ) denotes the estimand of ATE when X ′ is the covariates that we choose to control (i.e., calculate the expected difference in each stratum of X ′ = x ′ ). If X ′ = ∅ , DEV ( ∅ ) represents the naive estimator that directly calculates the expected difference of the outcome Y between the treatment group and the non-treatment group. With the causal estimand DEV ( X ′ ) defined, we then derive the true causal effects with the covariates X ′ when it scrambles both latent confounders and post-treatment variables according to the generative process described in Assumption 1:

Definition 2. Under Assumption 1, we define the Conditional Average Treatment Effect (CATE) for individuals with observed covariates X = x by controlling only the confounder part in X as:

with the Average Treatment Effect (ATE) of treatment T defined as:

Please note that we only consider the latent confounder component of the observed features X in the definition of CATE in Eq. (2). This is because the causal relationship between the post-treatment variables M and the outcome Y is indeterminate. 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 [5, 14].

3.3 Latent Post-Treatment Bias

With DEV ( X ′ ) (the ATE estimator that control for the covariates X ′ ), CATE, and ATE defined in Section 3.2, in this section, we analyze the latent post-treatment bias of existing proxy-of-confounderbased causal inference methods, such as CEVAE, that control for 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 as Assumption 1. In our analysis, Lemma 3.1 will be frequently used.

Lemma 3.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 C.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 ˆ Z = ˜ 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 3.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 proxybased causal inference algorithm that controls latent variables ˆ Z inferred from X via ˜ f = ¯ f ◦ f † : R K X → R K C + K M to estimate the ATE can be formulated as follows:

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 ] → ˆ Z is bijective, so we can use Lemma 3.1 to remove ¯ f ◦ f † ◦ f in the condition.

3.4 Examples in the Linear Case

Generally, the latent post-treatment bias defined in Eq. (4) cannot be simplified, because (i) the causal relationship between M and Y are indeterminate, and (ii) the causal influence of C , M , and T on Y can be arbitrary. However, for linear structural causal models with determined causal relationships 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 3.3. ( Mixed Latent Mediator ). For the linear Structural Causal Model (SCM) defined as: 170

where the mixture function f = A ∈ R K X × ( K C + K M ) is a full column-rank matrix, the CATE, ATE, 171 and the bias of proxy-of-confounder-based causal inference model that controls the latent variables 172 ˆ Z inferred via ˆ Z = ˜ f ( X ) = B T X can be formulated as follows: 173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

where B ∈ R K X × ( K C + K M ) is another full column-rank matrix. Since ∑ γ j · θ j is arbitrary, the estimator DEV ( ˆ Z ) = E [ DCEV ( B T X )] is arbitrarily biased for ATE estimation.

The proof of Eq. (6) is provided in Appendix C.2. In addition, we show that 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 C.3, we provide another example (i.e., Mixed Latent Correlator) in the linear case where M is correlated with Y through unobserved confounders U in Corollary C.1.

4 Methodology

In this section, we introduce the proposed Confounder-identifiable Variational Auto-Encoder ( CiVAE ) in detail. Specifically, we first prove that if the prior distribution of the true latent variables Z = [ C , M ] satisfies certain weak assumptions, CiVAE individually identify [ C , M ] up to bijective transformations. Then, utilizing the causal relations between C , M , and T , we novelly transform the challenging confounder-identifiability problem into a tractable pair-wise conditional independence test problem, which can be effectively solved with kernel-based methods. The generalization of CiVAE to address the interactions among [ C , M ] are discussed in Section D of the Appendix.

4.1 Generative Process

The fundamental work on the identifiability of deep variational inference, i.e., the identifiable V AE (iVAE) [19], 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. However, since both C and M form fork structures with the outcome Y (see Fig. 1-(c)) [22], C i , C j , M i , and M j are not independent given Y . Recently, Non-Factorized iVAE (NF-iVAE) [26] 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 is not necessarily invertible, which is risky 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.

In contrast to NF-iV AE, CiV AE guarantees the individual and bijective identifiability of Z by putting a general exponential family 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:

where Q ( Z ) is the base measure, and C ( Y, T ) is the normalizing constant independent of Z .

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

We justify that assumption 2 is weak and practical as follows. (i) Neural networks with ReLU activation have universal approximation ability of distributions [27]. Therefore, Eq. (7) can model arbitrary dependence between true latent confounders C and post-treatment variables M conditional on T and Y . (ii) 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 statistics 2 .

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 as 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 .

4.2 Optimization Objective

Combining Assumptions 1 and 2, the generative process assumed by CiVAE can be formulated as:

where θ = ( f, λ , S ) ∈ Θ are the parameters of the generative distribution. 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 [4], where we introduce an approximate posterior q ϕ ( Z | X , Y, T ) parameterized by a deep neural network with a trainable parameter ϕ , and in q ϕ ( Z |· ) finds the one closest to p θ ( Z |· ) measured by KL divergence. The minimization of KL is equivalent to maximization of the evidence lower bound (ELBO):

Since the normalization constant C in Eq. (7) is generally intractable, it is infeasible to directly learn S , λ by optimizing Eq. (9). Therefore, we substitute the KL term in Eq. (9) with the widely-used score matching [13] to learn unnormalized densities instead as follows:

4.3 Identifiability of CiV AE

With the generative process and optimization objective of CiV AE discussed in previous sub-sections, we are ready to introduce the final assumption of CiV AE, which, combined with Assumptions 1 and 2, leads to the main Theorem of this paper, which states the identifiability of CiV AE.

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. (8). (ii) The sufficient statistics, S i in S f are all twice differentiable. (iii) The mixture function f in Eq. (8) has all second-order cross derivatives. (iv) There exist k + 1 distinct points ( Y, T ) 0 , · · · , ( Y, T ) k s.t. 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 ) .

Here, we note that Assumptions (i) - (iii) are trivial for differentiable neural networks. The Assumption (iv) can be intuitively understood as independent samples of ( Y, T ) are required to identify C and M . The identifiability theorem of CiV AE can be formulated as follows.

Theorem 4.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 .

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

The proof of Theorem 4.1 non trivially extends the NF-iVAE paper [26] by incorporating the new assumption introduced in CiVAE (i.e., each S i has at least one invertible dimension) to ensure that the transformation of each Z i is bijective. The detailed proof is provided in Appendix C.4 for reference.

4.4 Identification of Latent Confounders

Theorem 4.1 ensures that the latent variables ˆ Z inferred by CiVAE cannot (i) mix confounders and post-treatment variables in each dimension, or (ii) collapsing of different values of the latent confounders into the same value. To further determine the dimensions of confounder and posttreatment variable in ˆ Z , we rely on the causal relations between latent variables ˆ Z and the treatment T and the associated marginal/conditional dependence properties, which 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 confounder exists, which is highly probable when covariates X are highdimensional, we can conduct independence test Ind ( ˆ Z i , ˆ Z j ) and CInd ( ˆ Z i , ˆ Z j | T ) for all pairs of inferred latent variables, which can be implemented via kernel-based methods as [43], and select the pairs where the 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 ˆ Z are determined, CiVAE has the same complexity as CEVAE for the estimation of CATE and ATE in the test phase.

4.5 ATE Estimator with Transformed Confounders

Finally, we demonstrate that controlling the transformed confounders ˆ C inferred by CiVAE provides an unbiased estimation of ATE. Specifically, we have the final Theorem show the unbiasedness.

Theorem 4.2. Controlling bijective of confounders is equivalent to original confounders in ATE estimation, i.e., DEV ( ˜ C ) = DEV ( g ( C )) = ATE , if the transformation function g is bijective.

The proof of Theorem 4.2 for discrete C is trivial (where ˆ C = 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 C.5. With Theorem 4.2, we can control the identified latent confounders as true confounders, providing an unbiased estimate of ATE.

5 Empirical Study

In this section, we provide and analyze the experiments we conduct on both simulated and real-world datasets, where a code demo written in PyTorch and Pyro is provided in this anonymous URL.

5.1 Datasets

Simulated Datasets. We first establish two simulated datasets, i.e., LatentMediator and LatentCorrelator , that consider two types of post-treatment variables, i.e., (i) mediators and (ii) correlators, i.e., variables that are correlated with the outcome Y via latent confounders U , where the causal generative process is under the full control of the experimenter. The generative process of the two datasets can be referred to in Corollary 3.3 and Corollary C.1 in the Appendix, respectively. In our experiments, C are generated from Gaussian distribution as C ∼ Gaussian (0 , I K C ) . For

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

0

1

1.0

0

1

1.0

1

2

1.0

1.0

2

1.0

0

1

1.0

0.8

0.8

0.8

0.8

8°0

0.8

0.6

0.6

0.6

0.6

0.6

0.6

0.4

0.4

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.2

0.2

0.0

0.0

0.0

0.0

0.0

0.0

(a)

Case 1:

Intra-Confounder

LatentMediator , γ is set as [ -1 , -1 , -1] , θ is set as [1 , 1 , 1] , and τ is set as 2, which results in ATE = -1 . For the LatentCorrelator dataset, we set the same γ and θ as the LatentMediator dataset, where parameters ϕ and τ are set to 1, which results in an overall ATE of 1 .

Real-world Datasets. In addition, we build real-world datasets from the Company to estimate the ATE of switching a job from onsite to online work mode to the statistics of the applicants . The average age and the variance of gender of the applicants are two outcomes of interest. Covariates X ∈ { 0 , 1 } K X include the required skills of the job. Specifically, we establish a cohort of 3,228 jobs from the Bay Area in the US, where a preliminary study shows that DEV ( ∅ ) ≈ 2 years 3 (i.e., online job applicants are two years younger than onsite job applicants in the collected data), and DEV ( ∅ ) ≈ -0 . 015 (i.e., online jobs exhibit 0.015 more gender variance than onsite jobs in the collected data). To simulate C and M , we first learn a generative model as follows:

where Multi represents multinomial distribution, NN f is a neural network with softmax activation, Z , w ∈ R K Z , K Z = 8 , and ⊙ represents the element-wise product operator, respectively. We then treat the first K C = 5 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 NN f and w according to Eq. (11), 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 : ) + τ .

5.1.1 Disentangle Confounders and Post-treatment Variables

We first show the p -value of the kernel-based 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 distinction of the intra-confounder case from the other two cases discussed in Subsection 4.4 is significant. Here, we should note this 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 [39] 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 (Experiments see Section D) in Appendix.

5.2 Baselines

The baselines we include for comparisons can be categorized into three classes. (i) 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 [15] (LR0 equals to DEV ( ∅ ) , i.e., the difference of the average outcome between the treatment and non-treatment group). (ii) ControlX , 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 ) (iii) ControlZ , which controls latent variables Z learned from the covariates X . Methods from this class include the CEVAE [25] and covariate disentanglement methods, such as DR-CFR [12], TEDVAE [44], NICE [38], and AFS [41].

5.2.1 Results and Analysis

From Table 1, we can find that for all four datasets, CEV AE is worse than the naive LR0 estimator. In addition, for the LatentMediator and Company (Age) dataset, all methods except CiVAE fail to predict the negativity of the ATE. Covariates disentanglement-based methods, i.e., DR-CFR and TEDVAE, inherit the latent post-treatment bias of CEVAE. The reason is that, these methods disentangle latent confounders C from latent instrumental variables I and latent adjusters A by

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

2:6

3:5

4:4

5:3

6:2

Ratio

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Error

(a) Company (Age)

CiVAE

TEDVAE

2:6

3:5

4:4

5:3

6:2

Ratio

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Error

(b) Company (Gender)

CiVAE

TEDVAE

utilizing their causal relations with T and Y , i.e., I is predictive only for T , A is predictive only for Y , whereas C is predictive for both T and Y . For example, TEDVAE includes three encoders to infer three sets of latent variables ˆ I , ˆ A , ˆ C from X and adds classification losses p ( T | ˆ I , ˆ C ) and p ( Y | T, ˆ C , ˆ A ) on the CEVAE loss. However, since both latent confounders C and latent posttreatment variables M are correlated with both T and Y , these methods cannot disentangle C from M . An exception is NICE [38], which uses invariant risk minimization (IRM) [3] to find all causal parents of the outcome Y as the confounders, which makes it more robust in the LatentCorrelator case. However, since mediators M are also the causal parent of Y , the performance degrades substantially on the LatentMediator dataset. Although AFS [41] considers the existence of posttreatment variables M in the proxy X , it assumes that they can be separated from other variables in X in the observational space, and no relationship exists between the post-treatment variables and the outcome, so it still has poor performance in our setting since both assumptions are violated.

5.3 Sensitivity Analysis

In this part, we vary the number of confounders and post-treatment variables that generate proxy X in the Company (Age) and Company (Gender) datasets and compare CiVAE with the baseline TEDVAE in Fig. 3. Fig. 3 shows that the error is consistently lower for CiV AE. In addition, the error is comparatively higher when the number of confounders is low since the misidentification of latent post-treatment variables as confounders can have a com- paratively larger influence on the ATE estimation. In addition, when the number of confounders becomes larger, the performance gap between CiVAE and TEDVAE gracefully shrinks.

6 Conclusions

In this paper, we systematically investigate the latent post-treatment bias in causal inference from 365 observational data. We first prove that unresolved latent post-treatment variables scrambled in the 366 proxy of confounders can arbitrarily bias the ATE estimation. To address the bias, we proposed 367 the Confounder-identifiable VAE (CiVAE), which, utilizing a mild assumption regarding the prior 368 of latent factors, guarantees the identifiability of latent confounders up to bijective transformations. 369 Finally, we show that controlling the latent confounders inferred by CiV AE can provide an unbiased 370 estimation of the ATE. Experiments on both simulated and real-world datasets demonstrate that 371 CiVAE has superior robustness to latent post-treatment bias compared to state-of-the-art methods. 372

References 373

A. Acharya, M. Blackwell, and M. Sen. Explaining causal findings without bias: Detecting and 374 assessing direct effects. American Political Science Review , 110(3):512-529, 2016. 375

A. Anandkumar, R. Ge, D. Hsu, S. M. Kakade, and M. Telgarsky. Tensor decompositions for 376 learning latent variable models. Journal of Machine Learning Research , 15:2773-2832, 2014. 377

M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv 378 preprint arXiv:1907.02893 , 2019. 379

D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. 380 Journal of the American Statistical Association , 112(518):859-877, 2017. 381

L. Cheng, R. Guo, and H. Liu. Causal mediation analysis with hidden confounders. In 382 Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining , 383 pages 113-122, 2022. 384

T. D. Cook, D. T. Campbell, and W. Shadish. Experimental and quasi-experimental designs for 385 generalized causal inference . Houghton Mifflin Boston, MA, 2002. 386

387

388

P. Ding and L. W. Miratrix. To adjust or not to adjust? sensitivity analysis of m-bias and butterfly-bias. Journal of Causal Inference , 3(1):41-57, 2015.

389

390

391

392

393

394

395

J. K. Edwards, S. R. Cole, and D. Westreich. All your data are always missing: incorporating bias due to measurement error into the potential outcomes framework. International Journal of Epidemiology , 44(4):1452-1459, 2015.

F. Elwert and C. Winship. Endogenous selection bias: The problem of conditioning on a collider variable. Annual review of sociology , 40:31-53, 2014.

T. A. Glass, S. N. Goodman, M. A. Hernán, and J. M. Samet. Causal inference in public health. Annual Review of Public Health , 34:61-75, 2013.

N. Hassanpour and R. Greiner. Counterfactual regression with importance sampling weights. 396 In IJCAI , pages 5880-5887, 2019. 397

N. Hassanpour and R. Greiner. Learning disentangled representations for counterfactual 398 regression. In International Conference on Learning Representations , 2020. 399

A. Hyvärinen and P. Dayan. Estimation of non-normalized statistical models by score matching. 400 Journal of Machine Learning Research , 6(4), 2005. 401

K. Imai, L. Keele, and D. Tingley. A general approach to causal mediation analysis. Psycholog402 ical Methods , 15(4):309, 2010. 403

G. W. Imbens and D. B. Rubin. Causal inference in statistics, social, and biomedical sciences . 404 Cambridge University Press, 2015. 405

K. Jager, C. Zoccali, A. Macleod, and F. Dekker. Confounding: what it is and how to deal with 406 it. Kidney international , 73(3):256-260, 2008. 407

F. Johansson, U. Shalit, and D. Sontag. Learning representations for counterfactual inference. 408 In International Conference on Machine Learning , pages 3020-3029, 2016. 409

M. Kalisch and P. Bühlman. Estimating high-dimensional directed acyclic graphs with the 410 pc-algorithm. Journal of Machine Learning Research , 8(3), 2007. 411

I. Khemakhem, D. Kingma, R. Monti, and A. Hyvarinen. Variational autoencoders and nonlinear 412 ICA: A unifying framework. In International Conference on Artificial Intelligence and Statistics , 413 pages 2207-2217. PMLR, 2020. 414

G. King. A hard unsolved problem? post-treatment bias in big social science questions. In 415 Hard Problems in Social Science' Symposium, April , volume 10, 2010. 416

G. King and L. Zeng. The dangers of extreme counterfactuals. Political Analysis , 14(2):131-159, 417 2006. 418

D. Koller and N. Friedman. Probabilistic graphical models: principles and techniques . MIT 419 press, 2009. 420

421

422

M. Kuroki and J. Pearl. Measurement bias and effect restoration in causal inference. Biometrika , 101(2):423-437, 2014.

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

F. Li, K. L. Morgan, and A. M. Zaslavsky. Balancing covariates via propensity score weighting. Journal of the American Statistical Association , 113(521):390-400, 2018.

C. Louizos, U. Shalit, J. M. Mooij, D. Sontag, R. Zemel, and M. Welling. Causal effect inference with deep latent-variable models. Advances in Neural Information Processing Systems , 30, 2017.

C. Lu, Y. Wu, J. M. Hernández-Lobato, and B. Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In International Conference on Learning Representations , 2021.

Y. Lu and J. Lu. A universal approximation theorem of deep neural networks for expressing probability distributions. In Advances in Neural Information Processing Systems , pages 30943105, 2020.

D. Madras, E. Creager, T. Pitassi, and R. Zemel. Fairness through causal awareness: Learning causal latent-variable models for biased data. In Proceedings of the Conference on Fairness, Accountability, and Transparency , pages 349-358, 2019.

W. Miao, Z. Geng, and E. J. Tchetgen Tchetgen. Identifying causal effects with proxy variables of an unmeasured confounder. Biometrika , 105(4):987-993, 2018.

J. M. Montgomery, B. Nyhan, and M. Torres. How conditioning on posttreatment variables can ruin your experiment and what to do about it. American Journal of Political Science , 62(3):760-775, 2018.

J. Pearl. On measurement bias in causal inference. arXiv preprint arXiv:1203.3504 , 2012.

J. Pearl. Conditioning on post-treatment variables. Journal of Causal Inference , 3(1):131-137, 2015.

S. J. Pocock, S. E. Assmann, L. E. Enos, and L. E. Kasten. Subgroup analysis, covariate adjustment and baseline comparisons in clinical trial reporting: current practiceand problems. Statistics in Medicine , 21(19):2917-2930, 2002.

M. Prosperi, Y. Guo, M. Sperrin, J. S. Koopman, J. S. Min, X. He, S. Rich, M. Wang, I. E. Buchan, and J. Bian. Causal inference and counterfactual prediction in machine learning for actionable healthcare. Nature Machine Intelligence , 2(7):369-375, 2020.

K. J. Rothman, S. Greenland, T. L. Lash, et al. Modern epidemiology , volume 3. Wolters 451 Kluwer Health/Lippincott Williams & Wilkins Philadelphia, 2008. 452

U. Shalit, F. D. Johansson, and D. Sontag. Estimating individual treatment effect: generalization 453 bounds and algorithms. In International Conference on Machine Learning , pages 3076-3085, 454 2017. 455

C. Shi, D. Blei, and V. Veitch. Adapting neural networks for the estimation of treatment effects. 456 In Advances in Neural Information Processing Systems , 2019. 457

458

459

C. Shi, V. Veitch, and D. M. Blei. Invariant representation learning for treatment effect estimation. In Uncertainty in Artificial Intelligence , pages 1546-1555. PMLR, 2021.

P. Spirtes, C. N. Glymour, R. Scheines, and D. Heckerman. Causation, prediction, and search . 460 MIT press, 2000. 461

S. Wager and S. Athey. Estimation and inference of heterogeneous treatment effects using 462 random forests. Journal of the American Statistical Association , 113(523):1228-1242, 2018. 463

H. Wang, K. Kuang, H. Chi, L. Yang, M. Geng, W. Huang, and W. Yang. Treatment effect 464 estimation with adjustment feature selection. In Proceedings of the 29th ACM SIGKDD 465 Conference on Knowledge Discovery and Data Mining , pages 2290-2301, 2023. 466

L. Yao, S. Li, Y . Li, M. Huai, J. Gao, and A. Zhang. Representation learning for treatment effect 467 estimation from observational data. In Advances in Neural Information Processing Systems , 468 volume 31, 2018. 469

470

471

K. Zhang, J. Peters, D. Janzing, and B. Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv:1202.3775 , 2012.

W. Zhang, L. Liu, and J. Li. Treatment effect estimation with disentangled latent factors. In 472 Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10923-10930, 473

474

Appendix 475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

A Broader Impact

The proposed CiVAE is a universal model for causal effect estimation with observational data. Although we use the Company job data that estimate the causal effects of online working mode to applicant statistics as a real-world example, proxy-of-confounder-based methods have been heavily used in other observational studies, which may be susceptible to latent post-treatment bias. Therefore, we speculate that the proposed CiVAE will have a broader impact on causal inference community.

B Related Work

B.1 Post-Treatment Bias in Causal Inference

Bias due to accidentally controlling post-treatment variables, i.e., post-treatment bias , has long been recognized as dangerous in causal effect estimation [20]. Back at 2005, Pearl [32] cautioned that controlling more is not better, and uses the collider bias [9] and M-Bias [7] as two examples to show that bias can be increased when controlling the post-treatment variables. Furthermore, [30] show that indirect correlations between post-treatment variable M and outcome Y can still cause bias. Recent works prove that even if M has no causal relationship with Y , controlling it can still increase the variance of estimand [12]. However, most of these works study the post-treatment bias in the observational space, where latent post-treatment variables that are mixed with confounders to generate the observed covariates can be easily ignored by the researcher. Therefore, it motivates us to develop CiVAE, which is robust to the latent post-treatment bias under mild assumptions.

B.2 Covariate Disentanglement

Recently, researchers have realized that directly controlling proxy of confounders X may not be safe, as variables other than confounders could lurk in the proxy and ruin the ATE estimation [12]. Traditional methods assume that the variables that generate X are a mixture of confounders, adjusters, and influencers [36], where adjusters should not be controlled as it can increase the estimation variance [11]. Most methods rely on the fact that adjusters are correlated only with the treatment to separate them from other variables [12, 44] (see Fig. (1)). This can also be used to remove posttreatment variables that are not correlated with the outcome, which have similar statistics properties with adjustors [41]. Here, a different work is NICE [38], which uses the fact that confounders and influencers are direct causal parents of the outcome to find these variables with invariant learning as the control set [3]. However, since mediators are also direct parents of the outcome, NICE is still not robust to general post-treatment bias. Given that all above methods cannot satisfactorily address the latent post-treatment in general cases, it is imperative to design the CiV AE, where confounders can be identified and distinguished with latent post-treatment variables for unbiased adjustment.

C Theoretical Analysis

C.1 Proof of Lemma 3.1.

Proof. Let Z = f ( X ) and z = f ( x ) . If f is injective and differentiable a.e. , and f † is the left-inverse, we have:

where f · and f ·|· represent the marginal and conditional density function, respectively, and J f † ( z ) is the Jacobian matrix of function f † evaluated at z . Based on Eq. (12), we have:

C.2 Proof of Corollary 3.3. 515

Proof. For X = x , let [ c || m ] . = [ f † C ( x ) || f † M ( x )] . = f † ( x ) = A † ( x -α X ) , where A † is the left 516 inverse of the full column-rank matrix A in Eq. (2), we have: 517

where the first equality is due to the definition of CATE in Eq. (2). In addition, the causal estimand 518 and bias of a proxy-of-confounder-based causal inference model that controls the latent variable Z 519 inferred via Z = ¯ f ( X ) = B T X (where B is also a full column-rank matrix) can be formulated as: 520

521

522

523

524

525

526

where 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 3.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).

C.3 Another Case of Linear SCM with Latent Correlators

Corollary C.1. For another Linear Structural Causal Model defined as follows

where f = A ∈ R K X × ( K C + K M ) is a full column-rank matrix, the CATE, ATE, and the bias of 527 proxy-of-confounder-based causal inference model that controls the latent variable Z inferred via 528 Z = ¯ f ( X ) = B T X can be formulated as follows: 529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

where B ∈ R K X × ( K C + K M ) is another full column-rank matrix. Since ∑ θ j · γ j ϕ j is arbitrary, the estimator 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 Z inferred via Z = ¯ T

where step (a) and the rest of the proof follow the same logic as the proof in Section 3.3.

C.4 Proof of Theorem 4.1

The strict definitions of the exponential family, strong exponential (which is assumed for the factorized part of the conditional prior), and identifiability follow [19, 26], and can be referred to in Appendix E, F of [26], which we omit to avoid redundancy. The proof of Theorem 4.1 is largely based on the NF-iVAE paper [26], 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.

C.4.1 PART I

Step 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.

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

Step (a) is based on the rule of change-of-variable, where vol( A ) = √ det ( A T A ) . In step (b), we define ˜ p f, S , λ ,Y,T ( X ) ≜ p S , λ ( f † ( X ) | Y, T ) vol ( J f † ( X ) ) I X ( X ) . In step (c), we use F [ · ] to denote the Fourier transform. In step (d), we drop φ ε ( ω ) as it is non-zero a.e. (see Assumption 3).

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. (19), we have:

Let ( 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. (20) at these points, where the first equation is subtracted from the remaining ones, which leads to the following equation system:

Let 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:

where A = 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. (22) 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 [26].

C.4.2 PART II

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. (22) can be rewritten into:

We aim to prove that A in Eq. (23) 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. (23) with respect to Z s and Z t , where s = t , we have:

Note that for the factorized part of the sufficient statistics S , i.e., S f , all cross-derivatives are zero, 569 and for the non-factorized part of S , i.e., S nf , which is a neural network with ReLU activation (i.e., 570 linear a.e. ), all second-order derivatives are zero. Therefore, the second order cross-derivatives on 571 the LHS. of Eq. (24) are zero, which leads to the following equality: 572

Eq. (25) can be written into the matrix-vector product form as follows: 573

576

577

578

579

580

581

582

583

584

585

where and

If we denote the concatenation as ˜ S ′′′ ( Z ) = [ ˜ S ′′ ( Z ) , ˜ S ′ ( Z ) ] ∈ R k × 2 K Z and v ′′ s,t ( Z ) = 574 [ v ′ s,t ( Z ) T , v ′′ s,t ( Z ) T ] T ∈ R 2 K z , we have: 575

Finally, 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 [19] and the assumption that k ≥ 2 K Z , we have that the rank of ˜ S ′′′ f ( Z ) is 2 K Z . Since k ≥ 2 K Z , the rank of ˜ S ′′′ f ( Z ) is also 2 K Z . Since the rank of A is k , the rank of A ˜ S ′′′ ( Z ) is 2 K Z , which implies that v ′′′ s,t ( Z ) ∈ R 2 K 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 v -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. (23) can be reformulated as S ( Z ) = A S ( v ( Z )) . We then apply v -1 to Z on both sides, which gives

̸

Let t be the index of an entry in S that corresponds to the factorized part S f . For all s = t , we have: 586

̸

Since the entries of ˜ S are linearly independent, a tj is zero for any j such that ∂ S ( Z ) j ∂Z s = 0 . This 587 includes the entries S j that correspond to (1) the factorized part that does not depend on Z t ; and (2) 588 the non-factorized part S nf . Therefore, when t is the index of an entry in the sufficient statistics S 589 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 590 map between S f,i ( Z i ) and S f,i ( v i ( Z i )) . Therefore, we can construct an invertible submatrix A ′ i 591 with all non-zero elements a tj for all t that corresponds to factor i , such that 592

593

594

595

596

where c i denotes the corresponding elements of c . Eq. (30) 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 ˜ S f,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.

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

C.4.3 PART III

Let ˜ Z i = v i ( Z i ) = ˜ f † ( X ) i be the i th 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 ˜ Z i is the bijective transformation of the true latent variable. The proof is as follows.

Proof. Plugging ˜ Z i into Eq. (30), we have:

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. (31), which gives:

Since 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 ˜ S f,i is injective, ¯ S f,i is also injective, and therefore the composite map S † f,i A ′ i ¯ S f,i : R → R that applies on ˜ Z i is a bijective. This completes the proof that each inferred latent variable ˜ Z i 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.

C.4.4 Consistency

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. (9) (with the KL term replaced by the score matching loss defined in Eq. (10)) over its parameter ϕ , the score matching term will eventually vanish. Therefore, the ELBO term in Eq. (9) 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.

C.5 Proof of Theorem 4.2

Proof. Let C be the true latent confounders and ˜ C be the transformed confounders, where the transformation function f is bijective and differentiable a.e. Let f -1 denote its inverse. The ATE estimator that controls transformed confounders ˜ C can be formulated as:

Specifically, for the continuous case where density functions exist, for each term, we have:

For the marginal density f ˜ C (˜ c ) , the following equality holds:

As for the conditional density f Y | T, ˜ C ( y | t, ˜ c ) , since f is bijective, according to Eq. (12), we have:

Combining Eqs. (35) and (36), and given that d ˜ c = | J f ( c ) | d c , we have: 627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

where the term | J f -1 (˜ c ) | · | J f ( c ) | vanishes in step (a) as the two factors have the product of one. Therefore, if we plug Eq. (37) into Eq. (33), it leads to the following equality:

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.

D Extending CiVAE to address Latent Interactions

In this section, we extend CiVAE to more general cases where interactions exist among the latent confounders C and the latent post-treatment variables M . Here, we note that the identification of latent confounders C in CiVAE is achieved in two steps. (i) CiVAE individually identifies latent variables [ C , M ] that generate X in inferred Z (but which dims of Z correspond to C or M is unknown). (ii) pairwise independence test to identify C . Since Assumption 2 allows arbitrary dependence among C and M , step (i) still holds when interactions among [ C , M ] exist. To distinguish C in these cases, we can use more general causal discovery algorithms, e.g., the PC algorithm [18] in the second step. In this section, we consider two cases of interaction: (i) Intra-Interaction among mediators, and (ii) Inter-Interaction among mediators and confounders.

D.1 Intra-Interactions among Latent Mediators

In this subsection, we discuss the case where latent post-treatment variables M interact with each other. Since in this case, M cannot causally influence the latent confounders C (otherwise C will be post-treatment), and the PC algorithm orients edges in causal graphs via colliders, latent confounders can still be identified from the inferred Z as they form colliders with the treatment T .

To empirically verify the claim, we extend the simulated datasets described in Section 5.1, where we make (i) T directly affects M 1 , (ii) M 1 affects M 2 , and (iii) M 1 , M 2 affect M 3 . The coefficients are randomly sampled from N (0 , 1 / 3) . In step (ii) , we use the PC algorithm [18] to identify C from the inferred Z . The results in Table 2 demonstrate that the adapted CiV AE is still significantly more robust to latent post-treatment bias compared to CEV AE and TEDVAE, which empirically verify our claim that PC-adapted CiVAE can address the interaction among post-treatment variables.

D.2 Inter-Interactions between Latent Mediators and Latent Confounders

In this subsection, we discuss another case where inter-interactions exist between latent confounders C and latent post-treatment variables M . Since in this case, M still cannot causally influence C (otherwise C will be post-treatment), and the PC algorithm orients edges in causal graph via colliders, latent confounders C can still be identified from Z as they form colliders with the treatment T .

To verify the claim, we extend the simulated datasets described in Section 5.1 to allow each latent 658

confounder C i ∈ R 3 to determine M ∈ R 3 . The coefficients are randomly sampled from N (0 , 1 / 3) . 659

In step (ii) , we use the PC algorithm to identify C from the inferred Z . The results in Table 3 660

demonstrate that the PC-adapted CiVAE is still significantly more robust to latent post-treatment bias 661

compared to CEVAE and TEDVAE, which empirically verify our claim that PC-adapted CiVAE can 662

address the case where inter-interactions exist among latent confounders and post-treatment variables. 663

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The contribution of this paper can be summarized as: We study a critical but easily overlooked problem in causal effect estimation: latent post-treatment bias, and we propose a novel framework, i.e., CiVAE, to address the bias. The details are in Section 4.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have discussed the potential issue of the vanilla when interactions among the latent variables exists. However, in Section D we have addressed the issue by extendeding our framework.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We have introduced the three mild assumptions required for the identification of causal effects under latent post-treatment bias. In addition, we have provided the proof for all the theorems in the Appendix.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We have provided implementation details in Section 5.1. In addition, we have provided a code demo in an anonymous URL.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have reported the error of five independent run for both the proposed CiVAE and all the baselines in the main paper.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the com711 puter resources (type of compute workers, memory, time of execution) needed to reproduce 712 the experiments? 713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

Answer: [Yes]

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have carefully read the code of ethics and behaved strictly according to it.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: See Section A of the Appendix.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our model does not have a high risk for misuse.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited the papers of our baselines and honor their license of code.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide the Readme file along side the codes.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No human subjects are involved in our experiments.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] 755

Justification: No human subjects are involved in our experiments. 756

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n• Causal reasoning in the context of latent confounder and post-treatment variables is an important topic especially with observational data.\n\n• The authors clearly state the necessary assumptions for the identifiability of true latent variables, and the logic of determining the dimensions of $\\boldsymbol{C}$ and $\\boldsymbol{M}$ is well presented.\n\n• The paper has a well-established theoretical basis.\n\nWEAKNESSES:\n•\tFor the illustrative example in the introduction, it might be better to explicitly specify what the post-treatment variable is.\n\n•\tOther existing works [1-3] on identifying latent confounder/mediators based on the iVAE architecture should also be included in the related work.\n\n•\tThe 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$?\n\n•\tThe 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.\n\n•\tThe empirical evaluation consists of only one real-world dataset, which somehow limits the applicability of the proposed method.\n\nReferences:\n\n[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.\n\n[2]. Sorrenson, P., Rother, C., & Köthe, U. (2020). Disentanglement by nonlinear ica with general incompressible-flow networks (gin). arXiv preprint arXiv:2001.04872.\n\n[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.\n\nQUESTIONS:\n• 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.\n\n• 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?\n\n• 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?\n\nLIMITATIONS:\nThe authors do not include a paragraph discussing the limitations and potential societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-16"}

comments: 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.

STRENGTHS:
The paper deals with a very interesting problem.	 The proposed method appeared to be theoretically robust. The method is evaluated with proper experimental analysis on synthetic and real-world datasets and compared with multiple benchmarks.

WEAKNESSES:
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:
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:
The authors discussed a very few limitations of their paper but more discussion should be done.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* The paper investigates a critical question concerning the mitigation of post-treatment bias, which is essential in various practical scenarios.\n* The ideas presented in the paper are clear and easy to follow, and the theoretical analysis is well-established.\n\nWEAKNESSES:\n* 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.\n\n* 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.\n\n* The experiment lacks sufficient details on setup and implementation. Could the authors provide more specific information to enhance understanding of the empirical results?\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\n* The proposed method relies on very strong assumptions to ensure identifiability, which can be challenging to verify in practical applications.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nBeing able to recover latent variables (cofounders, post-treatment variables, or others) from observations is challenging and important. Ignoring latent variables or assuming non-existence of latent variables is unrealistic and can lead to the wrong conclusion and decisions. The authors further motivated the importance of recovering latent cofounders, post-treatment variables and the consequence of not doing so  (Theorem 3.2). The solution provided shows originality and quality.\n\nWEAKNESSES:\nThe presentation can be improved.\n\nQUESTIONS:\nIs 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:\n1. 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.\n2. 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.\n3. Latent Mediator Variables: These latent variables mediate the effect of the treatment (T) on the outcome (Y) and are not directly observed.\n4. 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.\n\nCan the proposed method handle these types too (with some extension), or some of the types are quite disruptive to the proposed methodology?\n\nLIMITATIONS:\nn.a.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}