Active Learning for Improving Decision-Making from Imbalanced Data

Iiris Sundin 1 Peter Schulam * 2 Eero Siivola * 1 Aki Vehtari 1 Suchi Saria 2 Samuel Kaski 1

Abstract
This extended abstract considers the reliabilityof prediction-based decision-making in a task ofdeciding which action a to take for a target unit
after observing its covariates x̃ and predicted outcomes p̂(ỹ | x̃, a). An example case is personalized medicine and the decision of which treatment
to give to a patient. A common problem whenlearning these models from observational data isimbalance, that is, difference in treated/control
covariate distributions. We propose to assessthe decision-making reliability by estimating themodel’s Type S error rate, which is the probabilityof the model inferring the sign of the treatment
effect wrong. Furthermore, we use the estimatedreliability as a criterion for active learning, inorder to collect new (possibly expensive) observations, instead of making a forced choice basedon unreliable predictions. We demonstrate the effectiveness of this decision-making aware active
learning in two decision-making tasks: in simulated data with binary outcomes and in a medical
dataset with synthetic and continuous treatment
outcomes.

1. Introduction
In this extended abstract, we consider a case where machinelearning is used to provide predictions of outcomes underalternative actions, to aid human decision-making (Schulam& Saria, 2017). To fit a model to this task, we need data
recording previous actions a, observed outcomes y, and any
features relevant to the context of the decision, x. Then, thegoal is to estimate p(Y | X = x, A = a) and, further, individual treatment effect (ITE), τ (x) = E[Y [1] − Y [0] | X =x], where Y [a] denotes the potential outcome of treatment
A = a (Rubin, 1978). ITE provides sufficient information
to choose between two actions.
*

Equal contribution 1 Department of Computer Science, Aalto
University, Espoo, Finland 2 Department of Computer Science,Johns Hopkins University, Baltimore, MD 21218, USA. Correspondence to: Iiris Sundin <iiris.sundin@aalto.fi>.
Real-world Sequential Decision Making workshop at ICML 2019.
Copyright 2019 by the author(s).

yResponse observations
from treatments •, and control N
Estimated treatment responses

p0
x

x̃
True treatment responses

y

x̃

x

Figure 1. An example decision-making task is to choose a treatment for a specific x̃ (red mark). Upper graph shows the posteriormeans of the potential outcome models given observations, andthe prior mean p0 . The true responses are in the lower graph,
which shows that there are three regions with different responsetypes (marked in green, blue and magenta). Lower y is better. The
problem is that there are no observations about treatment • in thegreen region, which causes it to appear to be the best choice for x̃,although that is incorrect.

The estimation of ITE is susceptible to many error
sources (Pearl, 2009; Schulam & Saria, 2017; Mitchell et al.,2018), of which we concentrate on imbalance (Gelman &Hill, 2007).
Imbalance is defined as the difference in covariate distributions in the treated and control groups. Imbalance
makes correct model specification essential for avoiding
bias in treatment effect estimates (Gelman & Hill, 2007).
Mis-specification of the model could be avoided by using
non-parametric models, but their variance increases quickly
when extrapolating. Recently, imbalance has been shownto increase the upper bound of the model error in estimation of ITE (Shalit et al., 2017). Furthermore, imbalance
becomes the more prevalent issue the higher-dimensional

Active Learning for Improving Decision-Making

the covariate space is (D’Amour et al., 2018).

2. Related Work

There are many existing ways to deal with imbalance whenlearning the average treatment effect (ATE). In causal inference, the most common methods are propensity scorematching or weighting (Rosenbaum & Rubin, 1984; Hirano
et al., 2003; Lunceford & Davidian, 2004), and modelingthe potential outcomes (Imbens & Rubin, 2015; Hernán &Robins, 2018), as well as doubly robust methods which
implement both (Bang & Robins, 2005; Funk et al., 2011).

