MissDiff: Training Diffusion Models on Tabular Data with Missing Values

We thank the reviewer for the comments, and we appreciate the time you spent on the paper.

Q: Missing related work: There are works that utilize the ability of gradient-boosted decision trees (GBDTs) to handle missing features in order to avoid the impute-then-generate framework. Examples include [ForestDiffusion 1, DiffPuter 2] that use diffusion modeling and [UnmaskingTrees 3] that uses autoregression modeling. In particular, it is not clear whether the diffusion modeling approaches in [1, 2] can be seen as instances of the authors' proposed framework. (Granted, I realize that [1, 2, 3] may be parallel works, and that these rely on GBDTs for missingness, which is not a very general solution. So while I think these are relevant related works, I would like to emphasize that these do not diminish the potential impact of MissDiff, which is the first to enable the leap to neural networks, which will be important for modeling textual fields, and which is the first to be rigorously shown to not rely on assumptions like MCAR.)

A: Thank you for your comments. We include a discussion of the related work in lines 310-317. Especially, [ForestDiffusion 1] adopts XGBoost to estimate the score, [DiffPuter 2] leverages the Expectation-Maximization that first learns the joint distribution of both the observed and currently estimated missing data and then re-estimates the missing data based on the conditional probability given the observed data, and [UnmaskingTrees 3] adopts tree-based autoregressive modeling of tabular data. These methods can be used for imputation together with generation. However, we believe our work still has the takeaway. We propose the masked score matching loss and prove the theoretical soundness of this technique.

Q: It would be preferable to report how MissDiff performs on the imputation benchmark of [ForestDiffusion 1], for the following reasons. First, it has a larger and more diverse set of 27 datasets. Second, unlike Zheng & Charoenphakdee (2022)'s benchmark, this one shows that MissForest is better than GAIN, which is frankly more plausible. Third, without any further reimplementation work, it would enable one to compare MissDiff to the current claimed SotA diffusion [2] and autoregression [3] methods. Fourth, it should be simple to run MissDiff on this benchmark during the rebuttal period.

A: Thank you very much for your insightful suggestions. We appreciate your detailed reasoning and agree that a broader evaluation could provide additional insights. We provide the results on the imputation benchmark of [ForestDiffusion 1] as follows:

Tabular data imputation (27 datasets, 3 experiments per dataset, 10 imputations per experiment) with 20% missing values; raw score - mean (standard error)

Q: Line 351 "MICE based on random forest (MissForest) (Stekhoven, 2015)". MissForest is not exactly following the MICE framework using random forest. Are you comparing to MICE-Forest (presumably the implementation of AnotherSamWilson/miceforest) or truly MissForest (presumably the implementation of epsilon-machine/missingpy)? Please clarify such details in the manuscript.

A: Thank you very much for pointing this out. We adopt truly the MissForest implementation of https://github.com/vanderschaarlab/hyperimpute. We revise the description of MissForest in line 360.

Q: While leaving the architectural details to the manuscript appendix is fine, it would be nice if the main text indicated that a neural network model architecture is used.

A: Thank you very much for your suggestion. We add the description of the architecture details in the revised manuscript (lines 425-427).

Q: I did not see any challenge in extending the proposed method to scale up for addressing missing data issues in high-dimensional cases, such as image imputation. Could you provide relevant results to demonstrate the method's effectiveness in these scenarios?

A: Thank you very much for your question. We think our paper mainly focuses on the tabular data. Existing literature including [1-7] do not contain the experiments in the image. We would like to show more evidence of the effectiveness of the proposed method. Therefore, we conduct the experiments in imputation tasks against misGAN, the experimental results are as follows:

The effectiveness of MissDiff on MNIST.

[1]: Kim, J., Lee, C.E., & Park, N. (2023). STaSy: Score-based Tabular data Synthesis. International Conference on Learning Representations.

[2]: Xu, L., Skoularidou, M., Cuesta-Infante, A., & Veeramachaneni, K. (2019). Modeling Tabular data using Conditional GAN. Neural Information Processing Systems.

[3]: Kotelnikov, A., Baranchuk, D., Rubachev, I., & Babenko, A. (2023). TabDDPM: Modelling Tabular Data with Diffusion Models. International Conference on Machine Learning.

