$\beta$-Intact-VAE: Identifying and Estimating Causal Effects under Limited Overlap
Authors: Pengzhou Abel Wu
Kenji Fukumizu
Venue: ICLR 2022 Poster (2022)
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 in causal inference, we discuss the identification and estimation of treatment effects (TEs) under limited overlap; that is, when subjects with certain features belong to a single treatment group. We use a latent variable to model a prognostic score 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 prognostic score, and the model identifies individualized treatment effects. The model is then learned as $\beta$-Intact-VAE––a new type of variational autoencoder (VAE). We derive the TE error bounds that enable representations balanced for treatment groups conditioned on individualized features. The proposed method is compared with recent methods using (semi-)synthetic datasets. 

Published as a conference paper at ICLR 2022
B-INTACT- VAE: IDENTIFYING AND ESTIMATING CAUSAL EFFECTS UNDER LIMITED OVERLAP
Pengzhou (Abel) Wu & Kenji Fukumizu
Department of Statistical Science, The Graduate University for Advanced Studies & The Institute of Statistical Mathematics Tachikawa, Tokyo {wu. pengzhou, fukumi zu)e 0 sm. ac. ip
ABSTRACT
As important problem in causal inference, discuss the identification and estimation of treatment effects (TEs) under limited overlap: that is, when subjects with certain features belong to a single treatment group. We use a latent variable to model a prognostic score which is widely used in biostatistics and sufficient for TEs; i.e., we build a generative prognostic model. We prove that the latent vari- able recovers prognostic score, and the model identifies individualized treatment effects. The model is then learned 3-Intact- VAE-a new type of variational au- toencoder (VAE). We derive the TE error bounds that enable representations bal- anced for treatment groups conditioned on individualized features. The proposed method is compared with recent methods using (semi-)synthetic datasets.
INTRODUCTION
Causal inference (Imbens& & Rubin 2015 Pearl 2009). , i.c, inferring causal effects of interventions, is fundamental field of research. In this work, we focus on treatment effects (TEs) based on a set of observations comprising binary labels T for treatment/control (non-treated), outcome Y, and other covariates X. Typical examples include estimating the effects of public policies or new drugs based on the personal records of the subjects. The fundamental difficulty of causal inference that we never observe counterfactual outcomes that would have been we had made the other decision (treatment or control). While randomized controlled trials (RCTs) control biases through randomization and are ideal protocols for causal inference, they often have ethical and practical issues, or suffer from expensive costs. Thus, causal inference from observational data is important.
Causal inference from observational data has other challenges as well. One confounding: there may be variables, called confounders, that causally affect both the treatment and the outcome, and spurious correlation/bias follows. The other the systematic imbalance (difference) of the distribu- tions of the covariates between the treatment and control groups-that is, X depends on T,, which introduces bias in estimation. A majority of studies on causal inference, including the current work, have relied on unconfoundednes:; this means that the confounding can be controlled by condition- ing on the covariates. The more covariates are collected the more likely unconfoundedness holds; however, more covariates tends to introduce stronger imbalance between treatment and control.
The current work studies the issue of imbalance in estimating individualized TEs conditioned on X. Classical approaches aim for covariate balance, X independent of T' by matching and re-weighting Stuart 2010; Rosenbaum 2020). Machine learning methods have also been exploited; there are semi-parametric methods e.g.. Van der Laan & Rose (2018 TMLE) improve finite sam- ple performance, as well as non- non-parametric methods e.g., Wager & Athey (2018) CF). Notably, from Johansson et al. (2016). there has been a recent increase interest in balanced representation learning (BRL) to learn representations Z of the covariates, such that Z independent of T.
The most serious form of imbalance the limited (or weak) overlap of covariates, which means that sample points with certain covariate values belong to single treatment group. In this case, straightforward estimation of TES is not possible at non-overlapping covariate values due to lack of data. There are works that provide robustness to limited overlap (Armstrong & Kolesár 2021). trim non-overlapping data points (Yang & Ding 2018), weight data points by overlap (Li & Li 2019). or study convergence rates depending overlap (Hong et al. 2020). Limited overlap particularly relevant to machine learning methods that exploit figh-dimensional covariates. This is because, with higher-dimensional covariates, overlap is harder satisfy and verify (D) Amour et al. 2020).

---

Published as a conference paper at ICLR 2022
To address imbalance and limited overlap, statistic of outcome predictors and is among the key concepts of sufficient scores for TE estimation. As function of covariates, it can map some non-overlapping values to overlapping value in space of lower-dimensions. For individualized TEs, we consider conditionally balanced representa- tion Z, such that Z is independent of 7 given X-which, a balanced prognostic score. Moreover, prognostic score modeling can benefit from methods in pre- dictive analytics and exploit rich literature, particularly in medicine and health (Hajage al. (017). Thus, it is promising to combine the predictive power of prognostic modeling and machine learning. With this idea, our method builds on generative prognostic model that models the prognostic score as latent variable and factorizes to the score distribution and outcome distribution.
a prognostic score (Hansen 2008 is sufficient
will see, is necessary condition for
As we consider latent variables and causal inference, identification is an issue that must be dis- cussed before estimation is considered. "Identification" means that the parameters of interest (in our case, representation function and TEs) are uniquely determined and expressed using the true observational distribution. Without identification, consistent estimator is impossible to obtain, and a model would fail silently; in other words, the model may fit perfectly but will return an estima- tor that converges to wrong one, does not converge all (Lewbel, 2019 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 (White & Chalak 2013). Thus, it is critical to specify the theoretical conditions for identification, and then the appli- cability of the methods can be judged by knowledge of an application domain.
A major strength of our generative model is that the latent variable is identifiable. This is because the factorization of our model naturally realized as combination of identifiable VAE (Khemakhem it al. 2020a iVAE) and conditional VAE (Sohn et al. 2013] CVAE). Based on model identifiability, we develop two identification results for individualized TEs under limited overlap. A similar VAE architecture was proposed in Wu & Fukumizu (020b); the current study is different in setting, theory, learning objective, and experiments. The previous work studies unobserved confounding but not limited overlap, with different set of assumptions and identification theories. The current study further provides bounds on individualized TE error, and the bounds justify a conditionally balancing term controlled by hyperparameter B, as interpolation between the two identifications.
In summary, we study the identification (Sec. 3 and estimation (Sec. 4 of individualized TEs under limited overlap. Our approach is based recovering prognostic scores from observed variables. To this end, our method exploits recent advances in identifiable representation represetationnppaaicllaaal iVAE. The code is in Supplementary Material, and the proofs in Sec. A Our main contributions are:
1) TE identification under limited overlap of X, via prognostic scores and identifiable model;
2) bounds on individualized TE error, which justify our conditional BRL;
3) new regularized VAE, B- B-Intect- VAE, realizing the identification and conditional balance;
4) experimental comparison to the state-of of---of-hhhh the-tharrt art methods on (semi-) msynh-isythhttic datasets.
1.1 RELATED WORK
Limited overlap. Under limited overlap, Luo et al. (2017 7) estimate the average TE (ATE) by
ducing covariates to a linear prognostic score. Farrell 2015) estimates a constant TE under a partial linear outcome model. D' Amour & Franks (202l) y study the identification of ATE by general class of scores, given the (linear) propensity score and prognostic score. Machine learning studies on this topic have focused on finding overlapping regions @Oersst et al. 2020 Dai & Stultz 2020), or indi- cating possible failure under limited overlap (Jesson et al. 2020). but not remedies. An exception is Johansson al. 2020). which provides bounds under limited overlap. To the best of our knowledge, our method is the first machine learning method that provides identification under limited overlap.
Prognostic scores have been recently combined with machine learning approaches, mainly in the biostatistics community. For example, Huang & Chan (2017) estimate individualized TE by reduc- ing covariates to a linear score which is a joint -yppensstyy-roonotttc score. Tarr & Imai] (202 l) use SVM to minimize the worst-case bias due to prognostic score imbalance. However, in the machine learning community, few methods consider prognostic scores; Zhang et al. (2020a) and Hassanpour & Greiner (2019) learn outcome predictors, without mentioning prognostic score Jerieewwillllooeeiee Johansson ct al. (2020) conceptually, but not formally, connects BRL to prognostic score. Our work is the first to formally connect generative learning and prognostic scores for TE estimation.
2

---

Published as a conference paper at ICLR 2022
Identifiable representation. Recently, independent component analysis (ICA) and representation
learning- both ill- ill-posed inverse problems- -eet together to yield nonlinear ICA and identifiable representation: for example, using VAEs (Khemakhem et al. 2020a), and energy models (Khe- makhem et al.] 2020b). The results are exploited in causal discovery (Wu & Fukumizu 2020a) and out-of-distriuution (OOD) generalization (Sun et al., 2020 This study is the first explore identifiable representations in TE identification.
BRL and related methods amount a major direction. Early BRL methods include BLR/BNN (Johansson et al. 2016) and TARnet/CFR (Shalit et al. (017). In addition. Yao et al. 20188 exploit the local similarity between data points. Shi et al. (2019) use similar architecture
TARnet, con- sidering the importance of treatment probability. There are also methods that use GAN Yoon et al.) 2018 GANITE) and Gaussian processes (Alaa van der Schaar 2017). Our method shares the idaa of BRI., and further extends to conditional flaancewwiiee is natural for individualized TE.
More. Our work lays conceptual and theoretical foundations of VAE methods for TEs (e.g., CEVAE Louizos et al.] 2017 Lu et al.) (020). See Ser. more related works, there we also make detailed comparisons to CFR and CEVAE, which are well-known machine learning methods.
SETUP AND PRELIMINARIES
2.1 COUNTERFACTUALS, TREATMENT EFFECTS, AND IDENTIFICATION
Following Imbens Rubin 2015), we assume there exist potential outcomes (t) R R (0. 1).
(t) is the outcome that would have been observed if the treatment value T = t was applied. We see (t) as the hidden variables that give the factual outcome Y under factual assignment T Formally, (I) is defined by the consistency of counterfuctuals: Y = (1) if = I; or simply Y Y(T). The fundamental problem of causal inference is that, for a unit under research, we can observe only one of Y (0) ( w.tt the treatment value applied. That is, "factual" refers to Y or T, which is observable; or estimators built on the observables. We also observe relevant covariate(s) X X' C RTh which is associated with individuals, with distribution D = YX,YT) ~ p(a; y.t),, y.t). We use upper-case (e.g. T) to denote random variables, and lower-case (e.g. f) for realizations.
The expected potential outcome is denoted by (e(x) (t) X x)) conditioned on X x. The estimands in this work are the conditional ATE (CATE) and ATE, defined, respectively, by:
T ( = μM (x) no (x), v = (()))).
CATE is seen as an individual- anddvddud-levell.esssnaaaaal treatment effect, given highly discriminative X.
Standard results (Rubin 2005) & Robins 2020 Ch. 3) show sufficient conditions for TE identification in general settings. They are Exchangeability: Y [t) and Overlap: p[t|x) for any æ E X. Both are required for & E {0, When & appears in statements without quantifi- cation, we always mean "for both t". Often, Consistency is also listed; however, as mentioned, better known as the well-definedness of counterfactuals. Exchangeability means, just as in RCTs, but additionally given X, that there is no correlation between factual T and potential (().. Note that the popular assumption Y (0), Y (1) JTT is stronger than (I) X and is not necessary for identification (Hernan & Robins, 2020 pp. 15). Overlap means that the supports of p(z = 0) and p(z/t 1) should be the same, and this ensures that there are data for pe (xxx) any (x,t).
will relax the condition of the overlap- ping covariate allow some non-overlapping values that is, covariate X is fimited-overlappnng In this paper, we also discuss overlapping variables other than X (e.g., prognostic scores), and pro- vide definition for any random variable V with support V follows:
We rely on consistency and exchangeability, but in Sec. 3.2
3.2
Definition 1. V is Overlapping if p(t|V = v) > 0 for any t E 10,1), V. If the condition is violated at some value v. then v is non-overlapping and is limited-overlapping.
2.2 PROGNOSTIC SCORES
Our method aims to recover prognostic score (Hansen 2008), ) adapted to account for both t as in Definition 2 On the other hand, balancing scores (Rosenbaum & Rubin 1983) b(X) are defined by T TLXX X ), of which the propensity score p(t =XX) is a special case. See Sec. B. for detail.
3

---

