Identifying and Estimating Causal Effects under Weak Overlap by Generative Prognostic Model
Authors: Pengzhou Abel Wu
Kenji Fukumizu
Venue: NeurIPS 2021 Submitted (2021)
Keywords: VAE
variational autoencoder
balanced representation Learning
treatment effects
causal inference
identifiability
identification
CATE
ATE
weak overlap
limited overlap
Prognostic Model
Prognostic score

Abstract: As an important problem of causal inference, we discuss the identification and estimation of treatment effects (TEs) under weak overlap, i.e., subjects with certain features all belong to a single treatment group. We use a latent variable to model a prognostic score (PGS), which is widely used in biostatistics and sufficient for TEs, i.e., we build a generative prognostic model. We prove that the latent variable recovers a PGS, and the model identifies individualized treatment effects. The model is then learned as the  Intact-VAE, a new type of variational autoencoder (VAE). We derive counterfactual generalization bounds which motivate representation balanced for treatment groups conditioned on individualized features. The proposed method is compared with recent methods using  (semi-)synthetic datasets. 

Identifying and Estimating Causal Effects under Weak Overlap by Generative Prognostic Model
Anonymous Author(s)
Affiliation Address
email
Abstract
As an important problem of causal inference, we discuss the identification and estimation of treatment effects (TEs) under weak overlap, i.e., subjects with certain features all belong to a single treatment group. We use latent variable model prognostic score (PGS), which is widely used in biostatistics and sufficient for TEs, i... we build generative prognostic model. We prove that the latent variable recovers PGS, and the model identifies individualized treatment effects. The model is then learned as the Intact- VAE, new type of variational autoencoder (VAE). We derive counterfactual generalization bounds which motivate represen- tation balanced for treatment groups conditioned on individualized features. The proposed method is compared with recent methods using (semi- myytttiyyt datasets.
3
5
7
Introduction
i.e, inferring causal effects of interventions, a fundamental problem. In 13. this work, we focus treatment effects (TEs), such as effects of public policies new drug, based 14 on a set of observations consisting of binary labels for treatment control (non-treated), outcome, and other covariates (e.g, patients' personal records). The fundamental difficulty of causal inference is
12 Causal inference
16 that we never observe counterfactual outcomes, which would have been if we had made the other decision (treatment or control). While the ideal protocol for causal inference is randomized controlled
18 trials (RCTs), they often have ethical and practical issues, or suffer from prohibitive costs. Thus, causal inference from observational data is indispensable. It introduces other challenges, however. The most crucial one is confounding: there may be variables (called confounders) that causally affect both the treatment and the outcome, and spurious correlation follows.
20
21
22 large majority of works, including this work, rely on the unconfoundedness, which means that appropriate covariates are collected so that the confounding can be controlled by conditioning on covariates. That is, all the confounders are essence observed. This is still challenging, due to systematic imbalance (difference) of the distributions of the covariates between the treatment and control groups, introducing bias in estimation. Among classical ways of dealing with imbalance Machine learning methods also exploited; there non non-paraen parametrie, tree-based methods, c.g.,
23
24
26
matching and re- re-weighting 44 35
28 semi parametric methods, c.g. 48
TMLE), which have better finite sample performance, and also
49 Causal Forests (CF)). Notably, starting from (23), there a recent rise of interest in learning representation of covariates, which independent of treatment
30
31 groups, i.e., balanced representation learning (BRL).
32 The most serious form of imbalance that sample points with certain values of covariate are all belong to a single treatment group, which is called weak overlap of the covariate. Causal effects are 34 not directly estimable non-overlapped covariate values. There are lines work that give robustness to weak overlap 3 trim non-overlapped sample points [52), or study convergence rate depending
36 on overlap [9]. Weak overlap is particularly relevant to machine learning methods exploiting rich covariates, because, with digeer-dimensional covariates, overlap is harder to satisfy and verify [10].
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.

---

38
Our approach to the weak overlap issue is based on the prognostic score (PGS) [4] which is among the important concepts of sufficient scores. While the most well-known score is the propensity score (PPS) [36). PGSs have also long been known to improve methods using PPS 37 S. and interests last 40
39
41 in biostatistics 44 2. Prognostic modeling can benefit more from predictive systems and exploit
42 richer literature than propensity modeling, particularly in Medicine and Health. A comparative study
43 in [3] shows PGS- methods perform better, or as well as, PPS methods. Thus, it is promising
44
to combine the predictive powers of prognostic modeling and machine learning.
45
To solve the inverse problem of recovering PGSs, our method exploits also the recent advance of identifiable representation, particularly of VAE [26] iVAEJ. Identification means parameters of interest (for us, representation function and causal effects) are uniquely determined and given by true observational distribution. Identification logically precedes estimation and inference. Without identification there is no hope of a consistent estimator, and a model would fail silently; may fit perfectly but return an estimator that converges to the wrong one or does not converge (29) particularly Sec. 8). Identification is even more important for causal inference, because, unlike usual (non-causal) model misspecification, causal assumptions are often unverifiable through observables [50]. Thus, it
48
47
48
49
50
51
53
is critical to specify theoretical conditions for identification, and then the applicability of methods can be judged by knowledge of an application domain.
54
55 In this work, we study identification (Sec. 3) and estimation (Sec. of TEs under weak overlap. We particularly discuss individualized treatment effects, conditioned on the covariates. Code and proofs are in Supplementary Materials. The main contributions of this paper are:
58
57
58
1) theory of TE identification under weak overlap of covariates, using PGS and identifiable model;
59
2) counterfactual generalization bounds on TE error, which motivates our conditional BRL;
3) new regularized VAE to estimate TEs, with connections to identification and balancing;
61 4) experimental comparison to state-of-the-ee methods on (semi- ceni-)syythheicc datasets.
2
Setup and motivation
62
68
2.1
Counterfactuals, treatment effects, and identification
64 Following (21]. we introduce potential outcomes (POs, or counterfactual outcomes) y(t) E Rd Rdd t {0, 1)- (() is the outcome that would have been observed, if treatment value t t had been applied. Formally, this is the consistency of counterfactuals: y = y(t) i t, or simply = y().. We see t. The fundamental problem of causal inference is that, for a unit under research, we could observe only one of y y(0) or (1). corresponding the treatment value applied. That is, "factual" refers to y or t that is in principle observable in data, or statistical entities (e.g, estimators) built on them. We also observe relevant
66
67
y y!) as the hidden variables that give factual y under factual assignment
68
69
70
71 covariate(s) x E X C RTH. which is associated with individuals, with distribution D ~ p(x,y, t).
72 Note, we use Roman fonts for random variables (c.g., t) and italic for realization (c.g., t).
x = x)), conditioned on 74 this work are the Conditional Average TE (CATE) and Average TE (ATE), defined respectively by
The expected PO is denoted by (xx) (t)(
The estimands in
T xx) M1 (a:) po(a)) V = B(T()))
75
CATE is an individual-level, personalized, treatment effect, given highly discriminative covariate.
76
Standard results (38 [6] Ch. 3] give sufficient conditions for identification under general setting. They are Exchangeability: [()__x and Overlap: p(1|æc) > O for any x E X'. Both are required for t t {0. 13. When t appears in statement without quantification, we always mean "for both t". Often, Consistency is also listed, but, as above, it is better known as well-definedness of counterfactuals. Exchangeability means, just as in RCTs but additionally given x, that there no correlation between factual treatment and counterfactual outcomes y(t). Overlap means that the supports of p(x = 0) and p(x]t= = 1) should be the same, and this ensures it is valid to condition on any (ax, t).
77
78
79
80
81
82
83 We relax overlapped covariate in Sec. 3.2 to allow some non-overlapped values x, i.e., covariate x is
84 weakly overlapped. In Sec. 2.2 we introduce a condition which gives exchangeability (and PGSs). In
85 this paper, we also discuss overlap of variables other than x (e.g. PGSs), and we give a definition for
86 any random variable with support V as following.
87 Definition 1. Overlap of means p() v) > 0 for all I. E (0, 1), u E V. We also say v is overlapped. If the condition is violated at some value v, then is non-overlapped and v is weakly overlapped.

