Bayesian EDDI: Sequential Variable Selection with Bayesian Partial VAE

Chao Ma 1 * Wenbo Gong 1 * Sebastian Tschiatschek 2 Sebastian Tschiatschek 1 2 Sebastian Nowozin 2 3
José Miguel Hernández-Lobato 1 2 Cheng Zhang 2* Equal contribution

Abstract
ST: 1
Obtaining more relevant information enables better decision making, but may be costly. Optimal
sequential decision making allows us to trade
off the desire to make good decisions by acquiring further information with the cost of performing that acquisition. To this end, we propose aprincipled framework, named EDDI (Efficient
Dynamic Discovery of high-value Information),based on the theory of Bayesian experimental
design. In EDDI we propose a novel partial variational autoencoder (Partial VAE), to efficiently
handle missing data with different missing patterns. Additionally, we extend the VAE based
framework to Bayesian treatment of the weights,
which obtains better performance in small data
regime. EDDI then combines it with an acquisition function that maximizes expected information gain on a set of target variables at each step.
1

Introduction

Imagine that a person walks into a hospital with a brokenarm. The first question from health-care personnel would be:“How did you break the arm?” instead of “Do you have acold?”, because the answer reveals relevant information forthis patient. Human experts dynamically acquire information based on the current understanding of the situation in asequential manner . Automating this human expertise of asking relevant questions is difficult. In other applications suchas online questionnaires, for example, most existing online
questionnaire system either present exhaustive questions(Lewenberg et al., 2017; Shim et al., 2018) or use extremely
1 Department of Engineering, University of Cambridge, Cambridge, UK 2 Microsoft Research Cambridge, Cambridge, UK
3 now at Google AI, Berlin, Germany (contributed to this paper during his stay at Microsoft Research Cambridge). Correspondence
to: Cheng Zhang <Cheng.Zhang@microsoft.com>.

Real-world Sequential Decision Making workshop at ICML 2019.
Copyright 2019 by the author(s).

time-consuming human labeling work to manually builda decision tree for a reduced number of questions (Zakim
et al., 2008). This wastes the valuable time of experts orusers (patients). An automated solution for personalized
dynamic acquisition of information has great potential tosave much of this time in many real-life applications.
What are the technical challenges to building a sequential
intelligent information acquisition system? Missing data isa key issue: taking the questionnaire scenario as an example,
there are two types of missing data problems. Firstly, foreach subject, at any point in time we only observe a small
subset of answers yet have to reason about possible answers
for the remaining questions. We thus need an accurate probabilistic model that can perform inference given a variable
subset of observed answers. Secondly, in many real-lifescenarios, the number of initially available subjects is often limited. Can the system perform robustly under small
data? Another key problem is deciding what to ask next: thisrequires assessing the value of each possible question ormeasurement, the exact computation of which is intractable.
However, compared to current active learning methods weselect individual features, not instances; therefore, existing
methods are not applicable. In addition, these traditional
methods are often not scalable to the large volume of data
available in many practical cases (Settles, 2012).
We propose the EDDI (Efficient Dynamic Discovery ofhigh-value Information) framework as a scalable sequential
information acquisition system. We assume that informationacquisition is always associated with cost. Given a task, we
dynamically decide which piece of information to acquire
next. We contribute:• A new partial amortized inference method for generative modeling under partial observations.We name it
Partial VAE (Partial Variational AutoEncoder). (Section 3.1).
• A novel fully-Bayesian treatment of Partial VAE toincorporate small datasets, where variational inference
and MCMC are combined for inference (Section 3.1).
• An information theoretic acquisition function with anefficient approximation, yielding a novel variable-wise

EDDI

active learning method (Section 3.2).
φ

2

Z

X

Our proposed framework - EDDI - requires the missing data
model to handle partial observations and sequential decision
making objective to acquire information actively. We thus
review highly related work on these two topics.
Missing data imputation. Prediction based methods haveshown advantages for missing value imputation (Rubin,
1976; Dempster et al., 1977; Scheffer, 2002). Efficient matrix factorization (MF) based methods are often applied(Keshavan et al., 2010; Jain et al., 2010; Salakhutdinov &Mnih, 2008). The probabilistic MF has been used in the active variable selection framework (Lewenberg et al., 2017).
However, they do not scale to large volumes of data andthus are usually not applicable in real-world applications.
Nazabal et al. (2018) use zero imputation (ZI) for amortized inference for both training and test sets with missing
data entries. The drawback of ZI is that it introduces bias
when the data are not missing completely at random. It
also produces poor uncertainty estimates since it does notdistinguish between observed data and imputed data.
Active Learning. Traditional active learning aims to selectthe next data point to label. Information theoretical approaches have shown promising results with different acquisition functions (MacKay, 1992; McCallumzy & Nigamy,
1998; Houlsby et al., 2011). Little work has dealt withmissing values within instances. Zheng & Padmanabhan(2002) deal with missing data values by imputing with traditional non-probabilistic methods (Little & Rubin, 1987)
first. We note that these methods assume that the data are
fully observed, and the acquisition decision is instance wise.
Different from the approaches, our proposed framework
performs variable-wise active learning for each instance.

