Variational Autoencoder with Arbitrary Conditioning

1 INTRODUCTION . In past years , a number of generative probabilistic models based on neural networks have been proposed . The most popular approaches include variational autoencoder ( Kingma & Welling , 2013 ) ( VAE ) and generative adversarial net ( Goodfellow et al. , 2014 ) ( GANs ) . They learn a distribution over objects p ( x ) and allow sampling from this distribution . In many cases , we are interested in learning a conditional distribution p ( x|y ) . For instance , if x is an image of a face , y could be the characteristics describing the face ( are glasses present or not ; length of hair , etc . ) Conditional variational autoencoder ( Sohn et al. , 2015 ) and conditional generative adversarial nets ( Mirza & Osindero , 2014 ) are popular methods for this problem . In this paper , we consider the problem of learning all conditional distributions of the form p ( xI |xU\I ) , where U is the set of all features and I is its arbitrary subset . This problem generalizes both learning the joint distribution p ( x ) and learning the conditional distribution p ( x|y ) . To tackle this problem , we propose a Variational Autoencoder with Arbitrary Conditioning ( VAEAC ) model . It is a latent variable model similar to VAE , but allows conditioning on an arbitrary subset of the features . The conditioning features affect the prior on the latent Gaussian variables which are used to generate unobserved features . The model is trained using stochastic gradient variational Bayes ( Kingma & Welling , 2013 ) . We consider two most natural applications of the proposed model . The first one is feature imputation where the goal is to restore the missing features given the observed ones . The imputed values may be valuable by themselves or may improve the performance of other machine learning algorithms which process the dataset . Another application is image inpainting in which the goal is to fill in an unobserved part of an image with an artificial content in a realistic way . This can be used for removing unnecessary objects from the images or , vice versa , for complementing the partially closed or corrupted object . ∗Author is in DeepMind now . The experimental evaluation shows that the proposed model successfully samples from the conditional distributions . The distribution over samples is close to the true conditional distribution . This property is very important when the true distribution has several modes . The model is shown to be effective in feature imputation problem which helps to increase the quality of subsequent discriminative models on different problems from UCI datasets collection ( Lichman , 2013 ) . We demonstrate that model can generate diverse and realistic image inpaintings on MNIST ( LeCun et al. , 1998 ) , Omniglot ( Lake et al. , 2015 ) and CelebA ( Liu et al. , 2015 ) datasets , and works even better than the current state of the art inpainting techniques in terms of peak signal to noise ratio ( PSNR ) . The paper is organized as follows . In section 2 we review the related works . In section 3 we briefly describe variational autoencoders and conditional variational autoencoders . In section 4 we define the problem , describe the VAEAC model and its training procedure . In section 5 we evaluate VAEAC . Section 6 concludes the paper . Appendix contains additional explanations , theoretical analysis , and experiments for VAEAC . 2 RELATED WORK . Universal Marginalizer ( Douglas et al. , 2017 ) is a model based on a feed-forward neural network which approximates marginals of unobserved features conditioned on observable values . A related idea of an autoregressive model of joint probability was previously proposed in Germain et al . ( 2015 ) and Uria et al . ( 2016 ) . The description of the model and comparison with VAEAC are available in section 5.3 . Yoon et al . ( 2018 ) propose a GANs-based model called GAIN which solves the same problem as VAEAC . In contrast to VAEAC , GAIN does not use unobserved data during training , which makes it easier to apply to the missing features imputation problem . Nevertheless , it turns into a disadvantage when the fully-observed training data is available but the missingness rate at the testing stage is high . For example , in inpainting setting GAIN can not learn the conditional distribution over MNIST digits given one horizontal line of the image while VAEAC can ( see appendix D.4 ) . The comparison of VAEAC and GAIN on the missing feature imputation problem is given in section 5.1 and appendix D.2 . Rezende et al . ( 2014 ) [ Appendix F ] , Sohl-Dickstein et al . ( 2015 ) , Goyal et al . ( 2017 ) , and Bordes et al . ( 2017 ) propose to fill missing data with noise and run Markov chain with a learned transition operator . The stationary distribution of such chains approximates the true conditional distribution of the unobserved features . Bachman & Precup ( 2015 ) consider missing feature imputation in terms of Markov decision process and propose LSTM-based sequential decision making model to solve it . Nevertheless , these methods are computationally expensive at the test time and require fully-observed training data . Image inpainting is a classic computer vision problem . Most of the earlier methods rely on local and texture information or hand-crafted problem-specific features ( Bertalmio et al. , 2000 ) . In past years multiple neural network based approaches have been proposed . Pathak et al . ( 2016 ) , Yeh et al . ( 2016 ) and Yang et al . ( 2017 ) use different kinds and combinations of adversarial , reconstruction , texture and other losses . Li et al . ( 2017 ) focuses on face inpainting and uses two adversarial losses and one semantic parsing loss to train the generative model . In Yeh et al . ( 2017 ) GANs are first trained on the whole training dataset . The inpainting is an optimization procedure that finds the latent variables that explain the observed features best . Then , the obtained latents are passed through the generative model to restore the unobserved portion of the image . We can say that VAEAC is a similar model which uses prior network to find a proper latents instead of solving the optimization problem . All described methods aim to produce a single realistic inpainting , while VAEAC is capable of sampling diverse inpaintings . Additionally , Yeh et al . ( 2016 ) , Yang et al . ( 2017 ) and Yeh et al . ( 2017 ) have high testtime computational complexity of inpainting , because they require an optimization problem to be solved . On the other hand , VAEAC is a “ single-shot ” method with a low computational cost . 3 BACKGROUND . 3.1 VARIATIONAL AUTOENCODER . Variational autoencoder ( Kingma & Welling , 2013 ) ( VAE ) is a directed generative model with latent variables . The generative process in variational autoencoder is as follows : first , a latent variable z is generated from the prior distribution p ( z ) , and then the data x is generated from the generative distribution pθ ( x|z ) , where θ are the generative model ’ s parameters . This process induces the distribution pθ ( x ) = Ep ( z ) pθ ( x|z ) . The distribution pθ ( x|z ) is modeled by a neural network with parameters θ. p ( z ) is a standard Gaussian distribution . The parameters θ are tuned by maximizing the likelihood of the training data points { xi } Ni=1 from the true data distribution pd ( x ) . In general , this optimization problem is challenging due to intractable posterior inference . However , a variational lower bound can be optimized efficiently using backpropagation and stochastic gradient descent : log pθ ( x ) = Eqφ ( z|x ) log pθ ( x , z ) qφ ( z|x ) +DKL ( qφ ( z|x ) ‖p ( z|x , θ ) ) ≥ Eqφ ( z|x ) log pθ ( x|z ) −DKL ( qφ ( z|x ) ‖p ( z ) ) = LV AE ( x ; θ , φ ) ( 1 ) Here qφ ( z|x ) is a proposal distribution parameterized by neural network with parameters φ that approximates the posterior p ( z|x , θ ) . Usually this distribution is Gaussian with a diagonal covariance matrix . The closer qφ ( z|x ) to p ( z|x , θ ) , the tighter variational lower bound LV AE ( θ , φ ) . To compute the gradient of the variational lower bound with respect to φ , reparameterization trick is used : z = µφ ( x ) + εσφ ( x ) where ε ∼ N ( 0 , I ) and µφ and σφ are deterministic functions parameterized by neural networks . So the gradient can be estimated using Monte-Carlo method for the first term and computing the second term analytically : ∂LV AE ( x ; θ , φ ) ∂φ = Eε∼N ( 0 , I ) ∂ ∂φ log pθ ( x|µφ ( x ) + εσφ ( x ) ) − ∂ ∂φ DKL ( qφ ( z|x ) ‖p ( z ) ) ( 2 ) So LV AE ( θ , φ ) can be optimized using stochastic gradient ascent with respect to φ and θ . 3.2 CONDITIONAL VARIATIONAL AUTOENCODER . Conditional variational autoencoder ( Sohn et al. , 2015 ) ( CVAE ) approximates the conditional distribution pd ( x|y ) . It outperforms deterministic models when the distribution pd ( x|y ) is multi-modal ( diverse xs are probable for the given y ) . For example , assume that x is a real-valued image . Then , a deterministic regression model with mean squared error loss would predict the average blurry value for x . On the other hand , CVAE learns the distribution of x , from which one can sample diverse and realistic objects . Variational lower bound for CVAE can be derived similarly to VAE by conditioning all considered distributions on y : LCV AE ( x , y ; θ , ψ , φ ) = Eqφ ( z|x , y ) log pθ ( x|z , y ) −DKL ( qφ ( z|x , y ) ‖pψ ( z|y ) ) ≤ log pθ , ψ ( x|y ) ( 3 ) Similarly to VAE , this objective is optimized using the reparameterization trick . Note that the prior distribution pψ ( z|y ) is conditioned on y and is modeled by a neural network with parameters ψ . Thus , CVAE uses three trainable neural networks , while VAE only uses two . Also authors propose such modifications of CVAE as Gaussian stochastic neural network and hybrid model . These modifications can be applied to our model as well . Nevertheless , we don ’ t use them , because of their disadvantage which is described in appendix C . 4 VARIATIONAL AUTOENCODER WITH ARBITRARY CONDITIONING . 4.1 PROBLEM STATEMENT . Consider a distribution pd ( x ) over a D-dimensional vector x with real or categorical components . The components of the vector are called features . Let binary vector b ∈ { 0 , 1 } D be the binary mask of unobserved features of the object . Then we describe the vector of unobserved features as xb = { xi : bi=1 } . For example , x ( 0,1,1,0,1 ) = ( x2 , x3 , x5 ) . Using this notation we denote x1−b as a vector of observed features . Our goal is to build a model of the conditional distribution pψ , θ ( xb|x1−b , b ) ≈ pd ( xb|x1−b , b ) for an arbitrary b , where ψ and θ are parameters that are used in our model at the testing stage . However , the true distribution pd ( xb|x1−b , b ) is intractable without strong assumptions about pd ( x ) . Therefore , our model pψ , θ ( xb|x1−b , b ) has to be more precise for some b and less precise for others . To formalize our requirements about the accuracy of our model we introduce the distribution p ( b ) over different unobserved feature masks . The distribution p ( b ) is arbitrary and may be defined by the user depending on the problem . Generally it should have full support over { 0 , 1 } D so that pψ , θ ( xb|x1−b , b ) can evaluate arbitrary conditioning . Nevertheless , it is not necessary if the model is used for specific kinds of conditioning ( as we do in section 5.2 ) . Using p ( b ) we can introduce the following log-likelihood objective function for the model : max ψ , θ Epd ( x ) Ep ( b ) log pψ , θ ( xb|x1−b , b ) ( 4 ) The special cases of the objective ( 4 ) are variational autoencoder ( bi = 1 ∀i ∈ { 1 , . . . , D } ) and conditional variational autoencoder ( b is constant ) .

Summary: The paper presents a model for learning conditional distribution when arbitrary partitioning the input to observed and masked parts. The idea is to extend the conditional VAE framework such that the posterior is a function of an arbitrary subset of observed variables. Accordingly, reconstruction loss only penalizes the error in the reconstruction of masked (unobserved) variables. The method is compared against 1) classical approaches in missing data imputation on UCI benchmarks; 2) image inpainting against recently proposed GANS for the similar task, as well as; 3) against universal marginalizer, which learns conditional densities using a feedforward / autoregressive architecture.