Posterior Matching for Arbitrary Conditioning
Authors: Ryan Strauss
Junier Oliva
Venue: NeurIPS 2022 Accept (2022)
Keywords: arbitrary conditioning
variational autoencoders
density estimation
inpainting
unsupervised learning

Abstract: Arbitrary conditioning is an important problem in unsupervised learning, where we seek to model the conditional densities $p(\mathbf{x}_u \mid \mathbf{x}_o)$ that underly some data, for all possible non-intersecting subsets $o, u \subset \{1, \dots , d\}$. However, the vast majority of density estimation only focuses on modeling the joint distribution $p(\mathbf{x})$, in which important conditional dependencies between features are opaque. We propose a simple and general framework, coined Posterior Matching, that enables Variational Autoencoders (VAEs) to perform arbitrary conditioning, without modification to the VAE itself. Posterior Matching applies to the numerous existing VAE-based approaches to joint density estimation, thereby circumventing the specialized models required by previous approaches to arbitrary conditioning. We find that Posterior Matching is comparable or superior to current state-of-the-art methods for a variety of tasks with an assortment of VAEs (e.g.~discrete, hierarchical, VaDE).

Posterior Matching for Arbitrary Conditioning
Ryan R. Strauss Department of Computer Science UNC at Chapel Hill rrsêcs unc edu
Junier B. Oliva Department Computer Science UNC at Chapel Hill joli unc. edu
Abstract
Arbitrary conditioning is an important problem in unsupervised learning, where we seek to model the conditional densities p(x.. xx) that underly some data, for all possible non-intersecting subsets o, u C .l.....d]. However, the vast majority density estimation only focuses modeling the joint distribution p(x). in which important conditional dependencies between features are opaque. We propose a simple and general framework, coined Posterior Matching, that enables Variational Autoencoders (VAFs) to perform arbitrary conditioning, without modification to the VAE itself. Posterior Matching applies to the numerous existing VAE-based approaches to joint density estimation, thereby circumventing the specialized models required by previous approaches arbitrary conditioning. We find that Posterior Matching is comparable or superior to current state-of-the-art methods for variety of tasks with assortment of VAEs (e.g. discrete, hierarchical, VaDE).
Introduction
Variational Autoencoders (VAEs) 21 are a widely adopted class of generative model that have been successfully employed in numerous areas [5] 26 33 [6]. Much of their appeal stems from their ability to probabilistically represent complex data in terms of lower- snwer-dimensional latent codes.
Like most other generative models, VAEs typically designed to model the joint data distribution, which communicates likelihoods for particular configurations of all features at once. This can be useful for some tasks, such generating images, but the joint distribution is limited by its inability to explicitly convey the conditional dependencies between features. In many cases, conditional distributions, which provide the likelihood of an event given some known information, relevant and useful. Conditionals can be obtained in theory by marginalizing the joint distribution, but in practice, this is generally not analytically available and is expensive approximate.
more
Easily assessing the conditional distribution over any subset features is important for tasks where decisions and predictions must be made medical applications may require reasoning over: the distribution of blood pressure given age and weight; or the distribution of heart-rate and blood-oxygen level given age, blood pressure, and BMI; etc. For flexibility and scalability, it is desirable for single model to provide all such conditionals inference time. More formally, this task is known as arbitrary conditioning, where the goal is to model the conditional density p(x, x,) for any arbitrary subsets of unobserved features x,, and observed features X,. In this work, show, by way of simple and general framework, that traditional VAEs can perform arbitrary conditioning, without modification to the VAE model itself.
varied set of possible information. For example, some
Our approach, which call Posterior Matching, is to model the distribution p(z xx) that is induced by some VAE, where z is the latent code. In other words, we consider the distribution of latent codes given partially observed features. We do this by having a neural network output approximate partially observed posterior g(z x.). In order to train this network, we develop a straightforward maximum likelihood estimation objective and show that is equivalent to maximizing p(x,, x.).
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

the quantity of interest. Unlike prior works that use VAEs for arbitrary conditioning, we do not make special assumptions or optimize custom variational lower bounds. Rather, training via Posterior Matching is simple, highly flexible, and without limiting assumptions on approximate posteriors (e.g, q(z xo) need not be reparameterized and can thus be highly expressive).
We conduct several experiments in which we apply Posterior Matching to various types of VAES for myriad of different tasks, including image inpainting, tabular arbitrary conditional density estimation, partially observed clustering, and active feature acquisition. We find that Posterior Matching leads to improvements over prior VAE-based methods across the range of tasks we consider.
2 Background
Arbitrary Conditioning A core problem in unsupervised learning is density estimation, where we are given a dataset D {x(") Ji=1 of i..d. samples drawn from an unknown distribution p(x) and wish to learn a model that best approximates the probability density function p. A limitation of only learning the joint distribution p(x) that it does not provide direct access to the conditional dependencies between features. Arbilrary conditional density estimation [18] 24
38) is more
general task where we want to estimate the conditional density p(x,. xx) for all possible subsets of observed features
{l...., ... d) and unobserved features {L...., . d} such that and do not intersect. Here, X, E glol and x,, E I Estimation of joint marginal likelihoods is special case where 0 = 0. Note that, while not strictly necessary for arbitrary conditioning methods [2. 38 assume D is fully observed, requirement for training traditional VAEs.
class of genera- tive models that assume a generative process in which data likelihoods are represented as p(x) f p(x z) p()) dzz where z latent variable that typically has lower dimensionality than the data x. A tractable distribution that affords easy sampling and likelihood evaluation, such as standard Gaussian, is usually imposed on the prior p(z). These models learned by maximizing the evidence lower bound (ELBO) of the data likelihood:
Variational Autoencoders Variational Autoencoders (VAEs)
logp(x) 2 Ex-gu (1xx) llggllloosss 2)] (L x) p(z)),
where To (z x) and Px(x z) are the encoder (or approximate posterior) and decoder of the VAE, respectively. The encoder and decoder are generally neural networks that output tractable distributions (e.g.. a multivariate Gaussian). In order to properly optimize the ELBO, samples drawn from qu [z x) must be differentiable with respect to the parameters of the encoder (often called the reparameterizatiio trick). After training, new data point & can be easily generated by first sampling z from the prior, then sampling p z).
3 Posterior Matching
In this section we describe our framework, coined Posterior Matching, model the under- lying arbitrary conditionals in a VAE. In many respects, Posterior Matching cuts the Gordian knot to uncover the conditional dependencies. Following our insights, we show that our ap- proach is direct and intuitive. Notwithstanding, we are the first to apply this direct methodol- ogy for arbitrary conditionals in VAES and are the first to connect our proposed loss with arbi- trary conditional likelihoods p(x, Xo). Note that we are not proposing a new type of VAE, Rather, we are formalizing simple and intuitive methodology that can be applied to numerous existing (or future) VAEs.
Observ Right
36
Xa
Lusant Space
Conditionll Latent Codes
Figure A two-dimensional VAE latent space (left) that represents the distribution of codes of handwritten 3s. 5s, and 8s. Conditioning on the subset of pixels shown in the center should then result in distribution over latent codes that only correspond to 3s and 8s, as shown on the right.