3

θ

φ

Z

θ

φ

Z

θ

Related work(a) VAEX

X(b) Partial VAE (c) Bayesian PVAE

Figure 1: Graphical representation
where R(i | xO ) quantifies the merit of our prediction of f (·)
given x0 and xi . Furthermore, the reward can quantify otherproperties, e.g. the cost of acquiring xi .
3.1

Partial Amortization of Inference Queries

As discussed before, to enable sequential active feature selection, we need a scalable model to handle missing values.
We present Partial VAE (Figure 1(b)) which enables amortized inference for partially observed data and present a full
Bayesian extension, Bayesian Partial VAE (Figure 1(c)).
From VAE to Partial VAE. A VAE (Figure 1(a)) defines a
generative model in which the data x is generated from latent variables z, defined as p(x, z; θ ) = ∏n pθ (xn |zn )p(zn ),where pθ (xn |zn ), is realized by a deep neural network. To
approximate the posterior pθ (zn |xn ), VAEs use amortizedvariational inference. Specifically, it uses an encoder, which
is another neural network with the data xn as input to produce a variational approximation of the posterior q(zn |xn ; φ ).
As traditional variational inference, VAE is trained by maximizing an evidence lower bound (ELBO).
VAEs are not directly applicable when data points containmissing values. Non-armortized variational inference could
be applied, but will be too impractical due to the space ofpossible missing patterns is huge. However, note that in a
VAE, p(x|z) is factorized, i.e.
p(x|z) = ∏ pi (xi |z).

(2)iMethodProblem formulation In this paper we address the following sequential active variable selection problem. Let
x = [x1 , . . . , x|I| ] be a set of random variables with probability density p(x). At each step, let a subset of the variables
xO , O ⊂ I, be observed while the variables xU , U = I \ O,
are unobserved. We assume that we can query the valueof variables xi for i ∈ U. The goal of sequential active
variable selection at each step is to query a sequence ofvariables in U in order to predict a quantity of interest f (x),as accurately as possible, where f (·) can be any (random)function. This problem, in the simplified myopic setting,can be formalized as that of proposing the next variable xi∗
to be queried by maximizing a reward function R, i.e.
i∗ = arg max R(i | xO ),
i∈U

(1)

This implies that given z, the observed variables xO areconditionally independent of xU . Therefore,p(xU |xO , z) = p(xU |z),(3)

and inferences about xU can be reduced to inference aboutz. Hence, the key object of interest in this setting is p(z|xO ),i.e., the posterior over the shared latent variables z given theobserved variables xO . Once we obtain z, computing xU is
straightforward. To approximate p(z|xO ), we introduce an
auxiliary variational inference network q(z|xO ) and definea partial variational lower bound L partial ,log p(xO ) ≥ log p(xO ) − DKL (q(z|xO )kp(z|xO ))= Ez∼q(z|xO ) [log p(xO |z) + log p(z) − log q(z|xO )]≡ L partial .

(4)

EDDI

We call the the inference net for q(z|xO ) the partial inference net since it takes a set of partially observed variables
xO whose dimensionality may vary among different data
points. Specifying q(z|xO ) requires distributions for any
partition {O,U} of I.

• Variational inference stage: Conditioned on the previous MCMC sample of w, we optimize the L partial
w.r.t the encoder q(z|xO ), with w fixed.
• MCMC stage: We use Stochastic Gradient Hamiltonian Monte Carlo (SG-HMC) (Chen et al., 2014) tosample w.
This requires computing the gradient∇w log p(xO ,w), which is approximated by L partial +log p(w) from the variational inference stage.
Amortized Inference with partial observations. For both
Partial VAE and Bayesian Partial VAE , inference underpartial observations requires the inference net q(z|xO ) to becapable to handle arbitrary set of observed data, and sharing
parameters across these different sized sets of observations
for amortization.
Inspired by the Point Net (PN) approach for point cloudclassification (Qi et al., 2017; Zaheer et al., 2017), we specify the approximate distribution q(z|xO ) by a permutation
invariant set function encoding, given by:c(xO ) := g(h(s1 ), h(s2 ), ..., h(s|O| )),(5)