Active learning proposed by Javdani et al. (2014) aims at
improving automated decision-making by reducing model
uncertainty so that the remaining hypotheses are confined tothe same decision region. The importance of active learningfor decision-making tasks has been noted in other fields,
where e.g. Saar-Tsechansky & Provost (2007) developed aheuristic method for deciding which consumers to target inmarketing campaigns. Closest to our work is the work on
active learning with logged data (Yan et al., 2018), which
proposes a de-biasing query strategy for a classification
task. The difference to our work is that they assume thepropensities (probability of revealing the label) to be known.

Even though these methods can decrease bias in treatment
effect estimates, they will increase variance, and therefore
may make the decision-making less reliable. This is especially the case with ITE; For example, in the areas ofcovariate space where there are more control units thantreated, intuition is that the model for the treated outcome either has to generalize from less-representative observations(increasing bias) or extrapolate (increasing variance). Eitherway, there is higher uncertainty about the treated outcome,which makes reliable decision-making difficult. An extreme
case of this is illustrated in Fig. 1. A natural question thenis, could other data sources be exploited instead of making
a forced choice based on insufficient observational data.
This work has three main contributions, which are complementary to each other and can be used independently. First,
we describe how imbalance decreases decision-making performance by increasing Type S error rate (Gelman & Tuerlinckx, 2000), which is the probability of the model inferring
the sign of the treatment effect wrong. Second, we proposea Bayesian estimate for the Type S error rate, which allows quantifying the reliability of a decision-support model.
Third, we propose to alleviate the consequences of imbalance by actively collecting more data (e.g. from a simulator,
new experiment, or by asking a domain expert). To this end,
we introduce decision-making aware active learning criteriathat improve decision-making performance by minimizing
the estimated Type S error rate.
Finally, in many cases there are restrictions on what canbe measured. For example, in medicine it is in general not
ethical to do an experiment on a patient in order to get information to improve the treatment plan of another patient.
Therefore, regular active learning would not be possible and,instead, any new information has to be acquired indirectly.
For this reason, we introduce the idea of counterfactual elicitation which means soliciting indirect observations aboutcounterfactual outcomes, that is, what would have been
the outcome had x been treated with a0 instead of a. We
demonstrate the effectiveness of the proposed active learning criteria applied to counterfactual elicitation.

Using data to estimate the effect of interventions has been
extensively studied in the field of causal analysis (see e.g.
Pearl (2009); Morgan & Winship (2014); Imbens & Rubin(2015); Hernán & Robins (2018)). The issue of imbalanceis discussed by Johansson et al. (2016), and they propose anapproach based on empirical risk minimization and domain
shift to improve predictions. Xu et al. (2016) estimate theindividual treatment effects in a longitudinal setting whereindividual-specific treatment parameters are refined overtime as more observations are collected. Alaa & van der
Schaar (2017) use Gaussian processes to model individualspecific outcomes under alternative treatments and prove
minimax rates on the risk that the approach achieves. Our
work builds on these ideas, and is most closely related to
the works by Xu et al. (2016) and Alaa & van der Schaar(2017). Although we also use Gaussian processes to model
treatment effects, our work is unique in that it leverages
the probabilistic framework to design an active learningalgorithm to improve the decision-making process.

3. Methods
Preliminaries. Causal effect of a treatment is the difference between outcomes when a unit x ∈ X is treated andnot treated, where X is the population. Individualized treatment effect is defined as τ (x) = E[Y [1] − Y [0] | X = x],where a = 1 means treated and a = 0 not treated, i.e.
control. The estimated treatment effect is τ̂ (x).
Setup. The outcome model has been learned in retrospective from observational data, which may be imbalanced in
the observed actions. The observational data D is a setof n observations {yi , ai , xi }ni=1 , where yi is the observed
outcome of action ai for unit xi . Imbalance means that thecovariate distributions are different in the treated (ai = 1)
and control groups (ai = 0).
We assume that the unknown policy used to choose actionsin the training data only depends on the observed covariatesx ∈ X . This is equivalent to the no unmeasured confounders

Active Learning for Improving Decision-Making

