DiffPuter: Empowering Diffusion Models for Missing Data Imputation

How the missing data mechanism is implemented?

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)

We follow the methods proposed in TDM[4] to implement MAR and MNAR.

• "For MAR, we first sample a subset of features (columns in $X$) that will not contain missing values and then we use a logistic model with these non-missing columns as input to 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 set of columns. The difference is that after determining the missing probability of the second set, we 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.

We have also updated this content in Section D.3 of the revised paper's Appendix.

References:

[1] Tianyu Du, Luca Melis, and Ting Wang. Remasker: Imputing tabular data with masked autoencoding. In International Conference on Learning Representations, 2024.

[2] Daniel Jarrett, Bogdan C Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperimpute: Generalized iterative imputation with automatic model selection. In International Conference on Machine Learning, pp. 9916–9937. PMLR, 2022.

[3] Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular data. In NeurIPS 2022 First Table Representation Workshop, 2022.

[4] Zhao, H., Sun, K., Dezfouli, A., & Bonilla, E. V. (2023, July). Transformed distribution matching for missing value imputation. In International Conference on Machine Learning (pp. 42159-42186). PMLR.

Q3: The selection of hyperparameters for different baselines.

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?

Most of the deep learning baselines recommend the use of one set of hyperparameters for all datasets. For these methods, we directly follow their guidelines and use the default hyperparameters. For Remasker, GRAPE, and IGRM, where the model widths can be enlarged or reduced, we tried to align their model sizes but observed little impact on their performance. Below is the detailed introduction of how the hyperparameters of each baseline method is selected (and we've added a discussion part in Appendix D.6):

• ReMasker: we use the recommended hyperparameters provided in Appendix A.2 in the original paper [1]. This set of hyperparameters is searched by tuning on the Letter dataset and is deployed for all the datasets in the original paper; hence, we follow this setting.
• HyperImpute: since HyperImpute works by searching over the space of classifiers/regressors and their hyperparameters, it does not have hyperparameters itself except parameters related to the AutoML search budget. We adopt the default budget parameters of HyperImpute's official implementation for all datasets. The default budget parameters and AutoML search space are provided in https://github.com/vanderschaarlab/hyperimpute/blob/main/src/hyperimpute/plugins/imputers/plugin_hyperimpute.py and Table 5 in the original paper [2].
• MOT and TDM: 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 n_iter. While the default value is 3000 (https://github.com/BorisMuzellec/MissingDataOT/blob/6417eacfde1c1052a63568350dfec2d0373ac056/experiment.py#L42), 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 is controlled by max_iter and rr_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 $\sqrt{n}$, where $n$ is the number of samples in the dataset.
• GRAPE and IGRM: we adopt the recommended set of hyperparameters used in the original paper for all datasets. For a detailed explanation of the meaning of the parameters, please see https://github.com/maxiaoba/GRAPE/blob/0ea0c59272a977d0184a8fd95178f68211455ef5/train_mdi.py#L18 for GRAPE and https://github.com/G-AILab/IGRM/blob/5cfc955daa5d0f4bbdcbad1552cfd7493dfd5fd0/main.py#L17 for IGRM.
• MissDiff: since the original implementation is not available, and it is based on diffusion model, for fair comparison, we simply use the same set of hyperparameters with our DiffPuter.
• TabCSDI: We follow the guide for selecting hyperparameters in the original paper (Appendix B in [3]). Specifically, we use a large version of the TabCSDI model with a number of layers set to 4 (see more detailed hyperparameters about the large TabCSDI model at https://github.com/pfnet-research/TabCSDI/blob/main/config/census_onehot_analog.yaml). For batch size, we 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 we used have the number of samples between 20000 to 40000, we scale the batch size to 256 and use it for all datasets.
• MCFlow: we adopt the recommended hyperparameters provided in the official implementation for all datasets (https://github.com/trevor-richardson/MCFlow/blob/70fe137db79255bfbec07c5605ccd3fe0c52c789/main.py#L67).

For the remaining classical machine learning methods, including EM, GAIN， MICE，Miracle，MissForest, and Softimpute where hyperparameters might be important. Since we use the implementations from the 'hyperimpute' package, we tune the hyperparameters within the hyperparameter space provided in the package (e.g., https://github.com/vanderschaarlab/hyperimpute/blob/e9506c7b3a1f5089f00797534e0460fd28f9730c/src/hyperimpute/plugins/imputers/plugin_missforest.py#L73). 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.

l.197: could you specify the choice of $\sigma(t)$?

In this paper, we set $\sigma(t) = t$ that is linear w.r.t. the time. We apologize for this omission and have added a supplementary explanation below Equation (2). This is a widely adopted setting for diffusion models in order to achieve better sampling efficiency (faster sampling speed)[1].

l.225-227: the paragraph does not correspond to the equation: the negative log-likelihood is upper bound by the loss plus a constant, which does not imply that optimizing the first leads to optimizing the second.

Thank you for your correction. You are right; Remark 2 cannot state that optimizing the first necessarily optimizes the second. Instead, it is approximating the maximum likelihood estimation. The necessary condition is that the bound should be sufficiently tight (for example, become an equality). In fact, Theorem 2 of Song et al., 2021a [2] has extended Remark 2, showing that the equality can be achieved when score function $s_{\theta}$ perfectly aligns with the score function $\nabla_{\mathbf{x}}\log q_t(\mathbf{x})$ of a time-dependent reverse-time diffusion process with boundary distribution $q_T = \pi$ and $\mathbf{x}(0) \sim q_0$. i.e. $s_{\theta}(\mathbf{x}, t)\equiv \nabla_{\mathbf{x}}\log q_t(\mathbf{x})$.

Theorem 1 also claims this condition is hard to satisfy since the score-based model $s_{\theta}(\mathbf{x}, t)$ will not exactly match $\nabla_{\mathbf{x}}\log q_t(\mathbf{x})$ everywhere (due to the capacity limitation of neural networks), therefore it is only an approximation. However, the empirical results of [2] did demonstrate that the training is actually able to improve the log-likelihood of data across multiple datasets and model architectures.

We have revised the relevant paragraphs to make our claims more accurate. Thank you again for your correction.

References:

[1] Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, 26565-26577.

[2] Song, Yang, et al. "Maximum likelihood training of score-based diffusion models." Advances in neural information processing systems 34 (2021): 1415-1428.

W3: Understanding of the proof of Theorem 1

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. 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.

Thank you for your criticism and suggestions. In the updated paper, we have added a new section at the beginning of the Appendix (Appendix A) to introduce all symbols and notations used in this paper.

The proof of Theorem 1 in the Appendix indeed needs the help of some mild assumptions for approximation. We list these assumptions/approximation as follows:

• line 730: Small reverse process time intervals: $\Delta t \rightarrow 0$, such that the difference between $x_t^{\rm mis}$ and $x_{t-\Delta t}^{\rm mis}$ is negligible.
• line 731 should be '=' rather than '$\approx$' since Monte Carlo estimation is unbiased.
• line 750: We require a large maximum timestep $\sigma(T) \gg 1$, such that $p(\mathbf{x}_T | \mathbf{x}_0)\approx p(\mathbf{x}_T)$

The necessary assumptions have been added to the description of Theorem 1 in the revised paper. Thanks again for your criticism.

Note: Line 751, Theorem 1 does not need to assume the quality of the neural network $\epsilon_{\theta}$, because Theorem 1/Lemma 1 fundamentally explains that for any denoising network $\epsilon_{\theta}$, our method can sample/generate from its induced time-dependent distribution at arbitrary time step $t$, $p_{\theta}(\mathbf{x}_t | \hat{\mathbf{x}}^{obs})$.

The statement '''the score function $\nabla_{\mathbf{x}_t}\log p (\mathbf{x}_t)$

is replaced with $\varepsilon_{\theta}(\mathbf{x}_t, t)$''',

does not indicate approximation, but rather indicates that the distribution obtained by the reverse process is induced by the denoising neural network. Specifically, when $\varepsilon_{\theta}(\mathbf{x}_t, t)$ perfectly fits the score function $\nabla{\mathbf{x}_t}\log p (\mathbf{x}_t)$, the induced distribution is the ground-truth data distribution.

W2: The computational performance should be discussed in the main text.

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.

Thank you for the suggestion. We've moved this part to Section 5.3 in the main text. In addition, we consider a novel dataset with more rows and columns, the Covertype dataset. The original dataset consists of 581,012 instances and 54 features, and we subsample 200,000 records such that we can complete the experiments in the rebuttal phase. The following table compares the training time and our DiffPuter's performance advantage over representative baselines (we also present the results on Adult and California for reference):

On such a large dataset, graph-based methods like GRAPE and IGRM both face OOM (Out Of Memory) problems, while other methods that can perform mini-batch training can run normally.

From the perspective of training speed, the training time required by different methods generally shows a linear relationship with the number of rows and columns in the dataset. The training time required by our DiffPuter is basically at the same level as HyperImpute and Remasker. From the perspective of imputation performance, our DiffPuter shows even greater advantages compared to other methods, which may be attributed to the huge capacity and exceptional ability of Diffusion models in modeling large-scale and high-dimensional data distributions.

We will include the CoverType dataset as a benchmark dataset in our regular experiments, and add all its results to the updated paper once they are obtained.

We thank the reviewer for the time spent reviewing our paper and for providing constructive questions and suggestions. The following is our detailed response to every question.

W1: A novel metric, the imputation scores

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.

We thank the reviewer for recommending the novel metric, and we agree that comparing the performance of different methods under this metric will significantly improve our paper. Therefore, we have conducted additional experiments for this metric, and will add these results in the revised paper. Also, this metric has been integrated into our codebase.

When using this metric, we find that the computation of the imputation score seems rather inefficient and coul take a extremely long time for the entire dataset. Therefore, we randomly sample 500 rows/samples from the in-sample imputation results. The hyperparameters are set as follows:

• num.trees.per.proj to be 5
• The minimum node size in each tree is 10 (the default for a probability RF).
• We chose the number of projections (num.proj) adaptively to the dimension $d$ of the data set: for $d \le 6$ we used 50, for $7 ≤ d ≤ 14$ we used $100$ and for $d ≥ 15$ we used $200$.

We present the imputation scores of different methods in the following table.

As demonstrated in the table, in terms of the imputation score metric, our DiffPuter still significantly outperforms other methods.

Q2: The experimental section lacks fair comparison and clarity

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.

Thanks for your suggestion. We've conducted additional experiments combining EM algorithm with these DGMs. In the following table, we present the performance (MAE metric) of the proposed DiffPuter with EM+other DGMs, i.e., MIWAE[1] and HIWAE[2]. For comparion, we also present the performance of DiffPuter at different EM iteratons.

As demonstrated, even when considering only the first round of iteration, our DiffPuter's performance is far better than MIWAE and HIWAE, which demonstrates the irreplaceable ability of Diffusion models to reconstruct ground-truth data distribution compared to other DGMs. As the iteration steps increase, DiffPuter's performance shows significant improvement. MIWAE's performance also shows some improvement, but the magnitude is quite small. HIWAE's performance shows no significant improvement, and instead appears to be fluctuating. This shows that even with more EM iterations, if the data density is not correctly learned, the imputation performance remains difficult to improve.

The implementation code for these two models has been added to our codebase. We will include these DGM+EM as variants of DiffPuter in the ablation study, after all the new experiments are completed.

References:

[1] Mattei, Pierre-Alexandre, and Jes Frellsen. "MIWAE: Deep generative modeling and imputation of incomplete data sets." International conference on machine learning. PMLR, 2019.

[2] Nazabal, Alfredo, Pablo M. Olmos, Zoubin Ghahramani, and Isabel Valera. "Handling incomplete heterogeneous data using vaes." Pattern Recognition 107 (2020): 107501.

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.

The assertion "Traditional machine learning methods are still powerful imputers" is first observed from the empirical performance in our Figure 2. In Figure 2, we observe that:

• Simple machine learning model, such as simple vanilla EM, can already achieve performance exceeding many deep learning methods (such as GRAPE, IGRM, MOT, TDM).
• HyperImpute, an AutoML method that iteratively incorporates multiple machine learning-based (e.g., tree-based) imputation method ranks the second among all the baseline methods.

In addition, several recent works have also highlighted the of importance and efficacy of traditional ML methods in tabular data [1,2,3,4], and we have added them in the revised version of paper.

References:

[1] Jolicoeur-Martineau, Alexia, Kilian Fatras, and Tal Kachman. "Generating and imputing tabular data via diffusion and flow-based gradient-boosted trees." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

[2] McCarter, Calvin. "Unmasking Trees for Tabular Data." arXiv preprint arXiv:2407.05593 (2024).

[3] Lalande, Florian, and Kenji Doya. "Numerical data imputation: Choose kNN over deep learning." International Conference on Similarity Search and Applications. Cham: Springer International Publishing, 2022.

[4] Suh, Heajung, and Jongwoo Song. "A comparison of imputation methods using machine learning models." Communications for Statistical Applications and Methods 30.3 (2023): 331-341.

The claim that generative methods excel on continuous data requires clarification. Here, the diffusion model seems to assume Gaussianity across all dimensions, using as a proxy to obtain discrete outputs, which is not the optimal to model heterogeneous data.

This is another summary observation from our empirical results. We found that our method, DiffPuter's imputation performance on continuous data, often far exceeds discriminative methods (such as Remasker, GRAPE, and MOT), but for discrete data, the lead margin is very small.

Your statement is correct - the diffusion model indeed assumes Gaussian noise for each dimension. Therefore, it might only be suitable for continuous data and not for discrete data (and heterogeneous data). Our DiffPuter uses the simplest one-hot encoding to handle discrete data, and has already achieved good results. Studying different encoding methods is out of the scope of this paper, but we are interested in combining them with potentially better encoding methods if they are available.

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.

We apologize for writing this part too briefly. Here, we want to emphasize that several SOTA discriminative deep learning methods cannot be applied to out-of-sample imputation setting:

• TDM [1] treats the values of missing entries as learnable parameters within the model, therefore it can only perform imputation on missing data in the training set and cannot be applied to the test set.
• IGRM [2] needs to iteratively construct an implicit social network between all samples during training. For samples in the test set, they cannot be added to the existing network, therefore it doesn't work in the out-of-distribution setting.

We agree that general deep generative models can naturally be applied to out-of-distribution settings. However, considering their suboptimal performance in the in-sample imputation setting (Figure 1 and Table 1), we didn't make comparisons with them. We will add the full comparison in the revised version.

References:

[1] Zhao, He, et al. "Transformed distribution matching for missing value imputation." International Conference on Machine Learning. PMLR, 2023.

[2] Zhong, Jiajun, Ning Gui, and Weiwei Ye. "Data imputation with iterative graph reconstruction." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 9. 2023.

Q1: Motivation of the paper

Motivation appears to overlook recent work. The two main issues presented as motivation for this work are unclear. If multiple DGMs could, indeed, replace diffusion models within the EM framework, how is diffusion specifically justified for tabular data?

We apologize for the confusion caused by the lack of clarity in the motivation section. Our motivation is indeed based on "the necessity of using diffusion models as deep generative models to learn tabular distributions". This is because if a generative model cannot accurately learn the data distribution, even if it can perform accurate conditional sampling, the obtained imputation values will be inaccurate. And in the context of EM, such errors may further accumulate.

Regarding the performance comparison of different types of generative models on tabular data generation tasks, it has been studied in recent works[1,2]. Through extensive experiments, the conclusion is that early deep generative models, such as GANs and VAEs[3,4,5], are quite poor at learning tabular data distributions, and their generated samples struggle to faithfully recover the ground-truth distribution (poor capacity in density estimation). In contrast, Diffusion-based models[1,2] have shown performance far exceeding that of VAE, GAN, and other methods.

Although Diffusion models, as a type of deep generative model, have achieved SOTA performance in tabular data generation, they directly model the joint distribution of all columns, making it difficult to perform conditional inference. This forms our complete motivation.

The motivation in the initial version didn't clearly explain this point, so we made substantial modifications to the introduction to make our motivation clearer (which has been updated in the revised paper). The key points are as follows:

• Motivation of using EM algorithm: The biggest challenge in applying deep generative models to missing data imputation is the inability to observe the complete distribution. The EM algorithm addresses the incomplete likelihood issue by iteratively refining the values of the missing data.

• Challenges of combine EM algorithm and DGMs

  • EM algorithm consists of an M-step (density estimation) and E-step (conditional inference).
  • Achieving both a good E step and M-step is a dilemma.
  • Existing DGMs (Deep Generative Models) that can perform conditional inference (such as VAE) cannot capture tabular data distribution well. Meanwhile, SOTA methods that can faithfully recover data distribution (such as Diffusion) struggle to perform conditional inference.

• Our contribution: Our method explores a path that makes diffusion models compatible with the EM framework for missing data imputation.

Therefore, although other DGMs can technically replace the Diffusion model in our framework, in terms of empirical performance, the diffusion model is irreplaceable. In the following response to Q2, we have conducted supplementary experiments to demonstrate the importance of the Diffusion model in our framework.

References:

[1] Kotelnikov, A., Baranchuk, D., Rubachev, I., & Babenko, A. (2023, July). Tabddpm: Modelling tabular data with diffusion models. In International Conference on Machine Learning (pp. 17564-17579). PMLR.

[2] Zhang, H., Zhang, J., Shen, Z., Srinivasan, B., Qin, X., Faloutsos, C., ... & Karypis, G. Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space. In The Twelfth International Conference on Learning Representations.

[3] Xu, L., Skoularidou, M., Cuesta-Infante, A., & Veeramachaneni, K. (2019). Modeling tabular data using conditional gan. Advances in neural information processing systems, 32.

[4] Liu, T., Qian, Z., Berrevoets, J., & van der Schaar, M. (2023). GOGGLE: Generative modelling for tabular data by learning relational structure. In The Eleventh International Conference on Learning Representations.

[5] Richardson, T. W., Wu, W., Lin, L., Xu, B., & Bernal, E. A. (2020). Mcflow: Monte carlo flow models for data imputation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 14205-14214).

Q3: Discussion of Limitations is lacking (Training time)

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 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.

We thank the reviewer for the suggestion, and we have moved the comparison of training time in Section 5.3 in the main text in the revised paper.

The computational efficiency of this method when applied to high-dimensional data (such as images) could indeed be a problem, because training diffusion models on image data is already inefficient. And the EM algorithm will repeat training this diffusion model many times.

I am also curious if the M-step converges faster with higher values of $k$, as this could enhance efficiency.

Thank you for raising this interesting question. In our original experiments, for each M-step, the parameters of the diffusion model are randomly reinitialized, and the model is trained from scratch, so the training time for each step is actually about the same.

Your point is very reasonable. If in the new M-step, we could continue training based on the diffusion model parameters obtained from the previous M-step, the number of steps needed for model convergence could indeed be greatly reduced, thereby improving training speed.

To verify this, we have conducted additional experiments using the new training strategy. In the following table, we present the number of training epochs for the convergence (loss patience of 200 = 0.2k) of diffusion model at different EM iterations.

From the table results, we can see that as EM iterations increase, the convergence speed of the diffusion model indeed becomes faster and faster, and the training time of the second iteration's diffusion model can already be reduced by over 70% compared to the first iteration. By the 5th and 6th iterations, the diffusion model has almost converged right through the beginning a few steps, and the imputation results remain almost unchanged.

Regarding the overall training time, since the sampling time of the E-step remains constant, DiffPuter's overall training speed can be improved by 2 times (doubled), making it much faster than the other competitive baseline methods IGRM, Hyperimpute, and Remasker.

Q4: Other minor questions

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.

We feel sorry about causing your misunderstanding about Figure 6. Note that for the y-axis in Figure 6, higher values indicate lower MAE. Therefore, when the observed ratio decreases, MAE actually increases (indicating worse imputation performance). We've updated the corresponding paragraph to make it clearer.

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.

MissDiff fails on these datasets (much larger MAE/RMSE values than other methods). In fact, MissDiff's original paper[1] has contradictory descriptions of the model: in the first half of the paper, it states that $m = 1$ indicates observed entries, while in the training section, it states that $m = 1$ indicates missing entries. Finally, the diffusion score-matching loss is only calculated on entries where $m = 1$.

Since the paper does not provide implementation code, we can only reproduce the method based on the paper's textual description. Considering that calculating score-matching loss on missing entries is meaningless (since we don't know the ground-truth values at all), in our reproduction we calculate the loss on observed entries.

Since MissDiff only utilizes the information from partially observed data during training while completely ignoring the missing parts in the form of masks, the data distribution it learns is inherently incomplete (i.e., concentrating only on the observed part). Thus, it's not surprising that it shows such poor imputation performance on the missing part.

References:

[1] Ouyang, Y., Xie, L., Li, C., & Cheng, G. (2023). Missdiff: Training diffusion models on tabular data with missing values. arXiv preprint arXiv:2307.00467.

Thank you for your response. I appreciate the corrections and the effort to more comprehensively describe the related literature. While adding the suggested baselines would enhance the significance of the paper, I understand the timing constraints of the rebuttal process. In its current form, I believe the paper is acceptable.

As a result, I have upgraded my score and will consider raising it further if comparisons with the aforementioned methods are included in the final version. In any case, I will recommend the paper for acceptance.

We thank the reviewer for the further response. We acknowledge that we overlooked several important papers in our previous literature survey, which may have led to some inaccurate statements in the introduction section. In the latest revised paper, we have modified the relevant descriptions, added citations to these related works, and highlighted some of their advantages in missing data imputation.

Regarding the importance of diffusion models in our framework empirically, we are currently conducting additional experiments replacing diffusion models with the two advanced VAE-based methods you have mentioned. Once these experiments are completed, we will add comprehensive experimental explanations in the ablation studies section. In addition, we will add one-EM-step versions of these methods as the general baselines for comparison (in Figure 2 and Table 1).

Again, we thank you for your constructive comments.

I thank the reviewers for their detailed feedback and the additional experiments provided during the rebuttal process. I greatly appreciate the effort and consideration that went into addressing my concerns. While many of my initial questions have been clarified, some issues still remain, mainly about the motivation of the paper:

• I continue to find the justification for using diffusion models on tabular data unconvincing. Specifically, the claim that “previous generative models are poor at imputation” is unfair. There is extensive literature, not cited in the paper, demonstrating that models such as VAEs can achieve competitive results in imputation tasks on tabular data. For example, [2] and [3] build on [5] to successfully improve heterogeneous data modeling, achieving accurate imputation results. Furthermore, both TabDDPM and TABSYN, which apply a similar technique for adapting to heterogeneous data, also omits these critical references.

• In your rebuttal, you state: “even if it—a generative model—can perform accurate conditional sampling, the obtained imputation values will be inaccurate.” However, it’s important to note that the deep generative models referenced above are trained on $p(x_o)$, where $x_o$ is artificially constructed by adding random missingness to the available data during training. This key trick, assumed in these works [1-5], significantly enhances imputation accuracy.

As I emphasized in my initial review and again in this rebuttal, other deep generative models, such as the ones that demonstrated competitive imputation results, [2,3] could competitively replace the proposed Diffusion model for this task. I appreciate that the authors included some of the references as baselines, and I do not necessarily expect the rest to be included. However, although this comparison does not diminish the significance or novelty of DiffPutter, I simply find misleading to claim in the paper that VAEs and GANs perform poorly at imputation. This statement is inaccurate and should be revised to reflect the broader literature.

Other issues

Given the method's novelty isn't specific to tabular data, the related work should include other imputation methods (e.g., image inpainting)

Thank you for the suggestion! We've included several related works about image inpainting in Section 2 of the revised paper.

A simulation study with multiple modes would be valuable, particularly as diffusion-based models should excel in such scenarios.

We don't quite get what you mean by "simulation study of multiple modes". We would be very grateful if you could provide clearer hints about this.

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)

How to initialize missing values is indeed important. In the paper, we only adopted the simplest approach of using the mean of observed values for initialization and have already achieved good results. Considering the convergence property of the EM algorithm, developing novel initialization methods might not bring significant improvements in performance but might improve training speed (via improving the convergence speed).

We've also tried to initialize the missing values using KNN (which is a baseline in the experiments). The table below compares the training time and MAE score obtained without and with KNN as the base imputation method, respectively, on the adult dataset.

As demonstrated in the Table, using a KNN model as the initial imputer can greatly reduce the training time with marginal performance improvement.

Q1: How to address "missing data as limit of detection" (LOD) problem?

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.

Thank you for raising this interesting and worthy question for discussion. We are happy to share some of our thoughts on addressing this issue.

First, it is obvious that if we have only one column in our data, where both LOD data and missing data are forcibly observed as 0, this would indeed create a fundamental identification problem, specifically manifested in:

Assuming the real data follows a certain distribution (such as normal distribution) When we observe a value of 0, we cannot distinguish whether this 0 comes from:

• Real value ≤ 0 being truncated
• Completely random missing.

At the point of 0, the observed results produced by these two mechanisms are exactly the same. However, since our data is multivariate (having many different columns), and the values in other columns may influence whether the target column's value is LOD or missing. Therefore, we can obtain additional information from these to help us identify whether a column is LOD or missing.

For example, consider the following simple scenario. Besides the target column (denoted as $x_{j_1}$), we have another column $x_{j_2}$ that is fully observed, and $x_{j_1}$'s observed values have a strong correlation with $x_{j_2}$. For instance, when $y \ge 0$, $x_{j_1}$'s values tend to concentrate near LOD = 0, while when $y$ < 0, $x_{j_1}$'s values are typically much larger than 0. Then, if we obtain an $x_{j_1} = 0$ and $x_{j_2}$ > 1, we can reasonably conclude that the $x_{j_1}$'s value is indeed the observed LOD value, rather than a missing value.

Extending to more general cases, if the value of the target column $x$ (whether it is missing or not) is indeed conditioned on other columns, we can train a model to identify whether it is an observed LOD value or missing. We can assume that this model can output a probability $p$, representing a $p$ probability that the position is a missing value, and a $1-p$ probability that it is an observed LOD value.

Consider a tabular dataset $X \in \mathbb{R}^{N \times d}$, and $x_{ij}$ denotes the value at row $i$ and column $j$. Let's assume $LOD = 0$ for convenience. Then, we may resort to statistic models, e.g., mixture models, to handle this problem.

For example, we may consider the mixture of Gamma and Gaussian distributions. For each non-zero observation $x_{ij}$ in position $(i,j)$: $P(x_{ij}) = w_j P_{\text{gamma}}(x_{ij}|\alpha_j,\beta_j) + (1-w_j) P_{\text{gaussian}}(x_{ij}|\mu_j,\sigma_j^2)$

where: $P_{\text{gamma}}(x|\alpha,\beta) = \frac{\beta^\alpha}{\Gamma(\alpha)} x^{\alpha-1} e^{-\beta x}$, $P_{\text{gaussian}}(x|\mu,\sigma^2) = \frac{1}{\sqrt{2\pi\sigma^2}} e^{-\frac{(x-\mu)^2}{2\sigma^2}}$

The log-likelihood for column $j$:

$\mathcal{L}_j = \sum\_{i, x\_{ij} \neq 0} \log [w_j P _{gamma}(x _{ij}|\alpha _j,\beta _j)]$

$+ (1-w_j) P_{\text{gaussian}}(x_{ij}|\mu_j,\sigma_j^2)] + \sum\_{i:x_{ij}=0} \log[w_j P_{\text{gamma}}(\varepsilon|\alpha_j,\beta_j) + (1-w_j) P_{\text{gaussian}}(0|\mu_j,\sigma_j^2)]$

The optimal distribution parameters can be obtained via Maximum likelihood estimation: $\hat{\alpha} _j,\hat{\beta} _j, \hat{\mu} _j, \hat{\sigma} _j^2, \hat{w} _j = \arg\max _{\alpha _j,\beta _j,\mu _j,\sigma _j^2,w _j} \mathcal{L}_j$ subject to the following conditions:

• $\alpha_j, \beta_j, \sigma_j > 0$
• $0 \leq w_j \leq 1$

The above approach assumes that the missing probability only relates to the column and not to the row. To consider the influence of relationships between rows, we can first cluster the input data, and for each cluster, learn a cluster-specific mixture model.

Finally, for zero values, the missing probability is: $P(\text{miss}|x_{ij}=0) = \frac{(1-w_j)P_{\text{gaussian}}(0|\mu_j,\sigma_j^2)}{w_j P_{\text{gamma}}(\varepsilon|\alpha_j,\beta_j) + (1-w_j)P_{\text{gaussian}}(0|\mu_j,\sigma_j^2)}$

Our method can be effective based on this foundation. For each position marked as an LOD value ($=0$), there is a probability $p$ where we assume it is missing, and using our method we ultimately obtain a predicted imputation value, which we denote as $x_{pred}$. Additionally, there is a $1-p$ probability where we consider it to be the observed LOD value. Therefore, we mix these two scenarios according to their probabilities, meaning the final predicted value is $p * x_{pred} + (1-p) * 0$.

W4: Encoding for categorical data.

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).