where sd carries the information of the input of d-th observed variable, and |O| is the number of observed variables.
In particular, sd contains the information about the identity of the input ed and the corresponding input value xd .
There are many ways to define ed , such as the coordinatesof observed pixels for images. In this work, we treat e as an
unknown embedding, which is optimized during training.
There are also different ways to construct sd . A common
choice is concatenation, sd = [ed , xd ], which is commonly

e1
...
e|O|

×
x1
×

(1, K)hsharedgch

ENCODER

From Partial VAE to Bayesian Partial VAE Both Partial
VAE and VAE utilize deep neural networks for the generative model. It is flexible but requires a large volume of data
for training. In many real-world applications, only small
amounts of training data are available, resulting in that themodel overfits the training data. To address this, we propose
a fully Bayesian scheme for the generative model, naming
Bayesian Partial VAE. In particular, we assume the decoder weights, w, are also random variables with the priordistribution p(w). Approximate inference for both z and w
is required. Mean-field amortized inference is scalable but
less accurate comparing to MCMC due to the mean-fieldassumption. In this case, z is a local variable which scaleswith the number of data N, whereas w is a low dimensional
global variable. We thus propose to combine variational
inference for z and MCMC inference for w by performing
the following two steps recursively:(|O|, M)

x|O|

Figure 2: Illustration of Partial VAE encoder architecture(PNP setting)used in computer vision applications (Qi et al., 2017). We
refer to this setting using concatenation as the Pointnet (PN)
specification of Partial VAE. However, the construction of
sd can be more flexible.
We propose to construct sd = ed ∗ xd using element-wise
multiplication, shown in Figure 2. We show that this formulation generalizes naive Zero Imputation (ZI) VAE (Nazabal
et al., 2018) (cf. Appendix B.1). We refer to the multiplication setting as Pointnet Plus (PNP) specification of Partial
VAE. The theoretical consideration of relating ZI to PNP ispresented in Appendix B.1.
We can then use a neural network h(·) to map input sd to RK ,where and K is the latent space size. The key to the PNP
structure is the permutation invariant aggregation operationg(·), such as max-pooling or summation. In this way, themapping c(xO ) is invariant to permutations of elements of
xO , and xO can have arbitrary length. Finally, the fixed-size
code c(xO ) is fed into a ordinary neural network, that transforms the code into the statistics of a multivariate Gaussian
distribution to approximate p(z|xO ).
3.2

Efficient Dynamic Discovery of High-value
Information

We now cast the sequential active variable selection problem(1) as adaptive Bayesian experimental design at each step,
utilizing p(xU |xO ) inferred by the Partial VAE. Algorithm 1summarize the EDDI framework.
Information Reward. We designed a variable selection acquisition function in an information theoretic way following
Bayesian experimental design (Lindley, 1956; Bernardo,1979). For a given task, we may be interested in statistics
of some variables xφ , where xφ ⊂ xU . Given a new instance(user), assume we have observed xO so far for this instance,
and we need to select the next variable xi (an element of
xU\φ ) to observe. Following Bernardo (1979), We select xi
by maximizing:R(i, xO ) = Exi ∼p(xi |xO ) DKL p(xφ |xi , xO ) k p(xφ |xO ) . (6)In our paper, we mainly consider the case that a subset of interesting observations represents the statistics of interest xφ .
Sampling xi ∼ p(xi |xo ) is approximated by xi ∼ p̂(xi |xo ),EDDI

Algorithm 1 EDDI: Algorithm Overview
Require: Training dataset X, which is partially observed;
Test dataset X∗ with no observations collected yet; Indices φ of target variables.
1: Train (Bayesian) Partial VAE by optimizing partial
variational bound(cf. Section 3.1)
2: Actively acquire feature value xi to estimate x∗φ foreach test point (cf. Section 3.2)
for each test instance doxO ← 0/ (no variable value has been observed forany test point)
repeatChoose variable xi from U \ φ to maximize theinformation reward (Equation (9))
xO ← xi ∪ xO
until Stopping criterion reached (e.g. the time budget)end for
where p̂(xi |xo ) can be obtained by (Bayesian) Partial VAE.
It is implemented by sampling z ∼ q(z|xo ), then xi ∼ p(xi |z).
The same applies for p(xi , xφ |xo ) in Equation (8).
Efficient approximation of the Information reward. The
Partial VAE allows us to sample xi ∼ p(xi |xo ). However,the KL term in Equation (6),

