Posterior Matching for Arbitrary Conditioning
Anonymous Author(s)
AffiliationAddress email
Abstract
Arbitrary conditioning is an important problem in unsupervised learning, where weseek to model the conditional densities p ( x u | x o ) that underly some data, for allpossible non-intersecting subsets o, u ⊂ { 1 , . . .
, d } .
However, the vast majority ofdensity estimation only focuses on modeling the joint distribution p ( x ) , in whichimportant conditional dependencies between features are opaque.
We propose asimple and general framework, coined Posterior Matching, that enables Variational
Autoencoders (VAEs) to perform arbitrary conditioning, without modification tothe VAE itself.
Posterior Matching applies to the numerous existing VAE-basedapproaches to joint density estimation, thereby circumventing the specializedmodels required by previous approaches to arbitrary conditioning.
We find that
Posterior Matching is comparable or superior to current state-of-the-art methods fora variety of tasks with an assortment of VAEs (e.g. discrete, hierarchical, VaDE)
Introduction
Variational Autoencoders (VAEs) [19] are a widely adopted class of generative model that have beensuccessfully employed in numerous areas [4, 14, 24, 31, 15].
Much of their appeal stems from theirability to probabilistically represent complex data in terms of lower-dimensional latent codes.
Like most other generative models, VAEs are typically designed to model the joint data distribution,which communicates likelihoods for particular configurations of all features at once.
This can beuseful for some tasks, such as generating images, but the joint distribution is limited by its inabilityto explicitly convey the conditional dependencies between features.
In many cases, conditionaldistributions, which provide the likelihood of an event given some known information, are morerelevant 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 to approximate.
Easily assessing the conditional distribution over any subset of features is important for tasks wheredecisions and predictions must be made over a varied set of possible information.
For example, somemedical applications may require reasoning over: the distribution of blood pressure given age andweight ; 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 a single model to provide all such conditionalsat inference time.
More formally, this task is known as arbitrary conditioning , where the goal isto model the conditional density p ( x u | x o ) for any arbitrary subsets of unobserved features x uand observed features x o .
In this work, we show, by way of a simple and general framework, thattraditional VAEs can perform arbitrary conditioning, without modification to the VAE model itself.
Our approach, which we call Posterior Matching, is to model the distribution p ( z | x o ) that is inducedby some VAE, where z is the latent code.
In other words, we consider the distribution of latent codesgiven partially observed features.
We do this by having a neural network output an approximatepartially observed posterior q ( z | x o ) .
In order to train this network, we develop a straightforwardmaximum likelihood estimation objective and show that it is equivalent to maximizing p ( x u | x o ) , the quantity of interest.
Unlike prior works that use VAEs for arbitrary conditioning, we do not makespecial 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 | x o ) 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 amyriad of different tasks, including image inpainting, tabular arbitrary conditional density estimation,partially observed clustering, and active feature acquisition.
We find that Posterior Matching leads toimprovements over prior VAE-based methods across the range of tasks we consider.
Background
Arbitrary Conditioning A core problem in unsupervised learning is density estimation , wherewe are given a dataset D
= { x ( i ) }
Ni =1 of i.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 ofonly learning the joint distribution p ( x ) is that it does not provide direct access to the conditionaldependencies between features.
Arbitrary conditional density estimation [16, 22, 35] is a moregeneral task where we want to estimate the conditional density p ( x u | x o ) for all possible subsets ofobserved features o ⊂ { 1 , . . .
, d } and unobserved features u ⊂ { 1 , . . .
, d } such that o and u do notintersect.
Here, x o ∈ R | o | and x u ∈ R | u | .
Estimation of joint or marginal likelihoods is a specialcase where o = ∅ .
Variational Autoencoders Variational Autoencoders (VAEs) [19] are a class of generative models that assume a generative process in which data likelihoods are represented asp ( x ) =
(cid:82) p ( x | z ) p ( z )
d z , where z is a latent variable that typically has lower dimensionality than the data x .
A tractable distribution that affords easy sampling and likelihood evaluation, such as a standard Gaussian, is usually imposed on the prior p ( z ) .
These models are learned by maximizing the evidence lower bound (ELBO) of the data likelihood:log p ( x ) ≥ E z ∼ q ψ ( ·| x ) [log p ϕ ( x | z )] − KL ( q ψ ( z | x ) || p ( z )) , where q ψ ( z | x ) and p ϕ ( x | z ) arethe encoder (or approximate posterior) and decoder of the VAE, respectively.
The encoder anddecoder are generally neural networks that output tractable distributions (e.g., a multivariate Gaussian).
In order to properly optimize the ELBO, samples drawn from q ψ ( z | x ) must be differentiablewith respect to the parameters of the encoder (often called the reparameterization trick ).
Aftertraining, a new data point ˆ x can be easily generated by first sampling z from the prior, then samplingˆ x ∼ p ϕ ( · | z ) .
Posterior Matching
[Figure]
In this section we describe our framework,coined Posterior Matching, to model the underlying arbitrary conditionals in a VAE.
In manyrespects, Posterior Matching cuts the Gordianknot to uncover the conditional dependencies.
Following our insights, we show that our approach is direct and intuitive .
Notwithstanding,we are the first to apply this direct methodology for arbitrary conditionals in VAEs and arethe first to connect our proposed loss with arbitrary conditional likelihoods p ( x u | x o ) .
Notethat we are not proposing a new type of VAE.
Rather, we are formalizing a simple and intuitivemethodology that can be applied to numerousexisting (or future) VAEs.
3.1 Motivation
Let us begin with a motivating example, depicted in Figure 1.
Suppose we have trained a VAE onimages of handwritten 3s, 5s, and 8s.
This VAE has thus learned to represent these images in a low-dimensional latent space.
Any given code (vector) in this latent space represents a distributionover images in the original data space, which can be retrieved by passing that code through the VAE’sdecoder.
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 an interest in mapping from a given image x to adistribution over the latent codes that could represent that image, i.e., the posterior q ( z | x ) .
However,we can just as 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 1, we know the digit could bea 3 or an 8, but certainly not a 5.
Thus, the distribution over latent codes that could correspond tothe full image, that is p ( z | x o ) , should only include regions that represent 3s or 8s.
Decoding anysample from p ( z | x o ) 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 o changes the distribution overlatent codes without explicitly worrying about what the (potentially higher-dimensional and morecomplicated) conditional distribution over x u looks like.
Once we know p ( z | x o ) , we can easilymove back to the original data space using the decoder.
3.2 Approximating the Partially Observed Posterior
The true partially observed posterior is not readily available, as it is implicitly defined by the VAE:
[Equation]
1) where q ψ ( z | x o , x u ) =
q ψ ( z | x ) is the VAE’s encoder.
Thus, we introduce a neural network inorder to approximate it.
Given a network that outputs the distribution q θ ( z | x o ) , we now discuss our approach to training it.
Our approach is guided by the priorities of simplicity and generality.
We minimize the followinglikelihoods, where the samples are coming from our target distribution as defined in Equation 1:
[Equation]
We discuss how this is optimized in practice in Section 3.4.
Due to the relationship between negative log-likelihood minimization and KL-divergence minimization [3], we can interpret Equation 2 as minimizing:
[Equation]
We can directly minimize the KL-divergence in Equation 3 if it is analytically available between thetwo posteriors, for instance if both posteriors are Gaussians.
However, Equation 2 is more general inthat it allows us to use more expressive (e.g., autoregressive) distributions for q θ ( z | x o ) with whichthe KL-divergence cannot be directly computed.
This is important given that p ( z | x o ) is likely tobe complex (e.g., multimodal) and not easily captured by a Gaussian (as in Figure 1).
Importantly,there is no requirement for q θ ( z | x o ) to be reparameterized, which would further limit the class ofdistributions that can be used.
There is a high degree of flexibility in the choice of distribution for thepartially observed posterior.
Note that this objective does not utilize the decoder.
3.3 Connection with Arbitrary Conditioning
While the Posterior Matching objective from Equation 2 and Equation 3 is intuitive, it is not immediately clear how this approach relates back to the arbitrary conditioning objective of maximizingp ( x u | x o ) .
We formalize this connection in Theorem 3.1 (see Appendix for proof).
Theorem 3.1.
Let q ψ ( z | x ) and p ϕ
( x | z ) be the encoder and decoder, respectively, for some
VAE.
Additionally, let q θ ( z | x o ) be an approximate partially observed posterior.
Then minimizing
E x u
[Equation]
p ( ·| x o ) (cid:104) KL (cid:0) q ψ ( z | x o , x u ) ||
q θ ( z | x o ) (cid:1)(cid:105) is equivalent to minimizing
[Equation]
with respect to the parameters θ .
[Figure]
The first term inside the expectation in Equation 4 gives us the explicit connection back to the arbitraryconditioning likelihood p ( x u | x o ) , which is being maximized when minimizing Equation 4.
Thesecond term acts as a sort of regularizer by trying to 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 nothave access to the true distribution p ( x u | x o ) , but for a given instance x that has been partitionedinto x o and x u , we do have one sample from this distribution, namely x u .
So we approximate thisexpectation using x u as a single sample.
This type of single-sample approximation is common with
VAEs, e.g., when estimating the ELBO.
For the inner expectation, we have access to q ψ ( z | x ) , whichcan easily be sampled in order to estimate the expectation.
In practice, we generally use a singlesample for this as well.
This gives us the following Posterior Matching loss:
[Equation]
5) where o is the set of observed feature indices.
During training, o can be randomly sampled from aproblem-specific distribution for each minibatch.
Figure 2 provides a visual overview of our approach.
In practice, we represent x o as a concatenationof x that has had unobserved features set to zero and a bitmask b that indicates which features areobserved.
This representation has been successful in other arbitrary conditioning models [22, 35].
However, this choice is not particularly important to Posterior Matching itself, and alternativerepresentations, such as set embeddings, are valid as well.
As required by VAEs, samples from q ψ ( z | x ) will be reparameterized, which means that minimizing
L PM will influence the parameters of the VAE’s encoder in addition to the partially observed posteriornetwork.
In some cases, this may be advantageous, as the encoder can be guided towards learninga latent representation that is more conducive to arbitrary conditioning.
However, it might also bedesirable to train the VAE independently of the partially observed posterior, in which case we canchoose to stop gradients on the samples z ∼ q ψ ( · | x ) when computing L PM .
Similarly, the partially observed posterior can be trained against an existing pretrained VAE.
In thiscase, the parameters of the VAE’s encoder and decoder are frozen, and we only optimize L PM withrespect to θ .
Otherwise, we jointly optimize the VAE’s ELBO and L PM .
We emphasize that there is a high degree of flexibility with the choice of VAE, i.e. we have not imposedany unusual constraints.
However, there are some potentially limiting practical considerations thathave not been explicitly mentioned yet.
First, the training data must be fully observed, as withtraditional VAEs, since L PM requires sampling q ψ ( z | x ) .
However, given that the base VAE requiresfully observed training data anyway, this is generally not a relevant limitation for our purposes.
Second,it is convenient in practice for the VAE’s decoder to be factorized, i.e. p ( x | z ) =
(cid:81)
i p ( x i | z ) ,as this allows us to easily sample from p ( x u | z ) (sampling x u is less straightforward with othertypes of decoders).
However, it is standard practice to use factorized decoders with VAEs, so this is ordinarily not a concern.
We also note that, while useful for easy sampling, a factorized decoder isnot necessary for optimizing the Posterior Matching loss, which does not incorporate the decoder.
[Figure]
3.5 Posterior Matching Beyond Arbitrary Conditioning
[Figure]
The concept of matching VAE posteriors is quitegeneral and has other uses beyond the application of arbitrary conditioning.
We consider onesuch example, which still has ties to arbitraryconditioning, in order to give a flavor for otherpotential uses.
A common application of arbitrary conditioningis active feature acquisition [12, 21, 23], whereinformative features are sequentially acquiredon an instance-by-instance basis.
In the unsupervised case, the aim is to acquire as few featuresas possible while maximizing the ability to reconstruct the remaining unobserved features (see Figure 3 for example).
One approach to active feature acquisition is to greedily select the feature that will maximize theexpected amount of information to be gained about the currently unobserved features [21, 23].
For
VAEs, Ma et al.
[23] show that this is equivalent to selecting each feature according toargmax i ∈ u
[Equation]
For certain families of posteriors, such as multivariate Gaussians, the entropies in Equation 6 canbe analytically computed.
In practice, approximating the expectation in Equation 6 is done viathe
[Equation]
(cid:104)
q θ ( z | x o , x i ) (cid:105) , where samples from p θ,ϕ ( x i | x o ) areproduced by first sampling z ∼ q θ ( · | x o ) and then passing z through the VAE’s decoder p ϕ
( x i | z )(we call p ( i ) ( z | x o ) the “lookahead” posterior for feature i , since it is obtained by imagining whatthe posterior will look like after one acquisition into the future).
Hence, computing the resultingentropies requires one network evaluation per sample of x i to encode z , for i ∈ u .
Thus, if using ksamples for each x i , each greedy step will be Ω( k · | u | ) , which may be prohibitive in high dimensions.
In analogous fashion to the Posterior Matching approach that has already been discussed, we can traina neural network to directly output the lookahead posteriors for all features at once.
The Posterior
Matching loss in this case is
[Equation]
where ω is the parameters of the lookahead posterior network.
In practice, we train a single sharednetwork with a final output layer that outputs the parameters of all q ( i ) ω
( z | x o ) .
Note that given thedistributions q ( i ) ω ( z | x o ) for all i , computing the greedy acquisition choice consists of doing a forwardevaluation of our network, then choosing the feature i ∈
u
such that the entropy of q ( i ) ω ( z | x o ) isminimized.
In other words, we may bypass the individual samples of x i , and use a single sharednetwork for a faster acquisition step.
In this setting, we let q ( i ) ω ( z | x o ) be a multivariate Gaussian sothat the entropy computation is trivial.
See Appendix for a diagram of the entire process.
This useof Posterior Matching leads to large improvements in the computational efficiency of greedy activefeature acquisition (demonstrated empirically in Section 5.5).
Prior Work
A variety of approaches to arbitrary conditioning have been previously proposed.
ACE is an autoregressive, energy-based method that is the current state-of-the-art for arbitrary conditional likelihoodestimation and imputation, although it can be computationally intensive for very high dimensionaldata [35].
ACFlow is a variant of normalizing flows that can give analytical arbitrary conditional likelihoods [22].
Several other methods, including Sum-Product Networks [30, 6], Neural Conditioner[2], and Universal Marginalizer [9], also have the ability to estimate conditional likelihoods.
VAEAC is a previous VAE variant designed for arbitrary conditioning [16].
Unlike Posterior Matching,
VAEAC is not a general framework and cannot be used with typical pretrained VAEs.
EDDI is a
VAE-based approach to active feature acquisition and relies on arbitrary conditioning [23].
Theauthors introduce a “Partial VAE” in order to perform the arbitrary conditioning, which, similarly to
Posterior Matching, tries to model p ( z | x o ) .
Unlike Posterior Matching, they do this by maximizinga variational lower bound on p ( x o ) using a partial inference network q ( z | x o ) (there is no standard
VAE posterior q ( z | x ) in EDDI).
Gong et al.
[12] 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 ( z | x o )must be reparameterizable in order to optimize the lower bound (the authors use a diagonal Gaussian).
Thus, certain more expressive distributions (e.g., autoregressive) cannot be used.
Additionally, thesemethods cannot be applied to existing VAEs.
There are also several works that have considered learning to identify desirable regions in latentspaces.
Engel et al.
[10] start from a pretrained VAE, but then train a separate GAN [13] with specialregularizers to do their conditioning.
They only condition on binary vectors, y , that correspond toa small number of predefined attributes, whereas we allow for conditioning on arbitrary subsets ofcontinuous features x o (a more complicated conditioning space).
Also, their resulting GAN does notmake the likelihood q ( z | y ) available, whereas Posterior Matching directly (and flexibly) modelsq
( z | x o ) , 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 anadditional critic.
Whang et al.
[37] learn conditional distributions, but not arbitrary conditionaldistributions (a much harder problem).
They also consider normalizing flow models, which arelimited to invertible architectures with tractable Jacobian determinants and latent spaces that have thesame dimensionality as the data (unlike VAEs).
Cannella et al.
[7] similarly do conditional samplingfrom a model of the joint distribution, but are also restricted to normalizing flow architectures andrequire a more expensive MCMC procedure for sampling.
Experiments
In order to empirically test Posterior Matching, we apply it to a variety of VAEs aimed at differenttasks.
We find that our models are able to match or surpass the performance of previous specialized
VAE methods.
All experiments were conducted using JAX [5] and the DeepMind JAX Ecosystem[1].
A link to publicly available code will be included upon publication.
Our results are dependent on the choice of VAE, and the particular VAEs used in our experimentswere not the product of extensive comparisons and did not undergo thorough hyperparameter tuning— that is not the focus of this work.
With more carefully selected or tuned VAEs, and as new VAEscontinue to be developed, we can expect Posterior Matching’s downstream performance to improveaccordingly on any given task.
We emphasize that our experiments span a diverse set of task, domains,and types of VAE , wherein Posterior Matching was effective.
5.1 MNIST
In this first experiment, our goal is to demonstrate that Posterior Matching replicates the intuitiondepicted in Figure 1.
We do this by training a convolutional VAE with Posterior Matching on the
MNIST dataset.
The latent space of this VAE is then mapped to two dimensions with UMAP [25] andvisualized in Figure 4.
In the figure, black points represent samples from q θ ( z | x o ) , and for selectsamples, the corresponding reconstruction is shown.
The encoded test data is shown, colored bytrue class label, to highlight which regions correspond to which digits.
We see that the experimentalresults nicely replicate our earlier intuitions — the learned distribution q θ ( z | x o ) puts probabilitymass only in parts of the latent space that correspond to plausible digits based on what is observed.
5.2 Image Inpainting
One practical application of arbitrary conditioning is image inpainting, where only part of an imageis observed and we want to fill in the missing pixels (missing completely at random) with visually
[Figure] [Figure]
[Table]
coherent imputations.
We test Posterior Matching as an approach to this task by pairing it with bothdiscrete and hierarchical VAEs.
Vector Quantized-VAEs We first consider VQ-VAE [28], a type of VAE that is known to workwell with images.
VQ-VAE differs from the typical VAE with its use of a discrete latent space.
Thatis, each latent code is a grid of discrete indices rather than a vector of continuous values.
Becausethe latent space is discrete, Oord et al.
[28] model the prior distribution with a PixelCNN [27, 33]after training the VQ-VAE.
We similarly use a conditional PixelCNN to model q θ ( z | x o ) .
First, aconvolutional network maps x o to a vector, and that vector is then used as a conditioning input to the
PixelCNN.
More architecture and training details can be found in the Appendix.
We train VQ-VAEswith Posterior Matching for the MNIST, O MNIGLOT , and C ELEB A datasets.
Table 1 reports peaksignal-to-noise ratio (PSNR) and precision/recall [32] for inpaintings produced by our model.
Wefind that Posterior Matching with VQ-VAE consistently achieves better precision/recall scores thanprevious models while having comparable PSNR.
Hierarchical VAEs Hierarchical VAEs [20, 34, 36] are a powerful extension of traditional VAEsthat allow for more expressive priors and posteriors by partitioning the latent variables into subsetsz
= { z 1 , . . .
, z L } .
A hierarchy is then created by fractorizing the prior p ( z ) =
(cid:81)
i p ( z i
| z <i )and posterior q ( z | x ) = (cid:81)
i q ( z i |
z <i , x ) .
These models have demonstrated impressive performance on images and can even outperform autoregressive models [8].
Posterior Matching can benaturally applied to hierarchical VAEs, where the partially observed posterior is represented asq ( z | x o )
= (cid:81)
i q ( z i
| z <i , x o ) .
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 architecturedetails).
We note that due to our hardware constraints, we trained smaller models and for feweriterations than Child [8].
Inpainting results for our VDVAE models are given in Table 1.
We see thatthey achieve better precision/recall scores than the VQ-VAE models and, unlike VQ-VAE, are able toattain better PSNR than ACFlow for MNIST and C ELEB A. Figure 5 shows some example inpaintings, and additional samples are provided in the Appendix.
The fact that we see better downstreamperformance when using VDVAE than when using VQ-VAE is illustrative of Posterior Matching’sability to admit easy performance gains by simply switching to a more powerful base VAE.
[Figure] [Figure]
[Table]
5.3 Real-valued Datasets
We evaluate Posterior Matching on real-valued tabular data, specifically the benchmark UCI repositorydatasets from Papamakarios et al.
We follow the experimental setup used by Li et al.
[22]and Strauss and Oliva [35].
In these experiments, we train basic VAE models while simultaneouslylearning the partially observed posterior.
Given the flexibility that Posterior Matching affords, we usean autoregressive distribution for q θ ( z | x o ) .
Further details can be found in the Appendix.
Table 2 reports the arbitrary conditional log-likelihoods and normalized root-mean-square error(NRMSE) of imputations for our models (with features missing completely at random).
Likelihoodsare computed using an importance sampling estimate (see Appendix for details).
We primarilycompare to VAEAC as a 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 accurateimputations and higher likelihoods than VAEAC.
While our models don’t match the likelihoodsachieved 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, we may perform “partially observed” clustering, which clustersinstances based on a subset of observed features.
We consider VaDE, which uses a mixture of
Gaussians as the prior, allowing it to do unsupervised clustering by treating each Gaussian componentas one of the clusters [17].
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 F ASHION MNIST (see Appendix for experimental details).
Figure 6 shows the clustering accuracy of these models as the percentage of (randomly selected)observed features changes.
As a baseline, we train a supervised model where the labels are thecluster predictions from the VaDE model when all of the features are observed.
We see that Posterior
Matching is able to match the performance of the baseline, and even slightly outperform it forlow percentages of observed features.
Unlike the supervised approach, Posterior Matching has theadvantage 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 arbitraryconditioning.
Here, we consider the problem of greedy active feature acquisition.
We train a VAEwith a Posterior Matching network that outputs the lookahead posteriors described in Section 3.5,
[Figure]
r i
[Figure]
g A cc u r a c y
[Figure]
using the loss in Equation 7.
Note that we are also still using Posterior Matching in order to learnq θ ( z | x o ) and therefore to produce reconstructions.
Training details can be found in the Appendix.
[Figure]
We consider the MNIST dataset and compare to
EDDI as a baseline, using the authors’ publiclyavailable code.
We downscale images to 16 ×16 since EDDI has difficulty scaling to highdimensional data.
We also only evaluate on thefirst 1000 instances of the MNIST test set, asthe EDDI code was very slow when computingthe greedy acquisition policy.
EDDI also usesa particular architecture that is not compatiblewith convolutions.
Thus we train a MLP-based
VAE on flattened images in order to make a faircomparison.
However, since Posterior Matchingdoes not place any limitations on the type of
VAE being used, we also train a convolutionalversion.
For our models, we greedily select thefeature to acquire using the more expensive sampling-based approach (similar to EDDI) as well aswith the lookahead posteriors (which requires no sampling).
In both cases, imputations are computedwith an expectation over 50 latent codes, as is done for EDDI.
An example acquisition trajectory isshown in Figure 3.
Figure 7 presents the root-mean-square error, averaged across the test instances, when imputing x uwith different numbers of acquired features.
We see that our models are able to achieve lower errorthan EDDI.
We also see that acquiring based on the lookahead posteriors incurs only a minimalincrease in error compared to the sampling-based method, despite being far more efficient.
Computingthe greedy choice with our model using the sampling-based approach takes 68 ms ± 917 µ s (for asingle acquisition on CPU).
Using the lookahead posteriors, the time is only 310 µ s ± 15.3 µ s, aroughly 219x speedup .
Conclusions
We have presented an elegant and general framework, called Posterior Matching, that allows VAEsto perform arbitrary conditioning.
That is, we can take an existing VAE that only models the jointdistribution p ( x ) and train an additional model that, when combined with the VAE, is able to assessany likelihood p ( x u | x o ) for arbitrary subsets of unobserved features x u and observed features x o .
We applied this approach to a variety of VAEs for a multitude of different tasks.
We found that
Posterior Matching outperforms previous specialized VAEs for arbitrary conditioning with tabulardata and for image inpainting.
Importantly, we find that one can switch to a more powerful base
VAE 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 performclustering based on partially observed inputs and to improve the efficiency of greedy active featureacquisition by several orders of magnitude at negligible cost to performance.
With this work, we hope to make arbitrary conditioning more widely accessible.
Arbitrary conditioning no longer requires specialized methods, but can instead be achieved by applying one generalframework to common VAEs.
As advances are made in VAEs for joint density estimation, we canexpect to immediately reap
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 3.4.
(c) Did you discuss any potential negative societal impacts of your work?
[No] We do notbelieve this work has any particular negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See supplemen-
tal material.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running ex-
periments multiple times)?
[Yes] We include std.
dev.
for relevant experiments (e.g.
Table 2).
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix.4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data
you’re
using /curating?
[N/A ](e
Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A
) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]