Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data
Authors: Andrew Jesson
Panagiotis Tigas
Joost van Amersfoort
Andreas Kirsch
Uri Shalit
Yarin Gal
Venue: NeurIPS 2021 Poster (2021)
Keywords: Causal Inference
Active Learning
Bayesian Deep Learning

Abstract: Estimating personalized treatment effects from high-dimensional observational data is essential in situations where experimental designs are infeasible, unethical, or expensive. Existing approaches rely on fitting deep models on outcomes observed for treated and control populations. However, when measuring individual outcomes is costly, as is the case of a tumor biopsy, a sample-efficient strategy for acquiring each result is required. Deep Bayesian active learning provides a framework for efficient data acquisition by selecting points with high uncertainty. However, existing methods bias training data acquisition towards regions of non-overlapping support between the treated and control populations. These are not sample-efficient because the treatment effect is not identifiable in such regions. We introduce causal, Bayesian acquisition functions grounded in information theory that bias data acquisition towards regions with overlapping support to maximize sample efficiency for learning personalized treatment effects. We demonstrate the performance of the proposed acquisition strategies on synthetic and semi-synthetic datasets IHDP and CMNIST and their extensions, which aim to simulate common dataset biases and pathologies.

Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data
Andrew Jesson"
Panagiotis Tigas"
Joost van Amersfoort OATML University of Oxford
Andreas Kirsch OATML University of Oxford
OATML University of Oxford
OATML
University of Oxford
Uri Shalit
Yarin Gal OATML University of Oxford
Machine Learning and Causal Inference in Healthcare Lab Technion
Abstract
Estimating personalized treatment effects from nigh-dimeesional observational data is essential in situations where experimental designs are infeasible, unethical, or expensive. Existing approaches rely on fitting deep models on outcomes observed for treated and control populations. However, when measuring individual outcomes is costly, as the case of tumor biopsy, sample-efficient strategy for acquiring each result is required. Deep Bayesian active learning provides framework for efficient data acquisition by selecting points with high uncertainty. However, existing methods bias training data acquisition towards regions of non-overlapping support between the treated and control populations. These are not sample-efficient because the treatment effect is not identifiable in such regions. We introduce causal, Bayesian acquisition functions grounded in information theory that bias data acquisition towards regions with overlapping support to maximize sample efficiency for learning personalized treatment effects. We demonstrate the performance of the proposed acquisition strategies synthetic and semi-synthetie datasets IHDP and CMNIST and their extensions, which aim to simulate common dataset biases and pathologies.
Introduction
How will a patient's 's health be affected by taking a medication [37]? How will user's question be answered by 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 set of pre-treatment covariates (age, blood pressure, socioeconomic status), an assigned treatment (medication, no medication), and post- treatment outcome (severity of migraines). An ideal personalized treatment effect is the difference between the post-treatment outcome 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 either the treatment group (individuals that received the treatment) the control group (individuals who did not). The personalized treatment effect is the expected difference
Equal contribution. Correspondence (andrew. jeeson, pawagiotis.tigee..............t tigas)esss....... ox. ac..uk
35th Conference on Neural Information Processing Systems (NcurIPS 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. (middle pane).
The use of pre-treatment covariates assembled from high-dimensional, heterogeneous measurements such 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, problem in deep learning data efficiency. While modern methods are capable of impressive performance, they need significant amount of labeled data. Acquiring labeled data can be expensive, requiring specialist knowledge or an invasive procedure to determine the outcome. Therefore, is desirable to minimize the amount of labeled data needed to obtain a well verf-orrfon 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 biopsy for the selected patients, and the model is retrained and evaluated. This process repeats until either satisfactory performance level is achieved 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 the mode of the data 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 unavailable for observa- tional data. In observational data, variables that af- feet treatment assignment (an untestable condition) may be unobserved. Moreover, the relative propor- tion of treated controlled individuals varies across different sub-populations of the data. Fig. illus- trates 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
deruity
controlted
5
nutccoe
move to either the left or the right. In extreme cases, say a group described by some covariate values were systematically excluded from treatment, the treatment effect for that group cannot be known [38]. Fig. illustrates this in re- gion 1, where only controlled examples reside, and in region where only treated cases occur. In the language of causal inference, the necessity of seeing both treated and untreated examples for each sub- population corresponds to satisfying the overlap (or positivity) assumption (see 2.3). The data available in the pool set limits overlap when treatments cannot be assigned. In this setting, regions and 4 of Fig. I are very interesting because while either the treated or control group are underrepresented, there may still be sufficient coverage to estimate treatment effects. D' A'mou 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.
Pau [KConbra) (xtreattd)
Civi value:
Figure 1: Observational data. Top: data density of treatment (right) and control (left) groups. Middle: observed outcome re- sponse for treatment (circles) and control (x's s) groups. Bottom: data density for active learned training set after number of acquisi- tion steps.
We hypothesize that the efficient acquisition of unlabeled data for treatment effect estimation foeuses on only exploring regions with sufficient overlap, and uncertainty should be high for areas with non-overlapping support. The bottom pane Fig. 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, we show below. To this end, we develop epistemic uncertainty-awary 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 pBALD and upBALD 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-Effect
Personalized treatment-effect estimation seeks to know the effect of treatment T E T on the outcome Y E ) for individuals described by covariates X E X. In this work, we consider the random variable (I.v.) T to be binary (T {0, 1}), the r.v. Y to be part of bounded set V, and X to be a multi-variate r.v. of dimension d (X' = IR.). Under the Peyman-Rubil causal model [33, 41], the individual treatment effect (TTE) for a person is defined as the difference in potential outcomes (u) ( (u), where the r.v. Y represents the potential outcome were they treated, and the n.v. Y represents the potential outcome were they controlled (not treated). Realizations of the random variables X, T, Y. Y", and Y are denoted by x, t, y. y", and y², respectively.
The ITE 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): T(x) = E Y X x] [1]. The CATE is identifiable from an observational dataset i=
D = {(x;, l1, [= of samples (x,, yi) from the joint empirical distribution Po(X, T, Y Y
under the following three assumptions [41]:
Assumption 2.1. (Consistency)) y tty + (1- +(--tly' i.e an individual's observed outcome y given assigned treatment t is identical to their potential outcome y".
Assumption 2.2. (Unconfouuddedness) (Y,,Y²)
Y") 1 |
Assumption 2.3. (Overlap) 0 < (x) < 1 Vt ET,
where Te xxx P(T = t |
X = x) the propensity for treatment for individuals described by covariates X = x. When these assumptions are satisfied, f(x) = E TT= 1, = x] - -E
T= 0, = x] is an unbiased estimator of r(x) and identifiable from observational data.
variety of parametric [40, 46, 42] and non- parameeretic 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 (y x, x,L). given measured covariates x, observed treatment and model parameters w. For continuous outcomes, a Gaussian likelihood can be used: M( mw (x, (), T. (x, t)). For discrete outcomes, a Bernoulli likelihood can be used: Bern(y Pw (x, t)). In both cases, mw (x, t) is parametric estimator of E|Y T x]. which leads to: (x) = (x, 1) (x, 0), a parametric CATE estimator.
Jesson al. [20] have shown that Bayesian inference over the model parameters w, treated as stochastic instances of the random variable Ω E W. yields a model capable of quantifying when assumption 2.3 (overlap) does not hold. Moreover, they show that such models quantify when there is insufficient knowledge about the treatment effect 7 r(x) because the observed value x lies far from the support of Po (X, T, Y Y1). Such methods seek to enable sampling from the posterior distribution of the model parameters given the data, p(2 DD. Each sample, w ~ p(52 'D) induces unique CATE function F. (x). Jesson et al. [20] propose Varw>p(((pp) (Pw (x, 1) Tie (x, 0)) as a measure of epistemic uncertainty (i.c., how much the functions "disagree" with one another at given value x) [23] for the CATE estimator.
2.2 Active Learning
Formally, an active learning setup consists of an unlabeled dataset D Dpoot = {xx) a labeled training set Dtrain = (x,, J Peratn and a predictive model with likelihood (y x) parameterized by p(2 Durala. The setup also assumes that oracle exists to provide outcomes y for any
3

---

data point in D pool. After model training, batch of data {x; }:=1 is selected from D pool 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 = (x+, t. 1=1==1 s=1 labeled training set Dirain = (x;, 14, yx);= Rirain and a predictive model with likelihood Pw (y x, t) parameterized by w p Dirain). The acquisition function takes as input D pool and returns a batch of data fx, t2 )im=- which are labelled using an oracle and added to Draiin We are specifically examining the case when there is access to (t)) only the treatments observed in the pool data {t, ("poo: i..., scenarios where treatment assignment is
not possible.
An intuitive way to define informativeness using the estimated uncertainty of our model. 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 (Mf) between the unknown output and model parameters as a measure of disagreement:
I(Y: 1 x, Purain) H(Y x, Durain) - (St [H(Y) x, w)]
(1)
where H is the entropy function; straightforward estimand for discrete outcomes with Bernoulli or Categorical likelihoods.
The general acquisition function based BALD for acquiring batch of data points given the pool dataset and the model parameters is given by the joint mutual information between the set {Y;} and the model parameters [26]:
@BALD (D pool p(5 D Dtrain))
arg max
1( {x:), Dtrain).
(2)
(x. (x)) EPool
This batch acquisition function can be upper-bounded by scoring each point in DD Poool 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, 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 voftmax-BALD, which involves importance weighted sampling across Ppool with the individual importance weights given by BALD [25]. We use softmax- BALD 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.
3 Methods
In this section: we introduce several acquisition functions, we then analyze how they bias the acquisition of training data, and show the resulting CATE functions learned from such training data. We are interested in acquisition functions conditioned on realizations of both x and t:
((Ppool, ((2 Deruin))
arg max
I( [xx, [2), Diruin).
(x, (x,tt))
where T(
X, t, Dtrain) is a measure of disagreement between parametric function predictions given and over samples p(l D). We make assumptions 2.1 and 2.2 (consistency, and unconfoundedness). We relax assumption 2.3 (overlap) by allowing for its violation over subsets of the support of Pooo. We present all theorems, proofs, and detailed assumptions in Appendix A.
3.1 How do naive acquisition functions bias the training data?
To motivate Causal-BAL.D, we first look at set of naive acquisition functions. A random acquisition function selects data points uniformly at random from Pooo and adds them to Dusin- In Fig. 2a

---

(a) Random
(b) Propensity
~BALD
pBBLD
Naive acquisition functions: How the training set is biased and how this effects the CATE function with a fixed budget of 300 acquired points.
Figure 2:
have acquired 300 such examples from synthetic dataset and trained deep-kernel Gaussian process [47] on those labeled examples. Comparing the top two panels, we see that Dtrain (middle) contains an unbiased sample of the data in D pool (top). However, in the bottom panel, we see that while the CATE estimator accurate and confident near the modes of 'D Ppool. it becomes less accurate as we move lower-density regions. In this way, the random acquisition of data reflects the biases inherent in D pool and over-allocates resources to the modes of the distribution. If the mode were to coincide with region of non-overlap, the function would most frequently acquire uninformative examples.
Next, we look at using the propensity score to bias data acquisition toward regions where the overlap assumption is satisfied.
Definition 3.1. Counterfactual Propensity Acquisition
I( x, t, Dtrain) = 1 (x)
(4)
Intuitively, this function prefers points where the propensity for observing the counterfactual is high. We are considering the setup where Dpool contains observations of both X and T, so is straightforward to train an estimator for the propensity, मt (x). Figure 2b shows that while propensity score acquisition matches the treated and control densities in the train set, it still biases data selection towards the modes of D Prool.
The goal of BALD to acquire data (x, t) that maximally reduces uncertainty in the model parameters Ω used to predict the treatment effect. The most direct way to apply BALD is to use our uncertainty over the predicted treatment effect, expressed using the following information theoretic quantity: Definition 3.2. TBALD
I(Y1 Y Ω x, t, Dirain)
Var (SPp(())Perran))
(Dw 1)
Pow (x, Px, 0)).
(5)
Building off the result in [20], we show how the LHS measure about the unobservable potential outcomes is estimated by the variance over Ω of the identifiable difference in expected outcomes in Theorem of the appendix. Alaa and van der Schaar [4] propose a similar result is for non-parametric models. Intuitively, this measure represents the information gain for S2 if we observe the difference in potential outcomes Y Y for a given measurement x and Derain.
However, a fundamental flaw with this measure exists: observing labels for the random variable Y1 Y impossible. Thus, rBALD represents an irreducible measure of uncertainty. That is, rBALD will be high if it is uncertain about the label given the unobserved treatment t', regardless of its certainty about the outcome given the factual treatment t, which makes rBALD highest for low-density regions and regions with no overlap. Figure 2c illustrates these consequences. We see the acquisition biases the training data away from the modes of the Ppoot, where we cannot know the treatment effect (no overlap). In datasets where we have limited overlap, it leads to uninformative acquisitions.
One remedy to the issues of rBALD is to only focus on reducible uncertainty:
Definition 3.3. μBALD
I(Y Ω x, t, Derain)
Var (X,, t)).
(6)
wrop( Derma))
5

---

Prs Mikeetos
pol Sendy
Plo Infrrsted
control trenteh
Injestrot Iepretest
trih Senday
Auqurtd: 100
Acqured 393
Araaied
(a) BALD
(b) pBALD
(c) upBALD
Figure 3: Causal-BALD acquisition functions: How the training set is biased and how this effects the CATE function with a fixed budget of 300 acquired points.
This measure represents the information gain for the model parameters 12 if we obtain label for the observed potential outcome Y* given a data point (x, and Persin. We give proof for these results in Theorem of the appendix.
μBALD only contains observable quantities; however, does not account for our belief about the counterfactual outcome. As illustrated Fig. 2d, this approach can prefer acquiring (x, t) when are also very uncertain about (x, ('), even if (x, (') is in D Pooo Since we can neither reduce uncertainty over such (x, U') nor know the treatment effect, the acquisition function would not be optimally data efficient.
3.2 Causal-BALD.
In the previous section, we looked at naive methods that either considered overlap or considered information gain. In this section, we present three measures that account for both factors when choosing new point to acquire for model training.
A straightforward to combine knowledge about a data point's 's information gain and overlap is to simply multiply μBALD(6) by the propensity acquisition term (4):
Definition 3.4. μTBALD
[(uT x, t, t, Dtrain) -1 (x)) w Dresin)
Var
(Pw (x, t))
We in Fig. 3a that the acquisition of training data results in matched sampling that for propensity acquisition in Fig. 2b. However, the tails of the overlapping distributions extend further into the low-density regions of the pool set support where the overlap assumption is satisfied.
Alternatively, we can take an information-thhorreie approach to combine knowledge about a data point's information gain and overlap. Let fi be an instance of the random variable 1 E R corresponding to the expected outcome conditioned on t. Further, let T (x) be an instance of the random variable in Un ph corresponding to the CATE. Then,
Definition 3.5. pBALD
Varw (Xw t)) 2 Cov (Pw (x, (x, +1). Varw (Pw (x, t'))
I(Y"; in xx Durin)
2 log
(8)
This measure represents the information gain for the CATE TQ if we observe the outcome Y for datapoint (x, t) and the data we have trained on Dtrain- We give proof for this result Theorem 3.
In contrast to B-BALD, this measure accounts for overlap in two ways. First, p-BALD will be scaled by the inverse of the variance of the expected counterfactual outcome Po Pw, t'). This scaling biases acquisition towards examples for which we know about the counterfactual outcome, so we can assume that overlap is satisfied for observed (x, t). Second, p-BALD discounted by Cov (Pw t), (x, t')). This discounting concept that we will leave for future discussion.

---

In Fig. 3b we see that p-BALD has matched the distributions of the treated and control groups similarly to propensity acquisition in Fig. 2b. Further, we see that the CATE estimator more accurate over the support of the data.
There is, however, shorteoming of p-BALD that may lead to suboptimal data efficiency. Con- sider two examples in poool (x1, ti) and (x2, t2) where Varw (Pw (x1, (x,,tt) t1)) = Var (āw (x1, ti ) and (Nw t²)) Var (Pw bg))): for each point, we are as uncertain about the condi- tional expectation given the factual treatment as we are uncertain given the counterfactual treat- ment. Further, let (Pw 11). Plw (x1,
Cov (X2, Pw (x),tt)). Finally, let Var (Pw (x1,tz) t1)) > Varw (Pw (x2, t2)): we are more uncertain about the conditional expectation given the factual treatment for data point (x1, t1) than we are for data point (x2, t2). Under these three conditions. p p-BALD would rank these two points equally, and so this method would bias training data to the modes of Ppool when (X2; 12) is more frequent than (x1, f1). practice, may be more data-efficient to choose (x1, (1) over (x2, 12) as it would more likely be a point as yet unseen by the model.
To combine the positive attributes of p-BALD and p-BALD, while mitigating their shortcomings, we introduce ppBALD. Definition 3.6. poBALD
T(pp x, x, ttrriin) = Var (x,, t)) Varw (Fw (x))
Var (Th (x, t'))
(Fu Here, we scale Equation 8, which has equivalent expression Varw (x,1") by our measure for μBALD such that in the cases where the ratio may be equal, there is a preference for data points the current model is more uncertain about. We can see in Fig. 3c that the training data acquisition is distributed more uniformly over the support of the pool data where the overlap assumption is satisfied. Furthermore, the accuracy of the CATE estimator is highest over that region.
4 Related Work
Deng et al. [10] propose the use of Active Learning for recruiting patients to assign treatments that will reduce the uncertainty of an Individual Treatment Effect model. However, their setting is different from ours we assume that suggesting treatments are too risky or even potentially lethal. Instead, we acquire patients to reveal their outcome (e.g.. by having a biopsy). Additionally, although their method uses predictive uncertainty to identify which patients to recruit, it does not disentangle the sources of uncertainty; therefore, it will also recruit patients with high outcome variance. Closer to our proposal is the work from Sundin et al. [45]. They propose using Gaussian process (GP) to model the individual treatment effect and use the expected information gain over the S-type error rate, defined the error in predicting the sign of the CATE, as their acquisition function. Although GPs are suitable for quantifying uncertainty, they do not work well on high-dimensional input spaces. In this work, we use Neural network methods to obtain uncertainty: Deep Ensembles [28] and DUE [47], a Deep Kernel Learning GP, both of which work well even on high dimensional inputs. Additionally, the authors assume that noisy observations about the counterfactual treatments are available at training time where we make no such assumptions. We compare to this in our experiment by limiting the access to counterfactual observations (7 baseline) and adapting it to Deep Ensembles [28] and DUE [47] (we provide more details about the adaptation in Appendix B. 1). Recent work by Qin et al. [39] looks at budgeted heterogeneous effect estimation but does not factor weak or limited overlap into their acquisition function.
7

---

Experiments
In this section, we evaluate our acquisition objectives on synthetic and semi-synthetic datasets. Code to reproduce these experiments is available at https: //github. httssthu.cooaaaeeeeee
Control
Count
Treated
Datasets Starting from the hypothesis that differ- ent objectives can target different types of imbal- ances and degrees overlap, we construct a synthetic dataset [22] demonstrating the various biases. We
depict this dataset graphically in Fig. 1. We use this dataset primarily for illustrative purposes. By design.
Outcome
we have constructed a primary data mode and have study
regions of weak no overlap. Additionally, the performance of our acquisition functions on the
IHDP dataset [17, 42], which is standard bench- mark in causal treatment effect literature. Finally, we demonstrate that method is suitable for high di- mensional, large-sample datasets on CMNIST [21]. an MNIST [29] based dataset adapted for causal treat- ment effect studies. In Fig. 4, we see that CMNIST is an adaptation of the synthetic dataset. Model inputs are MNIST digits and assigned treatments, and the response surfaces are generated based on projection of the digits onto latent 1-dimensional manifold. The observed digits are high-dimensional proxies for the confounding covariate d. Detailed descriptions of each dataset are available in Appendix C.
-2
0 Covariate embedding ø
2
Figure 4: Visualizing CMNIST dataset. Model inputs are MNIST digits and assigned treatments. The MNIST digits are high- dimensional proxies for the latent confound- ing covariate Φb Digits are projected onto by ordering them first by image intensity and then by digit class (0 9). Methods must be able to implicitly learn this non-linear map- ping in order to predict the conditional ex- pected outcomes.
Model Our objectives rely on methods that are ca- pable of modeling uncertainty and handling high- dimensional data modalities. DUE [47] is an instance of Deep Kernel Learning [48] that uses deep feature extractor to transform the inputs and defines Gaussian process (GP) kernel over the extracted feature representation. In particular, DUE uses variational inducing point approximation [16) and a constrained feature extractor that contains residual connections and spectral normalization to enable reliable uncertainty. Due obtains SotA results on IHDP [47]. In DUE, we distinguish between the model parameters and the variational parameters w, and we are Bayesian only over the parameters. Since DUE is GP, we obtain a full Gaussian posterior over outcomes from which we can use the mean and covariance directly. When necessary, sampling is very efficient and only requires single forward pass in the deep model. We describe all hyperparameters in Appendix F.
Baselines We compare against the following baselines: Random. This acquisition function selects points uniformly at random. Propensity. An acquisition function based on the propensity score (Eq. 4). We train propensity model on the pool data, which we then use to acquire points based on their propensity score. Please note that this is valid assumption training propensity model does not require outcomes. y (S-type error rate) [45]. This acquisition function the S-type error rate based method proposed by Sundin et al. [45]. We have adapted the acquisition function to use with Bayesian Deep Neural Networks. The objective is defined as I(>: x, Durain), where
(-
and probit (.) is the cumulative distribution function of normal distribution. In contrast to the original formulation, we do not assume access to counterfactual observations training time.
(x) probit
/ (V
5.1 Experimental Results
For each of the acquisition objectives, dataset, and model we present the mean and standard error of empirical square root of precision in estimation of heterogenous effect (PEHE) 2, We summarize
VEPEHE= = V Σ. (f(x) r(x))2
8

---

Table 1: Summary of active learning parameters for each dataset.
Dataset
Warm-up size
Acquisition size
Acquisition steps
Pool Size
Valid Size
Synthetic
10
10
30
10k
lk
IHDP
100
10
38
471
201
CMNIST
250
50
55
35k
15k
Baselines:
random μuBAALD
propensity pBALD
BALD objectives:
UMBALD
-
μBAALD
TBALD
159
130 222
252
304
1900
1168
2003
252
3000 2002
Figure 5: VEREbE performance (shaded standard error) for DUE models. (left to right) synthetic (40 seeds), and IHDP (200 seeds). We observe that BALD objectives outperform the random, 7 and propensity acquisition functions significantly, suggesting that epistemic uncertainty aware methods that target reducible uncertainty can be more sample efficient.
each active learning setup in Table 1. The warm up size the number of examples in the initial pool dataset. Acquisition size is the number of examples labeled at each acquisition step. Acquisition steps the number of times we query a batch of labels. Pool size is the number of examples the pool dataset. Finally, valid size is the number of examples used for model selection when optimizing the model each acquisition step.
In Fig. 5, we see that epistemic uncertainty aware μpBALD outperforms the baselines, random, propensity, and S-Type error rate (y). As analyzed in section 3, we expect this improvement as our acquisition objectives target reducible uncertainty that is, epistemic uncertainty when there overlap between treatment and control. Additionally, BoBALDshows superior performance over the other objectives in the high dimensional dataset CMNIST verifying our qualitative analysis in Figure:: 3c.
Each (μBALD, pBALD, HTBALD, and μpBALD) outperform the baseline methods on these tasks. Of note, the performance /BALD improves as the dimensionality of the covariates increases. In contrast, the performance of the propensity score-based BALD worsens as the dimensionality of the covariates increases. Propensity score estimation is known be a problem in high-dimensions [12]. We see that both μBALD and μpBALD perform consistently as dimensionality increases, with μpBALD showing a statistically significant improvement over 4BALD on two of the three tasks. These improvements indicate that μpBALD is more robust for data with high sigh-dimensional covariates than um BALD moreover, upBALD does not need additional propensity score model.

---

6 Conclusion
We have introduced new acquisition function for active learning of individual-Ievel causal-treatmen treatment effects from high dimensional observational data, based on Bayesian Active Learning by Disagree- ment [18]. We derive our proposed method from an information-theoreiie perspective and compare it with acquisition strategies that either do not consider epistemic uncertainty (i.e., random or propensity- based) or target irreducible uncertainty in the observational setting (i.e., when we do not have access to counterfactual observations). We show that our methods significantly outperform baselines, while also studying the various properties of each of our proposed objectives in both quantitative and quali- tative analyses, potentially impacting areas like healthcare where sample efficiency in the acquisition of new examples implies improved safety and reductions in costs.
Broader Impact
Active I Learning for learning treatment effects from observational data is highly actionable research, and there are several sectors where our research can have an impact. Take, for example, hospital that needs to decide whom to treat, based some model. To these choices, the decision-maker needs to have a confident and accurate treatment effect prediction model. However, improving the performance of such a model requires data from patients, which might be costly and perhaps even unethical to acquire. With this work, we assume that we cannot assign new treatments to patients but only perform biopsies questionnaires post- pest-treatment reveal the outcome. We believe that this is an impactful and realistic scenario that will directly benefit from our proposal. However, our method can also impact fields like computational aaaiiiinnaaadeeeiiemeentt where the goal is to learn model to predict the captivate the attention of users, or policymaking where government wants to decide how to intervene for beneficial or malicious reasons.
Active learning inherently biases the acquisition of training data. We attempt to show how this can be beneficial or detrimental for learning treatment effects under different acquisition functions under the unconfoundedness assumption. We are not making guarantees on the overall unbiasedness of our methods. We guarantee only that results are conditional on the unconfoundedness assumption. Unobserved confounding can result in the biased sampling of training data concerning the hidden confounding variable. This bias can result in performance inequality between groups and biased estimate of the unconfounded CATE function. Further, the models' uncertainty estimates are not informative of when this may occur.
An anonymous reviewer writes, "one risk is that the method could, e.g., lead to biased, non- representative sampling in terms of ethnicity and other protected attributes particularly, if the unconfoundedness assumption this work based on is blindly trusted" Recent work by Andrus et al. [5] does a great job of highlighting the difficulties practitioners face when accounting for algorithmic bias across protected attributes. In such cases, model uncertainty is not enough to identify non-representative sampling concerning the protected attribute. They report about practitioner's use of structural causal models in concert with domain expert feedback as a means to inform clients of potential sources of bias. Perhaps such methodology could be used with causal sensitivity analysis for CATE [22, 50, 21) as a way to model beliefs about the protected attribute without observing it.
Impactful avenues for future work include relaxing the unconfoundedness assumption, incorporating beliefs about hidden confounding into the acquisition function. Furthermore, in addition to uncovering the outcome, we think it is interesting to revisit the active treatment assignment problem. On the active learning side, exploring more rigorous batch acquisition methods could yield improvements over the current stochastic sampling estimation we use. Finally, in this work, assume access to a validation set, which may not be available in practice, so exploration of active acquisition of validation data will also have an impact [27].
Acknowledgments and Disclosure of Funding
We would like to thank OATML group members and all anonymous reviewers for sharing their valuable feedback and insights. PT and AK are supported by the UK EPSRC CDT Autonomous Intelligent Machines and Systems (grant reference EP/L 58//1015997/1). JvA is grateful for funding by the EPSRC (grant reference EP/N50971 and by Google-DeepMind. U.S. was partially supported by the Israel Science Foundation (grant No. 1950/19).
10

---

References
[1] Jason Abrevaya, Yu- Yu-Chin Hsu, and Robert P Licli. Estimating conditional average treatment effects. Journal of Business & Economic Statistics, (3(4))485-505, 2015.
[2] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprini arXiv:1803.08375, 2018.
[3] Ahmed M Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In Advances in Neural Information Processing Systems, pages 3424-3432, 2017.
[4] Ahmed M Alaa and Mihaela der Schaar. Bayesian nonparametric causal inference: Infor- mation rates and learning algorithms. IEEE Journal of Selected Topics in Signal Processing, 12 (5)103311046,, 2018.
[5] McKane Andrus, Elena Spitzer, Jeffrey Brown, and Alice Xiang. What we can't measure, we can't tt understand: Challenges to demographic data procurement in the pursuit of fairesss. In Proceedings of the 2021 ACM Conference Fairness, Accountability, and Transparency, pages 249-260, 2021.
[6] James Bergstra, Daniel Yamins, and David Cox. Making science of model search: Hyperpa- rameter optimization hundreds of dimensions for vision architectures. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learn- ing. volume 28 of Proceedings of Machine Learning Research, pages 115-123, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL http: best:://procee http://pproceedings mlr press http///////////aaa/eee html.
[7] Djórk-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). In International conference on machine learning. pages 448-456 PMLR, 2016.
[8] David A Cohn, Zoubin Ghahramani, and Michael I Jordan. Active learning with statistical models. Journal of artificial intelligence research, 4:129-145, 1996.
[9] Alexander D' Amour and Alexander Franks. Deconfounding scores: Feature representations for causal effect estimation with weak overlap. arXiv preprint arXiv:2104.05762, 2021.
[10] Kun Deng, Joelle Pineau, and Susan Murphy. Active learning for personalizing treatment. In 2011 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pages 32-39. IEEE, 2011.
[11] Armen Der Kiureghian and Ove Ditlevsen. Aleatory or epistemic? does matter? Structural safery, (1(2):055-11, 2009.
[12] Alexander D' D'Amorr, Peng Ding, Avi Feller, Lihua Lei, and Jasject Sekhon. Overlap in observational studies with high- high-dimensional covariates. Journal of Econometrics, 221(2): 644 654, 2021.
[13] Sebastian Farquhar, Yarin Gal, and Tom Rainforth. On statistical bias in active learning: How and when to fix it. In International Conference on Learning Representations, 2021. URL
https://openreviier net/ JiX
[14] Zijun Gao and Yanjun Han. Minimax optimal nonparametric estimation of heterogeneous treatment effects. In Advances in Neural Information Processing Systems, 2020.
[15] Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael Cree. Regularisation of neural networks by enforcing lipschitz continuity. Machine Learning, 110(2):393 416, 2021.
[16] James Hensman, Alexander Matthews, and Zoubin Ghahramani. Scalable variational gaussian process classification. In Artificial Intelligence and Statistics, pages 351-360. PMLR, 2015.
[17] Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computa- tional and Graphical Statistics, 20(1):2172240, 2011.

---

[18] Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning. stat, 1050:24, 2011.
[19] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448 456. PMLR, 2015.
[20] Andrew Jesson, Sören Mindermann, Uni Shalit, and Yarin Gal. Identifying causal-effect inference failure with mucertainty-wwan models. Advances in Neural Information Processing Systems, 33, 2020.
[21] Andrew Jesson, Sören Mindermann, Yarin Gal, and Uri Shalit. Quantifying ignorance in individual-level causal-effect estimates under hidden confounding. In International Conference on Machine Learning, pages
PMLR, 2021.
[22] Nathan Kallus, Xiaojie Mao, and Angela Zhou. Interval estimation of individual-level causal effects under unobserved confounding. In The 22nd International Conference on Artificial Intelligence and Statistics, pages
PMLR, 2019.
[23] Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in Neural Information Processing Systems, 30:5574-5584, 2017.
[24] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[25] Andreas Kirsch. PowerEvaluationn Efficient evaluation-oriented deep (bayesian) active learning with stochastic acquisition functions. arXiv preprint arXiv:2101 100.05552,. 2021.
[26] Andreas Kirsch, Joost Van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. Advances in neural information processing systems, 32:7026-7037, 2019.
[27] Jannik Kossen, Sebastian Farquhar, Yarin Gal, and Tom Rainforth. Active testing: Sample- efficient model evaluation. arXiv preprint (03.::2033.0333) 2021.
[28] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, 30, 2017.
[29] Yann LeCun. The MNIST database of handwritten digits. http://yann. lecun. com/exdb/mnistl/ 1998.
[30] Richard Liaw, Eric Liang, Robert Nishihara, Philipp Moritz, Joseph E Gonzalez, and Ion Stoica. Tune: A research platform for distributed model selection and training. arXiv preprint 187 2018.
[31] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
[32] Philipp Moritz, Robert Nishihara. Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, and Ion Stoica. Ray: A distributed framework for emerging ai applications, 2018.
[33] Jersey Neyman. Sur les applications de la théorie des probabilités aux experiences agricoles: Essai des principes. Roczniki Nauk Rolniczych, 10:1-51, 1923.
[34] Safiya Umoja Noble. Algorithms of oppression: How search engines reinforce racism. NYU Press, 2018.
[35] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Bue, E. Fox, and Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8035. Curran Associates, Inc., 2019.
12

---

[36] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and F. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
[37] Caroline Criado Perez. Invisible women: Exposing data bias in a world designed for men. Random House, 2019.
[38] Maya L Petersen, Kristin E Porter, Susan Gruber, Yue Wang, and Mark J Van Der Laan. Diagnosing and responding to violations in the positivity assumption. Statistical methods in medical research, 21(1):31-54, 2012.
[39] Tian Qin, Tian-Zuo Wang, and Zhi-Hua Zhou. Budgeted heterogeneous treatment effect estimation. In International Conference on Machine Learning, pages
PMLR, 2021.
[40] James M Robins, Miguel Angel Hernán, and Babette Brumback. Marginal structural models and causal inference epidemiology. Epidemiology, 11(5):551, 2000.
[41] Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of educational Psychology, 66(5):688. 1974.
[42] Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pages 3076-3085. PMLR, 2017.
[43] Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in Neural Information Processing Systems, 32:2507 2517, 2019.
[44] Cathie Sudlow, John Gallacher, Naomi Allen, Valerie Beral, Paul Burton, John Danesh, Paul Downey, Paul Elliott, Jane Green, Martin Landray, et al. Uk biobank: an open access resource for identifying the causes of wide range of complex diseases of middle and old age. Plos med, e 2015.
[45] liris Sundin, Peter Schulam, Eero Siivola, Aki Vehtari, Suchi Saria, and Samuel Kaski. Active learning for decision-making from imbalanced observational data. In International Conference on Machine Learning, pages 6046-6053 PMLR, 2019.
[46] Lu Tian, Ash A Alizadeh, Andrew J Gentles, and Robert Tibshirani. A simple method for estimating interactions between treatment and large number of covariates. Journal of the American Statistical Association, (095508):11171111 2014.
[47] Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal. Improving deterministic uncertainty estimation in deep learning for classification and regression. arXiv preprint arXiv:2 1409, 2021.
[48] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep kernel learning. In Artificial intelligence and statistics, pages 370-378. PMLR, 2016.
[49] Yu Xie, Jennie E Brand, and Ben Jann. Estimating heterogeneous treatment effects with observational data. Sociological methodology, 42(1):314-347, 2012.
[50] Steve Yadlowsky, Hongseok Namkoong, Sanjay Basu, John Duchi, and Lu Tian. Bounds on the conditional and average treatment effect with unobserved confounding factors. arXiv preprint arXiv:1808.0m 09521, 2018
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Section
(c) Did you discuss any potential negative societal impacts of your work? [Yes] In Section
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions all theoretical results? [Yes] In Section 1
(b) Did you include complete proofs of all theoretical results? [Yes] In Section 3 and Appendix A.2, A.3
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] supplemental material
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] In Appendix sections C and F
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, cloud provider)? [Yes] In Appendix section E
4. If you using existing assets (e.g., code, data, models) or curating/releasing
assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]] [22]
(b) Did you mention the license of the assets? [No] All datasets and assets are open source and their license retrievable from the citation.
(c) Did you include any new assets either the supplemental material or as URL? [Yes] Described in Appendix C and code generate is provided
(d) Did you discuss whether and how consent was obtained from people whose data you't 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] IHDP section
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "pGxKopF3hrJ", "reviewer": "Reviewer_EmX1", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper was well-written and clear in its exposition.  The problem that the authors aim to tackle is interesting and important.  In many scenarios where you wish to estimate a treatment effect, you are also limited in the amount of data you can acquire - hence, applying Bayesian active learning here seems very natural. Additionally, as far as I know, there is not a lot of literature on using active learning when estimating treatment effects (aside from the two papers that the authors cite - namely Deng et al. (2011) and Sundin (2019)).  The authors' approach seems reasonable, and I liked the fact that the authors evaluated the performance of their method on not just 1, but 3 datasets.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The comments that I have relating to this paper are mostly minor:\n- It would have been nice to have had a fuller description of the CMNIST dataset in text.  Could you even have shown a few examples of the labels, y, generated for a few example digits?  It was hard, if not impossible, to conceive of the elements of CMNIST from the description given in the main text.  The math in the supplement helped, but I am still having trouble visualizing exactly what some elements of this dataset look like.\n- Do the authors have an explanation for why particular acquisition functions work better for some datasets than others?  E.g. it seems that $\\mu\\pi$-BALD works well on the synthetic dataset, but not so well on CMNIST\n- It would also have been nice to have had (even in the supplement) a histogram of true treatment effects associated with each of the datasets used for evaluation.  I am wondering if your method works equally well for estimating all sizes of treatment effects, or if it is better at estimating larger treatment effects \n- Finally, it would have been nice to have seen some analysis of the examples selected by the authors' best performing acquisition functions. \n This would have been particularly interesting in e.g. the CMNIST case, where one acquisition function vastly outperforms the other.  Which examples are being selected in that case, compared to for other acquisition functions?  ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "LqwDLQR1Eak", "reviewer": "Reviewer_Tcca", "summary": "The paper proposes several acquisition functions for performing active learning in a causal setting. The paper explores the limitations of more naive acquisition functions and motivates its own proposals. The performance of the proposed acquisition functions is empirically compared to that of more naive acquisition functions. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**: The authors do acknowledge other similar methods; however, I think they fairly outline the differences between them and why they matter. Furthermore, in comparison to other acquisition functions, I think a good case is made for why the proposed functions are far superior. Overall, not wildly original, but original in ways that I was convinced do matter. \n\n**Quality**: I found the work to be technically sound, with careful consideration of the problems that can arise with naive acquisition functions and a logical progression towards more intelligently constructed ones. I found the experimental section convincing. I would have appreciated some comment on the limitations of the method / avenues for future work. \n\n**Clarity**: I found the work to be clearly written and organized. Figures 2 and 3 were very effective.\n\n**Significance**: I believe the work is significant. Active learning is an important field of research, with undeniable potential for social impact. The proposed acquisition functions greatly outperform naive ones and the overall approach does differentiate itself in important ways from competing methods. The proposed acquisition functions can also be used in tandem with other BALD variants that may be developed in the future, furthering their applicability. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "EI_5UydT2z", "reviewer": "Reviewer_8K1A", "summary": "Update (12th Aug 2021) - increased score to reflect that my comments were satisfactorily addressed by the authors in their rebuttal.\n\n-----------\nThe authors propose new set of acquisition functions, dubbed Causal-BALD, based on Bayesian Active Learning by Disagreement (BALD) that consider overlap and information gain as a decision criterion for iterative improvement of a counterfactual estimator in settings in which direct treatment assignment by an algorithm is unethical or not feasible. The authors performed a clean evaluation in various settings and with multiple underlying models to show that their proposed approach outperforms the previous state-of-the-art.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**:\n- The authors study the application of active learning methods for interactively querying a data set for the improvement of an underlying counterfactual estimator (deep ensembles and DUE). To the best of my knowledge, this is a non-trivial, relevant and understudied setting, and originality is therefore clearly a strength of the presented work.\n\n**Quality**:\n- The performed experiments and their descriptions are comprehensive. The authors evaluate multiple relevant settings, underlying models and benchmark datasets to substantiate their claims.\n\n- The benefits of Causal-BALD over treatment-naive BALD (eq.s 1-2) are discussed at length in the manuscript. However, a treatment-naive BALD baseline is missing from the experimental comparison.\n\n**Clarity**:\n- The paper is well written and the text is easy to follow for a reader versed in the field of counterfactual inference. The exposition is clear and the paper follows a compelling narrative.\n\n- How experimental hyper-parameters were systematically chosen for the models is not described in the manuscript. Only the final hyper-parameters used in the experiments are presented in the supplementary materials. On review of the source code, I saw that a systematic approach was likely performed (see tuning.py) - this information should be included in the manuscript.\n\n- In terms of reproducibility, I particularly appreciate the comprehensive appendix that covers the details missing in the main body of the paper and the clear reference implementation provided by the authors.\n\n**Significance**:\n- Active learning for estimating individual treatment effects is an understudied, but relevant, problem with several potential practical applications. The proposed method seems to be generally applicable under the stated assumptions, and the experimental results support its use.\n\n- The presented acquisition functions, crucially, rely on the availability of meaningful uncertainty estimates. An additional source of uncertainty in the counterfactual estimation setting is the potential presence of unobserved confounding that can not be ruled out in practical scenarios and will not be reflected in the models' uncertainty estimates. Claims towards \"unbiasedness\" should therefore be toned down to reflect that we can not test the unconfoundedness assumption, and the proposed method therefore will not have guarantees to be unbiased in  real-world applications.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "\n**Societal Impact**: \n- As a general-purpose method for active learning in the counterfactual estimation setting with potentially broad applicability, I see no immediate reason for additional concern beyond that inherent in any active learning method. One risk is that the method could, e.g., lead to a biased, non-representative sampling in terms of ethnicity and other protected attributes - particularly, if the unconfoundedness assumption this work is based on is blindly trusted. The authors could briefly outline this risk and how one could preempt such a scenario with appropriate checks and balances.\n\n**Questions/suggestions for the authors**:\n- Please describe how the experimental hyper-parameters presented in the supplementary materials were chosen. In general, conducting experiments over multiple hyper-parameter sets would be preferable to ensure the results generalize to other choices of hyper-parameters for the underlying models.\n\n- From the appendix, it appears the authors used a variant of TARNet [1] as constituents of the deep ensemble models (at least in some experiments). However, as far as I can see, [1] is not disclosed/referenced anywhere in the paper, nor is there a discussion of deep counterfactual representation learning methods in the paper. Please add the missing information on what models were the constituents of the deep ensembles and any related discussion to the related work section.\n\n- Please include a short discussion on how the [standard in the field but relatively strict] assumptions (particularly the unconfoundedness assumption) and its potential impact - if violated - on the uncertainty estimates may affect your proposed acquisition function in practice so that readers are aware of this inherent risk.\n\n- Please add treatment-naive BALD (eq.s 1-2) to the experimental comparison to highlight and quantify the benefit of adapting the acquisition function to the counterfactual setting.\n\n**Conclusion**: A relevant and understudied setting, clear and compelling narrative and a comprehensive experimental evaluation. There are minor issues that I trust can be resolved during review. This paper ticks all the boxes for a clear accept.\n\nNote: I have not verified the proofs in the supplementary materials.\n\n[1] Shalit, Uri, Fredrik D. Johansson, and David Sontag. \"Estimating individual treatment effect: generalization bounds and algorithms.\" International Conference on Machine Learning. PMLR, 2017.\n", "review_text": "", "extra_scores": {"ethical_concerns": "I have no ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "74Uc04MoRUn", "reviewer": "Reviewer_wzW6", "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.\n\nThe proposed function is built on the existing Bayesian active learning by disagreement (BALD). \n\nFirst, \\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. \n\nSecond, 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. \n\nThird, \\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.\n\nExperimental 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well motivated. The problem setting of treatment effect estimation requires an adapted version of active learning algorithms that can focus on acquire samples from the overlap region. This work started from this intuition and proposed new acquisition function that can achieve both overlap and uncertainty reduction. The experimental results show the performance gain of the new acquisition function compared to existing baseline approaches.\n\nI don't have major complains on the method of this work. A few minor comments are listed below.\n\n1) A few variations of BALD are listed in this work. It could help to highlight the major contribution of this work in the abstract or introduction so that readers can clearly grasp this work's novelty.\n\n2) In Table 2, are these results based on within-sample evaluation or out-of-sample evaluation?\n\n3) Did the authors try to run experiments on other commonly used benchmark datasets for treatment effect estimation, such as Twins and Jobs?\n\n4) The font size in the legend in Figure 2 should be increased to become more readable.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "1.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to provide a sample efficient strategy to train deep learning models that can effectively predict individual-level treatment effects from high dimensional observation data. A unique challenge in estimating personalized treatment effect is the necessity to choose data samples from regions with support from both treated and control populations so that the treatment effect can be properly identified. As a result, most existing data acquisition functions that primarily focus on uncertainty become less effective.  The proposed model, referred to as Causal-BALD, addresses this challenge by integrating Bayesian active learning and causal inference in novel ways that cover the uncertainty and overlapping criteria simultaneously. Effectiveness of the proposed acquisition functions is clearly demonstrated using both synthetic and semi-synthetic datasets. 

As acknowledged by all the reviewers, the proposed acquisition functions are clearly motivated, well explained, sufficiently novel, and further supported by convincing evaluation results. It has the potential to impact both the scientific community and multiple application domains. 
The authors adequately addressed most of the questions raised by reviewers in their original reviews during the discussion period. In particular, since multiple acquisition functions are proposed and some show better performance than others on certain datasets, it is important for the authors to offer some deeper insight to better understand the behaviors of different acquisition functions. It is also interesting to show some sampled data instances and connect them with the corresponding acquisition function. The authors promised to move some results to the appendix and add necessary discussions to better highlight the key properties of the acquisition functions and their potential limitations. 

Decision: Accept (Poster)