Title: CONDITIONAL INFERENCE IN PRE-TRAINED VARIA-TIONAL AUTOENCODERS VIA CROSS-CODING

Abstract: Variational Autoencoders (VAEs) are a popular generative model, but one in which conditional inference can be challenging. If the decomposition into query and evidence variables is fixed, conditional VAEs provide an attractive solution. To support arbitrary queries, one is generally reduced to Markov Chain Monte Carlo sampling methods that can suffer from long mixing times. In this paper, we propose an idea we term cross-coding to approximate the distribution over the latent variables after conditioning on an evidence assignment to some subset of the variables. This allows generating query samples without retraining the full VAE. We experimentally evaluate three variations of cross-coding showing that (i) they can be quickly optimized for different decompositions of evidence and query and (ii) they quantitatively and qualitatively outperform Hamiltonian Monte Carlo.

1. INTRODUCTION:

Variational Autoencoders (VAEs) (Kingma & Welling, 2014) are a popular deep generative model with numerous extensions including variations for planar flow (Rezende & Mohamed, 2015) , inverse autoregressive flow (Kingma et al., 2016) , importance weighting (Burda et al., 2016) , ladder networks (Maaløe et al., 2016) , and discrete latent spaces (Rolfe, 2017) to name just a few. Unfortunately, existing methods for conditional inference in VAEs are limited. Conditional VAEs (CVAEs) (Sohn et al., 2015) allow VAE training conditioned on a fixed decomposition of evidence and query, but are computationally impractical when varying queries are made. Alternatively, Markov Chain Monte Carlo methods such as Hamiltonian Monte Carlo (HMC) (Girolami & Calderhead, 2011; Daniel Levy, 2018) are difficult to adapt to these problems, and can suffer from long mixing times as we show empirically.

To remedy the limitations of existing methods for conditional inference in VAEs, we aim to approximate the distribution over the latent variables after conditioning on an evidence assignment through a variational Bayesian methodology. In doing this, we reuse the decoder of the VAE and show that the error of the distribution over query variables is controlled by that over latent variables via a fortuitous cancellation in the KL-divergence. This avoids the computational expense of re-training the decoder as done by the CVAE approach. We term the network that generates the conditional latent distribution the cross-coder.

We experiment with two cross-coding alternatives: Gaussian variational inference via a linear transform (GVI) and Normalizing Flows (NF). We also provide some comparison to a fully connected network (FCN), which suffers from some technical and computational issues but provides a useful point of reference for experimental comparison purposes. Overall, our results show that the GVI and NF variants of cross-coding can be optimized quickly for arbitrary decompositions of query and evidence and compare favorably against a ground truth provided by rejection sampling for low latent dimensionality. For high dimensionality, we observe that HMC often fails to mix despite our systematic efforts to tune its parameters and hence demonstrates poor performance compared to cross-coding in both quantitative and qualitative evaluation.

Algorithm 1 Conditional Inference via Cross-coding. Input (a) Pre-trained VAE p(z)p θ (t|z) with p θ (t|z) based on Decoder θ (z). (Encoder ignored.) (b) Single query x (any subset of t) for which to predict y. (Rest of t.) Optimize Define q( )q ψ (z| ) with q ψ (z| ) based on XCoder ψ ( ). Find ψ to maximize C-ELBO[q ψ (Z) p θ (Z, x)] (Defined in Theorem 2). Estimate stochastic gradients by drawing random ∼ q( ) and using the reparameterization trick. Predict Draw a sample {z m } M m=1 ∼ q ψ (z) by setting z m = XCoder ψ ( m ) for m ∼ q( ). Predict p θ (y|x) ≈ 1 M M m=1 p θ (y|z m ). (Justified since the optimization phase tightened a bound (Lemma 1) on the divergence between q ψ (z)p θ (y|z)dz and p θ (y|x) )

Figure 1 : Proposed Cross-coding framework for conditional inference with variational auto-encoders. Arrow and text colors are aligned with the description in Algorithm 1, where t = (x, y).

2. BACKGROUND 2.1 VARIATIONAL AUTO-ENCODERS:

One way to define an expressive generative model p θ (t) is to introduce latent variables z. Variational Auto-Encoders (VAEs) (Kingma & Welling, 2014) model p(z) as a simple fixed Gaussian distribution. Then, for real t, p θ (t|z)is a Gaussian with the mean determined by a "decoder" network as p θ (t|z) = N (t; Decoder θ (z), σ 2 I).

(1)

If t is binary, a product of independent Bernoulli's is parameterized by a sigmoidally transformed decoder. If the decoder network has high capacity, the marginal distribution p θ (t) can represent a wide range of distributions. In principle, one might wish to train such a model by (regularized) maximum likelihood. Unfortunately, the marginal p θ (t) is intractable. However, a classic idea (Saul et al., 1996) is to use variational inference to lower-bound it. For any distributions p θ and q φ ,

EQUATION

Since the KL-divergence is non-negative, the "evidence lower bound" (ELBO) lower bounds log p θ (t). Thus, as a surrogate to maximizing the likelihood over θ one can maximize the ELBO over to θ and φ simultaneously.

VAEs define q φ (z) as the marginal of q(t)q φ (z|t) where q(t) is simple and fixed and q φ (z|t) = N (z; Encoder φ (t)) is a Gaussian with a mean and covariance both determined by an "encoder" network.

2.2. THE CONDITIONAL INFERENCE PROBLEM:

In this paper, we assume a VAE has been pre-trained. Then, at test time, some arbitrary subset x of t is observed as evidence, and the goal is to predict the distribution of the non-observed y where the decomposition t = (x, y) is unpredictable. If this decomposition of t into evidence and query variables is fixed and known ahead of time, a natural solution is to train an explicit conditional model, the approach taken by Conditional Variational Autoencoders (Sohn et al., 2015) . We focus on supporting arbitrary queries, where training a conditional model for each possible decomposition t = (x, y) is infeasible.

3. CONDITIONAL INFERENCE ON VARIATIONAL AUTO-ENCODERS:

We now turn to the details of conditional inference. We assume we have pretrained a VAE and now wish to approximate the distribution p θ (y|x) where x is some new "test" input not known at VAE training time. Unfortunately, exact inference is difficult, since computing this probability exactly would require marginalizing out z.

3.1. EXPLOITING FACTORIZATION IN THE OUTPUT:

One helpful property comes from the fact that in a VAE, the conditional distribution over the output (Eq. 1) has a diagonal covariance, which leads to the following decomposition:

Observation 1 The distribution of a VAE can be factorized as p θ (x, y, z) = p(z)p θ (x|z)p θ (y|z).

Since x and y are conditionally independent given z, the conditional of y given x can be written as

EQUATION

Here, p θ (y|z) can easily be evaluated or simulated. However p θ (z|x) is much more difficult to work with since it involves "inverting" the decoder. This factorization can also be exploited by Markov chain Monte Carlo methods (MCMC), such as Hamiltonian Monte Carlo (HMC) (Girolami & Calderhead, 2011; Daniel Levy, 2018) . In this case, it allows the Markov chain to be defined over z alone, rather than z and y together. That is, one can use MCMC to attempt sampling from p θ (z|x), and then draw exact samples from p θ (y|z) just by evaluating the decoder network at each of the samples of z. The experiments using MCMC in Section 4 use this strategy.

3.2. VARIATIONAL INFERENCE BOUNDS:

The basic idea of variational inference (VI) is to posit some distribution q ψ , and optimize ψ to make it match the target distribution as closely as possible. So, in principle, the goal of VI would be to minimize KL[q ψ (Y) p θ (Y|x)]. For an arbitrary distribution q ψ this divergence would be difficult to work with due to the need to marginalize out z in p θ as in Eq. 3.

However, if q ψ is chosen carefully, then the above divergence can be upper-bounded by one defined directly over Z. Specifically, we will choose q ψ so that the dependence of y on z under q ψ is the same as under p θ (both determined by the "decoder").

Lemma 1. Suppose we choose q ψ (z, y) = q ψ (z)p θ (y|z).

Then

EQUATION

This is proven in the Appendix. The result follows from using the chain rule of KL-divergence (Cover & Thomas, 2006) to bound the divergence over y by the divergence jointly over y and z. Then the common factors in q ψ and p θ mean this simplifies into a divergence over z alone.

