DiffPuter: Empowering Diffusion Models for Missing Data Imputation
Authors: Hengrui Zhang
Liancheng Fang
Qitian Wu
Philip S. Yu
Venue: ICLR 2025 Spotlight (2025)
Keywords: Diffusion models
missing data imputation
Primary area: generative models
TL;DR: This paper combines EM algorithm and a Diffusion model for missing data imputation

Abstract: Generative models play an important role in missing data imputation in that they aim to learn the joint distribution of full data. However, applying advanced deep generative models (such as Diffusion models) to missing data imputation is challenging due to 1) the inherent incompleteness of the training data and 2) the difficulty in performing conditional inference from unconditional generative models. To deal with these challenges, this paper introduces DiffPuter, a tailored diffusion model combined with the Expectation-Maximization (EM) algorithm for missing data imputation. DiffPuter iteratively trains a diffusion model to learn the joint distribution of missing and observed data and performs an accurate conditional sampling to update the missing values using a tailored reversed sampling strategy. Our theoretical analysis shows that DiffPuter's training step corresponds to the maximum likelihood estimation of data density (M-step), and its sampling step represents the Expected A Posteriori estimation of missing values (E-step). Extensive experiments across ten diverse datasets and comparisons with 17 different imputation methods demonstrate DiffPuter's superior performance. Notably, DiffPuter achieves an average improvement of 8.10\% in MAE and 5.64\% in RMSE compared to the most competitive existing method.

Published as a conference paper at ICLR 2025
DIFFPUTER: EMPOWERING DIFFUSION MODELS FOR MISSING DATA IMPUTATION
Hengrui Zhang' Liancheng Fang'
Qitian Wu² Philip S. Yu'
'Computer Science Department, University of Illinois at Chicago "Eric and Wendy Schmidt Center, Broad Institute of MIT and Harvard (hzhan55, 1 psyu) @ui mic... edu wugitiaan@brrccdissittteee.
ABSTRACT
Generative models play an important role in missing data imputation that they aim to learn the joint distribution of full data. However, applying advanced deep generative models (such as Diffusion models) to missing data imputation is chal- lenging due 1) the inherent incompleteness of the training data and 2) the difficulty in performing conditional inference from unconditional generative mod- els. To deal with these challenges, this paper introduces DIFFPUTER, a tailored diffusion model combined with the Dnpectation-aximmiaan for missing data imputation. DIFFPUTER iteratively trains a diffusion model learn the joint distribution of missing and observed data and performs accu- rate conditional sampling to update the missing values using a tailored reversed sampling strategy. Our theoretical analysis shows that DIFFPUTER' training step corresponds to the maximum likelihood estimation data density (M-step), and its sampling step represents the Expected A Posteriori estimation of missing values (E-step). Extensive experiments across ten diverse datasets and comparisons with 17 different imputation methods demonstrate DIFFPUTER's superior performance. Notably, DITEPUTER achieves an average improvement of 6.94% in MAE and 4.78% RMSE compared the most competitive existing method.
(EM) algorithm
INTRODUCTION
In the field of data science and machine learning. missing data in tabular datasets is common issue that can severely impair the performance of predictive models and the reliability of statistical analysis. Missing data can result from various factors, including data entry errors, non-responses in surveys, and system errors during data collection (Barnard & Meng, 1999; Lillard al., 1986; Eckert et al., 2020). Properly handling missing data is essential, as improper treatment can lead to biased estimates, reduced statistical power, and invalid conclusions.
A plethora of work proposed over the past decades has propelled the development of missing data imputation research. Early classical methods often relied on partially observed statistical features to impute missing values were based on conventional machine learning techniques, such as KNN (Pujianto et al., 2019), or simple parametric models, such as Bayesian models (Ma & Chen, 2018) or Gaussian Mixture Models (Garcfa-Laencina et al., 2010). With the advent of deep learning, recent research has primarily focused on predictive (Stekhoven & Bühlmann, 2012; Van Buuren & Karin, 2011; Kyono et al., 2021) or generative deep models (Yoon al., 2018; Mattei & Frellsen, 2019; Richardson et al., 2020) for missing data imputation. Predictive models learn to predict the target entries conditioned on other observed entries, guided by masking mechanisms (Du et al., 2024) or graph regularization techniques (You et al., 2020; Zhong et al., 2023). By contrast, generative methods learn the joint distribution of missing entries and observed entries and aim to impute the missing data via conditional sampling (Richardson et al., 2020: Nazabal et al., 2020b; Mattei & Frellsen, 2019; Yoon al., 2018; Ouyang al., 2023; Zheng & Charoenphakdee, 2022). However, generative imputation methods still fall short compared to predictive methods even when employing state-of-the-ar generative models (Kingma Welling, 2013; Rezende & Mohamed,
"Corresooniigg author.

---

Published as a conference paper at ICLR 2025
2015; Goodfellow al., 2014; Ho et al., 2020). We think this is primarily due to the incomplete likelihoud nature issue in missing data imputation: generative models need to estimate the joint distribution of missing data and observed data. However, since the missing data itself is unknown, there is an inherent error in the estimated data density. The classical Expectation- Expectation-Mximmiaaii Maximization (EM) algorithm (Dempster et al., 1977) offers elegant route handle this issue, being capable of addressing the incomplete likelihood issue by iteratively refining the values of the missing data.
Integrating EM algorithms with generative models has been extensively studied (Richardson et al., 2020), however, its combination with Diffusion models (Ho et al., 2020), currently the most powerful generative models, is still unexplored. In the M-step, the diffusion models have been shown able to faithfully reconstruct the ground -round-truth distribution. However, in the E-step, it is usually considered challenging to utilize them to perform condition inference (e.g., predicting missing entries based on observed entries). This is because diffusion models directly model and generate the complete joint distribution of data across all dimensions simultaneously, lacking the flexibility found in VAE and GAN-based approaches (Ma et al., 2020; Peis et al., 2022b; Ma et al., 2018).
This paper introduces DIFFPUTER, a principled generative method for missing data imputation. DIFFPUTER explores a path that makes diffusion models compatible with the EM framework, allowing both E-step and M-step to be effective. Specifically: 1) In the M-step, DIFFPUTER employs diffusion model to learn the joint distribution of the missing and observed data. With the powerful ability of diffusion models to learn tabular data distributions, in the M-step. DIFFPUTER learns parameterized distribution that faithfully recovers the joint distribution of observed and missing entries. 2) In the E-step, DIFT-UTER uses the learned diffusion model to perform flexible and accurate conditional sampling by mixing the forward process for observed entries with the reverse process for missing entries. Theoretically, we show that DIFFPUTER'S 's M-step corresponds to the maximum likelihood estimation of the data density, while its E-step represents the Expected A Posteriori (EAP) estimation of the missing values, conditioned on the observed values.
We conduct experiments on nine benchmark tabular datasets containing both continuous and discrete features under various missing data scenarios. We compare the performance of DIFFPUTER with 17 competitive imputation methods from different categories. Experimental results demonstrate the superior performance of DIFFPUTER across all settings and on almost all datasets.
2 RELATED WORKS
Iterative Methods for Missing Data Imputation. Iterative imputation is a widely used approach due to its ability to continuously refine predictions of missing data, resulting in more accurate imputation outcomes. This iterative process is especially crucial for methods requiring an initial estimation of the missing data. The Expectation-Maximle Maximization (EM) algorithm (Dempster et al., 1977), a classical method, can be employed for missing data imputation. However, earlier applications often assume simple data distributions, such as mixtures of Gaussians for continuous data or Bernoulli and multinominal densities for discrete data (Garefa-Lacncina et al., 2010). These assumptions limit the imputation capabilities of these methods due to the restricted density estimation of simple distributions. The integration of EM with deep generative models remains underexplored. A closely related approach is MCFlow (Richardson et al., 2020), which iteratively imputes missing data using normalizing flows (Rezende & Mohamed, 2015). However, MCFlow focuses on recovering missing data through maximum likelihood rather than expectation, and its conditional imputation is achieved through soft regularization instead of precise sampling based on the conditional distribution. Beyond EM, the concept of iterative training is prevalent in state-of-the -hrt--rrtttttttte deep learming-basee imputation methods. For instance, IGRM (Zhong et al., 2023) constructs graph from all dataset samples and introduces the concept of friend networks, which are iteratively updated during the imputation learning process. Hyperlmpute (Jarrett et al., 2022) proposes an AutoML imputation method that iteratively refines both model selection and imputed values.
Diffusson Models for Missing Data Imputation. We are not the first to utilize diffusion models for missing data imputation. In computer vision, diffusion models have been widely applied in image inpainting, either in the data space (Lugmayr et al., 2022) or the latent space (Corneanu et al., 2024). In the tabular data area, TabCSDI (Zheng& & Charoenphakdee, 2022) employs conditional diffusion
'he code is available at htct:::///gitubb.cc/////gg/ 28ang98 28/11 wtt.t///ppfffeeeeeeeee/e ffPuter.
2

---