[4]: Zhang, H., Zhang, J., Srinivasan, B., Shen, Z., Qin, X., Faloutsos, C., Rangwala, H., & Karypis, G. (2024). Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space. International Conference on Learning Representations.

[5]Jolicoeur-Martineau, A., Fatras, K., & Kachman, T. (2023). Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees. International Conference on Artificial Intelligence and Statistics.

[6]Zhang, H., Fang, L., & Yu, P.S. (2024). Unleashing the Potential of Diffusion Models for Incomplete Data Imputation. ArXiv, abs/2405.20690.

[7]McCarter, C. (2024). Unmasking Trees for Tabular Data. ArXiv, abs/2407.05593.

To help the reviewer understand the proof of Theorem 3.2 well, we summarize the main procedure in the proof as follows,

1. we relate the score on the incomplete data is equivalent to the score on the complete data when performing element-wise multiplication with mask, i.e., when Gaussian transition distribution with the isotropic covariance matrix, we have $\nabla_{\mathbf{x}^{\text{obs}}(t)} \log p(\mathbf{x}^{\text{obs}}(t) |\mathbf{x}^{\text{obs}}(0)) \odot \mathbf{m} = \nabla_{\mathbf{x}(t)} \log p(\mathbf{x}(t) |\mathbf{x}(0)) \odot \mathbf{m}$.
2. Then, we can get the optimal solution of the Denosing Score Matching objective on missing data admits the same solution of the Denosing Score Matching objective on the complete data. $\underset{\boldsymbol{\theta}}{\arg \min} \ \mathbb{E} _{p(\mathbf{x}^{\text{obs}}(0),\mathbf{m})} \mathbb{E} _{p(\mathbf{x}^{\text{obs}}(t) | \mathbf{x}^{\text{obs}}(0))} [\|(\mathbf{s} _{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}(t), t)-\nabla _{\mathbf{x}^{\text{obs}}(t)} \log p(\mathbf{x}^{\text{obs}}(t) |\mathbf{x}^{\text{obs}}(0))) \odot \mathbf{m}\|_2^2]$ has the same solution as $\underset{\boldsymbol{\theta}}{\arg \min} \ \mathbb{E} _{p(\mathbf{x}(0),\mathbf{m})} \mathbb{E} _{p(\mathbf{x}(t) | \mathbf{x}(0))} [\|(\mathbf{s} _{\boldsymbol{\theta}}(\mathbf{x}(t), t)-\nabla _{\mathbf{x}(t)} \log p(\mathbf{x}(t) |\mathbf{x}(0))) \odot \mathbf{m}\|_2^2]$. That is because when we can decompose the least square loss into three terms, two of which are the squares of their respective components, and the remaining one is a cross-term. By utilizing the previous conclusion, we can get equality.
3. Finally, we take the conditional expectation of mask inside the least squares loss since $\mathbf{m}_i\in\{0,1\}$ we have $\mathbb{E}[\mathbf{m}_i^2]=\mathbb{E}[\mathbf{m}_i]$ for any distribution of $\mathbf{m}$.

Illustration: The derivation is correct no matter what value we use to impute the missing value in $\mathbf{x}^{\text{obs}}(0)$. Let $\boldsymbol{\theta}^* _{\text{obs}}$ denotes the optimal solution for the Denosing Score Matching objective on missing data, i.e., $\boldsymbol{\theta}^* _{\text{obs}}= \underset{\boldsymbol{\theta}}{\arg \min} \ \mathbb{E} _{p(\mathbf{x}^{\text{obs}}(0),\mathbf{m})} \mathbb{E} _{p(\mathbf{x}^{\text{obs}}(t) | \mathbf{x}^{\text{obs}}(0))} [\|(\mathbf{s} _{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}(t), t)-\nabla _{\mathbf{x}^{\text{obs}}(t)} \log p(\mathbf{x}^{\text{obs}}(t) |\mathbf{x}^{\text{obs}}(0))) \odot \mathbf{m}\| _2^2]$, and let $\boldsymbol{\theta}^*$ denotes the optimal solution for the Denosing Score Matching objective on complete data, i.e., $\boldsymbol{\theta}^*= \underset{\boldsymbol{\theta}}{\arg \min} \ \mathbb{E} _{p(\mathbf{x}(0),\mathbf{m})} \mathbb{E} _{p(\mathbf{x}(t) | \mathbf{x}(0))} [\|(\mathbf{s} _{\boldsymbol{\theta}}(\mathbf{x}(t), t)-\nabla _{\mathbf{x}(t)} \log p(\mathbf{x}(t) |\mathbf{x}(0))) \odot \mathbf{m}\|_2^2]$, we are not saying that $\mathbf{s} _{\boldsymbol{\theta}^* _{\text{obs}}}(\mathbf{x}^{\text{obs}}(t), t)=\mathbf{s} _{\boldsymbol{\theta}^*}(\mathbf{x}(t), t)$. We show that $\mathbf{s} _{\boldsymbol{\theta}^* _{\text{obs}}}(\mathbf{x}(t), t)=\mathbf{s} _{\boldsymbol{\theta}^*}(\mathbf{x}(t), t)$.

