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: The paper proposed a CATE estimation method for graph-structured treatment by Robinson Decomposition. It is an extension of Robinson Decomposition, and the experiment shows the effectiveness of GIN on two small world datasets. The writing is good, but some contents are missing, so that the paper is not that easy to read. My detailed comments and concerns are described in the following.