---

3.1 Motivation
Let us begin with a motivating example, depicted in Figure I| Suppose we have trained VAE on images of handwritten 3s, 5s, and 8s. This VAE has thus learned to represent these images in low-dimensional latent space, Any given code (vector) this latent space represents a distribution over images in the original data space, which can be retrieved by passing that code through the VAE's decoder. Some regions in the latent space will contain codes that represent 3s, some will represent 5s, and some will represent 8s. There is typically only interest in mapping from a given image x to distribution over the latent codes that could represent that image, i.e., the posterior q(z x). However, can just easily ask which latent codes are feasible having only observed part of an image.
For example, if we only see the right half the image shown in Figure we know the digit could be 3 or an 8, but certainly not a 5. Thus, the distribution over latent codes that could correspond to the full image, that (z X.) (where & is the encoder's parameters), should only include regions that represent 3s or 8s. Decoding any sample from Ps Z x,) will produce an image of a 3 or an 8 that aligns with what has been observed.
The important insight is that we can think about how conditioning on x, changes the distribution over latent codes without explicitly worrying about what the (potentially higher-dimensional and more complicated) conditional distribution over X looks like. Once we know (z xx), we can easily move back the original data space using the decoder.
3.2
Approximating the Partially Observed Posterior
The partially observed approximate posterior of interest is not readily available, as it is implicitly defined by the VAE:
Pel xx) = Ex (xx_mp(()xx (z Xo, Xo, xu)]
(1)
where qu (z Xo, x.) = 2 (z x) is the VAE's encoder. Thus, we introduce neural network in order to approximate
Given a network that outputs the distribution go x.) (i.e. the partially observed encoder in Our approach is guided by the priorities of simplicity and generality. We minimize (with respect to 0) the following likelihoods, where the samples are coming from our target distribution as defined in Equation I
Figure 2). we now discuss our approach to training
Ex.=p(()xx))) (xx)) ((x.,,..) | log (z xx)]]
We discuss how this is optimized in practice in Section 3.4
Due to the relationship between negative log-likelihood minimization and KL-divergence minimiza- tion 3 we can interpret Equation 2 as minimizing:
E Ex,, ppp xx) KI. (yro (z X.) go xx.)]
We can directly minimize the KL-divergence in Equation 3 if it is analytically available between the two posteriors, for instance if both posteriors are Gaussians. However, [Equation 2 is more general in that it allows us to use more expressive (e.g., autoregressive) distributions for qo x.) with which the KL-divergence cannot be directly computed. This is important given that Pz. x.) is likely to
be complex (e.g., multimodal) and not easily captured by a Gaussian (as in Figure 1p. Importantly,
there is requirement for go(z] x.) to be reparameterized, which would further limit the class of distributions that can be used. There is a high degree of flexibility in the choice of distribution for the partially observed posterior. Note that this objective does not utilize the decoder.
3.3 Connection with Arbitrary Conditioning
While the Posterior Matching objective from Equation and Equation 3 is intuitive, it is not imme- diately clear how this approach relates back to the arbitrary conditioning objective of maximizing p(X,, x.). We formalize this connection in Theorem 3.1 (see Appendix for proof).
3

---

Messing!
loger's xx)
meghtt
Deceeer
Palx [x))
Figure 2: Overview of the Posterior Matching framework. The partially observed encoder can be appended to any existing VAE enable arbitrary conditioning. At inference time, we can decode samples from qo [ x.) to perform imputation or compute p(x.. xx).
Theorem 3.1. Let (z x) and Ps(x z) be the encoder and decoder. respectively, for some VAE. Additionally, let go(z xx) be an approximate partially observed posterior. Then minimizing
Ex.xp(-(xxo) KL (go(z) (z) Xx, X.. |l go(z xx)) is equivalent minimizing
Ex.-p(((xx))
log o xx) KL (g0(z) X Xx. X.)
go(z X xu))
(4)
with respect to the parameters 0.
The first term inside the expectation in Equation 4 gives us the explicit connection back to the arbitrary conditioning likelihood p(x,, X.). which is being maximized when minimizing Equation 4 The second term acts sort of regularizer by trying make the partially observed posterior match the VAE posterior when conditioned on all of x intuitively, this makes sense as a desirable outcome.
3.4 Implementation
A practical training loss follows quickly from Equation 2. For the outer expectation, we do not have access to the true distribution p(x,, xx). but for given instance that has been partitioned into x. and X., we do have one sample from this distribution, namely x,, So approximate this expectation using X as single sample. This type single-sample approximation is common with VAEs, e.g., when estimating the ELBO. For the inner expectation, we have access to (z x), which can easily be sampled in order to estimate the expectation. practice, we generally use a single sample for this well. This gives the following Posterior Matching loss:
( o, v = (-1x [log go(z xx)]]
(5)
where o is the set of observed feature indices. During training, can be randomly sampled from problem-speeific distribution for each minibatch.
Figure 2 provides a visual overview of our approach. In practice we represent x, as a concatenation of that has had unobserved features set to zero and bitmask that indicates which features observed. This representation has been successful in other arbitrary conditioning models [24] 38) However, this choice is not particularly important Posterior Matching itself, and alternative representations, such as set embeddings, are valid as well.
As required by VAEs, samples from (z x) will be reparameterized, which means that minimizing LPM will influence the parameters the VAE's encoder in addition to the partially observed posterior network. In some cases, this may be advantagcous, the encoder can be guided towards learning a latent representation that is more conducive to arbitrary conditioning. However, it might also be desirable train the VAE independently of the partially observed posterior, in which choose to stop gradients on the samples z ~ (- x) when computing LPM-
Similarly, the partially observed posterior can be trained against an existing pretrained VAE. In this case, the parameters of the VAE's encoder and decoder are frozen, and we only optimize Lpm with respect to 0. Otherwise, we jointly optimize the VAE's ELBO and LPM-

