Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data

1 Introduction . How will a patient ’ s health be affected by taking a medication [ 37 ] ? How will a user ’ s question be answered by a search recommendation [ 34 ] ? We can gain insight into these questions by learning about personalized treatment effects . Estimating personalized treatment effects from observational data is essential when experimental designs are infeasible , unethical , or expensive . Observational data represent a population of individuals described by a set of pre-treatment covariates ( age , blood pressure , socioeconomic status ) , an assigned treatment ( medication , no medication ) , and a posttreatment outcome ( severity of migraines ) . An ideal personalized treatment effect is the difference between the post-treatment outcome if the individual receives treatment and the post-treatment outcome if they do not receive treatment . However , it is impossible to observe both outcomes for an individual ; therefore , the difference is estimated between populations instead . In the setting of binary treatments , data belong to either the treatment group ( individuals that received the treatment ) or the control group ( individuals who did not ) . The personalized treatment effect is the expected difference ∗Equal contribution . Correspondence to { andrew.jesson , panagiotis.tigas } @ cs.ox.ac.uk 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in outcomes between treated and controlled individuals who share the same ( or similar ) measured covariates ; as an illustration , see the difference between the solid lines in Fig . 1 ( middle pane ) . The use of pre-treatment covariates assembled from high-dimensional , heterogeneous measurements such as medical images and electronic health records is increasing [ 44 ] . Deep learning methods have been shown capable of learning personalized treatment effects from such data [ 42 , 43 , 21 ] . However , a problem in deep learning is data efficiency . While modern methods are capable of impressive performance , they need a significant amount of labeled data . Acquiring labeled data can be expensive , requiring specialist knowledge or an invasive procedure to determine the outcome . Therefore , it is desirable to minimize the amount of labeled data needed to obtain a well-performing model . Active learning provides a principled framework to address this concern [ 8 ] . In active learning for treatment effects [ 10 , 45 , 39 ] a model is trained on available labeled data consisting of covariates , assigned treatments , and acquired outcomes . The model predictions decide the most informative examples from data comprised of only covariates and treatment indicators . Outcomes are acquired , e.g. , by performing a biopsy for the selected patients , and the model is retrained and evaluated . This process repeats until either a satisfactory performance level is achieved or the labeling budget is exhausted . At first sight , this might seem simple ; however , active learning induces biases that result in divergence between the distribution of the acquired training data and the distribution of the pool set data [ 13 ] . In the context of learning causal effects , such bias can have both positive and negative consequences . For example , while random acquisition active learning results in an unbiased sample of the training data , we demonstrate how it can lead to over-allocation of resources to the mode of the data at the expense of learning about underrepresented data . Conversely , while biasing acquisitions toward lower density regions of the pool data can be desirable , it can also lead to outcome acquisition for data with unidentifiable treatment effects , which leads to uninformed , potentially harmful , personalized recommendations . To see how training data bias can benefit treatment effect estimation , consider one difference between experimental and observational data : the treatment assignment mechanism is unavailable for observational data . In observational data , variables that affect treatment assignment ( an untestable condition ) may be unobserved . Moreover , the relative proportion of treated to controlled individuals varies across different sub-populations of the data . Fig . 1 illustrates the latter point , where there are relatively equal proportions of treated and controlled examples for data in region 3 . However , the proportions become less balanced as we move to either the left or the right . In extreme cases , say if a group described by some covariate values were systematically excluded from treatment , the treatment effect for that group can not be known [ 38 ] . Fig . 1 illustrates this in region 1 , where only controlled examples reside , and in region 5 , where only treated cases occur . In the language of causal inference , the necessity of seeing both treated and untreated examples for each subpopulation corresponds to satisfying the overlap ( or positivity ) assumption ( see 2.3 ) . The data available in the pool set limits overlap when treatments can not be assigned . In this setting , regions 2 and 4 of Fig . 1 are very interesting because while either the treated or control group are underrepresented , there may still be sufficient coverage to estimate treatment effects . D ’ Amour and Franks [ 9 ] have described such regions as having weak overlap . Training data bias towards such regions can benefit treatment effect estimation for underrepresented data by acquiring low-frequency data with sufficient overlap . We hypothesize that the efficient acquisition of unlabeled data for treatment effect estimation focuses on only exploring regions with sufficient overlap , and uncertainty should be high for areas with non-overlapping support . The bottom pane of Fig . 1 imagines what a resulting training set distribution could look like at an intermediate active learning step . It is not trivial to design such acquisition functions : naively applying active learning acquisition functions results in suboptimal and sample inefficient acquisitions of training examples , as we show below . To this end , we develop epistemic uncertainty-aware methods for active learning of personalized treatment effects from high dimensional observational data . In contrast to previous work that uses only information gain as the acquisition objective , we propose ρBALD and µρBALD as “ Causal BALD ” objectives because they consider both the information gain and overlap between treated and control groups . We demonstrate the performance of the proposed acquisition strategies using synthetic and semi-synthetic datasets . 2 Background . 2.1 Estimation of Personalized Treatment-Effects . Personalized treatment-effect estimation seeks to know the effect of a treatment T ∈ T on the outcome Y ∈ Y for individuals described by covariates X ∈ X . In this work , we consider the random variable ( r.v . ) T to be binary ( T = { 0 , 1 } ) , the r.v . Y to be part of a bounded set Y , and X to be a multi-variate r.v . of dimension d ( X = Rd ) . Under the Neyman-Rubin causal model [ 33 , 41 ] , the individual treatment effect ( ITE ) for a person u is defined as the difference in potential outcomes Y1 ( u ) − Y0 ( u ) , where the r.v . Y1 represents the potential outcome were they treated , and the r.v . Y0 represents the potential outcome were they controlled ( not treated ) . Realizations of the random variables X , T , Y , Y0 , and Y1 are denoted by x , t , y , y0 , and y1 , respectively . The ITE is a fundamentally unidentifiable quantity , so instead we look at the expected difference in potential outcomes for individuals described by X , or the Conditional Average Treatment Effect ( CATE ) : τ ( x ) ≡ E [ Y1 −Y0 | X = x ] [ 1 ] . The CATE is identifiable from an observational dataset D = { ( xi , ti , yi ) } ni=1 of samples ( xi , ti , yi ) from the joint empirical distribution PD ( X , T , Y0 , Y1 ) , under the following three assumptions [ 41 ] : Assumption 2.1 . ( Consistency ) y = tyt + ( 1 − t ) y1−t , i.e . an individual ’ s observed outcome y given assigned treatment t is identical to their potential outcome yt . Assumption 2.2 . ( Unconfoundedness ) ( Y0 , Y1 ) ⊥ T | X . Assumption 2.3 . ( Overlap ) 0 < πt ( x ) < 1 : ∀t ∈ T , where πt ( x ) ≡ P ( T = t | X = x ) is the propensity for treatment for individuals described by covariates X = x . When these assumptions are satisfied , τ̂ ( x ) ≡ E [ Y | T = 1 , X = x ] − E [ Y | T = 0 , X = x ] is an unbiased estimator of τ ( x ) and is identifiable from observational data . A variety of parametric [ 40 , 46 , 42 ] and non-parametric estimators [ 17 , 49 , 3 , 14 ] have been proposed for CATE . Here , we focus on parametric estimators for compactness . Parametric CATE estimators assume that outcomes y are generated according to a likelihood pω ( y | x , t ) , given measured covariates x , observed treatment t , and model parameters ω . For continuous outcomes , a Gaussian likelihood can be used : N ( y | µ̂ω ( x , t ) , σ̂ω ( x , t ) ) . For discrete outcomes , a Bernoulli likelihood can be used : Bern ( y | µ̂ω ( x , t ) ) . In both cases , µ̂ω ( x , t ) is a parametric estimator of E [ Y | T = t , X = x ] , which leads to : τ̂ω ( x ) ≡ µ̂ω ( x , 1 ) − µ̂ω ( x , 0 ) , a parametric CATE estimator . Jesson et al . [ 20 ] have shown that Bayesian inference over the model parameters ω , treated as stochastic instances of the random variable Ω ∈ W , yields a model capable of quantifying when assumption 2.3 ( overlap ) does not hold . Moreover , they show that such models can quantify when there is insufficient knowledge about the treatment effect τ ( x ) because the observed value x lies far from the support of PD ( X , T , Y0 , Y1 ) . Such methods seek to enable sampling from the posterior distribution of the model parameters given the data , p ( Ω | D ) . Each sample , ω ∼ p ( Ω | D ) induces a unique CATE function τ̂ω ( x ) . Jesson et al . [ 20 ] propose Varω∼p ( Ω|D ) ( µ̂ω ( x , 1 ) − µ̂ω ( x , 0 ) ) as a measure of epistemic uncertainty ( i.e. , how much the functions “ disagree ” with one another at a given value x ) [ 23 ] for the CATE estimator . 2.2 Active Learning . Formally , an active learning setup consists of an unlabeled dataset Dpool = { xi } npool i=1 , a labeled training set Dtrain = { xi , yi } ntraini=1 , and a predictive model with likelihood pω ( y | x ) parameterized by ω ∼ p ( Ω | Dtrain ) . The setup also assumes that an oracle exists to provide outcomes y for any data point in Dpool . After model training , a batch of data { x∗i } bi=1 is selected from Dpool using an acquisition function a according to the informativeness of the batch . By including the treatment , we depart from the standard active learning setting . For active learning of treatment effects , we define Dpool = { xi , ti } npool i=1 , a labeled training set Dtrain = { xi , ti , yi } ntrain i=1 , and a predictive model with likelihood pω ( y | x , t ) parameterized by ω ∼ p ( Ω | Dtrain ) . The acquisition function takes as input Dpool and returns a batch of data { xi , ti } bi=1 which are labelled using an oracle and added to Dtrain . We are specifically examining the case when there is access to only the treatments observed in the pool data { ti } npool i=1 : i.e. , scenarios where treatment assignment is not possible . An intuitive way to define informativeness is using the estimated uncertainty of our model . In general , we can distinguish two sources of uncertainty : epistemic and aleatoric uncertainty [ 11 , 23 ] . Epistemic ( or model ) uncertainty , arises from ignorance about the model parameters . For example , this is caused by the model not seeing similar data points during training , so it is unclear what the correct label would be . We focus on using epistemic uncertainty to identify the most informative points for label acquisition . Bayesian Active Learning by Disagreement ( BALD ) [ 18 ] defines an acquisition function based on epistemic uncertainty . Specifically , it uses the mutual information ( MI ) between the unknown output and model parameters as a measure of disagreement : I ( Y ; Ω | x , Dtrain ) = H ( Y | x , Dtrain ) − Eω∼p ( Ω|Dtrain ) [ H ( Y | x , ω ) ] , ( 1 ) where H is the entropy function ; a straightforward estimand for discrete outcomes with Bernoulli or Categorical likelihoods . The general acquisition function based on BALD for acquiring a batch of data points given the pool dataset and the model parameters is given by the joint mutual information between the set { Yi } and the model parameters [ 26 ] : aBALD ( Dpool , p ( Ω | Dtrain ) ) = argmax { xi } bi=1⊆Dpool I ( { Yi } ; Ω | { xi } , Dtrain ) . ( 2 ) This batch acquisition function can be upper-bounded by scoring each point in Dpool independently and taking the top b ; however , this bound ignores correlations between the samples . In fact , for datasets with significant repetition , this approach can perform worse than random acquisition , and computing the joint mutual information ( introduced as BatchBALD ) rectifies the issue [ 26 ] . Estimating the joint mutual information is computationally expensive , as evaluating the joint entropy over all possible outcomes ( for classification ) or a covariance matrix over all inputs ( for regression ) is required . An alternative approach is to use softmax-BALD , which involves importance weighted sampling across Dpool with the individual importance weights given by BALD [ 25 ] . We use softmaxBALD for batch acquisition because it is computationally more efficient and performs competitively with BatchBALD . We discuss how BALD maps onto epistemic uncertainty quantification in CATE and the arising complications stemming from the question of overlap in Section 3 .