assumption (Hernán & Robins, 2018) and implies that all
confounders are included in x. We further assume consistency of potential outcomes, which means that the potential
outcomes p(y[a] | x) = p(y | X = x, A = a) can be directly estimated from the observed outcomes in the training
data. Regardless of this, imbalance will still cause issues.
In this work, the task is to decide which action a to choose
for x̃. Decision-making performance is measured as theprobability of correct decision, or equivalently, the proportion of correct decisions in repeated decision-making tasks.
Definition: Type S error rate γ is the probability of themodel inferring the sign of the treatment effect wrong;γ: M × PX ,Y → [0, 1] where M is the model space and
PX ,Y is the true treatment effects (distributions over X ×Y).
The expected Type S error rate in X isγ = EPX ,Y [I(sign(τ̂ ) 6= sign(τ ))], where I(A) = 1 if
condition A is true, and 0 otherwise.
The expected proportion of correct decisions in population
X is then 1 − γ, which makes Type S error rate a natural
measure of the decision-making performance.
We assume that the decision-maker’s policy is to choose
the action with the highest expected utility of the outcome.
Without loss of generality we will assume that the utility is
directly the outcome y (higher better).
3.1. Effect of Imbalance on Type S Error Rate
In this section we give intuition on how imbalance increases
the error rate in decision-making. First, assume a probabilistic model of potential outcomes, with broad prior distributions. This implies that when the sample size is small,posteriors will be wide. Then, under certain assumptions, it
can be shown that imbalance decreases the expected number of samples locally, which makes posteriors wider andtherefore increases the Type S error rate locally.
Then the question remains whether the error rate increases
globally as well, or do the local effects cancel out eachother. The intuition is that since the error rate increases
exponentially with decreasing number of samples, if there
is high-enough density in areas with high imbalance, thenthe local increase in the error rate cannot be compensatedelsewhere.
3.2. Estimated Type S Error Rate
In the Bayesian sense, a model p̂(y[a] | x, D) captures ourcurrent understanding of the problem, and therefore theestimated Type S error rate is p̂(y[1] < y[0] | x, D) if the
expected effect is positive, that is, if Ep̂(y[1]|x,D) [y[1]] >Ep̂(y[0]|x,D) [y[0]]. Respectively, if the expected effect isnegative, then the estimated error rate is p̂(y[1] > y[0] |x, D).

For example, when the predicted treatment effect τ = y[1]−y[0] is normally distributed, then it is easy to show thatthe estimated Type S error rate in a test unit x̃ is the tailprobability|Ep̂(τ̃ |x̃,D) [τ̃ ]|−1γ̂(x̃) = probit
−
,
(1)
1
Var(p̂(τ̃ | x̃, D)) 2
where probit−1 is the cumulative distribution function of
normal distribution, and the expectations and variances areover the posterior predictive distribution of τ̃ | x̃. From (1)we see that the estimated Type S error rate will increase if
the estimated treatment effect decreases, or if posterior uncertainty (variance) increases. Intuitively this makes sense.
3.3. Decision-Making Aware Active Learning
Our hypothesis is that active learning criteria that reduce theestimated Type S error rate will result in higher decisionmaking performance. We call active learning criteria that
reduce the estimated Type S error decision-making aware.
We propose to minimize the estimated Type S error rate ineq. (1) with active learning. The criterion is to maximize
the estimated reliability of a decision at x̃, that is, 1 − γ̂(x̃).
Directly minimizing the error can be interpreted as only exploiting what we already know, which lacks in exploration.
There is an easy fix, though, which is to add explorationon the error. The exploration-exploitation trade-off is managed by maximizing the expected information gain on thepredictive distribution of Type S error: Bernoulli(γ̂(x̃))(technically, its relative entropy). The maximization of theinformation gain is equivalent to minimizing the posterior
entropy and consequently the log-loss (Settles, 2012).
3.4. Counterfactual Elicitation
Assuming it is possible to acquire (noisy) observations aboutcounterfactuals in the training data, we can do more, as discussed in the introduction. Denote by D the set of training
examples {xi , ai , yi }ni=1 for which the factual outcomes yi
have been observed, and denote by U the counterfactual
examples {xi , 1 − ai }ni=1 , for which the outcomes are unknown. Then we can use active learning to construct a set
L that contains the new observations, and L will be data
which would not normally be available. At each iteration,
the algorithm selects {x∗ , a∗ } ∈ U to solicit a counterfactual outcome y ∗ . After this, {x∗ , a∗ , y ∗ } is added to L, andremoved from U. So the optimization problem at each query
iteration k becomesx∗ , a∗ = arg min Ep̂(y|x,a,D,L) [γ̂k+1 (x̃)] , where{x,a}∈Uγ̂k+1 (x̃) = p̂ (y[ax̃ ] < y[1 − ax̃ ] | x̃, D, L, {x, a, y}) ,and ax̃ is the treatment with the highest expected outcome
for x̃: ax̃ = arg maxa0 Ep̂(y[a0 ]|x̃,D,L,{x,a,y}) [y[a0 ]].

