EMFlow: Data Imputation in Latent Space via EM and Deep Flow Models

1 INTRODUCTION . Missing values are often encountered for real-world datasets and are known to adversely impact the validity of down-stream analysis . Most machine learning ( ML ) algorithms and statistical tools often ignore the problem of partially observed features by simply dropping entire cases with missing values , which could result into biased estimation and underestimation of parameter uncertainty ( Schmitt et al. , 2015 ; Somasundaram & Nedunchezhian , 2011 ; Lall , 2016 ) . There have been some recent attempts to develop ML methods that properly handle missing values . Early works in this field performed imputation in a supervised manner where a complete training set is needed to learning the correlation between missing and observed entries ( e.g . Garcı́a-Laencina et al. , 2010 ; Rezende et al. , 2014 ; Bertalmio et al. , 2000 ; Xie et al. , 2012 ; Yeh et al. , 2016 ) . However , it is common that a large collection of fully observed data is hard to acquire , which makes those data-greedy supervised algorithms less effective . On the other hand , early attempts in unsupervised imputation methods , such as the collaborative filtering ( Sarwar et al. , 2001 ) , usually learn the correlation across dimensions by embedding the data into a lower dimensional latent space linearly ( e.g . Little & Rubin , 2019 ; Audigier et al. , 2016 ) . To improve the limited representation power of those linear models , kernel-based methods are also proposed but at the cost of excessive computational load when dealing with large datasets ( e.g . Sanguinetti & Lawrence , 2006 ; Liu et al. , 2016 ) . More recently , several imputation methods based on deep generative models have been proposed , providing much more accurate estimates of missing values especially for high-dimensional image datasets . In this work , we introduce EMFlow that integrates the normalizing flow ( NF ) ( Dinh et al. , 2016 ; Rezende & Mohamed , 2015 ; Dinh et al. , 2014 ; Kingma & Dhariwal , 2018 ) with an online version of Expectation-Maximization ( EM ) algorithm ( Cappé & Moulines , 2009 ) . The proposed framework is motivated by the strength and weakness of EM . As a class of iterative algorithm designed for latent variable models , EM can be applied to data imputation in an interpretable way ( Little & Rubin , 2019 ) . Additionally , the learning of EM is usually numerically stable and its convergence property has been studied extensively ( e.g . Ma et al. , 2000 ; Zhao et al. , 2020 ; Meng et al. , 1994 ; Wu , 1983 ) . However , the E-step and M-step are only traceable with simple underlying distributions including multivariate Gaussian or Student ’ s t distributions as well as their mixtures ( e.g . Di Zio et al. , 2007 ; xian Wang et al. , 2004 ) . On the other hand , NF is capable of latent representation and efficient data sampling , which makes it a convenient bridge between the data space and the latent space . Therefore , we let EM perform imputation in the latent space where simple inter-feature dependency is assumed ( i.e . the underlying distribution is multivariate Gaussian ) . Meanwhile , NF is used to recover the mapping between the complex inter-feature dependency in the data space and the one in the latent space learned by EM . The inference of EMFlow adopts an iterative learning strategy that has been widely used in modelbased multiple imputation methods where the initial naive imputation is refined step by step until convergence ( e.g . Gondara & Wang , 2017 ; Buuren & Groothuis-Oudshoorn , 2010 ; Stekhoven & Bühlmann , 2012 ) . Specifically , three steps are performed alternatively : • update the density estimation of complete data including the observed and current imputed values ; • update the base distribution ( i.e . µ and Σ ) in the latent space by EM ; and • update the imputation in the data space . Note that the first update corresponds to optimizing the complete data likelihood as well as the reconstruction error computed on the observed entries . We also derive an online version of EM such that it only consumes a batch of data at a time for parameter updates and thus can work with deep generative models smoothly . We will show that such learning schema has simpler implementation and leads to faster convergence compared to other competing methods . The main contributions of this work are ( i ) an imputation framework combining an online version of Expectation-Maximization ( EM ) algorithm and the normalizing flow ; ( ii ) an iterative learning schema that alternatively updates the inter-feature dependency in the latent space ( i.e . the base distribution ) and density estimation in the data space ; ( iii ) a derivation of online EM in the context of missing data imputation ; and ( iv ) extensive experiments on multiple image and tabular datasets to demonstrate the imputation quality and convergence speed of the proposed framework . 2 RELATED WORK . Recently , the applications of deep generative models like Generative Adversarial Networks ( GAN ) ( Goodfellow et al. , 2014 ) have been extended to the field of missing data imputation under the assumption of Missing at Random ( MAR ) or Missing Completely at Random ( MCAR ) . GAIN ( Yoon et al. , 2018 ) designs a compete data generator that performs imputation , and a discriminator to differentiate imputed and observed components with the help of a hint mechanism . MisGAN ( Li et al. , 2019 ) introduces another pair of generator-discriminator that aims to learn the distribution of missing mask . However , training GAN-based models is a notoriously challenging for its excessively complex structure and non-convex objectives . For example , MisGAN optimizes three objectives jointly involving six neural networks for imputation . Furthermore , GAN-based models do not obtain explicit density estimation that could be critical for down-stream analysis ( e.g . Ferdosi et al. , 2011 ; Zambom & Ronaldo , 2013 ) . Some imputation techniques based on Variational Autoencoders ( VAE ) ( Kingma & Welling , 2013 ) are also developed . For example , MIWAE Mattei & Frellsen ( 2019 ) leverages the importanceweighted autoencoder ( Burda et al. , 2015 ) and optimizes a lower bound of the likelihood of observed data . But the zero imputation adopted by it can be problematic since observed entries can also be zero . It also needs quite large computational power to make the bound to be tight . EDDI ( Ma et al. , 2018 ) avoids zero imputation by introducing a permutation invariant encoder . Imputation under Missing not at Random ( MNAR ) has also been explored by explicitly modeling the missing mechanism ( Ipsen et al. , 2020 ; Collier et al. , 2020 ) . However , all VAE-based approaches only permit approximate density estimation no matter how expressive the inference model is . Our proposed framework builds on the work of MCFlow ( Richardson et al. , 2020 ) that utilizes NF to learn exact density estimation on incomplete data via an iterative learning schema . The core component of MCFlow is a feed forward network that operates in the latent space and attempts to find the likeliest embedding vector in a supervised manner . Although MCFlow achieves impressive performance compared to other state-of-art methods , it remains ambiguous that how the feed forward network exploits the correlation in the latent space . A key difference between the MCFlow and our proposed method is that we exploit the correlation in the latent space which in turn induces correlation in ambient space via NF . We also find that the inference of MCFlow often has slow convergence speed and can be unstable . 3 APPROACH . 3.1 PROBLEM DEFINITION . We define the complete dataset X as a collection of p-dimensional vectors { x1 , . . . , xn } that are independent and identically distributed ( i.i.d . ) samples drawn from pX ( · ; θ ) in a p-dimensional data space X . The missing pattern of xi is described by a binary mask mi ∈ { 0 , 1 } p such that xij is missing if mij = 1 , and xij is observed if mij = 0 . Let xo and xm be the observed and missing parts of x , and pM ( m|x ) = pM ( m|xo , xm ) be the conditional distribution of the mask . Based on dependency between m and ( xo , xm ) , the missing mechanism can be classified into three classes ( Little & Rubin , 2019 ) : • MCAR : pM ( m|xo , xm ) = pM ( m ) • MAR : pM ( m|xo , xm ) = pM ( m|xo ) • MNAR : The probability of missing depends on both xo and xm . Throughout this paper , we only focus on MCAR or MAR where the missing mechanism can be safely ignored . The typical objective of learning a latent model is to maximize the observed data log-likelihood defined as Lobs ( θ ) = n∑ i=1 log ∫ pX ( x o i , x m i ; θ ) dx m i . ( 1 ) However , our goal in this work is to estimate pX from incomplete data as well as obtain accurate imputation under pX . Therefore , we attempt to learn the reconstructed data X̂ = ( x̂1 , . . . , x̂n ) T and the estimated model parameter θ̂ via ( X̂ , θ̂ ) = arg max xi∈X ′i , θ n∑ i=1 log pX ( xi|θ ) ( 2 ) where X ′i ⊆ X is the search space for xi where the observed locations have fixed values . 3.2 NORMALIZING FLOWS . To make it possible to optimize equation 2 , pX ( x ; θ ) needs to be specified in a parametric way that should be expressive enough , as the density is potentially complex and high-dimensional . To this end , we use NF to model pX ( x ; θ ) as an invertible transformation fψ of a base distribution pZ ( z ; φ ) in the latent space Z . Under the change of variables theorem , the complete data log-density is specified as pX ( x ; θ ) = pZ ( f −1 ψ ( x ) ; φ ) ∣∣∣∣∣det ( ∂f−1ψ ( x ) ∂xT ) ∣∣∣∣∣ ( 3 ) where θ = ( ψ , φ ) . fψ is usually composed by a sequence of relatively simple transformations to approximate arbitrarily complex distributions with high representation power . In this work , we choose Real NVP ( Dinh et al. , 2016 ) based on affine coupling transformations as the flow model1 . Recently , Teshima 1See appendix A for the details of Real NVP . et al . ( 2020 ) has shown that flow models constructed from affine coupling layers can be universal distributional approximators . Data imputation relies on the inter-feature dependency that is usually intractable and hard to capture in the data space X with the presence of missing entries . Therefore , we make the following two assumptions related to NF . Assumption 1 : The inter-feature dependency in the latent space Z is simple and can be characterized by a multivariate Gaussian density , that is : pZ ( z ; φ ) = N ( z ; µ , Σ ) ( 4 ) where φ = ( µ , Σ ) consists of latent parameters ; mean vector µ and covariance matrix Σ . Note that the base distribution is usually chosen as a standard Gaussian distribution ( with µ = 0 and Σ = Ip ) in the literature for simplicity . However , the covariance Σ is essential in our imputation work to represent the inter-feature dependency . Assumption 2 : Given that the transformations involved in NF are feature-wise ( e.g . Dinh et al. , 2016 ; 2014 ) , we expect NF to learn the mapping between the complex inter-feature dependency in the data space X and the simple one in the latent space Z .

Summary: This paper presents a novel imputation method for high-dimensional datasets that typically serve as benchmarks for machine learning methods. This method (EMFlow) innovates by training a normalizing flow network to map input data samples to a multivariate Gaussian, where imputation is performed via an online version of expectation-maximization (EM), which is commonly used for missing data imputation. EMFlow is applied across regression tasks in datasets in the UCI machine learning dataset repository, and standard image classification datasets, MNIST and CIFAR-10. Empirical results show strong performance for missing data imputation, as well as downstream classification from these imputations, and the model design choices and well-constructed architecture make for easy training and fast optimization convergence.