---

89
2.2
Prognostic scores
90
Our method is motivated by PGSs [14), adapted as Pt-score and P-score in Definition in this paper, related to balancing scores b(x), which is defined by tilx ((x) [36.. The PPS p(t [xx) is special case of this. Both are sufficient scores for identification; PGSs are sufficient statistics of outcome predictors and b(x) is for the treatment (see Appendix for details).
92
94
Definition& 2. A Pt-score (P(S) is two functions Prr (x) (t 0, 1) such that y (x([xxx A PIS is called P-score (PS) if Pu P1.
95
Note that, a PtS is by definition two functions, thus overlapped (x(x) means both (x( and (x( 96
overlapped. Why PtS (PGS)? PtS is more applicable than balancing score b(x) under weak overlap. Overlapped b(x) implies overlapped x, which in turn implies overlapped PGS [[0]. Lower- dimensional than x, PtS is likely more overlapped than and, moreover, there is evidence that PIS maximizes overlap among all sufficient scores for ATE
96
100
Below is direct corollary of Proposition 5 in [14]. Both of PtS and PS give CATE, but, as we will see, PS better as a conditionally balanced representation, since P, (x ||| only when Fo Pr.
102
Proposition 1 (CATE by PtS). P, is PtS, then CATE can be given by
(x) = (t) x))
t)) p(ulP. = Px())) t)ydy
(2)
With the knowledge of P,, we choose of Po, Pı corresponding to the counterfactual outcome of interest. This ability of counterfactual assignment resolves the problem of non-overlap at æ.
105
PiSs exist under general settings when y() follows an additive noise model (ANM).
107
(G1) (ANM) the data generating process (DGP) for y y (M((x,, t) + e where f* M are functions and e is zero-mean exogenous (external) noise.
108
The DGP defines (t) by setting = t in the equation. And it also specifies how other variables causally affect y. For example, x affects y through M. so M(xx is the effect modifier [14), which is
110
often components of x affecting, y directly. Note (G1) also implies exchangeability given x, through ((t) LIMM(x) ANMs are also commonly used in nonparametric regression methods for TEs [6].
112
Under (G1) 1) P. == ft (M(()) (x Pist not PS, 2) M is a PS (x is trivial PS), and 3) PP (No (x), ini (x)) is PS. We use the same symbol to denote PtS and the random variable
114
defined by it, when appropriate.
3 Identification under generative prognostic model
116
introduce our generative prognostic model and VAE based on ply, z|x,t) t) and prove identifiability of our model. In Sec. 3.2 we prove identification of CATEs, one of our main contributions. The theoretical analysis involves only our generative model (i.e., prior and decoder), but property of model, and
In Sec.
*(r/s,e)
118
y
119
pinie)
120
121
not encoder, because model identifiability causal identification is about DGP and model. The encoder is involved in estimation which is studied in Sec.
p(eix)
p(yle)
122
123
p/yle.3)
3.1
Intact-VANE model, architecture, and identifiability
x(xx,t) W
125
Generative models are useful to solve the inverse problem of recovering Pt-score. Our goal is to build model that can be learned by VAE from observational data to obtain a PtS, or more ideally PS, via the latent variable That is, a generative prognostic model.
Figure I: Graphical models of the decoders. From top: CVAE, iVAE, and Intact- VAE. The encoders are sim- ilar, taking all observables and build approximate pos- teriors, and thus are omitted. The first factor is our decoder which models p(v[P,, t) in 22) and the second factor is our conditional prior which models (x). Conditioning on x in the joint model
126
127
128
129
With the above goal, the generative model of our VAE is built as
ply,z|x.1 t) = p(y|z, t)p(z) x, t).
(3)
130
131
I We often write t function argument subscript, which indicates possible counterfactual assignment.
My is the most common PGS, to the extent that some call it the PGS (e.g. [39 g 47 even without ANMs.
3

---

132 ply, z|x, t) reflects that our estimand is CATE given x. Modeling the score by a conditional distri- bution rather than deterministic function is more flexible. We parameterize our model by ANM outcome and factorized Gaussian prior as
133
134
Ps (y (2. - - PA()xx) ~ (z; (xx),,,, diag((((x)))
-
(4)
where e = (f. h. k) are functional parameters and is noise. X == diag k (x)) -2)
135
136
is the natural parameter of the Gaussian prior, and we also use it as a shorthand for both h, k.
137
The ELBO of our model can be derived from standard variational lower bound togp(y/x; x, t) 2 p(y x, t) DKL (g(z)x. y, ()llp())x, Bgz log p(y|=,t) - (g x, t)llp(z)x,)) t)).
(5)
138
Our encoder g. which conditions on all the observables, standard, and we will see its importance later. We name this architecture Intact- VAE (/dentifiable greatmnntcooddiionn
139
VAE).
140
We naturally have identifiable conditional VAE (CVAE), as the name suggests. Note that (3) has similar factorization with the generative model of iAAE [26]. that is p(y, 2 = p(y 2 p))(z the in both the factors of 3) our VAE architecture is a combination of iVAE and CVAE [42 28). with as the conditioning variable. See Figure [ for the comparison in terms of graphical models. The core idea of iVAE is reflected in our model identifiability (Lemmal below). See Appendix for the basics of VAEs.
141
142
first factor does not depend on Further, since have the conditioning
143
144
145
148
The following conditions on the model are used in theoretical analysis.
(M1) i) ft injective. ii) fo is differentiable, and iii) n == dim(z) = d(= dim(y)).
147
148
Lemma (Model identifiability). Given model and under (MM)errMMi)i)) and ii), for t = t, assume
149
(D1) (Linear independence of A) there exist 2n+11 points vo, ......... D2n E X such that the 2n-square matrix L, : [7t,1, ... Yt,2n] is invertible, where Yt := (ass) (2
150
151
Then, given t = l, the family is identifiable up to an equivalence class. That is, if pe (y|x, 1/xx -
152
Pa (y x, t). we have the relation between parameters: for any yt in the image of fü.
f = diag(a, (a) f(( (y:) + b =: A(f
(6)
153
where diag(a) is an invertible n n-dlaagonal malrix and b is n-vector, both depend on X
154
The conditions are inherited from iVAE. (D1) holds easily in practice, if the components of (x)
155
are linearly independent; if (DI) fails, then the support of >x is in (2n 1)-dimensional space.
156
The essence of the result is f! = fi o Ar, that is, f. can be identified (learned) up to an affine transformation defined by X This is achieved by combining the techniques from (26 and [43] and
157
158
essentially the same results can be proved for other exponential family priors [4 In this paper, symbol (prime) always indicates another parameter (variable, etc.).
159
160
3.2
Nonparametric identifications under weakly-overlapped covariate
161
In this subsection, we give two identification results based on (partial) recovery of PS or PtS. respectively. Since Ptss are functions of x, the recovery is achieved by a noiseless prior, that is. (xx) = 0; the prior 2x,t ~ Px (z|x, = t) degenerates to deterministic function (x).
162
163
164
PiSs with dimensionality lower than equal to y are essential work for weak overlap of x, the extent that, from now on, we simply say PSs / PiSs when referring to this kind of low-dimensional PSs/ / PtSs, unless particularly indicated. (MI) jii). i.e. n d, is not restrictive because μ4, a PtS of the same dimension as y under (GI) Also, in practice, n = d means that we seek a low-dimensional representation of x. In fact, to make the dimensionality explicit in (GT) we introduce an alternative (G1') which includes (G1) with P.- = fe and ji is identity.
185
166
187
168
189
(G1') (Low-dimensional PtS) Under (ader(G1) (G1)
(x(x) (x), for some P. and injective jr-
171
We use GG)) afterwards. Clearly, P, in (G1') is a PtS, and injectivity and n d ensure n dim(y) (P.).. Similarly, the next (G2) reduces unverifiable 2 dim((P) to n = d, for PS.
172
173
(G2) (Low-dimensional PS) Under (11)
(x(x) = 1 (F(x)) for some IP and injective Jr.
174
means that CATES are given by nio and an invertible function é := ÿ1 jo See Appendix for more discussion and a (closely related) real world example. In Sec. 4.1 we argue that there often
G2
175
exist equivalent PSs under (GI)) at least approximately.

---

177
With (G1') or (G2) overlapped x can be relaxed to overlapped P(tSS plus the following.
(M2) (Score partition preserving) For any x; x' E X, (x = Px('')
he(2)) h. h.(x")
178
179
Note that (M2) is in fact required for optimal h, in the sense specified in Proposition aand Theorem | below. The intuition is that, IP, maps non-overlapped to an overlapped value, and h; preserves this property, through learning. In fact, (M2) is trivially satisfied if P, and h2 are linear, and this is still challenging and considered by many works [32 9].
181
182
or some with linear outcome models [1 (3].
183
Our first identification, Proposition 2 relies on (G2) and our generative model, without model identifiability (so differentiable f: is not needed). This is a nonparametric i iidetificcation under shape
184
185
restriction 77. because f. h are functional parameters, and injectivity is monotonicity if je is on R.
186
Proposition 2 (Identification with PS). Given [G2) and model 3 and (4) under (MT) i) and ii). and (M3) (PS matching) ho (x) = ha (x) and k x 0. Then, if EEp (y x. E6y x, U). we have
187
188
1) (Recovery of PS) xx.: h,(x)) (P())) on overlapped x.
189
where v : + 12 is an injective function and == {P(z) [e(=))ooerlpppedd x}
190
2) (Identification) if I" in [G2) is overlapped, and (M2) is satisfied, then μ4 (z(x)) = fie
for any E (0, }), a E X, where (x)) == (x p, (y/z, = ft (1)
191
192
The essence is, i) the true DGP is identified up to an invertible mapping v, so that f. = jr v and h = v Pr, and ii) P. is recovered up to v and yt) [t) [PP,i is preserved, with same vo for both t.
193
194
PS is preferred since it satisfies overlap more easily and (M2) than PtS which refers to two functions.
195
However, the existence of diwddmeniional PS is uncertain in practice when our knowledge of the
196
DGP is limited. Thus, we need Theorem to work under PtS which generally exists.
197
Theorem (Identification with PIS). Given the DGP (GT) and model 3&44)) under M11 and
(M3') (Noise matching) Pe- Pe and k(x) kk' (x), k + O, assume (DI)
and
199
(D2) (Spontaneous balance) There exist 2n + 1 points Zy, ..., I2n X, 2n-square matrix C. and 2n-vector d, such that L L C and Bn C-TB: d/k for optimal X (see
200
below), where L, is defined in (D1) B, = (a,( or (x ..., (x(2))) a=(x)))" and
202
(x; A²) is the log-partition function of the prior in (4)
203
Then, if po ( t) = p(y|x, t), conclusions 1) and 2) in Proposition 2 hold with P replaced with F, in (GI') and the domain of v becomes P == Po P1, P. = (F)(2) overlapped æ).
204
205
Theorem
also achieves the two essential points, but in different and complementary ways. Proposi- tion 2 starts from the prior by Po = IP, and setting ho h1 Conversely, Theorem starts from the decoder with Pe- = Pe and strengthens model identifiability (6) by (D2) (D2) restricts the discrepancy between Au \ on 2n +1 1 points of x, thus is relatively easy to satisly with high-dimensional x.
207
209
We see more reasons to prefer PS. In general, to identify the mean function (x).. a regression is enough, and Pe Pe is unnecessary as in Proposition 2 Also, (D2) is trivial if we have PS and set An AI. See Appendix for more on the complementarity between the two identifications.
210
211
Estimation by 3-Intact- VAE
212
213
4.1 Prior as PS, posterior as PtS, and as regularization strength
214
In this subsection, we discuss our focus on balanced PtSs and give an estimation method to realize it.
215
In learning the Intact VAE with data, we assume that there is PtS and the decomposition of (G1') holds. Such decomposition is not unique in general, however. Among possible PtSs, we wish to learn a balanced PtS. which is close to PS. This is based on the observations in Sec. 3.2 we saw that existence of PS is preferable in identifying the true DGP up to equivalent expression. Here, we introduce the notion of balanced PtS in a non-rigorous way: a PtS P, called balanced if the value of some measure for the conditional independence P. (x) | x is small. The idea is common in practice. For example, in a real-world nutrition study, 20 reduces 11 covariates to a 1-dimensional línear PS.
216
217
218
219
220
221
"Soe references (e.g. [6] only refer to identification without models as "nonparametric"" However, with recent advances, we have identification under nonparametric models [29]. also the case in this paper.
"Same results hold without "under (G1) in (G2) (or (G1" for Theorem I). except that zx, is not necessarily a PS (or PtS for Theorem I). This is because (GI) is required ensure μt is PtS.
5

---

222
Assuming that there is balanced PtS. we consider two ways for estimating it with Intact VAE. One
223
to exploit a prior that does not depend on t. Namely, we set Ån = X =: A in Q The other is
224
to introduce a hyperparameter / in the ELBO as in 3-VAR [17. More specifically, with the prior
225
PA (z)x), we use factorized Gaussian for the decoder and encoder:
(y,g t) ~ (y; fz),,, (z))( (2)2)xx
~ Wi mi(, y), (xi,g((((x,)))) (7)
226
The modified ELBO with B. up to additive constant, is derived as
- Exwg log (z)]] (8) Here, for convenience, we omit the summation (also in L.y in Sec. 4.2), as if y was univariate. The approximate posterior (or encoder) (z(z)x, .tt) depends on t, which can realize a PtS. With B, we control the trade-off between the first and second term: the former is the divergence of the posterior from the balanced prior, and the latter is the reconstruction of the outcome. By choosing 3 in an
E (1s x, y, ()),))))) p)
Ex-,ply f((z))
227
228
229
230
231
appropriate way, such as by validation, the ELBO can find a solution that explains the outcome while
232
keeping the balancedness of the posterior. Also note that, the parameters g and kc, which models the outcome noise and expresses uncertainty of the prior, respectively, are both learned by the ELBO.
233
234
This deviates from the theoretical conditions in Sec. 3.2] but more practical and gives better results in experiments. See Appendix for much more on ideas and connections behind the ELBO.
235
236
Once the encoder qs is learned the estimate of the expected POs is given by
(9) where (z|xx) == 1p (z| (z) y, t) is the aggregated posterior and D p(x. y, t). In estimation, consider the case where is observed in the data, and the sample of (y, .) are taken from the data given x = x (when x is not in the data, we replace ds with px in (9), see Appendix for details and results). Note that in ( indicates counterfactual assignment which may not be the same as factual = / in the data. That is, we set = i in the decoder. The assignment not applied to the encoder, but it is learned from factual x, y. (see also Sec. 4.2 the explanation for ECF.)). The overall
(x(x)) = Ey(aix m) fz(z)
E2f (z) y, E {0.1),
239
240
241
242
243
algorithm steps are i) we train VAE by (8) ii) infer CATE (x) = fi (x) pi (z) by
4.2
Conditional balanced representation learning
We formally justify our ELBO 8 from the viewpoint of BRL. Usually, particularly in ATE estimation, 245
balance means covariate balance, i.e., x 1t 44. Influenced by this, most BRL methods learn balanced
247
covariate representation z such that z. ll 44 [31 and usually z is a function of x. From Sec. 4.1 we
248
understand that larger B in ELBO (8) encourages z.Lx(x which is given by the prior, corresponding
249
to [t)xLt|x for a balanced PtS. Here, we show that, this conditional balance of representation z
250
is natural for CATE estimation, and CATE error due to bad recovery of j. in (G1') is controlled by
251
EL BOBO 88. In Appendix, we detail novel implications of our bounds, compared to those in [40 31).
Using to estimate CATE, ły fz) -f((2) is marginalized on {zz The bounds below
253
motivate both prior and posterior balaneing. Let us first consider errors defined by the aggregated
prior x(z := E. p(t (z) x)) (denote Pr (z x) == PA (z/x, t)), and (z( afterwards. We
255
introduce the objective we bound. The true CATE, given covariate x = x or score z, is
T(x)= = m2 m1 (z)( mo (Po(())
(z) = m1 mo(z)
(10)
256
where mz(z) = (() of prior CATE, with without knowing P., is naturally defined as
2) and P. is a balanced PtS in (G)) Accordingly, given x, the error
257
(xr) = Eppe (F)(z) - (=)))? (xx)) = (f)(() - Tm (2))
(11)
258
We bound of instead of t because the error between 7xx) and Tm T. (z) is small if balanced P, is recovered (then z Po( R Pi (x) in (10). see Appendix for details). Instead, we consider the error between fy and Tm below. We define the risks of outcome regression, into which e" is decomposed.
259
260
261
Definition 3 (PO Risks). The expected loss of PO at (z. factual risk, and counterfactual risk are
(z 9: ( Bpir()) (y((t) - f((2)) (x()) 2 [ 1 f(=))) = x)dy; (x) = Cn
=
= (
*As usual, we expect variational inference and optimization procedure are (near) optimal, i.c., Consistency of VAE (see Appendix for formal statement). Consistent estimation using the prior is direct corollary of consistent VAE. Under Gaussian models, it is possible to prove consistency of posterior estimation, as shown in 4.

---

262
With yt) involved, Ly is PO enror of f weighted by 9. Factual and counterfactual counterparts are defined accordingly, w.r.t factual pt learned from data. Note, in Cp.t, unit u = (x. y. t) involves in the learning of pa (2lar) (and qe (Z]ar) afterwards), and also in L,(2, () since y(1 = y for the unit. In CFF, however, y(t) # y' = y( -t) for u' = (x. y' -) (particularly relevant for the posterior).
263
264
265
Thus, the regression error (second) term in ELBO 8 controls e via via factual data. On the other hand, &Cr, is not estimable due to unobservable y t), but is bounded as below.
286
257
Lemma 2 (Counterfactual risk bound). Assume L (z, ()l S M, we have
(Cr(æ)) S E, - 1/x) (a) + (x)
(12)
289
where (x)( ((p((1-())CC (x)) and (* = E. v
=
Por (a:) bounded by Eg,, plus MD)( which measures the imbalance between pa (z/ar) and is 270
271
symmetric for t. We can implicitly control by making D small. Again, this means PS is preferred as a conditional balanced representation, and justifies our balanced prior pA (z|x). Moreover, the results (including Steerempeeetee below) also hold for posterior extimation, that is, replace pe (z) x) with
272
273
274
qo (z)x) == q(]|x,1 1) Epyy (x))))(())))) (,") the results and proofs hold it was. This implies that the imbalance between de should also be controlled. Correspondingly, the symmetric KL term in
275
278
ELBO (8) balances qe (z)x) by encouraging z ||| for the posterior.
277
2 in turn bounds eyy. by decomposing it to elp, &CF,, and Vg. Theorem (Generalization bound). Assume Lz(2, f)] S M and [ge(z)] S G. then, <z(x)= S 2( + (x)) + MD(xx) Vy(x))p
Theorem
278
(13)
where vg(x) == Ep(x(=) E, Eptrtt)) (y m4 (z))² and "Ip" collects all superseripts q.
279
280
The new term (2) reflects the intrinsic variance in the DGP and can not be controlled, and is negative because of
d is defined by mean functions fo and my not yt). The other two terms, as we indicated, is estimated by our ELBO.
281
282
283
Estimating G, MM is nontrivial. Instead, similarly to 40 we rely on 3 in the ELBO to weight the two terms in (13). We do not need two hyperparameters since G is implicitly controlled by the third term in ELBO (8) which is a norm constraint. As in matching methods, 33 is trade-off between conditional balance of learned PtS (affected by fi) and precision / effective sample size of outcome
284
285
286
regression, and can be seen as the probabilistic counterpart of 47 25
288
Finally, we note the bounds do not directly address non-overlap; in Lemma 2 when p(1 - +|xx)
289
e --r in the r.h.s is unbounded since p1-e(z(z0 x)) can not be learned from data. However, as we argued in Sec. 3.2. with more balanced IP, recovered as representation, overlap is more easily satisfied.
290
Related work
291
Weak overlap. Under (respective versions of) weak overlap, 32 estimates ATE by reducing covariates to a linear PGS, [ estimates homogeneous (constant) TE under partial linear outcome model, and [9] studies identification of ATE by general class of scores, given (linear) PPS and PGS. In machine learning, current focus is on finding overlap regions B3 g. or indicating possible failure under weak overlap (22, but not remedies. An exception is 24 which provides bounds without overlap. 44 31 are in line of (24) and have similar bounds to ours, without relating to overlap. Our method is the first in machine learning that gives identification without overlap.
293
295
297
298
299
Prognostic scores are recently combined with machine learning, mainly in biostatistics. For example, trains a flexible PGS and fits a linear regression on the PGS among others, for constant TEs, and models PGS in its Bayesian regression tree for CATE. More related, 20 estimates CATE by reducing covariates to a linear score that is joint PPS and PGS, and [47 uses SVM to minimize
39
301
12
303
the worst-case bias due to PGS imbalance. However, in machine learning, few methods consider
PGSs; [5] learn outcome predictors, without connection to PGS, while 24] conceptually, but not formally, connects BRL to PGS. Our work follows the recent boom in biostatistics and is the first to formally connect generative learning, PGS, and BRL (see below on BRL) for TE estimation.
305
306
307
Identifiable representation. Recently, independent component analysis (ICA) and representation
learning, both ill-posed inverse problems, meet together to give nonlinear ICA and identifiable
7

---

309
representation, e.g., using VAEs 126]. and energy models (27]. The results are exploited in causal
310
discovery [5T] and out-of-distribution generalization (46] This work is the first to explore identifiable representations in TE identification.
312
BRL and related methods amount major direction. Early BRL methods are BLR/BNN 23
313
and TARnet/CFR [40]. Adding to this, [53] also exploits the local similarity of between data points. uses similar architecture to TARnet, considering the importance of treatment probability. There
314
315
are also methods using GAN (54] GANITEJ and Gaussian process M. Our method shares the idea of BRL, and further extends to conditional balancing more suitable for CATE.
316
317
Our work hopefully lays conceptual and theoretical foundations of VAE methods for TEs (e.g.. BI), under unconfoundedness. Also, monotonicity, which is injectivity on R. is important in See Appendix for details.
30
319
causal inference, and some works consider it together with overlap 24 56
Experiments
320
321
We compare the proposed method with existing methods on three types of datasets. Here we present two experiments, and the rest one, on the Pokec social network dataset, can be found in Appendix. As
322
323
in previous works [40] 30). we report the absolute error of ATE ATE := [ED(y(1) -y
and, as a surrogate of CATE, the empirical PEHE 18 CPEHE == Ex((v(1) y(0) =(x))²
324
325
Unless otherwise indicated, for each function f. k. in 4 we use a multilayer perceptron (MIP) that has 3*200 hidden units with ReLU activation, and X (h,,k) k)) depends only on x. We fix
326
327
(x) = 1 because the datasets have fixed noise scale, and results with learned on synthetic dataset with dependent noise is in Appendix. The Adam optimizer with initial learning rate 10-1 and batch
328
size 100 is employed. All experiments use early-stopping of training by evaluating the ELBO on a validation set, and results are reported on testing set. Each set of running on synthetic dataset (a line in the figure) is within I hour on an 8-CPU machine, and it is within day for IHDP. More details on hyper-parameters and settings are given in each experiment and Appendix.
329
330
331
332
333
6.1
Synthetic dataset
334
We generate synthetic datasets following (14). Both x, w are factorized Gaussians. μt, a are randomly sampled. The functions h, k I are linear. Outcome models fo, J1 are built by NNs with invertible activations. y univariate, dim(xx
335
336
30, and dim(w) (w) ranges from I to 5. w is a PS, but is nor low-dimensional when (im) > 1. We control overlap by which multiplies the logit value, and have 5 different overlap levels from strong overlap to very weak overlap. See Appendix for details.
337
338
x NN O); w|x N (h(x,, t(x)): ~ Bern Bern(Togi(wl(x)):: y/w,t ((((w) 1). (14)
With the same (dim(w) w), we evaluate method and CFR on 10 random DGPs, with different sets of functions f h, k, I in (1)). For each DGP, we sample 1500 data points, and split them into 3 equal sets for training, val- idation, and testing. We show our results for different hyperparameter 3. For CFR, we try different balancing parameters and present the best results (see Appendix for details). We report EPEHE, E, see Appendix for ATE results.
1.0
0.328
340
0.300
341
342
0.276
0.6
343
0.250
344
0.225
0.4
345
0,200
348
0.2
0.176
In each panel of Figure 2 we adjust one of w, dim (w) respectively, with the other fixed to the lowest. In the left panel, both our method and CFR are quite robust overlap level, supporting respective theories ((24) gives bounds for
348
non-overtap (omega))
dmm(w)
349
ctr
bota:1.6
beta:2.5
bets-1.0
b.5s-2.0
beta-3.0
350
351
CFR under weak overlap). Too large 3 seems to worsen the results, possibly because we already have an apparent
Figure 2: VEPEHE on synthetic dataset. Er- ror bar on 10 random DGPs.
352
353
PS (w with dim(w) = 1) and large ,3 incurs sub-optimal ELBO with no gain.
354 In the right panel, when dim(w) (w) > 1, ft in (14) is non-injecitve and learning of PtS is necessary. Thus, larger 3 has negative effect, and particularly, B = 1 is significantly better than 3 = 3. The drop of error for dim(w) > 3 is due to the randomness of f in (14). In Sec. Sec. we saw that the 2- 2-dimsensional PS P == (Ao(x)) Mn (x)) always exists under ANMs. Thus, when dim(w)= 2, our method tries to recover that P. and generally performs not worse than under dim(w) = 2, but still
356
358
359
not better than under tim(w) 1.
8

---

380
Our method is much robuster against different DGPs than CFR (see error bars), though it is worse than CFR when dim(w) > 1. This is unsurprising because our model has I-dimensional z, while CFR uses 200-dimensional representation. Thus, the results already show the power of identification and recovery of scores (see Figure Jaloo) In
Intact VAE
CEVAE
361
363
364
365
fact, we observed that our method outperforms or matches
Figure 3: Plots of recovered true latent.
367
CFR with higher-dimensional z (see Appendix). Thus, we believe the performance gap with dim (z) = 1 is due to the capacity of NNs in Attact-VAE.
Blue: L = Orange: 1.
369
When dim (w) 1, there are no better PSs than w, because ft is invertible and no information can be dropped from w. Thus, as shown in Figure 3 our method learns z as an approximate affine transformation of the true w, showing identification. For comparison, we run (30 CEVAF] which is also based on VAE but without identification, and it shows much lower quality of recovery. As expected, both recovery and estimation are better with balanced prior pA (z) x), and we can see an example of bad recovery using pA (z|x,1) t) in Appendix. More latent plots can also be found there.
370
371
372
373
6.2
IHDP benchmark dataset
376
This experiment shows our conditional balancing matches state-of-the-art BRL methods, and does not overfit to PEHE. The IHDP dataset [8]] widely used to evaluate machine learning based methods, e.g. 40 41. It is also used in 244 which considers weak overlap, because the covariates are weakly overlapped due to their correlation to the artificial treatment assignment. Finally, there is a linear PS
377
378
379
(linear combination of the covariates). See Appendix for details.
381
Note, most of covariates are binary, so the support of the PS is often on small and separated intervals and is possibly discrete. Thus, Gaussian latent z is misspecified. We use multivariate z in model to address this, similarly to [30]. We set 3 = 1 since it works well on synthetic dataset with weak overlap. To see our balancing property clearly, we modify our method and add two components for unconditional balancing from 440 (see Appendix). and compare this modified version to the original.
382
383
384
385
Table 1: Errors IHDP over 1000 random DGPs. We report results with dim(z) 10. Bold indicates method(s) that are significantly better. The results are taken from (40), except GANITE (54) and CEVAE (0)
Method
TMLE
BNN
CFR
CF
CEVAE
GANITE
Ours
Ours Mod.
CATE
.30 01 .37 03 25
.18
01 .34..01
43 o0
.178 006
.167 005
VEPRUR
2.01.2
2 .71 o2 38
2.7
1.9.4
.89
27772.026
386
As shown in Table I Intact- outperforms or matches the state-of-the-ar methods. Particularly, our method has the best ATE estimation, and is slightly worse than CFR for PEHE. This is possibly due to the fitting capacity (recall Sec. 6.11) and also we do not tune 3.. Notably, our method
387
388
309
outperforms other generative models (CEVAE and GANITE) by large margins. The modified
version is slightly improved, but we should note that the improvement for CATE is barely significant.
380
391
This indicates overfitting to PEHE. In fact, PEHE estimates the marginalized error E<(x) where
(T(x) - =(x))², and, compared with CATE, it focuses on values x with high probability and / or large (ar)) The balancing [40] is based on bounding Ex(x,, and thus tends to overly focus on
(x)
393
394
the above values of x, resulting in sub-optimal estimation of CATE and even of ATE. This tendency
395
is more apparent with sub-optimal hyperparameter for the unconditional balancing (see Appendix).
7 Conclusion
397
In this work, we proposed a method for CATE estimation, under weak overlap. Our method exploits identifiable VAE, a recent advance in generative models, and is fully motivated and theoretically justified by causal considerations: identification, PGS, and balaneing. We show that VAEs are suitable for principled causal inference thanks to its probabilistic nature, if not compromised by ad hoc heuristics. We believe it is possible to extend the bounds in Sec. t4. weak overlap, just as
398
399
400
401
402
24 extends 40] to weak overlap, and leave this for future. Experiments show evidence that the injectivity of f in our model is possibly unnecessary because diin(z) > din((y) often gives better results. Theoretical study of this is an interesting future direction. To avoid potential negative societal impact (e.g, bad prescriptions), practitioners should judge the conditions of the proposed method by their domain expertise, and careful trials are always recommended.
403
405

---

References
407
[1] Ahmed M Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In Advances in Neural Information Processing Systems, pages 3424-3432, 2017.
409
410
[2] Joseph Antonelli, Matthew Cefalu, Nathan Palmer, and Denis Agniel. Doubly robust matching estimators for high dimensional confounding adjustment. Biometrics, 74(4):1171-1179, 2018.
412
413
[3] Timothy B. Armstrong and Michal Kolesár. Finite-sample optimal estimation and inference on average treatment effects under unconfoundedness. arXiv preprint axXiv:171..0459455, 2021.
414
415
[4] Stéphane Bonhomme and Martin Weidner. 1rXi::1906.0033005,, 2021.
Posterior average effects.
arXiv preprint
[5] M Alan Brookhart, Sebastian Schneeweiss, Kenneth Rothman, Robert Glynn, Jemry Avor, and Til Stürmer. Variable selection for propensity score models. American journal of epidemi- ology, (3311))1149-11555 2006.
418
420
[6] Alberto Caron, Ioanna Manolopoulou, and Gianluca Baio. Estimating individual treatment effects using non-parametric regression models: a review. arXiv preprint N0Xiv:200900677. 2020.
421
423
Denis Chetverikov, Andres Santos, and Azeem M Shaikh. The econometries of shape restrictions. Annual Review of Economics, 10:31-63, 2018.
424
425
[8] Wangzhi Dai and Collin M Stultz. Quantifying common support between multiple treatment groups using contrastive vae. In Machine Learning for Health, pages 41-52. PMLR, 2020.
426
427
99]
Alexander D' Amour and Alexander Franks. Deconfounding scores: Feature representations for causal effect estimation with weak overlap. arXiv preprint arXiv:2104.05762, 2021.
[10] Alexander D' Amour, Peng Ding, Avi Feller, Lihua Lei, and Jasject Sekhon. Overlap in observational studies with high- high-dimensional covariates. Journal of Econometrics, 2020.
430
431
Max Farrell. Robust inference on average treatment effects with possibly more covariates than observations. Journal of Economelrics, 189(1):1-23, 2015.
432
433
[12]
Richard Hahn, Jared S Murray, Carlos M Carvalho, al. Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion). Bayesian Analysis, 15()):955-1056 2020.
434
435
436
[13] David Hajage, Yann De Rycke, Guillaume Chauvet, and Florence Tubach. Estimation of conditional and marginal odds ratios using the prognostic score. Statistics in medicine,33(44)677.
437
438
716, 2017.
439
[14] Ben Hansen. The prognostic analogue of the propensity score. Biometrika, 95(2):481-48) 2008.
[15] Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2019.
442
443
[16] Miguel A. Hernan and James M. Robins. Causal Inference: What If. CRC Press, 1st edition, 2020.
445
[17]
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a con- strained variational framework. In 5th International Conference on Learning Representations, 2017.
446
447
448
449
[18] Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computa- tional and Graphical Statistics, 20(1):2172240, 2011.
450
451
[19] Han Hong, Michael Leung, and Jessie Li. Inference on Enite-populatte treatment effects under limited overlap. The Econometrics Journal, 23(1):32 2020.
10