Published as a conference paper at ICLR 2022
Definition 2. iG {p(X, (o. such that (t()), where paz, t) (pe(z) hereafter)
is a function defined on X' x (0, A PGS is called balanced (and a bPGS) if po = P1-
We say PGS is overlapping, if both Po(X) and (x(X) are overlapping. Obviously, a bPGS p(X) is conditionally balanced representation (defined as Z in Introduction) and is thus named. We often write f of the function argument in subscripts.
We use bPGS or PGS to construct representations for CATE estimation. Why not balancing scores?
While balancing scores b(X) have been widely used in causal inference, PGSs are more suitable for discussing overlap. Our purpose is to recover an overlapping score for linited-overlappingg X. is known that overlapping b(X) implics overlapping X (D' D'moour et al. 2020), which counters our purpose. In contrast, overlapping bPGS does not imply overlapping b(X). Example. Let T + 0) and Y = x[ T) + e, where is the indicator function, e and e are exogenous zero-mean noises, and the support of X is on the entire real line while is bounded. Now, X itself is a balancing score and X is bPGS; and [X] is overlapping but X is not. Moreover, with theoretical and experimental evidence, is recently conjectured that PGSs maximize overlap among a class of sufficient scores, including b(X) (D' Amour & Franks 2021 In general, Hajage et al. (2017) show that prognostic score methods perform better or as well as propensity score methods.
Below is a corollary of Proposition 5 in Hansen (2008); note that (X) satisfies exchangeability. Proposition 1 (Identification via PGS). (X)) a PGS and (x;, ~ PYIP; P,) where E {0, }} is a counterfactual assignment, then CATE and ATE are identified, using ① and (z(x) (Y (xX), x) E(Plp; T =
f)ddy
(2)
With the knowledge of Pt and py),; we choose one of Po- P1 and set t in the density function, w.r.t the μ4 of interest. This counterfactual assignment resolves the problem of non-overlap at æ. Note that a sample point with X may not have T = f.
We consider additive noise models for (t). which ensures the existence of PGSs.
(G1) 1 Additiie noise model) the data generating process (DGP) for Y is Y = (m(XX e where f", are functions and is a zero-mean exogenous (external) noise.
The DGP is causal and defines potential outcomes by Y(t) = fi (m, + e, and specifies m( (X, T), T, and e as the only direct causes of Y. Particularly, mx(X) is a sufficient statistics of X for (t). For example, 1) m. (X) can be the component(s) of X that affect (t) directly, or 2)
if ( follows a generalized linear model, then m, m,(X) can be the linear predictor of (
Under (G1) 1) (X) is a PGS; 2) (X) = (m, is PGS; 3) X is a (trivial) bPGS; and 11 (X)) is bPGS. The essence of our method is to recover the PGS m, (X) than Y and approximately balanced. Note that (x))), ), our final target, is a low-dimensional PGS but not balanced, and we estimate conditioning on the approximate bPGS (X).
4) u( = (X),
as representation, assuming m, (X) is not
IDENTIFICATION UNDER GENERATIVE PROGNOSTIC MODEL.
(a.c)
specify the generative prognostic model p(y. 2 ax, t). In Sec. 3.1
and show its identifiability. In Sec. 3.2, we prove the identification of CATES, which is one of our main contributions. The theoretical analysis involves only our generative model (i.e., prior and decoder), but not the encoder. The encoder is not part of the generative model and is involved as an approximate posterior in the estimation, which is studied in Sec. 4
s(elo)
p(aix
p(y(a)
rlyls.))
3.1 MODEL,, ARCHITECTURE. AND IDENTIFIABILITY
Our goal is to build a model that can be learned by VAE from obser- vational data to obtain a PGS, better, bPGS, via the latent variable
Figure 1: CVAE, iVAE, and
Z. The generative prognostic model of the proposed method is in ( Intact-VAR Graphical mod- els of the decoders.
'The labels G. M, or D mean Generating process (of Y), probabilistic Model, or Distribution (of X). We introduce assumptions when appropriate but compile them in one place in Sec. C.1
4

---

Published as a conference paper at ICLR 2022
where 0 == (f.h. k) contains the functional parameters. The first factor Py(y),,,) 2,))),,), decoder, models Ty)p,.T(ypp: P. in (2) and is an additive noise model, with e ~ Pe as the exogenous noise. The second factor Px ax t). our conditional prior, models Pr(X) and is a factorized Gaussian, with (x = diag (X)(()))( (h))))((X) -x)))-))* -1)² as its natural parameter the exponential family, where diag() gives diagonal matrix from a vector.
(v. 2|2, t) - (y ())x2 pr(y|z,t)= = P2 fz)),, pA (2 t) (a); diag(k;(x))))
(3)
We denote n = dim(m)))) For inference, the ELBO given by the standard variational lower bound loy 2 Bgzg Po (y) 2. (wt(((/lea)))
(g(=/ y, (=) t)).
(4)
Note that the encoder q conditions on all the observables (X, Y, T); this fact plays important role in Sec. 4.1] Full parameterization of the encoder and decoder is also given Sec. 4.1 This
architecture called Intacı-VAL (Identifiable tredtment-condiiional VAE). See Figure for com-
parison in terms of graphical models (which have not causal implications here). See Sec. more expositions and Sec. B.2 for basics of VAEs.
C.2 for
Our model identifiability extends the theory of iVAE, and the following conditions are inherited.
(M1) i) ft injective, and ii) ft is differentiable.
(D1) (X) is non-degenerate, i.e., the linear hull of its support is 2n-dimensional.
Under (MI) and (D1), obtain the following identifiability of the parameters in the model: if po xz 1) = pa (y æt 1), we have, for any in the image of fr:
(5) where diag(a) is invertible n-diagonal matrix and b is n-vector, both of which depend on (x(x) and (x). The essence of the result is that f; =
(v:) = diag(a) (v.) + b= A. (g:))
Ar; that is, fr can be identified (learned) up an affine transformation A,. See Sec. A for the proof and relaxation of (D1) In
this paper, symbol (prime) always indicates another parameter (variable, etc.): O' (f', X))
3.2 IDENTIFICATIONS UNDER LIMITED-OVERLAPPEN COVARIATE
In this subsection, we present two results of CATE identification based on the recovery of equivalent bPGS and PGS, respectively. Since PGSs are functions of X, the theory assumes a noiseless prior for simplicity, i.c., ((X)= O; the prior Zx, px(=/xx t) degenerates to function (x.)).
PGSs with dimensionality lower than or equal to = dim(YY are essential to address limited overlapping, shown below. We set = d because μ4, a PGS of the same dimension Y under (G1) In practice, n = d means that we seek a low-dimensional representation of X. We introduce (G1') (Low-dimensional PGS)
G1 true, and μr j. Pt for some Pt and injective j..
which is equivalent to (G1) because μ4, = j: o P: is trivially satisfied with je is identity and p. = ur-- (G1') is used instead in this subsection. First, it explicitly restricts dim (p:) via injectivity, which ensures that n = dim (Y) 2 dim(p.). Second, it reminds us that, possibly, the decomposition is not unique; and, clearly, all p. that satisfy(G1)) are PGSs. For example, if f; is injective, then j. fi j: o p.. Finally, it is then natural to introduce
and p. = m, satisfies
(G2) (Low-dimensional bPGS) G11 true, and 14 o p for some p and injective jt.
which is stronger than (11)
(G2) is satisfied (G2) implies μ1 = dA where = j1 o jo in words, CATEs are given by Ho and invertible function. See Sec. C for real-world examples and more discussions.
gives bPGS pXX), and ensures that n 2 dim(p).
f² is injective and my mi
With or G2) overlapping can be relaxed to overlapping bPGS or PGS plus the following:
(M2) (Score partition preserving) For any x, x' X if px(''), then (x: h,(x''.
Note that (M2) is only required for the optimal h specified in Proposition or Theorem | The intu- ition is that P. maps each non-overlapping
an overlapping value, and hr preserves this property through learning. This is non-trie con-trivial because, for a given t, some values of X are unobserved due to limited overlap. Thus, (M2) can be seen weak form of OOD generalization: the NNs for h can
5

---

Published as a conference paper at ICLR 2022
learn the OOD score partition. While unnecessary for us, linear P2 and h2 trivially imply (M2) and
are often assumed, e.g., in Huang & Chan (017); Luo et al. (2017); ) D' Ammur & Franks 2021
Our first identification, Proposition] 2 relies on (G2) and our generative model, without model iden- tifiability (so differentiable f. is not needed).
Proposition 2 (Identification via recovery of bPGS). Suppose we have DGP (G2) and model (3)
with ni = d. Assume (MI) i) and (M3) (PS matching) let ho(X)= = h1 (X) and k = 0. Then, if Epo (Y (XX T") = E T) we have
1) (Recovery of bPGS) ZA,t = h₂(2) v(p(z)) on overlapping x. where + In is an injective function, and = {p(a loverlapping x):
2) (CATE identification) ((X) in (G2) is overlapping, and (M2) is satisfied, then
[e(x)) = [x)x)) == (p) Ep, ZTt) fh (1)()))), for any t {0, 1} and æ E X.
In essence, i) the true DGP is identified up to an invertible mapping v, such that ft = Je o v- and h = v p; and ii) is recovered up to v, and [t)) (X) is preserved with same for both Theorem [ below also achieves the essence i) and ii), under Po + pl.
The existence of bPGS is preferred, because it satisfies overlap and (M2) more easily than PGS which requires the conditions for each of the two functions of PGS. However, the existence of low- dimensional bPGS is uncertain in practice when our knowledge of the DGP is limited. Thus, depend on Theorem based on the model identifiability to work under PGS which generally exists.
Theorem 1 (Identification via recovery of PGS). Suppose have DGP (GI') and model (3) with n = d For the model, assume (MI) and (M3') (Noise matching) let Pe = Pe and O. Assume further that (DI) and (D2) (Balance from data) Ao = Ai in pe(y/a (ylxtt) f) (2 t); conclusions 7) and 2) in Proposition 2 hold with replaced with Pa in (GI') and the domain of becomes p {P=(x))p((t
kXX)= kk' (X),
(5) Then, if pe
> 0}.
Theorem Timplies that, without bPGS, we need to know or learn the distribution of hidden noise to 2 and Theorem achieve recovery and identification in a complementary manner, the former starts from the prior by Po = P1 and ho = h1, while the latter starts from the decoder by Ao Ar and Pe = Pe- We see that Ag Aı acts as a kind of balance because replaces Pu = P1 in Proposition We show in Sec. A sufficient and necessary condition (D2) on data that ensures Ao = Ai Note that the singularities due to k + 0 (e.g., À - 0) cancel out in (5) See Sec. C.4 for more on the complementarity between the two identifications.
have Pe Pe. Proposition
4 ESTIMATION BY E-INTACT-VAE
4.1 PRIOR AS BPGS, POSTERIOR AS PGS, AND B REGULARIZATION STRENGTH
In Sec. 3.2 we see that the existence of bPGS (Proposition 2 is preferable in identifying the true DGP up to an equivalent expression-we wwile Theorem [ allows us to deal with PGS by adding other conditions. In learning our model with data, we formally require (GT) and further expect that (G2) holds approximately; the latter is true when f² is injective and mo R my (m,(X) is an approximate bPGS). Instead of the trivial regression Mr (X) = B(Y imate bPGS my (X). This idea is common in practice. For example, in a real-world nutrition study (Huang& & Chan 2017), reduction of 11 covariates recovers 1-dimensional linear bPGS.
t), we want to recover the approx-
We consider two ways to recover an approximate bPGS by a VAE. One is to use a prior which does not depend on I, indicating a preference for bPGS. Namely, we set A Ai denote A (X) == X(X)
and have p^ (z|=) as the prior in (3). The decoder and encoder are factorized Gaussians:
Pf.e P (y 2, t) = (3y: fi (z), (tag(g((2))), (0(=|x y, t) = A( (1 y), diag(s: (x)))))) (6)
where Ø- (r. s). The other is introduce a hyperparameter in the ELBO as in 3-VAE (Higgins)
et al. 2017). The modified ELBO with 3, up to the additive constant, is derived as:
E-(-3Dkr. (g llx) Ez-gy [yy - f()())) ((z)))?/g.(
Ez-ag log (=)]
(7)
For convenience, here and in Cy in Sec. 4.2 we omit the summation as if Y is univariate. The encoder 9s depends on and can realize PGS. With 3 we control the trade-off between the first and second terms: the former is the divergence of the posterior from the balanced prior, and the latter is the reconstruction of the outcome. Note that larger / encourages the conditional

---

