Year: 2021

Abstract:
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 flow (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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 APPROACH

heading: 3.1 PROBLEM DEFINITION
text: 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.

heading: 3.2 NORMALIZING FLOWS
text: 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 .

heading: 3.3 ONLINE EM
text: 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 {z1, . . . , zn} where zi = f−1ψ (xi), and the corresponding missing mask {m1, . . . ,mn}, EM aims to estimate (µ,Σ) in an iterative way:
µ̂(t+1) = g µ(µ̂ (t), Σ̂(t); {zi,mi}ni=1)
Σ̂(t+1) = gΣ(µ̂ (t), Σ̂(t); {zi,mi}ni=1)
(5)
where (µ̂(t), Σ̂(t)) are the estimates at the tth iteration, and {gµ(·), gΣ(·)} denote the mappings between two consecutive iterations 2.
Given estimates (µ̂, Σ̂), the missing part zmi is imputed by its conditional mean given the observed part zoi :
ẑmi = E(z m i |zoi ; µ̂, Σ̂) = µ̂mi + Σ̂mioi ( Σ̂oioi )−1 (zoi − µ̂oi) (6)
where oi is the observed mask (i.e. the complement of mi), and the subscripts of (µ̂, Σ̂) 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 Cappé & 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 first obtains local estimates from a batch of data:
µ̂local = gµ(µ̂ (t), Σ̂(t); {zi,mi}i∈B)
Σ̂local = gΣ(µ̂ (t), Σ̂(t); {zi,mi}i∈B)
(7)
The global estimates are then updated in the fashion of weighted average:
µ̂(t+1) = ρt+1µ̂local + (1− ρt+1)µ̂(t) Σ̂(t+1) = ρt+1Σ̂local + (1− ρt+1)Σ̂(t) (8)
2See appendix B for the details of gµ and gΣ.
where (ρ1, ρ2, . . .) are a sequence of step sizes satisfying
0 < ρt < 1, ∞∑ i=1 ρi =∞ and ∞∑ i=1 ρ2i <∞ (9)
In this work, we use a step size schedule defined by
ρt = Ct −γ , t = 1, 2, . . . (10)
where C is a positive constant and γ ∈ (0.5, 1].

heading: 3.4 ARCHITECTURE AND INFERENCE
text: 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 pX . 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 Xin are imputed naively (e.g. median imputation for tabular datasets) at the very beginning to get the initial current imputed data 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 X̂ stay fixed, while the parameter estimates of NF (i.e. ψ̂) and base distribution (i.e. µ̂ and Σ̂) are updated in different ways.
First of all, given the current estimated base distribution N (·; µ̂, Σ̂), the flow model fψ are learned by minimizing the negative log-likelihood of a batch of the current imputed data X̂B :
L1(ψ) = − 1 |B| ∑ i∈B log pX(x̂i;ψ, µ̂, Σ̂) (11)
where |B| denotes the batch size. The computation of L1 requires exact likelihood evaluation that is equipped by NF. Once the flow model parameters get updated, we obtain the embedding vectors in the latent space:
zi = f −1 ψ̂ (x̂i), i ∈ B (12)
Although the embedding vectors are complete, they are treated as incomplete using the missing masks in the data space {mi}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 (µ̂, Σ̂):
ẑmi = E(z m i |zoi ; µ̂, Σ̂), i ∈ B (13)
which results in new embedding vectors {ẑ}i∈B where ẑi consists of the observed part zoi and the imputed part ẑmi . After the imputation, the global estimates (µ̂, Σ̂) are also updated following equation 7 and equation 8.
Since the base distribution has been changed, it’s necessary to update the flow model fψ again by optimizing a composite loss:
L2(ψ) = − 1 |B| ∑ i∈B [ log pX(x̃i;ψ, µ̂, Σ̂)− αLrec(x̃i, x̂i,mi) ] (14)
where x̃i = fψ(ẑi), and Lrec(x̃i, x̂i,mi) is the reconstruction error only for non-missing values:
Lrec(x̃i, x̂i,mi) = p∑ j=1 (1−mij)(x̃ij − x̂ij)2 (15)
In this composite loss, the first term forces the reconstructed data vectors {x̃}i∈B to have high likelihood in the data space, while the second term encourages {x̃}i∈B to match the observed parts of the incomplete training data. And since {x̃}i∈B are transformed from {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.
Algorithm 1 Training Phase
1: Input: Current imputation: X̂ = (x̂1, . . . , x̂n)T , missing masks M = (m1, . . . ,mn)T , initial estimates of the base distribution (µ̂(0), Σ̂(0)), online EM step size sequence: ρ1, ρ2, . . . , ρt, . . . 2: for t = 1 to Tepoch do 3: Get a mini-batch X̂B = {x̂i}i∈B 4: # update the flow model 5: Compute L1 in equation 11 6: Update ψ via gradient descent 7: # update the base distribution 8: zi = f −1 ψ (x̂i), i ∈ B
9: Impute in the latent space with (µ̂(t−1), Σ̂(t−1)) to get {ẑi}i∈B via equation 13 10: Obtain updated (µ̂(t), Σ̂(t)) via equation 7 and equation 8 11: # update the flow model again 12: x̃i = fψ(ẑi), i ∈ B 13: Compute L2 via equation 14 14: Update ψ via gradient descent
Re-Imputation Phase After the training phase, the re-imputation phase is executed to update the current imputation X̂. This procedure is similar to that in the training phase, except that all model parameters are kept fixed. As shown by the last line of Algorithm 2, the missing parts of X̂ are replaced with those of the reconstructed data vectors, while the observed parts are kept the same.
Algorithm 2 Re-imputation Phase

heading: 4 EXPERIMENTS
text: In this section, 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).
To make the comparison more objective, both models use the same normalizing flow with six affine 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.

heading: 4.1 MULTIVARIATE DATASETS
text: Ten multivariate datasets from the UCI repository (Dheeru & Taniskidou, 2017) are used for evaluation. For all of them, each feature is scaled to fit 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 first 70% features3.
The initial imputation is performed by randomly sampling from the observed entries of each feature. All experiments are conducted using five-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 significantly faster than MCFlow. In fact, EMFlow converges within three iterations for most of the UCI datasets.

heading: 4.2 IMAGE DATASETS
text: 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
3See appendix E for the details of how MAR is simulated.
[0, 1]. 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 filled 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 efficiency of EMFlow, We also compare EMFlow and MCFlow with respect to the accuracy of post-imputation classification. 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 significant on CIFAR-10. We note that these findings 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 first 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.

heading: 5 CONCLUSION
text: We propose a novel architecture EMFlow for missing data imputation. It combines the strength of the online EM and the normalizing flow to learn the density estimation in the presence of incomplete data while performing imputation. Various experiments with multivariate and image datasets show that EMFlow significantly 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 classification on image datasets also demonstrates the superior EMFlow’s ability of recovering semantic structure from incomplete data.
REPRODUCIBILITY STATEMENT
The code and data to reproduce the results in this work are included in the supplementary materials.

heading: A AFFINE COUPLING LAYERS
text: The basic idea behind normalizing flow is to find 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) ∼ Np(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 Uj = gj(X) ∼ U(0, 1) for j = 1, . . . , p, which are known as universal residuals. Thus, choosing Zj = f−1j (X) ≡ Φ−1(gj(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−11 (x), . . . , f −1 p (x)) that enables to transform X to Z. However, finding 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 efficiency.
Following the framework of Real NVP (Dinh et al., 2016), the NF used throughout this work consists of a sequence of affine coupling layers faff : x 7→ y defined as
y1:d = x1:d
yd+1:p = xd+1:p exp (s (x1:d)) + t (x1:d) (16)
where the scale and shift parameters s(·) and t(·) are usually implemented by neural networks, and is the element-wise product. Therefore, the input x are split into two parts:the first d dimensions stay the same, while the other dimensions undergo an affine transformation whose parameters are the functions of the first d dimensions.

heading: B ONLINE EM FOR MISSING DATA IMPUTATION
text: 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 Cappé & Moulines (2009) can be easily applied here.
B.1 VANILLA EM
Given n i.i.d data points {zi = (zoi , zmi )}ni=1 distributed under N (µ,Σ), the E-step of each EM iteration evaluates the conditional expectation of the complete data likelihood:
Q(φ; φ̂(t)) = 1
n n∑ i=1 Qi(φ; φ̂ (t)) = 1 n n∑ i=1 Eφ̂(t) [logN (zi;φ)|z o i ] (17)
where φ = (µ,Σ), and φ̂(t) = (µ̂(t), Σ̂(t)) are the estimates at the tth iteration. We also follow the treatment by Cappé & 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 zmi given z o i is still Gaussian:
zmi |zoi ; φ̂(t) ∼ N (µ̃ m(t) i , Σ̃ m(t) i ) (18)
where
µ̃ m(t) i = µ̂ (t) mi + Σ̂ (t) mioi ( Σ̂(t)oioi )−1 ( zoi − µ̂(t)oi ) ,
Σ̃ m(t) i = Σ̂ (t) mimi − Σ̂ (t) mioi ( Σ̂(t)oioi )−1 Σ̂(t)oimi ,
(19)
mi and oi are the missing and observed masks respectively, and the subscripts of (µ̂(t), Σ̂(t)) represent the slicing indexes.
Then it’s easily easy to find
Qi(φ; φ̂ (t)) = −1
2
[( ẑ (t) i − µ )T Σ−1 ( ẑ (t) i − µ ) +Tr ([( Σ−1 ) mimi ] Σ̃ m(t) i )] − 1 2 log |2πΣ|
(20)
where Tr(·) denotes the matrix trace, and ẑ(t)i is in fact the imputed data vector whose missing part is replaced by the conditional mean µ̃m(t)i .
When it comes to the M-step, the derivatives of Q(φ;φ(t)) are calculated as
∂Q ∂µ = 1 n Σ−1 n∑ i=1 ( ẑ (t) i − µ ) ∂Q
∂Σ−1 = − 1 2n n∑ i=1 ( ẑ (t) i − µ )( ẑ (t) i − µ )T + 1 2 [ Σ− 1 n n∑ i=1 Σ̃ (t) i ] (21)
where Σ̃(t)i is p× p matrix satisfying Σ̃ (t) i,mimi = Σ̃ m(t) i and all other elements equal to 0.
Therefore, the maximizer of Q(φ; φ̂(t)) are
µ̂(t+1) = 1
n n∑ i=1 ẑ (t) i ,
Σ̂(t+1) = 1
n n∑ i=1 [( ẑ (t) i − µ (t+1) )( ẑ (t) i − µ (t+1) )T + Σ̃ (t) i ] .
(22)
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 = (zo, zm) ∼ N (µ,Σ), its likelihood is maximized at zm = E [zm|zo].
To prove it, first note that the log-likelihood of z can be written as
logN (zm, zo|µ,Σ) = −1 2 [ zm − µm zo − µo ]T Σ−1 [ zm − µm zo − µo ] + C (23)
Denote the conditional mean and covariance of zm as µ̃ = µm + Σmo (Σoo) −1
(zo − µo) and Σ̃ = Σmm −Σmo (Σoo)−1 Σom. The matrix block-wise inversion gives
Σ−1 = [ Σmm Σmo Σom Σoo ]−1 = [ Σ∗11 Σ ∗ 12 Σ∗21 Σ ∗ 22 ] (24)
where
Σ∗11 = Σ̃ −1 Σ∗12 = −Σ̃−1ΣmoΣ−1oo
Σ∗21 = −Σ−1oo ΣmoΣ̃−1
Σ∗22 = Σ −1 oo + Σ −1 oo ΣmoΣ̃ −1ΣmoΣ −1 oo .
(25)
After a series of linear algebra, we have
logN (zm, zo|µ,Σ)
= −1 2 (zm − µ̃)T Σ̃−1 (zm − µ̃)− 1 2 (zo − µo)T Σ−1oo (zo − µo) + C (26)
As Σ̃−1 is positive definite, it is obvious that the likelihood reaches the maximum at zm = µ̃.
B.2 ONLINE 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, Cappé & Moulines (2009) propose to replace the reestimation functional in the E-step by a stochastic approximation step:
Q̂(t+1)(φ) = (1− ρt+1)Q̂(t)(φ) + ρt+1 · 1 |B| ∑ i∈B Eφ̂(t) [log f (zi;φ) | z o i ] (27)
where f is the probability density of complete data, {ρi}∞i=1 are a series of step sizes satisfying the conditions specified 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 Cappé & Moulines (2009), equation 27 boils down to a stochastic approximation of the estimates of sufficient statistics:
ŝ(t+1) = (1− ρt+1)ŝ(t) + ρt+1 1 |B| ∑ i∈B Eφ̂(t) [S(z)|z o i ] (28)
where S(z) is the sufficient statistic of f .
Note that the sufficient statistics for multivariate Gaussian are just sample mean and sample covariance. Therefore, equation 8 follows immediately form equation 28.
C INTERPRET THE OPTIMIZATION
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,
(x̂, ψ̂, µ̂, Σ̂) = arg max x∈X ′,ψ,µ,Σ
log pX(x|ψ,µ,Σ) (29)
While it’s hard to optimize such objective simultaneously, it’s possible to optimize it alternatively. For example, learningL1(ψ) 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.,
N (ẑ; µ̂, Σ̂) ≥ N (z; µ̂, Σ̂) (30)
where z and 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 flow model is updated again by optimizing L2(ψ) in equation 14 to ensure that x̃ = fψ(ẑ) has a high likelihood in the data space X . Note that although x̃ doesn’t necessarily belong to X ′, the reconstruction penalty Lrec defined in equation 15 forces x̃ to be close to X ′. Finally, in the re-imputation phase, the imputed data x̂ are updated by projecting x̃ onto X ′.
D IMPLEMENTATION DETAILS
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 flow model fψ are reinitialized after each iteration to learn the new data density faster.
The base distribution estimate N (z; µ̂, Σ̂) is initialized at the very beginning when the online EM encounters the first batch of embedding vectors. Specifically µ̂ and Σ̂ are initialized to be the sample mean and sample covariance. By default, the base distribution is also reinitialized along with the flow model after each iteration.
Stabilization of Online EM Online EM can be unstable in some cases where the covariance estimate Σ̂ 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 significantly; (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:
Σ̂Rob = Σ̂ + β · Diag(Σ̂) (31)
where Diag(Σ̂) is a diagonal matrix with the same diagonal entries as Σ̂, and β a positive hyperparameter. In practice, β would be decreased gradually to 0 during the first 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 modified 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 Ssuper. 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.

heading: E SIMULATE MAR ON UCI DATASETS
text: To simulate MAR, the first 70% features of each data point xi is retained, and the remaining 30% features are removed with probability:
sigmoid b0.7pc∑ j=1 xij  (32) where b·c 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.

heading: F HYPERPARAMETER ROBUSTNESS
text: EMFlow has a relatively simpler architecture than GAN-based models, and there are just two most important two hyperparameters specific 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 first 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 fixed. The RMSEs in all configurations are well below that obtained by MCFlow, and their fluctuations are almost negligible.
F.1 HYPERPARAMETERS FOR UCI DATASETS
We use α = 106 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 difficult under MAR where the marginal observed density is distorted more. Therefore, the robust covariance estimation in equation 31 is adopted by default under MAR. Specifically, we use β = 10−2 for the first two iterations, β = 10−3 for the next two iterations, and β = 0 for the remaining ones.
F.2 HYPERPARAMETERS FOR IMAGE DATASETS
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×108 and 1×106 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.

heading: G ADDITIONAL EXPERIMENT
text: G.1 MODEL CONGENIALITY
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 coefficients 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.
G.2 COMPARISON 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 4. On the other hand, the implementation of MIWAE on image datasets are not provided, and thus we tried our best to tune MIWAE on MNIST.
As shown in Table 6, EMFlow and MIWAE have mixed performance on UCI datasets, 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.
G.3 EFFECT OF INITIAL IMPUTATION
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, 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.
4https://github.com/pamattei/miwae

Review summaries:
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.

This paper presents a model named EMFlow, which performs data imputation in the latent space using the online EM algorithm together with the normalizing flow models. The normalizing flow models aim to capture the complete data density $p_X$ and the bidirectional mapping between the data space and the latent space, even when the data is only partially observed. The parameters in the latent space are updated using an online EM algorithm. Thanks to the feature-wise mapping, the dependency between features in the data space is carried over to the latent space and hence the imputation can be done. Evaluation using ten UCI datasets, MNIST, and CIFAR-10 datasets show impressive improvement against baseline models and the convergence is faster than MCFlow.

The paper aims at imputing missing data which are MCAR and MAR.  For modeling the observed data distribution, it utilizes the framework of normalizing flow of which the latent variable/source variable space is Gaussian. By assuming the consistency of inter-feature dependencies in the latent/source variable space, it applies online EM for the imputation of the latent space variables. In the experiments, the proposed method, EMFlow is compared with GAIN, MisGAN, and MCFlow.

The authors propose a novel architecture EMFlow for missing data imputation. The authors also show the results of various experiments with multivariate and image datasets. Finally, the authors report the accuracy of post-imputation classification on image datasets.