---

453
[20] Ming- Yueh Huang and Kwun Chuen Gary Chan. Joint sufficient dimension reduction and estimation of conditional and average treatment effects. Biometrika, 104(3):533-566, 2017.
454
495
[21] Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge University Press, 2015.
457
[22] Andrew Jesson, Sören Mindermann, Uri Shalit, and Yarin Gal. Identifying causal-effect inference failure with uncertainty-aware models. Advances in Neural Information Processing Systems, 33, 2020.
458
459
460
[23] Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pages 3020-3029, 2016.
[24] Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. arXiv preprint arXiv:2001.07400 2020.
463
465
[25] Nathan Kallus, Brenton Pennicooke, and Michele Santacatterina. More robust estimation of sample average treatment effects using kernel optimal matching in an observational study of spine surgical interventions. arXiv preprint arXiv: 181 /. ..4274, 2018.
466
467
488
[26] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational au- toencoders and nonlinear ica: A unifying framework. In International Conference Artificial Intelligence and Statistics, pages 2207-2217, 2020.
470
471
[27] Ilyes Khemakhem, Ricardo Monti, Diederik Kingma, and Aapo Hyvarinen. Ice-beem: Iden- tifiable conditional energy-based deep models based on nonlinear ica. Advances in Neural Information Processing Systems, 33, 2020.
472
473
(28) Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling, Semi-supervised learning with deep generative models. In Advances in neural information processing systems, pages 3581-3589, 2014.
475
477
[29]
Arthur Lewbel. The identification zoo: Meanings of identification in econometrics. Journal of Economic Literature, 57(4))835-903, 2019.
478
479
[30] Christos Louizos, Uri Shalit, Joris M Mooij. David Sontag, Richard Zemel, and Max Welling, Causal effect inference with deep variabee models. In Advances in Neural Information Processing Systems, pages 6446-6456, 2017.
480
[31] Danni Lu, Chenyang Tao, Junya Chen, Fan Li, Feng Guo, and Lawrence Carin. Reconsidering generative objectives for counterfactual reasoning. Advances in Neural Information Processing Systems, 33, 2020.
483
485
[32] Wei Luo, Yeying Zhu, and Debashis Ghosh. On estimating regression-based causal effects using sufficient dimension reduction. Biometrika, 104(1):51-63 2017.
[33] Michael Oberst, Fredrik Johansson, Dennis Wei. Tian Gao, Gabriel Brat, David Sontag, and Kush Varshney. Characterization of overlap in observational studies. In International Conference on Artificial Intelligence and Statistics, pages 788-798 PMLR, 2020.
488
480
[34] Judea Pearl. Causality: models, reasoning and inference. Cambridge University Press, 2009.
491
[35] Paul R Rosenbaum. Modern algorithms for matching in observational studies. Annual Review of Statistics and Its Applicaation, 7: 7:143-176, 2020.
[36] Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55 1983.
494
485
[37]
Donald B Rubin. Estimating causal effects from large data sets using propensity scores. Annals of internal medicine, 127(8 Part 2)7777763, 1997.
496
497
[38] Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-3311 2005.
11