DKL p(xφ |xi , xo )||p(xφ |xo )
(7)
=−

Z
xφ

p(xφ |xi , xo ) log

p(xφ |xo ),p(xφ |xi , xo )is intractable since both p(xφ |xi , xo ) and p(xφ |xo )are intractable.
For high dimensional xφ , entropy
estimation
could
be difficult.
The entropy term
R
p(x
|x
,
x
)
log
p(x
|x
,
x
)
depends
on i hence cannot
φ i o
φ i o
xφ
be ignored. In the following, we show how to approximate
this expression.
Our proposal is based on the observation that analytic solutions of KL-divergences are available under specific variational distribution families of q(z|xO ) (such as the Gaussian
distribution commonly used in VAEs). Instead of calculating information reward in x space, we have shown that
one can equivalently perform calculations in z space (cf.
Appendix A.1):
R(i, xo ) =Exi ∼p(xi |xo ) DKL [p(z|xi , xo )||p(z|xo )] −(8)
Exφ ,xi ∼p(xφ ,xi |xo ) DKL p(z|xφ , xi , xo )||p(z|xφ , xo ) .
Note that Equation (8) is exact. Additionally, we use thepartial VAE approximation p(z|xφ , xi , xo ) ≈ q(z|xφ , xi , xo ),p(z|xo ) ≈ q(zi |xo ) and p(z|xi , xo ) ≈ q(zi |xi , xo ). This leads
to the final approximation of the information reward:R̂(i, xo ) =Exi ∼ p̂(xi |xo ) DKL [q(z|xi , xo )||q(z|xo )] −(9)
Exφ ,xi ∼ p̂(xφ ,xi |xo ) DKL q(z|xφ , xi , xo )||q(z|xφ , xo ) .

With this approximation, the divergence between q(z|xi , xo )and q(z|xo ) can often computed analytically in the(Bayesian) Partial VAE setting, for example, under Gaussian parameterization. The only Monte Carlo sampling
required is the one set of samples xφ , xi ∼ p(xφ , xi |xo ) thatcan be shared across different KL terms in Equation (9).

4Experiments

In this section, we evaluate our EDDI framework on bothUCI benchmark and real-life intensive care data. We use
EDDI to refer to our framework with Partial VAE and use
Bayesian EDDI (BEDDI) when Bayesian Patial VAE isused in this section. Both EDDI and BEDDI demonstratesignificant superior performance on prediction accuracy at
all steps in the sequential information acquisition.
4.1

Active variable selection with small UCI set

In this section, we demonstrate the performance of EDDI
and BEDDI with Boston Housing dataset from UCI benchmark datasets (Dheeru & Karra Taniskidou, 2017). In
this experiment, we treat different variables of the same
instances as different data points that can be acquired separately. To compare the performance under different amountof training data, We train Partial VAE and Bayesian Partial VAE with random subset of training data with different
sizes. After training under each subset of the training
data , the EDDI and Bayesian EDDI are used to perform
active variable selection. To compare the different strategies, we use the area under the information curve (AUIC),∑t − log p(xφ |xOt ), where the information curve p(xφ |xOt )is defined as predictive likelihood of target variables aftert acquisitions. Smaller AUIC value (could be or negative)
indicates better performance. The AUIC performances of
EDDI and Bayesian EDDI are compared as the number ofobservations are gradually increased.
Results are shown in Figure 3. In Figure 3 (a), we plot
curves of active learning AUIC on Boston Housing dataset
against the size of data set used in the training phase. BothEDDI and Bayesian EDDI consistently outperforms random
active feature selection strategy (denoted as RAND). More
importantly, Fig 3 shows that Bayesian EDDI is more robust
and always outperforms EDDI across different sizes of the
training dataset. Specifically, we further visualize the active
variable selection process under two sizes of training set.
We track the predictive negative likelihoods during active
acquisitions. When the training set is relatively small (Fig 3(b)), EDDI suffers from severe overfitting, while Bayesian
EDDI performs better. When the training set is sufficient(Figure 3 (c)), Bayesian EDDI and EDDI obtains better
performance. In Figure (c), we can see that using EDDI,apersonalized sequential feature selection framework, we canobtain the same prediction performance using only about
30% of the data.

2.0
1.8
1.6
1.4
1.2
1.0
0.8
0.6

avg. neg. test likelihood

BEDDI
EDDI
RAND

AUIC

55.0
52.5
50.0
47.5
45.0
42.5
40.0
37.5
35.0

