Causal Effect Inference for Structured Treatments

1 Introduction . Estimating feature-level causal effects , so-called conditional average treatment effects ( CATEs ) , from observational data is a fundamental problem across many domains . Examples include understanding the effects of non-pharmaceutical interventions on the transmission of COVID-19 in a specific region [ 12 ] , how school meal programs impact child health [ 13 ] , and the effects of chemotherapy drugs on cancer patients [ 52 ] . Supervised learning methods face two challenges in such settings : ( i ) missing interventions , the fact that we only observe one treatment for each individual means models must extrapolate to new treatments without access to ground truth , and ( ii ) confounding factors that affect both treatment assignment and the outcome means that extrapolation from observation to intervention requires assumptions . Many approaches have been proposed to overcome these issues [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 9 , 10 , 15 , 18 , 19 , 21 , 22 , 23 , 25 , 27 , 29 , 33 , 39 , 41 , 42 , 45 , 52 , 56 , 57 , 60 , 64 , 67 ] . In many cases , treatments are naturally structured . For instance , a drug is commonly represented by its molecular structure ( graph ) , the nutritional content of a meal as a food label ( text ) , and geographic regions affected by a new policy as a map ( image ) . Taking this structure into account can provide several advantages : ( i ) higher data-efficiency , ( ii ) capability to work with many treatments , and ( iii ) generalizing to unseen treatments during test time . However , the vast majority of prior work operates on either binary or continuous scalar treatments ( structured treatments are rarely considered , a notable exception to this trend is Harada & Kashima [ 16 ] which we describe in Section 2 ) . To estimate CATEs with structured interventions , our contributions include : • Generalized Robinson decomposition ( GRD ) : A generalization of the Robinson decomposition [ 47 ] to treatments that can be vectorized as a continuous embedding . This GRD reveals a learnable ∗Correspondence to jean.kaddour.20 @ ucl.ac.uk 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . pseudo-outcome target that isolates the causal component of the observed signal by partialling out confounding associations . Further , it allows one to learn the nuisance and target functions using any supervised learning method , thus extending recent work on plug-in estimators [ 42 , 29 ] . • Quasi-oracle convergence guarantee : A result that shows that given access to estimators of certain nuisance functions , as long as the estimates converge at an O ( n−1/4 ) rate , the target estimator for the CATE achieves the same error bounds as an oracle who has ground-truth knowledge of both nuisance components , the propensity features , and conditional mean outcome . • Structured Intervention Networks ( SIN ) : A practical algorithm using GRD , representation learning , and alternating gradient descent . Our PyTorch [ 43 ] implementation is online.2 • Evaluation metrics designed for structured treatments . Since previous evaluation protocols of CATE estimators have mostly focused on binary or scalar-continuous treatment settings , we believe that our proposed evaluation metrics can be useful for comparing future work . • Experimental results with graph treatments in which SIN outperforms previous approaches . 2 Related Work . Closest to our work is GraphITE [ 16 ] , a method that learns representations of graph interventions for CATE estimation . They propose to minimize prediction loss plus a regularization term that aims to control for confounding based on the Hilbert-Schmidt Independence Criterion ( HSIC ) [ 14 ] . This technique suffers from two drawbacks : ( i ) the HSIC requires multiplication of kernel matrices and scales quadratically in the batch size ; ( ii ) selecting the HSIC kernel hyper-parameter is not straightforward , as ground-truth CATEs are never observed , and empirical loss does not bound CATE estimation error [ 1 ] . We discuss other related work not on structured treatments in Appendix A . 3 Preliminaries . 3.1 Conditional Average Treatment Effects ( CATEs ) . Imagine a dataset where each example ( xi , ti , yi ) ∈ D represents a hospital patient ’ s medical history record xi , prescribed drug treatment ti , and health outcome yi , as illustrated in Figure 1 ( Center ) . Further , we wish to understand how changing the treatment changes a patient ’ s health outcome . The CATE , τ ( t′ , ti , xi ) , describes the expected change in outcome for individuals with history xi , when treatment ti is replaced by t′ , depicted in Figure 1 ( Right ) . In real-world scenarios , we only observe one outcome for each patient at one treatment level . Further , the patient ’ s pre-treatment health conditions xi influence both the doctor ’ s treatment prescription and outcome , thereby confounding the effect of the treatment on the outcome . Formally , we have the dataset D= { ( xi , ti , yi ) } n i=1 sampled from a joint distribution p ( X , T , Y ) , where Y = f ( X , T ) + ε , as depicted in Figure 1 ( Left ) . We define the causal effect of fixing 2https : //github.com/JeanKaddour/SIN treatment variable T ∈ T to a value t on outcome variable Y ∈ R using the do-operator [ 44 ] as E [ Y | do ( T = t ) ] . Crucially , this estimate differs from the conditional expectation E [ Y | T = t ] in that it describes the effect of an external entity intervening on T by fixing it to a value t ( removing the edge X → T ) . We further condition on pre-treatment covariates X to define the conditional causal estimand E [ Y | X = x , do ( T = t ) ] . The conditional average treatment effect ( CATE ) is the difference between expected outcomes at different treatment values t , t′ for given covariates x , τ ( t′ , t , x ) , E [ Y | X = x , do ( T = t′ ) ] ︸ ︷︷ ︸ = : µt′ ( x ) −E [ Y | X = x , do ( T = t ) ] ︸ ︷︷ ︸ = : µt ( x ) , ( 1 ) where µt ( x ) is defined as the expected outcome for a covariate vector x under treatment t. Because we do not observe both treatments t , t′ for a single covariate x , we need to make assumptions that allow us to identify the CATE from observational data . Assumption 1 . ( Unconfoundedness ) There are no confounders of the effect between T and Y beyond X . Therefore , Pr ( Y ≤ y | x , do ( t ) ) = Pr ( Y ≤ y | x , t ) , for all ( x , t , y ) . Assumption 2 . ( Overlap ) It holds that 0 < p ( t | x ) < 1 , for all ( x , t ) . Assumption 2 means that all sub-populations have some probability of receiving any value of treatment ( otherwise , some τ ( t′ , t , x ) may be undefined or impossible to estimate . ) These assumptions allow us to estimate the causal quantity τ ( t′ , t , x ) through statistical estimands : τ ( t′ , t , x ) = µt′ ( x ) − µt ( x ) = E [ Y | X = x , T = t′ ] − E [ Y | X = x , T = t ] . ( 2 ) While one can model µt ( x ) with regression models , such approaches suffer from bias [ 9 , 26 , 29 ] due to two factors : ( i ) associations between X and T , due to confounding , makes it hard to identify the distinct contributions of X and T on Y , and ( ii ) regularization for predictive performance can harm effect estimation . Mitigating these biases relies on exposing and removing nuisance components . This transforms the optimization into a ( regularized ) regression problem that isolates the causal effect . 3.2 Robinson Decomposition . One way to formulate such nuisance components is via the Robinson decomposition [ 47 ] . Originally a reformulation of the CATE for binary treatments , it was used by the R-learner [ 42 ] to construct a plug-in estimator . The R-learner exploits the decomposition by partialling out the confounding of X on T and Y . It also isolates the CATE , thereby removing regularization bias . Let the treatment variable be T ∈ { 0 , 1 } and the outcome model p ( y | x , t ) parameterized as Y = f ( X , T ) + ε ≡ µ0 ( X ) + T × τb ( X ) + ε , ( 3 ) where we define error term ε such that E [ ε | x , t ] = E [ ε | x ] = 0 , and τb ( x ) , τ ( 1 , 0 , x ) . Define the propensity score [ 48 ] e ( x ) , p ( T = 1 | x ) and the conditional mean outcome as m ( x ) , E [ Y | x ] = µ0 ( x ) + e ( x ) τb ( x ) . ( 4 ) From model ( 3 ) and the previous definitions , it follows that Y −m ( X ) = ( T − e ( X ) ) τb ( X ) + ε , ( 5 ) allowing us to define the estimator τ̂b ( · ) = arg min τb { 1 n n∑ i=1 ( ỹi − t̃i × τb ( xi ) ) 2 + Λ ( τb ( · ) ) } , ( 6 ) where ỹi , yi−m̂ ( xi ) and t̃i , ti−ê ( xi ) are pseudo-data points defined through estimated nuisance functions m̂ ( · ) , ê ( · ) , which can be learned separately with any supervised learning algorithm . 4 The Generalized Robinson Decomposition . Our goal is to estimate the CATE τ ( t′ , t , x ) for structured interventions t′ , t ( e.g. , graphs , images , text ) while accounting for the confounding of X on T and Y . Inspired by the Robinson decomposition , which has enabled flexible CATE estimation for binary treatments [ 6 , 9 , 33 , 42 ] , we propose the Generalized Robinson Decomposition from which we extract a pseudo-outcome that targets the causal effect . We demonstrate the usefulness of this decomposition from both a theoretical view ( quasi-oracle convergence rate in Section 4.2 ) and practical view ( Structured Intervention Networks in Section 5 ) . For details on its motivation and derivation , we refer the reader to Appendix B . 4.1 Generalizing the Robinson Decomposition . To generalize the Robinson decomposition to structured treatments , we introduce two concepts : ( a ) we assume that the causal effect is a product effect : the outcome function f∗ ( X , T ) can be written as an inner product of two separate functionals , one over the covariates and one over the treatment , and ( b ) propensity features , which partial out the effects from the covariates on the treatment features . Similar techniques have been previously shown to add to the robustness of estimation [ 9 , 42 ] . Assumption 3 . ( Product effect ) We consider the following partial parameterization of p ( y | x , t ) , Y = g ( X ) > h ( T ) + ε , ( 7 ) where g : X → Rd , h : T → Rd and E [ ε | x , t ] = E [ ε | x ] = 0 , for all ( x , t ) ∈ X × T . This assumption is mild , as we can formally justify its universality . The following asserts that provided we allow the dimensionality of g and h to grow , we may approximate any arbitrary bounded continuous functions in C ( X × T ) where X × T is compact . Proposition 1 . ( Universality of product effect ) LetHX×T be a Reproducing Kernel Hilbert Space ( RKHS ) on the set X × T with universal kernel k. For any δ > 0 , and any f ∈ HX×T , there is a d ∈ N such that there exist two d-dimensional vector fields g : X → Rd and h : T → Rd , where ‖f − g > h‖L2 ( PX×T ) ≤ δ . ( Proof in Appendix C ) This assumption allows us to simplify the expression of the CATE for treatments t′ , t , given x , τ ( t′ , t , x ) = g ( x ) > ( h ( t′ ) − h ( t ) ) . ( 8 ) Define propensity features eh ( x ) , E [ h ( T ) | x ] and m ( x ) , E [ Y | x ] = g ( x ) > eh ( x ) . Following the same steps as in Section 3.2 , the Generalized Robinson Decomposition for eq . ( 7 ) is Y −m ( X ) = g ( X ) > ( h ( T ) − eh ( X ) ) + ε . ( 9 ) Given nuisance estimates m̂ ( · ) , êh ( · ) , we can use this decomposition to derive an optimization problem for h ( · ) , g ( · ) ( note êh ( · ) implicitly depends on h ( · ) , we address this dependence in Section 5 ) . ĝ ( · ) , ĥ ( · ) , arg min g , h 1n n∑ i=1 ( Yi − m̂ ( Xi ) − g ( Xi ) > ( h ( Ti ) − êh ( Xi ) ) ) 2 + Λ ( g ( · ) ) ( 10 )

Summary: This paper studies the estimation of conditional average treatment effect (CATE) with multiple treatments that are graph structured. The proposed method extends the R-learner (Nie and Wager, 2020), which is designed for binary treatment, by learning a representation that reduces the dimension of the treatments. The representation is parametrized by a graph neural network that takes the graph structure into account. Unlike the standard R-learner that learns the nuisance components $E[Y|X]$ and $E[T|X]$ separately from the CATE, the proposed method only estimates $E[Y|X]$ separately while learns the representation $h(T)$, the propensity features $E[h(T)|X]$, and CATE jointly via gradient-based methods.