Active Learning for Improving Decision-Making

4. Experiments and Results

GPy1 (version 1.9.2), and use mixed noise likelihood tolearn the noise in the observations acquired by active learning. We use an exponentiated quadratic kernel with a separate length-scale parameter for each variable, and optimize
the hyperparameters using marginal likelihood. We use
Gauss-Hermite quadrature of order 32 to approximate theexpectations in D-M aware, Targeted-IG, and EIG.

We evaluate the performance of decision-making aware
active learning (D-M Aware) using counterfactual elicitation in two cases. We compare the performance of DM aware to two widely-used earlier active learning approaches, uncertainty sampling and maximum expected
information gain (Culotta & McCallum, 2005; Roy & McCallum, 2001) (EIG), and also include results for a previously introduced special case of decision-making aware
active learning, targeted-IG (Sundin et al., 2018). In orderto make the methods comparable, we use the variant of D-M
aware that also explores (see Section 3.3), as do the EIG andtargeted-IG. The performance of the non-exploring variantof D-M aware is slightly lower (not shown) in all cases.

The results in the IHDP data are similar as in the simulated
data; The decision-making performance improves fastest
with D-M aware and Targeted-IG, compared to EIG and uncertainty sampling (see Fig. 2(b)). D-M aware and TargetedIG achieve statistically significant improvement in decisionmaking performance already with one query (based on 95%
bootstrap confidence intervals).

4.1. Simulated Data with Binary Outcomes

4.3. Comparative Feedback

In this experiment, we study the proposed active learning approach in simulated data, where binary outcome y indicatesthe occurrence of an adverse event, and the decision-making
task is therefore to choose the treatment that results in alower probability of the adverse outcome. The setting issimilar to that in the Fig. 1.

Last, we demonstrate the use of comparative observations
for counterfactual elicitation, that is, indication of which
potential outcome is higher. The motivation is that thistype of feedback may be more natural for human experts
to give. The setting is the same as in Section 4.1, withthe difference that the query is about which treatment has
lower Bernoulli parameter. We fit the outcomes with GPs
using Stan (Stan Development Team, 2017; Carpenter et al.,2017), which allows the model to learn both from direct
and comparative observations. The results in Fig. 3 showthat comparative observations increase decision-making performance efficiently in the simulated data setting. We notethat the results with comparative feedback are better thanthose in Fig. 2(a), because here the queries give information
about θx,1 > θx,0 , thus providing more information than the
direct observations on outcome y[a] | x ∼ Bernoulli(θx,a ).