One-hot encoding is a conventional method for handling discrete/categorical tabular data, and has been used in many previous papers [1,2,3], although it may not be the most appropriate one. We acknowledge that how to more effectively handle categorical data is indeed a major challenge in tabular deep learning, but it is out-of-scope of this paper. Essentially, one-hot encoding is just the simplest choice for processing categorical data, and if there are other more reasonable methods, they can be directly combined with our approach.

References:

[1] Kim, Jayoung, Chaejeong Lee, and Noseong Park. "STaSy: Score-based Tabular data Synthesis." The Eleventh International Conference on Learning Representations.

[2] Zheng, Shuhan, and Nontawat Charoenphakdee. "Diffusion models for missing value imputation in tabular data." arXiv preprint arXiv:2210.17128 (2022).

[3] Liu, Tennison, et al. "GOGGLE: Generative modelling for tabular data by learning relational structure." The Eleventh International Conference on Learning Representations. 2023.

Also, the use of mean is inherently problematic due to common multi-modality in the data.

For numerical data, mean imputation is natural. For categorical data, our mean imputation based on one-hot encoding essentially estimates the missing column values according to the marginal distribution (prior distribution) of each category. For example, for the gender column, if the data contains 80% male and 20% female, then for a data point with a missing gender column, we naturally assume it has an 80% probability of being male and a 20% probability of being female, therefore assigning its corresponding (one-hot) embedding vector as [0.8, 0.2]. This is a natural and reasonable processing approach.