Summary: This paper proposed a new acquisition function for applying active learning to observation data consisting of covariates, treatments and outcomes. The motivation of developing a new acquisition function is the existing ones could either query too many samples close to the distribution mode or query samples outside the overlapping region of the treatment and control group. Therefore, the key contribution of this work is to propose a acquisition function that 1) query samples in the overlapping region of case and control; and 2) reduce uncertainity of the model parameters.  The proposed function is built on the existing Bayesian active learning by disagreement (BALD).   First, \tauBALD aims to maximize the conditional mutual information between CATE and model parameters \Omega. However, since the control outcome and treatment outcome cannot be observed simultaneously, leading to irreducible uncertainity.   Second, to overcome this drawback, \muBALD was considered by maximizing the conditional mutual information between outcome Y^t and the model parameters \Omega. The proposed Causal-BALD improved \muBALD by considering both model uncertainty reduction and overlapping. Specifically, \mu\piBALD directly multiplies the propensity acquisition function and \muBALD.   Third, \rhoBALD can also handle overlap by maximizing the mutual information between outcome Y^t and the CATE computed from model parameter \Omega. \rhoBALD can also be combined to \muBALD to form \mu\rhoBALD, which can handle both model uncertainity reduction and overlap.  Experimental results on synthetic dataset and semi-synthetic data showed the proposed \mu\rhoBALD outperforms existing acquisition function for the application active learning to individual treatment effect estimation.