Learning from Partially-Observed Multimodal Data with Variational Autoencoders

1 INTRODUCTION . Learning from data is an integral part of machine learning and artificial intelligence . Modern deep learning techniques rely heavily on extracting information form large scale datasets . While such frameworks have been shown to be effective on various down-stream tasks such as classification , regression , representation learning , and prediction , it is typically crucial to have access to clean and complete training data . Complete data in this case can be either labeled data ( for classification ) , or time-series data with no missing values ( for regression ) , or simply image with no missing pixels ( for generation ) . As such , if a model can only access partially-observed data , the performance will likely be much worse than those trained with fully-observed data , if not completely failing . In practical scenarios , however , it is usually costly to acquire clean and complete data due to the limited human resources and time . Having a model designed to learn and extract information from partially-observed data will not only largely increase the application spectrum of deep learning based models , but also provide benefit to new down-stream tasks , for example , data imputation . Data imputation with deep generative models has been an active research area ( Yoon et al. , 2018 ; Ivanov et al. , 2019 ; Nazabal et al. , 2018 ) . Despite promising progress , there are still challenges in learning effective models . First , some prior works focus on learning from fully-observed data while performing imputation on partially-observed data during test phase ( Suzuki et al. , 2016 ; Ivanov et al. , 2019 ) . Second , they usually have strong assumptions on missingness mechanism ( see A.1 ) such as data is missing completely at random ( MCAR ) ( Yoon et al. , 2018 ) . Third , mostly unimodal imputation such as image in-painting has been explored for high-dimensional data ( Ivanov et al. , 2019 ; Mattei & Frellsen , 2019 ) . Unimodal data refers to data with only one modality such as image , video , or text . Modeling any combination of data modalities is not well-established yet , which apparently limits the potential of such models , since raw data in real-life is usually acquired in a multimodal manner ( Ngiam et al. , 2011 ) with more than one source of data gathered to represent a practical scenario . In practice , one or more of the modalities maybe be missing , leading to a challenging multimodal data imputation task . In this work , we propose Variational Selective Autoencoder ( VSAE ) for multimodal data generation and imputation . Our proposed VSAE tries to address the challenges above by learning from partiallyobserved training data . By constructing an encoder for each modality independently , the latent representation selectively takes only the observed modalities as input , while a set of decoders maps the latent codes to not only full data ( including both observed and unobserved modalities ) , but also a mask representing the missingness scheme . Thus , it can model the joint distribution of the data and the mask together and avoid limiting assumptions such as MCAR , and is optimized efficiently with a single variational objective . In our experimental validation , we evaluate our proposed VSAE on both synthetic high-dimensional multimodal data and challenging low-dimensional tabular data , and show that VSAE can outperform state-of-the-art baseline models for data imputation task . The contributions are summarized as follows : ( 1 ) A novel framework VSAE to learn from partially-observed multimodal data . ( 2 ) The proposed VSAE is capable of learning the joint distribution of observed and unobserved modalities as well as the imputation mask , resulting in a unified model for various down-stream tasks including data generation and imputation with relaxed assumptions on missigness mechanism . ( 3 ) Evaluation on both synthetic high-dimensional and challenging low-dimensional multimodal datasets shows improvement over the state-of-the-art data imputation models . 2 RELATED WORK . Our work is related to literature on data imputation and multi-modal representation learning . In this section , we briefly review recent models proposed in these two domains . Data Imputation . Classical imputation methods such as MICE ( Buuren & Groothuis-Oudshoorn , 2010 ) and MissForest ( Stekhoven & Bühlmann , 2011 ) learn discriminative models to impute missing features from observed ones . With recent advances in deep learning , several deep imputation models have been proposed based on autoencoders ( Vincent et al. , 2008 ; Gondara & Wang , 2017 ; Ivanov et al. , 2019 ) , generative adversarial nets ( GANs ) ( Yoon et al. , 2018 ; Li et al. , 2019 ) , and autoregressive models ( Bachman & Precup , 2015 ) . GAN-based imputation method GAIN proposed by Yoon et al . ( 2018 ) assumes that data is missing completely at random . Moreover , this method does not scale to high-dimensional multimodal data . Several VAE based methods ( Ivanov et al. , 2019 ; Nazabal et al. , 2018 ; Mattei & Frellsen , 2019 ) have been proposed in recent years . Ivanov et al . ( 2019 ) formulated VAE with arbitrary conditioning ( VAEAC ) which allows generation of missing data conditioned on any combination of observed data . This algorithm needs complete data during training and can not learn from partially-observed data only . Nazabal et al . ( 2018 ) and Mattei & Frellsen ( 2019 ) modified VAE formulation to model the likelihood of the observed data only . However , they require training of a separate generative network for each dimension thereby increasing computational requirements . In contrast , our method aims to model joint distribution of observed and unobserved data along with the missing pattern ( imputation mask ) . This enables our model to perform both data generation and imputation even under relaxed assumptions on missingness mechanism ( see Appendix A.1 ) . Learning from Multimodal Data . A class of prior works such as conditional VAE ( Sohn et al. , 2015 ) and conditional multimodal VAE ( Pandey & Dukkipati , 2017 ) focus on learning the conditional likelihood of the modalities . However , these models requires complete data during training and can not handle arbitrary conditioning . Alternatively , several generative models aim to model joint distribution of all modalities ( Ngiam et al. , 2011 ; Srivastava & Salakhutdinov , 2012 ; Sohn et al. , 2014 ; Suzuki et al. , 2016 ) . However , multimodal VAE based methods such as joint multimodal VAE ( Suzuki et al. , 2016 ) and multimodal factorization model ( MFM ) ( Tsai et al. , 2019 ) require complete data during training . On the other hand , Wu & Goodman ( 2018 ) proposed another multimodal VAE ( namely MVAE ) can be trained with incomplete data . This model leverages a shared latent space for all modalities and obtains an approximate joint posterior for the shared space assuming each modalities to be factorized . However , if training data is complete , this model can not learn the individual inference networks and consequently does not learn to handle missing data during test . Building over multimodal VAE approaches , our model aims to address the shortcomings above within a flexible framework . In particular , our model can learn multimodal representations from partially observed training data and perform data imputation from arbitrary subset of modalities during test . By employing a factorized multimodal representations in the latent space it resembles disentangled models which can train factors specialized for learning from different parts of data ( Tsai et al. , 2019 ) . 3 METHOD . In this section , we introduce a novel VAE-based framework named Variational Selective Autoencoder ( VSAE ) to learn from partially-observed multimodal data . We first formalize our problem and then provide a detailed description of our model . 3.1 PROBLEM STATEMENT . Let x = [ x1 , x2 ... , xM ] be the complete data with M modalities , where xi denotes the feature representation for the i-th modality . The size of each xi varies and can be very high-dimensional ( e.g . multimedia data ) or low-dimensional ( e.g . tabular data ) . We define an M -dimensional binary mask variable m ∈ { 0 , 1 } M to represent the observed and unobserved modalities : mi = 1 if the i-th modality is observed and 0 if unobserved . Thus we have the set of observed modalities O = { i|mi = 1 } , and the set of unobserved modalities U = { i|mi = 0 } . O and U are complementary subsets of all modalities . Accordingly , we denote the representation for the observed and unobserved modalities with xo = [ xi|mi = 1 ] and xu = [ xi|mi = 0 ] , respectively . In this paper , we assume the data x and the mask m are dependent , and aim to model the joint distribution of them together . As a result of such joint modeling , VSAE has higher capacity and can be used for both data imputation and data/mask generation . We encoder the multimodal data to a latent space factorized with respect to the modalities . To handle training and test with partially-observed data , the variational latent variable of each modality is modeled selectively to choose between a unimodal encoder if the corresponding modality is observed , or a multimodal encoder if the modality is unobserved . In addition , all the modalities and mask are reconstructed by decoding the aggregated latent codes through decoders . 3.2 BACKGROUND : VARIATIONAL AUTOENCODER . VAE ( Kingma & Welling , 2013 ) is a probabilistic latent variable model to generate a random variable x from a latent variable z with a prior distribution p ( z ) according to the marginalized distribution p ( x ) = Ez∼p ( z ) p ( x|z ) = ∫ p ( x|z ) p ( z ) dz . However , this is computationally intractable , so the likelihood log p ( x ) is approximated by variational lower bound ( ELBO ) Lθ , φ ( x ) : log p ( x ) ≥ Lθ , φ ( x ) = Ez∼qφ ( z|x ) [ log pθ ( x|z ) ] −DKL [ qφ ( z|x ) ||p ( z ) ] . ( 1 ) In this equation , qφ ( z|x ) is a proposal distribution to approximate intractable true posterior p ( z|x ) and parameterized by an inference network ( a.k.a encoder ) . pθ ( x|z ) is the conditional likelihood parameterized by another generative network ( a.k.a decoder ) . DKL is the Kullback-Leibler ( KL ) divergence between the prior and the proposal distribution and functions as a regularizer term , DKL [ qφ ( z|x ) ||p ( z ) ] = Ez∼qφ ( z|x ) [ log qφ ( z|x ) − log p ( z ) ] . To train this model Lθ , φ ( x ) is optimized over all training data with respect to the parameters θ and φ . For more details see Appendix A.2 . 3.3 PROPOSED MODEL : VARIATIONAL SELECTIVE AUTOENCODER . Our goal is to model the joint distribution p ( x , m ) = ∫ p ( x , m|z ) p ( z ) dz where x = [ xo , xu ] . Following VAE formulation , we construct a proposal distribution q ( z|x , m ) to approximate the intractable true posterior . See the architecture in Figure 1 , we denote the parameters of encoder by { φ , ψ } , and decoders of data and mask by θ and respectively . A lower bound of log p ( x , m ) can be derived as : Lφ , ψ , θ , ( x , m ) = Ez∼qφ , ψ ( z|x , m ) [ log pθ , ( x , m|z ) ] −DKL [ qφ , ψ ( z|x , m ) ||p ( z ) ] = Ez∼qφ , ψ ( z|x , m ) [ log pθ ( x|m , z ) + log p ( m|z ) − log qφ , ψ ( z|x , m ) + log p ( z ) ] . ( 2 ) We assume the variational latent variables can be factorized with respect to the modalities z = [ z1 , z2 , ... , zM ] , which is a standard assumption for multimodal data ( Tsai et al. , 2019 ) : p ( z ) = M∏ i=1 p ( zi ) , q ( z|x , m ) = M∏ i=1 q ( zi|x , m ) . ( 3 ) Given this , we define the proposal distribution parameterized by φ and ψ for each modality as qφ , ψ ( zi|x , m ) = { qφ ( zi|xi ) if mi = 1 qψ ( zi|xo , m ) if mi = 0 ( 4 ) This is based on the intuitive assumption that the latent space of each modality is independent of other modalities given its data is observed . But , if the data is missing for some modality , its latent space is constructed from the other observed modalities . We call this selective proposal distribution . In the decoder , the probability distribution also factorizes over the modalities assuming that the reconstructions are conditionally independent given the complete set of latent variables of all modalities : log pθ ( x|m , z ) = log pθ ( xo , xu|m , z ) = ∑ i∈O log pθ ( xi|m , z ) + ∑ j∈U log pθ ( xj |m , z ) ( 5 ) To summarize , the ELBO in Equation 2 can be rewritten as Lφ , ψ , θ , ( xo , xu , m ) =Ez ∑ i∈O log pθ ( xi|m , z ) + ∑ j∈U log pθ ( xj |m , z ) + Ez [ log p ( m|z ) ] − M∑ i=1 Ezi [ log qφ , ψ ( zi|x , m ) − log p ( zi ) ] , ( 6 ) where zi ∼ qφ , ψ ( zi|x , m ) according to the selective proposal distribution given in Equation 4 . For training the model , the ELBO should be maximized over training data . However under partiallyobserved setting , xu is missing and unavailable even during training . Thus , we define the objective function for training by taking expectation over xu L′φ , ψ , θ , ( xo , m ) = Exu [ Lφ , ψ , θ , ( xo , xu , m ) ] ( 7 ) Only one term in Equation 6 is dependent on xu , so the final objective function is obtained as L′φ , ψ , θ , ( xo , m ) =Ez ∑ i∈O log pθ ( xi|m , z ) + ∑ j∈U Exj [ log pθ ( xj |m , z ) ] + Ez [ log p ( m|z ) ] − M∑ i=1 Ezi [ log qφ , ψ ( zi|x , m ) − log p ( zi ) ] , where zi ∼ qφ , ψ ( zi|x , m ) ( 8 ) In the proposed algorithm , we approximate Exj [ log pθ ( xj |m , z ) ] , j ∈ U using reconstructed unobserved data sampling from the prior network . Our experiments show that even a single sample is sufficient to learn the model effectively . In fact , the prior network can be used as a self supervision mechanism to find the most likely samples which dominate the other samples when taking the expectation . In Equation 8 , pθ ( xi|m , z ) is the decoding term of corresponding modality xi and the type of distribution depends on the data . The mask decoding term pθ ( m|z ) is factorized Bernoulli distribution modeling the binary mask variable . The prior is standard normal distribution p ( z ) = ∏M i=1 p ( zi ) = ∏M i=1N ( zi ; 0 , I ) which is fully-factorized .

Summary: The paper proposes a novel training method for variational autoencoders that allows using partially-observed data with multiple modalities. A modality can be a whole block of features (e.g., a MNIST image) or just a single scalar feature. The probabilistic model contains a latent vector per modality. The key idea is to use two types of encoder networks: a unimodal encoder for every modality which is used when the modality is observed, and a shared multimodal encoder that is provided all the observed modalities and produces the latent vectors for the unobserved modalities. The whole latent vector is passed through a decoder that predicts the mask of observed modalities, and another decoder that predicts the actual values of all modalities. The “ground truth” values for the unobserved modalities are provided by sampling from the corresponding latent variables from the prior distribution once at some point of training.