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 goal of this paper is to use deep generative models for missing data imputation. This paper proposes learning a latent variable deep generative model over every randomly sampled subset of observed features. First, a masking variable is sampled from a chosen prior distribution. The mask determines which features are observed. Then, the likelihood of the observed features is maximized via a lower bound. Inference in this latent variable model is achieved through the use of an inference network which conditions on the set of "missing" (to the generative model) features.