kNNSampler: Stochastic Imputations for Recovering Missing Value Distributions
Authors: Parastoo PASHMCHI
Jérôme Benoit
Motonobu Kanagawa
Venue: Accepted by TMLR ()

Abstract: We study a missing-value imputation method, termed kNNSampler, that imputes a given unit's missing response by randomly sampling from the observed responses of the k most similar units to the given unit in terms of the observed covariates. This method can sample unknown missing values from their distributions, quantify the uncertainties of missing values, and be readily used for multiple imputation. Unlike popular kNNImputer, which estimates the conditional mean of a missing response given an observed covariate, kNNSampler is theoretically shown to estimate the conditional distribution of a missing response given an observed covariate. Experiments illustrate the performance of kNNSampler. The code for kNNSampler is made publicly available (https://github.com/SAP/knn-sampler).

Published in Transactions on Machine Learning Research (12/2025)
KNNSampler: Stochastic Imputations for Recovering Miss- ing Value Distributions
Parastoo Pashmchi
parastoo- wwwwoccsssssc.ccommm.
SAP Labs France E- E-Mobility Research EURECOM, Sophia Antipolis, France
Jérôme Benoit
Berome..enooitssapp. com
SAP Labs France E-Mobility Reseorch
Motonobu Kanagawa
EURECOM, Sophia Antipolis, France
Reviewed OpenReview: bttps://openreet net/ forum? id=ACDDIIACCQG
Abstract
We study a missing-value imputation method, termed KNNSampler, that imputes a given unit's missing response by randomly sampling from the observed responses of the most similar units the given unit terms of the observed covariates. This method can sample unknown missing values from their distributions, quantify the uncertainties of missing values, and be readily used for multiple imputation. Unlike popular kNNImputer, which estimates the conditional mean of missing response given an observed covariate, kNNSampler is theoretically shown to estimate the conditional distribution of a missing response given an observed covariate. Experiments illustrate the performance of kNNSampler. The code for kNNSampler made publicly available
Keywords: missing values imputation, nearest neighbours, conditional distri- bution, kernel mean embedding
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 common preprocessing step used to address missing data. Over the docades, various imputation methods have been proposed, ranging from
simple statistical techniques to machine learning algorithms (e.g., Rubin 1976 Schafer 1997 Schafer and
Graham 2002 Little and Rubin 2002 Mattei and Frellson 2019 Enders 2022
kNNImputer (Troyanskaya al. 2001 is one of the most widely used imputation methods, owing to its simplicity and availability in popular software packages such as scikit-learr" (Pedregosa et al. 2011). It imputes a missing response variable (c.g., customer satisfaction level) of given unit (c.g., a customer) 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 nearest neighbours (kNN) regression (Stone 1977) so the imputation is an estimate of the conditional eapectation of the missing response given covariate. The method has been widely usod in science and engincering, and many extensions
have been proposed (e.g., Larcia-Laennina et al. 2009 Tutz and Ramzan. 2015 De Silva and Perera 2016
Huang a..
2017 Faisal and Tutz 2021).
An issue of kNNImputer, shared by other regression-bases imputers, is that the distribution of imputations can be significantly different from the distribution of true (hidden) missing values. This is because, as
https://github... www.so////aaaacce...c S't/SAAA////n/// gampler
https://scikit- http://wwww..cbeceeaaaaaaa/aaaaa http:///wwwaaaacooooooooooooooon
htt.alaett..tutttteeiltttttttr KWWImptter..ettt html

---

Published in Transactions on Machine Learning Research (12/2025)
KNNImputer
KNNSampler
1.0
1.0
0.5
0.5
0.0
0.0
005
-0.5
-1.0
-1.0
-1.0
-0.5
0.0
0.5
1.0
1.5
-1.0
-0.5
0.0
0.5
1.0
1.5
X
X
Figure 1: Comparison of imputations by KNNImputer (left) and kNNSampler (right). In each figure, and are the covariate and response, respectively. Blue points are observed covariate-response pairs, green points are true missing values and rod points are imputed values. For dctails, see Section
Conditional Distribution vs KNN Y
2.00
True X=0.5 KNNSampler Y
1.75
KNNimputer
1.50
1.25
Density
1.00
0.75
0.50
0.25
0.00
-1.0
-0.5
0.0
0.5
1.0
Figure 2: Comparison of the samples of the true conditional distribution P(y(æ) of missing response of a unit with covariate x = 0.5 (blue) and the kNN conditional distribution P(y(æ) with = 1, 000 (orange) on the noisy ring data Figure [ with sample size 10, 000. The imputations by kNNImputer with are shown as the green dotted vertical line.
mentioned, imputation of KNNImputer is estimate of the conditional expectation of missing response, thus tending to be 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 illustrated in Figures I and 2 where the true conditional distribution of a missing response is bimodal when the covariate small, but the distribution of imputations unimodal and many imputations take values never realized by the true missing values. A substantial bias can occur an analysis of such distorted imputed dataset, for example, when estimating the variance, quantiles and modes in the population. (See Sections 2.1 and 2.2 for more formal discussions.)
The above issue of kNNImputer may be addressed by estimating the conditional distribution of missing response given a covariate, and randomly sampling imputations from it. This idea was investigated by Lalande and Doya 2023), who proposed the "kNN*KDE" approach that combines a soft version kNN and kernel density estimation (KDE). For given unit, the conditional density of missing response estimated as weighted average of Gaussian densities centered observed responses, where the weights are computed that units similar, in terms of covariates, the given unit receive larger weights. kNNxKDE.
2

---

Published in Transactions on Machine Learning Research (12/2025)
demonstrated to have good empirical performance in recovering the distribution of missing values, compared to established imputation methods, including KNNImputer, missForest (Stekhoven and Bühlmann 2012) Softlmpute (Hastie et al. 2015), and Gain (Yoon et al. 2018). However, no theoretical guarantee exists for kNNx such as its statistical consistency, i.c., 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 hyperparameterss should be chosen. While kNNxKDE 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 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 mcan of the observed responses of k nearest neighbours, kNNSampler simply samples one of those k observed responses. It is thus simpler than <NNxKDE as it does not involve 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 Kanagawa (2024). Figures and 2 describe imputations by KNNSampler, which align much better with the distribution of true missing values than imputations by KNNImputer. More systematic experiments are provided in Section !
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 's Current Population Survey and the National Center for Education Statistics (e.g.. Andridge and Little 2010). 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 nearesl-neighhour hot deck if nearest neighbours define the adjustment cell (Little and Rubin 2002 Example 4.9). 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 (Andridge and Little 2010
Our contribution is to establish KNNSampler and thus the rearest-neighhont random hot deck, as a theoreti- cally principled missing-valuu imputation method. To this end, we analyze the kNN conditional distribution, i.e. the empirical distribution of l 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 nunber k of the nearest neighbours, and other problem-speciffe constants. The error is measured by the maximum mean dis- crepancy (MMD) (Gretton et al. 2012). distance metric on probability distributions that metrizes the weak convergence (Simon- -Simon-Gabriel Gabriel et al.] 2023), , between the kNN and true conditional distri- butions. 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 infinity, if the number & of nearest neighbours increases to infinity at rate slower than This offers a theoretical foundation of the kNNSampler and thus the nearest-neighhour 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 Muuandet et al., 2017 Chapter 4), which is the RKHu-valued regression function (Grünewälder et al. 2012). The RKHS distance between these two embeddings is equivalent to the MMD between the kNN and true conditional
3

---

Published in Transactions on Machine Learning Research (12/2025)
distributions. Our bound leads to the consistency and convergence rates for the novel kNN-based estimator of the conditional mean embedding.
Our amalysis extends the error analysis by Kpotufe 2011 on real-valued kNN regression to RKHS- valued regression in which the response variable is infnito-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 Faniionnanddd its theory in Section 3 We report experimental results on synthetic data in Section 4 and on real solar-power data in Section 5
2 Proposed Approach
This section describes the proposed approach. Section 2.1 introduces the setting. Section 2.2 explains the KNNTmputer and its issue as preliminary. We describe k KNNSampler in Section 2.3 uncertainty quantifi- cation 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 X' and be measurable spaces representing the covariate space and the response space, respectively. For example, the covariate space may be the d-dimensional Euclidean space, X' = Rd, in which case covariate X consists of d features (c.g.; a person's 's age, weight, height), and the response space may be the real line y R, in which case response E y is real-valued (e.g., the person's blood pressure).
We assume that our dataset consists of n + m units (c.g., persons), where n units have both covariate x; E X and response yi E y observed, while m units have only covariate ã. X' observed and response Ymiss. y missing:
D, = (111,(1)))))))))) (Ins Un)), Dmiss := ((01, )1,miss)...... (7mm Frnsmiss))
(1)
For each of the n units with observed responses, we assume that the covariate follows marginal distribution P(x) and the response given the covariate follows the conditional distribution P(ylæ: in an independently and identically distributed (i.i.d.) manner:
(ii,,))...,,,,, Sid P(ulx)P(x)
(2)
On the other hand, for the units with missing responses, the covariate is assumed to follow marginal distribution Q(²), which can be different from P(x), while the conditional distribution of the missing response given the covariate remains the same:
(T D1. (Tm miss) 1i.. P(Imisa (T)
(3)
This assumption implies that the probability of unit missing its response is determined by the unit's covariate and is not affected by the response. Therefore, it is instance of the Missing-AttRRaddom (MAR) assumption RRbiin 1976.. In the special case where the covariate distributions for the two cases are the same, Q(7) P(ī), the assumption can be interpreted as the Missing-Compllteele Commpleteyy-Al-Rarddae (MCAR) assumption where missingness occurs completely randomly.
Under this setup, missing responses may be imputed by estimating the unknown conditional distribution P(ylar) of response given a covariate and sampling from it. This is what the kNNSampler does.
4

---

Published in Transactions on Machine Learning Research (12/2025)
Remark 1. Sampling imputations from the conditional distribution is needed for unbiased estimation of quantåly of interest and ils well-calibrited uncerlainly quanlifieation. We informally describe this from the Bayesian perspective of Ruubi (1987), where the quantity of interest, denoted here by 0, is treated as a random variable. For the frequentist perspective, see Rubin (1987; 1996); Murray 2018 A Bayesian analysis done by computing the posterior distribulion of given the observed dala in
P(0 Dn, T1. Tm)
= -f P(0) Da, (Z1, B1)..., (Zm: Fmm) P. Im Dnn 81,.... Fm) din ... dÿms (B)
(A)
where (A) is the posterior distribution of 0 given observed dalasel D, (111111),- v1)...., (I NnD3 and impuled dataset (Z1,01),... U1)...., (Ims Dm) and is computed by standard Bayesian analysis, treating the imputations observed data; (B) is the conditional distribution of missing responses I1..... given Dr. and observed covariates $1,..., Em for missing responses, and can written as
112
P 4m Dn: F1...., 3m) II P().).)).
where we used (3). Thus, by estimating the conditional distribution P(ğ of a response given a covariate and sampling from it, the posterior distribution of can be approzimalely compuled, using, e.g., the mullipple
imputation approach [Rubin) 1987 1996 Murray 2018).
2.2 Issue with KNNImputer and Regression-based Imputers
Before describing the proposed kNNSampler, we discuss an issue with the widely used kNNImputer (Troy- anskaya el al. 2001 and other regression-based imputation methods.
Suppose that the covariate space X' equipped with a distance metric dx x²) that quantifies the distance between any two points æ, x' X. For example, if X' is the Euclidean space, then dx(2, may be the Fuclidean distance between two vectors and x'. Tet Xn be the set of covariates for the n units with observed responses:
X, (., un]
For given covariate è and a number of nearest neighbours, let NN(6, k. Xn) be the indices of the k units whose covariales the most similar to z in terms of the distance metric among the units with observed responses
&, = f E (I....,m.}
x ... S dx,
(4)
I dx(ū, for all j E fI....,n] (fi,...,fk11,
That is, NN(2, X..) the indices of the & nearest neighbours in X,.
&NNImputer (Troyanskaya et al. 2001) imputes the missing response Vimiss of the unit with observed covariate ã; as the average of the observed responses yi.....,bh Uj. of its k-ncarest neighbors Tj.....;
Ds,imp JENN(2,
Uj
This is kNN regression (e.g., Györfi et al. 2002 and thus estimates the conditional mean of the missing response mise given the observed covariate ā;:
Kiimp s f(ā:) = / dP(j)a.),
3TF there tie the distances dx (8, xi), break randomly.

---

Published in Transactions on Machine Learning Research (12/2025)
Algorithm 1: KNNSampler
Input: Number of nearest neighbors k, observed covariates $1...., 3m X with missing responses, observed covariate-response pairs (I1 v1)...., ...., (Ins yn) X y Output: Imputed responses 0i,imp:-... Un..mip.......im y.
for i to m do
Vi.imp := V/j where (l.....,h n is uniformly sampled from NN k, Xn) in equalion of the k-nearest neighbors of ā; in Xn = (Ii,..,,Tah.
the indices
end
where f : X V is the regression function. In this case, the observed covariate and the imputed response (it, (m.mp) approximately follow the degenerate joint distribution
ő - (z))
where (ā) denotes the Dirac distribution at the conditional mean f()), i.c., the degencrate distribution whose mass concentrates at (()). This differs from the joint distribution the observed covariate and the true missing response (øt, mi.mss)::
P( ())
(5)
unless the conditional distribution P .)) is the Dirac distribution 8(ğ - f())), i.e., unless the missing response is the deterministic function of observed covariate. The same issue occurs with other single impu- tation methods based on regression, because they impute the missing response by estimating the conditional mcan.
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 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(ğ I
Remark 2. Consider the Bayesian analysis in Remark 1 If the missing responses are imputed by deter- ministic, regression estimates im......... im. the posterior distribution becomes that of the quantity of interest 0 given observed and imputed datasets, both treated as observed:
P(O Da, (ī1, (1..mp))..... .... (Tme 1m. (mm.imp))
This ignores uncerlainties in the missing responses, leading to final uncerlainly estimates for that are overconfident (a prediction interval may be much narrower than the actual error of a point estimate).
2.3 KNNSampler
We now describe kNNSampler (Algorithm 1 Consider imputing the missing response Jmiss of a unit with observed covariato a. KNNSampler estimates the conditional distribution P(Imiss ()) of Umiss given & as the empirical distribution of the observed responses Vi....., Mix of the nearest neighbours T,,...., of is
P(ymiss
T) Rd PTimis
Σ
8(Fmiss - -y;),
(6)
K
X...
which is the discrete distribution where each of -. Ui. has probability mass 1/k. An imputation fimp for the missing response is randomly sampled from this empirical distribution:
Imm
P t).
Algorithmically, this to randomly sample one of the kNN observed responses y/....., Uia Algorithm 1 independently applies this procedure the observed covariate é; to generate an imputation imp of missing value Wimisss for each unit i m,....

---

Published in Transactions on Machine Learning Research (12/2025)
Choice of k The mumber of nearest neighbors & is hyperparameter of kNNSampler. The theoretical and empirical results below indicate that k should not be fixed to 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 (I1.81).. (Zn. yo) and select & among candidate values that minimizes the mcan-square error on held-out observed responses, averaged over different training-validan
splits. In particular, the present work uses Leave-One-Out Valid-aaiiotoion (LOOCV) using the fast computation method recently proposed by Kanagawa (024
2.4 Uncertainty Quantification of Missing Values
Quantifying the uncertainty in missing values important for several reasons, including assessing the re- liability of imputations and the adequacy of the covariates used, well determining how to perform imputations (e.g., single multúple) 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.
Conditional Probability Estimation KNNSampler can be used to estimate the conditional probability of missing response Imiss belonging to a specified (measurable) subset S of the response space y, given observed covariate ă:
Pr(Tmiss E ES [| 8] (P(P)))))
where 1py E S] is the indicator function that outputs 1 if jE E S and otherwise. By replacing the unknown conditional distribution P(ğ 8) by the kNN conditional distribution P(j | 8) in (6) this conditional probability is approximated
Pr(Amisa ES 10 e $ a
Σ
Myy E S].
JENN(F,&,X.))
In other words, the conditional probability is estimated as the observed frequency of the kNN response values that fall in S.
Interval Estimation Let us focus on real-valued missing response Jmiss E y = R. The conditional probability of the missing response belonging to a given (finite or infinite) interval S = (e, u), where < 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 al a specified frequency << 1 a (e.g., a = 0.05, 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- a.
Such an interval uL.u) is ructed by defining its lower bound é and upper bound as, respectively, the lower and upper a/2 empirical quantiles of the kNN responses, i.., the ka/22-mmalest and the ka/2 2arrgesst kNN responses (c.g., if 200 and a = 0.03, the 5th smallest and the 5th largest kNN responses):
Pr(f < Amiss < u ( R 1-a
Conditional Standard Deviation Estimation The conditional standard deviation of 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 = I, 1,..,,,, B, kNNSampler is independently applied to impute the missing responses in the dataset Dmiss to create an imputed dataset.
[
D(b) n+m = D, U (.() where ((0) ((51, ((51,11.1))) ... (T
wilh Jmm.imp).

---

Published in Transactions on Machine Learning Research (12/2025)
(b6) where Dimmp is an imputation for the i-th unit with missing response Di,miss covariates Σz. This results in B imputed datasets:
D
D(B)
ptmm:.... am
An analysis can then be made based on the standard procedure of multiple imputation
Rubin 1987).
For example, suppose that we want to estimate a population quantity 0* (e.g., the mean customer satisfaction level of a population). Let Sa+m be a function of a dataset of size + m that outputs an estimate 4n+m of the unknown 0* (e.g., the empirical average of n + values): 4n+m S(Dn+m). Apply this function to each of the B imputed datasets, one obtains B estimates of 8":
1 2
S(D n+m); b= 1..,..., B.
(b)
The empirical average of these B estimates gives multiple-imuutation estimate of 0", The empirical standard deviation of the B estimates - m 0(4) quantifies the uncertainty due to the missingness in the original data. Combined with the standard error of each 0 a4m> this standard deviation be used to quantify the overall uncertainty of the estimate using Rubin's ' rule.
3 Theory
We describe theory for KNNSampler's conditional distribution (6) as an estimator of the true conditional distribution. We shall show that, as the number h 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) (Gretton et al. 2012), which implies the convergence in distribution (Sriperumbudur al. 2010 Section 5). We prove this by adapting the proof of the convergence rates of real-valued kNN regression by Kpotufe [201] Theorem 1) to Hilbert space- valued kNN regression
We use the framework of kernel mean embedding (Muander et al. 2017) in which every probability distribu- tion is represented as a distinct point in an ddfiiitedimensionn 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 the MMD, quantifies the estimation error. An upper bound on this distance is obtained terms of the sample size, the number of nearest neighbours, and other relevant quantities.
3.1 RKHS Embeddings of Conditional Distributions
Let first define an RKHS on the response space V. As before, y is a measurable space such the p-dimensional Euclidean space, y RP. A Hilbert 5 H consisting of functions on y is called RKHS if there exists map
Φ: H
called feature map, such that the value f(y) of any function f any point y in y can be written the inner product bet ʃ and the feature map pgy) of y:
e
f(y) = (f. ((v)) for all y E V,
where (())x denotes the inner product of H. The (y)) may be called feature vector of y, and H the feature space, which can be infinite-dimenssionall
The inner product between the feature maps (2)), (3y') of any two points v/ defines the kernel function
(l.,v) (0(v). ( ))x for all y, E.
'Hilbert space-valued KNN regression also analyzed in Lian 2011). but their results not directly applicable to case. This is because Lian 2011 assumees that Hilbert space-valued noises independent of input variables, but this assumption is too strong
A Hilbert space vector space in which inner product defined, the induced from the inner product, and the limit point of any convergent sequence this norm helongs the vector space.
8

---

Published in Transactions on Machine Learning Research (12/2025)
This is called reproducing kernel of the RKHS. The RKHS and the reproducing kernel are one-to-one, RKHS can be induced by defining kernel. For example, if y = IR,, the Gaussian kernel {v, v/) exp(-aily - e for > 0 is the reproducing kernel of a certain RKHS H. and there exists an infinite-
dimensional feature map that induces the Gaussian kernel as 7 See e.g. Steinwart and Christmann 2008 Kanagawa
22025 for details on RKHSs
Every probability distribution on y is represented the expected feature map:
(PP) = / (())d(()) H.
This is called mean embedding of P. the RKHS H is large enough, any two different probability distributions P and Q are mapped two distinct mean embeddings:
P#Q
((P) +(Q).
In this case, the RKHS is called characteristic (Sriperumbudur et al. 2010). For example, Gaussian, Matérn and Laplace kernels induce characteristics RKHSs.
The true and KNN conditional distributions in 2 and
66
are represented as their mean embeddings:
( = (1))d((ylz) and + (P
x))
E
p for all X.
k
JENN(Z.A.X.)
Here, the dot is used in the notation of the conditional distributions emphasize that they are probability distributions on )' and do not depend on specific value of y E y. The RKHS distance bet 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 known as conditional mean embeddding (Song et al. 2009 2013) and its estimator based on a regularized least-squares algorithm has been studied extensively (e.g.. Grünewälder al.) 2012 Li et al. 2022 2024). The mean embedding of the kNN conditional distribution in S is to the RKHS literature and may be of independent interest.
estimator of the conditional mean embedding. Its analysis below is thus new contribution
3.2 Assumptions
We describe key assumptions for the analysis, which follow Kpotufe (2011 with appropriate modifications,
The conditional mean embedding in (8) is the conditional expectation of the response feature vector p given a covariate x E X; thus, it is the RKHS-valued regression function (Grünewälder et al.) 2012). We assume that the map from a covariate x to the conditional mean embedding (PP(- x)) is smooth in the sense that it is Lipschitz continuous.
Assumption 1. There erists a constant >> 0 such that the RKHS distance between the conditional mean embeddings for any two inpuls x, x' E X is bounded by the dislance between x and x' limes X:
(PP x)) (P(( xl)) (x,)) for all x E X,
where || l| the norm of the RKHS H.
Our next assumption is that the reproducing kernel is bounded y. This is mild assumption satisfied by many commonly used kernels such Gaussian, Matérn and Laplace kernels.
Assumption 2. There exists constant Cker > 0 that upper-bounds the value of the reproducing kernel
0 S Ely.y.) Cher for all y.y/ e..
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:
[| x)) +(v) S 20 for all x E X and y E y.
(9)
9

---

Published in Transactions on Machine Learning Research (12/2025)
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(z) on the covariate space, which can be much smaller than the covariate's dimension p if I E RR. The error of the kNN conditional distribution shall 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) C X denote the ball of center x X' and radius
B(z,r) := fx' & dx(x,²)) r)
Assumption 3. For the marginal distribution P(x) on the covariate space X, there are positive constants Cdist > 0, max > O, and d > 0 such that
P(B(x,r) & Cdisa C (()(x,er) for all 0 <r<< < Tmax and all 0 0<<<1.
This assumption states that if the radius of ball is increased by a factor of e- the probability mass of the ball increases by at most a factor of (e-l)d. Therefore, the constant d is interpreted as the intrinsic dimension of the covariate distribution, and can be much lower than the ambient dimension if X' = RR. For example, if the distribution P(x) is supported on a line in two-dimensional space, then d while p= 2. If P(x) is supported on plane in three- dimensional space, then d = 2 and p 3 and so forth.
Lastly, we need the following technical condition.
Assumption 4. The covariate space X is a metric space with distance metric dx such that the class of all balls = (B(z,r) E X, r > 0} has a finite Vapnik-Cherrone Chervonenkis (VC) dimension VB >
This assumption is satisfied, for example, X R2 with 2 1, in which case VB $ p (e.g.. Mohri et al. 2018 Excrcise 3.17).
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 A is an adaptation of the proof of Kpotufe (2011 Theorem 1), which is an upper error bound on real-valued kNN regression, to our setting of RKHS-vall valued regression.
Theorem 1. Let (I1, v1)........ (Ins Un) sid. P(ula:)((a) and P(yla) be the kNN conditional distribation
6
with h nearest neighbours. Suppose that Assumptions 7 and hold. Let O < 8 < 1. Then, with probability at least - 28, the bound
( (2) x)) 2(p(- 2))||² S 4C (Vgl(()) In())) 1 223,2 k
k 2/d P(B(arr) ) (10)
3Crist
holds simultaneously for all x X, E fI, and 0 < < Tmax satisfying
k z (2 + (8(66) and k (B((r)))
(11)
n 3C
From Theorem [ the following observations can be made.
Consistency. Focusing on the dependence on the sample size n and the number k nearest neighbours, the bound 110) can be written as
In(n) (P(( <(p( m)) Ci
k
2/d
+ C2
(12)
k
where C and C2 are constants independent of and k. The first and sccond terms correspond the variance and bias, respectively, of the kNN-based conditional mean embedding estimator @(P(- x)). The overall error decreases to increases if both the variance and bias decrease to zero; this requires that
10

---

Published in Transactions on Machine Learning Research (12/2025)
k increases increases that the variance goes to zero, In(n)/k O, while it should not decrease "too fast" so that the bias also goes to zero, k/nn + 0:
(PP( (a)) ((P(. ((P((11))]
x))
0 as - 00 (with k/n - 0 and In(n)/k 0).
(13)
H
On the other hand, l is fixed to a constant value (e.g., h = 1), the variance term does not decrease even if the sample size increases. These observations are well known for real-valued kNN regression (e.g.. Györfi ltt al. 2002).
Convergence in Distribution. The above consistency 13 implies the convergence in distribution (or weak: converyence) of the kNN conditional distribution P(. x)) to the true one P(. x) if the response space y is compact metric space (e.g., y is a bounded closed subset of an Euclidean space) and H. is a universal R such as the RKHSs of Gaussian, Matérn and Laplace kernels (Sriperumbudrr et al. 2010 Theorem 23); see Camon-Gabbiel et al. (2023) for more generic conditions. That is, under these conditions, the expectation of any continuous bounded function f ::y- - R under the kNN distribution P(. (x) converges to the expectation under the true distribution P(. 1r):
/ flu)dP()la) -
f flwddP(y x) 20000 (with k/n 0 and In(n)/k - 0).
This supports using the approximate conditional distribution in multiple imputation of missing values.
Convergence Rates. An asymptotically optimal choice of & that minimizes the bound 12 up to the
In(n) factor, can be obtained by balancing the variance and bias terms. If we set l ,oत obtain the convergence rate
[P((( x)) (P(( ((P(11=))[/
|0>)| - C3 In(n) on zंत
(14)
where C3 is a constant independent n and k.
The rate 14] 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 X, which is captured by the VC dimension VB8 of all the balls in 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 Kpotufe (011h on real-valued kNN regression, and we extend it to RKHS-values kNN regression.
The rate 14 is the same as the minimax optimal rate for estimating a Lipschitz-contimmous valued regression function when the covariate distribution P(x) has the intrinsic dimension d [Kpotufe 2011 Theorem 2). 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 infiniteddimensional. Similar observations have been made for RKHS-valued kernel ridge regression Li et al. 2022 2024
The second inequality in the condition 11 implies that, for successful recovery of the missing value distribution, the support of the covariate distribution Q(z) for units with missing responses (see (3) should be reasonably covered by the support of the covariate distribution ((x) for units with observed responses. To explain this, suppose that missing-response unit has covariate a', i.e., x' is in the support of Q(x), but x' is not in the support of P(z) so that there exists some r' > 0 with P(B(x',r') = 0; then the condition [1] is not satisfied for any n and k.
Implication to Missing Value Imputation.
4 Synthetic Data Experiments
We describe experiments to assess 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.
6 An RKHS 76 consisting functions on a metric set Y called universal if any continuous bounded function f: can be approximated arbitrarily well in terms the supremum norm by functions H.
11

---

Published in Transactions on Machine Learning Research (12/2025)
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, be the number of units with missing responses, and N n + m be the total number of units.
Setup 1 (Linear with Chi-square noise). For each unit i = 1,..., N, covariate my is uniformly randomly generated on the interval [-2, 2]. Response yi is the sum of covariate z; and noise generated randomly from the chi-square distribution with degree freedom 2:
y x; Ei where x2 unif 2)), 6 x2)).
(15)
Since chi-square noises are positive, this setup enables assessing the capability of imputation methods recover non-Gausian,, asymmetric data distributions.
Setup 2 (Noisy 2D ring). This model, considered by Lalande and Doya 2023), randomly generates covariate I3 and response y/ for each unit i = ,....,N N from noisy two-dimensional ring of unit radius perturbed with an additive Gaussian noise of variance 0.1:
V4 (1 + G) sin().), r; z; (1 + c.) cos(B.), where 0: unif 2m), < ~ ~00,00)
(16)
The conditional distribution of response y: given covariate x; bi-modal when æ, is between about 00.5 and 0.5. Thus, this setup enables the assessment of imputation methods in recovering a multi-modal missing-valus distribution.
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 E {2800, 4800, 6800, 8800, 10800).
Missing Data Mechanism We consider the MAR (missing at random) setting 0
4.1.2 Performance Metric: Energy Distance
To quantify the performance an imputation method in recovering the missing value distribution, we compute the energy distance (Szćkely and Rizzo 2013) between the empirical distributions of the complete and imputed datasets. We use the energy distance as it a proper distance between distributions, can be easily computed from samples based on their Euclidean distances without the need for optimization (as compared with, e.g., Wasserstein distance whose computation requires optimization to solve optimal transport (Peyré et al., 2019)), and is parameter-free (in contrast to the MMD defined with, c.g., Gaussian kernel, which depends on the bandwidth parameter). The energy distance is a canonical instance of MMD defined with a distance- based kernel (Sejdinovic et al. 2001)); it is "canonical" in the sense that it is both scale-invariant (the distance scales linearly with the scale of data) and rotation-invariant (Szćkely and Rizzo 2013. Section 3).
Let 21,..., &m be the covariates of the m units whose responses U1...., Fm are missing, and ....., Un be their imputations. For each unit i, let 2, (øt, ūe) be the pair of the covariate and the true (missing) response, and 2 = (ă,, if) be the pair of the covariate and the imputation. We compute the energy distance between the empirical distributions of Dm == { am) and D = {zi,... :m.. as
E(Dm. m. Din) = m² 1,==
2 27 m
1
Σ
22 اارl n(n 1) k
1
Σ 1n1
-
This is an unbiased estimate of the squared energy distance between the two joint distributions Q(x, y) = ((z)))((z) and Q* (xxx)) = (())) where P(y|z) is the true conditional distribution of true response
We also performed the experiments under the MCAR (missing completely random) setting, but the results similar and thus omitted.
12

---

Published in Transactions on Machine Learning Research (12/2025)
y given covariate 2, (y|22) is the conditional distribution of imputed response y given covariate 2, and (x) is the covariate distribution of missing units:
(QQ := 2 z- =01 EE| - z'11 - Ell=* - -='11
=
where 2, x i.d.d. Q and x", 2 F.L.d. Q"
A lower energy distance indicates that the two joint distributions are more similar, implying better recovery of the missing-value distribution. A higher energy distance indicates that the imputed distribution is dissimilar to the true data distribution.
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 lincar and imputes missing response by its linear prediction applied to observed covariate. It should be regarded as benchmark slightly more sophisticated than naive methods such as mean imputation.
Random Forest (Stekhoven and Bühlmann 2012): 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 (e.g., Shah et al. 2014 Tang and Ishwaran 2017). We use the default configuration in scikit-learn.
KNNImputer (Troyankkaya al. 2001): 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 sctting in scikit-learn and widely used in practice.
kNNxKDE (Lalande and Doya. 2023 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 softmax function applied to covariate distances. We the authors' recommended settings: inverse temperature 50 and kernel bandwidth h 0.03.
As suggested earlier, the number h of nearest neighbours for kNNSampler is determined by the fast leave- one-out cross-validation method of Kanagawa (0024) using the observed covariate-response pairs.
4.2 Results
4.2.1
Qualitative Comparisons
Figures 3 and 4 descríbe imputation results by the different methods on datasets generated from the linear chi-square model (15) and the noisy ring model (16), respectively, with sample size N 10, 000 and 30% missing rate under the MAR mechanism. The results under the MCAR mechanism 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 kN- Nimputer 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 4 for which the im- puted responses lic inside the ring, which outside the support of the missing value distribution. This happens because these imputation methods estimate the conditional mean of the missing response given covariate.
kNNSampler and KNN*KII recover the distribution of missing values much better than the above imputa- tion methods. However, kNNx KDE generaled imputations for the linear chi-square model (Figure outside the support of the missing value distribution. This is because the noises in this dataset are asymmetric and non-Gaussian, while kNNxKDE uses Gaussian noises for generating imputations. In contrast, kNNSampler appears recover the missing-ver value distributions accurately. We will next quantitatively compare these methods.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Linear
Random Forest
KNNImputer
17.5
17.5
17.5
15.0
15.0
15.0
12.5
10.0
10.0
10.0
7.5
7.5
5.0
5.0
0.0
0.0
0.0
20
1.0
05
00
05
10
15
20
KNNKKDE
KNNSampler
17.5
15.0
12.5
12.5
10.0
10.0
50
5.0
25
255
-2
2.5
Figure 3: Missing value imputations by different methods for dataset from the linear chi-square model 15) with sample size N = 10, 000 with 30% missing rate under the MAR mechanism. True missing responses are shown in green, imputations in red, and the rest in blue.
Random
KNNImputer
1.0
0.5
0.5
0.0
0.0
0.0
-03
-05
-05
-10
-1.0
-1.0
-1.0
-05
0.0
0.5
1.0
13
-1.0
-0.5
0.0
0.5
10
155
-1.0
-0.5
0.0
0.5
1.0
1.5
X
KNNSampler
0.5
0.0
-05
0.5
-10
-10
1.5
-1.0 -0.5
X
Figure 4: Missing value imputations by different methods for dataset from the noisy ring model (16 with sample size N 10, 000 with 30% missing rate under the MAR mechanism. True missing responses are shown in green, imputations in red, and the rest in blue.
4.2.2 Quantitative Comparisons
Each experiment, consisting of data generation, imputations by each method, and the calculation of the evaluation metric, was independently repeated 10 times, and the mean and standard deviation of the eval- uation metric are reported. Tables and report the results the energy distance between the empirical distributions of the imputed and true missing values. See Section 4.1.2 for details.
kNNSampler and kNNxKDE yielded significantly smaller energy distances than the other methods, which suggests that their imputations are distributed more similarly with the true missing values and align with Figures 3 and The energy distance for the linear imputer the highest among the different mothods, quantifying the large discrepancy between the distributions of the imputations and true missing values, visually observed in Figures 3 and The energy distances for kNNImputer and Random Forest lower
14

---

Published in Transactions on Machine Learning Research (12/2025)
Table I: The energy distance between the empirical distributions of imputations and true missing values on the linear chi-square dataset (1).. For each method and sample size, the average and standard deviation over 10 independent runs are shown.
Sample Size
kNNSampler
Random Forest
KNNImputer
kNNxKDE
Linear
3000
0.027 + 0.031
0.076 0.023
0.200 = 0.038
0.036
0.033
0.585 1 0.053
5000
0.027 = 0.009
0.077 = 0.030
0.199 = 0.041
0.033 + 0.019
0.598 = 0.025
7000
0.027 + 0.021
0.080 = 0.018
0.219 = 0.024
0.028 + 0.018
0.589 = 0.034
9000
0.017 + 0.009
0.076 + 0.023
0.183 = 0.031
0.016 + 0.007
0.605 + 0.054
11000
0.018 + 0.011
0.080 + 0.033
0.198 + 0.040
0.026 + 0.021 0.584 + 0.034
Table 2: The energy distance between the empirical distributions of imputations and true missing values the noisy ring dataset (16). For each method and sample size, the average and standard deviation over 10
independent runs shown.
Sample Size
kNNSampler
Random Forest.
KNNImputer
kNN*KDE
Linear
3000
0.021 + 0.015
0.076 + 0.025
0.181 = 0.032
0.033 + 0.017
0.584 + 0.038
5000
0.019 + 0.015
0.069 + 0.023
0.216 = 0.055
0.024 : 0.013
0.576 + 0.042
7000
0.028 = 0.009
0.087 = 0.031
0.189 = 0.032
0.028 + 0.015
0.612 = 0.044
9000
0.028 + 0.022
0.074 + 0.027
0.197 + 0.043
0.020
0.013
0.593 + 0.033
11000
0.019 + 0.012
0.075 = 0.027
0.194 = 0.064
0.035 0.040
0.606 = 0.062
Table 3: The root mcan squared error of cach imputation method for different sample sizes on the lincar
chi-square dataset 15
The mean and standard deviation over 10 independent runs are shown for each
setting.
Sample Size
kNNSampler
Random Forest
kNNImputer
kNNxKDE
Linear
3000
2.691 + 0.151
2.338 + 0.154
2.117 + 0.158
2.876 = 0.126
1,885 + 0.195
5000
2.710 + 0.134
2.273 0.113
2.092 + 0.123
2.726 + 0.190
1.914
0.088
7000
2.729 + 0.102
2.307 0.118
2.100 + 0.185
2.789 + 0.135
1.895
0.121
9000
2.786 0.228
2.308
0.076
2.065
0.097
2.812 = 0.095
1.945
0.076
11000
2.793 0.188
2.388 0.127
2.055 0.116
2.708 + 0.184
1.913
0.154
Table 4: The root mean squared error of each imputation method for different sample sizes on the noisy ring The mcan and standard deviation over 10 independent runs are shown for cach setting.
dataset (16).
Sample Size
kNNSampler
Random Forest
kNNImputer
kNNxKDE
Linear
3000
2.680 + 0.238
2.309 0.133
2.073 + 0.169
2.811 + 0.112
1.951
0. 10108
5000
2.818 + 0.195
2.322 0.141
2.079 + 0.157
2.698 +: 0.130
1.870
0.123
7000
2.733 0.216
2.307 0.141
2.133 0.163
2.799 = 0.186
1.959
0.179
9000
2.638 + 0.146
2.281 + 0.103
2.138
0.141
2.637 + 0.177
1.923
0.157
11000
2.672 + 0.137
2.281 + 0.152
2.024 + 0.075
2.763 + 0.164
1.885 + 0.152
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.
For comparison, also report the root mean squared error (RMSE) for each method's imputations in Tables and 4 RMSE is expected to be smaller for regression-basee based methods, which estimate the conditional mcans of missing values and thereby minimize RMSE. A smaller RMSE docs not imply better recovery of the missing-value distribution. Indeed, imputations from the linear imputer have the lowest. RMSEs, but their distribution significantly differs from the distribution true missing values, as quantified in Figures [ and 2 and visually observed in Figures 3 and This result demonstrates that the RMSE is not a good metric for evaluating the distributional similarity between imputations and missing values. See Naf et al. 2023) for related discussion.
15

---

Published in Transactions on Machine Learning Research (12/2025)
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 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.
MR=0.1
MAR 80% CI
Ring MAR 90% C1
Ring MAR 95%
0.95
0.95
0.95
0.90
0.85
0.85
0.60
0.80
0.75
0.75
0.70
0.70
0.65
0.65
0.65
Linear- MAR B0% CI
Linear MAR 90%
Linear MAR-
0.95
0.95
0.95
0.90
0.90
0.65
0.85
0.80
0.80
0.75
0.75
0.75
0.70
0.70
0.70
1000
10000
100000
1000
10000
100000
1000
10000
100000
Sample Size
Sample Size
Sample Size
Figure Coverage probabilities of kNN prediction intervals different missing rates (MR) for different sample sizes. The mean and standard deviation over 10 independent runs are shown for each setting. The top three figures are on the noisy ring data, and the bottom three are on the linear chi-square data.
5 Real Data Experiments
Lastly, we present real-data experiments on solar power generated by photovoltaic panels, where miss- ing values are common due sensor failures and other factors (e.g., Phan al., 2023 Costa al. 2024). We use Kaggle datase that contains solar panel DC powers (responses) and the correspond- ing irradiations (covariates). totaling 67,698 covariate-response pairs. We randomly select subset of N covariate-response pairs from the full dataset. In this subset, we select randomly 30% of the units whose covariates are between 0.4 and 0.6 and set their responses to missing. These missing responses are imputed based on the remaining observed covariate-response pairs in the subset. We consider each of N {10, 000, 20, 000, 30, 000, 40, 000, 50, 000, 60, 0003. The configuration of each method follows Scction 4.1.3
This experiment is repeated 10 times independently for each setting, and the and standard devi- ation of the energy distance between imputations and true missing values are reported in Table 5 (see Scction 4.1.2 kNNSampler consistently gives lower energy distances than the other methods, this time in-
httpps:/wwwx.aggt http://wwww.aga..eeaaa/aaaaaaaaa
16

---

Published in Transactions on Machine Learning Research (12/2025)
cluding kNNx XDE. Moreover, KNNSampler's energy distance decreases the sample size increases, which aligns with its theoretical consistency in recovering missing-value distributions.
To understand the results, Figure G describes imputations by KNNSampler, kNNxKDE, and kNNImputer based on the full dataset. kNNx KDD's imputations do not capture well the heterogeneity and non-ne non-negativity of the missing-value distribution, as the imputations are sampled from Gaussian distributions with fixed, common variance. NNNputess's imputations are distributed more narrowly than the missing-value distribu- tion. In contrast, kNNSampler's imputations distributed similarly to the true missing values, successfully recovering the drissingvvalue distribution.
KNNSampler
KNNxKDE
Original Data
1400
Originel Data
1200
1200
Imputed Data
Imputed Data
DC PO.OOEERRR
1000
DC POWWRR
1000
800
800
600
600
400
400
200
200
0
0.2
0.4
0.6
0.8
1.0
0.0
0.4 IRRADIATION
0.8
IRRADIATION
KNNimputer
1400
Originel Dala
True Data
1200
Imputee
DC DC_ POWER
1000
800
600
400
200
0.0
0.2
0.4
0.6
0.8
1.0
IRRADIATION
Figure 6: Missing value imputations by kNNSampler, ANNxKDE, and KNNImputer on the full solar panel dataset in Section 5 with 30% missing rate under the MAR mechanism. True missing responses are shown in green, imputations in red, and the rest in blue.
Table 5: Comparison of the energy distance between the empirical distributions of imputations and true missing values across different sample sizes of the real solar panel dataset in Section 5 For each method and sample size, the average and standard deviation of the energy distance over 10 independent runs are shown.
Sample Size
kNNSampler 1.855 + 1.474
Random Forest
kNNImputer
kNNxKDE
Lincar
10000
5.544 = 1.916
10.619 = 2.958
3.333 = 2.609
190.546 + 16.701
20000
0.687 + 0.671
4.980 + 0.977
7.389 + 1.763
2.634 + 1.238
195.623 + 9.364
30000
0.500 + 0.309
5.112 + 1.207
4.874 + 1.144
2.273 + 0.990
195.412 + 3.468
40000
0.373 + 0.317
5.543 + 0.684
4.584 + 0.890
2.293 + 1.165
194.081 + 5.721
50000
0.190 0.138
5,667 0.968
4.161 + 0.881
2.559 + 1.133
198.261
4.039
60000
0.148 0.077
6.230 0.812
4.634 + 1.053
1.927 + 0.942
194.808 3.526
6 Conclusion and Discussion
We studied kNNSampler, a stochastic missing-vall value imputation method that imputes a missing response of given unit by searching for its k most similar units in terms of covariates and by randomly sampling one
17

---

Published in Transactions on Machine Learning Research (12/2025)
of the associated l observed responses. This method is interpreted sampling from an approximate KNN- based conditional distribution of missing response given a covariate. Assuming 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 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 conditional distribution. Empirical results demonstrate the capability of kNNSampler in recovering the distributions of missing values.
We discuss limitations of the current work, some of which stem from hot-dock methods in general (sce Andridge and Little 2010), and potential future directions. (i) While the experiments show the promising performance of &NNSampler, they are limited to low-dimensional covariates. Experiments with higher- dimensional covariates are needed to fully characterize the KNNSampler's practical performance. (ii) For higher-dimensional covariates, the choice of the distance function itsolf influences KNNSampler's performance, which should be invest investigated (iii) Leave-one-out cross-validation for selecting the number of nearest neigh- bours uses the mean square error, which should be modified to distributional error metrie. (iv) A further theoretical analysis is needed to understand how the distributional imputation quality affects subscquent analysis of quantity of interest, such as the well-calibratedness of uncertainty estimates.
Acknowledgments
We thank Frédéric Coutellier, Michele Bezzi and colleagues at SAP Labs France for their support and discussion. We also thank the reviewers and the Action Editor for their time and constructive feedback.
References
Andridge, R. R. and Little, R. J. (2010). A review of hot deck imputation for survey non-response. Inter- national Statistical Review, 78(1):40 64.
Costa, T., Falcão, B., Mohamed, M. A., Annuk, A., and Marinho, M. (2024). Employing machine learning for advanced gap imputation in solar power generation databases. Scientific Reports, 14(1):23801.
De Silva, H. and Perera, A. S. (2016). Missing data imputation using evolutionary k-nearest neighbor algorithm for gene expression data. In 2016 Sizteenth International Conference on Advonces in ICT for Emerying Regions (ICTer), pages 141-146. IEEE.
Enders, C. K. (2022). Ayplied Missing Dala Analysis. Guilford Publications.
Faisal, S. and Tutz, G. (2021). Multiple imputation using nearest neighbor methods. Information Sciences, 570:500-516,
Garefa-Lacncina, P. J., Sancho-Gómoz, J.-L., Figuciras Vidal, A. R., and Verleysen, M. (2009). nearest neighbours with mutual information for simultaneous classification and missing data imputation. Neuro- computing, 72((-9):1483 1493.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel two-sample test. Journal of Machine Learning Research, 13(1):723 773.
Grünewälder, S., Lever, G., Baldassarre, L., Patterson, S., Gretton, A., and Pontil, M. (2012). Conditional mean embeddings as regressors. In Proceedings of the 29th International Coference on International Conference on Machine Learning, pages 1503-1810.
Györfi, L., Kohler, M., Krzyzak, A., Walk, H., et al. (2002). Distribution-free Theory of Nonporumetric Regression. Springer.
Hastie, T., Mazumder, R., Lee, J. D., and Zadch, R. (2015). Matrix completion and low-rank SVD via fast alternating least squares. Journal of Machine Learning Research, (6(1):3377-3400
18

---

Published in Transactions on Machine Learning Research (12/2025)
Huang, J., Keung, J. W., Sarro, F., Li, Y.-F., Yu, Y.TT., Chan, W., and Sun, H. (2017). Cross-validatio based k nearest neighbor imputation for software quality datasets: an empirical study. Journal of Systems and Software, 132:226 252.
Kanagawa, M. (2024). Fast computation of leave-one-out cross-validation for k-NN regression. Transactions on Machine Learning Research.
Kanagawa, M., Hennig, P., Sejdinovic, D., and Sriperumbudur, B. K. (2025). Gaussian processes and reproducing kernels: Connections and equivalences. arXiv preprinl or Xiv:22506. 17366.
Kpotufe, S. (2011). k-NN regression adapts to local intrinsic dimension. Advances in Neural Information Processing Systems, 24.
Lalande, F. and Doya, K. (2023). Numerical data imputation for multimodal data sets: A probabilistic nearest-neighbor kernel density approach. Transactions on Machine Learning Research. Reproducibility Certification.
Li, 2., Meunier, D., Mollenhauer, M., and Gretton, A. (2022). Optimal rates for regularized conditional mean embedding learning. Advances in Newral Information Processing Systems, 35:4433-44 4445.
Li, Z., Meunier, D., Mollenhauer, M., and Gretton, A. (2024). Towards optimal sobolev norm rates for the vector-valued regularized least-squares algorithm. Journal of Machine Leurning Research, 25(181):1 51.
Lian, H. (2011). Convergence of functional k-nearest neighbor regression estimate with functional responses. Electronic Journal of Statistics, 5:31-40.
Little, R. J. and Rubin, D. B. (2002). Statistical Analysis with Missing Data. John Wiley Sons, 2nd edition.
Mattei, P.-. and Frellsen, J. (2019). MIWAE: Deep generative modelling and imputation of incomplete data sets. In International Conference on Machine Leurning, pages 4413 4423. PMLR.
Mohri, M., Rostamizadeh, A., and Talwalkar, A. (2018). Foundations of Machine Learning. MIT Press, second edition.
Muandet, K., Fukumizu, K., Sriperumbudur, B., Schölkopf, B., et al. (2017). Kernel mean embedding of distributions: A review and beyond. Foundations and Trends6 in Machine Learning, 10(1-2):1-141.
Murray, J. S. (2018). Multiple imputation: review of practical and theoretical findings. Statistical Science, 33(2):142 159.
Näf, J., Spohn, M.-L., Michel, L., and Meinshausen, N. (2023). Imputation scores. The Annals of Applied Statistics, 17(3):2452-244
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapcau, D., Brucher, M., Perrot, M., and Duchesnay, É. (2011). Scikit-learn: Machine learning in Python. Journal of Machine Learning Research,
12:2825 2830.
Peyré, G., Cuturi, M., et al. (2019). Computational optimal transport: With applications data science. Foundations and Trends0 in Machine Learning, 11(5-6))355 607.
Phan, Q.TT,, Wu, Y.-K., and Phan, Q-D. (2023). Enhancing ahedd probabilistic solar power forecast with hybrid transformer-lub model and missing data imputation. IEEE Transactions on Industry Applications, 60(1):1396-144) 1408.
Rubin, D. (1987). Multiple Imputation for Nonsesponse in Surveys. John Wiley Sons.
Rubin, D. B. (1976). Inference and missing data. Biometrika, 63(3):581-55) 592.
19

---

Published in Transactions on Machine Learning Research (12/2025)
Rubin, D. B. (1996). Multiple imputation after 18+ years. Journal of the American Statistical Association, (1444)):473-45 489..
Schafer, J. L. (1997). Analysis of Incomplete Multivoriate Data. CRC press.
Schafer, L. and Graham, J. W. (2002). Missing data: our view of the state of the art Psychological Methods, 7(2):147.
Sejdinovic, D., Sriperumbudur, B., Gretton, A., and Fukumizu, K. (2013). Equivalence of distance-basod and rkhs-based statistics in hypothesis testing. The Annals of Statistics, pages 2263 2291.
Shah, A. D., Bartlett, J. W., Carpenter, J., Nicholas, O., and Hemingway, (2014). Comparison of random forest and parametric imputation models for imputing missing data using MTCE: CALIBER study. American Journal of Epidemiology, 179(6):764 774.
Simon-Gabriel, C.-J., Barp, A., Schölkopf; B., and Mackey, (2023). Metrizing weak convergence with maximum mean discrepancies. Journal of Machine Learning Research, 24(184):1-22
Song, L., Fukumizu, K., and Gretton, A. (2013). Kernel embeddings of conditional distributions: A unified kernel framework for nonparametric inference in graphical models. IEEE Signal Processing Magazine, 30(4):98 111.
Song, L., Huang, J., Smola, A., and Fukumizu, K. (2009). Hilbert space embeddings of conditional distribu- tions with applications to dynamical systems. Proceedings of the 26th Annual International Conference on Machine Learning, pages 961-968.
Sriperumbudur, B. K., Gretton, A., Fukumizu, K., Schölkopf, B., and Lanckrict, R. (2010). Hilbert space embeddings and metrics on probability measures. Journal of Machine Learning Reseorch, 11:1517 1561.
Steinwart, I. and Christmann, A. (2008). Support Veclor Machines. Springer.
Stekhoven, D. J. and Bühhmann, P. (2012). type data. Bioinformatics, 2811):11
missing value imputation for mixed-
Stone, C. J. (1977). Consistent nonparametric regression. Annals of Statistics, pages 595-620.
Székely, G. J. and Rizzo, M. L. (2013). Energy statisties: A class of statistics based on distances. Journal of Statistical Planning and Inference, 14298):124-
Tang, F. and Ishwaran, H. (2017). Random forest missing data algorithms. Statistical Analysis and Data Mining: The ASA Data Science Journal, 10(6):363-377 377-377777...
Troyanskaya, O., Cantor, M., Sherlock, G., Brown, P., Hastic, T., Tibshirani, R., Botstcin, D.,, and Altman, R. B. (2001). Missing value estimation methods for dna microarrays. Bioinformalics, 17(6):520 525.
Tutz, G. and Ramzan, S. (2015). Improved methods for the imputation of missing data by nearest neighbor methods. Computational Statistics & Data Analysis, 90:84-99.
Yoon, J., Jordon, J., and Schaar, M. (2018). Gain: Missing data imputation using generative adversarial nets. Tn International Conference Machine Learning, pages 5689 5698. PMLR.
A Proof of Theorem 1
Proof. We proceed as the proof of Kpotufe 2011 Theorem 1) on real-valued kNN regression, with adapta- tions to our RKHS-valued kNN regression setting.
20

---

Published in Transactions on Machine Learning Research (12/2025)
The RKHS distance between the mean embeddings of the true and kNN conditional distributions is decom- posed into the "bias" and "variance" terms:
((P( x)) (P(: =)))
= =P(( =2)) - E x)) Xn]]
B x)) [x,] ((P( x=)) H² + 2 2 xx) Xm) Φ(P(. x0) 14
2 (PP( x)) B[ x)) X
(17)
Blas
Variance
where (P( =)) x is the conditional expectation of $(P(. x)) given Xn (a1,-.., the expectation being taken for the n output values 21....,,M Mun
E[ x)) x,
Σ
E(D(y)) X
Σ
$(P(: [zy)),
(18)
JEN.(z,R.X.)
JENN(S.
where the last identity follows from Vj ~ P(- xzy).
Lemma 2 in Section A. and Lemma 3 in Section A.2 respectively provide probabilistic upper bounds of the
bias and variance terms in the upper bound [17 each holding simultaneously for all X, k E (I...., n and > 0 satisfying the condition probabilistic bounds in 11
11 with probability at least - 5. The claim follows from using these
A.1 Bias Bound
Lemma below is from Kpotufe (0111 Lemma 1).
Lemma Suppose that Assumption 4 holds. Let D F1...., En iid. P(x) be an i.i.d. sample of size from E 8z, be the empirical distribution. Let 1. Then, probability distribution P on X, and Pn n
1
1 Σ Maa E B aa
(B
holds simultaneously for all balls B and for all constants a satisfying
P(B) 2 3a and 2 VB In2n) + In(8/8)
with probability at least 6.
Lemma 2. Suppose that Assumptions and 4 hold. Let (x1, ..... (Ens 3n) iid. P())z)P(z). Let
0 < < 1. Then the following bound holds with probability at least 1 8 simaltaneously for all X, 3 fI. and 0 < < Tmax satisfying the condition [1
3Cdisk Xm] S Ar (nP(B(a,r))
1/d
Proof. By using the triangle incquality and the Lipschitz continuity of the mapping ((P(. x)) in we obtain
Assumption D
((P(: ((P()1x)) - ((0(P(- Xm] H
(PP( (2)) -
Σ
@(P(- xx,))
Σ
( x)) ((P(- |x,1)}
K
jENN(z.k,X.))
1
Σ
()P((- x)) @PP(- (2;))
Σ Adx(z,zz) S Nmm((1),
(19)
JENN(a
21

---

Published in Transactions on Machine Learning Research (12/2025)
where (x) is the distance between and its k-th nearest neighbour in X,, This distance is bounded in the proof of Kpotufe (001 Lemma 2), which leads to the claimed bound. For completeness, prove it here.
The first inequality in the condition implies that
VB (2n) lm(8/6)
Define constant 0 < 6 < as
3C 3 nP((Bez r))
1/a
where follows from the second inequality in the condition
11 Then, Assumplion implies that
k 3a
P( er)) 2 Caist ed (P(B(x,
71
Thus, Lemma T with this choice of a implies that the following holds simultaneously for all a E X, e [I,..., n) and 0 < Tmax satisfying the condition (11) with probability least 1 6:
k = n
Pa (B
where the second identity follows from that ( is the distance between x and its k-nearest neighbour, so the ball of center æ and radius (x) contains h points from ,n....,In- This implies that
(x) er Sr ( 3Caisk 1/d nP(((axr))
simultancously holds for all E X, k E [l....,n} and 0 < < Tmax satisfying the condition probability at least. -. The claim obtained by using this and the bound (19
11 with
A.2 Variance Bound
Lemma 3. Suppose that Assumptions and d hold. Let (x1,91),... (Fns Un) Lid. P(y)x))(x). Let < 6 <
1. The following bound simallaneously holds for all x E X and l (I.....n)} with probabilily at least 1 8:
(P( x)) Xal - (P(( m))|(?
+ (Vs In(n) In()))
<202er
(20)
k
Proof. Denote by (NNN(I, Xw)) the left hand side of the inequality (00
without the square:
((NN(I,k,x,)) [ x)) X.ll - X.l- (((((-1))] #
(21)
@(P(- ( xz)) - <P xx)
.ENN(
H
Σ
$(P(- x3)) - 2(1,)
.E..(Z.
H
where the last expression follows from the definition of (P x:) 8
The notation (<NN(a, k; X,))
emphasizes that it depends only on the subset of training data (I1, ......., (Zn, By)) associated with the indices NN(x,, k X,) of the k-nearest neighbours of in X,, (z1,-..,
22

---

Published in Transactions on Machine Learning Research (12/2025)
Because of the bound 9 changing yi for any NN(x, k, X,) to any different value of y changes the value of (,, k; Xn)) at most 2/2Cker/ kk This can be shown as follows. Let write the last expression of 21 with the original ys and the one with yi replaced by y as
(NN(a,, k x, 11A Bl
K, 1| + Bl
where
1
A = (P(( xz)) - ((v)),
A' : k
( (x,) -(())),
k =
Σ
((P( xz)) ((1)).
(E, and jitt
The triangle inequality implies that
|| + Bl S ||Albb + ||
||4' + BlIx 2 ||B|| - |1
Therefore,
(NN k x.)) (,NN(,, k x,)D,, = =IA 1l - '1A' + Bllu
$ ||Al|/ + |1 (11B112 - [|A')bx) ||Al| ||AA'1l
Similarly,
(N ,INN(,, S 111A|x + ||Al
Hence,
4 (NN(o, k. X.m),,, A,, S 1|AA|| + ||A
1 xx)) (v)))l +
1 || x.))) +100l S
2 2v20ker
where the last inequality follows from the bound 9
On the other hand, the output yi associated with any non-k-nearest neighbours i & NN(x, k X,) does not appear in (NNN(r, k X,.)), so changing the value of y in this case does not change v (NN(x, ki X,)).
Thus, for fixed Xn, the probability that the random variable (NNN(z, kt Xn)) exceeds its expectation ([N((N(( k x,)) plus any positive constant e 0 is upper bounded by using McDiarmid's incquality as
(B X.)) (I,N((,, [x,]+ X I exp (-
ek 4C2 ker
(22)
This is a bound for fixed r and k.
Next, for fixed Xn, we consider the probability that the statement ((NN(I,k,X,)) > E [NNN((z, k X.)) X,] te
(23)
holds for some x E X' and & E (1.....n). The number of distinct such statements is identical to the number of distinct index sets ncarest neighbours NN(x, k Xn), since the random variable (NN(2,, k, X,)) depends (1,,),,..., (Im, yn) associated with NN(x, k, X.), as mentioned previously. In other
only on the subset words, if there are other x' X and k' [I....,n] that give the identical index set of nearest neighbours as for and k i.c.,
NN(x', X,) = NN(x, k. X,).
then the random variable (NN((Z) R, X.)) for x' and k is identical to that for x and
(NN((' K X.)) (NN((,, k X.)).
23

---

Published in Transactions on Machine Learning Research (12/2025)
The number of distinct index sets of nearest neighbours is identical to the mumber of distinct ways the set Xn of points is intersected by balls B(r, (k,)())) of center and radius Tk,n (2) being the distance of the k-th nearest neighbour from 2. This number is upper-bounded by the number of distinct ways Xn is intersected by the class B = {B(x, r)
E X. r > 0} of all balls, which further upper-bounedd by n' with the VC dimension ViS of (K potulle 2011 p.6). Therefore, by using the union bound, the probability that the statement 23 holds for some x and l is upper bounded by the bound (22) times Vr:
Pr (KNN(22 XX,) Ele(NN(z, k X.)) X for some a X and {I....,n. n2 Xn)
S n exp (-
22 for all 6> 0.
(24)
40 ker
Now, set
e
2 4Chor (VBln(n) - In()))
n Ve exp
4C
For any value of 0 8 < 1, there is a corresponding 0. Then, the bound (24 implics that, for fixed
Xn, the following upper bound on the random variable (2, X,)) squared holds for all X and & nI with at least probability 1- 6:
(NNN(,, x2 (NN( k X.)) xx) 2e 2NN k. X.))) [x]] 26
where the second inequality follows from Jensen's 's inequality. Replacing ei (NN(,, k, xx,)) )) by its definition (21) and e² by the above expression, we obtain the following bound on the variance term that holds for all x and is with probability at least 1 5:
Σ jE,N(
(P((1 (Pp( x))
1
8C2 ker
(VBs In(n) In()))
s 2E
Σ
((P(-120)) - (x)) X
(25)
K
h
Define H-valued random variables
2 (P(- xz)) - (wv) for all j EE NN(II,
These random variables conditionally independent given Xn. The conditional expectation of each 2j
given X,, is zero, and the conditional variance is uniformly upper bounded due to the bound (
E 2
Xa]
= E (P( (3))) + X.. = क (PP( xz)) - E (15)) [xg] = 0,
||
2x
Xx, [ 1 - (v;) 11 |وح < 2C
Therefore, the first term in the bound 25 can be expressed as (see also the definition of ( x)) in
E
((P(. (2,()12,)) (31)
X
E
2j
Xn
k
K
X.,
JENN(a,k,X.))
H
Σ
1
Σ
(2j, 2m) H Xn
12
JENN(Z.k,..)
Σ
E
[xn]
Σ
(E [zy Xx.] 1E [zm
K
(ENN(2.R.X.))
202 ker
k
24

---

Published in Transactions on Machine Learning Research (12/2025)
The proof completes by using this expression in the bound 25 and noting that this bound is independent of X,.
25


Reviews:
[{"review_id": "xPW52AanV9", "reviewer": "Reviewer_c5Gg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents kNNSampler, a nearest-neighbour oriented random hot deck data imputation method and conducts a theoretical analysis to establish that it is a principled approach. \n\nThe theoretical contributions include an error bound, in terms of maximum mean discrepancy, between the kNN and true conditional distributions for a given, fixed covariate. The bound captures the impact of the number of observed response-covariate pairs (n) and the number of nearest neighbours (k). The bound implies statistical consistency of the kNN conditional distribution, provided that the number of nearest neighbours increases to infinity at a rate slower than n. \n\nThe bound is derived by analyzing the mean embedding of the kNN conditional distribution in a RKHS. This can be viewed as a novel estimator for the mean embedding of the true conditional distribution (the conditional mean embedding). The bound provides consistency and convergence rates for the introduced kNN estimator. \n\nThe methodology conducts an error analysis on RKHS-valued regression where the response variable is infinite-dimensional. The analysis establishes that the required sample size for a given level of precision increases exponentially with the covariate’s intrinsic dimension, as opposed to its ambient dimension.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The main claims of the paper are theoretical in nature. The theoretical results are based on reasonable assumptions, the theorem statements are clear, and I did not detect any errors in the proofs. These claims are supported by accurate, convincing, and clear evidence. \n\nThe remaining claim (in the abstract) is that “Experiments demonstrate its effectiveness in recovering the distribution of missing values”. \n\nThe paper provides results for two data settings and compares to relative basic baselines (apart from the kNN x KDE). The support for this claim, as stated, is weak. There are only two investigated datasets, and both are synthetic and low-dimensional. To support this general claim, the experiments need to encompass a much broader range of datasets, pushing to high dimensions and preferably investigating performance in real-data settings where data is removed in various ways. Alternatively, since the primary contribution of the paper is theoretical, the claim could be modified to “illustrate… for low-dimensional synthetic data”. In my view, this wouldn’t diminish the contribution of the paper, because the value is the theory, and the experiments are primarily providing insights and some support.", "audience_yn": "Yes", "audience_explanation": "Imputation is a very important problem and the kNN-based imputation is widely used, so the findings of the paper would be of use and interest to multiple individuals in TMLR's audience.", "requested_changes": "The paper makes a strong theoretical contribution and provides some welcome support for kNN-based and hot deck methods. The following changes are requested:\n\n(1)\tAs discussed above, the final sentence of the paper claims that “Experiments demonstrate its effectiveness in recovering the distribution of missing values.” While the experiments are illustrative and provide useful insights, they are not sufficiently extensive or thorough to support the claim that they “demonstrate [the] effectiveness”. Either the claim should be modified or more comprehensive experimentation should be included. Since this claim only appears in the abstract, and the experiments are not a critical component of the contribution, I would suggest the former – “Experiments illustrate…” or “Experiments provide insights into the behaviour and performance…”\n\n(2)\tThe second sentence of the paper reads “Imputation — filling in missing values with their estimates — is a common preprocessing step…” Since it is (almost always) a preprocessing step, it is important to keep in mind that what is really important is its effect on the subsequent processing. But the paper largely ignores this, focusing on the output of the pre-processing step. The theoretical results focus on the conditional distribution, with analysis conducted in terms of its mean embedding. While it is natural to assume that the improved matching of the distributions leads to a meaningful improvement in downstream analysis, the paper doesn’t really establish this. There isn’t theory or experiments to quantify the gains that emerge in subsequent regression or estimation of population-based quantities. The paper should be modified to include either a more compelling discussion or to include some theory/experiments that establish the value of targeting the conditional distribution. Lalande and Doya (2023) provide imputation results. The paper would be more compelling if there were results establishing that the imputation quality (either in terms of NRMSE or likelihood) isn’t inferior to kNN x KDE. \n\n(3)\tIt would be helpful to include a section discussing limitations of the approach. Some of these stem from random hot-deck methods in general and I don’t think they are addressed by the approach. There are settings where suitable neighbours (donors) are not available because of the relationship between the covariates and missingness. In these settings it is possible for parametric (and possibly the KDE) approach to perform better. While these settings are probably ruled out by the missingness and smoothness assumptions for the presented theory, it would be helpful to have a discussion by the authors that clarifies when the technique may prove less effective.", "broader_impact_concerns": "None", "additional_comments": "It isn’t obvious to me how the 2/d in (10) becomes d/2 in (12). If this is correct, then I would recommend some text be added to explain the transition.", "extra_scores": {}, "extra_text": {}}, {"review_id": "hW0dvsGDp0", "reviewer": "Reviewer_gEd5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a theoretical analysis of kNNSampler, a data imputation method that fills in missing values by randomly selected a response from the k nearest neighbors based on covariates. \n\n### Contributions\n\n1.\tThe authors prove weak convergence between the estimated and true conditional distributions of missing values given the covariates. Using maximum mean discrepancy as the distance metric between distributions, they demonstrate that this distance converges to zero when both k/n and ln(n)/k converge to zero, where n is the total number of units used, and k is the number of nearest neighbors.\n2.\tThe proof employs the kernel mean embedding framework, which maps probability distributions to a reproducing kernel Hilbert space (RKHS). The authors adapt the proof by Kpotufe (2011) for real valued kNN regression to their specific imputation context.\n\n3.\tThe experimental results demonstrate that kNNSampler outperforms commonly used imputation methods in capturing the conditional distributions of mission values when the conditional distributions of the response given the covariates are asymmetric or multi-modal.\n\n### Strengths\n\n1.\tClear high-level message: The paper clearly demonstrates that regression-based imputation methods are inadequate when the goal is to recover the true distribution of missing data through imputation. The figures in the introduction provide digestible visual arguments for why this is the case.\n2.\tPractical algorithm: The algorithm being analyzed is simple, straightforward algorithm, and can easily be implemented in practice.\n\n### Weaknesses\n\n1.\tProof issue in Lemma 3: see detailed discussion in the next section", "claims_and_evidence_yn": "No", "claims_explanation": "Overall, the paper provides complete proofs in the appendix and effectively explains the key ideas of the methods and the proof techniques in the main text. The empirical evidence also supports the paper arguments about kNNSampler.\n\nHowever, I believe there is an error in the proof of Lemma 3, specifically in bounding $|\\psi(\\text{NN}(x,k,X_n))\\mid_{y_i} -\\psi(\\text{NN}(x,k,X_n))\\mid_{y'_i} |$. The issue appears in the final equality where the terms for $j \\in \\text{NN}(x,k,X_n)$ with $j \\neq i$ should be included but appear to be omitted. Please correct me if I’m mistaken. While I believe the statement itself is correct and this error is fixable, it should be addressed in the final version.", "audience_yn": "Yes", "audience_explanation": "The results of this paper are relevant to the TMLR's audience. The paper provides a theoretical convergence analysis for a simple randomized variation of k-Nearest Neighbor imputation and makes a compelling case for when and why this approach should be used.", "requested_changes": "1.\tProof: Correct the error in Lemma 3.\n2.\tTypos and notation: \n\n    a.\tThe power of 2/d in equations (10) and (12) is inconsistent.\n\n    b.\t The use of curly brackets in the proof of Lemma 3 might confuse the reader.\n\n3.\tFormatting: Figure 9 currently appears in the middle of the references section and should be moved before the references.", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hXAeyUckGe", "reviewer": "Reviewer_PmYc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a simple kNN-based method for the imputation of missing data. The method, called kNNSampler, estimates the conditional distribution of missing values using a k-nearest neighbor (kNN)–based empirical distribution, rather than imputing single-point estimates such as the mean or median of the kNN samples.\n\nMain Contributions:\n\nImputation framework: Instead of returning a single deterministic imputed value, kNNSampler approximates the conditional distribution P(y|x) by averaging the responses of the k nearest neighbors. This allows for stochastic imputations that are more realistic when p(y|x) is multimodal, as well as uncertainty quantification for each missing value.\n\nTheoretical convergence analysis: The authors show that the empirical kNN conditional distribution converges to the true conditional distribution in Maximum Mean Discrepancy (MMD) as the sample size n and the number of neighbors k increase suitably. They adapt classical kNN regression convergence proofs to the Hilbert space–valued setting, providing a bias–variance decomposition and rates of convergence under standard smoothness assumptions.\n\nKey Strengths:\n\nClear presentation: The paper is well-structured and readable, clearly stating the problem, assumptions, and theoretical claims.\n\nTheoretically sound: The convergence analysis is rigorous and grounded in established nonparametric regression theory and RKHS embedding literature.\n\nReasonable empirical validation: The experiments are consistent with the theoretical claims and demonstrate practical effectiveness on two artificial datasets.\n\nKey Weaknesses:\n\nMotivation for the energy distance used in the experiments: While the energy distance is a valid metric, the paper could better justify why it was chosen over alternatives (e.g., MMD, Wasserstein distance) and discuss its interpretability in the imputation context. For instance, is there a theoretical connection between the energy distance and the Maximum Mean Discrepancy used in the analytical part of the paper?\n\nLimited empirical scope: The experiments are conducted primarily on two artificially generated (toy) datasets. While these are useful for illustrating scenarios in which kNNSampler is more advisable than kNNImputer, they may not fully demonstrate the method’s performance in realistic scenarios. Including results on standard benchmark datasets would provide stronger evidence of the method’s practical applicability.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The approach is clearly motivated, simple, and supported by a convincing theoretical convergence analysis.", "audience_yn": "Yes", "audience_explanation": "The paper proposes a simple procedure based on kNN that is more appropriate for multimodal conditional distributions p(y|x) than the classic imputation approach based on the mean or median of the kNN samples, and its convergence is theoretically proved.", "requested_changes": "In the introduction, it would be beneficial to better motivate the stochastic sampling approach compared to the mean/median approach by explaining the differences between regression and imputation. In regression, the goal is typically to minimize the expected loss between the prediction y' and the true (but unknown) value y given x (e.g., squared error). In contrast, in imputation, the objective is to generate samples that follow the true conditional distribution p(y|x). The introduction could start by clearly distinguishing these two goals.\n\nSome claims about the superiority of kNNSampler over kNN x kDE should be moderated. Both methods yield very similar results, and the experiments have been conducted only on two toy datasets.\n\nIt is not clear how the parameter k is selected in the experiments for kNNSampler. Is this value different from the one used in kNNImputer (k=5)? Do the authors use the theoretical guidance suggesting k ≈ n^(2/(2+d))?\n\nFigures 5–9 should be replaced with tables, as tables occupy less space and provide more precise numerical information. In any case, Figure 9 should be moved from the references section to the main text.\n\nI would remove the section about p-values, since the results based on the energy distance already provide overlapping and sufficient evidence.\n\nIn Section 4.3, the uncertainty quantification analysis should include a similar procedure for kNNImputer, using it as a benchmark for comparison.", "broader_impact_concerns": "No concern", "additional_comments": "No comment", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)