---

499
[39] Alejandro Schuler, David Walsh, Diana Hall, Jon Walsh, and Charles Fisher. Increasing the efficiency of randomized trial estimates via linear adjustment for a prognostic score. arXiv preprint arXiv:2012.09933 2020.
500
[40] Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pages 3076-3085. PMLR, 2017.
503
[41] Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In Advances in Neural Information Processing Systems, pages 2507-2517, 2019.
506
507
508
[42] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in neural information processing systems, pages 1193-3491,20 2015.
509
510
511
[43] Peter Sorrenson, Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ica with general incompressible-llow networks (gin). In International Conference on Learning Repre- sentations, 2019.
513
514
[44] Elizabeth A. Stuart. Matching Methods for Causal Inference: A Review and a Look Forward. Statistical Science, 25(1):1 21, 2010.
515
516
[45]
Elizabeth A Stuart, Brian K Lee, and Finbarr P Leacy. Prognostic ccore-based balance measures can be a useful diagnostic for propensity score methods in comparative effectiveness research. Journal of clinical epidemiology, 6688):884-990 2013.
517
518
519
[46] Xinwei Sun, Botong Wu, Chang Liu, Xiangyu Zheng. Wei Chen, Tao Qin, and Tie-yan Liu. Latent causal invariant model. arXiv preprini 201i::001100220:: 2020.
[47] Alexander Tarr and Kosuke Imai. Estimating average treatment effects with support vector machines. arXiv preprint arXiv:2102.1 12.2.11992 11926, 2021.
522
529
[88] Mark J Van der Laan and Sherri Rose. Targeted learning: causal inference for observational and experimental data. Springer Science & Business Media, 2011.
525
[49] Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1136623):122--12222 2018.
528
527
528
[50] Halbert White and Karim Chalak. Identification and identification failure for treatment effects using structural systems. Econometric Reviews, 32(3):273-317, 2013.
[5I] Pengzhou Wu and Kenji Fukumizu. Causal mosaic: Cause-effect inference via nonlinear ica and ensemble method. In International Conference on Artificial Intelligence and Statistics, pages 1157-1167. PMLR, 2020.
531
[52] S Yang and P Ding. Asymptotic inference of causal effects with observational studies trimmed by the estimated propensity scores. Biometrika, 105(2):487-993, 03 2018.
534
535
[53] Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representa- tion learning for treatment effect estimation from observational data. In Advances in Neural Information Processing Systems, pages 2633 (633-2643,22 2018
536
537
538
[54] Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: Estimation of individ- ualized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018.
539
[55] Weijia Zhang, Lin Liu, and Jiuyong Li. Treatment effect estimation with disentangled latent factors. arXiv preprint arXiv:2001. 2020.
542
543
[56] Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the estimation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pages 1005-1014. PMLR, 2020.
544
545
12

---

Checklist
548
1. For all authors...
548
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
550
(b) Did you describe the limitations of your work? [Yes] Particularly in Conclusion.
551
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We men- tioned possible misuse of our method and suggestions for practitioners in Introduction and Conclusion.
552
553
554
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
555
558
2. If you are including theoretical results...
557
(a) Did you state the full set of assumptions all theoretical results? [Yes]
558
(b) Did you include complete proofs of all theoretical results? [Yes] In Appendix.
559
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] supplemental material.
561
663
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Possibly in Appendix.
565
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
567
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
569
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
570
(a) If your work uses existing assets, did you cite the creators? [Yes]
571
(b) Did you mention the license of the assets? [Yes] In code project.
572
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
573
In supplemental material.
574
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [Yes] In Appendix. And we refer to the original paper where this is discussed in details.
575
578
577
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] Refer to the original paper. And we believe there are no privacy issues.
578
579
580
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
582
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
584
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
586
13