avg. neg. test likelihood

EDDI

BEDDI
EDDI
RAND

2.0
1.8
1.6
1.4

0

500 1000 1500 2000 2500
Number of training data points

1.2

0

2 4 6 8 10 12
Number of acquisitions

(a)

BEDDI
EDDI
RAND

0

2 4 6 8 10 12
Number of acquisitions

(b)

(c)

Figure 3: (a): Curves of active learning AUIC on Boston Housing dataset against the size of training set. BEDDI: EDDI with Bayesianpartial VAE; EDDI: EDDI with partial VAE; RAND: random selection strategy (with parital VAE). (b)-(c) Information curve with 50, and2500 training data values, respectively. The x-axis shows the number of features that has been acquired and the y-axis shows the negative
log likelihood for the prediction.

Risk assessment with MIMIC-III

We apply EDDI to risk assessment tasks using the Medical
Information Mart for Intensive Care (MIMIC III) database((Johnson et al., 2016)). MIMIC III is the most extensive
publicly available clinical database, containing real-world
records from over 40,000 critical care patients with 60,000
ICU stays. The risk assessment task is to predict the final
mortality. We preprocess the data for this task following
Harutyunyan et al. (2017) 1 . This results in a dataset of21139 patients. We treat the final mortality of a patient as
a Bernoulli variable. For our task, we focus on variable selection, which corresponds to medical instrument selection.
We thus further process the time series variables into static
variables based on temporal averaging.
Since MIMIC-III contains sufficient data, we only focus on
EDDI. We compare the performance of EDDI, using four
different Partial VAE settings (our PNP-partial VAE, theoriginal PointNet approach, Zero Imputing, and Zero Imputing with masks as input), with three baselines. The firstbaseline is the random feature selection strategy (denoted as
RAND) which randomly picks the next variable to observe.
The second baseline is the single best strategy (denoted as
SING) which finds a single fixed global optimal order of
selecting variables. This order is then applied to all data
points. SING uses the objective function as in Equation (9)to find the optimal ordering by averaging over all the data.
Figure 4 shows the information curve of different strategies,
using PNP based Partial VAE as an example. Table 1 shows
the average ranking of AUIC with different settings. In thisapplication, EDDI significantly outperforms other variableselection strategies in all different settings of Partial VAE,
and PNP based setting performs best.
1 https://github.com/yerevann/mimic3-benchmarks

avg. neg. test likelihood

4.2

0.39
0.38
0.37
0.36
0.35
0.34
0.33
0.32

PNP+Ours
PNP+RANDPNP+SING

0

2

4

6

8 10 12 14 16
Steps

Figure 4: Information curves of active variable selection on risk
assessment task on MIMIC III with PNP setting.

MethodEDDI

RandomSingle best

ZI
ZI-m
PN
PNP

8.83 (0.01)
4.91 (0.01)
4.96 (0.01)
4.39 (0.01)

7.97 (0.02)
7.00 (0.01)
6.62 (0.01)
6.18 (0.01)

9.83 (0.01)
5.91 (0.01)
5.96 (0.01)
5.39 (0.01)

Table 1: Average ranking on AUIC of MIMIC III

5

Conclusion

In this paper, we present EDDI, a novel and efficient framework for dynamic active variable selection for each instance.
Within the EDDI framework, we propose Partial VAE which
performs amortized inference to handle missing data. We
also extend it to Bayesian Partial VAE which can be usedeven when a small amount of training data is available.
Based on (Bayesian) Partial VAE, we design a variable wise
acquisition function for EDDI and derive corresponding
approximation method. EDDI has demonstrated its effectiveness on active variable selection tasks across multiple
real-world applications. In the future, we would extend the
EDDI framework to handle more complicated scenarios,
such as time-series, or the cold-start situation.

EDDI

References
Bernardo, J. M. Expected information as expected utility.
The Annals of Statistics, pp. 686–690, 1979.
Chen, T., Fox, E., and Guestrin, C. Stochastic gradienthamiltonian monte carlo. In International conference on
machine learning, pp. 1683–1691, 2014.
Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum
likelihood from incomplete data via the em algorithm.
Journal of the Royal Statistical Society. Series B (methodological), pp. 1–38, 1977.
Dheeru, D. and Karra Taniskidou, E. UCI machine learning
repository, 2017. URL http://archive.ics.uci.
edu/ml.
Harutyunyan, H., Khachatrian, H., Kale, D. C., and Galstyan, A. Multitask learning and benchmarking with clinical time series data. arXiv preprint arXiv:1703.07771,
2017.

