β -I NTACT -VAE: I DENTIFYING AND E STIMATING C AUSAL E FFECTS UNDER L IMITED O VERLAP
Anonymous authors
Paper under double-blind review
A BSTRACT
As an important problem in causal inference, we discuss the identiﬁcation 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 sufﬁcient for
TEs; i.e., we build a generative prognostic model.
We prove that the latent variable recovers a prognostic score, and the model identiﬁes individualized treatment effects.
The model is then learned as β -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.
I NTRODUCTION
Causal inference (Imbens & Rubin, 2015; Pearl, 2009), i.e, inferring causal effects of interventions, is a fundamental ﬁeld 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 difﬁculty of causal inference is that we never observe counterfactual outcomes that would have been if 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 is confounding : there may be variables, called confounders, that causally affect both the treatment and the outcome, and spurious correlation/bias follows.
The other is the systematic imbalance (difference) of the distributions 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 unconfoundedness; this means that the confounding can be controlled by conditioning on the covariates.
The more covariates are collected the more likely unconfoundedness holds; however, more covariates tends to introduce a 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)––which improve ﬁnite sample performance, as well as non-parametric methods––e.g., Wager & Athey (2018, CF).
Notably, from Johansson et al.
, there has been a recent increase in 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 is the limited (or weak) overlap of covariates , which means that sample points with certain covariate values belong to a single treatment group.
In this case, a 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´ar, 2021), trim non-overlapping data points (Yang & Ding, 2018), weight data points by overlap (Li & Li, 2019), or study convergence rates depending on overlap (Hong et al., 2020).
Limited overlap is particularly relevant to machine learning methods that exploit high-dimensional covariates.
This is because, with higher-dimensional covariates, overlap is harder to satisfy and verify (D’Amour et al., 2020).
To address imbalance and limited overlap, we use a prognostic score (Hansen, 2008); it is a sufﬁcient statistic of outcome predictors and is among the key concepts of sufﬁcient scores for TE estimation.
As a function of covariates, it can map some non-overlapping values to an overlapping value in a space of lower-dimensions.
For individualized TEs, we consider conditionally balanced representation Z , such that Z is independent of T given X ––which, as we will see, is a necessary condition for a balanced prognostic score.
Moreover, prognostic score modeling can beneﬁt from methods in predictive analytics and exploit rich literature, particularly in medicine and health (Hajage et al., 2017).
Thus, it is promising to combine the predictive power of prognostic modeling and machine learning.
With this idea, our method builds on a generative prognostic model that models the prognostic score as a latent variable and factorizes to the score distribution and outcome distribution.
As we consider latent variables and causal inference, identiﬁcation is an issue that must be discussed before estimation is considered.
“Identiﬁcation” means that the parameters of interest (in our case, representation function and TEs) are uniquely determined and expressed using the true observational distribution.
Without identiﬁcation, a consistent estimator is impossible to obtain, and a model would fail silently; in other words, the model may ﬁt perfectly but will return an estimator that converges to a wrong one, or does not converge at all (Lewbel, 2019, particularly Sec. 8).
Identiﬁcation is even more important for causal inference; because, unlike usual (non-causal) model misspeciﬁcation, causal assumptions are often unveriﬁable through observables (White & Chalak, 2013).
Thus, it is critical to specify the theoretical conditions for identiﬁcation, and then the applicability 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 identiﬁable.
This is because the factorization of our model is naturally realized as a combination of identiﬁable VAE (Khemakhem et al., 2020a, iVAE) and
conditional VAE (Sohn et al., 2015, CVAE).
Based on model identiﬁability, we develop two identiﬁcation results for individualized TEs under limited overlap.
A similar VAE architecture was proposed in Wu & Fukumizu (2020b; 2021a); 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 identiﬁcation theories.
The current study further provides bounds on individualized TE error, and the bounds justify a conditionally balancing term controlled by hyperparameter β , as an interpolation between the two identiﬁcations.
In summary, we study the identiﬁcation (Sec. 3) and estimation (Sec. 4) of individualized TEs under limited overlap.
Our approach is based on recovering prognostic scores from observed variables.
To this end, our method exploits recent advances in identiﬁable representation––particularly iVAE.
The code is in Supplementary Material, and the proofs are in Sec.
A.
Our main contributions are:
1) TE identiﬁcation under limited overlap of X , via prognostic scores and an identiﬁable model; 2) bounds on individualized TE error, which justify our conditional BRL; 3) a new regularized VAE, β -Intact-VAE, realizing the identiﬁcation and conditional balance; 4) experimental comparison to the state-of-the-art methods on (semi-)synthetic datasets.
1. R ELATED WORK
Limited overlap.
Under limited overlap, Luo et al.
(2017) estimate the average TE (ATE) by reducing covariates to a linear prognostic score.
Farrell (2015) estimates a constant TE under a partial linear outcome model.
D’Amour & Franks (2021) study the identiﬁcation of ATE by a general class of scores, given the (linear) propensity score and prognostic score.
Machine learning studies on this topic have focused on ﬁnding overlapping regions (Oberst et al., 2020; Dai & Stultz, 2020), or indicating possible failure under limited overlap (Jesson et al., 2020), but not remedies.
An exception is
Johansson et al.
(2020), which provides bounds under limited overlap.
To the best of our knowledge, our method is the ﬁrst machine learning method that provides identiﬁcation 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 reducing covariates to a linear score which is a joint propensity-prognostic score.
Tarr & Imai (2021) 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––while Johansson et al.
(2020) conceptually, but not formally, connects BRL to prognostic score.
Our work is the ﬁrst to formally connect generative learning and prognostic scores for TE estimation.
Identiﬁable representation.
Recently, independent component analysis (ICA) and representation learning––both ill-posed inverse problems––meet together to yield nonlinear ICA and identiﬁable representation; for example, using VAEs (Khemakhem et al., 2020a), and energy models (Khemakhem et al., 2020b).
The results are exploited in causal discovery (Wu & Fukumizu, 2020a) and out-of-distribution generalization (Sun et al., 2020).
This study is the ﬁrst to explore identiﬁable representations in TE identiﬁcation.
BRL and related methods amount to a major direction.
Early BRL methods include BLR/BNN (Johansson et al., 2016) and TARnet/CFR (Shalit et al., 2017).
In addition, Yao et al.
(2018) exploit the local similarity between data points.
Shi et al.
(2019) use similar architecture to TARnet, considering 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 idea of BRL, and further extends to conditional balance––which is natural for individualized TE.
Others.
Our work can lay conceptual and theoretical foundations of VAE methods for TEs (e.g.,
Louizos et al., 2017; Lu et al., 2020).
In addition, some studies consider monotonicity, which is injectivity on R , together with overlap (Johansson et al., 2020; Zhang et al., 2020b).
See Sec.
D.
S ETUP AND PRELIMINARIES 2. C OUNTERFACTUALS , TREATMENT EFFECTS , AND IDENTIFICATION
Following Imbens & Rubin (2015), we assume there exist potential outcomes Y ( t ) ∈ R d , t ∈ { 0 , 1 } .
Y ( t ) is the outcome that would have been observed if the treatment value T = t was applied.
We see Y ( t ) as the hidden variables that give the factual outcome Y under factual assignment T = t .
Formally, Y ( t ) is deﬁned by the consistency of counterfactuals :
Y = Y ( t )
if T = t ; 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) or Y (1) ––w.r.t.
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 ⊆ R m , which is associated with individuals, with distribution D :=
( X, Y, T ) ∼ p ( x , y , t ) .
We use upper-case (e.g. T ) to denote random variables, and lower-case (e.g. t ) for realizations.
The expected potential outcome is denoted by µ t ( x ) =
E ( Y ( t ) | X = x ) conditioned on X = x .
The estimands in this work are the conditional ATE (CATE) and ATE, deﬁned, respectively, by:
[Equation]
CATE is seen as an individual-level , personalized, treatment effect, given highly discriminative X .
Standard results (Rubin, 2005)(Hernan & Robins, 2020, Ch. 3) show sufﬁcient conditions for TE identiﬁcation in general settings.
They are Exchangeability : Y ( t ) | =
T | X , and Overlap : p ( t | x ) >for any x ∈ X .
Both are required for t ∈ { 0 , 1 } .
When t appears in statements without quantiﬁcation, we always mean “for both t ”.
Often, Consistency is also listed; however, as mentioned, it is better known as the well-deﬁnedness of counterfactuals.
Exchangeability means, just as in RCTs, but additionally given X , that there is no correlation between factual T and potential Y ( t ) .
Note that the popular assumption Y (0) , Y (1) | =
T | X is stronger than Y ( t ) | =
T | X
and is not necessary for identiﬁcation (Hernan & Robins, 2020, pp. 15).
Overlap means that the supports of p ( x | t = 0) and p ( x | t = 1) should be the same, and this ensures that there are data for µ t ( x ) on any ( x , t ) .
We rely on consistency and exchangeability, but in Sec. 3.2, will relax the condition of the overlapping covariate to allow some non-overlapping values x ––that is, covariate X is limited-overlapping .
In this paper, we also discuss overlapping variables other than X (e.g., prognostic scores), and provide a deﬁnition for any random variable V with support V as follows:
Deﬁnition 1.
V is Overlapping if p ( t | V = v ) > 0 for any t ∈ { 0 , 1 } , v ∈ V .
If the condition is violated at some value v , then v is non-overlapping and V is limited-overlapping .
2. P ROGNOSTIC SCORES
Our method aims to recover a prognostic score (Hansen, 2008), adapted to account for both t as in Deﬁnition 2.
On the other hand, balancing scores (Rosenbaum & Rubin, 1983) b ( X ) are deﬁned by T | =
X | b ( X ) , of which the propensity score p ( t = 1 | X ) is a special case.
See Sec.
B.1 for detail.
Deﬁnition 2.
A PGS is { P ( X, t ) } t ∈{ 0 , 1 } such that Y ( t ) | =
X | P ( X, t ) , where P ( x , t ) ( P t ( x ) hereafter) is a function deﬁned on X × { 0 , 1 } .
A PGS is called balanced (and a bPGS ) if P 0 = P 1 .
We say a PGS is overlapping, if both P 0 ( X ) and P 1 ( X ) are overlapping. Obviously, a bPGS P ( X ) is a conditionally balanced representation (deﬁned as Z | = T | X in Introduction) and is thus named.
We often write t [-( X ) means both P 
 ( X ) and P 1 ( X ) are overlapping.-]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 limited-overlapping X .