Published as a conference paper at ICLR 2025
model to learn the distribution of masked observed entries conditioned on the unmasked observed entries. MissDiff (Ouyang et al., 2023) uses a diffusion model to learn the density of tabular data with missing values by masking the observed entries. Although Miss was not originally intended for imputation, it can be easily adapted for this task. Other methods, despite claiming applicability for imputation, trained on complete data and evaluated incomplete testing data (Zhang et al., 2024; volicocur-Martined et al., 2024). This approach contradicts the focus of this study, where the training data itself contains missing values. Additionally, all the aforementioned methods use one-step imputations, which overlook the issue that missing data in the training set can lead to inaccurate data density estimation. By contrast, the proposed DIFFPUTER is the first to integrate a diffusion- based generative model into the EM framework. Additionally, we achieved accurate conditional sampling by mixing the forward and reverse processes of diffusion and demonstrated the effectiveness of this approach through theoretical analysis.
3 PRELIMINARIES
3.1 MISSING VALUE IMPUTATION FOR INCOMPLETE DATA
This paper addresses the missing value imputation task for incomplete data, where only partial data entries are observable during the training process. Formally, let the complete d-dimensional data be denoted as ~ Pdata (x) E IRd. For each data sample, x, there is binary mask m E {0, 1]dd indicating the location of missing entries for x. Let the subscript & denote the k-th entry of vector, then m& 1 stands for missing entries while my = 0 stands for observable entries.
We further use xobs and mis to denote the observed data and missing data, respectively (i.e., (xobs x) Note that ob is the fixed ground-truth observation, while x x is conceptual and unknown, and we we aim to estimate it. The missing value imputation task aims to predict the missing entries x based on the observed entries xois
In-sample vs. Out-of-sample imputation. The missing data imputation task can be categorized into two types: in- un-sample and out-of-sample. In-sample imputation means that the model only has to impute the missing entries in the training set, while out-of-sample imputation requires the model's s capacity to generalize to the unseen data records without fitting its parameters again. Not all imputation methods can generalize to out-of-sand out-of-sample imputation tasks. For example, methods that directly treat the missing values as learnable parameters (Muzellec et al., 2020; Zhao et al., 2023) are hard to apply to unseen records. A desirable imputation method is expected to perform well on both in-sample and oft-ol- imputation tasks, and this paper studies both of the two settings.
3.2 FORMULATING MISSING DATA IMPUTATION WITH EXPECTATTON-MAXXIIIAAUU
Treating hhs as observed variables and mis as latent variables, with the estimated density of the complete data distribution parameterized as po(x), ), we can formulate the missing value imputation (EM) algorithm. Specifically, when the complete data distribution pa (x) is available, the optimal estimation of missing values is given by x mise Py oos 0). Conversely, when the missing entries x mis are known, the density parameters can be optimized via maximum likelihood estimation: 0 arg max p(x ( Consequently, with the
problem using the pnpeetaaioonMxximizatee
0
initial estimation of missing values mis the model parameters 0 and the missing value mis can be optimized by iteratively applying M-step and E-step:
Maximization-stop: Fix mis update 0 = arg max p (x = arg max pe obe xml))).
Expectation-step: Fix 0, update x miss
-Egeii p mis ob a 0).
4 METHODOLOOGY
In this section, we introduce DIFFPUTER- Iterative Missing Data Imputation with Diffusion. Based on the Expectation- Maximizaiinnniiiiiionnnn (EM) algorithm, DIFFPUTER updates the density parameter and hidden variables x mis in iterative manner. Fig. shows the overall architecture and training
3

---