Furthermore, from the experimental results, our simple processing method has already achieved quite good results (outperforming other baseline methods).

W5: Clarify the novelty of the proposed method in the related works section.

The novelty of the method compared to other approaches is not clearly articulated in the related works section.

We thank the Reviewer for their suggestion. In the last paragraph of Related Works, we implicitly pointed out the novelty of our method by summarizing the limitations of existing methods. In the revised paper, we have added the following sentence, which clearly and explicitly states the novelty of our method compared with existing ones:

"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."

W2 Importance of missing data imputation. It's effects on downstream tasks.

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.

Thank you for your suggestion. We agree with your point that applying well-imputed data to downstream-specific tasks is also very important. Therefore, we considered using the imputed data to perform classification and regression tasks for the target column. We have also added the implementation of this task to our codebase.

Specifically, we first split the complete data into training and test sets, train an XGBoost classifier or regressor on the training set, and test it on the test set. Then, we add masks to the training set to create missing data, and we use different missing value imputation models to obtain imputed data. Next, we train the same XGBoost classifier/regressor using the imputed data and test it on the test set. The effectiveness of imputation can be measured by comparing the performance differences obtained from these two tests.

In the table below, we show the performance on the test set when training on complete data, as well as training on imputed data obtained through different imputation methods.

As demonstrated in the table, our method still outperforms other SOTA imputation methods on this task. However, we find that the performance differences between different methods are not very large, therefore this task might not be sufficient to evaluate the quality of imputation comprehensively. We will appreciate it if you can suggest more appropriate downstream tasks.

