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: This paper proposes a generative prognostic model for identifying and estimating causal effects under weak overlap. The proposed method uses a latent variable to model a prognostic score (PGS). The authors prove that the latent variable recovers a PGS, and the model identifies individualized treatment effects. The model is learnt as an intact-VAE. The authors provide a theory of treatment effect identification and counterfactual generalization bounds on the treatment effect error. The generalization bounds motivate conditional balanced representation learning (BRL). Then they propose a new regularized VAE to estimate treatment effects. In experiments, they compare the proposed method with some state-of-the-art benchmarks on some synthetic and semi-synthetic datasets.