Published as a conference paper at ICLR 2025
Kiterstions
E stop
Cne hot chceding
Bose mon reputotte
cecent exaree
Hear
Diffusion Modal: pa(a)
Syn
Diffusion Model pu(a)
Hule Eide
1ite
Fix gh estiiatte
Fix B. estimate ait
Fest
(s nil)
(a .nn
Syn
DR
EEC
33
Figure I: An overview of the architecture of the proposed DIFFPUTER. DIFFPUTER utilizes one-hot encoding to transform discrete variables into continuous ones and use the mean of observed values to initialize the missing entries. The EM algorithm alternates the process of 1) fixing x and estimate diffusion model parameter 0, 2) fixing and estimate x for K iterations. The final imputation result x" is returned from the E-step of the last iteration.
process of DIFFPUTER: 1) The M-step fixes the missing entries mis then a diffusion model is trained estimate the density of the complete data distribution po (x) po x 2) The E-step fixes the model parameters 0, then we update the missing entries x via the reverse process of the learned diffusion model (x))). The above two steps are executed iteratively until convergence. The following sections introduce the M-step and E-step of DIFTPUTER, respectively. To avoid confusion, we use x, X xbbs etc. to denote samples from real data, x, X, gobs etc. denote samples obtained by the model 0, while X, X, ghhs etc. to denote the specific values of variables.
4.1 M-STEP: DENSITY ESTIMATION WITH DIFFUSION MODELS
Given an estimation of complete data x = (xobs, x x M-step aims to learn the density of x, parameterized by model 0, i.e., (x). Inspired by the impressive generative modeling capacity of diffusion models (Song et al., 2021b; Karras et al., 2022), DIFFPUTER learns po through diffusion process, which consists of forward process that gradually adds Gaussian noises of increasing scales to x, and a reverse process that recovers the clean data from the noisy one:
xn + oo)e, E N.O, dx, = -26(()((() logg(x;)dl 2/ (2
(Forward Process)
(1)
(Reverse Process)
In the forward process, Xp = = (xobe ( is the currently estimated data at time 0, and x, is the diffused data at time t. (tt) = =tit is the noise level (and o(t) its derivative w.r.t. t). i.c., the standard deviation of Gaussian noise, at time f. The forward process has defined a series of data distribution p(x,)= = fx, p(xx) ( )dxg and p(xo) = p(x). Note that when restricting the mean of X, to and keeping the variance of x. small (e.g., via standardization), p(xx) approaches a tractable prior distribution πt (xx) at t T when o (T') is large enough, meaning p(xr) IT (xx (Song al., 2021a). In our formulation in Eq. 1, p(x)) RS T (xx) (O, o2 (T)I).
In the reverse process, Vx, log p(x;) is the gradient of 's log-probability w.r.t., x,, and is also known as the score function. W; the standard Wiener process. The model is trained by (conditional) score matching (Song et al., 2021b), which utilizes a neural network ex(x,, () (called denoising/score network) to approximate the conditional score-function Vx, ((x,( (x)), which in expectation approximates Vx, log p(x:):
LSM = Bre
Vx, log ((; (o) where x; = xo)e (3)
Since the score of conditional distribution has analytical solutions. i.c.. Vx, log = Vx,p(x,/<o)) (Xo)
X.-Xi
-- Eq. 3 can be interpreted training neural network (x(,,) to
0²(()
4

---

Published as a conference paper at ICLR 2025
approximate the scaled noise. Therefore, co(X,, t) is also known as the denoising function, and in this paper, it is implemented as a five-layer MLP (see Appendix D.4).
Remark 1. Starting from the prior distribution XT ~ π (xx), and apply the reverse process in Eq. with Vx, log x(x)) replaced with (x(x,, 1), we obtain a series of distributions po po(x)), VI [0, 7]
Remark 2. (Corollary in (Song et al., 2021a)) Let p(x) = x(x)) be the data distribution, and (x( = po(Xo) be the marginal data distribution obtained from the reverse process, then the score-matching loss function in Eq. 3 is an upper bound of the negative- -og-tiie--o-llkeehhooodddd of real data ~ p(z) over 0. Formally,
pXx (og [o s LSM (0) + Const,
(4)
where Const is constant independent of O.
Remark 2 indicates that the 9 optimized by Eq. 3 approximates the maximum likelihood estimation of data distribution p(x). Consequently, pe (a) approximates p(x) accurately with enough capacity. The algorithmic illustration of DIFFPUTER M-step is summarized in Algorithm 1.
4.2 E-STEP: MISSING DATA IMPUTATION WITH A LEARNED DIFFUSION MODEL
Given the current estimation of data distribution po (x), the E-step aims to obtain the distribution of complete data, conditional the observed values, i.e. po (x ubb),, such that the estimated complete data x" can be updated by taking the expectation, i.e., x Ex pe (x) (x))
When there is an explicit density function for (x),,, when the conditional distribution pa (x x)x) is tractable (c.g., can be sampled), computing Ex p(x|x*b<
0) becomes feasible. While most deep generative models such VAEs and GANs support convenient unconditional sampling from pe (x), they do not naturally support conditional sampling, e.g., X po (x((xx">>) Luckily, since the diffusion model preserves the size and location of features in both the forward diffusion process and the reverse denoising process, it offers a convenient and accurate solution to perform conditional sampling po (xxxxxaing from an unconditional model (x) = pe xxxxbbe x
Specifically, let x be the data to impute, x xghs be the values of observed entries, m be the 0/1 indicators of the location of missing entries, and x, be the imputed data at time I. Then, can obtain the imputed data at time t - At via combining the observed entries from the forward process on X = x, and the missing entries from the reverse process on X; from the prior step t (Lugmayr et al., 2022; Song et al., 2021b):
foward x t-Ar
= x + ol - At) e, where e ~ I),
(Forward for observed entries) (5)
d-
x A = =x,+ +
dx,, where dx; is defined in Eq. 2 (Reverse for missing entries) (6)
X (1
forward m (-At Xt
(Merging)
Based on the above process, starting at random noise from the maximum time T, i.e., Xy \( (T))), we can obtain reconstructed Xo. such that the observed entries of 8a are the same as those of Xg. i.e., xgbs = &gbe, In Theorem 1, we prove that via the algorithm above, the obtained Xo is sampled from the desired conditional distribution, i.e.. Xn ~ po(x = =ggh)..
Theorem 1.
Let &T be sample from the prior distribution (x) = =N OT x be the data to impute, and the known entries of x are denoted by ohs gohs. The score function Vx, log p(x:) is approximated by neural network e so(x,.) Applying Eq. 5, Eq. 6, and Eq. 7 iteratively from t= T > until
with At -> O, then Xo is sample from po po(x), under the condition that its Sohis Formally,
observed entries babs
Xo ~ Po (s =x0bg
(8)
See proof in Appendix B.I. Theorem demonstrates that with a learned diffusion model O, we are able to obtain samples exactly from the conditional distribution pe (x xb))) through the aforementioned imputation process. For inference, we use Monte Carlo estimation to compute the expectation of the Σ (()) /N.
missing values, i.e., x" Expo(x (x
5

---

Published as a conference paper at ICLR 2025
Discretization. The above imputing process involves recovering x, continuously from t T to t= O, which infeasible in practice. In implementation, we discretize the process via M + discrete descending timesteps t.. t tt. to. where tm T, to 0. Therefore, starting from XIM ~ (t o = x0be), ), we obtain po xt, xobs = xgbe) from = M to 0.
Algorithm 1: M-step: Density Es-
Algorithm 2: E-step: Missing Data Imputation
timation using Diffusion Model
Input: Score network co (x t). Data with missing
Input: Data samples from p(x).
values x E Rd mask m E (0, 13". Number of samples N, Number of sampling steps M.
Output: Score network
Output: Imputed data x
while not converging do
Sample x ~ p(x)= E Rdd
for j 1 N do Sample x ~N(0,, ((0,002 (ta,)t for i t M to 1 do
Sample ~ p(t)
Sample (O, E Rd
Xo x
forward, 4-11
((4-))) (i)
X₁ = Xo + o(1)
reverse,1) (j)
X + pli pla-1 dx
(())- = |x c(t)
5
(
forward, (5) + +(1 +11-m) reverse,
Update via Adam optimizer
6
X = -m 1 -mox
1
N 1== N = - (55) /N
x
Since the desired imputed data x xiis is the expectation, i.c., x Exmis px pix bhs x we sample Xyii for N times, and take the average value as the imputed x miss. The algorithmic illustration of DIFFPUTER'S E-step summarized in Algorithm 2.
4.3 IMPLEMENTATIONS
Data Processing. Diffusion models are inherently suited for continuous data but not for discrete data. Given that real-world tabular data often contains both types of data, we use one-hot encoding to convert each dimension of discrete data into multi-dimmnsional 0/1 encoding, thereby treating as continuous data After that, we perform standardization on each column of data, ensuring that has mean of 0 and variance of 1. The final imputation performance is measured by the difference between the predicted value and the ground- tround-truth missing value after standardization.
Initialization of Missing Data. The execution of the EM algorithm requires the initialized values of missing entries m0), which might have huge impact on the model's convergence. For simplicity, we initialize the missing entries of each column to be the mean of the column's observed values, cquivalent to setting x mis( everywhere (since the data has been standardized).
Training. To obtain more accurate estimation of complete data DIFFPUTER iteratively executes the M-step and E-step. To be specific, let x
(x xx)s(()) ) be the estimation of complete data at k-th iteration, (k) be the diffusion model's parameters at k-th iteration, we write ((++1) as a function of ((k), i.e., ((k+1) = M-step (x(*)), and (k+1) as function of (k++1) and (A), i.e., (k+1)) = E-step (x()) ((++1)). Therefore, with the initialized (0) and the maximum iteration K, we are able to obtain (k) from = to K.
(x*hh(k)), (xnhh(()))bbb(())
Inference. For ia-sa imputation, the imputed values obtained iteratively during the training process. For out-of-sample imputation, with the observed incomplete data x, the mask m, and the trained score network eg. we directly apply the E-step once in Algorithm 2 to obtain the imputation values x*.
5 EXPERIMENTS
In this section, we conduct experiments to study the efficacy the proposed DIFFPUTER in missing data imputation tasks.

---

Published as a conference paper at ICLR 2025
5.1 EXPERIMENTAL SETTINGS
Datasets. We evaluate the proposed DIFFPUTER ten public real-world datasets of varying scales. We consider five datasets of only continuous features: California, Letter, Gesture, Magic, and Bean, and four datasets of both continuous and discrete features: Adult, Default, Shoppers, and News. The detailed information of these datasets is presented in Appendix D.2. Following previous works (Muzellec et al., 2020; Zhao et al., 2023), we study three missing mechanisms: MCAR, MAR, and MNAR. Differences between the three settings are in Appendix D.3. In this section, we only report the performance in the MCAR setting, while the results of the other two settings are in Appendix E. In the main experiments, we set the missing rate as = 30%. For each dataset, we generate 10 masks according to the missing mechanism and report the mean and standard deviation of the imputing performance.
Baselines. We compare DIFFPUTER with 16 powerful imputation methods from different categories:
1) Distribution-mate mattrbbiionnmtccingg matching methods based on optimal transport, including TDM (Zhao et al., 2023) and MOT (Muzellec et al., 2020). 2) Graph-based imputation methods, including GRAPE (You et al. 2020): pure bipartite graph-based framework for data imputation, and IGRM (Zhong et al., 2023): a graph-based imputation method that iteratively reconstructs the friendship network. 3) Iterative methods, including EM with multivariate Gaussian priors (Garefa-Laencine et al., 2010), MICE (Van Buuren& & Karin, 201 1), MIRACLE (Kyono et al., 2021), SoftImpute (Hastic et al., 2015), and MissForest (Stekhoven & Bühlmann, 2012). 4) Deep generative models, including MIWAE (Mattei & Frellsen, 2019), GAIN (Yoon et al., 2018), MCFlow (Richardson et al., 2020), MissDiff (Ouyang et al., 2023) and TabCSDI (Zheng & Charoenphakdee, 2022). It is worth noting that MissDiff' and TabCSDI are also based on diffusion. We also compare with two recent SOTA imputing methods ReMasker (Du et al., 2024) and Hyperlmpute (Jarrett et al., 2022).
Evaluation Protocols. For each dataset, we use 70% as the training set, and the remaining 30% as the testing set. All methods are trained on the training set. The imputation is applied both the missing values in the training set and the testing set. Consequently, the imputation of the training set is the 'in-sample" setting, while imputing the testing set is the "out-of-sample" setting. The performance of DIFFPUTER is evaluated by the divergence between the predicted values and ground-truth values of missing entries. For continuous features, we use Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE), while for discrete features, we use Accuracy. Note that both the MAE and RMSE are calculated based on the input data after standardization (zero mean and unit variance). The implementation details and hyperparameter settings are presented in Appendix D.4.
5.2 MAIN RESULTS (MASK COMPLETRLY AT RANDOM)
We first evaluate DIFFPUTER's performance in the in-sample imputation task. Figure 2 compares the performance of different imputation methods regarding continuous columns using MAE and RMSE. Table I compares the performance of different methods regarding discrete columns using accuracy. Our observations on these experimental results are summarized as follows.
Superior performance of DIFFPUTER. Across all datasets, DIFFPUTER provides high-quality imputation results, matching the best methods on some datasets and significantly outperforming the second-best best methods on the remaining datasets. Compared with the other two diffusion-baatd based methods, MissDiff and TabCSDI, the substantial performance advantage of DIFFPUTER convincingly demonstrates the correctness and superiority of combining the EM algorithm with Diffusion.
Traditional machine learning methods are still powerful imputors. A well-recognized claim in tabular data machine learning is that traditional machine learning methods can sometimes be better than deep learning methods (Lalande & Doya, 2022; Suh & Song, 2023; Mlicoeur-Martines et al., 2024), and we have similar observations in missing data imputation. For example, the simple EM algorithm under the mixture-Gaussiar assumption and the KNN algorithm give decent imputation performance and outperform a lot of early deep generative imputation methods.
The preference of different types of algorithms for heterogeneous tabular data. Another in- teresting finding is that discriminative methods (e.g.. Remasker, GRAPE, MOT) and generative methods (e.g., DIFFPUTER) have different preferences for continuous data and discrete data. Overall, generative models appear to be more effective at imputing continuous columns. As shown in Figure 2, model significantly outperforms these discriminative models. Conversely, discrininative models
7

---

Published as a conference paper at ICLR 2025
DiffPuter Remaskerg Hyperimputes
GRAPE
TabCSDI
MCFlow Softimpute
Missforest
EM
TDM
MIWAE
MIRACLE
Miss
MOT S
GAIN
MICE
KNN
1.0
0.8
0.6
MAE
0.5
0.4
California
Letter
Gesture
Magic
Bean
Adult
Default
Shappers
News
Average
1.0
RMSE
0.8
0.5
0.6
0.0
0.4 Average
Califomia
Letter
Gesture
Magic
Bean
Adult
Default
Shoppers
News
Dataset
Figure 2: MCAR, In-sample imputation performance on continuous columns: Comparing DIFF- PUTER with 17 baselines on imputing continuous data all the nine datasets. A blank column indicates that the method fails gets out-of-memor for that dataset. DIFFPUTER outperforms the most competitive baseline method by 6.94% (MAE score) and 4.78% (RMSE score) by average. The circled number after the model name denotes its ranking among all methods.
Ifomia
Lotter
Gesture
California
1.0
3.0
0.6
SPNGS
0.5
AMGE
0.75
MAE
2.5
0.5
0.50
0.4
0.4
10
Defaut
Shoppers
Adult
196
fnct
0.75
3.00
MAi
ANGE
1.0
0.2
MAE
0.35
0.6
0.50
0.6
0.25
Number samples:
Number of samples: N
Figure 3: Impacts of the number of EM iterations. DIFFPUTER' performance steadily improves as the number of EM interactions increases.
Figure 4: Impacts of the number of sampled imputations per iteration. A very small N leads to poor performance and large variance.
are more proficient at imputing discrete data: Table 1, the best discriminative model, Remasker, can achieve results comparable to our method on almost all datasets.
Next, we compare the performance the out-of satmlle imputation tasks. Noting that some im- putation methods are specifically designed for in-sample imputation and cannot be applied to the out-of-sample setting, the number baselines this task significantly reduced. Table in Ap- pendix E compares the MAEs and RMSEs in the OOS imputation task. Comparing it with the results of in-sample imputation, we can easily observe that some methods exhibit significant perfor- mance differences between the two settings. For example, craph-based methods GRAPE and IGRM perform well in in-sample imputation, but their performance degrades significantly out-of-sample imputation. IGRM even fails on all datasets in the out-of-sample imputation setting. In contrast, DIFFPUTER demonstrates excellent performance in both in-sample and out-of-sample imputation. This highlights DIFFPUTER' 's superior performance and robust generalization capabilities.
5.3
ABLATION STUDIES
Impacts of the number of EM iterations. In Fig. 3, we present the performance of DIFFPUTER'S imputation results from increasing EM iterations. Note that represents the imputation result of randomly initialized denoising network, and & represents the performance of pure diffusion model without iterative refinement. clearly observed that single diffusion imputation achieves only suboptimal performance, while DIFFPUTER'S performance steadily improves the number of
8

---

Published as a conference paper at ICLR 2025
Table 1: MCAR, In-sample imputation performance on discrete columns. Comparison of the imputation accuracy on discrete columns on five datasets. DIFFPUTER ranks the first among 19 imputation methods. Mean/Median/MF means using the mean/median/most frequent value as the imputation, which will give the same imputation result in this case.
Method
Adult
Default
Shoppers
News
Average
Rank
Statistical
Mean/Median/ME
55. 55.20
63.37
50.60
18.602o
44.44
15
Traditional iterative
EM (Garefa-Laen et al., 2010)
61.27+00.2
57 37.80=0.22
50.91.0.27
39.37=0.06
52.35
(Van Buuren & Karin, 2011)
50.31
51.88go
43.43 1o
30.05
43.92
MIRACLEF (Kyono 2021)
62.2810.31
55-79
45.74.10.0
39.33..
50.79
SoftImpute (Hastic al., 2015)
56.1819.08
56.01 20.07
50.94.0..7
18.70....00
45.46
MissForest (Stekhoven & Bühlmann, 2012)
63.51.1..3
56.97.1.225
51.1010..1
37.5522..4
52.28
Dist. match
MOT (Muzellec ul., 2020)
63.8620.16
71.95
56.91 (0 2.0.1
38. 22.662.0.0
57.83
TDM (Zhao et al., 2023)
62.64.1.36
63 =3.41
53.00
30.62
52.92
GNN
GRAPE (You et al., 2020)
69.844.0.30
73.85
57.322..11
40.83 44
60.46
IGRM (Zhong et al., 2023)
69.21.10.29
OOM
57.45.10.1
0OM
Generative models
MIWAE (Matti & Frellsen, 2019)
51.3310.0
48.1
49: 49.23310.01
17.32.0
41.50
18
GAIN (Yoon al., 2018)
53.7319.38
54.49 1.49.1.99 90
47.25
34.20 20.000..00
47.42
MCFlow (Richardson al., 2020)
60.31
68.3610.4
52.73110.0
30.08.3.00
52.87
TabCSDI (Zheng & Charoenphakdee, 2022)
58.72+0.25
60.57=0..0
43.24-+0.06
17.19==014
42. 12.51
MissDiff (Ouyang et al., 2023)
57.39+1.84
66. 00.01=1..s
48. 10.1.0
41.39==0.50
53.2 53.22
Other
KNN (Pujianto et al., 2019)
63.2 19.52
71.64.0.5
53.93
38.59
56.86
Remasker (Du et al., 2024)
89).18g0.
76.8829.62 62
17.86..0.0
44.33_.o.o
62.06
Hyperlmpute (Jarrett et al., 2022)
64.90 64.9920.00 0.00
2..39.2.07
59.19
40.29
59.72
DITPPUTER
70.12=0.11
64:00.a
58 28.8220.00
44 69.69
62.82
Training Time
MAE
0.500
tiwrr time
1800
2010
1830
2.550
1630
0.55
07
20
0:33 1630
15"
30
30
10
Califorii: Sampligg Step
Adul: Sampling
Figure 5: Impacts of sampling steps on the training time and imputation performance (MAE). Reducing diffusion sampling steps greatly reduces the training time at the cost of a slight performance drop.
Figure 6: Impacts of extremely large miss- ing ratios. The performance of DIFFPUTER is upper-bounded by the initialized missing values via mean imputation
iterations increases. Additionally, we observe that DIFFPUTER does not require a large number of iterations converge. In fact, iterations are sufficient for DIFFPUTER to converge stable and satisfying state.
Impacts of the number of sampled imputations per iteration. To obtain the expected value of missing entries, we have to sample a number of examples from the conditional distribution (via the reverse diffusion process). Therefore, the number N should have huge impact on DIFFPUTER performance. In Figure 4, study the impacts of the number of samples the imputation results. As expected, a very small N leads to poor performance and large variance, while increasing the sampling number consistently improves the performance and reduces the variance. An optimal and stable performance can be achieved at N 2 10. Since the sampling time is linear w.r.t to the sample number, N 10 by default.
Impacts of the number of sampling steps in diffusion. Another hyperparameter impacting the performance and training speed is the number of sampling steps M in the reverse process of diffusion.

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of the real training time on an Nvidia RTX 4090 GPU. DIFFPUTER has similar training cost to SOTA methods, yet brings from 8% to 25% performance improvement.
Datasets
MOT
TDM
GRAPE
IGRM
Hyperimpute
Remasker
DIFFPUTER
Time: California
446.47s 396.60 688
216.91s
635.7s
1267.5e
1277.3s 1806.9s
1s
1927 2s
Time: Adult
4.566
2347.1s
3865. 18
1902.4s
2142.9s
DIFFUTER'S perf. improve.
21.47%
21.37%
25.94%
20.97%
8.44%
10.65%
By default, DITPPUTER set M 50, and we study the impact of reducing M in Figure 5, As demonstrated, reducing the diffusion sampling steps can greatly reduce the training time at the cost of only a slight performance drop. More detailed, reducing the sampling step from 50 to 20 reduces about 25% of training time cost, with only 3% performance drop. Therefore, when time costs are high, and precision requirements are not stringent, the training speed can be further increased by reducing the sampling steps.
Impacts of large missing ratios. In Figure 6, study the performance of DIFFPUTER with missing ratio from 30% to 99% (observed ratio from 70% to 1%). Consistent with the intuition, when the observed ratio decreases (i.e.. missing ratio increases), there is significant increase in the imputation MAE score, indicating that the imputation performance becomes increasingly worse. However, we can observe a clear performance lower bound when the missing ratio approaches 99%, which is close to the performance of directly imputing using the mean of observed values. This is because, in PUFF'UTER's 's first E-step, each column's missing values are always initialized as the mean of the observed ones. Therefore, DIFFPUTER can still give reasonable imputation result even if the missing ratio is extremely large.
Comparison of Training Time. Intuitively, the training of DIFFPUTER would be very time- consuming, as it requires iteratively alternating between the training and sampling processes of the diffusion model. However, in practice. DIFFPUTER's 's training efficiency is significantly im- proved through the following design: 1) weightee MLP-architecturee denoising function; 2) reduced number of sampling steps (M = 50 < 1000 in traditional diffusion methods). To demon- strate this, we compare the training time of DIFFPUTER with other SOTA imputation methods. As demonstrated in Table 2, DIFFPUTER has similar training cost (similar scale) these representa- tive SOTA methods, yet brings from 8% to 25% performance improvement, which is acceptable.
Incorporating EM with other DGMs. Finally, we
Table 3: Effects of combining EM with other Deep Generative Models.
try to combine EM algorithms with other deep gener-
ative models that are easy to perform conditional sam-
pling, including MIWAE (Mattei & Frellsen, 2019), HIWAE (Nazabal et al., 2020a), VAEM (Ma et al., 2020) and HH-M (Peis et al., 2022a). Based on Table 3, our experimental results demonstrate that combining EM with other Deep Generative Models leads to performance improvements. However, it's noteworthy that despite these improvements, our pro- posed DIFFPUTER still outperforms all combinations, achieving superior results, where lower scores indicate better performance. This demonstrates the effectiveness and robustness of our approach compared to existing methods, even when they are enhanced through combination with EM.
Methods
Adult
Default
Shoppers
News
MIWAE
0.5763
0.5194
0.5047
0.6349
HIWAE
0.6155
0.3989
0.4707
0.5052
VAEM
0.5568
0.4292
0.4626
0.5204
HH-VAEM
0.5673
0.4589
EM MIWAE
0.5681
0.40833
0.4554
0.4171
HHWAE
0.5974
EM + VAEM
0.5492
0.4121
0.4362
0.5045
RM HH-
0.5402
0.4262
PUREP
0.3425
0.2661
0.3485
0.2855
6 CONCLUSIONS
In this paper, we have proposed DIFPPUTER for missing data imputation. DIFFPUTER is an iterative method that combines the ExpectationMMaximizztioo algorithm and diffusion models, where the diffusion model serves as both the density estimator and missing data imputer. We demonstrate theoretically that the training and sampling process of a diffusion model precisely corresponds to the M-step and E-step of the EM algorithm. Therefore, we can iteratively update the density of the complete data and the values of the missing data. Extensive experiments have demonstrated the efficacy of the proposed method.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported in part by NSF under grants I11-2106758, and POSE-2346158.
REFERENCES
John Barnard and Xiao-Li Meng. Applications of multiple imputation in medical studies: from aids to nhanes. Statistical methods in medical research, 8 8(1):17-36, 1999.
Ciprian Corneanu, Raghudeep Gadde, and Aleix M Martinez. Latentpaint: Image inpainting in latent space with diffusion models. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 4334-4343, 2024.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the algorithm. Journal of the royal statistical society: series B (methodological), 39(1): 1-22, 1977.
Tianyu Du, Luca Melis, and Ting Wang. Remasker: Imputing tabular data with masked autoencoding. In International Conference on Learning Representations, 2024.
Fabian Eckert, Teresa C Fort, Peter K Schott, and Natalic Yang. Imputing missing values in the us census bureau's 's county business patterns. Technical report, National Bureau of Economic Research, 2020.
Pedro Garcfa-Laencine José-Luis Sancho-Gómez, and Anibal R Figueiras- Vidal. Pattern classifica- tion with missing data: a review. Neural Computing and Applications, 19:263-282, 2010.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Proceedings of the 27th International Conference on Neural Information Processing Systems, pp. 2672-2680, 2014.
Trevor Hastie, Rahul Mazumder, Jason D Lee, and Reza Zadeh. Matrix completion and low-rank svd via fast alternating least squares. The Journal of Machine Learning Research, 16(1):3367-3402 2015.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, pp. 6840 6840-6851, 2020.
Daniel Jarrett, Bogdan C Cehere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperim- pute: Generalized iterative imputation with automatic model selection. International Conference on Machine Learning, pp. 9916-9937. PMLR, 2022.
Alexia Jolicoeur-Martineau, Kilian Fatras, and Tal Kachman. Generating and imputing tabular data via diffusion and flow-based gradient-booster trees. In International Conference on Artificial Intelligence and Statistics, pp. 1288- 1296. PMLR, 2024.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Proceedings of the 36th International Conference on Neural Informa- tion Processing Systems, pp. 26565-26577 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114 2013.
arXiv preprint
Akim Kotelnikov, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko. Tabddpm: Modelling tabular data with diffusion models. International Conference on Machine Learning, pp. 17564- 17579. PMLR, 2023.
Trent Kyono, Yao Zhang, Alexis Bellot, and Mihaela der Schaar. Miracle: ctusally-award imputation via learning missing data mechanisms, Advances Neural Information Processing Systems, 14:23806-2388 2021.

---

Published as a conference paper at ICLR 2025
Florian Lalande and Kenji Doya. Numerical data imputation: Choose knn over deep learning. In International Conference on Similarity Search and Applications. pp. 3-10. Springer, 2022.
Lee Lillard, James P Smith, and Finis Welch. What do we really know about wages? the importance of nonreporting and census imputation. Journal of Political Economy, 94(3, Part 1):489-506, 1986.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. ap..1146. 11461-11471, 2022.
Chao Ma, Sebastian Tschiatschek, Konstantina Palla, José Miguel Hernnneez- Sebastian Nowozin, and Cheng Zhang. Eddi: Efficient dynamic discovery of high-value information with partial vae. arXiv preprint arXiv: /110, 2018.
Chao Ma, Sebastian Tschiatschek, Richard Turner, José Miguel Hernández-Lobato, and Cheng Zhang. Vaem: a deep generative model for heterogeneous mixed type data. Advances in Neural Information Processing Systems, 33:1 33:11237-112444 11247, 2020.
Zhihua Ma and Guanghui Chen. Bayesian methods for dealing with missing data problems. Journal of the Korean Statistical Sociery, 47:297-313, 2018.
Alexandre Mattei and Jes Frellsen. Miwae: Deep generative modelling and imputation of incomplete data sets. In International conference on machine learning, pp. 4413-4423, PMLR, 2019.
Boris Muzellec, Julie Josse, Claire Boyer, and Marco Cuturi. Missing data imputation using optimal transport. In International Conference on Machine Learning, pp. 7130-7140. PMLR, 2020.
Alfredo Nazabal, Pablo M Olmos, Zoubin Ghahramani, and Isabel Valera. Handling incomplete heterogeneous data using vaes. Pattern Recognition, 107:107501, 2020a.
Alfredo Nazabal, Pablo M Olmos, Zoubin Ghahramani, and Isabel Valera. Handling incomplete heterogeneous data using vaes. Pattern Recognition, 107:107501, 2020b.
Yidong Ouyang, Liyan Xie, Chongxuan Li, and Guang Cheng. Missdiff: Training diffusion models tabular data with missing values. arXiv preprint arX1v:2307.00467, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: an imperative style, high- performance deep learning library. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 8026-8037, 2019.
Ignacio Peis, Chao Ma, and José Miguel Hemmndezz-Lbatt Missing data imputation and acquisition with deep hicrarchical models and hamiltonian monte carlo. Advances in Neural Information Processing Systems, 35:35839-3555 35851, 2022a.
Ignacio Peis, Chao Ma, and José Miguel Hemández-LLohat Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo. Advances in Neural Information Processing Systems, 35:35839-358511 2022b.
Utomo Pujianto, Aji Prasetya Wibawa, Muhammad Igbal Akbar, et al. K-nearest neighbor (k-nn) based missing data imputation. 2019 5th International Conference on Science in Information Techholooy (ICSITech). pp. 83-88. IEEE, 2019.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Trevor W Richardson, Wencheng Wu, Lei Lin, Beilei Xu, and Edgar A Bernal. Mcflow: Monte carlo flow models for data imputation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 14205-14214, 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems, 2021a.
12

---

Published as a conference paper at ICLR 2025
Yang Song, Jascha Sohl-Dickstein. Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In The Ninth International Conference Learning Representations, 2021b.
Daniel Stekhoven and Peter Bühlmann. nonssforet--oo--paaammttit missing value imputation for mixed-type data. Bioinformatics, (8(1):112-113 2012.
Heajung Suh and Jongwoo Song. A comparison of imputation methods using machine learning models. Communications for Statistical Applications and Methods, 30(3)3331-34, 2023.
Stef Van Buuren and Karin. mice: Multivariate imputation by chained equations in r. Journal of
statistical software, 45:1-67, 2011.
Jinsung Yoon, James Jordon, and Mihaela Schaar. Gain: Missing data imputation using generative adversarial nets. In International Conference on Machine Learning, pp. 5689-5698, PMLR, 2018.
Jiaxuan You, Xiaobai Ma, Yi Ding, Mykel J Kochenderfer, and Jure Leskovec. Handling missing data with graph representation learning. Advances in Neural Information Processing Systems, 33: 19075-19087, 2020.
Hengrui Zhang, Jiani Zhang, Balasubramaniam Srinivasan, Zhengyuan Shen, Xiao Qin, Chris- tos Faloutsos, Huzefa Rangwala, and George Karypis. Mixed-type tabular data synthesis with score-based diffusion in latent space. In The twelfth International Conference Learning Repre- sentations, 2024.
He Zhao, Ke Sun, Amir Dezfouli, and Edwin V Bonilla. Transformed distribution matching for missing value imputation. In International Conference on Machine Learning. pp. 42159 42189-42186. PMLR, 2023.
Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular data, In NeurIPS 2022 First Table Representation Workshop, 2022.
Jiajun Zhong, Ning Gui, and Weiwei Ye. Data imputation with iterative graph reconstruction. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pp. 11399 1407, 2023.
13

---

Published as a conference paper at ICLR 2025
A
DEFINITION OF SYMBOLS
In Table 4, we explain the definition of every symbol used in this paper.
Table 4: Explanations of the symbols used in this paper.
Symhol
Explanation
d
the number of columns of tabular data random variable, representing a row tabular data the distribution of the complete data hinary mask vector
E R²
p(x)
= Pdita(()
E {0, 1]d
oba
the observed part of x, indicated by my 0 the missing part of indicated by my
mis
the timestep of diffusion process
(1)
the noise level timestep the maximum timestep standard Gaussian noise unforwarded x (at timestep 0)
6 I)
Xg
X xbb
observed part missing part of x
X x mis
x (t)E
forwarded timestep t
x abe x xbs + (t)
observed part of x missing part
xo (t)e
0
diffusion model's parameters
Pe (x()
density data induced by the diffusion model sampled random noise for the reverse process observed part Ky missing part of R sampled data at timestep t in the reverse process observed part of X+ missing part of x:
X
x (x) = NN (T)1)
x7
8
X,
X
The exact ground-truth value(s) observed part
B PROOFS
B. PROOF FOR THEOREM
Proof. First, it is obvious that the observed entries from Eq. 7, i.c., goba satisfy
gobs p(x x =20b9), VL.
Then, we introduce the following Lemma.
Lemma 1. Let Xe ~ po (x ohs xgh0 If &I- de is a sample obtained from Eq. 5, Eq. 6, and Eq. then St. po (x: ael obs >2gb)) when Al 0+.
Proof. First, note that x, (x--1t) obtained via adding dimennsinnlll-inleeeedenn Gaussian noises Xo x (see the forward process in Eq. 1), we have
Po Po(x1_Ac(x*
20 gobs)
Po (x A1--A11 x obs
&o0s
pl (xpoe obs = & Do (x mis
ohs
(10)
Conespond to sobe in Eq.
14

---

Published as a conference paper at ICLR 2025
Therefore, Xi-Ai's observed entries X; is sampled from distribution p(x Ar x Then, we turn to the missing entries, where we have the following derivation:
Po(x xxp(1)
po (x;: Xx, x xeph"")) po Xx x x
xobs
E po(x x s seght (xx Ael x2, xobs Segbr)
(11)
po(x xovr = Po (x)
= po (x (Xx),
where &; is random sample from po pxxxx (x xg)t)). The first "Rv' holds because when At - tiis is almost predictable via without x0 Therefore, when X ~ po(x, xxobs xg()) po( (oox-=eexxx = =0) is approximately tractable via po (x
Note that A denotes the missing entries of Xr-Ar. According to the reverse denoising process in Eq. 2, given X; = X1, X is obtained via integrating dx, from to - At, iee.
- Xe
dx,.
(12)
Therefore, miss = xiss obtained from Eq. 6 is a sample from po po( (x))) And, approximately, x ~ pa (x, Ar =x0ba).
Since xo0e ~ p(x xobs xgge), xt- -r Po (x; |xobs po px - = sggh), according Eq. 10. Therefore, the proof for Lemma is completed.
&gbe), ) we have x,-
With Lemma we are able to prove Theorem via induction, as long as Ëy is also sampled from p(xT|x 2gbs). This holds because p(x7 (Xo) x p(XT), given the condition that Xo x has zero mean and unit variance, and (TT) >> 1.
Note that the score function Vx, lox:(x;) in Eq. is intractable, it is replaced with the output of the score neural network Eø (x,,t) Therefore, the finally obtained distribution can be rewritten as po xbx >g Therefore, the proof for Theorem is complete.
C FURTHER EXPLANATION OF THE DIFFUSION MODEL
The diffusion model we adopt in Section 4.1 is actually a simplified version of the Variance-Exploding SDE proposed in (Song et al., 2021b).
Note that (Song et al., 2021b) has provided unified formulation via the Stochastic Differential Equation (SDE) and defines the forward process of Diffusion as
dx = fx, f)dt + g(t) dur::
(13)
where f(-) and g(-) are the drift and diffusion coefficients and are selected differently for different diffusion processes, e.g., the variance preserving (VP) and variance exploding (VE) formulations. the standard Wiener process. Usually, J(.)) of the form f(x, () = f(1) x. Thus, SDE can be equivalently written as
dx = f(t) x df + g(t) dw:.
(14)
Let x be function of the time f, i.e., x(t), then the conditional distribution of x, given X (named as the perturbation kernel of the SDE) could be formulated as:
p(xx|xo) (xx (()xx, 8² (1)
(15)
where
<(1) exp ((ME)) and o(1)
I.
gP(5) s²(E)
de.
(16)
15

---

Published as a conference paper at ICLR 2025
Therefore, the forward diffusion process could be equivalently formulated by defining the perturbation kernels (via defining appropriate s(t) and (t))..
Variance Exploding (VE) implements the perturbation kernel Eq. 15 by setting s(t) indicating that the noise is directly added to the data rather than weighted mixing. Therefore, The noise variance (the noise level) is totally decided by (t). The diffusion model used in DIFFPUTER belongs VE-
SDE, but use linear noise level (i.e., (t) t) rather than (t) v in the vanilla VE-SDE (Song et al., 2021b). When (() = 1, the perturbation kernels become:
p xo) (xx 0, o²)
xx + ot)e
(17)
which aligns with the forward diffusion process in Eq. 1. The sampling process of diffusion SDE is given by:
dx [f(x, - 9 ( log): + g(()dur.
(18)
For VE-SDE, s(t) = + f(x,t) t) = f(t) O, and
(1)) V V. (P(() = ( (()2())S o1())
(t) dt²(t) (t)
(19)
dt
s(1) = (1)(()((1)
Plugging g(t) into Eq. 18, the reverse process in Eq. 2 is recovered:
dx; (2) logp(x,)dt (2)())
(20)
D EXPERIMENTAL DETAILS
D.1 CONFIGURATIONS.
We conduct all experiments with:
Operating System: Ubuntu 22.04.3 LTS
CPU: Intel 13th Gen Intel(R) Core(TM) 19-13900K
GPU: NVIDIA GeForce RTX 4090 with 24 GB of Memory
Software: CUDA 12.2, Python 3.9.16, PyTorch (Paszke al., 2019) 1.12.1
D.2 DATASETS
We real-world datasets of varying scales, and all of them available Kaggle² or the UCI Machine Learning repository'. We consider five datasets of only continuous features: California', Letter', Gestur Magic?, and Bean", and five datasets of both continuous and discrete features: Adult,, 10 Shoppers" and 12 The statistics of these datasets are presented in Table 5.
com/
16

---

Published as a conference paper at ICLR 2025
Table Statistics datasets. # Num stands for the number of numerical columns, and Cat stands for the number of categorical columns.
Dataset
Rows
Num
Cat
Train (In-sample)
Test (Out-of-Sample)
California Housing
20, 640
14,303
6,337
Letter Recognition
20, 000
14,000
6,000
Gesture Phase Segmentation
9,522
6, 66665
2, 2,7
Magic Gamma Telescope
19, 020
13, 314
5, 706
Dry Bean
13, 610
9,527
4, 083
Adult Income
32, 561
8
22, 792
9, 9,769
Default of Credit Card Clients
30, 000
14
10
21,000
9,000
Online Shoppers Purchase
330
10
7
8,631
3,699
Online News Popularity
39,644
2
27, 790
11,891
D.3 MISSING MECHANISMS
According to how the masks generated. there are three mainstream mechanisms of missingness, namely missing patterns: 1) Missing completely at random (MCAR) refers to the case that the probability of entry being missing is independent of the data, i.e., p(m [x] p(m). 2) In missing at random (MAR), the probability of missingness depends only on the observed values, i.e., p(m]|x p(m] xbbs) 3) All other cases are classified missing not random (MNAR), where the probability of missingness might also depend on other missing entries.
We follow the methods proposed in Zhao al. (2023) to implement MAR and MNAR.
"For MAR, we first sample subset of features (columns in X) that will not contain missing values, and then we use a logistic model with these non-missing columns input determine the missing values of the remaining columns, and we employ line search of the bias term to get the desired proportion of missing values."
For MNAR, we use the first approach proposed in [4], "Using a logistic model with the input masked by MCAR". Specifically, similar to the MAR setting, we first divide the data columns into two groups, with one group serving as input for the logistic model, outputting the missing probability for the other of columns. The difference is that after determining the missing probability of the second set, apply MCAR to the input columns (the first set). Hence, missing values from the second set will depend on the masked values of the first set.
The code for generating masks according to the three missing mechanisms is also provided at ht buttps://anony
4open. www.ecceeeeeeeeeec.eecceeeeeeeee wtw.eeeneee///IIIiEIEEEEe-
D.4 IMPLEMENTATION AND HYPERPARAMETER
We fixed set of hyperparameters, which will save significant efforts in hyperparamettr--uniin when applying DIFFPUTER to more datasets. For the diffusion model, we set the maximum time T 80, the noise level a (t) f, which is linear The score/denoising neural network c(x,, f) is implemented as a 5-layer MLP with hidden dimension 1024. t is transformed to sinusoidal timestep embeddings and then added to Xt, which is subsequently passed to the denoising function. When using the learned diffusion model for imputation, we set the number of discrete steps M 50 and the number of sampling times per data sample N 10. DIFFPUTER is implemented with Pytorch, and optimized using Adam (Kingma & Ba, 2015) optimizer with learning rate of 1 x 10
Architecture of denoising neural network. We use the same architecture of denoising neural network in recent tabular diffusion models for tabular data synthesis (Kotelnikov al., 2023: Zhang et al., 2024). The denoising MLP takes the current time step and the feature vector X, Ixd input. First, X; fed into a linear projection layer that converts the vector dimension to be dhidden = 1024:
hg = FCin (xx) E TR dxdhiidoon
(21)
where ho is the transformed vector, and dridden the output dimension of the input layer.
17

---

Published as a conference paper at ICLR 2025
Then, following the practice in TabDDPM (Kotelnikov et al., 2023), the sinusoidal timestep embed- dings temb R xdiiidee is added to hg to obtain the input vector haidden-
hnn
hg + temb-
(22)
The hidden layers three fully connected layers of the size diddeen 2+dhidden -2sdhiddeen with SiLU activation functions:
hidden
h1 SiLU(FC1 E R (2ddhidde))
h = SiLU(FC((h:) E x2ndhadden
(23)
hg = SiLU(FC;(hg) E R Xdhidden).
The estimated score is obtained via the last linear layer:
(x(X;,) hout FCout (h3) E Rxdd
(24)
Finally, co(x,, t) is applied to Eq. 3 for model training.
D.5 IMPLEMENTATIONS OF BASELINES
We implement most of the baseline methods according to the publicly available codebases:
Remasker (Du et al., 2024): http:///ttapagaaa...o.ooo///////
TDM (Zhao al., 2023): https: httppp/s//ggcchb.....///eeg/g///
MOT
(Muzellec
al., 2020): httpes//ggittuub.cc////bbbbbbbb
Mi
MissingDataOT
GRAPE (You et al., httg://ttppss//ttt.u....o///////
IGRM (Zhong et al., 2023): https:: Com.I-AAIA//AAIAIIGG
TabCSDI (Zheng
reeettrrseercchTaabbe Charoenphakdee, 2022): https://github.com
MCFlow (Richardson et al., http://://ttoo.ooo////ooooooooa/
MCFlow.
For Hyperlmputer (Jarrett et al., 2022), MissForest (Stekhoven & Bühlmann, 2012), MICE (Van Buuren & Karin, 2011), SoftImpute (Hastie et al., 2015), EM (Gareía- Laencina et al., 2010), GAIN (Yoon et al., 2018), MIRACLE (Kyono et al., 2021), and MIWAE (Mattei & Frellsen, 2019), we use the implementations at: https: / thub. conn/vanderschaee htwp:hhhaaaaaaaahaaraiaeeeee htpph/mpperr......e//e
MissDiff (Ouyang al., 2023) does not provide its official implementations. Therefore, we obtain its results based on our own implementation.
The codes for all the methods are available at https://anooymont 4open. science/r/ ICLR-DiffPuter
D.6 HYPERPARAMETER SETTINGS OF BASELINES
Most of the deep learning baselines recommend the use of one set of hyperparameters for all datasets. For these methods, directly follow their guidelines and use the default hyperparameters:
the recommended hyperparameters provided in Ap- pendix A.2 the original paper (Du et al., 2024). This set of hyperparameters is searched by tuning the Letter dataset and is deployed for all the datasets in the original paper; hence, we follow this setting.
ReMasker (Du et al., 2024):
HyperImpute (Jarrett et al., 2022): since HyperImpute works by searching over the space of classifees/regressorss and their hyperparameters, it does not have hyperparameters itself except parameters related to the AutoML search budget. We adopt the default budget parameters of Hyperlmpute's official implementation for all datasets. The default budget
18

---

Published as a conference paper at ICLR 2025
parameters and AutoML search space are provided and Table 5 in the original paper (Jarrett et al., 2022).
MOT (Muzellec al., 2020) and TDM (Zhao al., 2023): There is a main hyperparameter representing the number of subset pairs sampled from the dataset for computing optimal transport loss. Sinkhorn algorithm and TDM are controlled by hyperparameter jitr. While the default value is 3000, we set it as 12000 for all datasets to ensure the algorithm converges sufficiently. For the round-robin version of the algorithm, the number of sampled pairs controlled by max_ iter and rr r_iter; we adopt the default value 15, which is enough for the algorithm to converge. For the remaining hyperparameters related to network architectures, we use the default ones for all datasets.
KNN: we follow the common practice of selecting the number of nearest neighbors as vn. where is the number of samples the dataset.
GRAPE (You et al., 2020) and IGRM (Zhong et al., 2023): we adopt the recommended set of hyperparameters used in the original paper for all datasets. For detailed explanation of the meaning of the parameters, please see the github repos: GRAPE and IGRM
Mifs since the original implementation is not available and is based on the diffusion model, for a fair comparison, we simply use the same set of hyperparameters with our DiffPuter.
TabCSDI (Zheng & Charoenphakdee, 2022): we follow the guide for selecting hyperparam- eters in the original paper (Appendix B in [3]). Specifically, we use large version of the TabCSDI model with a number of layers set to 4 (see more detailed hyperparameters about the large TabCSDI model For batch size,
take the official choice of batch size (8) for the breast dataset 700 samples) as a base and scale the batch size accordingly with the sample size of our datasets: since most of the datasets used have the number of samples between 20000 to 40000, we scale the batch size to 256 and use it for all datasets.
MCFlow (Richardson et al., 2020): we adopt the recommended hyperparameters provided in the official implementation for all "atasets
For the remaining classical machine learning methods, including EM, GAIN, MICE, Miracle, MissForest, and Softimpute where hyperparameters might be important. Since implementations from the 'hyperimpute' package, we tune the hyperparameters within the hyperparameter space provided in the package To be specific, we set the maximum budge as 50, then we sample 50 different hyperparameter combinations according to the hyperparameter space. Finally, we report the optimal performance over the 50 trials.
the
E ADDITIONAL RESULTS
E. MCAR: OUT-OF-SAMPLE IMPUTATION
We present the out-of-sample imputation performance comparison under MCAR setting in Table 6.
E.2 MISSING AT RANDOM (MAR)
We present the performance comparison under MAR setting in Table and Table 8.
E.3 MISSING NOT RANDOM (MNAR)
We present the performance comparison under MNAR setting in Table 9 and Table 10.
httpr////..ttiicccooaaaooacoaoea "
http:p///gaaatacccccccaaaaaa///a onehot_analog.yyan
n.p.y
19

---

Published as a conference paper at ICLR 2025
Table 6: MCAR, Out-of-sample imputation performance MAE and RMSE metries (using base 10-2 2 for better presentation). DIFFPUTER outperforms the most competitive baseline methods by 13.37% MAE, and by 4.43% RMSE.
Method
California
Letter
Gesture
Adult
Default
Shoppers
Average
MCFlow
60.74+5.09
66.35eo.82
82.72+2-49
77.54+0.98
64.12+2.10
73.90+3.34
70.90
IGRM
138.22+4.2
121.5921.18
149.8813.92
149.64 299.6420.18
OOM
132.28:0.01
GRAPE
36,54+0.12
42. 10.11+
92.50+o.07
50.25+0.68
$7.5L+0.00
51.91+1.20
56.69
MOT
44.20=0.81
46. 16.7540.18
36.74+0.01
51.17+0.52
31.42+0.0
41.55+0.42
41.97
Remasker
33.97.0.10
34.0320.35
35.2924.11
49.2710.53
37.0213.56
41.59
38.53
DIFEPUTER
+3.47+0.25
30.60 30.69:0..e.
28.53
49.14+0.e2
24.63-+0.24
34.46-0044
33.49
Improv.
1.47%
9.81%
19.16%
0.26%
21.61%
17.06%
13.09%
MCFlow
84.57.7.87
87.7320.58
119.7321.10
220.6820.66
101.3322.62
115.75:2.82
104.97
IGRM
170.00+4.72
151.16=2.05
172.25++2.09
179.17==0.88
0OM
174.42
RMSE
GRAPE
56.0520.05
58. 18.1120.21
120.31
99.19#0.50
91.9420.8
92.37
86.33
MOT
71.0221.19
64.69
71.6420.81
94.6721.50
71.35:2 .08
82.6312.19
76
Remasker
52.96.20.00
48.6022..4
68.7813..0
90.38.1141
76.94
80.0010.86
69.61
DIPEPUTER
54.1820.25
47.7 20.59
60.90
91.31
70.50 83.
74.01:0.40
66.41
Improv.
1.77%
11.46%
1.19%
7.49%
4.60%
Table 7: MAR, In-sample imputation MAE.
Method
California
Letter
Gesture
Magic
Bean
Adult
Default
Shoppers
News
Traditional iterative
EM
39.17
58.05
35.43
48.13
15.41
62.54
33.78
43.01
39.52
MICE
58.56
82.32
65.14
72.81
22.17
98.44
63.92
75.70
61.12
MIRACLE
47.13
71.10
62.56
53.87
29.06
73.06
43.72
56.89
38.21
SoftImpute
59.47
67.01
53.95
58.55
31.55
70.39
42.66
61.91
64.85
MissForest
45.90
62.51
37.69
47.70
29.21
73.25
38.09
41.97
41.18
Dist. match MOT
44.52
50.86
44.38
46.95
30.41
58.00
36.48
40.33
49.92
TDM
39.27
46.86
42.69
45.31
29.11
59.59
37.93
45.00
63.25
GNN
GRAPE
35.14
42.45
38.088
39.87
18.81
60.73
48. 48.43
53.73 73
50.85
IGRM
35.23
41.31
36.45
40.09
18.97
61.35
59.86
Generative models
MIWAE
76.27
72.30
57.30
77.31
62.64
57.99
51.10
46.74
63.26
GAIN
69.68
67.50
73.70
66.11
53.46
92.81
60.04
46.69
57.49
MCFlow
60.00
68.23
65.59
55.27
30.40
84.36
67.87
72.99
56.36
TabCSDI
83.17
77.42
54. 64
76.93
75.34
57.11
46.
58.55 58.53
66.41
MissDiff
119.49
56.69
229. 19
55.82
45.62
55.01
50.62
48. 19
54.41
Other
KNN
49.25
49.41
45.52
49.70
30.81
50.83
35.15
46.81
48.53
ReMasker
34.26
33.93
36.52
50.21
16.85
52.60
40.58
38.24
30.37
Hyperlmpute
36.22
39.74
36.44
42.23
16.93
49.91
29.66
36.93
26.85
DIFPUTER
32.66
32.22
30.39
29.66
15.92
49.51
26.61
35.19
28.22
20

---

Published as a conference paper at ICLR 2025
Table 8: MAR, In-sample imputation RMSE.
Method
California
Letter
Gesture
Magic
Bean
Adult
Default
Shoppers
News
Traditional iterative
EM
60.28
77.77
67.68
74.17
39.37
88. 71
79.48
65.76
82.32
MICE
85.13
106.80
92.85
102.74
49.96
126.50
109.18
89.10
105.35
MIRACLE
105.54
96.64
104.10
105.48
100.19
102.18
107.70
105.52
105.53
SoftImpute
83.23
88.33
99.13
87.11
56.32
95.44
100.26
94.04
90.61
MissForest
72.08
84.20
75.70
74.75
51.24
106.97
83.22
70.15
93.03
Dist. match
MOT
74.86
69.71
86.57
76.53
57.92
85.96
88.74
81.01
92.78 02.78
TDM
67.88
64.66
80.44
80.77
62.96
86.98
91.34
95.22
95.02
GNN
GRAPE
55.73
57.71
72.88
06.88
42.33
86.67
96.19
80.85
96.50
IGRM
56.13
50.44
67.45
41.61
85.55
97.30
0.00
Generative models
MIWAE
105.54
96.64
104.10
105.48
100.19
102.18 18
107.70
105.52
105.53
GAIN
98.01
88.89
112.56
93.54
79.02
123.26
103.13
90.26
117.09
MCFlow
82.84
90.07
109.49
89.02
52.83
112.48
117.18
86.55
112.21
TabCSDI
111.64
100.27
101.72
101.65
101.15
87.37
101.74
100.07
104.73
Difs
256.09
90.38
98.17
121.62
676.10
1464.76
Other
KNN
80.05
68.50
92.64
75.52
62.92
90.24
85.50
74.04
80.12
ReMasker
56.08
47.97
70.55
77.09
36.87
78.08
78.43
58.94
93.74
Hyperlmpute
61.94
57.50
71.98
69.71
39.01
88. 73
79.84
58.96
71.46
DIFEPUTER
57.12
48.12
68.59
67.29
35.19
79.42
77.52
60.11
65.53
Table 9: MNAR, In-sample imputation MAE.
Method
California
Letter
Gesture
Magic
Bean
Adult
Default
Shoppers
News
Traditional iterative
EM
38.71
56.10
37.49
50.75
10.48
56.17
30.06
42.57
38.93
57.63
81.26
66.33
77.06
17.12
98.34
59.55
76.30
62.27
MIRACLE
42.42
71.83
45.59
15.27
59.32
37.41
78.56
SoftImpute
64.13
65.58
55.31
60.76
28.86
61.11
43.39
60.12
63.48
MissForest
43.11
58.61
39.74
51.15
25.37
57.40
34.49
42.58
39.57
Dist. match
MOT
42.48
44.78
48.05
49.39
35.84
42.20
48.14
TDM
35.93
45.22
43.43
46.24
20.91
50.29
41.07
43.48
53.17
GNN
GRAPE
35.82
40.46
40.68
13.24
57.97
44.70
55.06
49.36
IGRM
34.89
40.18
40.82
40.92
13.57
54.52
0.00
58.14
0.00
Generative models
MIWAE
75.91
83.18
60.93
79.34
62.09
59.03
56.06
49.34
64.40
79.60
67.04
83.73
64.75
39.68
123.87
85.44
55.89
57.33
MCFlow
59.23
65.53
67.03
60.37
26.30
81.16
67.03
77.87
54.96
TabCSDI
72.26
77.71
58.16
75.46
75.69
62.09
58.88
60.29
65.88
MissDiff
52.37
49.11
325.12
113. 14
38.52
98.34
63.11
168.17
75.23
Other
KNN
54.23
53.21
48.79
49.26
24.00
62.99
46.27
47.83
ReMasker
33.22
33.65
38.98
75.05
12.94
47.66
39.44
39.89
29.24
Hyperlmpute
33.96
40.20
38.30
42.88
11.96
50.19
28.36
39.80
31.74
DIFFPUTER
34.08
33.13
31.54
40.82
11.42
48.59
26.94
37.25
28.51
21

---

Published as a conference paper at ICLR 2025
Table 10: MNAR, In-sample imputation RMSE.
Method
California
Letter
Gesture
Magic
Bean
Adult
Default
Shoppers
News
Traditional iterative
EM
59.75
75.78
74.88
76.01
28.90
92.30
83.26
66.16
71.03
MICE
84.36
105.91
98.00
106.59
41.91
132.01
112.55
91.16
96.87
MIRACLE
73.98
109.85
125.48
74.93
50.36
106.10
150.95
82.04
100.59
SoftImpute
102.74
87.02
103.04
90.21
48.27
99.56
95.59
83.84
MissForest
65.80
80.53
81.17
76.66
42.72
99.58
91.87
70.15
85.34
match
MOT
88.63
91.16
93.92
83.08
86.44
TDM
59.81
63.17
86.36
78.16
47.59
91.81
96.36
88.88
94.51
GNN
GRAPE
55.78
56.43
78.33
63.91
32.66
94.56
99.99
81.17
84.97
IGRM
55.34
55.32
78.
64.30
32.57
93.64
97.58
Generative models
MIWAE
102.77
107.54
111.03
107.10
96.50
102.70
113.33
107.99
114.09
GAIN
106.38
88.99
128.97
89.24
58.92
184.39
116.56
93.92
145.92
MCFlow
82.16
86.87
114.77
91 91.89
47.98
117.01
1233.68
88.87
105.53
TabCSDI
93.59
107.65
100.20
101.99
102.92
107.07
MissDiff
137.09
68,98
6116.33
227.21
61.13
229.50
893.48
449.46
829.78
Other
KNN
83.15
73.01
101.66
75.44
53.33
103.03
86.35
85.52
98.56
ReMasker
54.34
48.02
75.66
102.23
32.15
85.60
82.65
58.61
80.17
Hyperlmpute
58.93
57.84
75.84
69.87
31.24
89.04
86.96
61.48
76.42
DIFFPUTER
57.48
48.54
69.54
67.66
32.59
85.26
78.82
59.83
66.74
22


Reviews:
[{"review_id": "G61C2r3UkE", "reviewer": "Reviewer_y7yQ", "summary": "The paper addresses the imputation of data missing completely at random (MCAR) in tabular data, handling both continuous and categorical variables. The authors propose an EM procedure with a conditional diffusion model for imputation, featuring a novel adaptation of the annealing process for better conditioning on observed values. The paper demonstrates strong results across multiple datasets in comparison with leading methods.", "questions": "In biology, missing values are often represented as 0 (or another \"limit of detection\" (LOD) value), making it difficult to distinguish between actual LOD values and data missing at random (which can comprise 30% of data in cases like proteomics and single-cell analysis). Do you have any ideas about how this problem could be addressed? Note that the fraction of missing values might be known and could potentially be conditioned on.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written with a clear, thorough, and concise introduction that effectively summarizes key points from previous works\n* The authors specifically address the challenges of the problem and provide clever solution to mitigate them\n* The paper's main novelty is supported by theoretical proof\n* The evaluations are comprehensive, with thorough and convincing ablation studies", "weaknesses": "* A major concern regarding evaluations - while the paper claims to use a single hyperparameter setting throughout, it's unclear how hyperparameters for other methods were selected and their sensitivity to these HP. For me, this concern significantly impacts the overall assessment of the paper.\n\n* While the results are impressive, their importance is not clear. A more convincing evaluation would include the effect on downstream tasks, given imputation is only a first step in most pipelines. \n\n* Another point regarding evaluations, is the sole focus on data missing completely at random. While the MER assumption is important, it is the MNER which is a primary focus in many imputation methods.\n\n* The 0/1 continuous encoding of categorical data is unusual, given that binary data is a known challenge for diffusion models (for example in fields like graph generation). Also, the use of mean is inherently problematic due to common multi-modality in the data\n\n* The novelty of the method compared to other approaches is not clearly articulated in the related works section\n\n### Smaller Issues\n* Given the method's novelty isn't specific to tabular data, the related work should include other imputation methods (e.g., image inpainting)\n* A simulation study with multiple modes would be valuable, particularly as diffusion-based models should excel in such scenarios\n* Despite highlighting the importance of initialization in the EM procedure, the paper doesn't address this point. (Particularly relevant given the naive initial imputation approach)\n* It would be interesting to analyze the relationship between delta_t size and ML solution approximation.\n* Figure 4 lacks clarity", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "fbRHHecko3", "reviewer": "Reviewer_8jb1", "summary": "The missing value imputation is a very important problem in both machine learning and statistics. Although deep generative imputation methods have shown promise, there is still substantial room for improvement, even for matching the performance of some traditional machine learning approaches. This paper introduces DIFFPUTER, a tailored diffusion model combined with the Expectation-Maximization (EM) algorithm for missing data imputation, and shows its promising performance on a variety of datasets,", "questions": "This paper is well-written and easy to understand. This work combines EM with the diffusion model to improve the potential inconsistency caused by missing values in the training process of diffusion models. Since EM is used with K iterations and N number of samples in the E-step, it would be beneficial to also compare the computational complexity (time complexity) of the proposed method with other diffusion-type methods, either in the discussion of the number of operations or comparing the running time in some of the numerical experiments. This would offer more insights into the proposed methods' performance from different perspectives.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Theoretical analysis: DIFFPUTER’s training step corresponds to the maximum likelihood estimation of data density (M-step), and its sampling step represents the Expected A Posteriori estimation of missing values (E-step). \n2. Extensive experiments that demonstrate the good performance, as compared with existing baselines, of the proposed method across various datasets.", "weaknesses": "the computational complexity is not explicitely discussed or compared on the numerical experiments, see details below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "efZp13tSqy", "reviewer": "Reviewer_TxpG", "summary": "This paper introduces an adaptation of the EM algorithm for missing data imputation, leveraging advanced diffusion-based models to perform precise density estimation in the M-step and provide robust imputations in the E-step, inspired by the RePaint algorithm. The authors reference theoretical analyses from prior work to support the use of diffusion models for density estimation and prove a theorem demonstrating that E-step samples can be drawn from the true conditional distribution. Extensive empirical evaluations highlight the proposed method’s robustness and superiority over various baseline approaches, many of which do not incorporate the EM algorithm.", "questions": "### Motivation appears to overlook recent work\n\n- The two main issues presented as motivation for this work are unclear. The paper claims that generative imputation methods (i) require joint estimation of observed and missing data distributions and (ii) struggle with conditional inference. I find both statements questionable. Numerous studies adapt deep generative models to estimate only the observed data distribution [1-5], which could serve in the M-step of an EM algorithm. Some of these are even referenced in this paper. Moreover, all of these methods allow for straightforward Monte Carlo estimation of $\\mathbb{E}[p(\\mathbf{x}_m | \\mathbf{x}_o)]$ for the E-step, similar to the proposed diffusion-based model. For instance, a more robust importance-weighted estimator is proposed in [4] (see Eq. (12)).\n\n- This brings me to a second point: if multiple DGMs could, indeed, replace diffusion models within the EM framework, how is diffusion specifically justified for tabular data? This approach might be advantageous for high-dimensional data, where diffusion models effectively approximate $p(\\mathbf{x})$ and avoid lossy compression (as in VAEs). However, given the lower-dimensional datasets studied here, it remains unclear why a VAE-based approach, for example, wouldn’t perform as well as diffusion.\n\n### Experimental section lacks fair comparison and clarity\n\n- Based on my earlier points, I would expect an ablation study comparing different DGMs within the EM algorithm. The baselines in the current experiments appear to rely on simple placeholder values for missing data (e.g., zero or mean imputation), effectively completing only one M-step. This is likely to produce suboptimal results, so a performance gap seems unsurprising.\n\n- The assertion \"Traditional machine learning methods are still powerful imputers\" would benefit from supporting references. I am skeptical, as optimal validation could be harder to achieve in probabilistic settings.\n\n- The claim that generative methods excel on continuous data requires clarification. Here, the diffusion model seems to assume Gaussianity across all dimensions, using $argmax$ as a proxy to obtain discrete outputs, which is not the optimal to model heterogeneous data [2, 3, 5]. \n\n- The statement \"imputation methods are specifically designed for in-sample imputation and cannot be applied to the out-of-sample setting\" also needs elaboration. As mentioned, many DGMs designed for missing data can perform out-of-sample imputation.\n\n- In Figure 2, MissDiff appears to fail or encounter out-of-memory issues. This is surprising, as MissDiff’s architecture is similar to the diffusion network used here.\n\n### Discussion of limitations is lacking\n\n- The main text does not discuss limitations, particularly the high computational cost. A brief note is found in the Appendix, but this isn’t referenced within the primary text. DiffPuter’s approach requires retraining the diffusion model $k$ times, so application to high-dimensional data (e.g., images) would be computationally intense relative to alternatives. I am also curious if the M-step converges faster with higher values of $k$, as this could enhance efficiency.\n\n### Other minor questions\n\n- Figure 6: Why does error decrease as observed data ratio drops? I found the final paragraph of Section 5 somewhat unclear; further clarification here would be helpful.\n\n### Typos\n- Line 515: Change *\", Reducing\"* to *\", reducing\"*.\n\n\n### References\n\n[1] Ma, Chao, et al. \"EDDI: Efficient Dynamic Discovery of High-Value Information with Partial VAE.\" International Conference on Machine Learning. PMLR, 2019.\n\n[2] Ma, Chao, et al. \"VAEM: a deep generative model for heterogeneous mixed type data.\" Advances in Neural Information Processing Systems 33 (2020): 11237-11247.\n\n[3] Peis, Ignacio, Chao Ma, and José Miguel Hernández-Lobato. \"Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo.\" Advances in Neural Information Processing Systems 35 (2022): 35839-35851.\n\n[4] Mattei, Pierre-Alexandre, and Jes Frellsen. \"MIWAE: Deep generative modelling and imputation of incomplete data sets.\" International conference on machine learning. PMLR, 2019.\n\n[5] Nazabal, Alfredo, et al. \"Handling incomplete heterogeneous data using VAEs.\" Pattern Recognition 107 (2020): 107501.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Robust imputation method based on EM. \n- Well written and structured. \n- The method is theoretically grounded. \n- The empirical analysis is extensive.", "weaknesses": "- Motivation appears to overlook recent work.\n- Experimental section lacks fair comparison and clarity.\n- Discussion of limitations is lacking.\n- Given these weaknesses, the contribution is not strongly justified.\n\n\n------- Post rebuttal update -------\n\nAll the weaknesses were thoroughly addressed in the rebuttal provided by the authors. I appreciate their efforts and the detailed responses, which resolved all my concerns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "9cE8wFTCZH", "reviewer": "Reviewer_rwwt", "summary": "The authors combine diffusion processes and Expectation-Maximization (EM) algorithms to propose a novel way to impute missing data when both training and tests sets contain missing data. The proposed solution is shown to target the correct conditional distribution (distribution of missing data conditional on observed ones). Imputation values are computed by taken the expectation with respect to this distribution, which is approximated by the sample mean. Experiments on ten real-world data sets show the benefit of the proposed method, compared to various state-of-the-art imputation algorithms (machine and deep learning algorithms).", "questions": "- l.197: could you specify the choice of $\\sigma(t)$? \n- l.225-227: the paragraph does not correspond to the equation: the negative log likelihood is upper bounded by the loss plus a constant, which does not imply that optimizing the first leads to optimizing the second.\n- Section 5.1, how does the method behave when different masks are present in the training and test set? Does it degrade the performances? \n- Section 5.1, how were the hyperparameter chosen for the different baselines? Are these baselines comparable (in terms of number of parameters for example) with the proposed method? Could you add such a discussion in the Appendix? Could you also describe in details the missing data mechanisms used for the different settings (MAR and MNAR encapsulate a lot of data generating processes)?\n- l.366-367, Can you explain the good performances of the proposed method compared to MissDiff and TabCSDI?\n\n\nl.399 : \"Imputaing\"\nl. 468 : \"Gestrue\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors propose a new method to impute missing data for continuous and discrete inputs. The proposed method appears to be new, with excellent performances. An extensive literature review has been done to present and explain the previous approaches to deal with missing values via imputation. The method is clearly explained, the paper well-written, and the experiments show the benefit of the proposed method.", "weaknesses": "I only have three remarks: \n- RMSE and MAE are measures that encourage the imputation to target the (conditional) mean or median. In both cases, the target is not a distribution but a single quantity. Recent works (https://arxiv.org/abs/2002.03860) have shown that such measures do not properly evaluate the correctness of an imputation method. Imputation score (https://arxiv.org/pdf/2106.03742) can be used instead to assess the quality of imputations. As the proposed method generates a distribution and not a single point estimate, it is likely that its performance will be higher with respect to this metric, showing that it is able to recover the underlying distribution of the data. Presenting imputation scores in the tables would definitely improve the strength of the paper, in my opinion.\n- The computational performances of DiffPuter should be discussed in the main text. Table 4 is interesting, as it shows that the training time is larger, but not too important. However, the two considered data sets have few features. It would be appealing to consider larger data sets with (i) more observations and/or (ii) more variables to see how the predictive performances and the training time behave. \n- I have trouble understanding the proof of Theorem 1. Notations are confused to me. Adding a table of notations, with exact definitions at the beginning of the Appendix would help. Besides, many approximations are done in the proof : l.730, 731, 750, 753. This results in the theorem being imprecise. For example, nothing is assumed about the quality of the neural network $\\varepsilon_{\\theta}$. What type of convergence is required for Theorem 1 to be valid? Similarly, in Theorem 1, $\\sigma(T)$ is not assumed to be large, whereas it is required in the proof. Please clarify the different assumptions and the proof.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)