Nazabal, A., Olmos, P. M., Ghahramani, Z., and Valera,
I. Handling incomplete heterogeneous data using vaes.
arXiv preprint arXiv:1807.03653, 2018.
Qi, C. R., Su, H., Mo, K., and Guibas, L. J. Pointnet:
Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE Conference on
Computer Vision and Pattern Recognition, pp. 652–660,
2017.
Ranganath, R., Tran, D., and Blei, D. Hierarchical variational models. In International Conference on Machine
Learning, pp. 324–333, 2016.
Rubin, D. B. Inference and missing data. Biometrika, 63(3):581–592, 1976.
Salakhutdinov, R. and Mnih, A. Bayesian probabilistic
matrix factorization using markov chain monte carlo. In
International conference on Machine learning, pp. 880–
887. ACM, 2008.

Houlsby, N., Huszár, F., Ghahramani, Z., and Lengyel, M.
Bayesian active learning for classification and preference
learning. arXiv preprint arXiv:1112.5745, 2011.

Scheffer, J. Dealing with missing data. 2002.

Jain, P., Meka, R., and Dhillon, I. S. Guaranteed rank
minimization via singular value projection. In Advancesin Neural Information Processing Systems, 2010.

Shim, H., Hwang, S. J., and Yang, E. Joint active featureacquisition and classification with variable-size set encoding. In Advances in Neural Information Processing
Systems, 2018.

Johnson, A. E., Pollard, T. J., Shen, L., Li-wei, H. L., Feng,
M., Ghassemi, M., Moody, B., Szolovits, P., Celi, L. A.,
and Mark, R. G. Mimic-iii, a freely accessible criticalcare database. Scientific Data, 3:160035, 2016.
Keshavan, R. H., Montanari, A., and Oh, S. Matrix completion from noisy entries. Journal of Machine Learning
Research, 2010.
Lewenberg, Y., Bachrach, Y., Paquet, U., and Rosenschein,
J. S. Knowing what to ask: A bayesian active learning
approach to the surveying problem. In AAAI, pp. 1396–
1402, 2017.
Lindley, D. V. On a measure of the information provided
by an experiment. The Annals of Mathematical Statistics,
pp. 986–1005, 1956.
Little, R. and Rubin, D. Statistical analysis with missing
data. Technical report, J. Wiley, 1987.
MacKay, D. J. Information-based objective functions for
active data selection. Neural computation, 4(4):590–604,
1992.
McCallumzy, A. K. and Nigamy, K. Employing em andpool-based active learning for text classification. In International Conference on Machine Learning, pp. 359–367.
Citeseer, 1998.

Settles, B. Active learning. Synthesis Lectures on Artificial
Intelligence and Machine Learning, 6(1):1–114, 2012.

Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B.,
Salakhutdinov, R. R., and Smola, A. J. Deep sets. In
Advances in Neural Information Processing Systems, pp.
3394–3404, 2017.
Zakim, D., Braun, N., Fritz, P., and Alscher, M. D. Underutilization of information and knowledge in everyday
medical practice: Evaluation of a computer-based solution. BMC Medical Informatics and Decision Making, 8(1):50, 2008.
Zheng, Z. and Padmanabhan, B. On active learning for dataacquisition. In International Conference on Data Mining,pp. 562–569. IEEE, 2002.

EDDI

A

Additional Derivations

A.1

Information reward approximation

In our paper, given the VAE model p(x|z) and a partial inference network q(z|xo ), the experimental design problem isformulated as maximization of the information reward:
R(i, xo ) = Exi ∼p(xi |xo ) [DKL (p(xφ |xi , xo )||p(xφ |xo ))]R

R

