1 Introduction↩︎

Missing values occur in real-world datasets for various reasons, such as non-response in surveys and sensor failures. Imputation — filling in missing values with their estimates — is a common preprocessing step used to address missing data. Over the decades, various imputation methods have been proposed, ranging from simple statistical techniques to machine learning algorithms e.g., [1], [2]–[6].

kNNImputer [7] is one of the most widely used imputation methods, owing to its simplicity and availability in popular software packages such as scikit-learn² [8]. It imputes a missing response variable (e.g., customer satisfaction level) of a given unit (e.g., a customer) as the average of the observed responses of the \(k\) most similar units to the given unit in terms of observed covariates (e.g., age, gender, occupation). This is to predict the missing response by \(k\) nearest neighbours (kNN) regression [9] so the imputation is an estimate of the conditional expectation of the missing response given a covariate. The method has been widely used in science and engineering, and many extensions have been proposed [10]–[14].

An issue of kNNImputer, shared by other regression-based imputers, is that the distribution of imputations can be significantly different from the distribution of true (hidden) missing values. This is because, as mentioned, an imputation of kNNImputer is an estimate of the conditional expectation of a missing response, thus tending to be a deterministic function of the covariate. As a result, the distribution of imputed responses is concentrated around the regression curve, even when the distribution of missing responses has large variability. This is illustrated in Figures 1 and 2, where the true conditional distribution of a missing response is bimodal when the covariate is small, but the distribution of imputations is unimodal and many imputations take values never realized by the true missing values. A substantial bias can occur in an analysis of such a distorted imputed dataset, for example, when estimating the variance, quantiles and modes in the population.

The above issue of kNNImputer may be addressed by estimating the conditional distribution of a missing response given a covariate, and randomly sampling imputations from it. This idea was investigated by [15], who proposed the “kNN\(\times\)KDE” approach that combines a soft version of kNN and kernel density estimation (KDE). For a given unit, the conditional density of a missing response is estimated as a weighted average of Gaussian densities centered at observed responses, where the weights are computed so that units more similar, in terms of covariates, to the given unit receive larger weights. kNN\(\times\)KDE was demonstrated to have good empirical performance in recovering the distribution of missing values, compared to established imputation methods, including kNNImputer, missForest [16], SoftImpute [17], and Gain [18]. However, no theoretical guarantee exists for kNN\(\times\)KDE, such as its statistical consistency, i.e., whether the estimated conditional density converges to the true one as the sample size increases. Consistency is not only important as a minimal theoretical guarantee but also in understanding how hyperparameters should be chosen. While kNN\(\times\)KDE has two main hyperparameters (the “inverse temperature” in the softmax function used for weight computations, and the variance of Gaussian densities), no systematic selection procedure was proposed.

This paper studies a simpler kNN-based stochastic imputation method named kNNSampler. For a given unit whose response is missing, it estimates the conditional distribution of the missing response given the unit’s observed covariate as the empirical distribution of the observed responses of the \(k\) most similar units to that unit in terms of covariates; an imputation is randomly sampled from this empirical distribution, which we call kNN conditional distribution. kNNSampler is as simple as kNNImputer: instead of taking the mean of the observed responses of \(k\) nearest neighbours, kNNSampler simply samples one of those \(k\) observed responses. It is thus simpler than kNN\(\times\)KDE as it does not involve an intermediate step of density estimation and is free of any hyperparameter for responses. The number \(k\) of nearest neighbours in kNNSampler can be efficiently chosen by leave-one-out cross validation using the fast computation method recently proposed by [19]. Figures 1 and 2 describe imputations by kNNSampler with \(k\) selected in this way, which align much better with the distribution of true missing values than imputations by kNNImputer. More systematic experiments are provided in Section 4.

kNNSampler can be interpreted as an instance of hot deck, classic imputation methods widely used in practice for socio-economic and public health surveys, including the U.S. Census Bureau’s Current Population Survey and the National Center for Education Statistics [20]. In a hot deck method, a missing value of a given unit is imputed as one of the response values of the units belonging to the same “adjustment cell” as the given unit. The method is called random hot deck if the imputation is selected randomly from the adjustment cell; it is called nearest-neighbour hot deck if nearest neighbours define the adjustment cell [4]. kNNSampler is thus essentially a nearest-neighbour random hot deck method. However, while classic and widely used, hot deck methods have not been well established theoretically [20].

Our contribution is to establish kNNSampler, and thus the nearest-neighbour random hot deck, as a theoretically principled missing-value imputation method. To this end, we analyze the kNN conditional distribution, i.e., the empirical distribution of \(k\) nearest neighbour responses from which an imputation is sampled, as an estimator of the true conditional distribution of a missing response given a covariate (Section 3). Our theoretical contributions are summarized as follows.

• We derive an error bound between the kNN and true conditional distributions for any given, fixed covariate, in terms of the number \(n\) of observed response-covariate pairs, the number \(k\) of the nearest neighbours, and other problem-specific constants. The error is measured by the maximum mean discrepancy (MMD) [21], a distance metric on probability distributions that metrizes the weak convergence [22], between the kNN and true conditional distributions. It holds under a Lipschitz condition that the response’s conditional distribution changes smoothly when the covariate changes continuously. A consequence of the bound is the statistical consistency of the kNN conditional distribution, in that the error decreases to zero as the sample size \(n\) goes to infinity, if the number \(k\) of nearest neighbours increases to infinity at a rate slower than \(n\). This offers a theoretical foundation of the kNNSampler and thus the nearest-neighbour random hot deck.

• To derive the bound, we analyze the mean embedding of the kNN conditional distribution in a reproducing kernel Hilbert space (RKHS) as a novel estimator of the mean embedding of the true conditional distribution, known as conditional mean embedding [23], which is the RKHS-valued regression function [24]. The RKHS distance between these two embeddings is equivalent to the MMD between the kNN and true conditional distributions. Our bound leads to the consistency and convergence rates for the novel kNN-based estimator of the conditional mean embedding.

• Our analysis extends the error analysis by [25] on real-valued kNN regression to RKHS-valued regression in which the response variable is infinite-dimensional. As a byproduct, we prove that the required sample size to attain a given level of precision increases exponentially not with the covariate’s ambient dimension but with the intrinsic dimension of the covariate distribution. Therefore, the kNNSampler may not be severely affected by the curse of dimensionality if the covariate distribution has a low intrinsic dimension.

This paper is organised as follows. We describe the proposed approach in Section 2, its theory in Section 3, and experiments in Section 4.

2 Proposed Approach↩︎

This section describes the proposed approach. Section 2.1 introduces the setting. Section 2.2 explains the kNNImputer and its issue as a preliminary. We describe kNNSampler in Section 2.3, uncertainty quantification with the kNN conditional distribution in Section 2.4, and multiple imputation with kNNSampler in Section 2.5.

2.1 Setting↩︎

We first describe the problem setup. Let \(\mathcal{X}\) and \(\mathcal{Y}\) be measurable spaces representing the covariate space and the response space, respectively. For example, the covariate space may be the \(d\)-dimensional Euclidean space, \(\mathcal{X}= \mathbb{R}^d\), in which case a covariate \(x \in \mathcal{X}\) consists of \(d\) features (e.g., a person’s age, weight, height), and the response space may be the real line \(\mathcal{Y}= \mathbb{R}\), in which case a response \(y \in \mathcal{Y}\) is real-valued (e.g., the person’s blood pressure).

We assume that our dataset consists of \(n + m\) units (e.g., persons), where \(n\) units have both covariate \(x_i \in \mathcal{X}\) and response \(y_i \in \mathcal{Y}\) observed, while \(m\) units have only covariate \(\tilde{x}_j \in \mathcal{X}\) observed and response \(\tilde{y}_{{\rm miss}, j} \in \mathcal{Y}\) missing: \[\begin{align} \mathcal{D}_n := \{(x_1, y_1), \dots, (x_n, y_n)\}, \quad \mathcal{D}_{\rm miss}:= \{ ( \tilde{x}_1, \tilde{y}_{1, {\rm miss}}), \dots, (\tilde{x}_{m}, \tilde{y}_{m, {\rm miss}}) \} \label{eq:missing-dataset} \end{align}\tag{1}\]

For each of the \(n\) units with observed responses, we assume that the covariate follows a marginal distribution \(P(x)\) and the response given the covariate follows the conditional distribution \(P(y|x)\) in an independently and identically distributed (i.i.d.) manner: \[\label{eq:observed-pairs-dist} (x_1, y_1), \dots, (x_n, y_n) \stackrel{i.i.d.}{\sim} P(y|x)P(x)\tag{2}\]

On the other hand, for the \(m\) units with missing responses, the covariate is assumed to follow a marginal distribution \(Q(\tilde{x})\), which can be different from \(P(x)\), while the conditional distribution of the missing response given the covariate remains the same: \[\label{eq:missing-pairs-joint-dist} ( \tilde{x}_1, \tilde{y}_{1, {\rm miss}}), \dots, (\tilde{x}_{m}, \tilde{y}_{m, {\rm miss}}) \stackrel{i.i.d.}{\sim} P(\tilde{y}_{\rm miss}|\tilde{x})Q(\tilde{x}).\tag{3}\] This assumption implies that the probability of a unit missing its response is determined by the unit’s covariate and is not affected by the response. Therefore, it is an instance of the Missing-At-Random (MAR) assumption [1]. In the special case where the covariate distributions for the two cases are the same, \(Q(\tilde{x}) = P(\tilde{x})\), the assumption can be interpreted as the Missing-Completely-At-Random (MCAR) assumption where missingness occurs completely randomly.

Under this setup, missing responses may be imputed by estimating the unknown conditional distribution \(P(y|x)\) of a response given a covariate, and sampling from the estimated conditional distribution. This is what the kNNSampler does.

2.2 Issue with kNNImputer and Regression-based Imputers↩︎

Before describing the proposed kNNSampler, we discuss an issue with the widely used kNNImputer [7] and other regression-based imputation methods.

