Conditional Inference in Pre-trained Variational Autoencoders via Cross-coding

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 { zm } Mm=1 ∼ qψ ( z ) by setting zm = XCoderψ ( m ) for m ∼ q ( ) . Predict pθ ( y|x ) ≈ 1M ∑M m=1 pθ ( y|zm ) . ( Justified since the optimization phase tightened a bound ( Lemma 1 ) on the divergence between ∫ qψ ( z ) pθ ( y|z ) dz and pθ ( y|x ) ) 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 ) , σ2I ) . ( 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φ , log pθ ( t ) = log ∫ z pθ ( t , z ) dz = Eqφ ( Z ) log pθ ( Z , t ) qφ ( Z ) ︸ ︷︷ ︸ ELBO [ qφ ( Z ) ‖pθ ( Z , t ) ] +KL [ qφ ( Z ) ||pθ ( Z|t ) ] . ( 2 ) 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 pθ ( y|x ) = ∫ z pθ ( z , y|x ) pθdz = ∫ z pθ ( z|x ) pθ ( y|z ) dz . ( 3 ) 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 KL [ qψ ( Y ) ‖pθ ( Y|x ) ] ≤ KL [ qψ ( Z ) ‖pθ ( Z|x ) ] . ( 4 ) 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 log pθ ( x ) = Eqφ ( Z ) log pθ ( Z , x ) qψ ( Z ) ︸ ︷︷ ︸ C-ELBO [ qψ ( Z ) ‖pθ ( Z , x ) ] +KL [ qψ ( Z ) ||pθ ( Z|x ) ] , ( 5 ) 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 qψ ( , z , y ) = q ( ) qψ ( z| ) pθ ( y|z ) with qψ ( z| ) = δ ( z−XCoderψ ( ) ) , ( 6 ) 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 ) ] = Eq ( ) [ 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 .

Summary: Given a pre-trained VAE (e.g. over images), this paper is about inferring the distribution over missing variables (e.g. given half the pixels, what is a plausible completion?). The paper describes an approach based on variational inference with normalizing flows: given observed variables, the posterior over the VAE's latents is inferred (variationally) and plausible completions for missing variables are sampled from the VAE decoder.