Where p(xφ |xi , xo ) = z p(xφ |z)q(z|xi , xo ), p(xφ |xo ) = z p(xφ |z)q(z|xo ) and q(z|xo ) are approximate condition distributions given by partial VAE models. Now we consider the problem of directly approximating R(i, xo ).
Applying the chain rule of KL-divergence, we have:
DKL (p(xφ |xi , xo )||p(xφ |xo ))= DKL (p(xφ , z|xi , xo )||p(xφ , z|xo ))− Exφ ∼p(xφ |xi ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo )) ,Using again the KL-divergence chain rule on DKL (p(xφ , z|xi , xo )||p(xφ , z|xo )), we have:
DKL (p(xφ , z|xi , xo )||p(xφ , z|xo ))= DKL (p(z|xi , xo )||p(z|xo )) + DKL (p(xφ |z, xi , xo )||p(xφ |z, xo ))= DKL (p(z|xi , xo )||p(z|xo )) + DKL (p(xφ |z)||p(xφ |z))
= DKL (p(z|xi , xo )||p(z|xo )).
The KL-divergence term in the reward formula is now rewritten as follows,
DKL (p(xφ |xi , xo )||p(xφ |xo ))= DKL (p(z|xi , xo )||p(z|xo ))− Exφ ∼p(xφ |xi ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo )) .
One can then plug in the partial VAE inference approximation:p(z|xφ , xi , xo ) ≈ q(z|xφ , xi , xo ),p(z|xi , xo ) ≈ q(z|xi , xo ), p(z|xo ) ≈ q(z|xo )Finally, the information reward is now approximated as:
R(i, xo )≈ Exi ∼p(xi |xo ) [DKL (q(z|xi , xo )||q(z|xo ))]− Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) DKL (q(z|xφ , xi , xo )||q(z|xφ , xo ))= Exi ∼p(xi |xo ) [DKL (q(z|xi , xo )||q(z|xo ))]− Exφ ,xi ∼p(xφ ,xi |xo ) DKL (q(z|xφ , xi , xo )||q(z|xφ , xo )) = R̂(i, xo ).
This new objective tries to maximize the shift of belief on latent variables z by introducing xi , while penalizing theinformation that cannot be absorbed by xφ (by the penalty term DKL (q(z|xφ , xi , xo )||q(z|xφ , xo ))). Moreover, it is morecomputationally efficient since one set of samples xφ , xi ∼ p(xφ , xi |xo ) can be shared across different terms, and the KLdivergence between common parameterizations of encoder (such as Gaussians and normalizing flows) can be computed
exactly without the need for approximate integrals. Note also that under approximationp(z|xφ , xi , xo ) ≈ q(z|xφ , xi , xo ), p(z|xi , xo ) ≈ q(z|xi , xo ), p(z|xo ) ≈ q(z|xo ), sampling xi ∼ p(xi |xo ) is approximated by xi ∼ p̂(xi |xo ), where p̂(xi |xo ) is defined by the following process in Partial
VAE. It is implemented by first sampling z ∼ q(z|xo ), and then xi ∼ p(xi |z). The same applies for p(xi , xφ |z).

EDDI

B

Additional Theoretical Contributions

B.1

Zero imputing as a Point Net

Here we present how the zero imputing (ZI) and PointNet (PN) approaches relate.
Zero imputation with inference netIn ZI, the natural parameter of λ (e.g., Gaussian parameters in variational autoencoders)is approximated using the following neural network:
L

(1)

(0)

f (x) := ∑ wl σ (wl xT )l=1,
where L is the number of hidden units, x is the input image with xi be the value of the ith pixel. To deal with partiallyobserved data x = xo ∪ xu , ZI simply sets all xu to zero, and use the full inference model f (x) to perform approximate
inference.
PointNet parameterization The PN approach approximates the natural parameter λ by a permutation invariant set function
g(h(s1 ), h(s2 ), ..., h(sO )),where si = [xi , ei ], ei is the I dimensional embedding/ID/location vector of the ith pixel, g(·) is a symmetric operation suchas max-pooling and summation, and h(·) is a nonlinear feature mapping from RI+1 to RK (we will always refer h as featuremaps ). In the current version of the partial-VAE implementation, where Gaussian approximation is used, we set K = 2H
with H being the dimension of latent variables. We set g to be the element-wise summation operator, i.e. a mapping from
RKO to RK defined by:g(h(s1 ), h(s2 ), ..., h(sO )) = ∑ h(si ).
i∈OThis parameterization corresponds to products of multiple Exp-Fam factors ∏i∈O exp{−hh(si ), Φi}.
From PN to ZI To derive the PN correspondence of the above ZI network we define the following PN functions:h(si ) := ei ∗ xi
Ig(h(s1 ), h(s2 ), ..., h(sO )) := ∑ θk σ ( ∑ hk (si )),
i∈O

k=1

where hk (·) is the kth output feature of h(·). The above PN parameterization is also permutation invariant; setting L = I,
(1)(0)θl = wl ,(wl )i = (ei )l the resulting PN model is equivalent to the ZI neural network.
Generalizing ZI from PN perspective In the ZI approach, the missing values are replaced with zeros. However, thisad-hoc approach does not distinguish missing values from actual observed zero values. In practice, being able to distinguishbetween these two is crucial for improving uncertainty estimation during partial inference. One the other hand, we havefound that PN-based partial VAE experiences difficulties in training. To alleviate both issues, we proposed a generalizationof the ZI approach that follows a PN perspective. One of the advantages of PN is setting the feature maps of the unobserved
variables to zero instead of the related weights. As discussed before, these two approaches are equivalent to each other only
if the factors are linear. More generally, we can parameterize the PN by:h(1) (si ) := ei ∗ xi
(1)

