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. Below we address the concerns and comments that you have provided.

Q: Why did the authors claim that Mattei & Frellsen (2019); Ipsen et al. (2020b) train multiple decoders?

A: Thank you very much for pointing this out. We agree that Mattei & Frellsen (2019); Ipsen et al. (2020b) only train one decoder and we revise the argument in Section 1.1 accordingly in the revised paper (updated on openreview). More specifically, different from Nazabal et al. (2018); Ma et al. (2020) that training a different VAE independently to each data dimension, Mattei & Frellsen (2019); Ipsen et al. (2020b) only need to train one decoder and sample multiple latent factors from the variational distribution. However, these two works model the distribution $p(x^{\text{obs}}|z)$ by a student $t$ distribution with location, scale, and degrees of freedom outputted by the decoder, which has limited representation power for directly generating complete new samples. Therefore, their methods still need a two-stage inference framework for generating complete new samples, i.e., generating the samples containing missing values and adopting the proposed single imputation or multiple imputation methods to obtain complete data.

Q: The authors said the references require MCAR assumption. However, I suppose most of them assumed the MAR case. Moreover, one of the contributions of Li et al. (2019) is to deal with NMAR data.

A: Thank you very much for your comments. We agree that the architecture in Li et al. (2019) can be easily extended to MAR and MNAR cases. However, the theoretical guarantees only held for the MCAR case, which is the same as Yoon et al. (2018a). In Li & Marlin (2020), the authors assume MCAR but can still be unbiased if the missing mechanism is MAR. In Ipsen et al. (2022) and Mattei & Frellsen (2019), the authors require MAR assumptions. We have modified the last sentence on page 3 to M(C)AR accordingly.

Q: The authors did not mention what kind of missing mechanism was used in Section 4.1. Moreover, it seems that they only compare with baselines under one missing mechanism in this section.

A: We follow the same experimental setup as Zheng & Charoenphakdee (2022) for the imputation tasks. The missing mechanism is MCAR and the missing ratio is 0.2. We clarify the experimental setup in Appendix B.1. Due to the time limit, we were unable to finish the comparison with six baseline methods on six datasets under different missing mechanisms. We believe the comparison under the same setting with previous work can demonstrate the effectiveness of MissDiff for imputation tasks.

Q: Why did not the authors compare with CSDI_T in Section 4.2? Why are the results of STaSy not included in Table 4?

A: Thank you very much for your suggestion. We have added the comparison with CSDI_T in Table 2, 4, 7, 9 and STaSy-delete and STaSy-mean in Table 4, 9. MissDiff outperforms CSDI_T on 8 out of 10 results and outperforms STaSy-delete and STaSy-mean by a large margin.

Q: While Tashiro et al. (2021) adopted conditional score matching, the work goes back to unconditional score matching. Can the authors provide more discussions about these two approaches and what are the advantages of using unconditional scores?

A: Thank you very much for raising this important question. First of all, there are various conditional scores (depending on which information is conditioned on), making them difficult to learn and analyze. In Tashiro et al. (2021) and its tabular variant CSDI_T, there were no theoretical guarantees on whether the learned conditional score satisfied the optimality condition. In our work, we provide theoretical guarantees on the unconditional score learning from incomplete data. Secondly, conditional score matching performs better in time series imputation tasks than unconditional score matching, which is not necessarily the case for tabular data. We believe that is because there may exist some complex or irregular dependencies between different columns in tabular data, e.g., some features might be redundant (uninformative). Therefore, MissDiff can achieve better performance than CSDI_T in the imputation tasks and generation tasks.

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

Q: How do you tune hyperparameters, especially for competing methods?

A: The hyperparameters of our method can be found in Appendix B.4 (revised version updated on openreview). In short, the optimizer, learning rate, batch size, and training epochs are the same as Zheng & Charoenphakdee (2022) for fair comparison. The hyperparameters of other baseline methods can also be found in Appendix B.4.

Q: I don't think that work like the VAEs for missing data can be described as "generate-than-impute". In (at least some) of these models, there is no separate imputation step, but they learn a joint density marginalizing over missing value.

A: Thank you very much for your question. We clarify this point below. VAE-based methods (especially Mattei Frellsen (2019)) maximize the likelihood of the observed value to train an encoder projecting $x^{\text{obs}}$ to feature $z$ and a decoder projecting feature $z$ to $x^{\text{obs}}$. For the imputation tasks, they learn the distribution $p(x^{\text{missing}}|x^{\text{obs}})$ by marginalizing over latent variable $z$. However, the previous VAE-based method cannot easily deal with the generation tasks without adopting "generate-then-impute" framework. The key reason is they model the distribution $p(x^{\text{obs}}|z)$ by a student $t$ distribution with location, scale, and degrees of freedom outputted by the decoder, which has limited representation power for the real distribution. We believe directly sampling from $p(x^{\text{obs}}|z)$ for generating tasks will have poor performance. Therefore, a practical solution for generating is first generating a set of missing data and then adopting VAE-based methods to impute them.

Q: Claiming that existing approaches are either biased or expensive seems very general.

A: We provide a detailed explanation of this claim in Section 3.1 and Section 1.1.

Q: The technical description is difficult to follow; typos; Stylistic critique

A: Thank you very much for your suggestion. The typos are carefully checked and corrected. To improve the readability of Section 2.2 and 3.2, we have added specific references to the main technical equations, including the forward and backward process, the score matching objective, and variance preserving SDE.

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

Q: The experimental setting for Experiment 1 should be detailed described in the main text or appendix.

A: Thank you for your suggestion. The experimental setting for Section 4.1 can be found in Appendix B.1 in the revised paper (updated on openrview).

Q: It is confusing to the reviewer how MissDiff’s RMSE and error rate are provided ... what is the training hyper parameters that are associated with MissDiff in this scenario?

A: Thank you very much for your question. We have added the imputation algorithm of MissDiff in Algorithm 2, which can be found in Appendix B.3. Since MissDiff models the score for complete data distribution, we can adopt the same model and training hyperparameters for imputation and generation tasks. The hyperparameters are described in Appendix B.4.

Q: It is also confusing how RMSE evaluation of MIssDiff on MINIC4ED dataset is acquired (see Table 3).

A: The RMSE in Table 3 is different from imputation tasks. As mentioned in "Evaluation Criterion" paragraph in section 4.2, we follow the evaluation criterion as Xu et al. (2019); Kim et al. (2023); Kotelnikov et al. (2022) that use {\it utility} to evaluate the performance of the generated complete data. We train a downstream model (e.g., XGBoost model) on generated data for regression tasks and report the RMSE of the predicted value against the oracle value on the real (test) data.

Q: Table 2 shows that MissDiff’s utility is significantly lower than Diff-mean for column missing scenario.

A: We believe the column missing mechanism described in Appendix B.1 is actually a special scenario. Most specifically, the mask $\mathbf m$ (indicator of missing values) for each row (sample) would depend on the masks of other rows as well, since the missing rate for each column is fixed. It leads to dependence between missing samples. We further note that in our population objective function eq (4), as a standard practice, we regard the sample pair (m,x) are iid and the expectation in (4) is taken with respect to this joint distribution. When the sample size of the dataset is relatively small, such sample dependence is more evident, and MissDiff is not as good as Diff-mean. However, when there is a sufficiently large number of samples, as in Table 3,8, and 11 for the MIMIC4ED dataset (22 times larger than the Census dataset), such dependence becomes very weak and the joint sample pairs (m,x) are close to independent; in these cases, MissDiff performs better.

Q: It is redundant to include Algorithm 2 in the main text.

A: Thank you for your suggestion. We have moved Algorithm 2 to Appendix B.3. For completeness, we have also added the algorithm for imputation in Appendix B.3.

Q: What do the numbers mean in the parenthesis for Table 1?

A: The numbers are the standard deviation of five independent trials. We add the explanation to the caption of Table 1.

Q: How missing values are defined in videos and images? How to evaluate MissDiff with more complicated data types such as video or language data?

A: Thank you very much for your question.

• For video and image data, we can either regard the missing frames or the missing pixels in each frame as the missing data. It is a good point of view that the amount of redundant information in video and images is high. We think MissDiff would still be useful against VAE-based or GAN-based methods since our approach can learn the complete data distribution from incomplete data (low-quality video data) and generate complete data in a unified framework.
• For the language data, missing values could be missing words or sentences. We would like to emphasize that missing values in language data are typically more challenging to deal with as compared with the tabular case we considered: (1) there are strong temporal dependencies in language data, (2) the length of the language data is unfixed.

As a special and simple example, our method can be extended for tabular data that contains text information (for instance, the diagnosis note from the doctor). Some of such descriptions might be missing. A language embedding model could potentially be incorporated into our proposed method to deal with the missingness of such short-length text information in certain columns. Nevertheless, this would need careful design case by case.

How do you input the missing x^{obs} into the score network whose input dimension is fixed, or taking 0 values at missing position? Thanks

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: Inconsistency of word "complete".

A: Thank you for your comments. Our proposed framework is learning from incomplete data and generating complete data. Therefore, in Section 2.1, we first define the complete (unobserved) data distribution $p_0(x)$ and then define the incomplete (observed) data samples $S^{\text {obs}}$. We have made this more clear in the revision (updated on openreview).

Q: More metrics need to be presented at the experiments section, eg RMSE is not necessarily representative in high-dimensional cases were the distribution is complex-multimodal.

A: Thank you for your good suggestion. For the generation task, we adopt two types of criteria, fidelity and utility, to evaluate the quality of the synthetic data generated. For the imputation task, we adopted RMSE in order to compare with previous works that evaluated the imputation performance mainly by RMSE.

Q: Could you provide some insights on why the method's performance improves as missingness rate increases in (0.1-0.6)?

A: For the slight performance improvement with missing rate increase in (0.1-0.6), it is the authors' conjecture that this is a phenomenon due to the unique structure of certain tabular datasets. For this simulated Bayesian network dataset, the dependencies between different columns are demonstrated in Figure 2: some features might be uninformative, for instance, variables C1, C2, and D1 are all uninformative to the value of D3, given that D2 is observed. This implies that for some rows with missing C1, C2, and D3 values, the model still has enough information to learn the full dependence between variables D3 and D2. Moreover, the model can potentially learn the distribution of D3|D2 better in such cases since other redundant variables are excluded. We conjecture this is the reason why we see a little improvement in the performance.

Nevertheless, we would like to emphasize that the improvement in performance is very small, after considering the observation noise, the authors believe it is more appropriate to conclude that the performance of the proposed method remains more ${\it stable}$ (as compared with other baselines) for missing rates in 0.1-0.6. Moreover, the performance starts to decrease when we increase the missing rate to 0.8, since in such case, we only have one variable left in each row and thus it is reasonable to expect worse performance.

Q: Could you add the case where missingness rate is 90% in the first case?

A: As mentioned in Appendix B.1 (which is Appendix B.2 in the revised version), there are only five variables (columns) in the data generated by Bayesian Network (three categorical variables and two continuous variables). Therefore, in the row missing mechanism, we only have the missing ratio is [0.2,0.4,0.6,0.8]. For the column missing or the independent missing mechanisms, we set the missing ratio to be [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9].