---

We emphasize that there is a high degree of flexibility with the choice of VAE, i.e. we have not imposed any unusual constraints. However, there are some potentially limiting practical considerations that have not been explicitly mentioned yet. First, the training data must be fully observed, as with traditional VAEs, since LPM requires sampling qv (z x). However, given that the base VAE requires fully observed training data anyway, this is generally not a relevant limitation for our purposes. Second, is convenient in practice for the VAE's decoder to be factorized, i.c. p(x z) = II, p(r; z). as this allows to easily sample from p(x,. z) (sampling x,, is less straightforward with other types of decoders). However, it is standard practice to use factorized decoders with VAEs, so this is ordinarily not concern. We also note that, while useful for easy sampling, factorized decoder is not necessary for optimizing the Posterior Matching loss, which does not incorporate the decoder.
3.5 Posterior Matching Beyond Arbitrary Conditioning
The concept of matching VAE posteriors is quite general and has other uses beyond the applica- tion of arbitrary conditioning. We consider one such example, which still has ties to arbitrary conditioning, in order to give a flavor for other potential uses.
Truth
Step 1
Step 5
Step 10
Step 20
Step 30
A common application of arbitrary conditioning active feature acquisition 13 23 25 where informative features are sequentially acquired on an instance-by by by-instance basis. In the unsuper- vised case, the aim is to acquire as few features as possible while maximizing the ability to re- construct the remaining unobserved features (see Figure 3 for example).
Figure Example acquisitions from our CNN model using the lookahead posteriors (see [Sec [ion 5.5). Top row shows x and blue pixels are unobserved. Bottom row shows imputation of x.-
One approach to active feature acquisition is greedily select the feature that will maximize the expected amount of information to be gained about the currently unobserved features [23 253 For
VAEs, Ma al. 25 show that this is equivalent selecting each feature according to
argmax H (z xx) (x. sEu
[z Xo. 2:)
argmin (xy-p(()xx.) H Xo. ...)
(6)
iEu
For certain families of posteriors, such as multivariate Gaussians, the entropies in EEquatin 6 can be analytically computed. In practice, approximating the expectation in Equation 6 done via entropies
of the posteriors (f1) (z x,) = Enp-poss x.) go xx. where samples from po, xx) are produced by first sampling z ~ (o x.) and then passing z through the VAE's decoder pa(a; z) (we call p(*) (z x,) the "lookahead" posterior for feature i, since it is obtained by imagining what the posterior will look like after one acquisition into the future). Hence, computing the resulting entropies requires one network evaluation per sample of to encode for i E u. Thus, if using k samples for each F;, each greedy step will be 12(k: -ul)). which may be prohibitive high dimensions.
In analogous fashion to the Posterior Matching approach that has already been discussed, we can train neural network directly output the lookahead posteriors for all features at once. The Posterior Matching loss in this case is
CPM--Lokkahead (x, o, u. w. 0) EBspnno.st
Baago (: (xx0.)))) log o (z x
(7)
VEN
where w is the parameters of the lookahead posterior network. In practice, train single shared network with a final output layer that outputs the parameters of ( (z) x.). Note that given the distributions g (z) x.) for all i, computing the greedy acquisition choice consists of doing a forward evaluation of our network, then choosing the feature i E u such that the entropy of q (z xx) is minimized. In other words, may bypass the individual samples of z,, and use single shared network for a faster acquisition step. In this setting, we let go (z xx) multivariate Gaussian so that the entropy computation is trivial. See Appendix for a diagram of the entire process. This use of Posterior Matching leads to large improvements in the computational efficiency of greedy active feature acquisition (demonstrated empirically in Section 5.5)
5

---

4
Prior Work
variety of approaches to arbitrary conditioning have been previously proposed. ACE is an autore- gressive, energy-based method that is the current oate-of-theearr for arbitrary conditional likelihood estimation and imputation, although it can be computationally intensive for very high dimensional data 38 ACFlow is a variant of normalizing flows that can give analytical arbitrary conditional likelihoods 244 Several other methods, including Sum-Product Networks [32 Neural Conditioner
2. and Universal Marginalizer [10], also have the ability to estimate conditional likelihoods.
Rezende et al. [34] were among the first to suggest that VAEs can be used for imputation. More recendy, VAEAC was proposed as a VAE variant designed for arbitrary conditioning [18]. Unlike Posterior Matching, VAEAC is not a general framework and cannot be used with typical pretrained VAEs. EDDI is VAE-based approach to active feature acquisition and relies on arbitrary conditioning 225. The authors introduce a "Partial VAE" in order to perform the arbitrary conditioning, which, similarly to Posterior Matching, tries to model p(z X.). Unlike Posterior Matching, they do this by maximizing a variational lower bound on p(x,) using a partial inference network q(z xx) (there is no standard VAE posterior g( x) in EDDI). Gong et al. I3
use a similar approach that is based on the Partial VAE of EDDI. The major drawback of these methods is that, unlike with Posterior Matching, q(2 x,) must be reparameterizable in order to optimize the lower bound (the authors use a diagonal Gaussian). Thus, certain more expressive distributions (c.g., autoregressive) cannot be used. Additionally, these methods cannot be applied to existing VAEs. The methods of Ipsen et al. [7] and Collier et al. [9] are also similar to EDDI, where the former optimizes an approximation of p(x.: b) and the latter optimizes lower bound on p(x, b). Ipsen et al. [17] also focuses on imputation for data that is missing "not at random", a setting that is outside the focus of our work.
There are also several works that have considered learning to identify desirable regions in latent spaces. Engel et al. [1]] start from a pretrained VAE, but then train a separate GAN [14] with special regularizers to do their conditioning. They only condition on binary vectors, y, that correspond to a small number of predefined attributes, whereas we allow for conditioning on arbitrary subsets of continuous features x, (a more complicated conditioning space). Also, their resulting GAN does not make the likelihood q(z y) available, whereas Posterior Matching directly (and flexibly) models q(z x.). which may be useful for downstream tasks (e.g. Section 5.5) and likelihood evaluation (see Appendix). Furthermore, Posterior Matching trains directly through KL, without requiring an additional critic. Whang et al. [40] learn conditional distributions, but not arbitrary conditional distributions (a much harder problem). They also consider normalizing flow models, which are limited invertible architectures with tractable Jacobian determinants and latent spaces that have the same dimensionality as the data (unlike VAEs). Cannella et al. 7 similarly do conditional sampling from model of the joint distribution, but are also restricted to normalizing flow architectures and require a more expensive MCMC procedure for sampling.
5 Experiments
In order to empirically test Posterior Matching, we apply to a variety of VAEs aimed at different tasks. We find that our models are able match surpass the performance of previous specialized VAE methods. All experiments were conducted using JAX (5] and the DeepMind JAX Ecosystem MI Code is available at httpp:///github. com lupalab/
posterior-matchiie cttp://apaaaaapppteiirro---teeeo
Our results are dependent on the choice of VAE, and the particular VAEs used in our experiments were not the product of extensive comparisons and did not undergo thorough hyperparameter tuming that not the focus of this work. With more carefully selected or tuned VAEs, and as new VAEs continue to be developed, we can expect Posterior Matching's 's downstream performance to improve accordingly on any given task. We emphasize that our experiments span a diverse set of task, domains, and types of VAF, wherein Posterior Matching was effective.
5.1 MNIST
In this first experiment, our goal is to demonstrate that Posterior Matching replicates the intuition depicted in Figure 1| We do this by training a convolutional VAE with Posterior Matching on the MNIST dataset. The latent space of this VAË is then mapped to two dimensions with UMAP (27) and visualized in Figure 4 In the figure, black points represent samples from go(z) xx). and for select