We model the probability of an adverse outcome given covariates x ∈ R and action a ∈ {0, 1} using logistic regression with interactions and radial basis functions for
local effects: p(y | x, a) ∼ Bernoulli(θx,a ).The model is
fit using a probabilistic programming language Stan (Stan
Development Team, 2017; Carpenter et al., 2017).
Fig. 2(a) shows that the decision-making aware active learning criteria, D-M aware and Targeted-IG, are the fastest to
improve decision-making performance, and achieve significant increase in correct decisions after just one query in
two cases out of three. After five solicited counterfactuals,
Expected Information Gain (EIG) has reached comparable
but still lower performance.
4.2. Medical Semi-Synthetic Data with Continuous
Outcomes
In the second experiment, the decision-making task is deciding on medical interventions in real medical data with
continuous-valued synthetic outcomes.
We use the Infant Health and Development Program (IHDP)
dataset from Hill (2011), also used e.g. by Shalit et al. (2017)
and Alaa & van der Schaar (2017), including synthetic outcomes, containing 747 observations of 25 features. We
evaluate the performance in leave-one-out cross-validation,
but in order to make the problem even more realisticallyhard, for each of the 747 target units we choose randomly
100 observations as training examples.
We fit separate GPs to the outcomes of each treatment with

5. Discussion and Conclusions
As machine learning systems are being integrated into human decision-making workflows, it is increasingly important that deployed models are correct and reliable. In
this paper we focused on the effect that imbalance in the
training data can have on the reliability of comparisons ofp̂(y[a] | x). We propose to improve the reliability by active
learning that aims at maximizing the estimated reliability. In
our experiments, this decision-making aware active learning
outperforms standard methods in decreasing the error rate
in decision-making.
One interesting application of the method is in personalized
medicine, where counterfactual elicitation could solicit practitioners’ knowledge to the model. For example, a clinician
in a hospital has access to medical records of previous patients, and may also have personal experience about some1

Toolbox available at: https://sheffieldml.github.io/GPy

Active Learning for Improving Decision-Making

(a)

(b)

Figure 2. Comparison of active learning criteria as a function of number of queries. (a) Simulated data (b) IHDP data. The topmost
panel in each shows imbalance, middle panel the estimated Type S error rate, and the lowest panel the proportion of correct decisions.
Information-gain-based approaches are more effective than uncertainty sampling, and the decision making-aware criteria D-M aware andTargeted-IG are the best. Shaded areas show the 95% bootstrap confidence interval.

References
Alaa, A. M. and van der Schaar, M. Bayesian inference ofindividualized treatment effects using multi-task Gaussian Processes. In Guyon, I., Luxburg, U. V., Bengio, S.,
Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R.
(eds.), Advances in Neural Information Processing Systems 30, pp. 3424–3432. Curran Associates, Inc., 2017.
Figure 3. Comparative observations on the Bernoulli parameterin the simulated data is effective in increasing the proportion ofcorrect decisions. The results are averaged over 3 target units ineach response type regions and 100 repetitions.

Bang, H. and Robins, J. M. Doubly robust estimation in
missing data and causal inference models. Biometrics, 61(4):962–973, 2005.
Carpenter, B., Gelman, A., Hoffman, M., Lee, D., Goodrich,
B., Betancourt, M., Brubaker, M., Guo, J., Li, P., andRiddell, A. Stan: A probabilistic programming language.
Journal of Statistical Software, 76(1):1–32, 2017. URL
https://www.jstatsoft.org/v076/i01.

of the previous patients. These data are rarely included intraining sets of the models, but active learning and counterfactual elicitation could allow leveraging this additionalsource of information to infer more accurately about the
future.

Acknowledgements
This work was supported by the Academy of Finland (grants
298742, 313122, 319264 and 294238), by Alfred KordelinFoundation and by Jenny and Antti Wihuri Foundation. We
acknowledge the computational resources provided by Aalto
Science-IT project, and the support from the Finnish Center
for Artificial Intelligence (FCAI).

Culotta, A. and McCallum, A. Reducing labeling effort
for structured prediction tasks. In AAAI, volume 5, pp.
746–751, 2005.
D’Amour, A., Ding, P., Feller, A., Lei, L., and Sekhon, J.
Overlap in observational studies with high-dimensional
covariates. arXiv preprint arXiv:1711.02582, 2018.
Funk, M. J., Westreich, D., Wiesen, C., Strmer, T.,
Brookhart, M. A., and Davidian, M. Doubly robust estimation of causal effects. American Journal of Epidemiology, 173(7):761–767, 2011.
Gelman, A. and Hill, J. Data analysis using regression andmultilevel/hierarchical models. Cambridge university
press, 2007.