Suppose that the covariate space \(\mathcal{X}\) is equipped with a distance metric \(d_{\mathcal{X}}(x,x')\) that quantifies the distance between any two points \(x, x' \in \mathcal{X}\). For example, if \(\mathcal{X}\) is the Euclidean space, then \(d_{\mathcal{X}}(x,x')\) may be the Euclidean distance between two vectors \(x\) and \(x'\). Let \(X_n\) be the set of covariates for the \(n\) units with observed responses: \[X_n := \{ x_1, \dots, x_n \}\] For a given covariate \(\tilde{x}\) and a number \(k\) of nearest neighbours, let \({\rm NN}(\tilde{x}, k, X_n)\) be the indices of the \(k\) units whose covariates are the most similar to \(\tilde{x}\) in terms of the distance metric among the \(n\) units with observed responses:³ \[\begin{align} \label{eq:nearest-neighbours} {\rm NN}(\tilde{x}, k, X_n) := \{j_1, \ldots, j_k \in \{1, \ldots, n\} \mid & ~~ d_\mathcal{X}(\tilde{x}, x_{j_1}) \le \ldots \le d_\mathcal{X}(\tilde{x}, x_{j_k})\\ &\le d_\mathcal{X}(\tilde{x}, x_j) \text{ for all } j \in \{1, \ldots, n\} \setminus \{j_1, \ldots, j_k\} \}. \nonumber \end{align}\tag{4}\] That is, \({\rm NN}(\tilde{x}, k, X_n)\) is the indices of the \(k\) nearest neighbours of \(\tilde{x}\) in \(X_n\).

kNNImputer [7] imputes the missing response \(\tilde{y}_{i, {\rm miss}}\) of the unit with observed covariate \(\tilde{x}_i\) as the average of the observed responses \(y_{j_1}, \dots, y_{j_k}\) of its \(k\)-nearest neighbors \(x_{j_1}, \dots, x_{j_k}\): \[\hat{y}_{i, {\rm imp}} = \frac{1}{k} \sum_{j \in {\rm NN}(\tilde{x}_i, k, X_n)} y_j.\] This is kNN regression [26] and thus estimates the conditional mean of the missing response \(\tilde{y}_{i, {\rm miss}}\) given the observed covariate \(\tilde{x}_i\): \[\hat{y}_{i, {\rm imp}} \approx f(\tilde{x}_i) := \int \tilde{y} ~dP(\tilde{y} | \tilde{x_i}),\] where \(f: \mathcal{X}\to \mathcal{Y}\) is the regression function. In this case, the observed covariate and the imputed response \((\tilde{x}_i, \hat{y}_{i, {\rm imp}} )\) approximately follow the degenerate joint distribution \[\delta(\tilde{y}-f(\tilde{x}))Q(\tilde{x}),\] where \(\delta(\tilde{y}-f(\tilde{x}))\) denotes the Dirac distribution at the conditional mean \(f(\tilde{x})\), i.e., the degenerate distribution whose mass concentrates at \(f(\tilde{x})\). This differs from the joint distribution of the observed covariate and the true missing response \((\tilde{x}_i, \tilde{y}_{i, {\rm miss}})\): \[\label{eq:true-dist-missing} P(\tilde{y} \mid \tilde{x})Q(\tilde{x})\tag{5}\] unless the conditional distribution \(P(\tilde{y} \mid \tilde{x})\) is the Dirac distribution \(\delta(\tilde{y}-f(\tilde{x}))\), i.e., unless the missing response is the deterministic function of observed covariate. The same issue occurs with other single imputation methods based on regression, because they impute the missing response by estimating the conditional mean.

To summarize, kNNImputer and other regression-based imputation methods do not generally recover the true distribution of the missing data. An analysis based on the imputed dataset may lead to a biased result. For instance, the variance of the imputed values may be much lower than the variance of the true missing values. kNNSampler alleviates this issue by imputing missing values by estimating the conditional distribution \(P(\tilde{y}\mid \tilde{x})\).

2.3 kNNSampler↩︎

We now describe kNNSampler (Algorithm 3). Consider imputing the missing response \(\tilde{y}_{{\rm miss}}\) of a unit with observed covariate \(\tilde{x}\). kNNSampler estimates the conditional distribution \(P(\tilde{y}_{{\rm miss}} \mid \tilde{x})\) of \(\tilde{y}_{{\rm miss}}\) given \(\tilde{x}\) as the empirical distribution of the observed responses \(y_{j_1}, \dots, y_{j_k}\) of the \(k\) nearest neighbours \(x_{j_1}, \dots, x_{j_k}\) of \(\tilde{x}\): \[\label{eq:cond-dist-est} P(\tilde{y}_{{\rm miss}} \mid \tilde{x}) \approx \hat{P}(\tilde{y}_{{\rm miss}} \mid \tilde{x}) := \frac{1}{k} \sum_{j \in {\rm NN}(\tilde{x}, k, X_n)} \delta(\tilde{y}_{{\rm miss}} - y_j),\tag{6}\] which is the discrete distribution where each of \(y_{j_1}, \dots, y_{j_k}\) has probability mass \(1/k\). An imputation \(\hat{y}_{{\rm imp}}\) for the missing response is randomly sampled from this empirical distribution: \[\hat{y}_{{\rm imp}} \sim \hat{P}(\tilde{y}_{{\rm miss}} \mid \tilde{x}).\] Algorithmically, this is to randomly sample one of the kNN observed responses \(y_{j_1}, \dots, y_{j_k}\). Algorithm 3 independently applies this procedure to the observed covariate \(\tilde{x}_i\) to generate an imputation \(\hat{y}_{i, {\rm imp}}\) of missing value \(y_{i, {\rm miss}}\) for each unit \(i = 1, \dots, m\).

2.3.0.1 Choice of \(k\)

The number of nearest neighbors \(k\) is a hyperparameter of kNNSampler. The theoretical and empirical results below indicate that \(k\) should not be fixed to a prespecified value (e.g., \(k=5\)), and should be chosen depending on the available data. One way is to perform cross-validation for kNN regression on the data \((x_1, y_1), \dots, (x_n, y_n)\) and select \(k\) among candidate values that minimizes the mean-square error on held-out observed responses, averaged over different training-validation splits. In particular, the present work uses Leave-One-Out Cross-Validation (LOOCV) using the fast computation method recently proposed by [19].

2.4 Uncertainty Quantification of Missing Values↩︎

Quantifying the uncertainty in missing values is important for several reasons, including assessing the reliability of imputations and the adequacy of the covariates used, as well as determining how to perform imputations (e.g., single or multiple) and how to use the imputations in subsequent analyses. We describe here how to perform uncertainty quantification of missing values with the kNN conditional distribution.

2.4.0.1 Conditional Probability Estimation

kNNSampler can be used to estimate the conditional probability of a missing response \(\tilde{y}_{\rm miss}\) belonging to a specified (measurable) subset \(S\) of the response space \(\mathcal{Y}\), given observed covariate \(\tilde{x}\): \[{\rm Pr}( \tilde{y}_{\rm miss}\in S \mid \tilde{x}) = \int ~\mathbb{I}[\tilde{y}\in S]~ dP( \tilde{y}\mid \tilde{x}),\] where \(\mathbb{I}[\tilde{y}\in S]\) is the indicator function that outputs \(1\) if \(\tilde{y}\in S\) and \(0\) otherwise. By replacing the unknown conditional distribution \(P(\tilde{y}~|~ \tilde{x})\) by the kNN conditional distribution \(\hat{P}(\tilde{y}~|~ \tilde{x})\) in (6 ), this conditional probability is approximated as \[\widehat{{\rm Pr}}( \tilde{y}_{\rm miss}\in S \mid \tilde{x}) = \int ~\mathbb{I}[\tilde{y}\in S]~ d\hat{P}( \tilde{y}\mid \tilde{x}) = \frac{1}{k} \sum_{j \in {\rm NN}(\tilde{x}, k, X_n)} ~\mathbb{I} [y_j \in S].\] In other words, the conditional probability is estimated as the observed frequency of the kNN response values that fall in \(S\).

2.4.0.2 Interval Estimation

Let us focus on a real-valued missing response \(\tilde{y}_{{\rm miss}} \in \mathcal{Y} = \mathbb{R}\). The conditional probability of the missing response belonging to a given (finite or infinite) interval \(S = (\ell,u)\), where \(\ell< u\), is estimated as the observed frequency of the k-NN responses belonging to that interval. This indicates that an interval to which the kNN responses belongs at a specified frequency \(0< 1-\alpha < 1\) (e.g., \(\alpha = 0.05\), in which case the 95% of the kNN responses belong to the interval) is an estimate of an interval to which the unknown missing response belongs at that probability \(1 - \alpha\).

Such an interval \((\ell, u)\) is constructed by defining its lower bound \(\ell\) and upper bound \(u\) as, respectively, the lower and upper \(\alpha/2\) empirical quantiles of the kNN responses, i.e., the \(k \alpha/2\)-smallest and the \(k \alpha/2\)-largest kNN responses (e.g., if \(k = 200\) and \(\alpha = 0.05\), the 5th smallest and the 5th largest kNN responses): \[{\rm Pr}( \ell < \tilde{y}_{{\rm miss}} < u \mid \tilde{x}) \approx 1 - \alpha\]

2.4.0.3 Conditional Standard Deviation Estimation

The conditional standard deviation of a missing response given observed covariate quantifies the variability of the missing response. This can be estimated by the empirical standard deviation of the kNN response values for the observed covariate.

2.5 Multiple Imputation with kNNSampler↩︎

kNNSampler can be used for multiple imputation by independently generating multiple imputed datasets. More precisely, let \(B\) be the number of multiple imputed datasets to be generated (e.g., \(B = 10\)). For each \(b = 1, \dots, B\), kNNSampler is independently applied to impute the missing responses in the dataset \(\mathcal{D}_{\rm miss}\) (1 ) to create an imputed dataset \[\mathcal{D}_{n+m}^{(b)} := \mathcal{D}_n \cup \mathcal{D}_{\rm imp}^{(b)} \quad \text{where} \quad \mathcal{D}_{\rm imp}^{(b)} := \{ ( \tilde{x}_1, \tilde{y}_{1, {\rm imp}}^{(b)}), \dots, (\tilde{x}_{m}, \tilde{y}_{m, {\rm imp}}^{(b)}) \},\] where \(\tilde{y}_{i, {\rm imp}}^{(b)}\) is an imputation for the \(i\)-th unit with a missing response \(\tilde{y}_{i,{\rm miss}}\) covariates \(\tilde{x}_i\). This results in \(B\) imputed datasets: \[\mathcal{D}_{n+m}^{(1)} , \dots, \mathcal{D}_{n+m}^{(B)}.\] An analysis can then be made based on the standard procedure of multiple imputation [27].

For example, suppose that we want to estimate a population quantity \(\theta^*\) (e.g., the mean customer satisfaction level of a population). Let \(S_{n+m}\) be a function of a dataset of size \(n+m\) that outputs an estimate \(\hat{\theta}_{n+m}\) of the unknown \(\theta^*\) (e.g., the empirical average of \(n+m\) values): \(\hat{\theta}_{n+m} = S(\mathcal{D}_{n+m})\). Apply this function to each of the \(B\) imputed datasets, one obtains \(B\) estimates of \(\theta^*\): \[\hat{\theta}^{(b)}_{n+m} = S(\mathcal{D}_{n+m}^{(b)}), \quad b = 1, \dots, B.\] The empirical average of these \(B\) estimates gives a multiple-imputation estimate of \(\theta^*\). The empirical standard deviation of the \(B\) estimates \(\hat{\theta}_{n+m}^{(1)}, \dots, \hat{\theta}_{n+m}^{(B)}\) quantifies the uncertainty due to the missingness in the original data. Combined with the standard error of each \(\hat{\theta}_{n+m}^{(b)}\), this standard deviation can be used to quantify the overall uncertainty of the estimate using Rubin’s rule.

3 Theory↩︎

We describe a theory for kNNSampler’s conditional distribution (6 ) as an estimator of the true conditional distribution. We shall show that, as the number \(k\) of nearest neighbors increases at an approximate rate as the increase of the number \(n\) of observed covariate-response pairs, the kNN conditional distribution converges to the true one in the Maximum Mean Discrepancy (MMD) [21], which implies the convergence in distribution [28]. We prove this by adapting the proof of the convergence rates of real-valued kNN regression by [25] to Hilbert space-valued kNN regression.⁴

We use the framework of kernel mean embedding [23] in which every probability distribution is represented as a distinct point in an infinite-dimensional feature space known as a reproducing kernel Hilbert space (RKHS). The true and kNN conditional distributions are represented as points in an RKHS, and the distance between them, which is the MMD, quantifies the estimation error. An upper bound on this distance is obtained in terms of the sample size, the number of nearest neighbours, and other relevant quantities.

3.1 RKHS Embeddings of Conditional Distributions↩︎

Let us first define an RKHS on the response space \(\mathcal{Y}\). As before, \(\mathcal{Y}\) is a measurable space such as the \(p\)-dimensional Euclidean space, \(\mathcal{Y}= \mathbb{R}^p\). A Hilbert space⁵ \(\mathcal{H}\) consisting of functions \(f\) on \(\mathcal{Y}\) is called RKHS if there exists a map \[\Phi: \mathcal{Y}\to \mathcal{H}\] called feature map, such that the value \(f(y)\) of any function \(f\) in \(\mathcal{H}\) at any point \(y\) in \(\mathcal{Y}\) can be written as the inner product between \(f\) and the feature map \(\Phi(y)\) of \(y\): \[f \in \mathcal{H}\quad \Longleftrightarrow \quad f(y) = \left< f, \Phi(y) \right>_{\mathcal{H}} \;\text{for all }\;y \in \mathcal{Y},\] where \(\left< \cdot, \cdot \right>_{\mathcal{H}}\) denotes the inner product of \(\mathcal{H}\). The \(\Phi(y)\) may be called feature vector of \(y\), and \(\mathcal{H}\) the feature space, which can be infinite-dimensional.

The inner product between the feature maps \(\Phi(y), \Phi(y')\) of any two points \(y, y'\) defines the kernel function \[\label{eq:kernel-inner-prod} \ell(y, y') := \left< \Phi(y), \Phi(y') \right>_{\mathcal{H}} \quad \text{for all }\;y, y', \in \mathcal{Y}.\tag{7}\] This is called reproducing kernel of the RKHS. The RKHS and the reproducing kernel are one-to-one, so an RKHS can be induced by defining a kernel. For example, if \(\mathcal{Y}= \mathbb{R}^p\), the Gaussian kernel \(\ell(y, y') = \exp( - \alpha \| y - y'\|^2 )\) for \(\alpha > 0\) is the reproducing kernel of a certain RKHS \(\mathcal{H}\), and there exists an infinite-dimensional feature map \(\Phi\) that induces the Gaussian kernel as (7 ). See e.g. [30], [31] for details on RKHSs.

Every probability distribution \(P\) on \(\mathcal{Y}\) is represented as the expected feature map: \[\Phi(P) := \int \Phi(y)dP(y) \in \mathcal{H}.\] This is called mean embedding of \(P\). If the RKHS \(\mathcal{H}\) is large enough, any two different probability distributions \(P\) and \(Q\) are mapped to two distinct mean embeddings: \[P \not= Q \quad \Longleftrightarrow \quad \Phi(P) \not= \Phi(Q).\] In this case, the RKHS is called characteristic [28]. For example, Gaussian, Matérn and Laplace kernels induce characteristics RKHSs.

The true and kNN conditional distributions in (2 ) and (6 ) are represented as their mean embeddings: \[\begin{align} \label{eq:mean-embeddings-cond-dists} \Phi(P(\cdot \mid x)) := \int \Phi(y) dP(y|x) \quad \text{and} \quad \Phi( \hat{P}( \cdot \mid x) ) := \frac{1}{k} \sum_{j \in {\rm NN}(x, k, X_n)} \Phi(y_j) \quad \text{for all }\;x \in \mathcal{X}. \end{align}\tag{8}\] Here, the dot “\(\;\cdot \;\)” is used in the notation of the conditional distributions to emphasize that they are probability distributions on \(\mathcal{Y}\) and do not depend on a specific value of \(y \in \mathcal{Y}\). The RKHS distance between the two conditional mean embeddings is the MMD between the true and kNN conditional distributions. It is used as an error metric of the kNN conditional distribution and theoretically analyzed in the following.

The mean embedding of the conditional distribution is known as conditional mean embedding [32], [33] and its estimator based on a regularized least-squares algorithm has been studied extensively [24], [34], [35]. The mean embedding of the kNN conditional distribution in (8 ) is a new estimator of the conditional mean embedding. Its analysis below is thus a new contribution to the RKHS literature and may be of independent interest.

3.2 Assumptions↩︎

We describe key assumptions for the analysis, which follow [25] with appropriate modifications.

The conditional mean embedding in (8 ) is the conditional expectation of the response feature vector \(\Phi(y)\) given a covariate \(x \in \mathcal{X}\); thus, it is the RKHS-valued regression function [24]. We assume that the map from a covariate \(x\) to the conditional mean embedding \(\Phi(P(\cdot \mid x))\) is smooth in the sense that it is Lipschitz continuous.

Our next assumption is that the reproducing kernel (7 ) is bounded on \(\mathcal{Y}\). This is a mild assumption satisfied by many commonly used kernels such as Gaussian, Matérn and Laplace kernels.

It can be easily shown that this assumption implies that the RKHS distance between the conditional mean embedding and any response’s feature vector is bounded: \[\label{eq:bound-output-noise} \left\| \Phi(P(\cdot \mid x)) - \Phi(y) \right\|_\mathcal{H}\leq \sqrt{2} C_{\rm ker} \quad \text{for all }\;x \in \mathcal{X}\;\text{and }\;y \in \mathcal{Y}.\tag{9}\] This implies that the “noise” in the RKHS-valued regression is bounded.

The next assumption is about the intrinsic dimension of the marginal distribution \(P(x)\) on the covariate space, which can be much smaller than the covariate’s dimension \(p\) if \(x \in \mathbb{R}^p\). The error of the kNN conditional distribution shall be shown to decrease as the sample size increases at a rate depending on the intrinsic dimension, not the covariate’s dimension. Let \(B(x,r) \subset \mathcal{X}\) denote the ball of center \(x \in \mathcal{X}\) and radius \(r > 0\): \[B(x,r) := \left\{ x' \in \mathcal{X}\mid d_\mathcal{X}(x, x') \leq r \right\}.\]

This assumption states that if the radius of a ball is increased by a factor of \(\epsilon^{-1}\), the probability mass of the ball increases by at most a factor of \((\epsilon^{-1})^d\). Therefore, the constant \(d\) is interpreted as the intrinsic dimension of the covariate distribution, and can be much lower than the ambient dimension \(p\) if \(\mathcal{X}= \mathbb{R}^p\). For example, if the distribution \(P(x)\) is supported on a line in a two-dimensional space, then \(d = 1\) while \(p = 2\). If \(P(x)\) is supported on a plane in a three-dimensional space, then \(d = 2\) and \(p = 3\) and so forth.

Lastly, we need the following technical condition.

This assumption is satisfied, for example, if \(\mathcal{X}= \mathbb{R}^p\) with \(p \geq 1\), in which case \(\mathcal{V}_\mathcal{B}\leq p + 2\) [36].

3.3 Error Bounds and Convergence Rates↩︎

Under the above assumptions, the distance between the true and kNN conditional distributions can be upper-bounded as follows. The proof, provided in Appendix 6, is an adaptation of the proof of [25], which is an upper error bound on real-valued kNN regression, to our setting of RKHS-valued regression.

From Theorem 1, the following observations can be made.

3.3.0.1 Consistency.

Focusing on the dependence on the sample size \(n\) and the number \(k\) of nearest neighbours, the bound (?? ) can be written as \[\label{eq:bound-constants-simplified} \left\| \Phi( P(\cdot \mid x)) - \Phi( \hat{P}(\cdot \mid x) ) \right\|_{\mathcal{H}}^2 \leq C_1 \frac{\ln(n)}{k} + C_2 \left( \frac{k}{n} \right)^{d/2},\tag{10}\] where \(C_1\) and \(C_2\) are constants independent of \(n\) and \(k\). The first and second terms correspond to the variance and bias, respectively, of the kNN-based conditional mean embedding estimator \(\Phi( \hat{P}(\cdot \mid x) )\). The overall error decreases to zero as \(n\) increases if both the variance and bias decrease to zero; this requires that \(k\) increases as \(n\) increases so that the variance goes to zero, \(\ln(n) / k \to 0\), while \(k\) should not decrease “too fast” so that the bias also goes to zero, \(k / n \to 0\): \[\label{eq:consitency-disp} \left\| \Phi( P(\cdot \mid x)) - \Phi( \hat{P}(\cdot \mid x) ) \right\|_{\mathcal{H}} \longrightarrow 0 \quad \text{as } \;n \to \infty \;\;\left(\text{with} \;k/n \to 0 \;\text{and} \; \ln(n) / k \to 0 \right).\tag{11}\] On the other hand, if \(k\) is fixed to a constant value (e.g., \(k = 1\)), the variance term does not decrease even if the sample size increases. These observations are well known for real-valued kNN regression [26].

3.3.0.2 Convergence in Distribution.

The above consistency (11 ) implies the convergence in distribution (or weak convergence) of the kNN conditional distribution \(\hat{P}(\cdot \mid x)\) to the true one \(P(\cdot \mid x)\) if the response space \(\mathcal{Y}\) is a compact metric space (e.g., \(\mathcal{Y}\) is a bounded closed subset of an Euclidean space) and \(\mathcal{H}\) is a universal RKHS⁶, such as the RKHSs of Gaussian, Matérn and Laplace kernels [28]; see [22] for more generic conditions. That is, under these conditions, the expectation of any continuous bounded function \(f: \mathcal{Y}\to \mathbb{R}\) under the kNN distribution \(\hat{P}(\cdot \mid x)\) converges to the expectation under the true distribution \(P(\cdot \mid x)\): \[\int f(y)d\hat{P}(y \mid x) \longrightarrow \int f(y)dP(y \mid x) \quad \text{as } \;n \to \infty \;\;\left(\text{with} \;k/n \to 0 \;\text{and} \; \ln(n) / k \to 0 \right).\] This supports using the approximate conditional distribution in multiple imputation of missing values.

3.3.0.3 Convergence Rates.

An asymptotically optimal choice of \(k\) that minimizes the bound (10 ), up to the \(\ln(n)\) factor, can be obtained by balancing the variance and bias terms. If we set \(k \propto n^{\frac{2}{2+d}}\), we obtain the convergence rate \[\label{eq:convergence-rate} \left\| \Phi( P(\cdot \mid x)) - \Phi( \hat{P}(\cdot \mid x) ) \right\|_{\mathcal{H}}^2 \leq C_3 \ln (n) \cdot n^{ - \frac{2}{2+d} },\tag{12}\] where \(C_3\) is a constant independent of \(n\) and \(k\).

The rate (12 ) shows that the required sample size \(n\) to attain a desired error level increases exponentially with respect to the intrinsic dimension \(d\) of the covariate distribution \(P(x)\), not the ambient dimension of the input space \(\mathcal{X}\), which is captured by the VC dimension \(\mathcal{V}_\mathcal{B}\) of all the balls in \(\mathcal{X}\). Therefore, even when the covariate’s dimension is large, the error can be small if the covariate features have strong correlations so that the intrinsic dimension \(d\) is small. This is the finding first made by [25] on real-valued kNN regression, and we extend it to RKHS-valued kNN regression.

The rate (12 ) is the same as the minimax optimal rate for estimating a Lipschitz-continuous real-valued regression function when the covariate distribution \(P(x)\) has the intrinsic dimension \(d\) [25]. An interesting point is that the same rate is attained with RKHS-valued kNN regression where the output space is an RKHS that can be infinite-dimensional. Similar observations have been made for RKHS-valued kernel ridge regression [34], [35].

3.3.0.4 Implication to Missing Value Imputation.

The second inequality in the condition (?? ) implies that, for successful recovery of the missing value distribution, the support of the covariate distribution \(Q(x)\) for units with missing responses (see (3 )) should be reasonably covered by the support of the covariate distribution \(P(x)\) for units with observed responses. To explain this, suppose that a missing-response unit has covariate \(x'\), i.e., \(x'\) is in the support of \(Q(x)\), but \(x'\) is not in the support of \(P(x)\) so that there exists some \(r' > 0\) with \(P(B(x',r')) = 0\); then the condition (?? ) is not satisfied for any \(n\) and \(k\).

4 Experiments↩︎

We describe experiments to asses the empirical performance of kNNSampler in recovering the distribution of missing values. Section 4.1 explains the settings, evaluation metrics and benchmark methods. Section 4.2 describes and discusses the results.

4.1 Settings, Evaluation Metrics and Benchmarks↩︎

4.1.1 Data Settings↩︎

We consider the following two models for data generation. As before, let \(n\) be the number of units with observed responses, \(m\) be the number of units with missing responses, and \(N = n+m\) be the total number of units.

4.1.1.1 Setup 1 (Linear with Chi-square noise).

For each unit \(i = 1, \dots, N\), covariate \(x_i\) is uniformly randomly generated on the interval \([-2, 2]\). Response \(y_i\) is the sum of covariate \(x_i\) and noise \(\epsilon_i\) generated randomly from the chi-square distribution with degree of freedom \(2\): \[\label{eq:synthetic-data-linear} y_i = x_i + \epsilon_i, \quad \text{where}~~ x_i \sim {\rm unif}([-2,2]), \quad \epsilon_i \sim \chi^2(2).\tag{13}\] Since chi-square noises are positive, this setup enables assessing the capability of imputation methods to recover non-Gaussian, asymmetric data distributions.

4.1.1.2 Setup 2 (Noisy 2D ring).

This model, considered by  [15], randomly generates covariate \(x_i\) and response \(y_i\) for each unit \(i = 1, \dots, N\) from a noisy two-dimensional ring of unit radius perturbed with an additive Gaussian noise of variance \(0.1\): \[\label{eq:synthetic-data-ring-data} y_i = (1+\epsilon_i)\sin(\theta_i), \quad x_i = (1+\epsilon_i)\cos(\theta_i), \quad \text{where}~~ \theta_i \sim \mathrm{unif}[0, 2\pi], \quad \epsilon_i \sim \mathcal{N}(0, 0.1).\tag{14}\] The conditional distribution of response \(y_i\) given covariate \(x_i\) is bi-modal when \(x_i\) is between about \(-0.5\) and \(0.5\). Thus, this setup enables the assessment of imputation methods in recovering a multi-modal missing-value distribution.

4.1.1.3 Missing Data Mechanism

We consider the MAR (missing at random) setting.⁷ We select \(m\) units uniformly randomly from the subset of the \(N\) units whose covariates lie on the interval \([0.5, 1.5]\) and make their responses missing. We set \(m = 200\), and vary \(n\) to assess the effect of training size on imputation performance. Specifically, we set \(n \in \{2800, 4800, 6800, 8800, 10800\}\).

4.1.2 Performance Metric: Energy Distance↩︎

To quantify the performance of an imputation method in recovering the missing value distribution, we compute the energy distance [37] between the empirical distributions of the complete and imputed datasets. The energy distance is a well-established distance metric between probability distributions and is a parameter-free special case of the MMD [38].

Let \(\tilde{x}_1, \dots, \tilde{x}_m\) be the covariates of the \(m\) units whose responses \(\tilde{y}_{1}, \dots, \tilde{y}_{m}\) are missing, and \(\tilde{y}_{1}^*, \dots, \tilde{y}_{m}^*\) be their imputations. For each unit \(i\), let \(z_i = (\tilde{x}_i, \tilde{y}_i)\) be the pair of the covariate and the true (missing) response, and \(z_i^* = (\tilde{x}_i, \tilde{y}^*_i)\) be the pair of the covariate and the imputation. We compute the energy distance between the empirical distributions of \(D_m := \{ z_1, \dots, z_m \}\) and \(D_m^* := \{ z_1^*, \dots, z_m^* \}\) as \[\mathcal{E}(D_m, D_m^*) := \frac{2}{m^2} \sum_{i,j=1}^m \|z_i - z_j^*\| - \frac{1}{m(m-1)} \sum_{i \ne j} \|z_i - z_j\| - \frac{1}{n(n-1)} \sum_{i \ne j} \|z_i^* - z_j^*\|.\] This is an unbiased estimate of the squared energy distance between the two joint distributions \(Q(x,y) = P(y|x)Q(x)\) and \(Q^*(x,y) = P^*(y|x)Q(x)\), where \(P(y|x)\) is the true conditional distribution of true response \(y\) given covariate \(x\), \(P^*(y|x)\) is the conditional distribution of imputed response \(y\) given covariate \(x\), and \(Q(x)\) is the covariate distribution of missing units. \[\mathcal{E}(Q, Q^*) := 2\mathbb{E}\|z - z^*\| - \mathbb{E}\|z - z'\| - \mathbb{E}\|z^* - z^{*'}\|,\] where \(z, z' \stackrel{i.i.d.}{\sim} Q\) and \(z^*, z^{*'} \stackrel{i.i.d.}{\sim} Q^*\).

A lower energy distance means that the two joint distributions are more similar, implying a better recovery of the missing value distribution. A higher energy distance implies that the imputed distribution is more different from the true data distribution.

4.1.2.1 P-values

Along with the energy distance itself, we report the p-value from a permutation two-sample test using the energy distance as a test statistic, where the null hypothesis is the identity of the two joint distributions \(Q(x,y)\) and \(Q^*(x,y)\). A smaller p-value suggests more significantly that the null hypothesis is false, implying that the imputed distribution \(Q^*(x,y)\) is dissimilar to the true missing value distribution \(Q(x,y)\). A larger p-value suggests that there is no evidence of the null hypothesis being false, suggesting that the imputed dataset is not statistically distinguishable from the dataset with true missing responses. Therefore, a larger p-value implies a better distribution recovery performance of the imputation method.

4.1.3 Benchmark Imputation Methods↩︎

We compare kNNSampler with the following kNN-based and other imputation methods.

Linear Imputation: This method models the response-covariates relation as linear and imputes a missing response by its linear prediction applied to an observed covariate. It should be regarded as a benchmark slightly more sophisticated than naive methods such as mean imputation.

Random Forest [16]: This method, widely used in practice, imputes a missing response by averaging its multiple predictions made by bootstrap-sampled tree regressors. It can learn a nonlinear relation between the response and covariate and handle the interactions among covariate features [39], [40]. We use the default configuration in scikit-learn.

kNNImputer [7]: See Section 2.2 for the description of the method. We set the number \(k\) of nearest neighbours as \(k =5\), which is the default setting in scikit-learn and widely used in practice.

kNN\(\times\)KDE [15]: As explained earlier, this method generates an imputation by sampling from an estimated conditional density of a missing response given a covariate. The conditional density is estimated by weighted Gaussian kernel density estimation over observed responses, with weights derived from a softmax function applied to covariate distances. We use the authors’ recommended settings: inverse temperature \(\tau = 50\) and kernel bandwidth \(h = 0.03\).

4.2 Results↩︎

4.2.1 Qualitative Comparisons↩︎

Figures 4 and 5 describe imputation results by the different methods on datasets generated from the linear chi-square model (13 ) and the noisy ring model (14 ), respectively, with sample size \(N = 10,000\) and \(30\%\) missing rate under the MAR mechanism. The results under the MCAR mechanism are similar and omitted.

The linear imputations ignore the variability in the missing responses and demonstrate the danger of naive imputation methods, such as mean and zero imputations. The imputations by Random Forest and kNNImputer appear to be better than the linear imputations, but are distributed more narrowly than the distribution of missing responses. This is evident for the noisy ring dataset (Figure 5), for which the imputed responses lie inside the ring, which is outside the support of the missing value distribution. This happens because these imputation methods estimate the conditional mean of the missing response given a covariate.

kNNSampler and kNN\(\times\)KDE recover the distribution of missing values much better than the above imputation methods. However, kNN\(\times\)KDE generated imputations for the linear chi-square model (Figure 4) outside the support of the missing value distribution. This is because the noises in this dataset are asymmetric and non-Gaussian, while kNN\(\times\)KDE uses Gaussian noises for generating imputations. In contrast, kNNSampler appears to recover the missing-value distributions accurately. We will next quantitatively compare these methods.

4.2.2 Quantitative Comparisons↩︎

Each experiment, consisting of data generation, imputations by each method, and the calculation of an evaluation metric, was independently repeated 10 times, and the mean and standard deviation of the evaluation metric are reported.

Figures 6 and 7 report the results on the energy distance between the empirical distributions of the imputed and true missing values; Figures 11 and 12 in Appendix 7 show the results for the best two methods for easier comparison. Figures 8 and 9 report the corresponding p-values. See Section 4.1.2 for details.

The energy distance for the linear imputer is the highest among the different methods, quantifying the large discrepancy between the distributions of the imputations and true missing values, as visually observed in Figures 4 and 5. In contrast, the root mean squared error (RMSE) for the linear imputer is the lowest (see Appendix [sec:exp-RMSE]), as the true regression function in either case is a linear function. This result demonstrates that the RMSE is not a good metric for evaluating the distributional similarity between imputations and missing values. See [41] for a further discussion.

The energy distances for kNNImputer and Random Forest are lower than those of the linear imputer, but they are still significantly higher than those of the two other methods. This is reasonable because they are estimating the conditional mean of the missing response given a covariate.

kNNSampler and kNN\(\times\)KDE resulted in small energy distances, implying that they produced imputations whose distribution is similar to that of true missing values. However, the standard deviations of the energy distance for kNN\(\times\)KDE tend to be higher than those of kNNSampler. Indeed, the p-values for kNN\(\times\)KDE are very small for some runs of the experiments (as implied by the error bars in Figures 8 and 9), in which case the distributions of imputations and true missing values are statistically significantly different. In contrast, the p-values of kNNSampler are stably higher, implying that kNNSampler stably generates imputations whose distribution is less distinguishable from that of true missing values.

4.3 kNNSampler Uncertainty Quantification↩︎

This section evaluates kNNSampler’s ability to quantify uncertainty in missing values, using the approach described in Section 2.4. Figure 10 shows the mean and standard deviation of the coverage probabilities of kNN prediction intervals over 10 independent runs, for each sample size and missing rate (MR). As the sample size increases, the coverage probabilities converge to the designed probabilities (80%, 90%, 95%) irrespective of the missing rate, supporting the validity of the prediction intervals.

5 Conclusion↩︎

We studied kNNSampler, a stochastic missing-value imputation method that imputes a missing response of a given unit by searching for its \(k\) most similar units in terms of covariates and by randomly sampling one of the associated \(k\) observed responses. This method is interpreted as sampling from an approximate kNN-based conditional distribution of a missing response given a covariate. Assuming a Lipschitz condition that the true conditional distribution changes continuously with covariates, we proved that the kNN conditional distribution converges to the true conditional distribution as the number \(k\) of nearest neighbours increases at a rate slower than the sample size increases. This analysis offers a theoretical justification for kNNSampler, and may be of independent as it analyzes a novel kNN-based estimator of the Hilbert space embedding of a conditional distribution. Empirical results demonstrate the capability of kNNSampler in recovering the distributions of missing values.

6 Proof of Theorem 1↩︎

6.1 Bias Bound↩︎

Lemma 1 below is from [25].

6.2 Variance Bound↩︎

7 Additional Experimental Results↩︎

Figures 11 and 12 show the results on the energy distance for the three best methods in Figures 6 and 7, respectively. Figure 13 shows the corresponding results on the root mean squared errors (RMSEs) between the imputations and true missing values. As mentioned in the main body, having low RMSEs does not imply good imputations in terms of recovering the distribution of missing values. Indeed, imputations from the linear imputer have the lowest RMSEs, but their distribution significantly differs from the distribution of true missing values, as quantified in Figures 6 and 7 and visually observed in Figures 4 and 5.

References↩︎