---

Test Data
Classes
8
3
Figure 4: UMAP 127 visualization of the latent space of a VAE trained on MNIST. Black dots represent samples from the distribution qe (z xx) learned via Posterior Matching. Images of x,, are to the right of their respective plot. Samples from each mode are decoded and shown.
Table 1:
Peak aggnal-tonoiise ratio (PSNR) and precision/recall scores 35 for image inpaintings. We report mean and standard deviation across five evaluations with different random masks. VAEAC and ACFlow results are taken from Li et al. 24 Higher better for all metrics.
MNIST
OMNIGLOT
CELERA
PSNR
Precision
Recall
PSNR
Precision
Recall
PSNR
Precision
Recall
VDVAE + PM cours)
21.603 + 0.022
0.996
0.996
18.256 = 0.038
0.995
0.994
27.190
0.049
0.995
0.995
VQ-VAE (ours)
19.981 0.021
0.989
0.993
17.954
0.046
0.979
0.973
25.531
0.036
0.982
0.984
VABAC
19.613 + 0.042
0,877
0.975
17,693 = 0.023
0.525
0.926
23.656
0.027
0.966
0.967
ACFlow
17.349 0.018
0.945
0.984
0.962
22.393
0.988
ACFlow+BG
20.828 + 0.031
0.947
0.983
18.838
0.009
0.967
0.970
25.723
0.020
0.964
0.987
samples, the corresponding reconstruction is shown. The encoded test data is shown, colored by true class label, to highlight which regions correspond to which digits. We see that the experimental results nicely replicate our earlier intuitions - the learned distribution (z x,) puts probability mass only in parts of the latent space that correspond plausible digits based on what is observed and successfully captures multimodal distributions (see the second column in Figure 4
5.2 Image Inpainting
One practical application of arbitrary conditioning is image inpainting, where only part of an image observed and we want to fill in the missing pixels with visually coherent imputations. As with prior works [18 24), assume pixels are missing completely at random. We test Posterior Matching as an approach to this task by pairing with both discrete and hierarchical VAEs.
Vector Quantized- VAEs We first consider VQ-VAE 30 a type of VAE that is known work well with images. VQ-VAR differs from the typical VAE with its use of discrete latent space. That is, each latent code is a grid of discrete indices rather than vector of continuous values. Because the latent space is discrete, Oord et al. 30 model the prior distribution with a PixelCNN 29 36
after training the VQ-VAE. We similarly use conditional PixeICNN model go xx). First, convolutional network maps X. to a vector, and that vector is then used as conditioning input to the PixeICNN. More architecture and training details be found in the Appendix. We train VQ-VAEs with Posterior Matching for the MNIST, OMNIGLOT, and CELEBA datasets. Table I reports peak signal-to to-il-e ratio (PSNR) and precision/recall [35] for inpaintings produced by our model. We find that Posterior Matching with VQ-VAK consistently achieves better precision/recall scores than previous models while having comparable PSNR.
Hierarchical VA Hierarchical VAEs 22 37 39 are a powerful extension of traditional VAEs that allow for more expressive priors and posteriors by partitioning the latent variables into subsets 2 [zi,..., ). A hierarchy is then created by fractorizing the prior p(z) = II. p(2; ziei and posterior g( x) = II. ( Zez x). These models have demonstrated impressive perfor- mance images and can even outperform autoregressive models [8]. Posterior Matching can be
7

---

18
10
10
2
R
2
2
2
2
2
0
3
0
6
5
5 5 5
5 5
(a) OMNIGLOT
(b) MNIST
(c) CELERA
Example image inpaintings from VDVAE with Posterior Matching. First two columns on the left are x and x,- Inpainting samples are on the right.
Figure 5:
Table 2: Test normalized root-mean
emror (NRMSE) and arbitrary conditional log- likeliiosss (LL) for UCI datasets. Lower is better for NRMSE, and higher is better for LL. Results for methods other than Posterior Matching are reproduced from Strauss and Oliva (38]. Mean and standard deviation are reported over 5 random observed masks for each instance.
IVASK
G-4
MIKIROONS
NRMSE
11.
NRMSE
12.
NRMSE
NRMSE:
11.
Macchng
VABAC:
0.890
0:001
-0.42: 0.002
0.534 1: 0M83
2 0..00
01896 0/00
-10.02
0.010
0,982
0.002
-3.432 0.000
0.615 a 0.000
/4350-2.000)
ACR
8:007
0661
0:007
8:077
9:64
0.990 0/001
-3.99
0.005
0-011 0.00
0.044
0.535
0000
ACE Proposd
0.00
0.32
0.301
-4417 = 0.005
0-36
0.535
0.009
0,561
8::27
ACFow BG
0.833 0.002
0.525
+ 0.00
0.369 9:016
7,503 e 0.011
0.861 = 0.001
-6.33
= 0.006
0,42 e: 0.000
.1.098 0.002
0,512 2 0.000
$1.399 0.000
naturally applied hierarchical VAEs, where the partially observed posterior is represented as q(z x.) - [I q(z Zez, Xo). We adopt the Very Deep VAE (VDVAE) architecture used by Child [8] and extend it to include the partially observed posterior (see Appendix for training and architecture details). We note that due to our hardware constraints. we trained smaller models and for fewer iterations than Child [8] Inpainting results for our VDVAE models are given in Table We see that they achieve better precision/recall scores than the VQ-VAE models and, unlike VQ-VAE, are able to attain better PSNR than ACFlow for MNIST and CELEBA. Figure 5 shows some example inpaint- ings, and additional samples are provided in the Appendix. The fact that we see better downstream performance when using VDVAE than when using VQ-VAE is illustrative of Posterior Matching's ability to admit easy performance gains by simply switching to more powerful base VAE.
5.3 Real-valued Datasets
We evaluate Posterior Matching on real-valued tabular data, specifically the benchmark UCI repository datasets from Papamakarios et al. (31). We follow the experimental setup used by Li et al. [24] and Strauss and Oliva 38 In these experiments, we train basic VAE models while simultaneously learning the partially observed posterior. Given the flexibiliiyy that Posterior Matching affords, use an autoregressive distribution for go x.). Further details can be found in the Appendix.
Table 2) reports the arbitrary conditional log-likelihoods and normalized root-mean-sque -qqure error (NRMSF) of imputations for our models (with features missing completely at random). Likelihoods are computed using an importance sampling estimate (see Appendix for details). We primarily compare to VAEAC as baseline in the VAE family, however we also provide results for ACE and ACFlow for reference. We see that Posterior Matching is able to consistently produce more accurate imputations and higher likelihoods than VAEAC. While our models don't 't match the likelihoods achieved by ACE and ACFlow, Posterior Matching is comparable to them for imputation NRMSE.
5.4 Partially Observed Clustering
Probabilistic clustering often views cluster assignments as a latent variable. Thus, when applying Posterior Matching in this setting, may perform "partially observed" clustering, which clusters instances based on a subset of observed features. We consider VaDE, which uses a mixture of Gaussians the prior, allowing it to do unsupervised clustering by treating each Gaussian component as one of the clusters [9)]. Despite differences in how VaDE is trained compared to a classic VAE, training a partially observed encoder via Posterior Matching remains exactly the same.
We train models on both MNIST and FASHION MNIST (see Appendix for experimental details). Figure 6 shows the clustering accuracy of these models as the percentage of (randomly selected)
8