W3: Evaluation on other missing settings (e.g., MNAR).

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.

Regarding MER and MNER mentioned by the Reviewer, we assume that these should be MAR (Missing At Random) and MNAR (Missing Not At Random). We humbly point out that our experimental section considered all three missing mechanisms: 1) Missing Completely At Random (MCAR), 2) Missing At Random (MAR), and 3) Missing Not At Random (MNAR). This is explained in Section 5.1, Datasets part. Due to length constraints in the main text, we only presented the experimental results for the MCAR setting. We placed the specific implementation of these three settings in Appendix C.3 (D.3 in the revised version of paper), and the experimental results in Appendix D.2 and Appendix D.3. (E.2 and E.3 in the revised paper).

We greatly appreciate the reviewer for taking the time to provide such constructive comments and suggestions. The following are detailed responses to your questions.

W1: How the hyperparameters of baseline methods are selected?

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.

Most of the deep learning baselines recommend the use of one set of hyperparameters for all datasets. For these methods, we directly follow their guidelines and use the default hyperparameters:

• ReMasker: we use the recommended hyperparameters provided in Appendix A.2 in [1]. This set of hyperparameters is searched by tuning on the Letter dataset and is deployed for all the datasets in the original paper; hence, we follow this setting.
• HyperImpute: since HyperImpute works by searching over the space of classifiers/regressors and their hyperparameters, it does not have hyperparameters itself except parameters related to the AutoML search budget. We adopt the default budget parameters of HyperImpute's official implementation for all datasets. The default budget parameters and AutoML search space are provided in https://github.com/vanderschaarlab/hyperimpute/blob/main/src/hyperimpute/plugins/imputers/plugin_hyperimpute.py and Table 5 in the original paper [2].
• MOT and TDM: 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 n_iter. While the default value is 3000 (https://github.com/BorisMuzellec/MissingDataOT/blob/master/experiment.py), 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 is controlled by max_iter and rr_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 $\sqrt{n}$, where $n$ is the number of samples in the dataset.
• GRAPE and IGRM: we adopt the recommended set of hyperparameters used in the original paper for all datasets. For a detailed explanation of the meaning of the parameters, please see https://github.com/maxiaoba/GRAPE/blob/master/train_mdi.py for GRAPE and https://github.com/G-AILab/IGRM/blob/main/main.py for IGRM.
• MissDiff: since the original implementation is not available, and it is based on diffusion model, for fair comparison, we simply use the same set of hyperparameters with our DiffPuter.
• TabCSDI: We follow the guide for selecting hyperparameters in the original paper (Appendix B in [3]). Specifically, we use a large version of the TabCSDI model with a number of layers set to 4 (see more detailed hyperparameters about the large TabCSDI model at https://github.com/pfnet-research/TabCSDI/blob/main/config/census_onehot_analog.yaml). For batch size, we 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 we used have the number of samples between 20000 to 40000, we scale the batch size to 256 and use it for all datasets.
• MCFlow: we adopt the recommended hyperparameters provided in the official implementation for all datasets (https://github.com/trevor-richardson/MCFlow/blob/master/main.py).

For the remaining classical machine learning methods, including EM, GAIN, MICE，Miracle，MissForest, and Softimpute where hyperparameters might be important. Since we use the implementations from the 'hyperimpute' package, we tune the hyperparameters within the hyperparameter space provided in the package (e.g., https://github.com/vanderschaarlab/hyperimpute/blob/main/src/hyperimpute/plugins/imputers/plugin_missforest.py). 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.

References:

[1] Tianyu Du, Luca Melis, and Ting Wang. Remasker: Imputing tabular data with masked autoencoding. In International Conference on Learning Representations, 2024.

[2] Daniel Jarrett, Bogdan C Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperimpute: Generalized iterative imputation with automatic model selection. In International Conference on Machine Learning, pp. 9916–9937. PMLR, 2022.

[3] Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular data. In NeurIPS 2022 First Table Representation Workshop, 2022.

We greatly appreciate the reviewer for taking the time to provide such constructive comments and suggestions. Below are our detailed responses to the questions you raised.

Comparison of complexity

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.

Thanks for your suggestion. Here we analyze the computational complexity of our DiffPuter and two other Diffusion models, TabCSDI[1] and MissDiff[2].

In the training of diffusion model, our DiffPuter and TabCSDI, MissDiff have the same level of time complexity, because they essentially all compute diffusion loss on the input data. The difference is that MissDiff and TabCSDI use 0/1 masks to compute loss only on certain entries, while our DiffPuter computes loss on all entries.

During the sampling phase, all methods have the same complexity, because both TabCSDI and MissDiff need to sample many samples and take their mean as the final imputation value.

Considering that our method needs to repeat this process K times, the overall computational complexity of our method is approximately K times that of MissDiff and TabCSDI.

The actual training time is influenced by many factors. For example, since our DiffPuter uses a lightweight MLP model as the denoising neural network, while TabCSDI uses Transformers, in practice DiffPuter's convergence speed per iteration is much faster than TabCSDI.

Considering the huge performance improvement that our DiffPuter brings compared to TabCSDI and MissDiff (approximately 50% improvement on MAE and RMSE), this increase in complexity is worthwhile.

The actual training time is influenced by many factors. For example, since our DiffPuter uses a lightweight MLP model as the denoising neural network, while TabCSDI uses Transformers, in practice DiffPuter's convergence speed per iteration is much faster than TabCSDI.

In Appendix D.1 (Section 5.3 in the revised version), we compared the training time of DiffPuter with other SOTA imputation methods. Overall, our method has similar training speed to other SOTA methods, but brings average performance improvements of 8% to 25%.

References:

[1] Zheng, S., & Charoenphakdee, N. (2022). Diffusion models for missing value imputation in tabular data. arXiv preprint arXiv:2210.17128

[2] Ouyang, Y., Xie, L., Li, C., & Cheng, G. (2023). Missdiff: Training diffusion models on tabular data with missing values. arXiv preprint arXiv:2307.00467.