Given this Lemma, it makes sense to seek a distribution q ψ such that the divergence on the right-hand side of Eq. 4 is as low as possible. To minimize this divergence, consider the decomposition

EQUATION

which is analogous to Eq. 2. Here, we call the first term the "conditional ELBO" (C-ELBO) to reflect that maximizing it is equivalent to minimizing an upper bound on KL[q ψ (Y) p θ (Y|x)].

3.3. INFERENCE VIA CROSS-CODING:

The previous section says that we should seek a distribution q ψ to approximate p θ (z|x). Although the latent distribution p(z) may be simple, the conditional distribution p θ (z|x) is typically complex and often multimodal (cf. Fig. 3 ).

To define a variational distribution satisfying the conditions of Lemma 1, we propose to draw from some fixed base density q( ) and then use a network with parameters ψ to map to the latent space z so that the marginal q ψ (z) is expressive. The conditional of y given z is exactly as in p. The full variational distribution is therefore

EQUATION

where δ is a multivariate delta function. We call this network a "Cross-coder" to emphasize that the parameters ψ are fit so that q ψ (Z) matches p θ (Z|x), and so that z, when "decoded" using θ, will predict y given x. Theorem 2. If q ψ is as defined in Eq. 6 and XCoder ψ ( ) is one-to-one for all ψ, the C-ELBO from Eq. 5 becomes

C-ELBO[q ψ (Z) p θ (Z, x)] = E q( ) [log p θ (XCoder ψ ( ), x) + log |∇XCoder ψ ( )|] + H[q( )],

where H[q( )] is the (fixed) entropy of q( ), ∇ is the Jacobian with respect to , and | • | is the determinant.

Informally, this result can be proven as follows: the C-ELBO was defined on z alone, while our definition of q ψ in Eq. 6 also involves y and . Marginalizing out y is trivial. Then, since z and are deterministically related under q ψ one can change variables to convert the expectation over z to one over , leaving the log-determinant Jacobian as an artifact of the entropy of q ψ (z).

This objective is related to the "triple ELBO" used by Vedantam et al. (2017) for a situation with a small number of fixed decompositions of t into (x, y). Algorithmically, the approaches are quite different since they pre-train a single network for each subset of t, which can be used for any x with that pattern, and a futher product of experts approximation is used for novel missing features at test time. We assume arbitrary queries and so pre-training is inapplicable and novel missing features pose no issue. Still, our bounding justification may provide additional insight for their approach.

3.4. CROSS-CODERS:

We explore the following two candidiate Cross-Coders.

Gaussian Variational Inference (GVI)::

The GVI XCoder ψ linearly warps a spherical Gaussian over into an arbitrary Gaussian z:

EQUATION

where ψ = (W, b) for a square matrix W and a mean vector b. While projected gradient descent can be used to maintain invertibility of W , we did not encounter issues with non-invertible W requiring projection during our experiments.

Normalizing Flows (NF): A normalizing flow (Rezende & Mohamed, 2015) projects a probability density through a sequence of easy computable and invertible mappings. By stacking multiple mappings, the transformation can be complex. We use the special structured network called Planar Normalizing Flow:

EQUATION

for all i, where h 0 = , i is the layer id, w and u are vectors, and the output dimension is exactly same with the input dimension. Using • for function composition, the XCoder ψ is given as

XCoder ψ ( ) = f k • f k-1 • • • f 1 ( ),

where

log |∇XCoder ψ ( )| = k i=1 log |∇f i |. (9)

The bound in Theorem 2 requires that XCoder ψ is invertible. Nevertheless, we find Fully Connected Networks (FCNs) useful for comparison in low-dimensional visualizations. Here, the Jacobian must be calculated using separate gradient calls for each ouput variable, and the lack of invertibility prevents the C-ELBO bound from being correct.

We summarize our approach in Algorithm 1. In brief, we define a variational distribution q ψ ( , z) = q( )q ψ (z| ) and optimize ψ so that q ψ (z) is close to p θ (z|x). The variational distribution includes a "CrossCoder" as q ψ (z| ) = δ(z -XCoder ψ ( )). The algorithm uses stochastic gradient decent on the C-ELBO with gradients estimated using Monte Carlo samples of and the reparameterization trick (Kingma & Welling, 2014; Titsias & Lázaro-Gredilla, 2014; Rezende et al., 2014) . After inference, the original VAE distribution q(y|z) = p θ (y|z) gives samples over the query variables.

4. EXPERIMENTS:

Having defined our cross-coding methodology for conditional inference with pre-trained VAEs, we now proceed to empirically evaluate our three previously defined XCoder instantiations and compare them with (Markov chain) Monte Carlo (MCMC) sampling approaches on three different pre-trained VAEs. Below we discuss our datasets and methodology followed by our experimental results.

4.1. DATASETS AND PRE-TRAINED VAES:

MNIST is the well-known benchmark handwritten digit dataset (LeCun & Cortes, 2010) . We use a pre-trained VAE with a fully connected encoder and decoder each with one hidden layer of 64 ReLU units, a final sigmoid layer with Bernoulli likelihood, and 2 latent dimensions for z. 1 The VAE has been trained on 60,000 black and white binary thresholded images of size 28 × 28. The limitation to 2 dimensions allows us to visualize the conditional latent distribution of all methods and compare to the ground truth through a fine-grained discretization of z.

Anime is a dataset of animated character faces (Jin et al., 2017) . We use a pre-trained VAE with convolutional encoder and deconvolutional decoder, each with 4 layers. The decoder contains respective channel sizes (256, 128, 32, 3) each using 5 × 5 filters of stride 2 and ReLU activations followed by batch norm layers. The VAE has a final tanh layer with Gaussian likelihood, and 64 latent dimensions for z. 2 The VAE has been trained on 20000 images encoded in RGB of size 64 × 64 × 3.

CelebA dataset (Liu et al., 2015) is a benchmark dataset of images of celebrity faces. We use a pre-trained VAE with a structure that exactly matches the Anime VAE provided above, except that it uses 100 latent dimensions for z. 3 The VAE has been trained on 200,000 images encoded in RGB of size 64 × 64 × 3.

4.2. METHODS COMPARED:

For sampling approaches, we evaluate rejection sampling (RS), which is only feasible for our MNIST VAE with a 2-dimensional latent embedding for z. We also compare to the MCMC method of Hamiltonian Monte Carlo (HMC) (Girolami & Calderhead, 2011; Daniel Levy, 2018) . Both sampling methods exploit the VAE decomposition and sampling methodology described in Section 3.1.

We went to great effort to tune the parameters of HMC. For MNIST, with low dimensions, this was generally feasible, with a few exceptions as noted in Figure 4 (b). For the high-dimensional latent space of the Anime and CelebA VAEs, finding parameters to achieve good mixing was often impossible, leading to poor performance. Section 6.4 of the Appendix discusses this in detail.

For the cross-coding methods, we use the three XCoder variants described in Section 3.3: Gaussian Variational Inference (GVI), Planar Normalizing Flow (NF), and a Fully Connected Neural Network (FCN). By definition, the latent dimensionality of must match the latent dimensionality of z for each pre-trained VAE. Given evidence as described in the experiments, all cross-coders were trained as described in Algorithm 1. We could not train the FCN XCoder for conditional inference in Anime and

CelebA due to the infeasibility of computing the Jacobian for the respective latent dimensionalities of these two VAEs.

In preliminary experiments, we considered the alternating sampling approach suggested by (Rezende et al., 2014, Appendix F) , but found it to perform very poorly when the evidence is ambiguous. We provide a thorough analysis of this in Section 6.3 of the Appendix comparing results on MNIST with various fractions of the input taken as evidence. In summary, Rezende's alternation method produces reasonable results when a large fraction of pixels are observed, so the posterior is highly concentrated. When less than around 40% are observed, however, performance rapidly degrades.

4.3. EVALUATION METHODOLOGY:

We experiment with a variety of evidence sets to demonstrate the efficiency and flexibility of our cross-coding methodology for arbitrary conditional inference queries in pre-trained VAEs. All crosscoding optimization and inference takes (typically well) under 32 seconds per evidence set for all experiments running on an Intel Xeon E5-1620 v4 CPU with 4 cores, 16Gb of RAM, and an NVIDIA GTX1080 GPU. A detailed running time comparison is provided in Section 6.5 of the Appendix.

Qualitatively, we visually examine the 2D latent distribution of z conditioned on the evidence for the special case of MNIST, which has low enough latent dimensionality to enable us to obtain ground truth through discretization. For all experiments, we qualitatively assess sampled query images generated for each evidence set to assess both the coverage of the distribution and the quality of match between the query samples and the evidence, which is fixed in the displayed images.

Quantitatively, we evaluate the performance of the proposed framework and candidate inference methods through the following two metrics.

C-ELBO::

As a comparative measure of inference quality for each of the XCoder methods, we provide pairwise scatterplots of the C-ELBO as defined in 5 for a variety of different evidence sets Query Marginal Likelihood: For each conditional inference evaluation, we randomly select an image and then a subset of that image as evidence x and the remaining pixels y as the ground truth query assignment. Given this, we can evaluate the marginal likelihood of the query y as follows: For conditional inference in MNIST, we begin with Figure 2 , which shows one example of conditional inference in the pre-trained MNIST model using the different inference methods. While the original image used to generate the evidence represents the digit 3, the evidence is very sparse allowing the plausible generation of other digits. It is easy to see that most of the methods can handle this simple conditional inference, with only GVI producing some samples that do not match the evidence well in this VAE with 2 latent dimensions.

To provide additional insight into Figure 2 , we now turn to Figure 3 , where we visually compare the true conditional latent distribution p(z|x) (leftmost) with the corresponding distributions of each of the inference methods. At a first glance, we note that the true distribution is both multimodal and non-Gaussian. We see that GVI covers some mass not present in the true distribution that explains its Likelihood for the same 50 evidence sets from (a), with each likelihood expectation generated from 500 samples. For both metrics, higher is better.

We now proceed to a quantitative comparison of performance on MNIST over 50 randomly generated queries. In Figure 4 (a), we present a pairwise comparison of the performance of each XCoder method on 50 randomly generated evidence sets. Noting that higher is better, we observe that FCN and NF perform comparably and generally outperform GVI. In Figure 4 (b), we examine the Query Marginal Likelihood distribution for the same 50 evidence sets from (a) with each likelihood expectation generated from 500 samples. Again, noting that higher is better, here we see that RS slightly edges out all other methods with all XCoders generally performing comparably. HMC performs worst here, where we remark that inadequate coverage of the latent z due to poor mixing properties leads to over-concentration on y leading to a long tail in a few cases with poor coverage. We will see that these issues with HMC mixing become much more pronounced as we move to experiments in VAEs with higher latent dimensionality in the next section.

4.5. CONDITIONAL INFERENCE ON ANIME AND CELEBA:

Now we proceed to our larger VAEs for Anime and CelebA with respective latent dimensionality of 64 and 100 that allow us to work with larger and more visually complex RGB images. In these cases, FCN could not be applied due to the infeasibilty of computing the Jacobian and RS is also infeasible for such high dimensionality. Hence, we only compare the two XCoders GVI and NF with HMC.

We now continue to a qualitative and quantitative performance analysis of conditional inference for the Anime and CelebA VAEs. Qualitatively, in Figure 5 for Anime, we see that inference for both the NF XCoder and HMC show little identifiable variation and seem to have collapsed into a single latent mode. In contrast, GVI appears to show better coverage, generating a wide range of faces that generally match very well with the superimposed evidence. For Figure 6 , HMC still performs poorly, but NF appears to perform much better, with both XCoders GVI and NF generating a wide range of faces that match the superimposed evidence, with perhaps slightly more face diversity for GVI.

Quantitatively, Figure 7 strongly reflects the qualitative visual observations above. In short for the XCoders, GVI solidly outperforms NF on the C-ELBO comparison. For all methods evaluated on Query Marginal Likelihood, GVI outperforms both NF and HMC on Anime, while for CelebA GVI performs comparably to (if not slightly worse) than NF, with both solidly outperforming HMC.

5. CONCLUSION:

We introduced Cross-coding, a novel variational inference method for conditional queries in pretrained VAEs that does not require retraining the decoder. Using three VAEs pre-trained on different datasets, we demonstrated that the Gaussian Variational Inference (GVI) and Normalizing Flows (NF) cross-coders generally outperform Hamiltonian Monte Carlo both qualitatively and quantitively, thus providing a novel and efficient tool for conditional inference in VAEs with arbitrary queries.

6. APPENDIX:

6.1 PROOFS Proof Of Lemma 1. To show this, we first note that the joint divergence over Y and Z is equivalent to one over Z only.

KL[q ψ (Y, Z) p θ (Y, Z|x)] = KL[q ψ (Z) p θ (Z|x)] + KL[q ψ (Y|Z) p θ (Y|Z, x)]

by the chain rule of KL-divergence = KL[q ψ (Z) p θ (Z|x)] + KL[q ψ (Y|Z) p θ (Y|Z)]

since Y⊥X|Z in both q ψ and p θ = KL[q ψ (Z) p θ (Z|x)] since q ψ (y|z) = p θ (y|z)

Then, the result follows just from observing (again by the chain rule of KL-divergence) that

KL[q ψ (Y) p θ (Y|x)] ≤ KL[q ψ (Y, Z) p θ (Y, Z|x)].

Proof of Theorem 2 . For the purpose of this proof, use c ψ to denote CrossEncoder ψ . Firstly, note that the marginal density of q ψ (z|x) is (via the standard formula for a change of variables (Kaplan, 1952) ) In this experiment, we do not use a VAE, but instead simply model a complex latent 2D multimodal distribution over z as a Gaussian mixture model to evaluate the ability of each conditional inference method to accurately draw samples from this complex distribution. In general, Figure 8 shows that while the XCoders NF and FCN work well here, GVI (by definition) cannot model this multimodal distribution and HMC draws too few samples from the disconnected mode compared to the true sample distribution, indicating slight failure to mix well. The first row shows the evidence observed, and the following rows show the mean of generated samples from the different algorithms. We note that with very high evidence, the posterior becomes extremely concentrated, meaning the rejection rates for rejection sampling become impractical. (b) The mean squared error between query variables of the original image and the generated samples of different algorithms. The results and standard deviations at each observation percentage come from 50 independent randomly selected queries.

q ψ (z = c ψ ( )) |∇c ψ ( )| = q( ) Thus, we can write C-ELBO[q ψ (Z) p θ (Z, x)] = E q ψ (Z) log p θ (Z, x) q ψ (Z) = E q( ) log p θ (Z = c ψ ( ), x) q ψ (Z = c ψ ( )) = E q( ) log p θ (Z = c ψ ( ), x) q( )/ |∇c ψ ( )| = E q( ) [log p θ (c ψ ( ), x) + log |∇c ψ ( )|] + H q [ ].

6.2. PRELIMINARY CHECK OF INFERENCE METHODS:



6.3. COMPARISON TO REZENDE ALTERNATION:

We compare to the alternating sampling approach of Rezende et al. ( 2014) (Appendix Section F) which is essentially an approximation of block Gibbs sampling. We call it the "Rezende method" in the following. This method does not asymptotically sample from the conditional distribution since the step sampling the latent variables given the query variables are approximated using the encoder.

Figure 9 (a) shows one experiment comparing all candidate algorithms including the Rezende method. We noticed that it fails to generate images that match the evidence when less than 40% of pixels are observed as evidence, while it makes reasonable predictions when the observation rate is higher. Figure 9 (b) shows this result is consistent over 50 randomly selected queries. While tuning HMC in lower dimensions was generally feasible for MNIST with a few exceptions noted in previous discussion of Figure 4 (b), we observed that HMC becomes very difficult to tune in the Anime and CelebA VAEs with higher latent dimensionality. To illustrate these HMC tuning difficulties, we present a summary of our systematic efforts to tune HMC on Anime and CelebA in Figure 10 with boxplots of the acceptance rate distribution of HMC for 30 Markov Chains vs different on (a) Anime and (b) CelebA. We ran each Markov chain for 10,000 burn-in samples with 10 leapfrog steps per iteration; we tried 3 different standard leapfrog step settings of {5, 10, 30}, finding that 10 leapfrog steps provided the best performance across a range of and hence chosen for Figure 10 .

In short, Figure 10 shows that only a very narrow band of lead to a reasonable acceptance rate for good mixing properties of HMC. Even then, however, the distribution of acceptance rates for any particular Markov Chain for a good is still highly unpredictable as given by the quartile ranges of the boxplot. In summary, we found that despite our systematic efforts to tune HMC for higher dimensional problems, it was difficult to achieve a good mixing rate and overall contributes to the generally poor performance observed for HMC on Anime and CelebA that we discuss next.

6.5. COMPARISON OF RUNNING TIME:

The running time of conditional inference with XCoding varies with the complexity of XCoders, the optimization algorithm used, and the complexity of the pretrained Decoder. We found that L-BFGS (Liu & Nocedal, 1989) consistently converged fastest and with the best results in comparison to SGD, Adam, Adadelta, and RMSProp. Table 1 , which follows, shows the computation time for each of the three candidate XCoders (FCN is only applicable to MNIST) as well as HMC and Rejection Sampling (RS is only applicable for MNIST). To assess the quality of the pre-trained VAE models, we show 100 samples from each in Figure 11 . 

:

https://github.com/kvfrans/variational-autoencoder

URL for pre-trained VAE suppressed for anonymous review.

https://github.com/yzwxx/vae-celebA

Figures and Tables:

Figure 2: One conditional inference example for MNIST. In all plots, the evidence subset has white replaced with orange and black replaced with blue. (a) The original digit t, the subset selected for evidence x, and the remaining ground truth query y. (b-f) Nine sample queries from method.

Figure 3: p(z|x) for the MNIST example in Figure 2. The contour plot (left) shows the true distribution. The remaining plots show samples from each method overlaid on the true distribution.

Figure 4: (a) Pairwise C-ELBO comparison of different XCoder methods evaluated over the 50 randomly generated evidence sets for MNIST. (b) Violin (distribution) plots of the Query MarginalLikelihood for the same 50 evidence sets from (a), with each likelihood expectation generated from 500 samples. For both metrics, higher is better.

Figure 5: One conditional inference example for Anime. (a) The original image t, the subset selected for evidence x, and the remaining ground truth query y. (b-d) 25 sample queries from each method with the evidence superimposed on each image. (c,d) NF and HMC demonstrate poor coverage. (x,y) y

Figure 8: Comparison of different inference methods on modeling a Gaussian mixture model distribution. The true distribution samples are directly sampled from a Gaussian mixture model. Maximum mean discrepancy (MMD) values given in the plot titles are generated relative to the true sample distribution.

Figure 9: Comparison of different conditional inference methods include the Rezende method on the MNIST dataset. (a) Shows one intuitive example. The first row shows the evidence observed, and the following rows show the mean of generated samples from the different algorithms. We note that with very high evidence, the posterior becomes extremely concentrated, meaning the rejection rates for rejection sampling become impractical. (b) The mean squared error between query variables of the original image and the generated samples of different algorithms. The results and standard deviations at each observation percentage come from 50 independent randomly selected queries.

Figure 10: Boxplots of acceptance rate distribution of HMC for 30 Markov Chains vs different on (a) Anime and (b) CelebA. Each Markov chain ran for 10,000 burn-in samples with 10 leapfrog steps per iteration.

Figure 11: Samples from each of the pre-trained VAE models.

Average running Time (in seconds) of experiments. We use L-BFGS for XCoders in this table. For HMC, we predefine the burn-in (optimization) iterations to be 1000 for all datasets. For all methods, the sample size is 500.

Table content: <table><tr><td>Period</td><td/><td/><td>MNIST</td><td/><td/><td/><td>Amine</td><td/><td/><td>CelebA</td><td/></tr><tr><td>-</td><td>GVI</td><td>NF</td><td>FCN</td><td colspan="2">HMC RS</td><td>GVI</td><td>NF</td><td>HMC</td><td>GVI</td><td>NF</td><td>HMC</td></tr><tr><td>Optimization</td><td>0.36</td><td>2.79</td><td>5.26</td><td>34.92</td><td>-</td><td>2.52</td><td>4.22</td><td>81.3</td><td>31.45</td><td>9.95</td><td>224.5</td></tr></table>

Peer review (iclr2019):
 (apologies for this belated review) Summary The authors consider the task of imputing missing data using variational auto-encoders. To do so, they assume a fixed pre-trained generative model, perform variational inference to infer a posterior on latent variables given a partial image, and then use this approximate posterior to predict missing pixels. They compare a variety of parameterizations of the variational distribution to HMC inference, and evaluate on MNIST, Celeb-A and the Anime data. Comments There are many things about this paper that I don’t understand. My main concern is that I fail to follow why the authors are interested in this task. In what settings would we be interested in performing non-autoencoding variational inference in order to impute missing data? Moreover, in cases where are interested in performing such imputations, what would we like to use the results for? This paper seems like a nice demo, but I’m not entirely convinced I see a compelling application. My second concern is about the baselines that are considered. If I were interested in carrying out this inference task, my inclination would not be to run an HMC chain to convergence, but instead to do something like annealed importance sampling (AIS), where at each step I run an iteration of HMC on a large batch of samples on a sequence of target densities that interpolate between the prior and full joint p(x, Z). If computational cost is a concern, I imagine this would not be more expensive than training a density estimator. Moreover, whereas HMC is generally not known to be a good method for estimating marginal likelihoods, AIS methods generally perform much better. Finally I find the language used in this paper confusing. Cross-coding seems a misnomer for the technique that the authors propose. Isn’t this simply a form of variational inference in which qψ(Z) approximates pθ(Ζ | x)? The term “-coding” suggests that we somehow define an encoder that accepts the query as input. Moreover, isn’t the XCoder network just a neural density estimator? Finally, Lemma 1 seems like a really roundabout way of deriving a lower bound. The authors could instead just write: log p(x) >= E_q(Z,Y)[log p(x, Y, Z) - log q(Z, Υ)] = E_q(Z,Y)[log p(x | Ζ) + log p(Y | Z) + log p(Z) - log q(Z) - log p(Υ | Z)] = E_q(Z,Y)[log p(x | Ζ) + log p(Z) - log q(Z)] = E_q(Z)[log p(x, Ζ) - log q(Z)] This avoids confusing terminology such as cross-coding, and shows that what the authors are doing is in fact just variational inference. Am I missing something here? I am also confused about how the comparison to HMC is set up. If you’re training qψ(Z), then you presumably need generate a certain number samples at training time. Shouldn’t you add this number of samples number of samples you generate in HMC, in order to get a more apples to apples comparison in terms of the amount of computation performed? As it stands, it is hard to evaluate whether these methods are given a similar number of samples. Finally, I am not quite sure what to make of the experimental evaluation. We see some scatter plots on MNIST with a 2D latent space, and some faces of celebrities in which there is arguably some sample diversity, although most of this diversity arises in blurry looking hairstyles. However, since the authors condition on the eyes, rather than, say, the nose or mouth, it is hard to know how good a job the network is doing at generalizing to multiple plausible faces. Overall, I find it difficult to judge the merit of this paper. Is this task in fact hard? Is it useful? Are the results good? Maybe the authors can give us some additional guidance on these questions. Questions - I’m a bit worried that not all the samples that we see in Figure 6 may have equally high probability under the posterior. Could the authors compute and report importance weights? W = p(x, Z) / q(Z) - Could the authors say something about the effective sample size that we obtain when using the learned distribution q(Z) as a proposal? ESS = (Σ_k w^k)^2 / (Σ_k (w^k)^2) - Should it be the case that the ESS is low, and the weights are high variance, could the authors generate a sufficient number of samples to ensure the the ESS = 25 (i.e. the number of images in the figure) and then show the 25 highest-weight samples (or resample 25 images with probability proportional to their weight)? Minor - Equation (3): There’s an extra p_θ in the first integral - In the proof in Appendix 6.1 KL[ qψ(Z) ‖ pθ(Z | x) ] + KL[ qψ(Y | Z) ‖ pθ(Y | Z, x)] it would be clearer to explicitly denote the expectation over qψ(Z) KL[ qψ(Z) ‖ pθ(Z | x) ] + E_qψ(Z)[ KL[ qψ(Y | Z) ‖ pθ(Y | Z, x)] ] (I had to google lecture notes to find out that this expectation is sometimes implicit, which as far as I know is not very standard).