---

Guserig Accurocy
10-
Clssoriig Accuracy
0.6
0.8
84
0.4
Pasterior Matchinn
Pastenar Matching
0.2
Supervied Bassline
02
Superiied Bassing
10
0.
0.4
0.6
08
16
Percent Observed Obsevved
Percent Observed
Figure 6: Partially observed clustering accuracy achieved by Posterior Matching with VaDE.
observed features changes. As baseline, we train supervised model where the labels are the cluster predictions from the VaDE model when of the features are observed. We see that Posterior Matching is able to match the performance of the baseline, and even slightly outperform it for low percentages of observed features. Unlike the supervised approach, Posterior Matching has advantage of being generative.
5.5 Very Fast Greedy Feature Acquisition
As discussed in Section 3.5 we can use Posterior Matching outside of the specific task of arbitrary conditioning. Here, we consider the problem greedy active feature acquisition. We train VAË with Posterior Matching network that outputs the lookahead posteriors described in [Section 3.5 using the loss [Equaatonn Note that we are also still using Posterior Matching order learn (z x,) and therefore produce reconstructions. Training details can be found in the Appendix.
We consider the MNIST dataset and compare to EDDI as a baseline, using the authors' publicly available code. We downscale images 16 x 16 since EDDI has difficulty scaling to high- dimensional data. We also only evaluate on the first 1000 instances of the MNIST test set, as the EDDI code was very slow when computing the greedy acquisition policy. EDDI also uses a particular architecture that not compatible with convolutions. Thus we train MI.P-based VAF on flattened images in order to make fair comparison. However, since Posterior Matching does not place any limitations on the type of VAE being used, we also train convolutional version. For our models, we greedily select the feature acquire using the more expensive sampling with the lookahead posteriors (which requires no sampling). In both cases, imputations are computed with expectation over 50 latent codes, is done for EDDI. An example acquisition trajectory is shown in Figure 3
0.3
022
0.0
08-
E201
Sending- Cill PV iso Londdd CMN
20
Mamber Acqui
Figure 7: Average RMSE of reconstructions for greedy active feature acquisition.
campling-based based approach (similar to EDDI) as well as
Figure 7 presents the root-mean-squar error, averaged across the test instances, when imputing with different numbers of acquired features. We see that our models are able to achieve lower error than EDDI. We also that acquiring based on the lookahead posteriors incurs only minimal increase in error compared to the sampling-based method, despite being far more efficient. Computing the greedy choice with our model using the sampling-basee approach takes 68 ms + 917 js (for single acquisition on CPU). Using the lookahead posteriors, the time only 310 μs = 15.3 jss roughly 219x speedup.
6 Conclusions
We have presented an elegant and general framework, called Posterior Matching, that allows VAEs to perform arbitrary conditioning. That is, we can take an existing VAE that only models the joint distribution p(x) and train an additional model that, when combined with the VAF, is able assess any likelihood p(x,, x,) for arbitrary subsets of unobserved features x, and observed features x.-

---

We applied this approach to variety of VAEs for a multitude of different tasks. We found that Posterior Matching outperforms previous specialized VAEs for arbitrary conditioning with tabular data and for image inpainting. Importantly, we find that one can switch to more powerful base VAF and get immediate improvements in downstream arbitrary conditioning performance "for free," " without making changes to Posterior Matching itself. We can also use Posterior Matching to perform clustering based on partially observed inputs and to improve the efficiency of greedy active feature acquisition by several orders of magnitude at negligible cost to performance.
With this work, we hope to make arbitrary conditioning more widely accessible. Arbitrary condi- tioning no longer requires specialized methods, but can instead be achieved by applying one general framework to common VAEs. As advances are made in VAEs for joint density estimation, we can expect to immediately reap the rewards for arbitrary conditioning.
Acknowledgments and Disclosure of Funding
We would like to thank Google's ' TPU Research Cloud program for providing free access to TPUs. This research was partly funded by NSF grant IIS2133595 and by NIH IR0IAA02689001AI.
References
[1] Babuschkin, K. Baumli, A. Bell, S. Bhupatiraju, Bruce, P. Buchlovsky, D. Budden, T. Cai, A. Clark, I. Danihelka, C. Fantacci, J. Godwin, C. Jones, T. Hennigan, M. Hes- sel, S. Kapturowski, T. Keck, I. Kemaev, M. King, L. Martens, V. Mikulik, Norman, J. Quan, G. Papamakarios, R. Ring. F. Ruiz, A. Sanchez, R. Schneider, E. Sezener, S. Spencer, Srinivasan, W. Stokowiec, and F. Viola. The DeepMind JAX Ecosystem, 2020. URL http://github. ctw/ deepmind
[2] M. Belghazi, M. Oquab, Y. LeCun, and D. Lopez-Paz. Learning about an exponential amount of conditional distributions. arXiv preprini 1r0iv:1902.08440 08401, 2019.
[3] C. M. Bishop. Pattern recognition. Machine learning, 128(9), 2006.
[4] S. R. Bowman, L. Vilnis, O. Vinyals, A. M. Dai, R. Jozefowicz, and S. Bengio. Generating sentences from continuous space. arXiv preprint arXiv:1511.0633 2015.
[5] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke,
J. VanderPlas, S. Wanderman-Milng and Q. Zhang. JAX: composable transformations of Python+NummPy programs, 2018 htth:tp:: /ttp///gtttbbbb com/ google/jax
[6] C. J. Butz, J. S. Oliveira, A. E. dos Santos, and A. L. Teixeira. Deep convolutional sum- sum-product networks. Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3248-3255, 2019.
[7] C. Cannella, M. Soltani, and V. Tarokh. Projected latent markov chain monte carlo: Conditional sampling of normalizing flows. arXiv preprint arXiv:2007.06140, 2020.
[8] R. Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv preprint 2tXi::20/1.106500 2020.
[9] M. Collier, A. Nazabal, and C. K. Williams. Vaes in the presence of missing data, arXiv preprint arXiv:2006.05301, 2020.
[10] I.. Douglas, I. Zarov, K. Gourgoulias, C. Lucas, C. Hart, A. Baker, M. Sahani, Y. Perov, and S. Johri. A universal marginalizer for amortized inference in generative models. arXiv preprint arXiv:1711.00663 2017.
[11] J. Engel, M. Hoffman, and A. Roberts. Latent constraints: Learning to generate conditionally from unconditional generative models. arXiv preprint arX1v:1711.0577,, 2017.
[12) H. Fu, C. Li, X. Liu, J. Gao, A. Celikyilmaz, and L. Carin. Cyclical annealing schedule: simple approach to mitigating kl vanishing. arXiv preprint 101145 2019.
10

---

[13] W. Gong, S. Tschiatschek, S. Nowozin, R. E. Turner. J. M. Liriindez-Lbbbtee and C. Zhang. Icebreaker: Element- wise efficient information acquisition with bayesian deep latent gaussian model. 2019.
[14] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
[15] K. Gregor, I. Danihelka, A. Graves, D. Rezende, and D. Wierstra. Draw: recurrent neural network for image generation. In International Conference on Machine Learning, pages 1462-1471. PMIR, 2015.
[16] K. Han, H. Wen, J. Shi, K.-H. Lu, Y. Zhang, D. Fu, and Z. Liu. Variational autoencoder: An unsupervised model for encoding and decoding finri activity in visual cortex. Neurolmage, 198: 125-136, 2019. ISSN 1053-8119. doi: https//dbb..og/////0oo/e/eeeennn
URL
https://wwy http://www.ceencerrett.tccce httpa//cwc...ct..........oo.///a
[17] N. B. Ipsen. P.-A. Mattei, and J. Frellsen. not-miwae: Deep generative modelling with missing not at random data, arXiv preprint arXiv:2006.12871, 2020.
[18] O. Ivanov, M. Figurnov, and D. Vetrov. Variational autoencoder with arbitrary conditioning. arXiv preprint arXiv: 1806.02062.22 2018.
[19] Z. Jiang. Y. Zheng, H. Tan, B. Tang, and H. Zhou. Variational deep embedding: An unsupervised and generative approach to clustering. arXiv preprint arXiv:1611.05144 2016.
[20] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980. 2014.
[21] D. P. Kingma and M. Welling. Auto-enn variational bayes. arXiv preprint arXiv: 2013.
[22] D. P. Kingma, T. Salimans, R. Jozefowicz, Chen, Sutskever, and M. Welling. Improved vari- ational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.
[23] Y. Li and J. Oliva. Active feature acquisition with generative surrogate models. In International Conference on Machine Learning, pages 6450 6459. PMLR, 2021.
[24] Y. Li, S. Akbar, and J. Oliva, Acflow: Flow models for arbitrary conditional likelihoods. In International Conference on Machine Learning, pages 5831-5841. PMLR, 2020.
[25] C. Ma, S. Tschiatschek, K. Palla, J. M. Herúández-Lobato, S. Nowozin, and C. Zhang. Eddi: Efficient dynamic discovery of high-value information with partial vae. arXiv preprint 1rXiv: 1142, 2018.
[26] L. Maaløe, C. K. Sønderby, S. K. Sønderby, and O. Winther. Auxiliary deep generative models.
In International conference on machine learning, pages 1445-1453. PMLR, 2016.
[27] L. McInnes, J. Healy, and J. Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426. 2018.
[28] C. Nash and C. Durkan. Autoregressive energy machines. In International Conference on Machine Learning, pages 1735 1744. PMLR, 2019.
[29] A. v. d. Oord, N. Kalchbrenner, O. Vinyals, L. Espeholt, A. Graves, and K. Kavukcuoglu. Conditional image generation with pixelenn decoders. arXiv preprint arXiv:1606.0 05328,
1606.05388 2016.
[30] A. v. Oord, O. Vinyals, and K. Kavukcuoglu. Neural discrete representation learning. arXiv preprint 1711v:1711.0090) 2017.
[31] G. Papamakarios, Pavlakou, and Murray. Masked autoregressive flow for density estimation. arXiv preprint arXiv: 1705000775, 2017.

---

[32] H. Poon and P. Domingos. Sum- product networks: A new deep architecture. In 2011 IEEE International Conference Computer Vision Workshops (ICCV Workshops), pages 689-690. IEFE, 2011.
[33] A. Razavi, A. van den Oord, and O. Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances neural information processing systems, pages 14866 14876, 2019.
[34] D. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference machine learning, pages 1278-1286. PMLR, 2014.
[35] M. S. Sajjadi, O. Bachem, M. Lucic, O. Bousquet, and S. Gelly, Assessing generative models 0003,, 2018.
via precision and recall. arXiv preprint
[36] Salimans, A. Karpathy, X. Chen, and D. Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv: 170.05/701.055/7. 2017.
[37] C. K. Sønderby, T. Raiko, L. Maaløe. S. K. Sønderby, and O. Winther. Ladder variational autoencoders. Advances in neural information processing systems, 29, 2016.
[38] R. R. Strauss and J. B. Oliva. Arbitrary conditional distributions with energy. arXiv preprini arXiv:2102.044266 2021.
[39] A. Vahdat and J. Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems, 33:19667-19677 2020.
[40] Whang, Lindgren, and Dimakis. Composing normalizing flows for inverse problems.
In International Conference on Machine Learning, pages 1 11158-11111 11169. PMLR, 2021.
12


Reviews:
[{"review_id": "duVFvJPgUhn", "reviewer": "Reviewer_HjaN", "summary": "This manuscript proposes a method \"Posterior Matching\" that trains Variational Autoencoders (VAEs) to be robust to missing data when they are trained with fully observed data.  The idea is straightforward to implement and include with a wide variety of VAEs, and leads to relatively straightforward algorithms.  Theory demonstrates that the method attempts to approximate arbitrary conditional distributions.  Limiting to a subset of VAE approaches, the posterior matching method demonstrates improved performance on a number of benchmarks.", "questions": "How feasible is it to extend this approach to a case where the data is not fully observed?\n\nCan the in-painting examples in fact capture the uncertainty in digits as motivating in Figure 1?\n\nHow well can this method handle missingness that isn't missing completely at random?\n\nWhat is the optimal strategy for picking the missingness patterns during training?  How would this be determined for a new dataset?", "limitations": "The authors clearly state that they assume that they have access to fully observed data for training, and they correctly identify that as a limitation; however, they state that \"this is generally not a relevant limitation for our purposes.\" This limitation is quite large and is downplayed too much in their manuscript.  For example, in the motivating medical case given in the introduction, this would render the method infeasible.  I work with medical data, and I am unaware of any problem within that domain where I could naturally apply their methods, whereas I can and do apply competing missing data approaches for VAEs.  I would suggest that the authors clarify this limitation upfront in the introduction, at least, and discuss it at greater length rather than making it an small comment in a technical section.\n\nThe other major limitation, which was not addressed, is that this method assumes a \"Missing Completely at Random\" model of missingness.  This model is incorrect for many medical problems, and would be an issue on the motivating example in the introduction.  I do not see how to adapt the Posterior Matching approach to handle this case, but it would be helpful for the authors to clearly state this limitation to give the readers a better sense of when this method will and will not work.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The biggest strength of this paper is that it is a theoretically motivated algorithm that is simple to include in many VAE models, and can especially work with more complex variants such as VADE.  While there are two glaring limitations (see next section), this is a big advantage and the empirical results suggest that the method is effective.  The simplicity of the active learning approach is quite nice as well.\n\nA weakness is the relative lack of comparisons.  While Posterior Matching is compared to multiple algorithms within the VAE space, there are many competing approaches for applications such as inpainting.  It seems clear that the proposed approach can improve in this cases, but it is unclear if I should consider their approach for an inpainting task.\n\nThere are several minor issues.  For example, the notation is a bit inconsistent.  As an example, in equation (1) the partially observed posterior is written in terms of the variational posterior; that is explicitly true as the variational distribution is an approximation and they are not exactly equal.  This seems like the authors meant that this term should be similar to $q_\\theta(z|x_0)$ instead.  Either way, this should be clarified and explained more.\n\nAs a comment, given the motivating case in Figure 1, it would have been nice to see the proposed approach maintain the multimodality in the generated examples of Figure 5(b).  I would appreciate getting stronger evidence that the multi-modality posterior is in fact maintained.\n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YtKWkPhcSXg", "reviewer": "Reviewer_hezX", "summary": "Authors propose a model and learning method, called posterior matching, to perform arbitrary conditioning with VAE models. The proposed method can be applied to a variety of VAE models without modifications to the VAE model itself, and can be used with a pre-trained VAE or can be trained alongside with the actual VAE model. ", "questions": "List below contains both minor and major comments. \n\n1. In Background section where you introduce the notation x_u and x_o, it would be very helpful to specifically mention that you assume fully observed datasets. A reader may easily get disoriented that you are dealing with partially observed data. \n\n2. Equation. 1 in Section 3.2: From the point of view of the generative model, Eq. 1 is not exact but rather already an approximation of p(z|x_o) = p(x_o|z)p(z)/p(x_o) because your expression involves the variational approbation q(z|x). \n\n3. Equation 2: it would be useful to more clearly define q_{\\theta}, e.g. by referring to Fig. 2 already at this stage. \n\n4. Equation 3: write already here that what are the parameters that you are optimizing. \n\n5. At a general level, as summarised in Fig. 2, the proposed model contains two branches: the \"original\" VAE and the \"proposed model part\" that includes the partially observed encoder and provides the posterior matching loss. The amortized variational approximation of the \"proposed model part\" takes as an input the \"observed\" variables x_o (used for conditioning) as well as a binary vector b that indicates which variables are missing/non-missing. This proposed model part seems very similar with the models proposed to learn VAEs from partially observed data in (Collier et al, 2020; Ipsen et al, 2021). These methods can learn approximations of the posteriors p(z|x_o,b) and thereby used to infer p(x_u|x_o). Please contrast your model with these two e.g. in Section 4, and include them in ablation studies (see below). \n\nREFS: \nMark Collier, Alfredo Nazabal, Christopher K.I. Williams. VAEs in the Presence of Missing Data pdf  Published on arXiv 13 July 2020. Presented at the first Workshop on the Art of Learning with Missing Values (Artemiss) hosted by the 37th International Conference on Machine Learning (ICML 2020).\n\nNB Ipsen, PA Mattei, J Frellsen, not-MIWAE: Deep Generative Modelling with Missing not at Random Data, International Conference on Learning Representations, 2021. \n\n6. Experiments and test metrics:\nThe general goal of the proposed method is to perform arbitrary conditioning p(x_u|x_o). I would have liked to see much more extensive experiments with variety of datasets that quantify performance directly using that metric, instead of devoting a large number of experiments to various other downstream tasks, where the contribution of the proposed model is less directly quantifiable. Current manuscript evaluate conditional log-likelihoods only in one experiment (Table 2). \n\nRelated to the comment on conditional log-likelihood evaluations: authors evaluate their method on various downstream tasks using a variety of evaluation metrics (peak SNR, precision, recall, clustering accuracy). While the metrics themselves are well-known, the way they are computed and applied here is not explained in detail, which leaves it unclear that what the numbers exactly tell about the proposed method itself. \n\n7. Ablation studies:\nI would have liked to see ablation studies. For example, it common to use the standard VAE or VAE with the Gaussian mixture prior and impute unobserved (missing) feature with zeros x_u = [0 0 ... 0], and then approximate the posterior directly using q(z|x_o,x_u), and then estimate p(x_u|x_o).  More recent methods explicitly model the missing features, such as the method by (Collier et al. 2020) and (Ipsen et al, 2021) (see comment above).  A comprehensive ablation study including e.g. these methods would provide more insight into the proposed model. \n\nOverall, this is an interesting study that addresses an important problem. The proposed method is intuitive and statistically well-motivated, and results appear promising. Regarding the novelty, I would like to see the proposed method to be compared by recent works that propose closely related methods.  An ablation study and more straightforward, direct comparisons on the arbitrary conditioning would strengthen the manuscript. ", "limitations": "Limitations are adequately addressed. ", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\nPrevious work has proposed deep generative models for arbitrary conditioning using e.g. autoregressive and flow based models.  VAE based have also been proposed for arbitrary conditioning.  Regarding originality and novelty, this work seems to extend previous VAE methods by allowing more expressive posteriors as well as by allowing applying the proposed method to a pre-trained VAE model. Authors do not contrast their method to recent methods that allow learning VAEs from partially observed data (see detailed comments below). \n\nQuality & Clarity:\nQuality and clarity appears generally good. I have comments regarding the clarity of the experiments (see below). \n\nSignificance: \nSignificance of the work is very high. Arbitrary conditioning with (more classical) statistical models is of central importance, and the same applies to arbitrary conditioning with deep generative models. Consequently, this work is well motivated and has important applications. ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "5T369ZoVhcx", "reviewer": "Reviewer_5R6L", "summary": "This paper presents a general approach for adding arbitrary conditioning to models based on the VAE architecture. An additional encoder which takes partially observed inputs is used to define a partially observed posterior which is trained to match the fully observed posterior by adding an additional term to the loss function. The simplicity of this approach allows one to train a partially observed posterior for any VAE-based model and even to learn a partially observed posterior for an existing trained VAE. \n\nThe authors demonstrate the effectiveness of their approach on a number of tasks, using MNIST to illustrate the model matches their motivating intuitions and then using image inpainting, imputation of tabular data, and partially observed clustering to demonstrate the effectiveness of their method on tasks involving arbitrary conditioning. The method presented in this paper does not achieve top performance on all tasks, but is competitive with the SOTA and improves on other VAE-based methods. Other methods are less flexible, in the sense that they cannot be immediately integrated with a new VAE architecture, one clear appeal of the method presented here.\n\nAdditionally, the authors apply their method can be used for greedy active feature acquisition, where a conditional distribution based on partially observed features is used to greedily select the most informative feature to add to the model. Here, the method is shown to be more accurate and faster than existing methods.", "questions": "No questions to the authors", "limitations": "One limitation, which the authors remark on, is that this method requires fully observed training data. While there may be cases where it is advantageous to be able to train a model on partially observed data, there are many cases where some fully observed data will be available at training time. ", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is based on a simple but novel idea, which is clearly explained and validated using well thought-out experiments. ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3O2F-M4xnwa", "reviewer": "Reviewer_cCKE", "summary": "This work proposes a simple method for the arbitrary conditioning problem, which trains an amortized inference network for the distribution $q(z | x_o)$. Compared with previous work, the method does not require reparameterizable inference networks, and thus enables the use of more flexible models.\n\n**Post-rebuttal update**: I am satisfied by the authors' response, and believe this work is suitable for publication.", "questions": "I do not have any major questions.  A few minor questions / comments:\n\n- You should cite the early works on data imputation with VAEs, e.g. Rezende, Mohamed & Wierstra (2014, Appendix F).\n\n- In Theorem 3.1 the definitions of $p_{\\theta,\\phi}, q_\\theta(z | x_o, x_u)$ should be moved to the text.  The latter distribution should also depend on $\\phi$ (the definition of $q_\\theta$ does not involve $x_u$).  As both terms depend on $\\theta$, ideally there should be more discussion about the optima, as mentioned in the above section.", "limitations": "Limitations and potential negative societal impact are adequately addressed.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\nThe method is simple, intuitive and appears to have competitive performance.\n\n**Weaknesses**:\n\nI do not have any major concerns.  My only gripe is that Theorem 3.1 appears to add little value -- the distributions to be optimized depend on $\\theta$ in a somewhat complicated way, and the only thing I can read out of Eq. (4) is that when both $p_\\phi$ and $q_\\psi$ are optimized, the imputation distribution determined by $q_\\theta$ will be correct; but this already follows from Eq. (3).  More discussions on the general case would be helpful.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept