GAIN: Missing Data Imputation using Generative Adversarial Nets

Introduction

Missing data is a pervasive problem. Data may be missing because it was never collected, records were lost or for many other reasons. In the medical domain, the respiratory rate of a patient may not have been measured (perhaps because it was deemed unnecessary/unimportant) or accidentally not recorded (Yoon et al. 2017; Alaa et al. 2018). It may also be the case that certain pieces of information are difficult or even dangerous to acquire (such as information gathered from a biopsy), and so these were not gathered for those reasons (Yoon et al. 2018b). An imputation algorithm can be used to estimate missing values based on data that was observed/measured, such as the systolic blood pressure and heart rate of the patient (Yoon et al. 2018c). A substantial amount of research has been dedicated to developing imputation algorithms for medical data (Barnard & Meng 1999; Mackinnon 2010; Sterne et al. 2009; Purwar & Singh 2015). Imputation algorithms are also used in many other applications such as image concealment, data compression, and counterfactual estimation (Rubin 2004; Kreindler & Lumsden 2012; Yoon et al. 2018a).

Missing data can be categorized into three types: (1) the data is missing completely at random (MCAR) if the missingness occurs entirely at random (there is no dependency on any of the variables), (2) the data is missing at random (MAR) if the missingness depends only on the observed variables A formal definition of MAR can be found in the Supplementary Materials., (3) the data is missing not at random (MNAR) if the missingness is neither MCAR nor MAR (more specifically, the data is MNAR if the missingness depends on both observed variables and the unobserved variables; thus, missingness cannot be fully accounted for by the observed variables). In this paper we provide theoretical results for our algorithm under the MCAR assumption, and compare to other state-of-the-art methods in this setting Empirical results for the MAR and MNAR settings are shown in the Supplementary Materials..

State-of-the-art imputation methods can be categorized as either discriminative or generative. Discriminative methods include MICE (Buuren & Oudshoorn 2000; Buuren & Groothuis-Oudshoorn 2011), MissForest (Stekhoven & Bühlmann 2011), and matrix completion (Mazumder et al. 2010a; Yu et al. 2016; Schnabel et al. 2016; Mazumder et al. 2010b); generative methods include algorithms based on Expectation Maximization (García-Laencina et al. 2010) and algorithms based on deep learning (e.g. denoising autoencoders (DAE) and generative adversarial nets (GAN)) (Vincent et al. 2008; Gondara & Wang 2017; Allen & Li 2016). However, current generative methods for imputation have various drawbacks. For instance, the approach for data imputation based on (García-Laencina et al. 2010) makes assumptions about the underlying distribution and fails to generalize well when datasets contain mixed categorical and continuous variables. In contrast, the approaches based on DAE (Vincent et al. 2008) have been shown to work well in practice but require complete data during training. In many circumstances, missing values are part of the inherent structure of the problem so obtaining a complete dataset is impossible. Another approach with DAE (Gondara & Wang 2017) allows for an incomplete dataset; however, it only utilizes the observed components to learn the representations of the data. (Allen & Li 2016) uses Deep Convolutional GANs for image completion; however, it also requires complete data for training the discriminator.

In this paper, we propose a novel imputation method, which we call Generative Adversarial Imputation Nets (GAIN), that generalizes the well-known GAN (Goodfellow et al. 2014) and is able to operate successfully even when complete data is unavailable. In GAIN, the generator’s goal is to accurately impute missing data, and the discriminator’s goal is to distinguish between observed and imputed components. The discriminator is trained to minimize the classification loss (when classifying which components were observed and which have been imputed), and the generator is trained to maximize the discriminator’s misclassification rate. Thus, these two networks are trained using an adversarial process. To achieve this goal, GAIN builds on and adapts the standard GAN architecture. To ensure that the result of this adversarial process is the desired target, the GAIN architecture provides the discriminator with additional information in the form of “hints”. This hinting ensures that the generator generates samples according to the true underlying data distribution.

Problem Formulation

Consider a $d$-dimensional space $\mathcal{X}=\mathcal{X}_{1}\times...\times\mathcal{X}_{d}$. Suppose that $\mathbf{X}=(X_{1},...,X_{d})$ is a random variable (either continuous or binary) taking values in $\mathcal{X}$, whose distribution we will denote $P(\mathbf{X})$. Suppose that $\mathbf{M}=(M_{1},...,M_{d})$ is a random variable taking values in $\{0,1\}^{d}$. We will call $\mathbf{X}$ the data vector, and $\mathbf{M}$ the mask vector.

Throughout the remainder of the paper, we will often use lower-case letters to denote realizations of a random variable and use the notation $\mathbf{1}$ to denote a vector of $1$s, whose dimension will be clear from the context (most often, $d$).

Generative Adversarial Imputation Nets

Then we define the random variables $\bar{\mathbf{X}},\hat{\mathbf{X}}\in\mathcal{X}$ by

2 Discriminator

As in the GAN framework, we introduce a discriminator, $D$, that will be used as an adversary to train $G$. However, unlike in a standard GAN where the output of the generator is either completely real or completely fake, in this setting the output is comprised of some components that are real and some that are fake. Rather than identifying that an entire vector is real or fake, the discriminator attempts to distinguish which components are real (observed) or fake (imputed) - this amounts to predicting the mask vector, $\mathbf{m}$. Note that the mask vector $\mathbf{M}$ is pre-determined by the dataset.

Formally, the discriminator is a function $D:\mathcal{X}\to^{d}$ with the $i$-th component of $D(\hat{\mathbf{x}})$ corresponding to the probability that the $i$-th component of $\hat{\mathbf{x}}$ was observed.

3 Hint

As will be seen in the theoretical results that follow, it is necessary to introduce what we call a hint mechanism. A hint mechanism is a random variable, $\mathbf{H}$, taking values in a space $\mathcal{H}$, both of which we define. We allow $\mathbf{H}$ to depend on $\mathbf{M}$ and for each (imputed) sample $(\hat{\mathbf{x}},\mathbf{m})$, we draw $\mathbf{h}$ according to the distribution $\mathbf{H}|\mathbf{M}=\mathbf{m}$. We pass $\mathbf{h}$ as an additional input to the discriminator and so it becomes a function $D:\mathcal{X}\times\mathcal{H}\to^{d}$, where now the $i$-th component of $D(\hat{\mathbf{x}},\mathbf{h})$ corresponds to the probability that the $i$-th component of $\hat{\mathbf{x}}$ was observed conditional on $\hat{\mathbf{X}}=\hat{\mathbf{x}}$ and $\mathbf{H}=\mathbf{h}$.

By defining $\mathbf{H}$ in different ways, we control the amount of information contained in $\mathbf{H}$ about $\mathbf{M}$ and in particular we show (in Proposition 1) that if we do not provide “enough” information about $\mathbf{M}$ to $D$ (such as if we simply did not have a hinting mechanism), then there are several distributions that $G$ could reproduce that would all be optimal with respect to $D$.

4 Objective

We train $D$ to maximize the probability of correctly predicting $\mathbf{M}$. We train $G$ to minimize the probability of $D$ predicting $\mathbf{M}$. We define the quantity $V(D,G)$ to be

where $\log$ is element-wise logarithm and dependence on $G$ is through $\hat{\mathbf{X}}$.

Then, as with the standard GAN, we define the objective of GAIN to be the minimax problem given by

Writing $\hat{\mathbf{M}}=D(\hat{\mathbf{X}},\mathbf{H})$, we can then rewrite (5) as

Theoretical Analysis

In this section we provide a theoretical analysis of (5). Given a d-dimensional space $\mathcal{Z}=\mathcal{Z}_{1}\times...\times\mathcal{Z}_{d}$, a (probability) density For ease of exposition, we use the term density even when referring to a probability mass function. $p$ over $\mathcal{Z}$ corresponding to a random variable $Z$, and a vector $\mathbf{b}\in\{0,1\}^{d}$ we define the set $A_{\mathbf{b}}=\{i:b_{i}=1\}$, the projection $\phi_{\mathbf{b}}:\mathcal{Z}\to\Pi_{i\in A_{\mathbf{b}}}\mathcal{Z}_{i}$ by $\phi_{\mathbf{b}}(z)=(z_{i})_{i\in A}$ and the density $p^{\mathbf{b}}$ to be the density of $\phi_{\mathbf{b}}(Z)$.

Throughout this section, we make the assumption that $\mathbf{M}$ is independent of $\mathbf{X}$, i.e. that the data is MCAR.

We will write $p(\mathbf{x},\mathbf{m},\mathbf{h})$ to denote the density of the random variable $(\hat{\mathbf{X}},\mathbf{M},\mathbf{H})$ and we will write $\hat{p}$, $p_{m}$ and $p_{h}$ to denote the marginal densities (of $p$) corresponding to $\hat{\mathbf{X}}$, $\mathbf{M}$ and $\mathbf{H}$, respectively. When referring to the joint density of two of the three variables (potentially conditioned on the third), we will simply use $p$, abusing notation slightly.

It is more intuitive to think of this density through its decomposition into densities corresponding to the true data generating process, and to the generator defined by (2),

Let $\mathbf{x}\in\mathcal{X}$. Let $p_{h}$ be a fixed density over the hint space $\mathcal{H}$ and let $\mathbf{h}\in\mathcal{H}$ be such that $p(\mathbf{x},\mathbf{h})>0$. Then for a fixed generator, $G$, the $i$-th component of the optimal discriminator, $D^{*}(\mathbf{x},\mathbf{h})$ is given by

All proofs are provided in Supplementary Materials. ∎

We now rewrite (4), substituting for $D^{*}$, to obtain the following minimization criterion for $G$:

where dependence on $G$ is through $p_{m}(\cdot|\hat{\mathbf{X}})$.

A global minimum for $C(G)$ is achieved if and only if the density $\hat{p}$ satisfies

for each $i\in\{1,...,d\}$, $\mathbf{x}\in\mathcal{X}$ and $\mathbf{h}\in\mathcal{H}$ such that $p_{h}(\mathbf{h}|m_{i}=t)>0$.

The following proposition asserts that if $\mathbf{H}$ does not contain “enough” information about $\mathbf{M}$, we cannot guarantee that $G$ learns the desired distribution (the one uniquely defined by the (underlying) data).

There exist distributions of $\mathbf{X}$, $\mathbf{M}$ and $\mathbf{H}$ for which solutions to (12) are not unique. In fact, if $\mathbf{H}$ is independent of $\mathbf{M}$, then (12) does not define a unique density, in general.

Let the random variable $\mathbf{B}=(B_{1},...,B_{d})\in\{0,1\}^{d}$ be defined by first sampling $k$ from $\{1,...,d\}$ uniformly at random and then setting

Let $\mathcal{H}=\{0,0.5,1\}^{d}$ and, given $\mathbf{M}$, define

Observe first that $\mathbf{H}$ is such that $H_{i}=t\implies M_{i}=t$ for $t\in\{0,1\}$ but that $H_{i}=0.5$ implies nothing about $M_{i}$. In other words, $\mathbf{H}$ reveals all but one of the components of $\mathbf{M}$ to $D$. Note, however, that $\mathbf{H}$ does contain some information about $M_{i}$ since $M_{i}$ is not assumed to be independent of the other components of $\mathbf{M}$.

The following lemma confirms that the discriminator behaves as we expect with respect to this hint mechanism.

Suppose $\mathbf{H}$ is defined as above. Then for $\mathbf{h}$ such that $h_{i}=0$ we have $D^{*}(\mathbf{x},\mathbf{h})_{i}=0$ and for $\mathbf{h}$ such that $h_{i}=1$ we have $D^{*}(\mathbf{x},\mathbf{h})_{i}=1$, for all $\mathbf{x}\in\mathcal{X}$, $i\in\{1,...,d\}$.

The final proposition we state tells us that $\mathbf{H}$ as specified above ensures the generator learns to replicate the desired distribution.

Suppose $\mathbf{H}$ is defined as above. Then the solution to (12) is unique and satisfies

for all $\mathbf{m}_{1},\mathbf{m}_{2}\in\{0,1\}^{d}$. In particular, $\hat{p}(\mathbf{x}|\mathbf{m})=\hat{p}(\mathbf{x}|\mathbf{1})$ and since $\mathbf{M}$ is independent of $\mathbf{X}$, $\hat{p}(\mathbf{x}|\mathbf{1})$ is the density of $\mathbf{X}$. The distribution of $\hat{\mathbf{X}}$ is therefore the same as the distribution of $\mathbf{X}$.

For the remainder of the paper, $\mathbf{B}$ and $\mathbf{H}$ will be defined as in equations (13) and (14).

GAIN Algorithm

Using an approach similar to that in (Goodfellow et al. 2014), we solve the minimax optimization problem (5) in an iterative manner. Both $G$ and $D$ are modeled as fully connected neural nets.

recalling that $\hat{\mathbf{m}}(j)=D(\hat{\mathbf{x}}(j),\mathbf{m}(j))$.

As can be seen from their definitions, $\mathcal{L}_{G}$ will apply to the missing components ($m_{i}=0$) and $\mathcal{L}_{M}$ will apply to the observed components ($m_{i}=1$).

$G$ is then trained to minimize the weighted sum of the two losses as follows:

The pseudo-code is presented in Algorithm 1.

Experiments

In this section, we validate the performance of GAIN using multiple real-world datasets. In the first set of experiments we qualitatively analyze the properties of GAIN. In the second we quantitatively evaluate the imputation performance of GAIN using various UCI datasets (Lichman 2013), giving comparisons with state-of-the-art imputation methods. In the third we evaluate the performance of GAIN in various settings (such as on datasets with different missing rates). In the final set of experiments we evaluate GAIN against other imputation algorithms when the goal is to perform prediction on the imputed dataset.

We conduct each experiment 10 times and within each experiment we use 5-cross validations. We report either RMSE or AUROC as the performance metric along with their standard deviations across the 10 experiments. Unless otherwise stated, missingness is applied to the datasets by randomly removing 20% of all data points (MCAR).

The potential sources of gain for the GAIN framework are: the use of a GAN-like architecture (through $\mathcal{L}_{G}$), the use of reconstruction error in the loss ($\mathcal{L}_{M}$), and the use of the hint ($\mathbf{H}$). In order to understand how each of these affects the performance of GAIN, we exclude one or two of them and compare the performances of the resulting architectures against the full GAIN architecture.

Table 1 shows that the performance of GAIN is improved when all three components are included. More specifically, the full GAIN framework has a 15% improvement over the simple auto-encoder model (i.e. GAIN w/o $\mathcal{L}_{G}$). Furthermore, utilizing the hint vector additionally gives improvements of 10%.

2 Quantitative analysis of GAIN

We use five real-world datasets from UCI Machine Learning Repository (Lichman 2013) (Breast, Spam, Letter, Credit, and News) to quantitatively evaluate the imputation performance of GAIN. Details of each dataset can be found in the Supplementary Materials.

In table 2 we report the RMSE (and its standard deviation) for GAIN and 5 other state-of-the-art imputation methods: MICE (Buuren & Oudshoorn 2000; Buuren & Groothuis-Oudshoorn 2011), MissForest (Stekhoven & Bühlmann 2011), Matrix completion (Matrix) (Mazumder et al. 2010a), Auto-encoder (Gondara & Wang 2017) and Expectation-maximization (EM) (García-Laencina et al. 2010). As can be seen from the table, GAIN significantly outperforms each benchmark. Results for the imputation quality of categorical variables in this experiment are given in the Supplementary Materials.

3 GAIN in different settings

