EMF LOW : D ATA I MPUTATION IN L ATENT S PACE VIA EM AND D EEP F LOW M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
The presence of missing values within high-dimensional data is an ubiquitous problem for many applied sciences.
A serious limitation of many available data mining and machine learning methods is their inability to handle partially missing values and so an integrated approach that combines imputation and model estimation is vital for down-stream analysis.
A computationally fast algorithm, called EMFlow, is introduced that performs imputation in a latent space via an online version of Expectation-Maximization (EM) algorithm by using a normalizing ﬂow (NF) model which maps the data space to a latent space.
The proposed
EMFlow algorithm is iterative, involving updating the parameters of online EM and NF alternatively.
Extensive experimental results for high-dimensional multivariate and image datasets are presented to illustrate the superior performance of the EMFlow compared to a couple of recently available methods in terms of both predictive accuracy and speed of algorithmic convergence.
I NTRODUCTION
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 ﬁeld 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 ﬁltering (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 ﬂow (NF)
(Dinh et al., 2016;
Rezende & Mohamed, 2015; Dinh et al., 2014; Kingma & Dhariwal, 2018) with an online version of Expectation-Maximization (EM) algorithm (Capp´e & 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 efﬁcient 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 reﬁned step by step until convergence (e.g. Gondara & Wang, 2017; Buuren & Groothuis-Oudshoorn, 2010; Stekhoven &
B¨uhlmann, 2012).
Speciﬁcally, 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 ﬁrst 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 ﬂow; (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.
R ELATED WORK
Recently, the applications of deep generative models like Generative Adversarial Networks (GAN) (Goodfellow et al., 2014) have been extended to the ﬁeld 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 ﬁnd 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 ﬁnd that the inference of MCFlow often has slow convergence speed and can be unstable.
A PPROACH 3. P ROBLEM D EFINITION
We deﬁne the complete dataset X as a collection of p -dimensional vectors { x 1 , . . .
, x n } that are independent and identically distributed (i.i.d.)
samples drawn from p X ( · ; θ ) in a p -dimensional data space X .
The missing pattern of x i is described by a binary mask m i ∈ { 0 , 1 } p such that x ij is missing if m ij = 1 , and x ij is observed if m ij = 0 .
Let x o and x m be the observed and missing parts of x , and p M ( m | x ) =
p M ( m | x o , x m ) be the conditional distribution of the mask.
Based on dependency between m and ( x o , x m ) , the missing mechanism can be classiﬁed into three classes (Little & Rubin, 2019):
• MCAR: p M ( m | x o , x m ) =
p M ( m ) • MAR: p M ( m | x o , x m ) =
p M ( m | x o ) • MNAR: The probability of missing depends on both x o and x m .
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 deﬁned as
[Equation]
However, our goal in this work is to estimate p X from incomplete data as well as obtain accurate imputation under p X .
Therefore, we attempt to learn the reconstructed data (cid:98) X
[Equation] [Equation]
x
[Equation]
x n ) T and the estimated model parameter (cid:98) θ via
[Equation]
where X (cid:48) i ⊆ X is the search space for x i where the observed locations have ﬁxed values.
3. N ORMALIZING F LOWS
To make it possible to optimize equation 2, p X ( x ; θ ) needs to be speciﬁed 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 p X ( x ; θ ) as an invertible transformation f ψ of a base distribution p Z ( z ; φ ) in the latent space Z .
Under the change of variables theorem, the complete data log-density is speciﬁed as
[Equation]
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 afﬁne coupling transformations as the ﬂow model 1 .
Recently, Teshima et al.
(2020) has shown that ﬂow models constructed from afﬁne 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:
[Equation]
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 Σ
= I p ) 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 .
3. O NLINE EM
EM is a class of iterative algorithms for latent variable models including missing data imputation(Little & Rubin, 2019).
In EMFlow, it works in the latent space Z where the underlying distribution is N ( z ; µ , Σ ) .
Given the embedding vectors { z 1 , . . .
, z n } where z i = f − 1 ψ ( x i ) , and the corresponding missing mask { m 1 , . . .
, m n } , EM aims to estimate ( µ , Σ ) in an iterative way:
[Equation]
where ( (cid:98) µ ( t ) , (cid:98) Σ ( t ) ) are the estimates at the t th iteration, and { g µ ( · ) , g Σ ( · ) } denote the mappings between two consecutive iterations 2 .
Given estimates ( (cid:98) µ , (cid:98) Σ ) , the missing part z mi is imputed by its conditional mean given the observed part z
oi :
[Equation]
where o i is the observed mask (i.e. the complement of m i ), and the subscripts of ( (cid:98) µ , (cid:98) Σ ) denote the slicing indexes.
When processing datasets of large volume, EM becomes impractical because it needs to read the whole data into the memory for each iteration.
Following the framework introduced by Capp´e &
Moulines (2009), we derive an online version of EM algorithm in the context of data imputation.
Let
B ⊂ { 1 , . . . , n } denote a mini-batch of sample indexes, the online EM ﬁrst obtains local estimates from a batch of data:
[Equation]
The global estimates are then updated in the fashion of weighted average:
[Equation]
where ( ρ 1 , ρ 2 , . . . ) are a sequence of step sizes satisfying
[Equation]
In this work, we use a step size schedule deﬁned by
[Equation]
where C is a positive constant and γ ∈ (0 . 5 , 1] .
3. A RCHITECTURE AND I NFERENCE
EMFlow is a composite framework that combines NF and online EM.
As illustrated in Fig 1, NF is the bidirectional tunnel between the data space and the latent space, aiming to learn the complete data density p X .
In the latent space, the online EM estimates the inter-feature dependency of the embedding vectors and performs imputation.
To address the issue that NF needs complete data vectors for computation, the incomplete data X in are imputed naively (e.g. median imputation for tabular datasets) at the very beginning to get the initial current imputed data (cid:98) X .
Afterwards, the objective in equation 2 is optimized in an iterative schema, where each iteration consists of a training phase and a re-imputation phase.
Training Phase
At this phase, the current imputed data (cid:98) X stay ﬁxed, while the parameter estimates of NF (i.e. (cid:98) ψ ) and base distribution (i.e. (cid:98) µ and (cid:98) Σ ) are updated in different ways.
First of all, given the current estimated base distribution N ( · ; (cid:98) µ, (cid:98) Σ) , the ﬂow model f ψ are learned by minimizing the negative log-likelihood of a batch of the current imputed data (cid:98) X B :
[Equation]
where | B | denotes the batch size.
The computation of L 1 requires exact likelihood evaluation that is equipped by NF.
Once the ﬂow model parameters get updated, we obtain the embedding vectors in the latent space:
[Equation]
Although the embedding vectors are complete, they are treated as incomplete using the missing masks in the data space { m i } i ∈ B , given that the invertible mapping parameterized by f ψ is featurewise.
Therefore, the online EM imputes the missing parts of the embedding vectors with the current global estimates ( (cid:98) µ , (cid:98) Σ ) :
[Equation]
which results in new embedding vectors { (cid:98) z } i ∈ B where (cid:98)
z
i consists of the observed part z oi and the imputed part (cid:98) z mi .
After the imputation, the global estimates ( (cid:98) µ , (cid:98) Σ ) are also updated following equation 7 and equation 8.
Since the base distribution has been changed, it’s necessary to update the ﬂow model f ψ again by optimizing a composite loss:
[Equation]
(14) where (cid:101) x i = f
ψ ( (cid:98) z i ) , and L rec ( (cid:101) x i , (cid:98) x i , m i ) is the reconstruction error only for non-missing values:
[Equation]
In this composite loss, the ﬁrst term forces the reconstructed data vectors { (cid:101) x } i ∈ B to have high likelihood in the data space, while the second term encourages { (cid:101) x } i ∈ B to match the observed parts of the incomplete training data.
And since { (cid:101) x } i ∈ B are transformed from { (cid:98) z } i ∈ B via f ψ , both terms are conditioned on the inter-feature dependency learned by EM.
A pseudocode for this training phase is presented in Algorithm 1, and the implementation details that aim to make the training phase more stable is presented in appendix D.
1: Input: Current imputation: (cid:98) X
= ( (cid:98) x 1 , . . .
, (cid:98) x n ) T , missing masks M
= ( m 1 , . . .
, m n ) T , initial estimates of the base distribution ( (cid:98) µ (0) , (cid:98) Σ (0) ) , online EM step size sequence: ρ 1 , ρ 2 , . .
. , ρ t , . . .
2: for t = 1 to T epoch do 3:
Get a mini-batch (cid:98) X B = { (cid:98)
x i } i ∈ B 4: # update the ﬂow model 5:
Compute L 1 in equation6:
Update ψ via gradient descent 7: # update the base distribution 8: z i = f − 1 ψ ( (cid:98) x i ) , i ∈ B 9:
Impute in the latent space with ( (cid:98) µ ( t − 1) , (cid:98) Σ ( t − 1) ) to get { (cid:98) z i } i ∈ B via equation10:
Obtain updated ( (cid:98) µ ( t ) , (cid:98) Σ ( t ) ) via equation 7 and equation11: # update the ﬂow model again 12: (cid:101) x i = f
ψ ( (cid:98) z i ) , i ∈ B 13:
Compute L 2 via equation14:
Update ψ via gradient descent
Re-Imputation Phase After the training phase, the re-imputation phase is executed to update the current imputation (cid:98) X .
This procedure is similar to that in the training phase, except that all model parameters are kept ﬁxed.
As shown by the last line of Algorithm 2, the missing parts of (cid:98) X are replaced with those of the reconstructed data vectors, while the observed parts are kept the same.
1: Input: Current imputation: (cid:98) X
= ( (cid:98) x 1 , . . .
, (cid:98) x n ) T , missing masks M
= ( m 1 , . . .
, m n ) T , estimates of the base distribution ( (cid:98) µ , (cid:98) Σ ) from the the previous training phase 2: for i = 1 to n do 3:
z i = f − 1 ψ ( (cid:98) x i ) 4:
Impute in the latent space with ( (cid:98) µ , (cid:98) Σ ) to get (cid:98) z i via equation5: (cid:101) x i = f
ψ ( (cid:98) z i ) 6: # update current imputation 7: (cid:98) x i = (cid:98) x i (cid:12) o i + (cid:101) x i (cid:12) m [Table] [Table] 131 ± .
024 ± .. 028 ± .. 029 ±
[Table]
Energy .
086 ± .. 092 ± .. 110 ± . 175 ± .. 176 ± .. 250 ± . CTG . 104 ± .. 140 ± .. 143 ± . 105 ± .. 153 ± ..
165 ± [Table] Song . 025 ± .. 030 ± .. 034 ± .
024 ± ..
031 ± .. 028 ± .
[Figure]
076 ± .
[Figure]
098 ± .. 097 ± .
102 ± .. 124 ± .. 135 ± .
[Figure] [Figure] ± [Figure]
002 ± .
[Figure]
003 ± .
002 ± .. 006 ± .. 006 ± .
[Figure]
E XPERIMENTS
In this section
[Figure]
we evaluate the performance of EMFlow on multivariate and image datasets in terms of the imputation quality and and the speed of model training.
Its performance is compared to that of MCFlow, the most related competitor that has been shown to be superior to other state of art methods (Richardson et al., 2020).
[Table] [Table]
make the comparison more objective, both models use the same normalizing ﬂow with six afﬁne coupling layers.
We also follow the authors’ suggestion for the the hyperparameter selection of
MCFlow throughout this section.
Additionally, we also present the benchmarks of other state-of-art models including GAIN (Yoon et al., 2018) and MisGAN (Li et al., 2019) for more comprehensive comparisons.
4. M ULTIVARIATE D ATASETS
Ten multivariate datasets from the UCI repository (Dheeru & Taniskidou, 2017) are used for evaluation.
For all of them, each feature is scaled to ﬁt inside the interval [0 , 1] via min-max normalization.
We simulate MCAR with a missing rate of 0.2 by removing each value independently according to a Bernoulli distribution.
We also simulate MAR scenario where the missing probability of the last 30% features depends on the values of the ﬁrst 70% features 3 .
The initial imputation is performed by randomly sampling from the observed entries of each feature.
All experiments are conducted using ﬁve-fold cross validation where the test set only goes through the re-imputation phase in each iteration.
The choices of hyperparameters are detailed in appendix
F, where we also show that EMFlow is not sensitive to the choice of hyperparameters.
Results The imputation performace is evaluated by calculating the Root Mean Squared Error (RMSE) between the imputed and true values.
As shown in Table 1, EMFlow performs constantly better than MCFlow under both MCAR and MAR settings for nearly all datasets.
Additionally, We trained EMFlow and MCFlow on the same machine with the same learning rate and batch size to compare the convergence speed.
Figure 2 shows the training loss and the test set
RMSE over time on three UCI datasets.
It shows that EMFlow converges signiﬁcantly faster than
MCFlow.
In fact, EMFlow converges within three iterations for most of the UCI datasets.
4. I MAGE D ATASETS
We also evaluate EMFlow on MNIST and CIFAR-10.
MNIST is a dataset of 28×28 grayscale images of handwritten digits (LeCun et al., 1998), and CIFAR-10 is a dataset of 32×32 colorful images from 10 classes (Krizhevsky et al., 2009).
For both datasets, the pixel values of each image are scaled to
[Figure]
In this section, we simulate MCAR where each pixel is independently missing with various probabilities from 0.1 to 0.9.
The initial imputation is performed by nearest-neighbor sampling where a missing pixel is ﬁlled by one of its nearest observed neighbors.
In our experiments, the standard 60,000/10,000 and 50,000/10,000 training-test set partitions are used for MNIST and CIFAR-10 respectively.
The choices of hyperparameters are detailed in appendix F.
Results Table 2 shows the RMSE of all considered methods on both image datasets.
In the case of
MNIST, EMFlow and MCFlow have similar RMSE and outperform other methods, while MCFlow starts to gain slight advantage under high missing rates.
In the case of CIFAR-10, EMFlow achieve much lower RMSE than all competing methods.
To further demonstrate the efﬁciency of EMFlow, We also compare EMFlow and MCFlow with respect to the accuracy of post-imputation classiﬁcation.
For this purpose, a LeNet-based model and a VGG19 model were trained on the original training sets of MNIST and CIFAR-10 respectively.
These models then made predictions on the imputed test sets under different missing rates.
Table 3 shows that EMFlow yields slightly better post-imputation prediction accuracy than MCFlow on
MNIST, while the improvement is much more signiﬁcant on CIFAR-10.
We note that these ﬁndings are in good agreement with the RMSE results.
To qualitatively compare the imputation quality of EMFlow and MCFlow, Figure 3 shows sample imputed images from CIFAR-10 with MCAR missing rates at 0.5 and 0.9.
The ﬁrst row includes the (complete) ground truth images for reference, while the second row includes the (incomplete) observed images on which the models were trained.
The last two rows showcase the reconstructed images by MCFlow and EMFlow, respectively.
It’s clear that EMFlow performs better than MCFlow by recovering more details and displaying sharper boundaries and cleaner background.
C ONCLUSION
We propose a novel architecture EMFlow for missing data imputation.
It combines the strength of the online EM and the normalizing ﬂow to learn the density estimation in the presence of incomplete data while performing imputation.
Various experiments with multivariate and image datasets show that EMFlow signiﬁcantly outperforms its state-of-art competitor with respect to imputation accuracy as well as the convergence speed under a wide range of missing rates and different missing mechanisms.
The accuracy of post-imputation classiﬁcation on image datasets also demonstrates the superior EMFlow’s ability of recovering semantic structure from incomplete data.
R EPRODUCIBILITY S TATEMENT
The code and data to reproduce the results in this work are included in the supplementary materials.
R
A A FFINE C OUPLING L AYERS
The basic idea behind normalizing ﬂow is to ﬁnd a transformation (of X ) that would be able to represent a complex density function as a function of simpler ones, such as a multivariate Gaussian random variable (to be denoted by Z ).
The quest for an invertible transformation f , such that Z = f − 1 ( X ) ∼ N p
0 , I ) can be challenging, but in theory such a transformation does exist.
Brockwell (2007) showed the existence of such a transformation which converts any p -dimensional random variable X to p independent uniform random variables U j = g j ( X ) ∼ U (0 , 1) for j = 1 , . .
, p , which are known as universal residuals.
Thus, choosing Z j = f − 1 j ( X ) ≡ Φ − 1 ( g j ( X )) ∼ N (0 , 1) , where Φ( · ) denotes the cumulative distribution function of a standard normal distribution, we can show the existence of Z = f ( X ) , where f − 1 ( x ) =
( f − 1 1 ( x ) , . . .
, f − 1 p ( x )) that enables to transform
X to Z .
However, ﬁnding such nonlinear transform based on observed data is challenging and so we use sequence of simpler transforms in the spirit of universal residuals, for simplicity and computational efﬁciency.
Following the framework of Real NVP (Dinh et al., 2016), the NF used throughout this work consists of a sequence of afﬁne coupling layers f aff : x (cid:55)→ y deﬁned as
[Equation]
where the scale and shift parameters s ( · ) and t ( · ) are usually implemented by neural networks, and (cid:12) is the element-wise product.
Therefore, the input x are split into two parts:the ﬁrst d dimensions stay the same, while the other dimensions undergo an afﬁne transformation whose parameters are the functions of the ﬁrst d dimensions.
B O NLINE EM FOR M ISSING D ATA IMPUTATION
In this section, the steps of vanilla EM in the context of missing data imputation are reviewed.
We then show how the online EM framework proposed by Capp´e & Moulines (2009) can be easily applied here.
B. V ANILLA EM
Given n i.i.d data points { z i =
( z oi , z mi ) } ni =1 distributed under N ( µ, Σ) , the E-step of each EM iteration evaluates the conditional expectation of the complete data likelihood:
[Equation] [Equation]
where φ = ( µ , Σ ) , and (cid:98) φ ( t ) =
( (cid:98) µ ( t ) , (cid:98) Σ ( t ) ) are the estimates at the t th iteration.
We also follow the treatment by Capp´e & Moulines (2009) to normalize the conditional expectation by 1 /n for easier transition to the online version of EM.
To calculate conditional expectation inside the summation, we can use the fact that the conditional distribution of z mi given z oi is still Gaussian:
[Equation]
where
[Equation]
m i and o i are the missing and observed masks respectively, and the subscripts of ( (cid:98) µ ( t ) , (cid:98) Σ ( t ) ) represent the slicing indexes.
Then it’s easily easy to ﬁnd
[Equation]
where Tr ( · ) denotes the matrix trace, and
(cid:98) z ( t ) i is in fact the imputed data vector whose missing part is replaced by the conditional mean (cid:101) µ m ( t ) i .
When it comes to the M-step, the derivatives of Q ( φ ; φ ( t ) ) are calculated as
[Equation]
(21) where (cid:101) Σ ( t ) i is p × p matrix satisfying (cid:101) Σ ( t ) i, m i m i = (cid:101)
Σ m ( t ) i and all other elements equal to 0 .
Therefore, the maximizer of Q ( φ ; (cid:98) φ ( t ) ) are
[Equation]
Here we provide a remark that helps to motivate the optimization of EMFlow described in Section 3.4.
Remark.
Given an incomplete data point z
=
( z o , z m ) ∼ N ( µ , Σ ) , its likelihood is maximized at z m =
E
[ z m | z
o ] .
To prove it, ﬁrst note that the log-likelihood of z can be written as
[Equation]
Denote the conditional mean and covariance of z m as (cid:101) µ = µ m + Σ mo ( Σ oo ) − 1 ( z o − µ o ) and (cid:101) Σ
= Σ
mm −
[Equation]
mo
[Equation]
Σ oo ) − 1 Σ om .
The matrix block-wise inversion gives
[Equation] [Equation]
where
[Equation] [Equation]
After a series of linear algebra, we have
[Equation]
As (cid:101) Σ − 1 is positive deﬁnite, it is obvious that the likelihood reaches the maximum at z m = (cid:101) µ .
B. O NLINE EM
It’s clear that the EM algorithm needs to process the whole dataset at each iteration, which becomes impractical when dealing with huge datasets.
To address this problem, Capp´e & Moulines (2009) propose to replace the reestimation functional in the E-step by a stochastic approximation step:
[Equation]
where f is the probability density of complete data, { ρ i } ∞ i =1 are a series of step sizes satisfying the conditions speciﬁed in equation 9, B ⊂ { 1 , . .
. , n } denotes the indexes of a mini-batch of samples, and | B | is the batch size.
Under some mild regulations such as the complete data model belongs to exponential family Capp´e & Moulines (2009), equation 27 boils down to a stochastic approximation of the estimates of sufﬁcient statistics:
[Equation]
where S ( z ) is the sufﬁcient statistic of f .
Note that the sufﬁcient statistics for multivariate Gaussian are just sample mean and sample covariance.
Therefore, equation 8 follows immediately form equation 28.
C I NTERPRET THE O PTIMIZATION
As described in Section 3.1, the estimated model parameters and the imputed data are supposed to maximize the complete data likelihood (only one data point is considered here for illustration purpose), namely,
[Equation]
While it’s hard to optimize such objective simultaneously, it’s possible to optimize it alternatively.
For example, learning L 1 ( ψ ) in equation 11 corresponds to updating ψ while keeping other variables unchanged.
The operations performed in the latent space needs more explanation.
First of all, it is shown in appendix B that the imputation increases the complete likelihood in the latent space, i.e.,
[Equation]
where z and (cid:98) z are the embedding vectors before and after the imputation, respectively.
However, the increase of likelihood in the latent space doesn’t guarantee the same thing in the data space.
Therefore, the ﬂow model is updated again by optimizing L 2 ( ψ ) in equation 14 to ensure that (cid:101) x
= f ψ ( (cid:98) z ) has a high likelihood in the data space X .
Note that although (cid:101) x doesn’t necessarily belong to X (cid:48) , the reconstruction penalty L rec deﬁned in equation 15 forces (cid:101) x to be close to X (cid:48) .
Finally, in the re-imputation phase, the imputed data (cid:98) x are updated by projecting (cid:101) x onto X (cid:48) .
D I MPLEMENTATION D ETAILS
Model Initialization Each inference iteration consists of a training phase and a re-imputation phase.
And since the latter phase updates the current imputation, the parameters of the ﬂow model f ψ are reinitialized after each iteration to learn the new data density faster.
The base distribution estimate N ( z ; (cid:98) µ , (cid:98) Σ ) is initialized at the very beginning when the online EM encounters the ﬁrst batch of embedding vectors.
Speciﬁcally (cid:98) µ and (cid:98) Σ are initialized to be the sample mean and sample covariance.
By default, the base distribution is also reinitialized along with the ﬂow model after each iteration.
Stabilization of Online EM Online EM can be unstable in some cases where the covariance estimate (cid:98) Σ becomes ill-conditioned during the inference.
Such instability can be traced back to two primary sources: (i) the initial naive imputation has really bad quality and the inter-feature dependency is distorted signiﬁcantly; (ii) the batch size is close to or less than the number of features, which makes it hard to estimate the covariance directly (e.g. Fan et al., 2008; Chen et al., 2011).
Fortunately, those two issues can be solved easily.
To reduce the impact of the initial naive imputation and make the covariance estimation more robust, we enlarge the diagonal entries of the original covariance estimate from equation 8 proportionally:
[Equation]
where Diag ( (cid:98) Σ ) is a diagonal matrix with the same diagonal entries as (cid:98) Σ , and β a positive hyperparameter.
In practice, β would be decreased gradually to 0 during the ﬁrst a few iterations as the quality of current imputation gets better and better.
To address the second issue, the training phase in Algorithm 1 can be modiﬁed to make the online
EM see more data points.
When updating the base distribution, the current batch of embedding vectors can be concatenated with previous ones to form a super -batch with a maximum size S super .
If the maximum size is exceeded, the most previous embedding vectors would be excluded.
And since the previous batches are already imputed, the super-batch brings nearly no additional computational overhead.
E S IMULATE MAR ON UCI D ATASETS
To simulate MAR, the ﬁrst 70% features of each data point x i is retained, and the remaining 30% features are removed with probability:
[Equation]
[Equation] where (cid:98)·(cid:99) denotes the integer part of a number.
Note that different data sets would exhibit different missing rates under our MAR setting.
Table shows the dimensions of UCI data sets used in our experiment as well as their MAR missing rates for the last 30% features.
[Table]
F H YPERPARAMETER R OBUSTNESS
EMFlow has a relatively simpler architecture than GAN-based models, and there are just two most important two hyperparameters speciﬁc to it.
One is the power index γ in the step size schedule of online EM (see equation 10), while the other is the strength of the reconstruction error α in the objective (see equation 14).
From our extensive experiments, γ ∈ [0 . 6 , 0 . 9] often leads to consistent results.
On the other hand, a practical guide for choosing α is to make the magnitudes of ﬁrst and second term of equation 14 don’t differ too much.
Figure 4 shows RMSE on the test set of Online News Popularity dataset versus different values of γ and α , while all other hyperparameters are kept ﬁxed.
The RMSEs in all conﬁgurations are well below that obtained by MCFlow, and their ﬂuctuations are almost negligible. [Figure] [Figure] EMFlow
MCFlow
F. H YPERPARAMETERS FOR UCI D ATASETS
We use α = 10 6 and choose ρ t = 0 . 99 · t − 0 .
8 as the step size schedule for all UCI datasets.
During the training of EMFlow, the batch size is 256 and the learning rate is 1 × 10 − 4 .
Compared to
MCAR, the initial imputation can be more difﬁcult under MAR where the marginal observed density is distorted more.
Therefore, the robust covariance estimation in equation 31 is adopted by default under MAR.
Speciﬁcally, we use β = 10 − 2 for the ﬁrst two iterations, β = 10 − 3 for the next two iterations, and β = 0 for the remaining ones.
F. H YPERPARAMETERS FOR I MAGE D ATASETS
For all experiments with the image datasets, the step size schedule of online EM is again ρ t = 0 .
99 · t − 0 . 8 .
For MNIST and CIFAR-10, we choose α to be 5 × 10 8 and 1 × 10 6 respectively.
During the training of EMFlow, the batch size is 512 and the learning rate is 1 × 10 − 3 .
Since image datasets have much larger dimensions than UCI datasets, the super-batch approach with a size of 3000 is adopted.
Additionally, similar robust covariance estimation used by UCI datasets is also applied to
CIFAR-10.
G A DDITIONAL E XPERIMENT G. M ODEL C ONGENIALITY
We also consider the congeniality of the imputation model that measures its ability to preserve the feature-label relationship (Meng, 1994; Burgess et al., 2013).
To quantify how the featurelabel relationship is changed after imputation, we calculated the euclidean distance between the coefﬁcients of two linear or logistic regression models, one learned from a complete test set and the other learned from the corresponding imputed version.
Three datasets with continuous or discrete targets were used in our experiments, and the results are shown in Table 5.
It shows that the gaps in the measured congeniality actually coincide with those in the measured RMSE.
[Table]
[Table]
G. C OMPARISON TO MIWAE
In this section, we compare EMFlow with MIWAE that introduces the importance-weighted autoencoder for imputation tasks.
For UCI datasets, we follow the implementation provided by the author.
On the other hand, the implementation of MIWAE on image datasets are not provided
[Table]
and thus we tried our best to tune MIWAE on MNIST.
As shown in Table
[Table]
EMFlow and MIWAE have mixed performance on UCI datasets
[Table]
and the differences in most cases are quite small.
Table 7 and 8 shows that EMFlow outperforms MIWAE on MNIST in terms of the RMSE on test set and post-imputation accuracy.
We also note that
EMFlow still converges much faster than MIWAE.
[Table]
[Table]
G. E FFECT OF I NITIAL I MPUTATION
EMFlow is an iterative imputation framework, and the nearest-neighbor (NN) imputation is performed for image datasets at the beginning as warm-start in all the previous experiments.
In this section, we compare NN imputation and zero imputation as the starting point of EMFlow and shows the convergence of the training loss, the RMSE on the test set
[Table]
and the Frobenius norm of the covariance matrix in the latent space.
As shown in Figure, 5, the difference between NN imputation and zero imputation is negligible and the convergence under both cases is fast.