Reviews:
[{"review_id": "vNNDJ6o81uz", "reviewer": "Reviewer_DRdR", "summary": "The paper proposed a novel method for conditional average treatment effect (CATE) estimation under weak overlap with theory of individualized treatment effect (TE) identification and counterfactual generalization bounds on TE errors.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well structured and written with decent clarity in theory explanations, experiment settings and results discussion. Code and supplementary materials are available to reproduce the experiment results. \n The proposed method is novel in terms of treatment effect identification without overlap, which is also supported by thorough reference and comparison to related work.\nTechnical analysis of synthetic dataset and benchmark dataset is sound with state-of-art baseline models and results with confidence intervals. It might be better to include more empirical experiment results from the supplementary materials in the main paper if feasible. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "ouxYuUjcYuf", "reviewer": "Reviewer_dsM9", "summary": "This work studies the identification and estimation of treatment effects (with special attention on the CATE) under \"weak\" overlap. \nIt proposes a theory of treatment effect identification when the overlap assumption does not hold for the entire domain of $\\mathbf{x}$ using the prognostic score. \nIt offers counterfactual generalization bounds on treatment effect estimates.\nIt develops the Intact-VAE as a regularized treatment effect estimation model that accounts for identification and balancing.\nIt offers proof that their generative model p(y, x | x, t) is identifiable.\nIt offers proof that CATE is identifiable via their generative model.\nThe methods are validated via a comparison to existing methods using synthetic and semi-synthetic datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality\n- The intact VAE is a novel combination of CVAE and iVAE.\n\n## Quality\n- [Line 64] Counterfactual outcome is not a synonym for potential outcome. The description of $\\mathbf{y}(t)$ corresponds to a potential outcome. Also, word choice on introduce, sounds like you are saying the definition of potential outcome is something you introduce. Could change to \"Following [22], the potential outcome of a treatment is given as ...\"\n- [equation 1] the expression for $\\tau(x)$ is the CATE given the sufficient conditions you describe later on line 77 and consistency. I would suggest moving the conditions before equation 1.\n- [line 77] I'm not sure that this is a valid way to express exchangeability. Is $\\mathrm{y}(1) \\perp \\mathrm{t} \\mid \\mathrm{x}$ and $\\mathrm{y}(0) \\perp \\mathrm{t} \\mid \\mathrm{x}$ equivalent to $\\\\{\\mathrm{y}(0), \\mathrm{y}(1)\\\\} \\perp \\mathrm{t} \\mid \\mathrm{x}$? Similarly for [line 112].\n- Section 3.1: the formulation of IntactVAE looks to be a sensible combination of CVAE and iVAE\n- Lemma 1: Is correct, as far as I know. It is not immediately obvious to me why D1 holds easily in practice, maybe an illustrative example would help.\n- [line 163] $k(x) \\to 0$? Lambda is infinite for $k(x) = 0$.\n\n\n## Clarity\nIn general I found this submission quite difficult to read. The majority of my time reviewing was spent on trying to parse the work, and I have not had the time to go deeply into the methods or results. Below, I have given some recommendations to make the work more readable. I hope that they are helpful.\n\n- In the introduction you define *weak overlap* on line 32 and give citations to related work on line 34. Each of those works use the term *limited overlap.* [9] do use the term weak overlap, but give a different definition. Why did you choose to use weak instead of limited?\n- In my opinion there are too many acronyms. BRL, PGS, PPS, PS, PtS, ANM, and DGP are non-standard acronyms and it is a big ask of the reader to keep in mind 7 of them. PGS, PPS, PS, PtS are all visually quite similar, which makes things challenging for the reader. I would suggest maybe 2-3 max (e.g. Pt-score, P-score, and Intact-VAE), and I would have had an easier time reading the paper if you just used the terms themselves for the rest.\n- What do M, G, and D stand for?\n- [line 16] \"counterfactual outcomes, which would have been\" -> counterfactual outcomes: outcomes that would have been\n- [line 32-33] \"is that\" -> is when. \"are all belong\" -> all belong\n- [line 36-37] What is a rich covariate? An explanation of the term or just descriptors of what you mean by rich covariate would be preferable as rich covariate is ambiguous. \n- [line 38-44] It is nice that you give an example of a prognostic score (the propensity score), but an intuitive definition of prognostic scores or  a description of how the propensity score is a prognostic score would be helpful.\n- [line 72] The difference between roman fonts and italic fonts is subtle. Another common option is to use capital roman fonts for r.v.s and lower-care roman fonts for realizations.\n- You use $\\mathbf{e}$ for noise on [line 107] and $\\mathbf{\\epsilon}$ on [line 134]. Please use one or the other consistently.\n- Equation (4) $\\theta$ is missing. Suggest an equation $p_{\\theta}(\\mathrm{y}, \\mathrm{z} \\mid \\mathrm{x}, \\mathrm{t}) = p_{f}(\\mathrm{y} \\mid \\mathrm{z}, \\mathrm{t}) p_{h, k}(\\mathrm{z} \\mid \\mathrm{x}, \\mathrm{t}),$ then define component distributions. Maybe keep in terms of h, k rather than lambda?\n- [line 134] density for $\\epsilon$?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have clearly defined their problem setting and the assumptions under which they are working. Namely that they are relaxing the overlap assumption, and assume unconfoundedness and consistency. They state the linear independence assumption of Lambda (though it is not clear to me why this holds easily in practive), and the consequences of this assumption failing.  There is a sentence in the conclusion recommending that practitioners judge whether the conditions hold under which their method is valid.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "jvdKO-2wI7g", "reviewer": "Reviewer_8ody", "summary": "The paper studies identification and estimation of treatment effects (TEs) under weak overlap of covariates (Definition 1, probability of treatment conditioned on covariate has different supports for some covariate, which is closer to the real world). For identification of TEs, the paper exploits the prognostic score (PGS) and the identiﬁable representation, and as a result, proposes a novel, VAE-based generative prognostic model. Prior studies have seen good practical performances of PGS in many settings and hence, the paper expects it to be a good choice and to outperform other candidates in joint use with additional techniques. The identifiable representation is essential to guarantee the identification of TEs. The estimation of the VAE model is via a $\\\\beta$-regularized-ELBO which helps to keep the posterior distribution of latent variables of VAE ''balanced'' (defined in line 219). Motivations for e.g. balanced representation are provided later in the theoretical analyses on the counterfactual risk bound and the generalization bound. The proposed method is compared with baselines on (semi-)synthetic datasets.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. Originality\n\nThe paper combines several recent techniques for TEs: generative learning (VAE), identifiable representation, PGS, and balanced representation learning. The combination was not explored before. The theoretical analyses are new under the setting of joint use of multiple methods.\n\n2. Quality\n\nThe VAE based method with theoretical guarantees and experimental validation is technically interesting. I haven't realized any technical issue in the main paper. However, it is slightly cumbersome to stack so many techniques.\n\n3. Clarity\n\nThe paper is dense and needs to provide e.g. examples or remarks to help readers understand concepts and assumptions. For example, the concept of ''balanced'' distribution non-rigorously introduced on line 219 is not easy to understand; the importance of the assumption G1' in the theorems needs explaining (if just to make dimensionality explicit, can G1 and G1' be merged as just one assumption in the paper?).\n\n4. Significance\n\nThe identifiable and balanced representation learning approach for TEs will be interesting for many readers in the field. The practical outperformance and the theoretical guarantee are the essential reasons.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper has proposed sound suggestions to solve the limitations and potential negative societal impact in the conclusion section.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}, {"review_id": "7ST8jhowO1", "reviewer": "Reviewer_r3t8", "summary": "This paper proposes a generative prognostic model for identifying and estimating causal effects under weak overlap. The proposed method uses a latent variable to model a prognostic score (PGS). The authors prove that the latent variable recovers a PGS, and the model identifies individualized treatment effects. The model is learnt as an intact-VAE. The authors provide a theory of treatment effect identification and counterfactual generalization bounds on the treatment effect error. The generalization bounds motivate conditional balanced representation learning (BRL). Then they propose a new regularized VAE to estimate treatment effects. In experiments, they compare the proposed method with some state-of-the-art benchmarks on some synthetic and semi-synthetic datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Poor presentation with inaccurate statements. I thought some statements in the introduction are misleading for readers who are not familiar with causal inference.  \n\nFor example, the description of RCTs in line 18 is incorrect. It is well known that RCTs play a crucial role in the process of drug development. RCTs are indeed expensive but not necessarily prohibitive. Is observational data free of ethical issues?  The statement in line 19, “causal inference from observational data is indispensable” exaggerates the importance of observational data. The advantage of RCT is that it controls certain biases through randomization. Experimental and observational studies have their own strengths and weaknesses. The authors should discuss both in short and provide references to support the claims.\n\nIn lines 20-21, you claim the most crucial challenge in observational data is confounding, spurious correlation arises from confounding. This statement is inaccurate. If the confounder is observed, there is no spurious correlation conditioning on the confounder. If the confounder is unobserved, there will be a spurious correlation. This is also the reason why we need the unconfoundedness assumption in observational studies. The explanation of unconfoundedness in lines 22-24, “the confounding can be control by conditioning on covariates” does not clearly capture the meaning of unconfoundedness.\n\nIn the second paragraph of the introduction, more works in semiparametric statistics and representation learning should be cited. In the third paragraph, you mention your approach to the weak overlap issue is a prognostic score (PGS) but do not clearly explain what PGS is and why it can address the overlap issue.  Then in the fourth paragraph, you start to discuss the problem of recovering PGSs without telling the reader what a PGS is. \n\nThere are existing works on each of the main contributions of the paper (Identification theory of treatment effects, (i)VAE, counterfactual generalization bounds and balanced representation learning). The introduction should make it clear what is the difference between your work and the existing works. Perhaps the related work section should come earlier for this purpose. \n\nSection 2. Please use upper-case letters for random variables and lower-case letters for fixed values. It is difficult for the readers to distinguish Roman and italic fonts because using Roman and italic fonts together is not a standard practice.\n\nIn lines 77 and 87, p(t|x) and p(t|v) should be also be less than 1.\n\nIn line 98, you mentioned overlapped b(x) implies overlapped x, overlapped x implies overlapped PGS. But to show PtS is more applicable, you also need to prove the other way around, PGS-> x -> b(x), does not hold. I don’t see why PtS is lower-dimensional and more overlapped? What is the definition of more overlapped? PtS is defined in lines 94-95. I can't see why it maximizes overlap.\n\nIn line 102, conditionally balanced representation is better for what? I am confused here, \nConditioning on x, the only randomness in P_t(x) is t. It is impossible that P_t(x) \\indep t | x.\nDo you mean x \\indep t | P_t(x)? I see this is what you have in Definition 2. A balanced representation may also lose important information of the outcomes. It is not better surely.\n\nIn Proposition 1, mu_t (x) is the outcome function, not the CATE function.\n\nSection 3. The comparison between Intact-VAE with CVAE and iVAE is clear. However, the novelty of Intact-VAE is limited.  Are the conditions in Lemma 1 and Theorem 1 all verifiable in practice? Is identifying f_t up to some affine transformation or some invertible mapping sufficient for consistent CATE estimation?\n\nSection 4. The generalization bounds in this section are not novel. Similar results are already obtained in the existing works [A, B, C]. More importantly, does minimize the  VAE objective function lead to minimizing all the empirical terms in the generalization bound so that the CATE estimation is consistent (i.e. the expected risk can be as small as the variance of the random noise in the ANM) as the sample size goes to infinity? BTW, why the error from random noises does not exist in the generalization bound? Rigorously, the generalization bound in Theorem 2 is not exactly a generalization bound in learning theory, because the terms in the upper bound are expected risks rather than empirical risks.\n\n[A] U Shalit, FD Johansson, D Sontag. Estimating individual treatment effect: generalization bounds and algorithms, ICML (2017).\n[B] Alaa, Ahmed, and Mihaela Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design, ICML (2018).\n[C] Yao Zhang, Alexis Bellot, Mihaela van der Schaar. Learning Overlapping Representations for the Estimation of Individualized Treatment Effects, AISTATS (2020).\n\nSection 6.  The results do not convince me that the proposed method is useful in practice. The method gives worse PEHE or \\sqrt{PEPE} than CFR, as shown in Figure 2 and Table 1. There is no experiment provided to back up the identifiability theory. For example, could you empirically show the function f_t in Lemma 1 is identified up to some affine transformation?\n\nOverall, I find the paper has a large room for improvement in terms of presentation, theory and experiments. \n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they do.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors study estimation of causal effects under weak overlap---when certain features have values that occur only for one treatment group. They propose a latent variable model of a prognostic score, a variant of the variational auto encoder, showing that it can identify conditional average treatment effects under certain conditions. A basic generalisation bound is provided and the approach is evaluated on a standard benchmark dataset. 

The reviewers found the proposed use of latent variable models for this purpose novel and appreciated the identification results. However, they found that proposed generalisation bound should be contextualised better in previous work and possibly renamed. This was partially addressed by the author response, but would require rewriting certain parts of the paper. Further, the empirical results could be improved by choosing more appropriate benchmarks for the proposed methods (which highlight the strengths and weaknesses of the idea) and by including further simulations which illustrate the proven identification results empirically. While the latter point was partially addressed by the authors in Figure 3, as pointed out in their response, this is a fairly anecdotal result. The synthetic results are also somewhat inconclusive with regards to the parameter \beta. In addition to this, several of the reviewers found that the text itself needs a large revision to reach the standards required for acceptance, beyond what can be expected for a camera-ready update. 
Decision: Reject