On Inductive Biases for Heterogeneous Treatment Effect Estimation

We investigate how to exploit structural similarities of an individual ’ s potential outcomes ( POs ) under different treatments to obtain better estimates of conditional average treatment effects in finite samples . Especially when it is unknown whether a treatment has an effect at all , it is natural to hypothesize that the POs are similar – yet , some existing strategies for treatment effect estimation employ regularization schemes that implicitly encourage heterogeneity even when it does not exist and fail to fully make use of shared structure . In this paper , we investigate and compare three end-to-end learning strategies to overcome this problem – based on regularization , reparametrization and a flexible multi-task architecture – each encoding inductive bias favoring shared behavior across POs . To build understanding of their relative strengths , we implement all strategies using neural networks and conduct a wide range of semi-synthetic experiments . We observe that all three approaches can lead to substantial improvements upon numerous baselines and gain insight into performance differences across various experimental settings . 1 Introduction . The advent of fields such as personalized medicine has led to rapid growth of the machine learning ( ML ) literature on heterogeneous treatment effect estimation in recent years [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ] . To further advance the understanding of how to incorporate insights from other areas of ML into treatment effect estimation , we revisit the well-established problem of estimating the conditional average treatment effect ( CATE ) of a binary treatment within the potential outcomes ( PO ) framework [ 9 ] . The PO framework allows to conceptualize the problem as estimating the expected difference between an individual ’ s expected ‘ potential ’ outcome with and without treatment , of which only one is observed in the factual world . This fundamental problem of causal inference [ 10 ] leads to the consensus that CATE estimation is not ‘ just another ’ supervised learning problem [ 6 ] . Under the standard assumption of ignorability – which precludes hidden confounding – we consider two statistical features central to estimating CATE : ( i ) the presence of confounding and ( ii ) CATE being a contrast between two PO functions , possibly exhibiting simpler structure than each PO separately . Much of the recent ML literature on CATE estimation has focused on the first feature , and treated confounding as a covariate shift problem . At this point , a range of sophisticated solutions exist which reduce the effect of confounding by balancing the covariate space [ 3 , 4 ] , importance weighting [ 11 , 12 , 13 , 14 ] or propensity drop-out [ 15 ] . How to exploit the second feature in an end-to-end manner , however , has received little explicit attention so far and is what we focus on here . We build on the intuition that the two tasks in the CATE problem – estimating the expected PO with and without treatment – are expected to be strongly related in practical applications ( regardless of the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . presence of confounding ) . In fact , under the common null hypothesis of no treatment effect , we expect them to be identical . Even when there is a non-zero treatment effect , we expect shared structure : In medicine , for example , one distinguishes between biomarkers that are prognostic of outcome regardless of treatment status – hence determining what is shared between the POs – and biomarkers that are predictive of treatment effectiveness – determining heterogeneity [ 16 , 17 ] . The induced difference between the POs is often expected to be small relative to the complexity of the PO functions themselves [ 7 ] , which , as in Fig . 1 , could manifest in terms of CATE being a much simpler function than each PO . One possible way of exploiting this is to estimate CATE directly using recently proposed modelagnostic multi-stage estimation strategies [ 7 , 8 , 18 , 19 ] . However , these estimators output an estimate of CATE only and do not include an estimate of the untreated PO , which is often of independent interest in practical decision support problems – e.g . to trade off outcomes at baseline with the effectiveness of high-risk treatments [ 20 ] . From this literature , we borrow the insight that explicitly targeting learning strategies towards CATE , instead of only the POs , can lead to better estimators . However , instead of relying on multi-stage strategies , we operate within an end-to-end learning framework similar to [ 4 ] and recent extensions . Due to a focus on confounding , this line of work did not explicitly investigate how to exploit the similarity of the PO functions beyond them sharing a jointly learned feature space . Combining the two lines of work , we investigate how to use end-to-end approaches to output better estimates of both the POs and CATE by incorporating the assumption that the POs share much structure ( which can result in a possibly simple CATE ) , as inductive bias . We focus on a fundamental question that has received little explicit attention so far : How can we best exploit the structural similarities of the POs for CATE estimation ? This question is crucial even in randomized experiments , making its solutions orthogonal to any of the sophisticated strategies developed to handle confounding . Therefore , we investigate approaches exploiting the shared structure of the POs which can be applied to modify existing modeling strategies , and thereby aim to provide guidance for improving existing CATE estimators along a new dimension . As such , the goal of this paper is not to promote the use of a specific approach , method or architecture . Rather , we aim to build systematic understanding and greater intuition of the ( dis ) advantages of different approaches . This is crucial in the context of CATE estimation , where model selection is notoriously difficult due to the absence of ground truth treatment effects in practice . We focus on gaining insight into the effect of different approaches relying on the same underlying ML method and use neural networks ( NNs ) due to their flexibility and popularity in related work , yet ( variants of ) the approaches we consider are applicable to many likelihood- or loss-based ML methods . Contributions We investigate three approaches incorporating inductive biases for shared structure into the estimation of the POs : ( 1 ) a soft approach , which relies on regularization to encourage the PO functions to be similar and is hence easy to combine with existing methods , ( 2 ) a hard approach , which hardcodes an assumption on similarity into the model specification by reparametrization of the PO functions and ( 3 ) a flexible approach , in which we build on ideas from multi-task learning to design a new architecture for CATE estimation ( FlexTENet ) , which adaptively learns what to share between the PO functions . We implement instantiations of all approaches using NNs and evaluate their performance across a wide range of semi-synthetic experiments , varying in the structural similarity of the PO functions . We empirically confirm that all approaches can improve upon baselines , including both end-to-end and multi-stage approaches , and present a number of insights into the relative strengths of each approach . We find that strategies significantly changing the model architecture – hard and flexible approaches – usually lead to the largest improvements , with FlexTENet performing best on average ; yet even the simple soft approach often leads to notable performance increases – an insight that can easily be incorporated into any existing method with treatment-specific parameters . 2 Problem Definition and Key Challenges . Assume we observe a sample D = { ( Yi , Xi , Wi ) } ni=1 , with ( Yi , Xi , Wi ) i.i.d.∼ P. Here , Yi ∈ Y is a continuous or binary outcome of interest , Xi ∈ X ⊂ Rd a vector of possible confounders ( i.e . pre-treatment covariates ) and Wi ∈ { 0 , 1 } is a binary treatment , assigned according to propensity score π ( x ) = P ( W = 1|X = x ) . Using the Neyman-Rubin potential outcomes ( PO ) framework [ 9 ] , our main interest lies in the individualized treatment effect : the difference between the PO Yi ( 1 ) if treatment is administered ( Wi = 1 ) and Yi ( 0 ) if individual i is not treated ( Wi = 0 ) . However , only one of the POs is observed as Yi = WiYi ( 1 ) + ( 1−Wi ) Yi ( 0 ) . Therefore , we focus on estimating the conditional average treatment effect ( CATE ) τ ( x ) = E [ Y ( 1 ) − Y ( 0 ) |X = x ] ( 1 ) which is the expected treatment effect for an individual with covariate values X = x . We operate under the standard identifying assumptions in the PO framework : Assumption 1 . [ Consistency , unconfoundedness and overlap ] Consistency : If individual i is assigned treatment wi , we observe the associated potential outcome Yi = Yi ( wi ) . Unconfoundedness : there are no unobserved confounders , so that Y ( 0 ) , Y ( 1 ) ⊥⊥ W |X . Overlap : treatment assignment is non-deterministic , i.e . 0 < π ( x ) < 1 , ∀x ∈ X . 2.1 The Key Challenges of CATE Estimation . As the ability to interpret a treatment effect estimate as causal ultimately relies on a set of untestable assumptions , the unique difficulty in making causal claims lies in using domain expertise to argue whether a treatment effect is identifiable [ 21 ] . Given identifiability , CATE estimation is a purely statistical problem – thus , if one is willing to rely on assumption 1 , CATE can be estimated using observed data . A simple strategy for doing so ( also known as the T-learner [ 7 ] ) obtains regression estimates µ̂w ( x ) of µw ( x ) = E [ Y |X = x , W = w ] , applying standard supervised learning methods using only observed data for which W = w , and finally sets τ̂ ( x ) = µ̂1 ( x ) − µ̂0 ( x ) . Yet , this seemingly straightforward solution is oblivious to two statistical challenges of CATE estimation : 1 . Confounding : If π ( x ) is not constant , then the distribution of covariates in treatment and control groups differs . Such imbalance can be the result of confounders , which are variables that affect both treatment selection and outcomes , and can be problematic when the PO functions are fit on the factual data using empirical risk minimization ( ERM ) because each problem is solved with respect to the wrong empirical distribution – namely X ∼ P ( ·|W = w ) instead of X ∼ P ( · ) . While this problem is not unique to CATE estimation – it is equivalent to the covariate shift problem encountered in e.g . domain adaptation [ 22 ] – it is usually emphasized as one of its main difficulties and motivated the literature on balanced representation learning [ 3 , 4 ] . 2 . CATE is the difference between two functions : While supervised learning usually targets a single function , the goal of CATE estimation is to estimate the difference between two related functions most accurately – which may require different considerations than estimating each function separately . To see this , consider the MSE of estimating CATE and let sq ( f̂ ( x ) ) = EX∼P [ ( f̂ ( X ) −f ( X ) ) 2 ] denote the MSE for an estimate f̂ ( x ) of f ( x ) . If we simply were to estimate τ̂ ( x ) = µ̂1 ( x ) −µ̂0 ( x ) as the difference between two separately learned functions , we would have that ( up to constants ) sq ( τ̂ ( x ) ) . sq ( µ̂1 ( x ) ) + sq ( µ̂0 ( x ) ) . Rateµ1+Rateµ0 . The convergence ratesRateµw depend on the used estimator and assumptions on e.g . smoothness or sparsity of the PO functions ; a well-known example would be [ 23 ] ’ s nonparametric minimax rate . If we had oracle access to both POs and could regress Y ( 1 ) −Y ( 0 ) on X directly , we would have sq ( τ̂ ( x ) ) . Rateτ . The assumption that τ ( x ) is often much simpler than each µw ( x ) separately [ 7 ] translates into Rateτ < maxw Rateµw , highlighting that targeting CATE directly could lead to faster convergence . Similarly , any shared structure across the PO regression tasks could also be exploited to improve upon the simple additive bound above . These observations motivated much of this paper , as they have largely been neglected in the literature on end-to-end learning for CATE . They have , however , been the motivation for some model-agnostic multi-stage learners which we discuss next .

Summary: This paper contributes an investigation of the inductive biases of various neural network architectures in the context of conditional average treatment effect estimation. Specifically, the authors put forth the hypothesis that $\mu_0(x) = E[Y|X=x, W=0]$ and  $\mu_1(x) = E[Y|X=x, W=1]$ have similar structure and thus $\tau(x) = \mu_1(x) - \mu_0(x)$ *can* be simpler than either $\mu_w(x)$.   The authors investigate three approaches: (1) a two path (one per treatment) neural network, where the parameters $\theta_w$ of both paths are regularized to be point-wise close, (2) reparametrization whereby they explicitly learn a $\mu_0$ and a $\tau$ and set $\mu_1$ as $\mu_0 + \tau$, and (3) the authors' proposed approach (FlexTENet) where there are "private" functions $\mu_w$ for each treatment $w$ and a "shared" subspace in each layer.