Published as a conference paper at ICLR 2022
balance Z X the posterior. By choosing 3 appropriately, e.g., by validation, the ELBO can recover an approximate bPGS while fitting the outcome well. In summary, we base the estimation on Proposition 2 and bPGS as much possible, but step into Theorem I and noise modeling required by Pe = Pc. when necessary.
Note also that the parameters g and k, which model the outcome noise and express the uncertainty of the prior, respectively, are both learned by the ELBO. This deviates from the theoretical conditions described in Sec. 3.2 but it is more practical and yields better results in our experiments. See Sec. cC5 for more ideas and connections behind the ELBO.
Once the VAE is learned the ELBO, the estimate of the expected potential outcomes is given by: [1(x)) = (x( f((2) = ED to fi fi()), f E {0,1),
(8)
where (=x(2=) == (p y. t) is the aggregated posterior. We mainly consider the case where is observed in the data, and the sample of (Y,T) is taken from the data given X x. When æ is not in the data, we replace 4s with PA in (8) (see Sec. for details and B for results). Note that in (8) indicates a counterfactual assignment that may not be the same as the factual T = t in the data. That is, we set T f in the decoder. The assignment is not applied to the encoder which is learned from factual X, Y, T (see also the explanation of CCF. in Sec. (42) The overall algorithm
steps are i) train the VAE using 77 and ii) infer CATE f(x) = أi (x) po ( by
4.2 CONDITIONALLY BALANCED REPRESENTATION LEARNING
We formally justify our ELBO 7 from the BRL viewpoint. We show that the conditional BRL via the KL (first) term of the ELBO results from bounding a CATE error; particularly, the error due to the imprecise recovery of j. in (GI') ii controlled by the ELBO, Previous works (Shalit al. 2017
Lu et al. 2020) instead focus on unconditional balance and bound PEHE which is marginalized on connect the bounds to identification and consider noise modeling through (z). Sec Sec. D.3 for detailed comparisons to previous works. In Sec. E.4 we empirically validate bounds, and, particularly, the bounds are more useful under weaker overlap.
X. Sec.
5.2 experimentally shows the advantage of our bounds and ELBO. Further,
We introduce the objective that we bound. Using (8) to estimate CATE, iz) fz) fo(z) is marginalized on (z(2z). On the other hand, the true CATE, given the covariate x or score 2, is:
r(x)) ji (p1 (xx)) (o())(z))) Ty(2) = 5=(=) fo(2),
where ji is associated with an approximate bPGS pe (say, m,) as the target of recovery by our VAE.
Accordingly, given æ, the error of posterior CATE, with or without knowing Pt, is defined
(x) := (g(=) - T(x))²; ef(x) = E (f - (z))²
(10)
We bound ey instead ez because the error between T r(X) and (2) is small-if the score recovery
works well, then z (o(=) (z)) in ( We consider the error between fy and ry below. We
define the risks of outcome regression, into which e is decomposed.
Definition (CATE risks). Let Y()) (X) ~ Pr(ôle; (v|P) and (=(2)x)) = q(=læ, = Bp( The potential outcome loss at (z, f) risk, and counterfactual risk are: (z) = (F,((x) == Cf (); (CF( = ((2,
f((2)) /g((z) = -1(())- f(v f((2)) Pr (vl dyy;
With (()) involved, Lg is a potential outcome loss on f, weighted by g. The factual and counter- factual counterparts, FF,, and ECF., are defined accordingly. In EF., unit u (ar, y, () is involved in the learning of q.(2) )) as well as in Ly t) since (t) = y for the unit. In ECF,t however, unit u (z. y/ t) is involved in (=(xx, but not in Ly since Y(t) + y Y -t)) t).
Thus, the regression error (second) term in ELBO controls eFF via factual data. On the other hand, ECF, is not estimable due to the unobservable Y (1 T)), but is bounded by ep plus M D(x) in Theorem 2 below-which. in turn, bounds cy by decomposing it to f CCF,1, and Vy.
AAs usual, we expect the variational inference and optimization procedure to be (near) optimal: that is, consistency of VAE. Consistent estimation using the prior is a direct corollary of the consistent VAE. See Sec. [C.6 for formal statements and proofs. Under Gaussian models, it is possible to prove the consistency of the posterior estimation, shown in Bonhomme & Weidner (20T
7

---

Published as a conference paper at ICLR 2022
Theorem (CATE enror bound). Assume (2, S M and [g((z)) G, then: cg(a)) = ((G(e + + (2) + (D)))) Vr(x))
(11)
where D(x)= = E. v DKL (9: (2)1 and (x) (((=/) Er Bprcoie jc()))
((x) measures the imbalance between qi (z((z)) and is symmetric for Conrespondingly, the KL Jerm in ELBO 7 is symmetric for L. and balances qe (z)az) by encouraging Z Z for the poste- rior. (ar) reflects the intrinsic variance in the DGP and can not be controlled. Estimating G, M is nontrivial. Instead, we rely on 3 in the ELBO 7 to weight the terms. We do not need two hyperparameters since G is implicitly controlled by the third term, a norm constraint, in ELBO.
5 EXPERIMENTS
We compare our method with existing methods on three types of datasets. Here, we present two experiments: the remaining one on the Pokec dataset is deferred to Sec. FF. As in previous works (Shalit et al.
2017 Louizos et al. 2017), we report the absolute error of ATF Gate == (v(
((0)) - Ept ax) and, as a surrogate of square CATE error Epate (x) = ED] (()()) -(00)-f((a)) f())]²
the empirical PEHE (pehe = Epccate (x) Hill 2011p, which is the average square CATE error.
Unless otherwise indicated, for each function f.g. h. k, r, in ELBO (7), we use multilayer perceptron, with 200 3 hidden units (width 200, layers), and ELU activations (Clevert et al.) 2015). A (h, k) depends only on X. The Adam optimizer with initial learning rate 10-2 and batch size 100 is employed. All experiments use early-stopping of training by evaluating the ELBO on validation set. More details on hyper-parameten and settings are given in each experiment.
5.1 SYNTHETIC DATASET
W|X ~ (((((X), ((X)); T|X (ern(Logi(uu(X))) VIW,T~NUT(W),...
(12)
We generate synthetic datasets following (12). Both X ~ N(,, o) and W are factorized Gaussians. μ., a are ran- domly sampled. The functions h. k. are linear. Outcome models fo: fi are built by NNs with invertible activations. Y is univariate, diim(X= = 30, and dim(W) ranges from 1 to W is bPGS, but the dimensionality not low enough to satisfy the injectivity in (G2) when tim(W) 1. We have different overlap levels controlled by w that multiplies the logit value. See Sec. E. for details and more results syn- thetic datasets.
0.8
0.5
0.3
022
ron-overiop (omega)
With the same (dim(W), w), we evaluate method and CFR on 10 random DGPs, with different sets of functions Sig, h, k in (12). For each DGP, we sample 1500 data points, and split them into 3 equal sets for training, valida- tion, and testing. We show our results for different hyperparameter B. For CFR, we try different balancing parameters and present the best results (see the Appendix for detail).
oft
beta-1 1.5
bata-2.5
bwtow?
bets- 20
betw-3.0
/Eehee synthetic datascts. Each error bar is 10 random DGPs
Figure 2:
In each panel of Figure 2. we adjust one of w, (W)) with the other fixed to the lowest. As implied by our theory, our method, with only I-dimensional Z. performs much better in the left panel (where dim(W in the right panel (when din(W) > 1). Although CFR uses 200-dimensional representation. in the left panel our method performs much better than CFR; moreover, in the right panel CFR is not much better than ours. Further, our method is much more robust against different DGPs than CFR (see the error bars). Thus, the results indicate the power of identification and recovery of scores. (see Figure3jalson
Intaci-VAE
CEVAE
1 satisfies (G2)) than
Figure 3: Plots of recovered true latent. Blue: T 0. Orange: T
Under the lowest overlap level (w = 22), large B(= 2.5, shows the best results, which accords with the intuition and bounds in Sec. 4 When dim(W) 1, fe in (12) is non-injecitve and learning
8

---

Published as a conference paper at ICLR 2022
of PGS is necessary, and thus, larger 3 has negative effect. In fact, 8 = 1 is significantly better than = when (im(W) > 2. We note that our method, with a higher-dimension or matches CFR also under clim(W) > 1 (see Appendix Figure 5 Thus, the performance gap under di)(W) > 1 in F should be due to the capacity of NNs in E-InntcttVAA. In Appendix Figure for ATE error, CFR drops performance w.r.t overlap levels. This is evidence that CFR and its unconditional balance overly focus on PEHE (see Sec. 5.2 for more explicit comparison).
Z, outperforms
When dim 1, there are no better PSs than W, because ft is invertible and no information can be dropped from W. Thus, our method stably learns Z as an approximate affine transformation of the true W, showing identification. An example is shown in Figure and more plots are in Appendix Figure [ For comparison, we run CEVAE, which is also based on VAE but without identification; CEVAE shows much lower quality of recovery. As expected, both recovery and estimation are better with the balanced prior PA (z|az), and we can see examples of bad recovery using (2) (z|ær, in Appendix Figure 10
5.2 IHDP BENCHMARK DATASET
This experiment shows our conditional BRL matches state-of-the-art BRL. methods and does not overly focus on PEHE. The IHDP (Hill 2011) is a widely used benchmark dataset; while it is less known, its covariates are limited-vverlappin
and thus it is used in Pohansson et al. (2020) which considers limited overlap. The dataset is based on an RCT, but Race is artificially introduced as a confounder by removing all treated babies with nonwhite mothers in the data. Thus, Race is highly limited-overlapping, and other covariates that have high correlation to Race, e.g, Birth weight (Kelly al. 2009) are also limited-overlappnng. See Sec. B.2 for detail and more results.
There is a linear bPGS (linear combination of the covariates). However, most of the covariates are binary, so the support of the bPGS often on small and separated intervals. Thus, the Gaussian dimee-idmmniional Z to address this, similar to
latent Z in our model is misspecified. We use Louizos et al. 2017). Specifically, we set dim Z) 50, together with NNs of 50 x 2 hidden units in the prior and encoder. We set B 1 since it works well on synthetic datasets with limited overlap.
As shown in Table [] 3-Intact- VAE outperforms or matches the state-of art methods; has the best performance measured by both Cate and Epehe and matches CF and CFR respectively. Also notably, our method outperforms other generative models (CEVAE and GANITE) by large margins.
To show our conditional balance is preferable, we also modify our method and add two components for unconditional balance from CFR (see the Appendix), which based on bounding PEHE and is controlled by another hyperparameter 7. In the modified version, the over-focus on PEHE of the unconditional balance is seen clearly-with different 7, it significantly affects PEHE, but barely affects ATE error. In fact, the unconditional balance, with larger 7 only worsens the performance. See also Appendix Figure 7 where CFR gives larger ATE errors with less overlap.
Table 1: Errors on IHDP over 1000 random DGPs. "Mod. *" indicates the modified version with unconditional balance of strength 7 Italic indicates where the modified version is significantly worse than the original. Bold indicates method(s) which is significantly better than others. The results other methods are taken from Shalit et al. (2017), except for GANITE and CEVAE, the results of which are taken from original works.
Method
TMLE BNN
CER
CF
CEVAB GANITE Ours
Mod.
Med.0.2
Mod. 0.1 Mod 0.05
Mod. 0.01
30..11 37+.03 .25 .18
34+.01 A34.08 .180- 185+
.85+.00s
.1864 .009
183+.00s .181+.00
Peeke
5.0.a 224.1 .71..00 38_2 2.7.11
194. .709 024 1.175 797 030
.748
029
.332 oas .7192.02
CONCLUSION
We proposed a method for CATE estimation under limited overlap. Our method exploits identifiable VAE, a recent advance in generative models, and is fully motivated and theoretically justified by causal considerations: identification, prognostic score, and balance. Experiments show evidence that the injectivity of fü in our model is possibly unnecessary because dim(2) dim(Y) yields better results. A theoretical study of this is an interesting future direction. We have evidence that Intact-VAE works under unobserved confounding and believe that VAEs are suitable for principled causal inference owing to their probabilistic nature, if not compromised by ad hoc heuristics (Wu & Fukumizu 2021)
9

---

Published as a conference paper at ICLR 2022
REFERENCES
Jason Abrevaya, Yu-Chin Hsu, and Robert Lieli, Estimating conditional average treatment effects. Journal of Business & Economic Statistics, 33(4):485-505, 2015.
Ahmed M Alaa and Mihaela der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In Advances in Neural Information Processing Systems, pp- 3424-3432, 2017.
Timothy B Armstrong and Michal Kolesár. Finite-sample optimal estimation and inference on aver- age treatment effects under unconfoundedness. Econometrica, 89(3):114111177, 2021.
Stéphane Bonhomme and Martin Weidner. Posterior average effects. Journal of Business & Eco- nomic Statistics, (just-accepted::1-22 1-38, 2021.
Victor Chernozhukov and Christian Hansen. Quantile models with endogeneity. Annu. Rev. Econ., 5(1):57-81, 2013.
Denis Chetverikov and Daniel Wilhelm. Nonparametric instrumental variable estimation under monotonicity. Econometrica, 85(4):1303-1320, 2017.
Denis Chetverikov, Andres Santos, and Azeem M Shaikh. The econometrics of shape restrictions. Annual Review of Economics, 10:31-63. 63, 2018.
Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv 2015.
Marco Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
Wangzhi Dai and Collin M Stultz. Quantifying common support between multiple treatment groups using a contrastive-vae. In Machine Learning for Health, pp. 41-52. PMLR, 2020.
Alexander Amour and Alexander Franks. Deconfounding scores: Feature representations for causal effect estimation with weak overlap. arXiv preprint arXiv:2104.05762, 2021.
Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05990 2016.
Alexander D' Amour, Peng Ding, Avi Feller, Lihua Lei, and Jasjeet Sekhon. Overlap in observational studies with high-dimensional covariates. Journal of Econometrics, 2020.
Max H Farrell. Robust inference on average treatment effects with possibly more covariates than observations. Journal of Econometrics, 189(1):1-22) 2015.
Joachim Freyberger and Joel Horowitz. Identification and shape restrictions in nonparametric instrumental variables estimation. Journal of Econometrics, (8:11), 2015.
Li Gan and Qi Li. Efficiency of thin and thick markets. Journal of Econometries, 192(1):40-54 2016.
Prem Gopalan and David M Blei. Efficient discovery of overlapping communities massive networks. Proceedings of the National Academy of Sciences,
2013.
Sander Greenland. The effect of misclassification in the presence of covariates. American journal of epidemiology, 114(4);564-569, 1980.
David Hajage, Yann De Rycke, Guillaume Chauvet, and Florence Tubach. Estimation of conditional and marginal odds ratios using the prognostic score. Statistics in medicine, 36(4):687-716, 2017.
Ben B Hansen. The prognostic analogue of the propensity score. Biometrika, 95(2):481 -488, 2008.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2019.
Miguel A. Hernan and James M. Robins. Causal Inference: What If. CRC Press, 1st edition, 2020.
ISBN 378-1144200-6111
10

---

Published as a conference paper at ICLR 2022
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 constrained variational framework. In 5th International Conference Learning Representations, 2017. URL
net/ formm? Fat
Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Compulational and Graphical Statistics, 20(1):217-240, 2011.
Han Hong, Michael P Leung, and Jessie Li. Inference on finite populationnn treatment effects under limited overlap. The Econometrics Journal, 23(1):32 47, 2020.
Ming- Yueh Huang and Kwun Chuen Gary Chan. Joint sufficient dimension reduction and estimation of conditional and average treatment effects. Biometrika, 104(3):583-596, 2017.
Martin Huber and Kaspar Wüthrich. Local average and quantile treatment effects under endogeneity: review. Journal of Econometric Methods, 8(1), 2018.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sci- ences. Cambridge University Press, 2015.
Dominik Janzing and Bernhard Scholkopf. Causal inference using the algorithmic markov condi- tion. IEEE Transactions on Information Theory, 56(10))5168-5119) 2010.
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.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual infer- ence. In International conference on machine learning, pp. 3020-3029, 2016.
Fredrik D Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 527-536. PMLR, 2019.
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.07426, 2020.
Nathan Kallus, Brenton Pennicooke, and Michele Santacatterina. More robust estimation of sample average treatment effects using kernel optimal matching an observational study of spine surgical interventions. arXiv preprint arXiv:181 (af:::811.0000 1.04274, 2018.
Yvonne Kelly, Lidia Panico. Mel Bartley, Michael Marmot, James Nazroo, and Amanda Sacker. Why does birthweight vary among ethnic groups in the uk? findings from the millennium cohort study. Journal of public health, 31(1):131-137 2009.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. In International Conference Artificial Intel- ligence and Statistics, 2207-2217, 2020a.
Ilyes Khemakhem, Ricardo Monti, Diederik Kingma, and Aapo Hyvarinen. Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. Advances in Neural Information Processing Systems, 33, 2020b.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint arXiv: 14, 2013.
Diederik Kingma, Max Welling, et al. An introduction variational autoencoders. Foundations and Trends@ in Machine Learning, 12(4):307-392, 2019. 2019.
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, pp. 3581-3589, 2014.
11

---

Published as a conference paper at ICLR 2022
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional net-
works. In 5th International Conference on Learning Representations, 2017. URL https:
//openzeview.co http:///ww......oorraaaaaaaaaaaa
Manabu Kuroki and Judea Pearl. Measurement bias and effect restoration in causal inference. Biometrika, 101(2):423 437, 2014.
Jure Leskovec and Andrej Krevl. Snap datasets: Stanford large network dataset collection, 2014.
Arthur Lewbel. The identification zoo: Meanings of identification in econometries. Journal of Economic Literature, 57(4):835-903, 2019.
Fan Li and Fan Li. Propensity score weighting for causal inference with multiple treatments. The Annals of Applied Statistics, ((4))2339-2411 2019.
Zheng Li, Guannan Liu, and Qi Li. Nonparametrie knn estimation with monotone constraints. Econometric Reviews, (6666-9)98881006, 2017.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In Advances in Neural Information Processing Systems, pp. 6446-6456, 2017.
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.
Wei Luo. Yeying Zhu, and Debashis Ghosh. On estimating regression-based causal effects using sufficient dimension reduction. Biometrika, 104(1);51-65, 2017.
Emile Mathieu, Tom Rainforth, Nana Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In International Conference on Machine Learning, pp. 4402-4412. PMLR, 2019.
Wang Miao, Zhi Geng, and Eric J Tchetgen Tchetgen. Identifying causal effects with proxy variables of an unmeasured confounder. Biometrika, 105(4):987-993, 2018.
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, pp. 788-798. PMLR, 2020.
Judea Pearl. Causality: models, reasoning and inference. Cambridge University Press, 2009.
Severi Rissanen and Pekka Marttinen. A critical look the identifiability of causal effects with deep latent variable models. NeurlPS 2021, to appear, 2021.
Paul R Rosenbaum. Modern algorithms for matching in observational studies. Annual Review of Statistics and Its Application, 7:143-176, 2020.
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.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-33) 2005.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effeet: gener- alization bounds and algorithms. In International Conference on Machine Learning, pp. 3076- 3085. PMLR, 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In Advances in Neural Information Processing Systems, pp. 2507-2517, 2019,
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, pp. 3483-3491, 2015.
12

---

Published as a conference paper at ICLR 2022
Peter Sorrenson, Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ica with general incompressible-foor networks (gin). In International Conference on Learning Representations, 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way prevent neural networks from overfitting. The journal of machine learning research, 15()):12291958, 2014.
Jennifer E Starling, Catherine E Aiken, Jared Murray, Annettee Nakimuli, and James Scott. Monotone function estimation in the presence of extreme data coarsening: Analysis of preeclamp- sia and birth weight in urban uganda. arXiv preprini arXiv: htti://92.000// 2.0066946, 2019.
Elizabeth A. Stuart. Matching Methods for Causal Inference: A Review and a Look Forward Statistical Science, 25(1):1 - 21, 2010. doi: 10.1214/09-STS113. URL httttps:/ddi.org//
10.1214/03 STS313.
Xinwei Sun, Botong Wu, Chang Liu, Xiangyu Zheng, Wei Chen, Tao Qin, and Tie-yan Liu. Latent causal invariant model. arXiv preprint (iii:22011002000000 2020.
Alexander Tarr and Kosuke Imai. Estimating average treatment effects with support vector ma-
chines. arXiv preprint arXiv:2102.1 1926, 2021.
Mark J Van der Laan and Sherri Rose. Targeted learning in data science: causal inference for complex longitudinal studies. Springer. 2018.
Victor Veitch, Yixin Wang, and David Blei. Using embeddings to correct for unobserved confound- ing in networks. Advances in Neural Information Processing Systems, pp. 13792-13802, 2019.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneoous treatment effects using random forests. Journal of the American Statistical Association, 1236622) 2018.
Shanshan Wang, Liren Yang, Li Shang. Wenfang Yang, Cuifang Qi, Liyan Huang, Guilan Xie, Ruiqi Wang. and Mei Chun Chung. Changing trends of birth weight with maternal age: a cross- cress-nectional study in xi' an city of northwestern china. BMC Pregnancy and Childbirth, 20(1):1-8, 2020.
Halbert White and Karim Chalak. Identification and identification failure for treatment effects using structural systems. Econometric Reviews, 32(3):273-317, 2013.
Pengzhou Wu and Kenji Fukumizu. Causal mosaic: Cause-effect inference via nonlinear ica and
ensemble method. In International Conference on Artificial Intelligence and Statistics, pp. 1157- 1167. PMLR, 2020a. httt:ttt::///cceeeeeeee ngs mm. r. wew.rppreeaee prwpsppss//..... /vsssvvll/ 08 wu2 ht.ml
Pengzhou Wu and Kenji Fukumizu. Towards principled causal effect estimation by deep identifiable models. arXiv preprint 1rX.v:2109.15002 2021.
Pengzhou Abel Wu and Kenji Fukumizu. Identifying treatment effects under unobserved con- founding by causal representation learning. submitted ICLR 2021, 2020b. URL https:
T/openzeview. l/ogeneeview.ne///ormmtrt http://fooummm----iiieeee D3TNqCapFpM
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. ISSN 0006-3444. doi: 10. 10.000///iomee//sy000... URL httpthttt///oiooootttttttt
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learn- ing for treatment effect estimation from observational data. In Advances in Neural Information Processing Systems, pp. (633-264)) 2018.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. In International Conference Learning Rep- resentations, 2018 8. URL fhtts::pps..... : roverrevevveeeeeeoeeee ew net/ Borrmmidd-ByKWeeee
Weijia Zhang, Lin Liu, and Jiuyong Li. Treatment effect estimation with disentangled latent factors. arXiv preprint arXiv:2001. 202220001.652, 2020a.
Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the esti- mation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1005-1014, PMLR, 2020b.
13

---

Published as a conference paper at ICLR 2022
A PROOFS
We restate our model identifiability formally.
Lemma (Model identifiability). Given model (3) under (M1) for T assume
(D1') (Non-degenerated data for A) there exist 2n + points 2o...., 22n X' such that the 2n-square matrix L, == [Ti,,, ...: 2n,2n] is invertible, where Yt,k (x(x) Ax
Then, given T = t, the family identifiable up to an equivalence class. That is, if po pol())) = pe ( x. t), we have the relation between parameters: for any yt in the image of ft.
f.² (3)) = diag(a) (y)) + bb= A.ft
(13)
where diag(a) an invertible n-diagonal matrix and is n-vector, both depend on X and X,.
Note, (D1) in the main text implies (D1') see Sec. B.2.3 in Khemakhem et al. (2020a). The main part of our model identifiability is essentially the same as that of Theorem in Khemakhem et al. (2200)). but now adapted to include the dependency on f. Here we give an outline of the proof, and the details can be easily filled by referring Khemakhem et al. (2020a). In the proof, subscripts are omitted for convenience.
Proof of Lemma ! Using (11) i) and ii), we transform Pf Pf.A 2, Pf (x into equality
of noiseless distributions, that is,
Of (v) = (F) =
() (1) TTy
(14)
where PA is the Gaussian density function of the conditional prior defined and vol(.A) =
det AAT. !f is defined similarly to Jf.A-
Then, apply model (3) to (14), plug the 2n points from (D1') into it, and re-arrange the resulting 2n + 1 equations in matrix form, we have
= F(Y) = L' - B
(15)
where t(Z) = (Z, ( is the sufficient statistics
factorized Gaussian, and B, == (ax(x1) -
or or(2z)); .... Or2 (ar2n) (xx))) where o, ax;X;A, Ax) a)) is the log-partition function of the conditional prior in 3 F is defined similarly to F, but with f ,',
Since L is invertible, have
1(ʃ (Y)= At(" +e
(16)
where A L and c L
The final part of the proof is to show, by following the same reasoning as in Appendix B of Sorrenson et al. (2019), that A is sparse matrix such that
A diag a) diag(u) u) diag(a?)
(17)
where is partitioned into four ne-quare matrices. Thus
+
(18)
where b is the first half of c.
Proof of Proposition Under
G2) and (33)
we have
Epo (") = LY IXT))
ft h(w)
je p(x) on (x, t) such that p(t, > 0. (19)
We show the solution set of (19) on overlapping æ is
((f.h)]f+ = j. A-
Aop,A. P Ra injective).
(20)
By (G2)
M1) and with injective ft, j. and dim (7) dim dYm > dim(p). for any A above, there
exists a functional parameter ft such that j. = f A.A. Thus, set (20) is non-empty, and any element
indeed solution because f. h j. A A p = j. p.
Any solution of (19) should be in (20). A solution should satisfy h((x) = f: o j. o p(x) for both t since æ is overlapping. This means the injective function f. j. should not depend on t, thus is one of the A in (20).
14

---

Published as a conference paper at ICLR 2022
We proved conclusion 1) with A. And, on overlapping æ, conclusion 2) is quickly seen from
[x((x) fo ((())) = je o (v o p(=)) = (:(()=)) = (xx(.))
(21)
We rely on overlapping p to work for non-overlapping æ. For any a, with p(1 - -t/x,) 0, to ensure
p(1 ((1-the =lp((z)) > there should exist - such that p(æ)-t) = p(x)) and p(1 1| > And
we also have h( = h(x,) due to (M2) Then, we have
i (x.) fi-i(((x;)) = f (h 21
j)-e(((m)-()) -)-r(p(x=))) mm))( (22)
The third equality uses 19 on (x1-1: t).
Below prove Theorem with (D2) replaced by
(D22) (Spontaneous balance) there exist 2n + points To, D2n E X, 2n-square matrix C,
and 2n-vector d, such that L L, C. and G0 C-TB, d/k for optimal X (see
below), where Lu is defined in (DI") Be == (a1(x11) a (x)))) ... (x(2))) ax(x))))
and (x; Xx) is the log-partition junction of the prior in (3)
(D2') restricts the discrepancy between Xo A1 2n. 1 values of X, thus is relatively easy to satisty with high-dimensional X. (D2') is general despite (or thanks to) the involved formulation. 0. Then, Lg 11 cl requires that, / (2x)) cho (2tx) ) is the same for 2n + 1 points Øx. This is easily satisfied except for > m where m is the dimension of X, which rarely happens in practice. And, 30
Let us see its generality even under highly special case: C c=aa cl and
31 -
becomes just 31 cB.. This is equivalent to a1 (æx) c0o(2x) same for 2n + points, again fine in practice. However, the high generality comes with price. Verifying (D2') using data is challenging, particularly with high-dimensional covariate and latent variable. Although we believe fast algorithms for this purpose could be developed, the effort would be nontrivial. This is another motivation to use the extreme case Xo X in Sec. 4.1 which corresponds to C' II and d d. 0.
for any injective function A P + R", there exists a functional parameter fü such that J. f A. Let hj A Pe, then, clearly from (M3') such parameters 0 = (f*, ; h² are optimal: Po- (y |z, () t)
Proof of Theorem 7 By (M1) and (G))
Since have all assumptions for Lemma 1 we have
Aoj j f = Ao (v) l, on (y. E {(i t)l((t, x) > 0}, where f is any optimal parameter, and "I," collects all subscripts t. Note, except for A, all the symbols should have subscript I.
(23)
Nevertheless, using (D2*) we can further prove Ag = Ai-
We repeat the core quantities from Lemma
here: A, L,"" L and c, L Bt).
From (D2') we immediately have
Ly L1
L.
1/ CC > Ao Ai
(24)
And also,
LØ L C + Ly L TT = -
(25)
-
CTT -
d/k + Cº (B, - (8) -Bi- 31-B3 -31
Multiply right hand sides of the two lines, we have Co = C₁ Now we have Ao = Ai == A. Apply this to 23), we have
f. = j. for any optimal parameters 0 = (f, h). Again, from (M3') we have
= A o A
(26)
(o( p(y) (z,
Pe (h, (x)) (x)) Pely (=(p)()))
(27)
where Pr. = Pe- And the above only possible when oo = j. Combined with f: j:00-1 we have conclusion 1).
And conclusion 2) follows from the same reasoning as Proposition] 2 applied to both po and P1-
Note, when multiplying the two lines of 25 the effects of l 0 cancel out, and c is finite and well-defined, Also, it is apparent from above proof that (D2") is a necessary and sufficient condition for An AI, if other conditions of Theorem are given.
15

---

Published as a conference paper at ICLR 2022
4.2 The definitions and results work for the prior; simply replace qe qe [)) with (z x) = PA ( z. t) in definitions and statements, and the proofs below hold f prevail, and the superscripts omitted. The arguments
Below, prove the results in Sec.
as the same. The dependence
sometimes also omitted.
Lemma (Counterfactual risk bound). Assume (z () S M, we have <CF(x) s E, q(1 (100)ep;((0)) (D)
(28)
where (CF(2) = E, p(1 -1-=)///));(x))) and D(x) = (1
Proof of Lemma
1p(1-
ECF
= =((0/))(cFF1 (FA1 -..1) p(1l=))ccFF- - cF,D) -0(0))) / Lz(z,
- (a( (((2)))( go(=l=))ddz
- 2MTV((1,40) & MD.
TV(,, = [ - {()) is the total variance distance between probability density The last inequality uses Pinsker's inequality TV(p, q) S (Du twice, get the symmetric D. Theorem 2 is a direct corollary of Lemma and the following.
Lemma 3. Define eF E, p(t]x)eFF We have
ef S 2(G(eF +CP) -vy)
(29)
Simply bound ECF in (29) by Lemma 2 we have Theorem 2 bias-variance decomposition eF and ECF.
To prove Lemma 3 we first examine a
ECF.1 En (z)))) (2) (1)z) (
2 GE - (=(m)) ( (z)=) (y- f((z))²
(30)
GE (x|m) Epr (v)z) ((y j=(z)) + ( f((2))
The second line uses (z)l < G. and the third line is virianer decomposition. Now we can define VCF,(x)) = En (= m) pp (1)z) jz))) and (CF,(()) = Ey (zx)))
fe(z))², and we have
Bcr,((x) = cCF 2 G(Vcr(x) + Bcr(x)) where VCF = E. - ( = E. (la) (w) (v)=)(y - fr(z))² and similarly Bcr E, () - f((2))2 Repeat the above derivation forer, er, we have
CCP, 2 G(V
(31)
EF G(V)(w)) + (x)))
(32)
where Vr Ee (z, F E (v)=) are ready to prove Lemma
(2)) (v)))) and Br E, (e f((2)) Now,
Proof of Lemma
f8 By(x Φ) ((fi - fo) (ji jo))2
Eg((f -11) + (jo fo))²
21, ((fi -i1)² + (Jo fo)
-2 / (fi -ii) 1/x) + (jo (21(000)z)+
(fi - 11) 0/z) + (jo fo) 2)), 1/x)d
2BB + BcF) 2(G(EF + Vy)
16

---

Published as a conference paper at ICLR 2022
The first inequality uses (a 6)2 < 2(a² + b2). The next equality splits g(z[x=) into g(z, 0/z) and q( 1/zz) and rearranges to get Bp and BeF. The last inequality uses the two bias variance decompositions, and Vy = VF + VcF.
B ADDITIONAL BACKGROUNDS
B.I PROGNOSTIC SCORE AND BALANCING SCORE
In the fundamental work of (Hansen 2008), prognostic score is defined equivalently Po (PO- score), but it in addition requires no effect modification to work for (1). Thus, a useful prognostic score corresponds our PGS. We give main properties of PGS as following.
Proposition 3. V gives exchangeabiliry, and (V)) is PGS, then (t)
The following three properties of conditional independence will be used repeatedly in proofs.
Proposition 4 (Properties of conditional independence). (Pearl 2009 Sec. 1.1.55) For random
variables W, X, Y, Z. We have:
Z A X YY, Z X XJL,W Y (Contraction)..
XHW, V = XJ,,WW, Z (Weak union).
X YIZ XJILYIZ (Decomposition).
Proof of Proposition From LLVV (exchangeability of V), and since P. is a function of V, we have Y (I) Pt. (1).
From (1) and (())_LV(pp(V) (definition of Pt-score), using contraction rule, we have ((I)LL,, (p. for both
Vp,
Prognostic scores are closely related to the important concept of balancing score (Rosenbaum & Rubin 1983). Note particularly, the proposition implies (t) (using decomposition rule). Thus, (VV) is P-score, then p also gives weak ignorability (exchangeability and overlap), which a nice property shared with balancing score, as we will see immediately.
Definition 4 (Balancing score). b(V), function of random variable V, is balancing score if b(V).
Proposition 5. Let b(V) be a function of random variable V. b(V) ) is a balancing score if and only if f(b(V)) p(T' = 1[V) = e(V) for some function f (or more formally, e(V) is b(V)- measurable). Assume further that V gives weak ignorability, then so does b(V).
Obviously, the propensity score e(V) := p(T 1|V), the propensity of assigning the treatment given V. balancing score (with f be the identity function). Also, given any invertible function v, the composition b is also balancing score since f (n b(V)) = f(b(V)) = e(V).
Compare the definition of balaneing score and prognostic score, we can say balancing score is sufficient for the treatment T (T (TV),) while prognostic score (Pt-score) is sufficient for the potential outcomes (t) (Y (L
They complement each other; conditioning on either deconfounds the potential outcomes from treatment, with the former focuses on the treatment side, the latter on the outcomes side.
B.2 VAE. CONDITIONAL VAE, AND IVAE
VAEs Kingma et al. 2019) are a class latent variable models with latent variable Z, and observ- able Y is generated by the decoder pe pe(y) y 2). In the standard formulation (Kingma & Welling 2013), the variational lower bound L(y; 0, Φ) the log- ligeliiil derived as:
log p(y) 2 log p(y) (x (p(()))))))))))
(33)
E - log po ( (xs ( llx())))
where DKL denotes KL divergence and the encoder ( by) introduced approximate the true posterior p(=ly). The decoder Po and encoder 9s are usually parametrized by NNs. We will omit the parameters 0, in notations when appropriate.
17