Active Learning for Improving Decision-Making

Gelman, A. and Tuerlinckx, F. Type S error rates for classical and Bayesian single and multiple comparison procedures. Computational Statistics, 15(3):373–390, 2000.

Schulam, P. and Saria, S. Reliable decision support using
counterfactual models. In Advances in Neural Information Processing Systems, pp. 1697–1708, 2017.

Hernán, M. and Robins, J. Causal Inference. Chapman &Hall, 2018.

Settles, B. Active learning: Synthesis lectures on artificial intelligence and machine learning. Long Island, NY:
Morgan & Clay Pool, 2012.

Hill, J. L. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics,
20(1):217–240, 2011.
Hirano, K., Imbens, G. W., and Ridder, G. Efficient estimation of average treatment effects using the estimatedpropensity score. Econometrica, 71(4):1161–1189, 2003.
Imbens, G. and Rubin, D. Causal inference in statistics,
social, and biomedical sciences. Cambridge University
Press, 2015.
Javdani, S., Chen, Y., Karbasi, A., Krause, A., Bagnell,D., and Srinivasa, S. S. Near optimal Bayesian active
learning for decision making. In AISTATS, volume 2, pp.
17, 2014.
Johansson, F., Shalit, U., and Sontag, D. Learning representations for counterfactual inference. In International
Conference on Machine Learning (ICML), 2016.
Lunceford, J. K. and Davidian, M. Stratification and weighting via the propensity score in estimation of causal treatment effects: a comparative study. Statistics in medicine,23(19):2937–2960, 2004.
Mitchell, S., Potash, E., and Barocas, S. Prediction-based
decisions and fairness: A catalogue of choices, assumptions, and definitions. arXiv preprint arXiv:1811.07867,
2018.
Morgan, S. and Winship, C. Counterfactuals and causal
inference. Cambridge University Press, 2014.
Pearl, J. Causality: models, reasoning and inference. Cambridge University Press, 2009.
Rosenbaum, P. R. and Rubin, D. B. Reducing bias in observational studies using subclassification on the propensity
score. Journal of the American statistical Association, 79(387):516–524, 1984.
Roy, N. and McCallum, A. Toward optimal active learning
through Monte Carlo estimation of error reduction. In
International Conference on Machine Learning (ICML),
pp. 441–448, 2001.
Rubin, D. Bayesian inference for causal effects: The role ofrandomization. The Annals of statistics, pp. 34–58, 1978.
Saar-Tsechansky, M. and Provost, F. Decision-centric active
learning of binary-outcome models. Information Systems
Research, 18(1):4–22, 2007.

Shalit, U., Johansson, F. D., and Sontag, D. Estimating
individual treatment effect: generalization bounds andalgorithms. In International Conference on Machine
Learning, pp. 3076–3085, 2017.
Stan Development Team. Pystan: the python interface tostan, version 2.16.0.0, 2017. URL http://mc-stan.
org.
Sundin, I., Peltola, T., Micallef, L., Afrabandpey, H., Soare,M., Mamun Majumder, M., Daee, P., He, C., Serim, B.,
Havulinna, A., Heckman, C., Jacucci, G., Marttinen, P.,
and Kaski, S. Improving genomics-based predictions for
precision medicine through active elicitation of expert
knowledge. Bioinformatics, 34(13):i395–i403, 2018.
Xu, Y., Xu, Y., and Saria, S. A Bayesian nonparametric approach for estimating individualized treatment-response
curves. In Machine Learning for Healthcare Conference(MLHC), pp. 282–300, 2016.
Yan, S., Chaudhuri, K., and Javidi, T. Active learning withlogged data. In International Conference on Machine
Learning, pp. 5517–5526, 2018.