It is known that overlapping b ( X ) implies overlapping X (D’Amour et al., 2020), which counters our purpose.
In contrast, overlapping bPGS does not imply overlapping b ( X ) (see Sec.
C.1 for an example).
Moreover, with theoretical and experimental evidence, it is recently conjectured that PGSs maximize overlap among a class of sufﬁcient 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 P t ( X ) satisﬁes exchangeability.
Proposition 1 (Identiﬁcation via PGS) .
If P t ( X ) is a PGS and Y | P ˆ t ( X ) , T ∼ p Y | P ˆ t ,T ( y | P, t ) where ˆ t ∈ { 0 , 1 } is a counterfactual assignment, then CATE and ATE are identiﬁed, using (1) and µ ˆ t ( x ) = E ( Y (ˆ t
[Equation]
With the knowledge of P t and p Y | P ˆ t ,T , we choose one of P 0 , P 1 and set t
= ˆ t in the density function, w.r.t the µ ˆ t of interest.
This counterfactual assignment resolves the problem of non-overlap at x .
Note that a sample point with X = x may not have T = ˆ t .
We mainly consider additive noise models for Y ( t ) , which ensures the existence of PGSs.
(G1) 1 (Additive noise model) the data generating process (DGP) for Y is Y =
f ∗ ( M ( X, T ) , T )+ e where f ∗ , M are functions and e is a zero-mean exogenous (external) noise.
The DGP deﬁnes potential outcomes by Y ( t ) :
= f ∗ t ( M t ( X )) + e , and speciﬁes M ( X, T ) , T , and e as the only direct causes of Y .
Particularly, M t ( X ) is a sufﬁcient statistics of X for Y ( t ) . For example, 1) M t ( X ) can be the component(s) of X that affect Y ( t ) directly, or 2) if Y ( t ) | X follows a generalized linear model, then M t ( X ) can be the linear predictor of Y ( t ) .
Under (G1) , 1) M t ( X ) is a PGS; 2) µ t ( X ) =
f ∗ t ( M t ( X )) is a PGS; 3) X is a (trivial) bPGS; and 4) U ( X ) :
= ( µ 0 ( X ) , µ 1 ( X )) is a bPGS.
The essence of our method is to recover the PGS M t ( X ) as a representation, assuming M t ( X ) is not higher-dimensional than Y and approximately balanced.
Note that µ t ( X ) , our ﬁnal target, is a low-dimensional PGS but not balanced, and we estimate it conditioning on the approximate bPGS M t ( X ) .
I DENTIFICATION UNDER GENERATIVE PROGNOSTIC MODEL
In Sec. 3.1, we specify the generative prognostic model p ( y , z | x , t ) , and show its identiﬁability.
In Sec. 3.2, we prove the identiﬁcation 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.
3. M ODEL , ARCHITECTURE , AND IDENTIFIABILITY
We present the necessary deﬁnitions and results, and defer more explanations to Sec.
C.2.
The contents of this subsection are essentially taken from Wu & Fukumizu (2021b), but included here for completeness.
Our goal is to build a model that can be learned by VAE from observational data to obtain a PGS, or better, a bPGS, via the latent variable Z .
The generative prognostic model of the proposed method is in (3), where θ :=
( f , h , k ) contains the functional parameters.
The ﬁrst factor p f ( y | z , t ) , our decoder, models p Y | P t ,T ( y | P, t ) in (2) and is an additive noise model, with (cid:15) ∼ p (cid:15) as the exogenous noise. The second factor p λ ( z | x , t ) , our conditional prior, models
P T ( X ) and is a factorized Gaussian, with λ T ( X ) : =
diag − 1 ( k T ( X ))
( h T ( X ) , − 12 ) T as its natural parameter in the exponential family, where diag() gives a diagonal matrix from a vector.
[Equation]
We denote n := dim( Z ) .
For inference, the standard argument derives the ELBO:
[Equation]
Note that the encoder q conditions on all the observables ( X, Y, T ) ; this fact plays an important role in Sec. 4.1.
Full parameterization of the encoder and decoder is also given in Sec. 4.1.
This architecture is called Intact-VAE ( I de n tiﬁable t re a tment- c ondi t ional VAE).
See Figure 1 for comparison in terms of graphical models.
See Sec.
B.2 for basics of VAEs.
Our model identiﬁability extends the theory of iVAE, and the following conditions are inherited.
(M1) i)
f t is injective, and ii) f t is differentiable.
(D1)
λ t ( X ) is non-degenerate, i.e., the linear hull of its support is 2 n -dimensional.
Under (M1) and (D1) , we obtain the following identiﬁability of the parameters in the model: if p θ ( y | x , t ) =
p θ (cid:48)
( y | x , t ) , we have, for any y t in the image of f t :
[Equation]
where diag( a ) is an invertible n -diagonal matrix and b is an n -vector, both of which depend on λ t ( x ) and λ (cid:48) t ( x ) .
The essence of the result is that f (cid:48) t = f t ◦ A t ; that is, f t can be identiﬁed (learned) up to an afﬁne transformation A t .
See Sec.
A for the proof and a relaxation of (D1) .
In this paper, symbol (cid:48) (prime) always indicates another parameter (variable, etc.): θ (cid:48)
=
( f (cid:48) , λ (cid:48) ) .
3. I DENTIFICATIONS UNDER LIMITED - OVERLAPPING COVARIATE
In this subsection, we present two results of CATE identiﬁcation based on the recovery of equivalent bPGS and PGS, respectively.
Since PGSs are functions of X , the theory assumes a noiseless prior for simpliﬁcation, i.e., k ( X ) = 0 ; the prior Z λ ,t ∼ p λ ( z | x , t ) degenerates to function h t ( X ) .
PGSs with dimensionality lower than or equal to d = dim( Y ) are essential to address limited overlapping, as shown below.
We set n = d because µ t is a PGS of the same dimension as Y under (G1) .
In practice, n = d means that we seek a low-dimensional representation of X .
We introduce (G1’)
(Low-dimensional PGS) (G1) is true, and µ t =
j ◦ P t for some P t and injective j , which is equivalent to (G1) because µ t = j ◦ P t is trivially satisﬁed with j is identity and P t = µ t .
(G1’) is used instead in this subsection.
First, it explicitly restricts dim( P t ) via injectivity, which ensures that n = dim( Y ) ≥ dim( P t ) .
Second, it reminds us that, possibly, the decomposition is not unique; and, clearly, all P t that satisfy (G1’) are PGSs.
For example, if f ∗ = f ∗ 0 = f ∗ 1 is injective, then j = f ∗ and P t = M t satisﬁes µ t = j ◦ P t . Finally, it is then natural to introduce (G2)
(Low-dimensional bPGS) (G1) is true, and µ t =
j t ◦ P for some P and injective j t , which is stronger than (G1) , gives bPGS P ( X ) , and ensures that n ≥ dim( P ) .
(G2) is satisﬁed if f ∗ t is injective and {+M 
 = M 1 . (G2) implies µ 
+}
=
i ◦ µ 0 where i :=
j 1 ◦ j − 1 0 ; in words, CATEs are given by µ 0 and an invertible function.
See Sec.
C.3 for real-world examples and more discussions.
With (G1’) or (G2) , overlapping X can be relaxed to overlapping bPGS or PGS plus the following: (M2) (Score partition preserving)
For any x x (cid:48) ∈ X , if P t x ) =
P t x
(cid:48) ) , then h t ( x ) =
h t ( x (cid:48) )
Note that (M2) is only required for optimal h that satisﬁes E p θ ( Y | X, T )
=
E ( Y | X, T ) in Proposition 2, or p θ ( y | x , t ) =
p ( y | x , t ) in Theorem 1.
The intuition is that P t maps non-overlapping x to an overlapping value, and h t preserves this property through learning.
Linear P t and h t imply (M2) and are often assumed, e.g., in Huang & Chan (2017)
Luo et al.
D’Amour & Franks (2021).
Linear outcome models (Farrell, Schuler et al., 2020) are also common.
Our ﬁrst identiﬁcation, Proposition 2, relies on (G2) and our generative model, without model identiﬁability (so differentiable f t is not needed).
Proposition 2 (Identiﬁcation via recovery of bPGS) .
Suppose we have DGP (G2) and model (3) with n = d . Assume (M1)
-i) and (M3) (PS matching)
h
( X ) =
h 1 ( X ) and k ( X ) = 0 .
Then, if E p θ ( Y | X, T ) =
E ( Y | X, T ) , we have
[Equation]
where v : P → R n is an injective function, and P :
= { P ( x ) | overlapping x } ; 2) (CATE Identiﬁcation) if P in (G2) is overlapping, and (M2) is satisﬁed, then µ
[Equation]
[Equation]
) , for any t ∈ { 0 , 1 } and x ∈ X .
In essence, i) the true DGP is identiﬁed up to an invertible mapping v , such that f t =
j t ◦ v − 1 and h = v ◦ P ; and ii) P t is recovered up to v , and Y ( t ) | =
X | P t ( X ) is preserved––with same v for both t .
Theorem 1 below also achieves the essence i) and ii), under P 0 (cid:54) =
P 1 and j 0 = j 1 .
The existence of bPGS is more preferred, because it satisﬁes 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, we depend on Theorem 1 based on the model identiﬁability to work under PGS which generally exists.
Theorem 1 (Identiﬁcation via recovery of PGS) .
Suppose we have DGP (G1’) and model (3) with n = d .
For the model, assume (M1) and (M3’)
(Noise matching)
p e = p (cid:15) and k ( X ) =
k k (cid:48)
( X ) , k → 0 .
Assume further that (D1) and (D2) (Balance from data)
A
=
A 1 in (5) .
Then, if p θ ( y | x , t ) = p ( y | x , t ) ; conclusions 1) and 2) in Proposition 2 hold with P replaced with P t in (G1’) ; and the domain of v becomes P :=
{ P t ( x ) | p ( t, x ) > 0 } .
Theorem 1 implies that, without bPGS, we need to know or learn the distribution of hidden noise (cid:15) to have p e
= p (cid:15) .
Proposition 2 and Theorem 1 achieve recovery and identiﬁcation in a complementary manner; the former starts from the prior by P
= P 1 and h
= h 1 , while the latter starts from the decoder by A
=
A 1 and p e = p (cid:15) .
We see that A
=
A 1 acts as a kind of balance because it replaces P
= P 1 in Proposition 2.
We show in Sec.
A a sufﬁcient and necessary condition (D2’) on data that ensures A
=
A 1 .
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 identiﬁcations.
E STIMATION BY β -I NTACT -VAE 4. P RIOR AS B PGS, POSTERIOR AS PGS, AND β AS 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––while Theorem 1 allows us to deal with PGS by adding other conditions.
In learning our model with data, we formally require (G1) and further expect that (G2) holds approximately; the latter is true when f ∗ t is injective and {+M 
 ≈ M 1 ( M t ( X ) is an approximate bPGS). Instead+} of the trivial regression µ t ( X )
= E ( Y | X, T = t ) , we want to recover the approximate bPGS M t ( X ) .
This idea is common in practice.
For example, in a real-world nutrition study (Huang & Chan, 2017), a reduction of 11 covariates recovers a 1-dimensional linear bPGS.
We consider two ways to recover an approximate bPGS by a VAE.
One is to use a prior which does not depend on t , indicating a preference for bPGS.
Namely, we set λ 0 = λ 1 , denote Λ ( X ) :
= λ ( X ) and have p Λ ( z | x ) as the prior in (3).
The decoder and encoder are factorized Gaussians:
[Equation]
where φ = ( r , s ) .
The other is to introduce a hyperparameter β in the ELBO as in β -VAE (Higgins et al., 2017).
The modiﬁed ELBO with β , up to the additive constant, is derived as:
[Equation]
For convenience, here and in L f in Sec. 4.2, we omit the summation as if Y is univariate.
The encoder q φ depends on t and can realize a PGS. With β , we control the trade-off between the ﬁrst 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 a larger β encourages the conditional balance Z | =
T | X on the posterior.
By choosing β appropriately, e.g., by validation, the ELBO can recover an approximate bPGS while ﬁtting the outcome well.
In summary, we base the estimation on Proposition 2 and bPGS as much as possible, but step into Theorem 1 and noise modeling required by p e = p (cid:15) 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.
C.5 for more ideas and connections behind the ELBO.
Once the VAE is learned 2 by the ELBO, the estimate of the expected potential outcomes is given by:
[Equation]
where q ( z | x ) :
= E p ( y ,t | x ) q φ ( z | x , y , t ) is the aggregated posterior.
We mainly consider the case where x is observed in the data, and the sample of ( Y, T ) is taken from the data given X = x .
When x is not in the data, we replace q φ with p Λ in (8) (see Sec.
C.7 for details and E for results).
Note that ˆ t 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
=
ˆ t in the decoder.
The assignment is not applied to the encoder which is learned from factual X, Y, T (see also the explanation of (cid:15) CF,t in Sec. 4.2).
The overall algorithm steps are i) train the VAE using (7), and ii) infer CATE ˆ τ ( x )
= ˆ µ 1 ( x ) − ˆ µ 0 ( x ) by (8).
4. C ONDITIONALLY BALANCED REPRESENTATION LEARNING
We formally justify our ELBO (7) from the BRL viewpoint.
We show that the conditional BRL via the ﬁrst term of the ELBO results from bounding a CATE error; particularly, the error due to the imprecise recovery of j t in (G1’) is controlled by the ELBO.
Previous works (Shalit et al., 2017;
Lu et al., 2020) instead focus on unconditional balance and bound PEHE which is marginalized on
X .
Sec.
5.2 experimentally shows the advantage of our bounds and ELBO.
Further, we connect the bounds to identiﬁcation and consider noise modeling through g t ( z ) .
Sec Sec. C.8 for detail.
We introduce the objective that we bound.
Using (8) to estimate CATE, ˆ τ f ( z ) :
= f 1 ( z ) − f
( z ) is marginalized on q ( z | x ) .
On the other hand, the true CATE , given the covariate x or score z , is:
[Equation]
where j t is associated with an approximate bPGS P t (say, M t ) as the target of recovery by our VAE.
Accordingly, given x , the error of posterior CATE , with or without knowing P t , is deﬁned as
[Equation]
We bound (cid:15) f instead of (cid:15) ∗ f because the error between τ ( X ) and τ j ( Z ) is small––if the balanced P t is recovered, then z ≈ P 0 ( x ) ≈ P 1 ( x ) in (9).
We consider the error between ˆ τ f and τ j below.
We deﬁne the risks of outcome regression, into which (cid:15) f is decomposed.
Deﬁnition 3 (CATE risks) .
Let Y (ˆ t ) | P ˆ t ( X ) ∼ p Y (ˆ t )
| P ˆ t ( y | P ) and q t ( z | x ) :
= q ( z | x , t ) =
E p ( y | x ,t ) q φ .
The potential outcome loss at ( z , t ) , factual risk , and counterfactual risk are:
[Equation]
With Y ( t ) involved, L f is a potential outcome loss on f , weighted by g .
The factual and counterfactual counterparts, (cid:15) F,t and (cid:15) CF,t , are deﬁned accordingly.
In (cid:15) F,t , unit u
=
( x , y , t ) is involved in the learning of q t ( z | x ) , as well as in L f ( z , t ) since Y ( t ) =
y for the unit.
In (cid:15) CF,t , however, unit u (cid:48) =
( x , y (cid:48) , 1 − t ) is involved in q 1 − t ( z | x ) , but not in L f
( z , t ) since Y ( t ) (cid:54) =
y (cid:48)
=
Y (1 − t ) .
Thus, the regression error (second) term in ELBO (7) controls (cid:15) F,t via factual data .
On the other hand, (cid:15) CF,t is not estimable due to the unobservable Y (1 − T ) , but is bounded by (cid:15) F,t plus M D ( x ) in Theorem 2 below––which, in turn, bounds (cid:15) f by decomposing it to (cid:15) F,t , (cid:15) CF,t , and V Y .
Theorem 2 (CATE error bound) .
Assume |L
f
( z , t ) | ≤ M and | g t ( z ) | ≤ G , then:
[Equation]
11) where D ( x ) := (cid:80) t (cid:112) D KL ( q t (cid:107) q 1 − t ) / 2 , and V Y ( x ) :
= E q ( z | x ) (cid:80) t E p Y ( t )
| P t ( y | z ) ( y − j t ( z )) 2 .
D ( x ) measures the imbalance between q t ( Z | x ) and is symmetric for t .
Correspondingly, the KL term in ELBO (7) is also symmetric for t and balances q t ( z | x ) by encouraging Z | =
T | X for the posterior.
V Y ( x ) reﬂects the intrinsic variance in the DGP and can not be controlled.
Estimating G, M is nontrivial.
Instead, we rely on β in the ELBO to weight the terms in (11).
We do not need two hyperparameters since G is implicitly controlled by the third term in ELBO (7), which is a norm constraint.
E XPERIMENTS
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.
E.3.
As in previous works (Shalit et al., 2017; Louizos et al., 2017), we report the absolute error of ATE (cid:15) ate := | E D ( y (1) − y (0))
− E D
ˆ τ ( x ) | and, as a surrogate of square CATE error (cid:15) cate ( x ) =
E D| x
[( y (1) − y (0))
− ˆ τ ( x )] 2 , the empirical PEHE (cid:15) pehe :=
E D (cid:15) cate ( x ) (Hill, 2011), which is the average square CATE error.
Unless otherwise indicated, for each function f , g , h , k , r , s in ELBO (7), we use a multilayer perceptron, with 3 ∗ 200 hidden units, and ReLU activations.
Further, Λ =
( h , k ) depends only on X .
The Adam optimizer with initial learning rate 10 − 4 and batch size 100 is employed.
All experiments use early-stopping of training by evaluating the ELBO on a validation set.
More details on hyper-parameters and settings are given in each experiment.
5. S YNTHETIC DATASET
[Equation]
; T | X ∼ Bern(Logi( ωl ( X )))
; Y | W, T ∼ N ( f T ( W ) , g T ( W )) .
We generate synthetic datasets following (12).
Both X ∼ N ( µ , σ ) and W are factorized Gaussians.
µ , σ are randomly sampled.
The functions h , k , l are linear.
Outcome models f 0 , f 1 are built by
NNs with invertible activations.
Y is univariate, dim( X ) = 30 , and dim( W ) ranges from 1 to 5.
W is a bPGS, but the dimensionality is not low enough to satisfy the injectivity in (G2) , when dim( W ) > 1 .
We have 5 different overlap levels controlled by ω that multiplies the logit value.
See
Sec.
E.1 for details and more results on synthetic datasets.
[Figure]
With the same (dim( W ) , ω ) , we evaluate our method and
CFR on 10 random DGPs, with different sets of functions f, g, h , k , l in (12).
For each DGP, we sample data points, and split them into 3 equal sets for training, validation, and testing.
We show our results for different hyperparameter β .
For CFR, we try different balancing parameters and present the best results (see the Appendix for detail).
In each panel of Figure 2, we adjust one of ω, dim( W ) , with the other ﬁxed to the lowest.
As implied by our theory, our method, with only 1-dimensional Z , performs much better in the left panel (where dim( W ) = 1 satisﬁes (G2) ) than in the right panel (when dim( 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 identiﬁcation and recovery of scores.
(see Figure 3 also).
Under the lowest overlap level ( ω = 22 ), large β (= 2 . 5 , 3) shows the best results, which accords with the intuition and bounds in Sec. 4.
When dim( W ) > 1 , f t in (12) is noninjecitve and learning of PGS is necessary, and thus, larger β has a negative effect.
In fact, β = 1 is signiﬁcantly better than β = 3 when dim( W ) > 2 .
We note that our method, with a higher-dimensional Z , outperforms or matches CFR also under dim( W ) > 1 (see Appendix Figure 5).
Thus, the performance gap under dim( W ) > 1 in Figure 2 should be due to the capacity of NNs in β - Intact-VAE.
In Appendix Figure 7 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 detail).
When dim( W ) = 1 , there are no better PSs than W , because f t is invertible and no information can be dropped from W .
Thus, our method stably learns Z as an approximate afﬁne transformation of the true W , showing identiﬁcation.
An example is shown in Figure 3, and more plots are in
Appendix Figure 9.
For comparison, we run CEVAE
(Louizos et al., 2017), which is also based on
VAE but without identiﬁcation; CEVAE shows much lower quality of recovery.
As expected, both recovery and estimation are better with the balanced prior p Λ ( z | x ) , and we can see examples of bad recovery using p λ ( z | x , t ) in Appendix Figure 10.
5. 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-overlapping, and thus it is used in Johansson et al.
(2020) which considers limited overlap.
The dataset is based on an RCT, but Race is artiﬁcially 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 et al., 2009), are also limited-overlapping.
See Sec.
E.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 is often on small and separated intervals.
Thus, the Gaussian latent Z in our model is misspeciﬁed.
We use 10-dimensional Z to address this, similar to Louizos et al.
We set β = 1 since it works well on synthetic datasets with limited overlap.
As shown in Table 1, β -Intact-VAE outperforms or matches the state-of-the-art methods.
Notably, our method outperforms other generative models (CEVAE and GANITE) by large margins.
Our method has the best ATE estimation and is only slightly worse than CFR for (cid:15) pehe .
This fact reﬂects that (cid:15) pehe is not a good criterion for CATE estimation because it is the marginalized CATE error––one expects less ATE error with overall less CATE error on individual-level, while PEHE focuses on those X
= x with high probability and/or large (cid:15) cate ( x ) .
Indeed, the unconditional balance in Shalit et al.
(2017) is based on bounding PEHE, thus results in sub-optimal ATE estimation (see also Appendix Figure 7 where CFR gives larger ATE errors with less overlap).
To examine conditional v.s unconditional balance clearly, we modify our method and add two components for unconditional balance from Shalit et al.
(2017) (see the Appendix), and compare the modiﬁed version to the original.
The strength of unconditional balance is controlled by hyperparameter γ .
The over-focus on PEHE of the unconditional balance from CFR is seen more clearly in the modiﬁed version.
With different values of the hyperparameter, unconditional balance does not improve (and barely affects) ATE estimation; it does affect PEHE more signiﬁcantly, but often gives worse PEHE unless γ is ﬁne-tuned ( γ = 0 . 1 ). [Table] [Table] [Table]
C ONCLUSION
We proposed a method for CATE estimation under limited overlap.
Our method exploits identiﬁable
VAE, a recent advance in generative models [Table] and is fully motivated and theoretically justiﬁed by causal considerations: identiﬁcation, prognostic score [Table]
and balance. Experiments show evidence that the injectivity of f t in our model is possibly unnecessary because dim( Z ) > dim( Y ) yields better results.
A theoretical study of this is an interesting future direction.
We believe that VAEs are suitable for principled causal inference owing to their probabilistic nature, if not compromised by ad hoc heuristics (see Sec. D.2) [Table]
A P ROOFS
We restate our model identiﬁability formally.
Lemma 1 (Model identiﬁability) .
Given model (3) under (M1) , for T = t , assume (D1’)
(Non-degenerated data for λ ) there exist 2 n + 1 points x 0 , ..., x 2 n ∈ X such that the 2 n -square matrix L t :=
[ γ t, 1 , ..., γ t, 2 n ] is invertible, where γ t,k :=
λ t ( x k ) − λ t ( x 0 ) .
Then, given T = t , the family is identiﬁable up to an equivalence class.
That is, if p θ ( y | x , t ) =
p θ (cid:48)
( y | x , t ) , we have the relation between parameters: for any y t in the image of f t ,
[Equation]
(13) where diag( a ) is an invertible n -diagonal matrix and b is a n -vector, both depend on λ t and λ (cid:48) t .
Note, (D1) in the main text implies (D1’) , see Sec.
B.2.3 in Khemakhem et al.
(2020a).
The main part of our model identiﬁability is essentially the same as that of Theorem 1 in Khemakhem et al.
(2020a), but now adapted to include the dependency on t .
Here we give an outline of the proof, and the details can be easily ﬁlled by referring to Khemakhem et al.
(2020a).
In the proof, subscripts t are omitted for convenience.
Proof of Lemma 1.
Using (M1) i) and ii) , we transform p f , λ
( y | x , t ) =
p f (cid:48) , λ (cid:48)
( y | x , t ) into equality of noiseless distributions, that is,
[Equation]
where p λ is the Gaussian density function of the conditional prior deﬁned in (3) and vol ( A ) :
= √ det AA T .
q f (cid:48) , λ (cid:48) is deﬁned similarly to q f , λ .
Then, apply model (3) to (14), plug the 2 n + 1 points from (D1’) into it, and re-arrange the resulting 2 n + 1 equations in matrix form, we have
[Equation]
where t ( Z ) :
=
( Z, Z 2 ) T is the sufﬁcient statistics of factorized Gaussian, and β t :=
( α t ( x 1 ) − α t ( x 0 ) , ..., α t ( x 2 n ) − α t ( x 0 ))
T where α t ( X ; λ t ) is the log-partition function of the conditional prior in (3).
F (cid:48) is deﬁned similarly to F , but with f (cid:48) , λ (cid:48) , α (cid:48)
Since L is invertible, we have
[Equation]
where
A =
L − T L (cid:48) T and c =
L − T ( β − β (cid:48) ) .
The ﬁnal part of the proof is to show, by following the same reasoning as in Appendix B of Sorrenson et al.
(2019), that A is a sparse matrix such that
[Equation]
where A is partitioned into four n -square matrices.
Thus
[Equation]
where b is the ﬁrst half of c .
Proof of Proposition 2.
Under (G2) , and (M3) , we have
[Equation]
that p
[Equation]
t
[Equation]
x ) >
[Equation]
We show the solution set of (19) on overlapping x is {
[Equation]
n is injective } .
[Equation]
By (G2)(M1) , and with injective f t , j t and dim( Z ) = dim( Y ) ≥ dim( P ) , for any ∆ above, there exists a functional parameter f t such that j t = f t ◦ ∆ .
Thus, set (20) is non-empty, and any element is indeed a solution because f t ◦ h =
j t ◦ ∆ − 1 ◦ ∆ ◦ P =
j t ◦ P .
Any solution of (19) should be in (20).
A solution should satisfy h ( x ) =
f − 1 t ◦ j t ◦ P ( x ) for both t since x is overlapping.
This means the injective function f − 1 t ◦ j t should not depend on t , thus it is one of the ∆ in (20).
We proved conclusion 1) with v := ∆ .
And, on overlapping x , conclusion 2) is quickly seen from
[Equation]
We rely on overlapping P to work for non-overlapping x .
For any x t with p (1 − t | x t ) = 0 , to ensure p (1 − t | P ( x t )) > 0 , there should exist x 1 − t such that P ( x 1 − t ) =
P ( x t ) and p (1 − t | x 1 − t ) > 0 .
And we also have h ( x 1 − t )
=
h ( x t ) due to (M2) .
Then, we have
[Equation]
The third equality uses (19) on ( x 1 − t , 1 − t ) .
Below we prove Theorem 1 with (D2) replaced by (D2’)
(Spontaneous balance) there exist 2 n + 1 points x 0 , ..., x 2 n ∈ X , 2 n -square matrix C , and 2 n -vector d , such that L − 1 0 L 1 = C and β
− C − T β 1 =
d /k for optimal λ t (see below), where L t is deﬁned in (D1’) , β t :=
( α t ( x 1 ) − α t ( x 0 ) , ..., α t ( x 2 n ) − α t ( x 0 ))
T , and α t ( X ; λ t ) is the log-partition function of the prior in (3).
(D2’) restricts the discrepancy between λ 0 , λ 1 on 2 n + 1 values of X , thus is relatively easy to satisfy with high-dimensional X .
(D2’) is general despite (or thanks to) the involved formulation.
Let us see its generality even under a highly special case: C =
c I and d = 0 .
Then, L − 1 0 L 1 =
c I requires that, h 1 ( x k )
− c h 0 ( x k ) is the same for 2 n + 1 points x k .
This is easily satisﬁed except for n (cid:29) m where m is the dimension of X , which rarely happens in practice.
And, β 0 − C − T β 1 = d becomes just β
= c β 0 .
This is equivalent to α 1 ( x k ) − cα
( x k ) same for 2 n + 1 points, again ﬁne 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 λ 0 =
λ 1 in Sec.
4.1, which corresponds to C
=
I and d = 0 .
Proof of Theorem 1.
By (M1) and (G1’) , for any injective function ∆ : P → R n , there exists a functional parameter f ∗ t such that j t = f ∗ t ◦ ∆ .
Let h ∗ t = ∆ ◦ P t , then, clearly from (M3’) , such parameters θ ∗ =
( f ∗ , h ∗ ) are optimal: p θ ∗ ( y | x , t )
= p ( y | x , t ) .
Since have all assumptions for Lemma 1, we have
[Equation]
where f is any optimal parameter, and “ | t ” collects all subscripts t .
Note, except for ∆ , all the symbols should have subscript t .
Nevertheless, using (D2’) , we can further prove A
=
A 1 .
We repeat the core quantities from Lemma 1 here: A t =
L − T t L (cid:48)
Tt and c t =
L − T t ( β t − β (cid:48) t ) .
From (D2’) , we immediately have
[Equation]
And also,
[Equation]
Multiply right hand sides of the two lines, we have c
=
c
Now we have A
=
A 1 := A .
Apply this to (23), we have
[Equation]
for any optimal parameters θ = ( f , h ) .
Again, from (M3’) , we have
[Equation]
where p (cid:15) =
p e .
And the above is only possible when f t ◦ h t =
j t ◦ P t .
Combined with f t =
j t ◦ v − 1 , we have conclusion 1).
And conclusion 2) follows from the same reasoning as Proposition 2, applied to both P 0 and P 1 .
Note, when multiplying the two lines of (25), the effects of k → 0 cancel out, and c t is ﬁnite and well-deﬁned.
Also, it is apparent from above proof that (D2’) is a necessary and sufﬁcient condition for A
=
A 1 , if other conditions of Theorem 1 are given.
Below, we prove the results in Sec.
The deﬁnitions and results work for the prior; simply replace q t ( x | x ) with p t ( z | x ) :
= p λ ( z | x , t ) in deﬁnitions and statements, and the proofs below hold as the same .
The dependence on f prevail, and the superscripts are omitted.
The arguments x are sometimes also omitted.
Lemma 2 (Counterfactual risk bound) .
Assume |L
f
( z , t ) | ≤ M , we have
[Equation]
28) where (cid:15) CF ( x ) :
= (cid:80) t p (1 − t | x ) (cid:15) CF,t ( x ) , and D ( x ) := (cid:80) t (cid:112) D KL ( q t (cid:107) q 1 − t ) / 2 .
Proof of Lemma 2.
[Equation]
TV ( p, q ) :
= 12 E | p ( z ) − q ( z ) | is the total variance distance between probability density p, q .
The last inequality uses Pinsker’s inequality TV ( p, q ) ≤ (cid:112) D KL ( p (cid:107) q ) / 2 twice, to get the symmetric D .
Theorem 2 is a direct corollary of Lemma 2 and the following.
Lemma 3.
Deﬁne (cid:15) F = (cid:80) t p ( t | x ) (cid:15) F,t .
We have
[Equation]
Simply bound (cid:15) CF in (29) by Lemma 2, we have Theorem 2.
To prove Lemma 3, we ﬁrst examine a bias-variance decomposition of (cid:15) F and (cid:15) CF .
[Equation]
The second line uses | g t ( z ) | ≤ G , and the third line is a bias-variance decomposition.
Now we can deﬁne V CF,t ( x ) :=
E q 1 − t ( z | x ) E p Y ( t )
| P t ( y | z ) ( y − j t ( z )) 2 and B CF,t ( x ) :
= E q 1 − t ( z | x ) ( j t ( z ) − f t ( z )) 2 , and we have [Equation]
[Equation]
where V CF :
= (cid:80) t p (1 − t | x ) V CF
[Equation] (cid:80)
[Equation] E [Equation] | x ) E p Y ( t ) | P t ( y | z ) ( y − j t ( z )) 2 and similarly
B CF =
(cid:80) t
E q ( z , 1 − t | x ) ( j t ( z ) − f t ( z )) 2 .
Repeat the above derivation for (cid:15) F , we have (cid:15)
[Equation]
32) where V F = (cid:80) t
E q ( z ,t | x ) E p Y ( t ) | P t ( y | z ) ( y − j t ( z )) 2 and B F = (cid:80) t
E q ( z ,t | x ) ( j t ( z ) − f t ( z )) 2 .
Now, we are ready to prove Lemma 3.
Proof of Lemma 3.
[Equation]
The ﬁrst inequality uses ( a + b ) 2 ≤ 2( a 2 + b 2 ) .
The next equality splits q ( z | x ) into q ( z , 0 | x ) and q ( z , 1 | x ) and rearranges to get B F and B CF .
The last inequality uses the two bias-variance decompositions, and V Y
=
V F + V CF .
B A DDITIONAL BACKGROUNDS B. P ROGNOSTIC SCORE AND BALANCING SCORE
In the fundamental work of (Hansen, 2008), prognostic score is deﬁned equivalently to our P 0 (P0score), but it in addition requires no effect modiﬁcation to work for Y (1) .
Thus, a useful prognostic score corresponds to our PGS. We give main properties of PGS as following.
Proposition 3.
If V gives exchangeability, and P t ( V ) is a PGS, then Y ( t ) | =
V, T | P t .
The following three properties of conditional independence will be used repeatedly in proofs.
Proposition 4 (Properties of conditional independence) .
(Pearl, 2009, Sec.
For random variables W, X, Y, Z .
We have:
[Equation]
Proof of Proposition 3.
From Y ( t ) | =
T | V ( exchangeability of V ), and since P t is a function of V , we have Y ( t ) | =
T | P t , V (1).
From (1) and Y ( t ) | =
V | P t ( V ) (deﬁnition of Pt-score), using contraction rule, we have
Y ( t ) | =
T, V | P t for both t .
Prognostic scores are closely related to the important concept of balancing score (Rosenbaum &
Rubin, 1983).
Note particularly, the proposition implies Y ( t ) | =
T | P t (using decomposition rule).
Thus, if P ( V ) is a P-score, then P also gives weak ignorability (exchangeability and overlap), which is a nice property shared with balancing score, as we will see immediately.
Deﬁnition 4 (Balancing score) .
b ( V ) , a function of random variable V , is a balancing score if
T | =
V | 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 , is a balancing score (with f be the identity function).
Also, given any invertible function v , the composition v ◦ b is also a balancing score since f ◦ v − 1 ( v ◦ b ( V )) =
f ( b ( V )) =
e ( V ) .
Compare the deﬁnition of balancing score and prognostic score, we can say balancing score is sufﬁcient for the treatment T ( T | =
V | b ( V ) ), while prognostic score (Pt-score) is sufﬁcient for the potential outcomes Y ( t )
( Y ( t ) | =
V | P t ( V ) ).
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. VAE, C ONDITIONAL VAE, AND I VAE
VAEs (Kingma et al., 2019) are a class of latent variable models with latent variable Z , and observable Y is generated by the decoder p θ ( y | z ) .
In the standard formulation (Kingma & Welling, 2013), the variational lower bound L ( y ; θ , φ ) of the log-likelihood is derived as:
[Equation]
where D KL denotes KL divergence and the encoder q φ ( z | y ) is introduced to approximate the true posterior p ( z | y ) .
The decoder p θ and encoder q φ are usually parametrized by NNs.
We will omit the parameters θ , φ in notations when appropriate.
The parameters of the VAE can be learned with stochastic gradient variational Bayes.
With Gaussian latent variables, the KL term of L has closed form, while the ﬁrst term can be evaluated by drawing samples from the approximate posterior q φ using the reparameterization trick (Kingma & Welling, 2013), then, optimizing the evidence lower bound (ELBO)
E y ∼D
( L ( y )) with data D , we train the
VAE efﬁciently.
Conditional VAE (CVAE)
(Sohn et al., 2015; Kingma et al., 2014) adds a conditioning variable C , usually a class label, to standard VAE (See Figure 1).
With the conditioning variable, CVAE can give better reconstruction of each class.
The variational lower bound is
[Equation]
The conditioning on C in the prior is usually omitted (Doersch, 2016), i.e., the prior becomes Z ∼
N ( 0 , I ) as in standard VAE, since the dependence between C and the latent representation is also modeled in the encoder q .
Moreover, unconditional prior in fact gives better reconstruction because it encourages learning representation independent of class, similarly to the idea of beta-VAE (Higgins et al., 2017).
As mentioned, identiﬁable VAE (iVAE) (Khemakhem et al., 2020a) provides the ﬁrst identiﬁability result for VAE, using auxiliary variable X .
It assumes Y | =
X | Z , that is, p ( y | z , x ) =
p ( y | z ) .
The variational lower bound is
[Equation]
z | x )) , (35) where Y =
f ( Z ) + (cid:15) , (cid:15) is additive noise, and Z has exponential family distribution with sufﬁcient statistics T and parameter λ ( X ) .
Note that, unlike CVAE, the decoder does not depend on X due to the independence assumption.
Here, identiﬁability of the model means that the functional parameters ( f , T , λ ) can be identiﬁed (learned) up to certain simple transformation.
Further, in the limit of (cid:15) → 0 , iVAE solves the nonlinear ICA problem of recovering Z = f − 1 ( Y ) .
C E XPOSITIONS
The order of subsections below follows that they are referred in the main text.
C. P ROGNOSTIC SCORE IS MORE APPLICABLE THAN BALANCING SCORE
Proposition 3 in D’Amour et al.
(2020) shows that overlapping balancing score implies overlapping
X , and Footnote 5 in D’Amour et al.
(2020) shows that overlapping X implies overlapping bPGS.
Here is a simple example showing overlapping bPGS does not imply overlapping balancing score.
Let T =
I ( X + (cid:15) > 0) and Y = f
( | X | , T )
+ e , where I is the indicator function, (cid:15) and e are exogenous zero-mean noises, and the support of X is on the entire real line while (cid:15) is bounded.
Now, X itself is a balancing score and | X | is a bPGS; and | X | is overlapping but X is not.
C. D ETAILS AND E XPLANATIONS ON I NTACT -VAE
Generative models are useful to solve the inverse problem of recovering Pt-score.
Our goal is to build a model that can be learned by VAE from observational data to obtain a PGS, or more ideally bPGS, via the latent variable Z .
That is, a generative prognostic model.
With the above goal, the generative model of our VAE is built as (3).
Conditioning on X in the joint model p ( y , z
| x , t ) reﬂects that our estimand is CATE given X .
Modeling the score by a conditional distribution rather than a deterministic function is more ﬂexible.
The ELBO of our model can be derived from standard variational lower bound as following:
[Equation]
[Equation]
We naturally have an identiﬁable conditional VAE (CVAE), as the name suggests.
Note that (3) has a similar factorization with the generative model of iVAE (Khemakhem et al., 2020a), that is p ( y , z | x ) =
p ( y | z ) p ( z | x ) ; the ﬁrst 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 et al., 2015; Kingma et al., 2014), with T as the conditioning variable.
See Figurefor the comparison in terms of graphical models.
The core idea of iVAE is reﬂected in our model identiﬁability (see Lemma 1).
C. D ISCUSSIONS AND EXAMPLES OF (G2)
We focus on univariate outcome on R which is the most practical case and the intuitions apply to more general types of outcomes.
Then, i , the mapping between µ 0 and µ 1 , is monotone, i.e, either increasing or decreasing.
The increasing i 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 & Scholkopf, 2010).
The decreasing i corresponds to another common interpretation when ICM does not hold.
Now, the treatment does change the way covariates affect Y , but in a global manner: it acts like a “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 µand µ 1 are both monotone increasing.
This, and also that both µ t are monotone decreasing, are natural and sufﬁcient conditions for increasing i , though not necessary.
The ﬁrst 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, t = receive healthcare program or not), both µ t should be monotone increasing in gestational age.
The next example is from economics.
(Gan & Li, 2016) shows that job-matching probability is monotone increasing in market size.
Then, we can imagine that, with t = receive training in job ﬁnding or not, the monotonicity is not changed.
Intuitively, the examples corresponds to two common scenarios: the causal effects are accumulated though time (the ﬁrst 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 et al., 2015) shows that, with t = mother smokes or not and X = mother’s age, the CATE τ ( x ) is monotone decreasing for 20 < x < 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 a surveyed population (Wang et al., 2020).
Thus, it is convince that, smoking changes the the tendency of birth weight w.r.t mother’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. C OMPLEMENTARITY BETWEEN THE TWO IDENTIFICATIONS
We examine the complementarity between the two identiﬁcations more closely.
The conditions (M3) / (M3’) and (G2) / (D2’) form two pairs, and are complementary inside each pair.
The ﬁrst pair matches model and truth, while the second pair restricts the discrepancy between the treatment groups.
In Theorem 1, (G2) ( P 0 = P 1 ) is replaced by (D2’) which instead makes A
=
A 1 :
=
A in (5).
And (D2’) is easily satisﬁed with high-dimensional X , even if the possible values of C , d are restricted to C
= c
I and d = 0 (see below).
On the other hand, p (cid:15) = p e in (M3’) is impractical, but it ensures that p θ ( y | x , t )
=
p ( y | x , 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 p (cid:15) =
p e can be relaxed.
C. I DEAS AND CONNECTIONS BEHIND THE ELBO (7)
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 ﬂexible Λ , and also the identiﬁcations, the ELBO encourages the recovery of an approximate bPGS as the posterior, which still learns the dependence on T if necessary.
Moreover, β expresses our additional knowledge (or, inductive bias) about whether or not there exist approximate bPGSs (e.g., from domain expertise).
In fact, β connects our VAE to β -VAE (Higgins et al., 2017), which is closely related to noise and variance control (Doersch, 2016, Sec.
2.4)(Mathieu et al., 2019).
Considerations on noise modeling.
In Theorem 1, with large and mismatched noises (then (M3’) is easily violated), the identiﬁcation of outcome model f t =
j t ◦ v − 1 would fail, and, in turn, the prior would learn confounding bias, by confusing the causal effect of T on P T and the correlation between T and X .
This is another reason to prefer λ 0 = λ 1 , besides balancing.
On the other hand, the posterior conditioning on Y provides information of noise e , and it is shown in (Bonhomme & Weidner, 2021) that posterior effect estimation has minimum worst-case error under model misspeciﬁcation (of the noise and prior, in our case).
Under large e , a relatively small β implicitly encourages g smaller than the scale of e , through stressing the third term in ELBO (7).
And the the model as a whole would still learn p ( y | x , t ) well, because the uncertainty of e can be moved to and modeled by the prior.
This is why k is not set to zero because learnable prior noise (variance) allows us to implicitly control g via β .
Intuitively, smaller g strengthens the correlation between Y and Z in our model, and this naturally reﬂects that posterior conditioning on Y is more important under larger e .
Hopefully, precise learning of outcome noise (M3’) is not required, as in Proposition 2.
Now, it is clear that β 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 1: relying on bPGS, or on model identiﬁability; learning loosely, or precisely, the outcome regression.
When the noise scale is different from truth, there would be error due to imperfect recovery of j .
Sec.
4.2 shows that this error and balancing form a trade-off, which is adjusted by β .
Importance of balancing from misspeciﬁcation view.
If we must learn an unapproximate bPGS, we have larger misspeciﬁcation under a balanced prior and rely more on Y in the posterior.
Both are bad because it is shown in (Bonhomme & Weidner, 2021) that posterior only helps under bounded (small) misspeciﬁcation, and posterior estimator has higher variance than prior estimator (see below for an extreme case).
Again, we want a regularizer to encourage learning of bPGS, so that we can explore the middle ground : relatively low-dimensional P , or relatively small e .
Example.
Assume the true outcome noise is (near) zero.
By setting (cid:15) → 0 in our model, the posterior p θ ( z | x , y , t ) =
p θ ( y , z | x , t ) /p
θ ( y | x , t ) degenerates to f − 1 T ( Y ) =
f − 1 T ( j T ( P T ))
= v − 1 ( P T ) , a factual PGS.
However, f − 1 1 − T ( Y ) =
f − 1 1 − T ( j T ( P T ))
= v − 1 ( j − 1 1 − T ◦ j T ( P T ))
(cid:54) = v − 1 ( P 1 − T ) , the score recovered by posterior does not work for counterfactual assignment !
The problem is, unlike 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. C ONSISTENCY OF VAE AND PRIOR ESTIMATION
The following is a reﬁned version of Theorem 4 in Khemakhem et al.
(2020a).
The result is proved by assuming: i) our VAE is ﬂexible enough to ensure the ELBO is tight (equals to the true log likelihood) for some parameters; ii) the optimization algorithm can achieve the global maximum of
ELBO (again equals to the log likelihood).
Proposition 6 (Consistency of Intact-VAE) .
Given model (3) & (6) , and let p ∗ ( x , y , t ) be the true observational distribution, assume
i) there exists (¯ θ , ¯ φ ) such that p ¯ θ ( y | x , t ) =
p ∗ ( y | x , t ) and p ¯ θ ( z | x , y , t ) =
q ¯ φ ( z | x , y , t ) ; ii) the ELBO E D∼ p ∗ ( L ( x , y , t ; θ , φ )) (4) can be optimized to its global maximum at ( θ (cid:48) , φ (cid:48) ) ;
Then, in the limit of inﬁnite data, p θ (cid:48)
( y | x , t ) =
p ∗ ( y | x , t ) and p θ (cid:48) ( z | x , y , t ) =
q φ (cid:48)
( z | x , y , t ) .
Proof.
From i), we have L ( x , y , t ; ¯ θ , ¯ φ ) = log p ∗ ( y | x , t ) .
But we know L is upper-bounded by log p ∗ ( y | x , t ) .
So, E D∼ p ∗ (log p ∗ ( y | x , t )) should be the global maximum of the ELBO (even if the data is ﬁnite).
Moreover, note that, for any ( θ , φ ) , we have D KL ( p θ ( z | x , y , t ) (cid:107) q φ ( z | x , y , t ) ≥ 0 and, in the limit of inﬁnite data, E D∼ p ∗ (log p θ ( y | x , t )) ≤ E D∼ p ∗ (log p ∗ ( y | x , t )) .
Thus, the global maximum of
ELBO is achieved only when p θ ( y | x , t ) =
p ∗ ( y | x , t ) and p θ ( z | x , y , t ) =
q φ ( z | x , y , t ) .
Consistent prior estimation of CATE follows directly from the identiﬁcations.
The following is a corollary of Theorem 1.
Corollary 1.
Under the conditions of Theorem 1, further require the consistency of Intact-VAE.
Then, in the limit of inﬁnite data, we have µ t ( X ) =
f t ( h t ( X )) where f , h are the optimal parameters learned by the VAE.
C. P RE /P OST - TREATMENT PREDICTION
Sampling posterior requires post-treatment observation ( y , t ) .
Often, it is desirable that we can also have pre-treatment prediction for a new 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 q φ with p Λ in (8) 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 p Λ gives the best balanced approximation of the target PGS.
The results of pre-treatment prediction are given in the experimental section E.
C.
A DDITIONAL NOTES ON NOVELTIES OF THE BOUNDS IN S EC . 4.
We give details and additional points regarding the novelties.
Lu et al.
(2020) also use a VAE and derive bounds most related to ours.
Still, our method strengthens Lu et al.
(2020), in a simpler and principled way: we distinguish true score and latent Z and show that identiﬁcation is the link; considering both prior and posterior, we show the symmetric nature of the balancing term and relate it to our KL term in (7), without ad hoc regularization; moreover, we consider outcome noise modeling which is a strength of VAE and relate it to hyperparameter β .
Particularly, in (Lu et al., 2020), latent variable Z is confused with the true representation ( P t up to invertible mapping in our case).
Without identiﬁcation, the method in fact has unbounded error.
Note that Shalit et al.
(2017) do not consider connection to identiﬁcation and noise modeling as well.
The error between ˆ τ f and τ j , which we bound, is due to the unknown outcome noise that is not accounted by our Theorem 1; thus, the theory in Sec.
4.2 is complementary to that in Sec. 3.2.
Finally, β is a trade-off between the conditional balance of learned PGS (affected by f t ), and precision/effective sample size of outcome regression––and can be seen as the probabilistic counterpart of Tarr & Imai (2021) and Kallus et al. D M ORE ON RELATED WORK D. CFR AND CEVAE
CFR and CEVAE are well-known machine learning methods for CATE estimation.
Here we make detailed comparisons to them. D.1. C OMPARISONS WITH CFR Our method is related to CFR in two ways. Theoretically, our bounds in Sec.
4.2 resemble those in Shalit et al.
But we bound CATE error, while CFR bounds PEHE; thus, our bounds give conditional balancing while CFR only has unconditional balancing.
See Sec. C.8 for more on the bounds. Conceptually, CFR is loosely related to our method because it also learns a representation as an outcome predictor, as mentioned in the follow-up Johansson et al.
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 our VAE.
D
Motivation CEVAE is motivated by exploiting proxy variables, and its intuition is 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 identiﬁes CATE.
There is no need to recover the hidden confounder in our framework. Architecture Our model is naturally based on (2), particularly the independence properties of PGS. And as a consequence, our VAE architecture is a natural combination of iVAE and CVAE (see
Figure 1).
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 we have a principled way for balancing. The encoder of CEVAE is even more complex.
To have post-treatment estimation, q ( T | X ) and q ( Y | X, T ) are added into the encoder.
As a result, the ELBO of CEVAE has two additional likelihood terms corresponding to the two distributions.
However, in our Intact-VAE, post-treatment estimation is given naturally by our standard encoder, thanks to the correspondence between our model and (2).
Justiﬁcation We have given the identiﬁcations and bounds of our method in this paper.
Moreover, we carefully distinguish assumptions on the DGP and assumptions on our model, and identify the assumptions that are important for causality.
There are few theoretical justiﬁcations for CEVAE.
Their Theorem 1 directly assumes the joint distribution p ( x , y , t, u ) including hidden confounder
U is recovered, then identiﬁcation 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 to derive conditions under which hidden confounder can be (partially) recovered (Greenland, 1980; Kuroki & Pearl, 2014; Miao et al., 2018).
In particular, Miao et al.
(2018) gives the most recent identiﬁcation result for proxy setting, which requires very speciﬁc two proxies structure, and other completeness assumptions on distributions.
Thus, it is unreasonable to believe that VAE, with simple descendant proxies, can recover the hidden confounder.
Moreover, the identiﬁability of VAE itself is a challenging problem.
As mentioned in Introduction and Sec.
B.2, Khemakhem et al.
(2020a) is the ﬁrst identiﬁability result for VAE, but it only identiﬁes an equivalence class, not a unique representation function.
Thus, it is also unconvincing that VAE can learn a unique latent distribution, without certain assumptions.
As we show in Sec. 5.1, for relatively 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).
D. VAE S FOR TE ESTIMATION
VAEs are suitable for causal estimation thanks to its probabilistic nature.
However, most VAE methods for TEs, e.g. (Louizos et al., 2017; Zhang et al., 2020a; Vowels et al., 2020; Lu et al., 2020), add ad hoc heuristics into their VAEs, and thus break down probabilistic modeling, not to mention identiﬁable representation.
Moreover, the methods rely on learning sufﬁcient representations from proxy variables, leading to either impractical assumptions or conceptual inconsistency, in causal identiﬁcation.
On identiﬁcation.
First, as to causal identiﬁcation, (Louizos et al., 2017) assumes unobserved confounder can be recovered, which is rarely possible even under further structural assumptions (Tchetgen et al., 2020), and (Rissanen & Marttinen, 2021) recently gives evidence that the method often fails.
Other methods (Zhang et al., 2020a; Vowels et al., 2020;
Lu et al., 2020) assume unconfoundedness but still rely on proxy at least intuitively; particularly, (Lu et al., 2020) factorizes the decoder as in the proxy setting.
However, unconfoundedness and proxy should not be put together .
The conceptual inconsistency is that, by deﬁnition, unconfoundedness means covariates fully control confounding, while the motivation for proxy is that unconfoundedness is often not satisﬁed in practice and covariates are at best proxies of confounding, which are non-confounders causally connected to confounders (Tchetgen et al., 2020).
Second, without identiﬁable representation, the empirical results of the methods lacks solid ground; under settings not covered by their experiments, the methods would silently fail to learn proper representations, as we show in Sec. 5.1.
On ad hoc heuristics.
Ad hoc heuristics break down probabilistic modeling and / or give ELBOs that do not estimate the probabilistic models.
For example, (Louizos et al., 2017) uses separated NNs for the two potential outcomes to mimic TARnet (Shalit et al., 2017).
And, to have pre-treatment estimation, q ( T | X ) and q ( Y | X, T ) are added into the encoder.
As a result, the ELBO of (Louizos et al., 2017) has two additional likelihood terms corresponding to the two distributions.
(Zhang et al., 2020a) is even more ad hoc because it splits the latent variable Z into three components, and applies the ad hoc tricks of (Louizos et al., 2017) to each of the component.
Particularly, when constructing the encoder, (Zhang et al., 2020a) implicitly assumes the three components of Z are conditional independent give X , which violates the intended graphical model.
Our method is motivated by the important concept of prognostic score, and is naturally based on (2).
As a consequence, our VAE architecture is a natural combination of iVAE and CVAE (see Figure 1).
Our ELBO (4) is derived by standard variational lower bound.
Moreover, in our β -Intact-VAE, pre-treatment prediction is given naturally by our conditional prior, thanks to the correspondence between our model and (2).
D. I NJECTIVITY , INVERTIBILITY , MONOTONICITY , AND OVERLAP
Let us note that any injective mapping deﬁnes an 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; a monotone mapping can be deﬁned by an injective and order-preserving mapping between ordered sets. Particularly, an injective and continuous mapping on R is monotone , and many works in econometrics give examples of this case.
Many classical and recent works (with many real world applications, see C.1) in econometrics are based on monotonicity.
Particularly, there is a long line of work based on monotonicity of treatment (Huber & W ¨ uthrich, 2018).
More related to our method is another line of work based on monotonicity of outcome , see (Chernozhukov & 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, 2017), where the structural equation of the outcome is assumed to be Y = f ( T ) + (cid:15) , and f is monotone and T (the treatment) is often 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 f is closely related to our method, but the difference is that T is replaced by a PGS in our method, and the PGS is recovered from observables.
Finally, as we mentioned in Sec. 3.2, monotonicity is a kind of shape restriction which also includes, e.g., concavity and symmetry and attracts recent interests (Chetverikov et al., 2018).
However, most of NPIV works focus on identifying f but not directly on TEs, and we do not know any works that use monotonicity to address limited overlap.
Recently in machine learning, (Johansson et al., 2019; Zhang et al., 2020b; Johansson et al., 2020) note the relationship between invertibility and overlap.
As mentioned, (Johansson et al., 2020) 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., 2020b) assumes overlap so that identiﬁcation is given, and then focuses on learning overlapping representation that preserves the overlapping the covariate.
However, it does not relate invertibility and overlap, but uses invertible representation function to preserve exchangeability given the covariate , and linear outcome regression to simply the model.
Related, our identiﬁcations required (M2) , of which linearity of PGS and representation function is a sufﬁcient condition, and our outcome model is injective, to preserve the exchangeability given the PGS .
Thus, our method works under more general setting, and arguably under weaker conditions.
E D ETAILS 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 a testing set.
See also Sec.
C. the pre/post-treatment distinction.
E. S YNTHETIC DATA
[Figure]
We detail how the random parameters in the DGPs are sampled.
µ i and σ i are uniformly sampled in range ( − 0 .
2 , 0 . 2) and (0 , 0 . 2) , respectively.
The weights of linear functions h , k , l are sampled from standard normal distributions.
The NNs f 0 , f 1 use leaky ReLU activation with α = 0 .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 f t is divided by C t :=
Var {D| T = t } ( f t ( Z )) .
When generating DGPs with dependent noise, the variance parameter g t for the outcome is generated by adding a softplus layer after respective f t , and then normalized to range (0 , 2) .
We use the original implementation of CFR 3 .
Very possibly due to bugs in implementation, the CFR version using Wasserstein distance has error 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 divergence, which still give reasonable results.
We search the balancing parameter alpha in [0.16, 0.32, 0.64, 0.8, 1.28], and ﬁx other hyperparameters as they were in the default conﬁg ﬁle.
[Figure]
We characterize the degree of limited overlap by examining the percentage of observed values x that give probability less than 0.001 for one of p ( t | x ) .
The threshold is chosen so that all sample points near those values x almost certainly belong to a single group since we have 500 sample point in total.
If we regard a DGP as very limited-overlapping when the above percentage is larger than 50%, then, as shown in Figure 4, non (all) of the
DGPs are very limited-overlapping with ω = 6 ( ω = 22 ).
For diversity of the datasets, we set g t ( W ) = 1 in DGPs in Appendix.
Figure 5 shows, with dim( Z )
= 200 , our method works better than CFR under dim( W ) =
1 and as well as CFR under dim( W ) > 1 .
As mentioned in
Conclusion, this indicates that the theoretical requirement of injective f t in our model might be relaxed.
Interestingly, larger β seems to give better results here, this is understandable because β controls the tradeoff between ﬁtting and balancing, and the ﬁtting capacity of our decoder is much increased with dim( Z )
= 200 .
Note that the above observations on dim( Z ) are not caused by ﬁxing g t ( W ) =(compare Figure 5 with Figure 6 below)
[Figure]
Figure 6 shows the importance of noise modeling.
Compared to Figure 2 in the main text, where g t ( W ) in DGPs is not ﬁxed, our method works worse here, particularly for large β , because now noise modeling ( g , k in the
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 our VAE, the changes become much insigniﬁcant.
The drop of error for dim( W ) > 3 is due to the randomness of f in (36).
In Sec. 2.2, we saw that the 2-dimsensional bPGS P :=
( µ 0 ( X ) , µ 1 ( X )) always exists under additive noise models.
Thus, when dim( W ) > 2 , our method tries to recover that P , and generally performs not worse than under dim( W )
= 2 , but still not better than under dim( W ) = 1 .
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 ).
This might be another evidence that CFR and its unconditional balancing overﬁt to PEHE (see Sec. 5.2).
Also note that, under dim( W ) = 1 , β = 3 gives the best results for ATE although it does not work well for PEHE, and we 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]
Figure 8 shows 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( W ) = 1 , there is no need to learn PGS. In the right panel, we also do not see signiﬁcant drop of performance compared to posttreatment.
This might be due to the hardness of learning approximate bPGS in this dataset, and posterior estimation does not give much improvements.
You can ﬁnd more plots for latent recovery at the end of the paper.
E. IHDP
IHDP is based on an RCT where each data point represents a baby with 25 features (6 continuous, 19 binary) about their birth and mothers.
Race is introduced as a confounder by artiﬁcially removing all treated children with nonwhite mothers.
There are 747 subjects 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 artiﬁcial by processing the data.
Also our results are only quantitative and we make no ethical conclusions.
[Figure]
The generating process is as following (Hill, 2011,
Sec. 4.1).
[Equation]
where a is a random coefﬁcient, b is a constant bias with all elements equal to 0 .
5 , and c is a random parameter adjusting degree of overlapping between the treatment groups.
As we can see, a T X is a true bPGS.
As mentioned 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 2 is hard, but our VAE works well.
The two added components in the modiﬁed version of our method are as following.
First, we build the two outcome functions f t ( Z ) , t = 0 , 1 in our learning model (3), using two separate NNs.
Second, we add to our ELBO (4) a regularization term, which is the Wasserstein distance (Cuturi, 2013) between
E D∼ p ( X | T = t ) p Λ ( Z | X ) , t ∈ { 0 , 1 } .
As shown in Table 2, best unconditional balancing parameter is 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 ﬁne tuned parameter, but possibly at the price of worse ATE estimation.
Table 3 shows pre-treatment results, All methods gives reasonable results.
[Table]
[Table]
E. P OKEC S OCIAL N ETWORK D ATASET
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 a real world social network dataset.
We experiment on a semisynthetic dataset based on Pokec, which was introduced in (Veitch et al., 2019), and use exactly the same pre-processing and generating procedure.
The pre-processed network has about 79, vertexes (users) connected by 1.3 × 10 6 undirected edges.
The subset of users used here are restricted to three living districts which are within the same region.
The network structure is expressed by 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 of pre-treatment ATE predictions.
We further separate the rest of users (in the other 9 folds) by 6 : 3 , for training and validation.
Each user has 12 attributes, among which district , age , or join date is used as a confounder
U to build 3 different datasets, with remaining 11 attributes used as covariate X .
Treatment T and outcome Y are synthesised as following:
[Equation]
where (cid:15) is standard normal.
Note that district is of 3 categories; age and join date are also discretized into three bins.
g ( U ) , which is a bPGS, maps these three categories and values to { 0 . 15 , 0 . 5 , 0 . 85 } .
β -Intact-VAE is expected to learn a bPGS from G , X , if we can exploit the network structure effectively.
Given the huge network structure, most users can practically be identiﬁed by their attributes and neighborhood structure, which means U can be roughly seen as a deterministic function of
G , X .
This idea is comparable to Assumptions 2 and 4 in (Veitch et al., 2019), which postulate directly that a balancing score can be learned in the limit of inﬁnite large network.
To extract information from the network structure, we use Graph Convolutional Network (GCN) (Kipf & Welling, 2017) in conditional prior and encoder of β -Intact-VAE.
The implementation details are given at the end of this subsection.
Table 4 shows the results.
The pre-treatment √ (cid:15) pehe for Age , District , and Join date confounders are 1.085, 0.686, and 0.699 respectively, practically the same as the ATE errors.
Note that,
Veitch et 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 β -Intact-VAE.
A difﬁculty is that, the network G and covariates X of all users are always needed by GCN, regardless of whether it is in training, validation, or testing phase.
However, the separation can still make sense if we take care that the treatment and outcome are used only in the respective phase, e.g., ( y m , t m ) of a testing user m is only used in testing.
[Table]
GCN takes the network matrix G and the whole covariates matrix X :=
( x T 1 , . . .
, x TM ) T , where
M is user number, and outputs a representation matrix R , 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 it is 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 ( R m, : ) T where x m was expected as input.
And we have respective selection operations for validation and testing.
We can still train β -Intact-VAE including GCN by Adam, simply setting the gradients of non-seleted rows of R to 0.
Note that GCN cannot be trained using mini-batch, instead, we 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 overﬁtting.
E. A DDITIONAL PLOTS ON SYNTHETIC DATASETS
See next pages.