---

Published as a conference paper at ICLR 2022
The parameters of the VAE can be learned with stochastic gradient variational Bayes. With Gaussian latent variables, the KL term L has closed form, while the first term can be evaluated by drawing samples from the approximate posterior 96 using the reparameterization trick (Kingma & Welling 2013). then, optimizing the evidence lower bound (ELBO) (By))(((y)) with data D, we train the VAE efficiently.
Conditional VAE (CVAE) (Sohn et al. 2015 Kingma et al. 2014) adds a conditioning variable C, usually class label, to standard VAE (See Figure D. With the conditioning variable, CVAE can give better reconstruction of each class. The variational lower bound is
log p(y c) 2 logpllo.
- (x y. c)llp(z|c)).
(34)
The conditioning on C' in the prior is usually omitted (Doersch) 2016). i.e., the prior becomes ~ (0,, I) in standard VAE, since the dependence between C' and the latent representation is also modeled the encoder q. Moreover, unconditional prior in fact gives better reconstruction because encourages learning representation independent of class, similarly to the idea of beta- VAE (Higgins cet al. 2017).
As mentioned, identifiable VAE (iVAE) (Khemakhem et al. 2020a) provides the first identifiability result for VAE, using auxiliary variable X. assumes Y 2 that is, p(y|z, x) = p(y|z). The variational lower bound is log log p(ylz))) 2 log p
x))x22/y x))
(35)
= Eznn logpy(y(z)) DRr. (g( ly, x)) A
where = f(Z) e, is additive noise, and Z has exponential family distribution with sufficient statistics T and parameter A(X). Note that, unlike CVAE, the decoder does not depend on X due to the independence assumption.
Here, identifiability of the model means that the functional parameters (f.T. A) can be identified (learned) up to certain simple transformation. Further, in the limit of - 0, iVVE solves the nonlinear ICA problem of recovering = f-)
C EXPOSITIONS
The order of subsections below follows that they are referred in the main text.
C.1 LIST OF ASSUMPTIONS
The following is a list of assumptions required by our identification theory, with comments on their roles and subtleties.
(G1)
additive noise model is needed to ensure the existence of PtSs. (G1) is equivalent to (G1) and
introduced for better presentation, e.g., it connects to (G2) and (M1) through injectivity.
(M1) and (D1) are inherited from iVAE and are required for model (parameter) identifiability (iden- tifying fz up affine mapping), which does not imply CATE identification in general. Arguably here the most important is that the mapping fü from latent Z to outcome Y is injective, or else some information of Z is in principle unrecoverable. These two conditions are not required by Proposition 2 which does not need model identifiability.
(M2) together with overlapping PiSs, is important to address limited overlap of X and can be seen weak form of OOD generalization.
(M3') means 1) we need to know or learn the distribution of hidden noise e and 2) noiseless prior. This simplifies the proof of identification, but when implementing the VAE as estimation method, both noises are learned.
(D2) in fact (D2') strengthens the model identifiability determine both fo and fi up to the same affine mapping, which replaces the balance PS.
(22) is required by Proposition but not Theorem I It is less important than GG)) because
the core intuition of our method is that (G2) should hold approximately. Sec. C,3 contains several
detailed real-world examples on (22
18

---

Published as a conference paper at ICLR 2022
C.2 DETAILS AND EXPLANATIONS ON INTACT- VAE
Our goal is to build model that can be learned by VAE from observational data to obtain a PGS, or more idcally bPGS, via the latent variable Z. That is, a generative prognostic model. Generative models useful to solve the inverse problem of recovering PGSs.
With the above goal, the generative model of our VAE is built as (3). Conditioning on X in the joint model ply, = 2/2, t) reflects that our estimand is CATE given X. Modeling the score by conditional distribution rather than a deterministic function is more flexible.
The ELBO of our model can be derived from standard variational lower bound as following:
log ar t) 2 log ( t) (x y, ()),p(z/ y. t))
(36)
Ex poy z. t) DKL(( y. 12((2))) t))
We naturally have an identifiable conditional VAE (CVAE), as the name suggests. Note that ( has a similar factorization with the generative model of iVAE (Khemakhem et al. 2020a), that is p(y. - p(y) z ))(==()); the first factor does not depend on X. Further, since we have the conditioning on T in both the factors of (3), our VAE architecture is a combination of iVAE and CVAE (Sohn al.) 2015] Kingma et al., (014). with T 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 (see Lemma Tp.
Please do not confuse the DGP (G1) and the generative model B3 of Intact- VAE. The former is the causal model, but the latter is not (at least before we show the TE identifications in Sec. 3.2. In our case, the generative model is built as a way to learn the scores through the correspondence to (
In particular, note that conditionally balanced representation Z is possible under the gen- erative model. This requires a violation of causal faithfulness, so that there are other conditional independence relations, which are not generally implied by the graphical model. Our method, based on iVAE, which achieves ICA, performs nonlinear ICA to recover the scores. In fact, ICA proce- dures often violate causal faithfulness, because requires finding causes from effects. Also, the violation of causal faithfulness is not caused by the generative model (which is shown in Figure I], because the representation is learned by the encoder, and Z X is enforced by B.
C.3 DISCUSSIONS AND EXAMPLES OF G2)
We focus on univariate outcome on RR which is the most practical case and the intuitions apply to more general types of outcomes. Then, i, the mapping between 140 and μl, is monotone, i.e, either increasing or decreasing. The increasing means, if a change of the value of X increases (decreases) the outcome in the treatment group, then it is also the case for the controlled group. This is often true because the treatment does not change the mechanism how the covariates affect the outcome, under the principle of "independence of causal mechanisms (ICM)" (Janzing & Scholkopi 2010) The decreasing corresponds to another common interpretation when ICM does not hold. Now, the treatment does change the way covariates affect Y, but in globul manner: it acts like "switch" on the mechanism: the same change of X always has opposite effects on the two treatment groups
We support the above reasoning by real world examples. First we give two examples where μo and μn are both monotone increasing. This, and also that both μ4 are monotone decreasing, are natural and sufficient conditions for increasing i, though not necessary. The first example is form Health. (Starling et al. 2019) mentions that gestational age (length of pregnancy) has a monotone increasing effect on babies birth weight, regardless of many other covariates. Thus, if we intervene on one of the other binary covariates (say, == receive healthcare program or not), both μ4 should be monotone increasing in gestational age. The next example from economics. (Gan & Lij 2016) shows that job-matching probability is monotone increasing in market size. Then, we can imagine that, with t = receive training in job finding or not, the monotonicity is not changed. Intuitively, the examples corresponds two common scenarios: the causal effects are accumulated though time (the first example), or the link between a covariate and the outcome is direct and/or strong (the second example).
Examples for decreasing i are rarer and the following is a bit deliberate. This example is also about babies' birth weight as the outcome. (Abrevaya al. 2015) shows that, with mother smokes
19

---

Published as a conference paper at ICLR 2022
or not and X mother's age, the CATE r(x) is monotone decreasing for 20 < 26 (smoking decreases birth weight, and the absolute causal effect is larger for older mother). On the other hand, it is shown that birth weight slightly increases (by about 100g) in the same age range in surveyed population (Wang et al. 2020). Thus, it is convince that, smoking changes the the tendency of birth weight w.r.t mother's s age from increasing to decreasing, and gives the large decreasing of birth weight (by about 300g) as its causal effect. This could be understood: the negative effects of smoking on mother's heath and in turn on birth weight are accumulated during the many years of smoking.
C.4 COMPLEMENTARITY BETWEEN THE TWO IDENTIFICATIONS
We examine the complementarity between the two identifications more closely. The conditions
(M3) (M3') and (G2) (D2' form two pairs, and are complementary inside each pair. The first
pair matches model and truth, while the second pair restricts the discrepancy between the treatment groups. In Theorem (G2) (Po P1) is replaced by (D2") which instead makes Ao A = A in
(5) And (D2") is easily satisfied with high-dimensional X. even if the possible values of C, d are restricted C cl and d 0 (see below). On the other hand, Pe Pe in (M3') is impractical,
but it ensures that po (z, t) = p(y æ. t) so that (5) can be used. In Sec. 4.1 we consider practical estimation method and introduce the regularization that encourages learning a PGS similar to bPGS so that Pa-
Pe can be relaxed.
C.5 IDEAS AND CONNECTIONS BEHIND THE ELBO
Bayesian approach is favorable to express the prior belief that bPGSs exist and the preference for them, and to still have reasonable posterior estimation when the belief fails and learning general PGS is necessary. This is the causal importance of VAE as an estimation method for us. By the un- conditional but still flexible A, and also the identifications, the ELBO encourages the recovery of an approximate bPGS as the posterior, which still learns the dependence on T if necessary. Moreover, B expresses our additional knowledge (or, inductive bias) about whether or not there exist approximate bPGSs (e.g., from domain expertise).
In fact, 3 connects our VAE to 3-VAE (Higgins et al. 2017). which is closely related to noise and variance control (Doersch 2016, Sec. M.thhet al. 2019).
Considerations on noise modeling. In Theorem [ with large and mismatched noises (then (M3' is casily violated), the identification of outcome model f.
vー would fail, and, in turn, the prior would learn confounding bias, by confusing the causal effect of on PT and the correlation between T and X. This is another reason to prefer Ao Y1, besides balancing. On the other hand, the posterior conditioning on Y provides information of noise c, and it is shown in (Bonhomme) & Weidner 2 that posterior effect estimation has minimum worst-case error under model mis- specification (of the noise and prior, in our case).
Under large e, relatively small / implicitly encourages 9 smaller than the scale of e, through stressing the third term in ELBO (7). And the the model as whole would still learn p (y/ f) well, because the uncertainty of e can be moved to and modeled by the prior. This is why to zero because learnable prior noise (variance) allows us to implicitly control via 8. Intuitively, smaller 9 strengthens the correlation between and Z in our model, and this naturally reflects that posterior conditioning on Y is more important under larger e. Hopefully, precise learning of outcome noise is not required, as Proposition 2
not set
Now, it is clear that 3 naturally controls at the same time noise scale and balancing. And the regularization can also be understood as an interpolation between Proposition 2 and Theorem relying on bPGS, or model identifiability; learning loosely, preciscly, the outcome regression. When the noise scale is different from truth, there would be error due to imperfect recovery of j. Sec. shows that this error and balancing form trade-off, which is adjusted by /3
D
Importance of balancing from misspecification view. we must learn an unapproximate bPGs, have larger misspecification under balanced prior and rely more on Y in the posterior. Both are bad because it is shown in (Bonhomme & Weidner (20T)) that posterior only helps under bounded (small) misspecification, and posterior estimator has higher variance than prior estimator (see below
20

---

Published as a conference paper at ICLR 2022
for extreme case). Again, want regularizer to encourage learning of bPGS, so that we explore the middle ground: relatively low-dimensional p, or relatively small e.
Example. Assume the true outcome noise is (near) zero. By setting - in model, the poste- rior (= xz, y, t) = Poly, x| t)/polyla t) degenerates to f (Y) (=))= fr (jr ( (pr), a factual PGS. However, -Y- (
" jr(pr)) + (p-rr)) the score recovered by posterior does work for counterfactual assignment! The problem is, un- like X, the outcome Y = Y(T) is affected by T. and, the degenerated posterior disregards the information of X from the prior and depends exclusively on factual (Y,T).
C.6 CONSISTENCY OF VAE AND PRIOR ESTIMATION
The following is refined version of Theorem in Khemakhem et al. (2020a). The result is proved by assuming: i) our VAE is flexible enough to ensure the ELBO is tight (equals to the true log likelihood) for some parameters; ii) the optimization algorithm achieve the global maximum of ELBO (again equals to the log likelihood).
Proposition 6 (Consistency of VAE)., Given model (3& (6) and let pa 1, be the true observational distribution, assume
(z)ae, t) and po po(z|æ, y. t) = gz(2 (z y.t): ii) the ELBO ED- (Clæ, y, t; O, x)) 4 can be optimized to its global maximum at (0', d');
i) there exists (0, 6) such that pa palyla
Then, the limit of infinite data po ( x. 1) = p (y az, 1) and po (z|az y, = 2a (2| y,
Proof. From i), we have C(x, y, 1; の) log p" (ylæ xe 1). But we know upper-bounded by log p" (z/lx, t). So, EDwp' p"ogp" (ylx, t)) should be the global maximum of the ELBO (even if the data is finite).
Moreover, note that, for any (0, Φ). we have DKL (Ko(p(= y. (. (=|æ, y.t) and, in the limit of infinite data, EDwg' (log po xz t)) I EDD p'ogp" (ylæ, t)). Thus, the global maximum of ELBO is achieved only when po (yl (x. p" (y az, and po (2/2x y, go (2 t).
Consistent prior estimation of CATE follows directly from the identifications. The following is corollary of Theorem
Corollary 1. Under the conditions of Theorem 1 further require the consistency of Intact- VAE.
Then, in the limit of infinite data, we have (X) fc(h(()) where f. are the optimal param- eters learned by the VAF.
C.7 PRE/POST-TREATMER PREDICTION
Sampling posterior requires post-treatment observation (y, t). Often, it is desirable that we can also have tre-treatment prediction for
subject, with only the observation of its covariate X x. To this end, we use the prior as a pre-treatment predictor for Z: replace d with PA in and get rid of the outer average taken on D; all the others remain the same. We also have sensible pre-treatment prediction even without true low-dimensional PSs, because PA gives the best balanced approximation of the target PGS. The results of pre-treatment prediction are given in the experimental section
D MORE ON RELATED WORK
D.1 CFR AND CEVAE
CFR and CEVAE are well-known machine learning methods for CATE estimation. Here we make detailed comparisons to them.
D.1.1 COMPARISONS WITH CFR
Our method is related to CFR in two ways. Theoretically, our bounds in Sec. 4.2 resemble those in Shalit al. (201)). But bound CATE error, while CFR bounds PEHE; thus, our bounds give
21

---

Published as a conference paper at ICLR 2022
conditional balancing while CFR only has unconditional balancing. See Sec. D.3 for more the bounds. Conceptually, CFR loosely related to our method because also learns a representation outcome predictor, mentioned in the follow-up Johanssson al. (2020). However, CFR does not have a generative model, so their representation is not formally related to PGSs. Moreover, CFR does not account the outcome noise, while the uncertainty due to the noise is accounted by VAE.
D.1.2 COMPARISONS WITH AND CRITICISMS OF CEVAE
Motivation CEVAB is motivated by exploiting proxy variables, and its intuition that the hidden confounder U can be recovered by VAE from proxy variables.
Our method is motivated by prognostic scores (Hansen 2008), and our model is directly based on equations (2) which identifies CATE. There is no need to recover the hidden confounder in our framework.
Architecture Our model naturally based on (2) particularly the independence properties of PGS. And consequence, our VAE architecture is natural combination of iVAE and CVAE (see Figure D. Our ELBO 4 is derived by the standard variational lower bound.
On the other hand, the architecture of CEVAE is more ad hoc and complex. Its decoder follows the graphical model of descendant proxy mentioned above, but adds an ad hoc component to mimic TARnet (Shalit et al. 2017): it uses separated NNs for the two potential outcomes. We tried this idea on the IHDP dataset, and, as we show in Sec. 5.2 it has basically no merits for our method, because have principled way for balancing.
The encoder of CEVAE is even more complex. To have post-treatment estimation, ( and q(||X, 7") added into the encoder. As result, the ELBO of CEVAE has two additional like- lihood terms corresponding to the two distributions. However, in our Intact-VAE, post-treatment estimation is given naturally by our standard encoder, thanks the correspondence between our model and
Justification We have given the identifications and bounds of method in this paper. Moreover, we carefully distinguish assumptions on the DGP and assumptions on our model, and identify the assumptions that important for causality. There few theoretical justifications for CEVAR. Their Theorem I directly assumes the joint distribution p(æ, y, ł, u) including hidden confounder U is recovered, then identification is trivial by using the standard adjustment equation.
However, the challenge is exactly that the confounder is hidden, unobserved. Many years of work have been done in causal inference derive conditions under which hidden confounder can be (partially) recovered (Greenland) 1980, Kuroki & Pearl 2014: Miao et al. 2018). In particular, Miao al. (2018) gives the most recent identification result for proxy setting, which requires very specific two proxies structure, and other completeness assumptions on distributions. Thus, it is unreasonable believe that VAE, with simple descendant proxies, can recover the hidden confounder. Indeed, Rissanen & Marttinen (2021) recently give evidence that the method often fails.
Moreover, the identifiability of VAE itself a challenging problem. As mentioned Introduction, Khemakhem et al. (2020a) is the first identifiability result for VAF, but it only identifies equiva- Ience class, not a unique representation function. Thus, it is also unconvincing that VAE can learn unique latent distribution, without certain assumptions. As simple synthetic datasets, CEVAE can not robustly recover the hidden confounder, even only up to transformation, while our method can (though, again, this is not needed for our method).
show in Sec. 5.1 for relatively
D.2 INJECTIVITY, INVERTIBILITY, MONOTONICITY, AND OVERLAP
Let us note that any injective mapping defines invertible mapping, by restrict the domain of the inverse function to the range of the injective mapping. Also note that injectivity is weaker than monotonicity; monotone mapping can be defined by an injective and order-preserving mapping between ordered sets. Particularly, an injective and continuous mapping on is monotone, and many works in econometries give examples of this case.
22

---

Published as a conference paper at ICLR 2022
Many classical and recent works (with many real world applications, C.1) in econometries based on monotonicity. Particularly, there is long line of work based on monotonicity of treatment (Huber & Wüthrich 2018). More related to our method is another line of work based monotonic- ity of outcome, see Chernozhukoy & Hansen, 2013) and references therein for early results. Some recent works apply monotonicity of outcome to nonparametric IV regression (NPIV) (Freyberger & Horowitz 2015; Li et al., 2017 Chetverikov & Wilhelm (017), where the structural equation of the outcome is assumed to be continuous. Particularly, (Chetverikov & Wilhelm 2017) combines monotonicity of both treatment and outcome, and (Freyberger & Horowitz 2015) considers discrete treatment (note continuity or differentiability is not necessary for monotonicity). NPIV with monotone / is closely related to our method, but the difference that replaced by a PGS in our method, and the PGS is recovered from observables. Finally, as we mentioned in Sec. 3.2 monotonicity is kind of shape restriction which also includes, c.g., concavity and symmetry and attracts recent interests (Chetverikov et al., 2018). However, most of NPIV works focus identifying f but not directly on TEs, and we do not know any works that use monotonicity to address limited overlap.
f(T) + e, and f is monotone and ' (the treatment) is often
Recently in machine learning, (Johansson al. 2019 Zhang et al. 2020b Johansson et al. 2020) note the relationship between invertibility and overlap. As mentioned, (Johansson et al. (0200 gives bounds without overlap, but the relationship between invertibility and overlap is not explicit in their theory. (Johansson et al. 2019) explicitly discuss overlap and invertibility, but does not focus on TEs. Zhang et al.) (2220) assumes overlap so that identification is given, and then focuses on learning overlapping representation that preserves the overlapping the covariate. However, it does not re- late invertibility and overlap, but uses invertible representation function preserve exchangeability given the covariate, and linear outcome regression to simply the model. Related, our identifications required (M2) of which linearity of PGS and representation function sufficient condition, and our outcome model is injective, to preserve the exchangeubility given the PGS. Thus, our method works under more general setting, and arguably under weaker conditions.
D.3 ADDITIONAL NOTES ON NOVELTIES OF THE BOUNDS IN SEC.
We give details and additional points regarding the novelties, Lu et (0220) also use a VAE and derive bounds most related to ours. Still, our method strengthens Lu et al. (2000). in a simpler and principled way: we distinguish true score and latent Z and show that identification the link; considering both prior and posterior, we show the symmetric nature of the balancing term and re- late it to our KL term in (7). without ad hoc regularization; moreover, we consider outcome noise modeling which is strength of VAE and relate it to hyperparameter B. Particularly, in (Lu et al.) 2020). latent variable Z is confused with the true representation (p. up invertible mapping in our case). Without identification, the method in fact has unbounded error. Note that Shalit et al. (2017) do not consider connection to identification and noise modeling as well. The error between fy and Ty, which bound, is due to the unknown outcome noise that is not accounted by our Theorem I thus, the theory in Sec. 4.2 is complementary to that in Sec. 3.2 Finally, 3 is a trade-off between the conditional balance of learned PGS (affected by ft), and precision/effeecive sample size of outcome regression and can be seen as the probabilistic counterpart of Tarr & Imail (021) l) and Kallus et al. (0118.
E DETAILS AND ADDITIONS OF EXPERIMENTS
We evaluate the post-treatment performance on training and validation set jointly (This is non-trivial. Recall the fundamental problem of causal inference). The treatment and (factual) outcome should not be observed for pre-treatment predictions, so we report them on testing set. See also Sec. C.7 the pre/post-
pre/post-treatment treatment distinction.
23

---

Published as a conference paper at ICLR 2022
E. SYNTHETIC DATA
We detail how the random parameters in the DGPs are sampled. μi and oj are uniformly sampled in range (-0.2, 0.2) and (0, 0.2), respectively. The weights of linear functions h, k, are sampled from standard normal distributions. The NNs fo, Ji use leaky ReLU activation with a 0.5
and are of 3 to 8 layers randomly, and the weights of each layer are sampled from (-1.1, -0.9). To have a large but still reasonable outcome variance, the output of fi is divided by C. = (DIT= (Z))). When generating DGPs with dependent noise, the variance parameter gr for the outcome is generated by adding softplus layer after respective ft, and then normalized to range 2).
We use the original implementation of CFR Very possibly due to bugs implementation, the CFR version using Wasserstein distance has er- ror of TensorFlow type mismatch on our synthetic dataset, and the CFR version using MMD diverges with very large loss value on one or two of the 10 random DGPs. We use MMD version, and, when the divergence of training happens, report the results from trained models before diver- gence, which still give reasonable results. We search the balancing parameter alpha in [0.16, 0.32, 0.64, 0.8, 1.28], and fix other hyperparameters as they were in the default config file.
omega
Figure 4: Degree of limited overlap w.r.t w.
We characterize the degree of limited overlap by examin- ing the percentage of observed values æ that give prob- ability less than 0.001 for one of ar). The threshold is chosen so that all sample points near those values x almost certainly belong
0.44
0.7
0.22
0.6
single group since we have 500 sample point in total. If we regard a DGP as very limited- -vvelaapiir when the above percentage than 50%, then, as shown in Figure 4 non (all) of the 10 DGPs are very limited-overrlapingg with ~ = 6 (w = 22).
0.20
0.4
larger
0.2
0.18
For diversity of the datasets, we set 9r (W) = 1 in DGPs in Appendix. Figure 5 shows, with dim(2) (Z) = 200, our method works better than CFR under di) = and as well as CFR under dim(W) > 1. As mentioned in Conclusion, this indicates that the theoretical requirement of injective ft our model might be relaxed. Interest- ingly, larger 3 seems give better results here, this is understandable because 3 controls the trade- off between fitting and balancing, and the fitting capacity of our decoder is much increased with dim 200. Note that the above observations on dim are not caused by fixing 9e (W)
10
14
18
22
2
3
4
dim(k)
betam
2.5a
bela-1 10
belae2.0
b.0
Figure 5: v Eeeke on synthetic dataset, with 9: (W) 1 DGPs, and dim(Z) = 200 in
our model. Error bar on 10 random DGPs.
(compare Figure 5 with Figure 6)
0.325
Figure S shows the importance of noise modeling. Com- pared to Figure the main text, where 9e (WW) in DGPs not fixed, our method works worse here, particularly for large 3, because now noise modeling (g, k in the
0.300
0.8
0.275
0.250
ELBO) only adds unnecessary complexity. The changes of performance w.r.t different should be unrelated to overlap levels, but to the complexity of random DGPs; compare to Figure 5 with larger NNs in VAE, the changes become much insignificant. The drop of er- for dim
0.228
0.4
0.200
0.175
18
22
2
3
5
non-overfap (omega)
dintw)
3 is due the randomness of in (36). In Sec. 2.2 we saw that the 2-dimsensional bPGS p = (Ho(X) X),, Mi (X)) always exists under additive noise models. Thus, when dlim(W') > 2, our method tries to re- cover that p. and generally performs not worse than under dim(W) = 2, but still not better than under dimmW) =1
bota:1.6
betam2.5
beta-1.0
bata-2.
beta-3.0
Figure Eppobe on synthetic dataset, with
(W) = 1 in DGPs. Frror har 10 ran- dom DGPs.
https://github.
m
24

---

Published as a conference paper at ICLR 2022
Figure 7 shows results of ATE estimation. Notably, CFR drops performance w.r.t degree of limited
overlap. Our method does not show this tendency except for very large 3 (3 = 3). This might be another evidence that CFR and its unconditional balancing overfit to PEHE (see Sec. 5.2) Also
note that, under (im = 1, B = 3 gives the best results for ATE although it does not work well for PEHE, and do not know if this generalizes to the conclusion that large gives better ATE estimation under the existence of bPGS, but leave this for future investigation.
Figure shoows results of pre-treatment prediction. In left panel, both our method and CFR perform only slightly worse than post-treatment. This is reasonable because here we have bPGS W with dim = 1, there is no need to learn PGS. In the right panel, we also do not see significant drop of performance compared to post- treatment. This might be due to the hardness of learning approximate bPGS in this dataset, and posterior estima- tion does not give much improvements.
0.4
0.3
0.08
0.2
0.04
You can find more plots for latent recovery at the end of the paper.
0.02
18
non-overlap (omaga)
dim(k)
di
betse115 1.5
beta-2.5
beta-1.0
betam2.0
Cate on synthetic dataset, with (W) in DGPs. Error bar on 10 ran- dom DGPs.
Figure 7:
E.2 IHDP
IHDP is based on an RCT where each data point repre- sents baby with 25 features (6 continuous, 19 binary) about their birth and mothers. Race is introduced as a confounder by artificially removing all treated children with nonwhite mothers. There 747 sub- jects left in the dataset. The outcome is synthesized by taking the covariates (features excluding Race) as input, hence unconfoundedness holds given the covariates. Following previous work, we split the dataset by 63:27:10 for training, validation, and testing. Note, there is no ethical concerns here, because the treatment assignment mechanism is artificial by processing the data. Also our results are only quantitative and we make no ethical conclusions.
The generating process is as following Sec. 4.1). (()) ~ N(ee" 1), N(a)~ 1,
(Hill
201T
1.0
0.326
0.800
0.275
(37)
where a is a random coefficient, b a constant bias with all elements equal to 0.5, and is random parameter adjusting degree of overlapping between the treatment groups. As can see, a X is true bPGS. As men- tioned in the main text, the bPGS might be discrete. Thus, this experiment also shows the importance of VAE, even if an apparent bPGS exists. Under discrete PSs, training an regression based on Proposition is hard, but our VAE works well.
0.250
0.225
0.4
0,200
0.2
0.178
non overvverappp (omegg))
dmmiw)
dt
beta11.5
betam2.5
beten1.0
bota:2.0
beta=3.0
Figure Pre-treatment /Epehe on synthetic dataset. Error bar on 10 random DGPs.
The two added components the modified version of our method are as following. First, we build the two outcome functions fr (Z), = 0, in our learning model (3). using two separate NNs. Second, we add to our ELBO X ( X), 10, As shown in Table 2 best unconditional balancing parameter 0.1. Larger parameters gives much worse PEHE and does not improve ATE estimation. Smaller parameters are more reasonable but still do not improve the results. The overall tendency is clear. Compared to ours, CFR with its unconditional balancing does not improve ATE estimation, it may improve PEHE results with fine tuned parameter, but possibly at the price of worse ATE estimation.
regularization term, which is the Wasserstein distance (Cuturi 2013) between
Table shows pre-treatment results, All methods gives reasonable results.
25

---

Published as a conference paper at ICLR 2022
Table 2: Performance of modified version with different unconditional balancing parameter, the values of which are shown after "Mod.".
Method
Mod.
Mod.
Mod.
Mod.
Mod.
CHR
0.2
0.1
0.05
0.01
477 007
1962. 008
.1774.00 007
.167 005
.177
.179
006
.25
v Vepake
843 050
1.979 4. osz 1.116 048 .777 = 026 8894
033
.84
029
.7122.02
Table 3: Pre-treatment Errors IHDP 1000 random DGPs. We report results with dim 10. Bold indicates method(s) which significantly better. The results are taken from Shalit al. (2017) except GANITE Yoon et al. 2018) and CEVAE [Louizos et al. 2017).
Method
TMLE
BNN
CTR
CF
CEVAR
GANITE
Ours
pre-Cate
NA
42
27..01
40
462.02
.49
.2111.011
pre-
2.1 J64.02 38
262.1
24 946
E.3 POKEC SOCIAL NETWORK DATASET
This experiment shows our method is the best compared with the methods specialized for networked deconfounding, a challenging problem in its own right. Thus, our method has the potential to work under unobserved confounding, but we leave detailed experimental and theoretical investigation to future.
Pokec (Leskovec & Krevl 2014 is real world social network dataset. We experiment semi- synthetic dataset based on Pokec, which was introduced in (Veitch et al.] (019). and use exactly the same pre-processing and generating procedure. The pre-processed network has about 79,000 vertexes (users) connected by 1.3 x 100 undirected edges. The subset of users used here are restricted three living districts which binary adjacency matrix G. Following (Veitch et al. 2019), we split the users into 10 folds, test on each fold and report the mean and std pre-treatment ATE predictions. We further separate the rest of users (in the other 9 folds) by 6: 3, for training and validation.
within the same region. The network structure is expressed by
Each user has 12 attributes, among which di strict, age, or join date is used as confounder U to build different datasets, with remaining 11 attributes used as covariate X. Treatment T and outcome are synthesised as following:
T (U)), Y Y= (0(gU)) -.5) +.,
(38)
where is standard normal. Note that district ct. is of 3 categories; age and join date also discretized into three bins. g(U), which is bPGS, maps these three categories and values to (0.15, 0.5, 0.853.
9-Intact-VAE is expected to learn a bPGS from G, X, if we can exploit the network structure effec- tively. Given the huge network structure, and neighborhood structure, which means U can be roughly seen as a deterministic function of G, This idea is comparable Assumptions and 4 in (Veitch et al. 2019). which postulate directly that balancing score can be learned in the limit of inlinite large network. To extract infor- mation from the network structure, we use Graph Convolutional Network (GCN) (Kipf Welling 2017) in conditional prior and encoder of I-Inctt VAE. The implementation details are given at the end of this subsection.
users can practically be identified by their attributes
Table 4 shows the results. The pre-treatment /fphhe for Age, District. and Join date con- founders are 1.085, 0.686, and 0.699 respectively, practically the same as the ATE errors. Note that, Veitch al. (2019) does not give individual-level prediction.
To extract information from the network structure, we use Graph Convolutional Network (GCN) (Kipf & Welling 2017) in conditional prior and encoder of 3-Intact-VAE. A difficulty is that, the network G and covariates testing phase. However, the separation can still make sense if we take care
all users are always needed by GCN, regardless of whether it is in
training. validation, that the treatment and outcome are used only in the respective phase, e.g., (ym. fm) of a testing user is only used in testing.
26

---

Published as a conference paper at ICLR 2022
Table 4: Pre-treatment ATE Pokec. Ground truth ATE
in (38) "Unadjusted" estimates
ATE by (y)) - Ep (yo). "Parametric" a stochastic block model for networked data (Gopalan Blei 2013). "Embed-" denotes the best alternatives given by (Veitch et al. 2019). Bold indicates method(s) which is significantly better than all the others. We report results with 20-dimensional latent Z. The results of the other
methods are taken from /Viichh et al. 2019).
Age
District
Join
Unadjusted
4.34
0.05
4.51 0.05
Parametric
4.06 1 0.01
3.22 0.01
3.73 1: 0.01
Embedding-Res Reg.
2.77 + 0.35
1.78 0.20
2.41 = 045
Embedding-
3.12
1.66 0.07
3.10
Ours
2.08 0.32
1.68 0.10
1.70 + 0.13
GCN takes the network matrix G and the whole covariates matrix X = M is user number, and oulputs representation matrix again for all users. During training, we select the rows in R that correspond to users in training set. Then, treat this training representation matrix as if the covariates matrix for a non-networked dataset, that is, the downstream networks in conditional prior and encoder are the same as in the other two experiments, but take (Rm.:)² where &m was expected as input. And we have respective selection operations for validation and testing. We can still train 3- Mintact-VAE including GCN by Adam, simply setting the gradients of non-seleted rows R to O.
.......:: .MM) (=?......2 where
Note that GCN cannot be trained using mini-batch, instead, perform batch gradient decent using full dataset for each iteration, with initial learning rate 10-2 We use dropout (Srivastava et al, 2014) with rate 0.1 to prevent overfitting.
E.4 EMPIRICAL VALIDATION THE BOUNDS SEC. 4.2
Here we focus on the D(X) term in Theorem because it is directly related to conditional balance.
In the Figure attached the end of the paper, the rows correspond overlap levels from strong to weak (w = 6, 14, 22 respectively). The first column shows the histograms of correlation coefficients between (XX) and (X) on 100 random DGPs. The vertical bars in the histograms are 25, 50, 75, 95 percentiles (the values are shown in the table below). The other 10 columns show the plots of distributions of (X), E (X)) for the first 10 DGPs. The correlation coefficient for each DGP is shown as corrcoef=* above each histogram. The plots are in log-log scale, because both D and Ef are single-sided, and most data points concentrate near (0, 0), making the plots bad-looking.
We have two important observations from the histograms: 1) the majority of DGPs, there positive correlations between D and ey: 2) the positive correlation is stronger with weaker overlap (the portion of large correlation increases, and the mean corrcoef O. 1000 0. 0.110, and 0.121, respectively).
Thus, our bounds and conditional balance have significance. Not all DGPs have positive correla- tions, and this is reasonable because our bound (1 has three other terms which can obscure the relation between D and ef. The DGPs 3, 6, 8, 10 show typical situations when there positive correlations.
Table 5: Percentiles of correlation coefficients between ((X) and e (X) on 100 random DGPs.
Pereentile
5
25
50
75
95
-0.289
-0.086
0.069
0.299
0.609
0.124
0.636
22
-0.274
40.128
0.067
0.341
0.634
27

---

Published as a conference paper at ICLR 2022
E.5
ADDITIONAL PLOTS ON SYNTHETIC DATASETS
Figure 9: Plots of recovered-true latent. Rows: first 10 nonlinear random models, columns: outcome noise level.
28

---

Published as a conference paper at ICLR 2022
Figure 10: Plots of recovered- true latent. Conditional prior depends on f. Rows: first 10 nonlincar random models, columns: outcome noise level. Compare the previous figure, we the transformations for 0, 1 are the same, confirming the importance of balanced prior.
29

---


Reviews:
[{"review_id": "p58NZYd9COq", "reviewer": "Reviewer_3gmW", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strengths:\n\nI believe the $\\beta$-Intact-VAE is novel and is useful in treatment effects estimation. \n\nThe authors have provided theoretical guarantees for the proposed method. \n\nExperimental results also demonstrate that $\\beta$-Intact-VAE outperforms other generative models. \n\n\nWeaknesses:\n\nI am able to understand the general idea of the method but have difficulty understanding some notations in the paper. The authors might want to define some of the symbols before using them; for example,  $f_t$, $k_t$, $h_t$, $diag$ and $diag^{-1}$ in Section 3.1 and $f_t$, $k_t$, $h_t$ in Section 4.1. The notation $f_t$ indications that there are two different functions for $t=0$ and $t=1$, respectively, is it true? \n\nI do not understand the differences between $p_\\theta(y|x, t)$ and $p(y|x, t)$ in Section 3.2. \n\nIn the caption of Table 1, What is the unconditional balancing hyperparameter?\n\nIn Section 4.2, are the parameters $\\Lambda$ learn from the data? In Equation (6), $ q_{\\phi} (x, y, t) $ is still a function of $t$. How is balanced PtS achieved in this case? In the second term of Equatnion (7), it looks like $g_t^2(z)$ should be $g_t(z)$.", "summary_of_the_paper": "This paper focuses on estimating treatment effects under a limited overlap, i.e., subjects with specific characteristics might only belong to one treatment group. To solve this problem, the authors propose a variational autoencoder (VAE) called $\\beta$-Intact-VAE, which extends the framework of CVAE and iVAE.  The authors prove that the identifies the treatment effects. The proposed model is compared with other methods on synthetic datasets.\n", "summary_of_the_review": "In general, I believe the $\\beta$-Intact-VAE is novel and is useful in treatment effects estimation. The authors have provided theoretical guarantees for the proposed method. Experimental results also demonstrate that $\\beta$-Intact-VAE outperforms other generative models. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "T-7GKkGxr5E", "reviewer": "Reviewer_ENDL", "summary": "", "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": "Strengths:\n\n1. Authors contribute a novel approach for limited overlapping of covariates. The setting of limited overlapping is interesting. The proposed approach combines the prognostic score and Intact-VAE, which are appropriate and novel under this setting. Although the \\beta variational lower bound was proposed in the previous methods, it plays a role of controlling balance and it can be viewed as a novel application.\n\n2. The experiment results are extensive. The authors follow some existing data generation processing or data and show that the balanced prior is important in practice compared to non-balanced ones and their models can obtain better performance.\n\n3. The theoretical analyses are novel under the setting of this paper. These results follow those of iVAE and adapt to the new settings. The authors elaborate on some conditions to help readers to get them, which is great.\n\nWeaknesses:\n\n1. The paper is dense and difficult to follow. It introduces many concepts from different areas. It takes much time to understand and check the paper. I suggest authors add some plots to help readers to faster understand some definitions, e.g. overlapping. The paper involves and requires readers to have wide knowledge, so I think that a good presentation is difficult. Maybe moving some contents to the appendix and illustrating some concepts could be a choice. In addition, characters with the same style are used for different types of subjects, such as $\\mathbb{R}$ for real numbers and $\\mathbb{M}$ for a function. I strongly suggest using consistent notations.\n\n2. Experiment results measured by pehe are worse than CFR. ATE and PEHE reflect different information. ATE focuses on the mean and PEHE reflects the stability because if there is a large CATE error, PEHE can be large (as mentioned on the page 9). Large CATE errors mean that the model performs unstably at some values. So, it is better to discuss the results more properly.\n", "summary_of_the_paper": "The paper considers the setting of limited overlap of covariates and studies identifying and estimating the causal effects. This setting is difficult to deal with because it is impossible to estimate causal effects at non-overlapping values due to a lack of data. To this end, the paper proposes an idea based on the prognostic score and identifiable representation (Intact-VAE). The prognostic score is an appropriate tool for limited overlapping because it can map some non-overlapping values to an overlapping value in a space of lower dimensions. Intact-VAE is a natural combination of iVAE and CVAE and is proposed to help to identify the causal effect. The paper implements this idea by proposing a new regularized VAE, called \\beta-Intact-VAE, and further analyzes the treatment effect error bounds. In the end, the paper compares the model with recent methods in the experiments.", "summary_of_the_review": "I vote for weak acceptance. The paper proposes a technically and theoretically sound approach and it studies an interesting setting -- limited overlapping, though it has some presentation issues.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "FTdq1u0bOCt", "reviewer": "Reviewer_ZUkW", "summary": "", "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": "Strengths:\n- With the exception of a few errors, the paper is well written and has an understandable language. The motivation of the work is clear and relevant.\n- The theoretical analysis in the paper is serious, as far as my understanding of the topic goes. \n- The proposed method seems reasonable, and justified by the theoretical analysis in the paper.\n\nWeaknesses:\n- I found the paper to be a bit information dense. In fact, I am unsure about what could be some of the most subtle assumptions of the method. Clearly there is the assumption that X is a valid adjustment set, but what about (G1, additive noise models)? If this is necessary, then why use the machinery of a neural network? It seems that in order to make the theoretical analysis easier you have to assume a noiseless prior. Is there any chance you can make a list of the required assumptions, and a sentence or two of why they are needed? Maybe you can consider having such an explicit version of the assumptions in the main document or a separate section in the appendix.\n\n- Along the same line of the previous point, I would invite the authors to think about scenarios in which such assumptions are fulfilled. The examples of this paper, in which we want to adjust a causal estimate using high dimensional covariates has always struck me as a bit unrealistic. As the authors might know, adjusting for all pre-treatment variables is not always the best option because that might introduce M-bias. I would like to say, though, that as the authors correctly point out, the more variables you have, the lower the chance of having overlap of the treatments. Can you please comment on what could be some possible scenarios in which the method could be applied?\n\n- Finally, if known to the authors, I would like to know the relation of this work between work on generalization of neural networks. It seems, from my point of view, that the problem of limited overlapping of conditioning variables can be fixed by having estimators of the counterfactuals that generalize well beyond the support of the combination of covariates, and treatment. Could it be possible that what is driving most of your results is the regularization of the estimates? In general regularization is a way of achieving out of support generalization. Can you comment on this?\n", "summary_of_the_paper": "The paper tackles the problem of estimating causal effects under partial overlap conditions. Overlap is a common assumption in causality, so the problem is hard and relevant. The authors present a theoretical analysis using PtS, an adaptation of prognostic scores. The estimation is carried out with a modified version of a Variational Autoencoder (VAE) called \\beta-intact-VAE. In addition, the authors present bounds on the error of the Conditional Average Treatment Effect (CATE) using their framework. Their paper ends with experimental results of the proposed method and a comparison with other recent methods for CATE estimation.", "summary_of_the_review": "I believe the paper solves the interesting problem of estimating CATE with limited overlap. It seems to me that most of the theory, and the estimation method were developed before this article, so in a certain way it is an incremental contribution. Nevertheless, both theoretical, and empirical analysis of the article seem serious.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6j_QRI5Veop", "reviewer": "Reviewer_sh3V", "summary": "", "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": "## Strong points\n- The authors extend the idea of prognostic scores to the context of machine learning and with the additional difficulty of limited overlap. Indeed this approach that can be seen as an alternative to the more common propensity score based approaches in causal inference is indeed interesting in cases of limited overlap which is relevant especially in high-dimensional settings.\n- The authors provide a broad range of theoretical results, ranging from treatment effect identification via prognostic scores to a novel estimation strategy with error bounds for the CATE estimation.\n- The article is well written, especially the first two sections lay out the context and motivation. The authors provide an adequately succinct yet sufficient list of references on related works to position their proposal.\n\n\n## Issues/Points that require clarification\n- As a general comment, the claims and statements of the main text, both in the identification and the estimation sections, are in part hard to follow without checking in the Appendix which makes the reading difficult. The content being very dense, I would suggest to either split the work into two articles, one for identification and the other on estimation or to submit this work under a more readable format or to defer the experiments to the appendix instead of the theoretical aspects and small examples which would help to follow the justifications in the identification and estimation sections.\n- In definition 1 the authors talk about overlapping random variables ($V$) and in definition 2 they talk about functions. Do they mean functions of $x\\in\\mathcal{X}$ and when they say overlapping functions $\\mathbb{P}_t(X)$, do they mean the random variables or do they make an implicit definition of overlapping functions?\n- The authors claim that $\\mathbb{M}(X)$ is an effect modifier. Since this relies on a definition of effect modifier from Hansen (2008) which is much less common than standard definitions of effect modifiers as interaction terms between covariates and the treatment assignment variable, it would be helpful to provide a concrete example (if not in the main text, then at least in the appendix) of such an effect modifier that does not involve $T$.\n- At the end of Section 2, the authors provide an informal definition of _balanced PtS_. Why this non-rigorous definition instead of a formal one?\n- Maybe I'm confusing the Intact VAE and the $\\beta$ Intact VAE but from the graphical model of Intact VAE in Figure 1 I don't understand the connection between $Z$ and a PS $\\mathbb{P}$. The definition of conditionally balanced representation requires that $Z\\perp T|X$ which does not seem to be possible under the model of Intact VAE. \n- The choice of the acronym \"PS\" for prognostic scores could be reconsidered since the target audience of this work seems to be rooted in causal inference where \"PS\" usually stands for the propensity score in many works.\n- I have several questions about the experiments\n    - Why the notation switch from $Z$ to $W$ in the experiments? Is there a conceptual difference between the two?\n    - Why the choice of linear functions $h,k,l$ and how are they chosen?\n    - Why isn't the proposed method from D'Amour and Franks (2021) also included in the comparative study? Even if their method relies on linear assumptions, it would be interesting to see a comparison in this context.\n    - Since the paper contains a derivation of error bounds on the CATE (in Theorem 2), it would be interesting to assess them in the simulations and compare the empirical results to these bounds.\n    - For the IHDP dataset, it is written that in the used model, $\\beta=1$ (referenced as _Ours_ in Table 1) but then there are also modified versions that are reported with other values of the hyperparameter $\\beta$. Shouldn't _Ours_ and _Mod.1_ show the same results or am I misunderstanding the table's description?\n\n\n### Minor comments (that did not impact the score)\n- p.1: Confounders are not necessarily high-dimensional. I would reformulate the last sentence of the 2nd paragraph to nuance the claim: rather say that the more covariates are collected the more likely unconfoundedness is to hold but that this can quickly lead to issues in separation between treatment and control groups.\n- p.1: I would suggest adding the work on overlap weights by Li and Li (2019)\n- p.3: In Section 2.1, I would recommend adding the SUTVA assumption for completeness (either by mentioning it directly or by adding the exclusion of interference in the assumptions).\n\n######################################################################################################\n### Post-rebuttal update\nI thank the authors for their detailed and timely responses to all reviewers. Their efforts to clarify and reformulated assumptions on the DGP, model and data, as well as their concession to adapt their notations following the reviewers' remarks and the additional experiments related to the theoretical bound have helped addressing most of my concerns. However, I still think the positioning of their work relative to related works is difficult to follow or find which makes it difficult to compare its contributions with previous works. I have therefore decided to only increase my rating by one level.\n\n######################################################################################################\n\n### References\n[1] Alexander D'Amour and Alexander Franks. Deconfounding Scores: Feature Representations for Causal Effect Estimation with Weak Overlap. _arXiv preprint arXiv:2104.05762_, 2021.\n\n[2] Ben B. Hansen. The prognostic analogue of the propensity score. _Biometrika_ 95(2): 481-488, 2008.\n\n[3] Fan Li and Fan Li. Propensity score weighting for causal inference with multiple treatments. _Annals of Applied Statistics_, 13(4):2389-2415, 2019.\n\n[4] Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In _Proceedings of the 31st International Conference on Neural Information Processing Systems_, pp. 6449-6459. 2017.\n\n[5] Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In _International Conference on Machine Learning_, pp. 3076-3085. PMLR, 2017.", "summary_of_the_paper": "The present paper proposes to address the issue of limited overlap in treatment effect estimation by investigating the identification assumptions of prognostic scores which, in certain contexts, are less restrictive on overlap than other methods such as propensity score based methods.\nThe main contributions of this work are presented in Sections 3.2 and 4.2, where they derive identifiability of the treatment effect via prognostic scores and propose a generative prognostic model based on variational auto-encoders (VAE).\nThe theoretical results are complemented with several synthetic and semi-synthetic experiments which also show that the proposed models can compete with and in certain cases improve upon state of the art.", "summary_of_the_review": "In summary, the aim and approach of this paper are interesting and I believe the approach could be a methodological gain for the causal machine learning audience, but its underlying assumptions are presented in a way that makes it difficult to link or compare them to other works and its significance for real-world examples is therefore difficult to judge. Especially, the details of the underlying causal model are not clear to me and make it difficult to conceptually compare it with other methods such as CFR (Shalit et al., 2017) or CEVAE (Louizos et al., 2017).\nI will read the rebuttal carefully and am willing to increase the score if the authors address the raised concerns.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors proposed a method for causal inference under limited overlap -- an important and understudied complication.  The authors propose to recover a prognostic score using a variational autoencoder, and thereby map a higher dimensional set of covariates with limited overlap to a lower dimensional set where overlap holds, and such that ignorability is maintained.

The paper was reviewed quite favorably by reviewers, and the authors updated the manuscript to address specific issues raised by reviewers.
Decision: Accept (Poster)