We thank the reviewer for the comments, and we appreciate the time you spent on the paper.

Q: The authors claim that the imputation-then-gernerate method may lead to inconsistent estimator. However, MISSDIFF the author proposed is trying to fill 0 to the observed data (the i-th element of $x_{\text{obs}}$ equals to 0, when the i-th element is missing) and use this imputed data to train the model. The only difference is that the author use a mask to relieve the effect of the loss function. As such, since MISSDIFF is also kind of imputation-then-gernerate method, why MISSDIFF does not suffer from the issue of imputation method? Is there any intuition to explain the superiority of MISSDIFF?

A: Thank you very much for your question. We clarify this point below. First of all, we prove the imputation-then-generate method is biased in Remark 3.1. We prove this line of methods no longer captures the data variability since $\underbrace{\mathbb{E} _{p_0\left(\mathbf{x}^{\mathrm{miss}} _{\text{}}| \mathbf{x}^{\mathrm{obs}}\right)}[p _{\mathbf{\phi}}(\mathbf{x}^{\text {obs}}, \mathbf{x}^{\text {miss}})]} _{\text{unbiased likelihood}}\neq \underbrace{p _{\mathbf{\phi}}(\mathbf{x}^{\mathrm{obs}}, \mathbb{E} _{p_0\left(\mathbf{x}^{\mathrm{miss}} _{\text{}}| \mathbf{x}^{\mathrm{obs}}\right)}[\mathbf{x}^{\mathrm{miss}}])} _{\text{imputation-then-generate likelihood}}$. However, MISSDIFF does not maximize the likelihood of imputed data $p _{\mathbf{\phi}}(\mathbf{x}^{\mathrm{obs}}, \mathbb{E} _{p_0\left(\mathbf{x}^{\mathrm{miss}} _{\text{}}| \mathbf{x}^{\mathrm{obs}}\right)}[\mathbf{x}^{\mathrm{miss}}])$. We proposed a masked score matching objective to learn the score function in the observed data. In Theorem 3.2, we prove our objective is the lower bound on the log-likelihood of observed data rather than the likelihood of imputed data (which is the objective of imputation-then-generate methods). Therefore, we don't regard MISSDIFF as a kind of imputation-then-generate method. Moreover, we verify the rationality of MISSDIFF by Theorem 3.2, i.e., our masked score matching objective can recover the oracle score on complete data, which leads to an unbiased method for generation tasks.

Q: The proof of Thereom 3.2 is not convincing. In the footnote 8, page 15, the author claim the score function of gaussian product the mask with the partial-observed data can recover the one with fully observed data. However, the property does not hold for $s_\theta(x(t), t)$. If one impute the i-th element of x(t) by 0, then possibly every element of $s_\theta(x(t), t)$ will change. In that case, how to show the equation line at 773-778? A: Thank you very much for your question. We clarify this point below. First of all, $\mathbf{x}(t)$ denotes the complete (unobservable) data in time $t$ and $\mathbf{x}^{\text{obs}}(t)$ denotes the observed data in time $t$. To keep the $\mathbf{x}^{\text{obs}}(t)$ has the same dimension over different missingness, we use 0 to replace $\mathrm{na}$ for continuous variables in $\mathbf{x}^{\text{obs}}(0)$ and a new category to represent $\mathrm{na}$ for discrete variables. We are going to show that this operation does not affect the correctness of the proof in Theorem 3.2.

• For footnote 8, page 15 Since the current diffusion model used in practice is the Gaussian transition distribution with the isotropic covariance matrix, we have $p(\mathbf{x}^{\text{obs}}(t) |\mathbf{x}^{\text{obs}}(0))=\mathcal{N}(\mathbf{x}^{\text{obs}}(t);\mathbf{\mu}^{\text{obs}},\Sigma)$ and $p(\mathbf{x}(t) |\mathbf{x}(0))=\mathcal{N}(\mathbf{x}(t);\mathbf{\mu},\Sigma)$, with $\Sigma=(1-\bar{\alpha}_t)\mathbb{I}$ and $\mathbf{\mu}^{\text{obs}} =\mathbf{\mu} \odot \mathbf{m}$. Then, we can get $\nabla _{\mathbf{x}^{\text{obs}}(t)} \log p(\mathbf{x}^{\text{obs}}(t) |\mathbf{x}^{\text{obs}}(0)) \odot \mathbf{m} = - \frac{1}{(1-\bar{\alpha}_t)}(\mathbf{x}^{\text{obs}}(t)-\mathbf{\mu}^{\text{obs}})\odot \mathbf{m} = -\frac{1}{(1-\bar{\alpha}_t)}(\mathbf{x}(t)-\mathbf{\mu})\odot \mathbf{m} = \nabla _{\mathbf{x}(t)} \log p(\mathbf{x}(t) |\mathbf{x}(0)) \odot \mathbf{m}$. This equality is correct no matter what value we use to impute the missing value in $\mathbf{x}^{\text{obs}}(0)$.
• For lines 773-778 (775-777 in the revised version): We take the conditional expectation of the binary mask inside the least square loss.

Q: Can any connections be built between (partial) identifiability and diffusion? If not, can authors demonstrate why the "general framework" considered in the paper is better than assuming missing data mechanisms (from a theoretical and utility perspective)? For example, when we don't have a user-defined assumption, how MissDiff can benefit users?

A: Thank you very much for your insightful question. We acknowledge that we did not explicitly consider the issue of identifiability in our work. To the best of our knowledge, the identifiability of diffusion model parameters remains an open problem without a definitive answer. This is primarily due to the architecture used in the diffusion model is $\mathcal{R}^d$ to $\mathcal{R}^d$. We also admit that we are not experts on identifiability and offer some relevant analysis:

To analyze the identifiability of the diffusion model without missing data, we decompose it into two parts:

1. the connection between data distribution $p_\theta(x)$ and its score $s_\theta(x)$;

2. the connection between $s_\theta(x)$ and $\theta$, in other words, if \theta^*=\underset{\theta}{\arg \min } \frac{T}{2}\mathbb{E}_t\Big\\{\lambda(t) \mathbb{E} _{x(0)} \mathbb{E} _{x(t) \mid x(0)} \left[\left\|\mathbf{s} _{\theta} (x(t), t)- \mathbf{s} ^*(x(t), t)\right\| _2^2\right]\Big\\}.

• 1. can be easily proved: Assume there are two different probability distributions $p_{\theta_1}(x)$ and $q_{\theta_1}(x)$ with the same score function, $s(x)=\nabla_x \log p_{\theta_1}(x)=\nabla_x \log q_{\theta_2}(x)$. $\log p_{\theta_1}(x)=\log q_{\theta_2}(x)+C$

Due to the integrals over the domain must equal 1, we have $p_{\theta_1}(x)=q_{\theta_2}(x)$.

• According to the similar data setting in [1], $p(x)=\Pi_{i=1}^{d}p(x_i)$ (the only difference with Data setting D1 in [1] is we did not consider the latent variable $z$), an immediate conclusion is score $s(x) = [s_1(x_1), ..., s_d(x_d)]$. Then the objective function becomes element-wise \underset{\theta}{\arg \min } \frac{T}{2} \mathbb{E} _t \Big\\{ \lambda(t) \mathbb{E} _{x(0)} \mathbb{E} _{x(t) \mid x(0)} \left[ \sum _{i=1}^{d} \left\| \mathbf{s} _{\theta_i}(x _i(t), t)- \mathbf{s}^*_i(x_i(t), t)\right\| _2^2 \right] \Big\\}. If oracle score $s^*(x)$ is unique, then $\theta^*$ is unique.