To better understand GAIN, we conduct several experiments in which we vary the missing rate, the number of samples, and the number of dimensions using Credit dataset. Fig. 2 shows the performance (RMSE) of GAIN within these different settings in comparison to the two most competitive benchmarks (MissForest and Auto-encoder). Fig. 2 (a) shows that, even though the performance of each algorithm decreases as missing rates increase, GAIN consistently outperforms the benchmarks across the entire range of missing rates.

Fig. 2 (b) shows that as the number of samples increases, the performance improvements of GAIN over the benchmarks also increases. This is due to the large number of parameters in GAIN that need to be optimized, however, as demonstrated on the Breast dataset (in Table 2), GAIN is still able to outperform the benchmarks even when the number of samples is relatively small.

Fig. 2 (c) shows that GAIN is also robust to the number of feature dimensions. On the other hand, the discriminative model (MissForest) cannot as easily cope when the number of feature dimensions is small.

4 Prediction Performance

We now compare GAIN against the same benchmarks with respect to the accuracy of post-imputation prediction. For this purpose, we use Area Under the Receiver Operating Characteristic Curve (AUROC) as the measure of performance. To be fair to all methods, we use the same predictive model (logistic regression) in all cases.

Comparisons are made on all datasets except Letter (as it has multi-class labels) and the results are reported in Table 3.

As Table 3 shows, GAIN, which we have already shown to achieve the best imputation accuracy (in Table 2), yields the best post-imputation prediction accuracy. However, even in cases where the improvement in imputation accuracy is large, the improvements in prediction accuracy are not always significant. This is probably due to the fact that there is sufficient information in the (80%) observed data to predict the label.

Prediction accuracy with various missing rates: In this experiment, we evaluate the post-imputation prediction performance when the missing rate of the dataset is varied. Note that every dataset (except Letter) has their own binary label.

The results of this experiment (for GAIN and the two most competitive benchmarks) are shown in Fig. 3. In particular, the performance of GAIN is significantly better than the other two for higher missing rates, this is due to the fact that as the information contained in the observed data decreases (due to more values being missing), the imputation quality becomes more important, and GAIN has already been shown to provide (significantly) better quality imputations.

5 Congeniality of GAIN

The congeniality of an imputation model is its ability to impute values that respect the feature-label relationship (Meng 1994; Burgess et al. 2013; Deng et al. 2016). The congeniality of an imputation model can be evaluated by measuring the effects on the feature-label relationships after the imputation. We compare the logistic regression parameters, $\mathbf{w}$, learned from the complete Credit dataset with the parameters, $\hat{\mathbf{w}}$, learned from an incomplete Credit dataset by first imputing and then performing logistic regression.

We report the mean and standard deviation of both the mean bias $(||\mathbf{w}-\hat{\mathbf{w}}||_{1})$ and the mean square error $(||\mathbf{w}-\hat{\mathbf{w}}||_{2})$ for each method in Table 4. These quantities being lower indicates that the imputation algorithm better respects the relationship between feature and label. As can be seen in the table, GAIN achieves significantly lower mean bias and mean square error than other state-of-the-art imputation algorithms (from 8.9% to 79.2% performance improvements).

Conclusion

We propose a generative model for missing data imputation, GAIN. This novel architecture generalizes the well-known GAN such that it can deal with the unique characteristics of the imputation problem. Various experiments with real-world datasets show that GAIN significantly outperforms state-of-the-art imputation techniques. The development of a new, state-of-the-art technique for imputation can have transformative impacts; most datasets in medicine as well as in other domains have missing data. Future work will investigate the performance of GAIN in recommender systems, error concealment as well as in active sensing (Yu et al. 2009). Preliminary results in error concealment using the MNIST dataset (LeCun & Cortes 2010) can be found in the Supplementary Materials - see Fig. 4 and 5.

Acknowledgement

The authors would like to thank the reviewers for their helpful comments. The research presented in this paper was supported by the Office of Naval Research (ONR) and the NSF (Grant number: ECCS1462245, ECCS1533983, and ECCS1407712).

References