Identifying and Estimating Causal Effects under Weak Overlap by Generative Prognostic Model

As an important problem of causal inference , we discuss the identification and1 estimation of treatment effects ( TEs ) under weak overlap , i.e. , subjects with certain2 features all belong to a single treatment group . We use a latent variable to model3 a prognostic score ( PGS ) , which is widely used in biostatistics and sufficient for4 TEs , i.e. , we build a generative prognostic model . We prove that the latent variable5 recovers a PGS , and the model identifies individualized treatment effects . The6 model is then learned as the Intact-VAE , a new type of variational autoencoder7 ( VAE ) . We derive counterfactual generalization bounds which motivate represen-8 tation balanced for treatment groups conditioned on individualized features . The9 proposed method is compared with recent methods using ( semi- ) synthetic datasets.10 1 Introduction11 Causal inference [ 21 , 34 ] , i.e , inferring causal effects of interventions , is a fundamental problem . In12 this work , we focus on treatment effects ( TEs ) , such as effects of public policies or a new drug , based13 on a set of observations consisting of binary labels for treatment / control ( non-treated ) , outcome , and14 other covariates ( e.g , patients ’ personal records ) . The fundamental difficulty of causal inference is15 that we never observe counterfactual outcomes , which would have been if we had made the other16 decision ( treatment or control ) . While the ideal protocol for causal inference is randomized controlled17 trials ( RCTs ) , they often have ethical and practical issues , or suffer from prohibitive costs . Thus,18 causal inference from observational data is indispensable . It introduces other challenges , however.19 The most crucial one is confounding : there may be variables ( called confounders ) that causally affect20 both the treatment and the outcome , and spurious correlation follows.21 A large majority of works , including this work , rely on the unconfoundedness , which means that22 appropriate covariates are collected so that the confounding can be controlled by conditioning on23 covariates . That is , all the confounders are in essence observed . This is still challenging , due to24 systematic imbalance ( difference ) of the distributions of the covariates between the treatment and25 control groups , introducing bias in estimation . Among classical ways of dealing with imbalance26 are matching and re-weighting [ 44 , 35 ] . Machine learning methods are also exploited ; there are27 semi-parametric methods , e.g , [ 48 , TMLE ] , which have better finite sample performance , and also28 non-parametric , tree-based methods , e.g. , [ 49 , Causal Forests ( CF ) ] . Notably , starting from [ 23 ] , there29 is a recent rise of interest in learning representation of covariates , which is independent of treatment30 groups , i.e. , balanced representation learning ( BRL ) .31 The most serious form of imbalance is that sample points with certain values of covariate are all32 belong to a single treatment group , which is called weak overlap of the covariate . Causal effects are33 not directly estimable at non-overlapped covariate values . There are lines of work that give robustness34 to weak overlap [ 3 ] , trim non-overlapped sample points [ 52 ] , or study convergence rate depending35 on overlap [ 19 ] . Weak overlap is particularly relevant to machine learning methods exploiting rich36 covariates , because , with higher-dimensional covariates , overlap is harder to satisfy and verify [ 10 ] .37 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . Our approach to the weak overlap issue is based on the prognostic score ( PGS ) [ 14 ] , which is among38 the important concepts of sufficient scores . While the most well-known score is the propensity score39 ( PPS ) [ 36 ] , PGSs have also long been known to improve methods using PPS [ 37 , 5 ] , and interests last40 in biostatistics [ 45 , 2 ] . Prognostic modeling can benefit more from predictive systems and exploit41 richer literature than propensity modeling , particularly in Medicine and Health . A comparative study42 in [ 13 ] shows PGS-based methods perform better , or as well as , PPS methods . Thus , it is promising43 to combine the predictive powers of prognostic modeling and machine learning.44 To solve the inverse problem of recovering PGSs , our method exploits also the recent advance45 of identifiable representation , particularly of VAE [ 26 , iVAE ] . Identification means parameters of46 interest ( for us , representation function and causal effects ) are uniquely determined and given by47 true observational distribution . Identification logically precedes estimation and inference . Without48 identification there is no hope of a consistent estimator , and a model would fail silently ; it may fit49 perfectly but return an estimator that converges to the wrong one or does not converge [ 29 , particularly50 Sec . 8 ] . Identification is even more important for causal inference , because , unlike usual ( non-causal ) 51 model misspecification , causal assumptions are often unverifiable through observables [ 50 ] . Thus , it52 is critical to specify theoretical conditions for identification , and then the applicability of methods53 can be judged by knowledge of an application domain.54 In this work , we study identification ( Sec . 3 ) and estimation ( Sec . 4 ) of TEs under weak overlap . We55 particularly discuss individualized treatment effects , conditioned on the covariates . Code and proofs56 are in Supplementary Materials . The main contributions of this paper are:57 1 ) theory of TE identification under weak overlap of covariates , using PGS and identifiable model ; 58 2 ) counterfactual generalization bounds on TE error , which motivates our conditional BRL ; 59 3 ) a new regularized VAE to estimate TEs , with connections to identification and balancing ; 60 4 ) experimental comparison to state-of-the-art methods on ( semi- ) synthetic datasets.61 2 Setup and motivation62 2.1 Counterfactuals , treatment effects , and identification63 Following [ 21 ] , we introduce potential outcomes ( POs , or counterfactual outcomes ) y ( t ) ∈ Rd , t ∈64 { 0 , 1 } . y ( t ) is the outcome that would have been observed , if treatment value t = t had been applied.65 Formally , this is the consistency of counterfactuals : y = y ( t ) if t = t , or simply y = y ( t ) . We see66 y ( t ) as the hidden variables that give factual y under factual assignment t = t. The fundamental67 problem of causal inference is that , for a unit under research , we could observe only one of y ( 0 ) or68 y ( 1 ) , corresponding the treatment value applied . That is , “ factual ” refers to y or t that is in principle69 observable in data , or statistical entities ( e.g , estimators ) built on them . We also observe relevant70 covariate ( s ) x ∈ X ⊆ Rm , which is associated with individuals , with distribution D ∼ p ( x , y , t ) .71 Note , we use Roman fonts for random variables ( e.g. , t ) and italic for realization ( e.g. , t ) .72 The expected PO is denoted by µt ( x ) = E ( y ( t ) |x = x ) , conditioned on x = x . The estimands in73 this work are the Conditional Average TE ( CATE ) and Average TE ( ATE ) , defined respectively by74 τ ( x ) = µ1 ( x ) − µ0 ( x ) , ν = E ( τ ( x ) ) . ( 1 ) CATE is an individual-level , personalized , treatment effect , given highly discriminative covariate.75 Standard results [ 38 ] [ 16 , Ch . 3 ] give sufficient conditions for identification under general setting.76 They are Exchangeability : y ( t ) |= t|x , and Overlap : p ( t|x ) > 0 for any x ∈ X . Both are required for77 t ∈ { 0 , 1 } . When t appears in a statement without quantification , we always mean “ for both t ” . Often,78 Consistency is also listed , but , as above , it is better known as well-definedness of counterfactuals.79 Exchangeability means , just as in RCTs but additionally given x , that there is no correlation between80 factual treatment t and counterfactual outcomes y ( t ) . Overlap means that the supports of p ( x|t = 0 ) 81 and p ( x|t = 1 ) should be the same , and this ensures it is valid to condition on any ( x , t ) .82 We relax overlapped covariate in Sec . 3.2 , to allow some non-overlapped values x , i.e. , covariate x is83 weakly overlapped . In Sec . 2.2 , we introduce a condition which gives exchangeability ( and PGSs ) . In84 this paper , we also discuss overlap of variables other than x ( e.g . PGSs ) , and we give a definition for85 any random variable v with support V as following.86 Definition 1 . Overlap of v means p ( t|v ) > 0 for all t ∈ { 0 , 1 } , v ∈ V . We also say v is overlapped.87 If the condition is violated at some value v , then v is non-overlapped and v is weakly overlapped.88 2.2 Prognostic scores89 Our method is motivated by PGSs [ 14 ] , adapted as Pt-score and P-score in Definition 2 in this paper,90 related to balancing scores b ( x ) , which is defined by t |= x|b ( x ) [ 36 ] . The PPS p ( t = 1|x ) is a special91 case of this . Both are sufficient scores for identification ; PGSs are sufficient statistics of outcome92 predictors and b ( x ) is for the treatment ( see Appendix for details ) .93 Definition 2 . A Pt-score ( PtS ) is two functions Pt ( x ) ( t = 0 , 1 ) such that y ( t ) |= x|Pt ( x ) . A PtS is94 called a P-score ( PS ) if P0 = P1.95 Note that , a PtS is by definition two functions , thus overlapped Pt ( x ) means both P0 ( x ) and P1 ( x ) 96 are overlapped . Why PtS ( PGS ) ? PtS is more applicable than balancing score b ( x ) under weak97 overlap . Overlapped b ( x ) implies overlapped x , which in turn implies overlapped PGS [ 10 ] . Lower-98 dimensional than x , PtS is likely more overlapped than x , and , moreover , there is evidence that PtS99 maximizes overlap among all sufficient scores for ATE [ 9 ] .100 Below is a direct corollary of Proposition 5 in [ 14 ] . Both of PtS and PS give CATE , but , as we will101 see , PS is better as a conditionally balanced representation , since Pt ( x ) |= t|x only when P0 = P1.102 Proposition 1 ( CATE by PtS ) . If Pt is a PtS , then CATE can be given by103 µt ( x ) = E ( E ( y ( t ) |Pt , x ) ) = E ( E ( y|Pt ( x ) , t = t ) ) = ∫ p ( y|Pt = Pt ( x ) , t ) ydy ( 2 ) With the knowledge of Pt , we choose one of P0 , P1 corresponding to the counterfactual outcome of104 interest . This ability of counterfactual assignment resolves the problem of non-overlap at x.105 PtSs exist under general settings when y ( t ) follows an additive noise model ( ANM ) .106 ( G1 ) ( ANM ) the data generating process ( DGP ) for y is y = f∗ ( M ( x ) , t ) + e where f∗ , M are107 functions and e is a zero-mean exogenous ( external ) noise.108 The DGP defines y ( t ) by setting t = t in the equation . And it also specifies how other variables109 causally affect y . For example , x affects y through M , so M ( x ) is the effect modifier [ 14 ] , which is110 often components of x affecting y directly . Note ( G1 ) also implies exchangeability given x , through111 y ( t ) |= t|M ( x ) . ANMs are also commonly used in nonparametric regression methods for TEs [ 6 ] .112 Under ( G1 ) , 1 ) Pt : = f∗t ( M ( x ) ) = µt ( x ) 1 is a PtS2 but not PS , 2 ) M is a PS ( x is a trivial PS ) , and113 3 ) P : = ( µ0 ( x ) , µ1 ( x ) ) is a PS . We use the same symbol to denote a PtS and the random variable114 defined by it , when appropriate.115 3 Identification under generative prognostic model116 In Sec . 3.1 , we introduce our generative prognostic model and VAE based117 on p ( y , z|x , t ) and prove identifiability of our model . In Sec . 3.2 , we prove118 identification of CATEs , one of our main contributions . The theoretical119 analysis involves only our generative model ( i.e. , prior and decoder ) , but120 not encoder , because model identifiability is a property of model , and121 causal identification is about DGP and model . The encoder is involved122 in estimation which is studied in Sec . 4.123 3.1 Intact-VAE : model , architecture , and identifiability124 Generative models are useful to solve the inverse problem of recovering125 Pt-score . Our goal is to build a model that can be learned by VAE from126 observational data to obtain a PtS , or more ideally PS , via the latent127 variable z . That is , a generative prognostic model.128 With the above goal , the generative model of our VAE is built as129 p ( y , z|x , t ) = p ( y|z , t ) p ( z|x , t ) . ( 3 ) The first factor is our decoder which models p ( y|Pt , t ) in ( 2 ) , and the130 second factor is our conditional prior which models Pt ( x ) . Conditioning on x in the joint model131 1We often write t of function argument in subscript , which indicates possible counterfactual assignment . 2µt is the most common PGS , to the extent that some call it the PGS ( e.g , [ 39 , 9 , 47 ] ) , even without ANMs . p ( y , z|x , t ) reflects that our estimand is CATE given x . Modeling the score by a conditional distri-132 bution rather than a deterministic function is more flexible . We parameterize our model by ANM133 outcome and factorized Gaussian prior as134 pf ( y|z , t ) = p ( y − ft ( z ) ) ; pλ ( z|x , t ) ∼ N ( z ; ht ( x ) , diag ( kt ( x ) ) ) ( 4 ) where θ = ( f , h , k ) are functional parameters and is a noise . λ ( x ) : = diag ( k−1t ( x ) ) ( h ( x ) , − 12 ) T135 is the natural parameter of the Gaussian prior , and we also use it as a shorthand for both h , k.136 The ELBO of our model can be derived from standard variational lower bound137 log p ( y|x , t ) ≥ log p ( y|x , t ) −DKL ( q ( z|x , y , t ) ‖p ( z|x , y , t ) ) = Ez∼q log p ( y|z , t ) −DKL ( q ( z|x , y , t ) ‖p ( z|x , t ) ) . ( 5 ) Our encoder q , which conditions on all the observables , is standard , and we will see its importance138 later . We name this architecture Intact-VAE ( Identifiable treatment-conditional VAE ) .139 We naturally have an identifiable conditional VAE ( CVAE ) , as the name suggests . Note that ( 3 ) has a140 similar factorization with the generative model of iVAE [ 26 ] , that is p ( y , z|x ) = p ( y|z ) p ( z|x ) ; the141 first factor does not depend on x . Further , since we have the conditioning on t in both the factors of142 ( 3 ) , our VAE architecture is a combination of iVAE and CVAE [ 42 , 28 ] , with t as the conditioning143 variable . See Figure 1 for the comparison in terms of graphical models . The core idea of iVAE is144 reflected in our model identifiability ( Lemma 1 below ) . See Appendix for the basics of VAEs.145 The following conditions on the model are used in theoretical analysis.146 ( M1 ) i ) ft is injective , ii ) ft is differentiable , and iii ) n : = dim ( z ) = d ( = dim ( y ) ) .147 Lemma 1 ( Model identifiability ) . Given model ( 3 ) and ( 4 ) under ( M1 ) i ) and ii ) , for t = t , assume148 ( D1 ) ( Linear independence of λ ) there exist 2n+1 points x0 , ... , x2n ∈ X such that the 2n-square149 matrix Lt : = [ γt,1 , ... , γt,2n ] is invertible , where γt , k : = λt ( xk ) − λt ( x0 ) .150 Then , given t = t , the family is identifiable up to an equivalence class . That is , if pθ ( y|x , t = t ) =151 pθ′ ( y|x , t = t ) , we have the relation between parameters : for any yt in the image of ft,152 f−1t ( yt ) = diag ( a ) f ′ t −1 ( yt ) + b = : A ( f ′t −1 ( yt ) ) ( 6 ) where diag ( a ) is an invertible n-diagonal matrix and b is a n-vector , both depend on λt.153 The conditions are inherited from iVAE . ( D1 ) holds easily in practice , if the components of λt ( x ) 154 are linearly independent ; if ( D1 ) fails , then the support of λt ( x ) is in a ( 2n− 1 ) -dimensional space.155 The essence of the result is f ′t = ft ◦ At , that is , ft can be identified ( learned ) up to an affine156 transformation defined by λt . This is achieved by combining the techniques from [ 26 ] and [ 43 ] , and157 essentially the same results can be proved for other exponential family priors [ 43 ] . In this paper,158 symbol ′ ( prime ) always indicates another parameter ( variable , etc . ) .159 3.2 Nonparametric identifications under weakly-overlapped covariate160 In this subsection , we give two identification results based on ( partial ) recovery of PS or PtS,161 respectively . Since PtSs are functions of x , the recovery is achieved by a noiseless prior , that is,162 k ( x ) = 0 ; the prior zλ , t ∼ pλ ( z|x , t = t ) degenerates to deterministic function ht ( x ) .163 PtSs with dimensionality lower than or equal to y are essential to work for weak overlap of x , to the164 extent that , from now on , we simply say PSs / PtSs when referring to this kind of low-dimensional165 PSs / PtSs , unless particularly indicated . ( M1 ) iii ) , i.e . n = d , is not restrictive because µt is a PtS of166 the same dimension as y under ( G1 ) . Also , in practice , n = d means that we seek a low-dimensional167 representation of x . In fact , to make the dimensionality explicit in ( G1 ) , we introduce an alternative168 ( G1 ’ ) which includes ( G1 ) with Pt = µt and jt is identity.169 ( G1 ’ ) ( Low-dimensional PtS ) Under ( G1 ) , µt ( x ) = jt ( Pt ( x ) ) for some Pt and injective jt.170 We use ( G1 ’ ) afterwards . Clearly , Pt in ( G1 ’ ) is a PtS , and injectivity and n = d ensure n =171 dim ( y ) ≥ dim ( Pt ) . Similarly , the next ( G2 ) reduces unverifiable n ≥ dim ( P ) to n = d , for PS.172 ( G2 ) ( Low-dimensional PS ) Under ( G1 ) , µt ( x ) = jt ( P ( x ) ) for some P and injective jt.173 ( G2 ) means that CATEs are given by µ0 and an invertible function i : = j1 ◦ j−10 . See Appendix for174 more discussion and a ( closely related ) real world example . In Sec . 4.1 , we argue that there often175 exist equivalent PSs under ( G1 ’ ) , at least approximately.176 With ( G1 ’ ) or ( G2 ) , overlapped x can be relaxed to overlapped P ( t ) S plus the following.177 ( M2 ) ( Score partition preserving ) For any x , x′ ∈ X , Pt ( x ) = Pt ( x′ ) =⇒ ht ( x ) = ht ( x′ ) .178 Note that ( M2 ) is in fact required for optimal h , in the sense specified in Proposition 1 and Theorem 1179 below . The intuition is that , Pt maps non-overlapped x to an overlapped value , and ht preserves this180 property , through learning . In fact , ( M2 ) is trivially satisfied if Pt and ht are linear , and this is still181 challenging and considered by many works [ 32 , 9 ] , or some with linear outcome models [ 11 , 39 ] .182 Our first identification , Proposition 2 , relies on ( G2 ) and our generative model , without model183 identifiability ( so differentiable ft is not needed ) . This is a nonparametric3 identification under shape184 restriction [ 7 ] , because f , h are functional parameters , and injectivity is monotonicity if jt is on R.185 Proposition 2 ( Identification with PS ) . Given ( G2 ) and model ( 3 ) and ( 4 ) under ( M1 ) i ) and iii ) ,186 and ( M3 ) ( PS matching ) h0 ( x ) = h1 ( x ) and k ( x ) = 0 . Then , if Epθ ( y|x , t ) = E ( y|x , t ) , we have4187 1 ) ( Recovery of PS ) zλ , t = ht ( x ) = v ( P ( x ) ) on overlapped x,188 where v : P → Rn is an injective function and P : = { P ( x ) |overlapped x } 189 2 ) ( Identification ) if P in ( G2 ) is overlapped , and ( M2 ) is satisfied , then µt ( x ) = µ̂t ( x ) 190 for any t ∈ { 0 , 1 } , x ∈ X , where µ̂t ( x ) : = Epλ ( z|x , t ) Epf ( y|z , t ) = ft ( ht ( x ) ) .191 The essence is , i ) the true DGP is identified up to an invertible mapping v , so that ft = jt ◦ v−1 and192 ht = v ◦ Pt , and ii ) Pt is recovered up to v and y ( t ) |= x|Pt is preserved , with same v for both t.193 PS is preferred since it satisfies overlap more easily and ( M2 ) than PtS which refers to two functions.194 However , the existence of low-dimensional PS is uncertain in practice when our knowledge of the195 DGP is limited . Thus , we need Theorem 1 to work under PtS which generally exists.196 Theorem 1 ( Identification with PtS ) . Given the DGP ( G1 ’ ) and model ( 3 ) & ( 4 ) under ( M1 ) and197 ( M3 ’ ) ( Noise matching ) p = pe and k ( x ) = kk′ ( x ) , k → 0 , assume ( D1 ) and198 ( D2 ) ( Spontaneous balance ) There exist 2n + 1 points x0 , ... , x2n ∈ X , 2n-square matrix C,199 and 2n-vector d , such that L−10 L1 = C and β0 − C−Tβ1 = d/k for optimal λt ( see200 below ) , where Lt is defined in ( D1 ) , βt : = ( αt ( x1 ) −αt ( x0 ) , ... , αt ( x2n ) −αt ( x0 ) ) T , and201 αt ( x ; λt ) is the log-partition function of the prior in ( 4 ) .202 Then , if pθ ( y|x , t ) = p ( y|x , t ) , conclusions 1 ) and 2 ) in Proposition 2 hold with P replaced with Pt203 in ( G1 ’ ) , and the domain of v becomes P : = P0 ∪ P1 , Pt : = { Pt ( x ) |overlapped x } .204 Theorem 1 also achieves the two essential points , but in different and complementary ways . Proposi-205 tion 2 starts from the prior by P0 = P1 and setting h0 = h1 . Conversely , Theorem 1 starts from the206 decoder with p = pe and strengthens model identifiability ( 6 ) by ( D2 ) . ( D2 ) restricts the discrepancy207 between λ0 , λ1 on 2n+ 1 points of x , thus is relatively easy to satisfy with high-dimensional x.208 We see more reasons to prefer PS . In general , to identify the mean function µt ( x ) , a regression is209 enough , and p = pe is unnecessary as in Proposition 2 . Also , ( D2 ) is trivial if we have PS and set210 λ0 = λ1 . See Appendix for more on the complementarity between the two identifications.211

Summary: The paper studies identification and estimation of treatment effects (TEs) under weak overlap of covariates (Definition 1, probability of treatment conditioned on covariate has different supports for some covariate, which is closer to the real world). For identification of TEs, the paper exploits the prognostic score (PGS) and the identiﬁable representation, and as a result, proposes a novel, VAE-based generative prognostic model. Prior studies have seen good practical performances of PGS in many settings and hence, the paper expects it to be a good choice and to outperform other candidates in joint use with additional techniques. The identifiable representation is essential to guarantee the identification of TEs. The estimation of the VAE model is via a $\\beta$-regularized-ELBO which helps to keep the posterior distribution of latent variables of VAE ''balanced'' (defined in line 219). Motivations for e.g. balanced representation are provided later in the theoretical analyses on the counterfactual risk bound and the generalization bound. The proposed method is compared with baselines on (semi-)synthetic datasets.