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 introduces "CausalBALD", a set of acquisition functions for use in Bayesian Active Learning when trying to estimate the Conditional Average Treatment Effect ("CATE") associated with an intervention, $\tau(\mathbf{x}) = \mathbb{E}[Y^{1}-Y^{0}|\mathbf{x}]$.  The authors introduce multiple information-theoretic acquisition functions, and then show the improved performance (relative to random sampling of examples) of their acquisition functions for estimating the true treatment effects present in three simulated/semi-synthetic datasets.  Through a toy problem (Figures 2/3), they also provide some intuition for why their acquisition functions outperform random sampling.