Therefore, we find that diffusion models have good properties and are identifiable when there is no missing data. When there is missing data, we prove the score network learned from our objective can recover the oracle score of complete data in Theorem 3.2. Therefore, we demonstrate the identifiability of the score function. To demonstrate the connection between $s_\theta(x)$ and $\theta$, we aim to prove \theta^*=\underset{\theta}{\arg \min } \frac{T}{2}\mathbb{E} _t\Big\\{\lambda(t) \mathbb{E} _{x^{\text{obs}}(0)} \mathbb{E} _{x^{\text{obs}}(t) \mid x^{\text{obs}}(0)} \left[\left\|\left(\mathbf{s} _{\theta}(x^{\text{obs}}(t), t)- \mathbf{s}^*(x^{\text{obs}}(t) \right) \odot \mathbf{m}\right\| _2^2\right]\Big\\}.

Under the same data setting, we can also decompose the score function into element-wise, i.e., \theta^*=\underset{\theta}{\arg \min } \frac{T}{2} \mathbb{E}_t \Big\\{ \lambda(t) \mathbb{E} _{x^{\text{obs}}(0)} \mathbb{E} _{x^{\text{obs}}(t) \mid x^{\text{obs}}(0)} \left[ \sum _{i=1}^d \left\|\left(\mathbf{s} _{\theta_i}(x_i^{\text{obs}}(t), t)- \mathbf{s}^*_i(x_i^{\text{obs}}(t) \right) \odot \mathbf{m} \right\| _2^2\right]\Big\\}. If oracle score $s^*(x)$ is unique and the supreme of missing rates $\rho _{\text{max}}<1$, then $\theta^*$ is unique.

Q: Have you considered how ignoring observations with missing elements impacts statistical efficiency and identifiability of the model? Could you provide analysis on this?

A: Thank you very much for your suggestion. While we did not provide a specific analysis on this aspect, we believe the impact would heavily depend on the underlying missing data mechanism and the missing ratio. The performance is likely to degrade due to the significant reduction in the size of the training data. For example, if the missing mechanism results in very few complete observations, there would be insufficient training data to effectively train the diffusion model.

We thank the reviewer for the comments, and we appreciate the time you spent on the paper.

Q: Missing data assumptions are crucial for statistical efficiency. And theoretical results presented in the paper ignoring these observations. Ignoring observations with missing elements may lead to invalid inference (imputation) and loss of statistical efficiency. As a result, MissDiff might have better experimental results and can't derive the identifiability result. Could authors clarify how do the theoretical guarantees change under different missing data mechanisms (e.g. MCAR, MAR, MNAR)? It seems to me MissDiff is a general imputation framework, what's the strong assumption under the hood?

A: Thank you very much for your insightful question. In Theorem 3.2, we prove that if the supreme of missing rates $\rho_{\text{max}}:=\max_{i=1,\ldots,d} \sup_{\mathbf{x}} \rho_i(\mathbf{x})<1$, then our objective can recover the score on complete data. This assumption is quite different from missing data mechanisms (e.g. MCAR, MAR, MNAR). Even in MCAR, we can have one dimension of the data completely missing, which avoids our assumption $\rho_{\text{max}}<1$. Even in MNAR, $\rho_{\text{max}}<1$ might be not a too strong assumption.

After learning the score of complete data, we use it for imputation. The same as the procedure in [1] page 5 last paragraph: the Bayesian inference problem $p_\theta(x^{\text{miss}}|x^{\text{obs}}) \propto p_\theta(x^{\text{miss}},x^{\text{obs}})$. Therefore, MissDiff requires $\rho_{\text{max}}<1$ for guarantee good performance.

Q: Results in table 4 and appendix C2 provide some analysis on Q1. But why MICE is not compared? From my experience, MICE is a pretty strong baseline for these two scenarios (MAR and NMAR).

A: Thank you very much for raising this question. We believe the experimental results in Table 1 are convincing to prove the effectiveness of the proposed method, which is the benchmark of previous imputation tasks. We further demonstrate the effectiveness against the existing diffusion-based method, therefore, we change the missing mechanism and provide the experimental results in Appendix C2. Additionally, we included comparisons with MICE as follows. We believe the additional results further substantiate the competitiveness of our approach. However, we compare the utility of the different methods for generating new samples. MICE cannot be directly applied to generate new samples. Therefore, we did not add the experimental results of MICE in Table 4.

The effectiveness of MissDiff on imputation tasks under MAR and NMAR.

Q: The idea of misGAN is also very similar to the proposed algorithm. On the high level, it tries to generate samples looks like the observed data after masking. It trains a generator by minimizing a divergence between observed data ($p_{\text{obs}}$) and generated samples ($p_\theta$). The KL divergence between $p_{\text{obs}}$ and $p_\theta$ is the same as the cross entropy in Theorem 3.3 up to a constant. I would appreciate discussions on the similarities and differences between the proposed method and misGAN.

A: Thank you for your question. on the high level, MissDiff minimizes the KL divergence between $p_{\text{obs}}$ and $p_\theta$ while misGAN minimizes the Jensen–Shannon (JS) between $p_{\text{obs}}$ and $p_\theta$, which shows the similarities between MissDiff and MisGAN. On the low level, we approximate the divergence by different modeling. MissDiff builds a trajectory from noise distribution to the data distribution, then the KL divergence can be calculated by the l2 norm of the velocity field (score function) of the trajectory. While misGAN calculates the JS divergence by using the domain classifier.

MissDiff builds upon diffusion models, which offer two major advantages over GAN-based generative models: (1) the ability to produce higher-quality synthetic data, as demonstrated by Dhariwal and Nichol (2021); and (2) a more stable training objective, which makes the training process significantly easier compared to GANs.

Ref: Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in neural information processing systems 34 (2021): 8780-8794.

Q: There is another work that similarly maximizes the observed likelihood using normalizing flow: https://arxiv.org/pdf/2003.12628 which I think could also be straightforwardly extended to generating new samples.

A: Thank you for pointing out the reference. We have added this reference to related work in lines 299-300. This work is based on normalizing flow-type generative models. We would like to briefly note that flow-based methods typically fall short in data quality compared to the GAN- and diffusion-based generative models employed in our work.

Q: Why does the imputation work? Algorithm 2 describes the imputation algorithm which is similar to the reverse process described in eq 2 using a joint score model. However, the training process only add noise to the observed coordinates, and the unobserved coordinates does not contribute to the loss (due to "odot M"), so the learned score model seems to be a "marginal score". This is different from training a conditional score model then using it to sample from a conditional probability distribution. What is the justification of authors' imputation process? Could authors please write down: What is the optimal solution of equation 5? what is the corresponding forward process and backward process that give rise to the imputation algorithm?

A: Thank you for your question. For each sample, our learning objective is to learn the "marginal score" for the particular missingness. In the population sense, our learning objective can learn the "joint score", which is proven in Theorem 3.2. We rely on the assumption that the supreme of missing rates $\rho_{\text{max}}<1$, then the optimal solution of our objective recover score on complete data. Therefore, by utilizing the existing literature that uses the score function estimated in complete data to impute data [1,2], we can verify the rationality of the imputation process.

Ref: [1]: Lugmayr, A., Danelljan, M., Romero, A., Yu, F., Timofte, R., & Gool, L.V. (2022). RePaint: Inpainting using Denoising Diffusion Probabilistic Models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 11451-11461.

[2]: Rout, L., Parulekar, A., Caramanis, C., & Shakkottai, S. (2023). A Theoretical Justification for Image Inpainting using Denoising Diffusion Probabilistic Models. ArXiv, abs/2302.01217.

Q: equation (12) in Li et al., 2019 defines the imputer of misGAN. I can just set m = [0, 0, 0... ] in (12) to generate new sample, right? The same question above. Equation (2) and (3) in Yoon et al., 2018 defines the imputer of GAIN. Can I just set m = [0, 0, 0, ...] to generate new sample straightforwardly from GAIN? If yes to question 2 and 3, what is the unique advantage of the proposed unified framework for imputing/generating samples comparing to misGAN and GAIN?

A: Thank you for your question. We can set m = [0, 0, 0... ] to use misGAN and GAIN for generating new samples. (We believe there is a small difference between misGAN and GAIN. GAIN takes addition $\tilde{x}$ and $m$ as input, which might be hard for it to generalize since during training, the network input $\tilde{x}$ and $m$ will not be all zero. The misGAN is more similar to our work that transfers the initial value $z$/$x_T$ to the output $x$/$x_0$. The smaller the missing ratio is, the smaller the discrepancy will be for misGAN and our method, while the discrepancy will be larger for GAIN.) The main advantage we claim in the paper is our proposed method does not have additional parameters to deal with missing values, while misGAN needs to train additional networks $G_m$ and $G_x$. There are two key advantages of using diffusion models to model data distributions:

Stability in Training: Diffusion models avoid adversarial training, which is a common source of instability in GAN-based models. The absence of adversarial interactions between a generator and a discriminator simplifies the training process and often leads to more stable convergence.

Modeling Uncertainty: Diffusion models inherently capture the uncertainty of generated data due to their stochastic nature. Since the generation process is governed by a stochastic differential equation, multiple generations can naturally reflect the underlying variability of the data distribution. In contrast, GAN-based methods are deterministic, lacking the ability to directly model this uncertainty.

We thank the reviewer for the comments and suggestions. We appreciate the time you spent on the paper. Below we address the concerns and comments that you have provided.

Q: This paper lacks a proper motivation for generating new samples in addition to imputing missing data. In the second paragraph in the Introduction, there is one sentence on "deep generative models can be used to... enhance the performance of image classification tasks". However, I suggest authors discuss why imputation alone is not enough, and why additional samples need to be generated and back it up with specific examples.

A: Thank you for your insightful suggestion. In the revised manuscript, we have expanded the discussion on the importance of generating additional samples alongside imputing missing data, which can be found in lines 40-42. We also provide the motivation in the following.

Two primary motivations for generating synthetic tabular data are: (1) protecting the privacy of original tabular data by releasing fake tabular data, and (2) augmenting the original tabular data with fake data for better training machine learning models [1]. These advantages have been demonstrated in recent generative models on tabular data, such as [1-4]. However, these methods are not designed to handle missing values, a common and long-standing challenge in various domains.

In contrast, most existing research on missing data focuses on imputation tasks. Although some methods are able to generate new samples, they often incur additional computational costs, due to two-stage inference processes or the need to train additional networks, as detailed in Appendix B.1.

To address these limitations, our main objective in this work is to provide a unified framework (our learning objective is the same for imputation and generation). By modeling the score of the complete data distribution, we can handle both imputation and generation tasks efficiently, by avoiding two-stage inference problems or additional network training.

Ref: [1]: Kim, J., Lee, C.E., & Park, N. (2023). STaSy: Score-based Tabular data Synthesis. International Conference on Learning Representations.

[2]: Xu, L., Skoularidou, M., Cuesta-Infante, A., & Veeramachaneni, K. (2019). Modeling Tabular data using Conditional GAN. Neural Information Processing Systems.

[3]: Kotelnikov, A., Baranchuk, D., Rubachev, I., & Babenko, A. (2023). TabDDPM: Modelling Tabular Data with Diffusion Models. International Conference on Machine Learning.

[4]: Zhang, H., Zhang, J., Srinivasan, B., Shen, Z., Qin, X., Faloutsos, C., Rangwala, H., & Karypis, G. (2024). Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space. International Conference on Learning Representations.

Q: The authors do not explain why other generator based imputation algorithm, such as GAIN or misGAN, could not be straightforwardly extended to generate new samples while imputing missing values. Both GAIN and misGAN train a generator G which outputs a complete sample. The generator is then "tweaked" to impute missing data (similar to the tweaking authors have done in line 9, algorithm 2). If they can be straightforwardly modified to generate new images, I think authors should also compare them in the experiments or at least acknowledge them as earlier unified imputation/generation algorithms.

A: Thank you for your suggestion. We provide the discussion about misGAN in Appendix B.1 in the previous version (GAN-based approaches can also be used for generation tasks, while misGAN trains two generator-discriminator pairs for the masks and data respectively, which increases the computational cost). In the revised manuscript, we have expanded the discussion of GAIN in lines 922-924.

Utility (classification accuracy) evaluation of MissDiff on Census dataset.

Thanks for the detailed reply.

For each sample, our learning objective is to learn the "marginal score" for the particular missingness. In the population sense, our learning objective can learn the "joint score", which is proven in Theorem 3.2.

1. I do not understand why the objective learn marginal score and the joint score at the same time. Do you mean, for each missing pattern, the neural network can learn a marginal score? Given 2^d possible missing patterns, it requires the neural network to have incredible amount of capacity.

2. I am still not sure why Theorem 3.2 justifies that the objective can learn the joint score? Theorem 3.2 is an upper bound of the likelihood. It is unclear how tight the bound is. Could the authors please elaborate?