(1)

h(2) (hi ) := NN1 (hi )
(2)

(1)

g(h(s1 ), h(s2 ), ..., h(sO )) := NN2 (σ ( ∑ hk (hi ))),
i∈O

where NN1 is a mapping from RI to RK defined by a neural network, and NN2 is a mapping from RK to R2H defined byanother neural network.

EDDI

B.2

Approximation Difficulty of the Acquisition Function

Traditional variational approximation approaches provide wrong approximation direction when applied in this case (resulting
in an upper bound of the objective Rφ (i, xO ) which we maximize). Justification issues aside, (black box) variational
approximation requires sampling from approximate posterior q(z|xO ), which leads to extra uncertainties and computations.
For common proposals of approximation:• Directly estimate entropy via sampling ⇒ problematic for high dimensional target variables• Using reversed information reward Exi ∼p(xi |xo ) [DKL (p(xφ |xo )||p(xφ |xo , xi ))], and then apply ELBO (KL-divergence)⇒ This does not make sense mathematically, since this will result in upper bound approximation of the (reversed)
information objective, this is in the wrong direction.
• Ranganath’s bound (Ranganath et al., 2016) on estimating entropy⇒ gives upper bound of the objective, wrong
direction.
• All the above methods also needs samples from latent space (therefore second level approximation needed).
B.3

Connection of EDDI information reward with BALDWe briefly discuss connection of EDDI information reward with BALD (Houlsby et al., 2011) and. MacKay’s work
(MacKay, 1992). Assuming the model is correct, i.e. q = p, we haveR(i, xo ) = Exi ∼p(xi |xo ) [DKL (p(z|xi , xo )||p(z|xo ))]− Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo )) .
Note that based on McKay’s relationship between entropy and KL-divergence reduction, we have:
Exi ∼p(xi |xo ) [DKL (p(z|xi , xo )||p(z|xo ))]=Exi ∼p(xi |xo ) [H(p(z|xi , xo )) − H(p(z|xo ))]] .
Similarly, we haveExi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo ))=Exφ ∼p(xφ |xo ) Exi ∼p(xi |xφ ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo ))=Exφ ∼p(xφ |xo ) Exi ∼p(xi |xφ ,xo ) H(p(z|xφ , xi , xo )) − H(p(z|xφ , xo ))=Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) H(p(z|xφ , xi , xo )) − Exφ ∼p(xφ |xo ) Exi ∼p(xi |xφ ,xo ) H(p(z|xφ , xo ))=Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) H(p(z|xφ , xi , xo )) − Exφ ∼p(xφ |xo ) H(p(z|xφ , xo )) ,
where MacKay’s result is applied to Exi ∼p(xi |xφ ,xo ) DKL (p(z|xφ , xi , xo )||p(z|xφ , xo )) . Putting everything together, we haveR(i, xo ) = Exi ∼p(xi |xo ) [H(p(z|xi , xo )) − H(p(z|xo ))]]

− Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) H(p(z|xφ , xi , xo )) + Exφ ∼p(xφ |xo ) H(p(z|xφ , xo ))no= Exi ∼p(xi |xo ) [H(p(z|xi , xo ))] − Exi ∼p(xi |xo ) Exφ ∼p(xφ |xi ,xo ) H(p(z|xφ , xi , xo ))no− Exi ∼p(xi |xo ) [H(p(z|xo ))] − Exφ ∼p(xφ |xo ) H(p(z|xφ , xo )) .

EDDI

We can show that
H(p(z|xi , xo )) − Exφ ∼p(xφ |xi ,xo ) H(p(z|xφ , xi , xo ))
=−

Z
z

Z

=
z,xφ

Z

p(z|xi , xo ) log p(z|xi , xo )dz +p(z, xφ |xi , xo ) log

z,xφp(z, xφ |xi , xo ) log p(z|xφ , xi , xo )p(z, xφ |xi , xo )p(z|xi , xo )p(xφ |xi , xo )=I z, xφ |xi , xo ,which is exactly the conditional mutual information I z, xφ |xi , xo used in BALD. Therefore, our chain rule representationof reward function leads us toR(i, xo ) = Exi ∼p(xi |xo ) I z, xφ |xi , xo − Exi ∼p(xi |xo ) I z, xφ |xo .