Paper: MissDiff: Training Diffusion Models on Tabular Data with Missing Values (ICLR-2024)

Reviewer comment (Reproducibility):
The experimental setting for Experiment 1 should be detailed described in the main text or appendix. It is confusing to the reviewer how MissDiff’s RMSE and error rate are provided (if the MissDiff is doing unconditional generation according to Algorithm 2), and what is the training hyper parameters that are associated with MissDiff in this scenario? It is also confusing how RMSE evaluation of MIssDiff on MINIC4ED dataset is acquired (see Table 3). The authors should consider providing more experimental details.

Author response:
**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.

Reviewer comment (Evaluation):
Table 2 shows that MissDiff’s utility is significantly lower than Diff-mean for column missing scenario, the authors should have provided an explanation for the inferior performance of MissDiff for column missing Census data.

Author response:
**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.

Reviewer comment (Presentation):
In terms of writing, Algorithm 2 is not different from the existing variance preserving sampling method. Therefore, it is redundant to include Algorithm 2 in the main text. The reviewer suggest to move the Algorithm 2 into the appendix.

Author response:
**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.

Reviewer comment (Presentation):
What do the numbers mean in the parenthesis for Table 1?

Author response:
**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.

Reviewer comment (Experiments):
The reviewer is wondering how MissDiff can handle the categorical features and nonnumerical features in Tabular Data. While the categorical features and non-numerical features could be treated as continuous features according to Equation 4, the performance is not as good as directly performing discrete diffusion models (D3PM [1], tauLDR[2]). Could MissDiff be adapted to discrete diffusion models?

Author response:
**Q**: *How MissDiff can handle the categorical features and nonnumerical features in Tabular Data?*

**A**: As mentioned in the second line on page 8, we adopt the same pre/post-processing as (Kim et al., 2023; Kotelnikov et al., 2022; Zheng & Charoenphakdee, 2022) for dealing with mixed-type data. It is a good point that incorporating discrete diffusion models (D3PM, tauLDR) may achieve better performance for modeling tabular data. We believe MissDiff can be developed upon D3PM as a masked version.

Reviewer comment (Writing):
The technical description is difficult to follow, e.g. in section 2.2. or 3.2. I think it can be well understood by someone who is already familiar with the material but it doesn't do a great job building this up for the less versed reader.

Author response:
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.

Reviewer comment (Novelty):
Claiming that existing approaches are either biased or expensive seems very general.

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

Reviewer comment (Writing):
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. I may be wrong. Can you clarify?

Author response:
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.

Reviewer comment (Theory):
In Section 1.1, why did the authors claim that Mattei & Frellsen (2019); Ipsen et al. (2020b) train multiple decoders? From my understanding, they proposed multiple imputation via sampling multiple latent factors, rather than training multiple decoders.

Author response:
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.

Reviewer comment (Theory):
In the last line of Page 3, 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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
Why did not the authors compare with CSDI_T in Section 4.2? Why are the results of STaSy not included in Table 4?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Writing):
inconsistency of word "complete" : page 3, paragraph 2.1 : "complete d-dimensional data ... m=(m_1, ...,m_d) in {0,1}^d",  page 4, "framework for learning on incomplete data"

Author response:
**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).

Reviewer comment (Evaluation):
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

Author response:
**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.

Reviewer comment (Experiments):
at table 2 it seems that MissDiff and Diff-Mean seem to have similar performance -- could you comment on that and explore more datasets?

Author response:
**Q**: *At table 2 it seems that MissDiff and Diff-Mean seem to have similar performance -- could you comment on that and explore more datasets?*

**A**: Table 2 demonstrates the experimental result for the Census dataset, under a relatively simpler setting (thus an easier task) with missing ratio 0.2 and the missing mechanism being MCAR. When changing the missing mechanisms to MAR and NMAR, MissDiff performs much better than Diff-mean. We also note that MissDiff performs better for the more complicated MIMIC4ED dataset in Table 3, which contains more features (5 times than Census) and training samples (22 times than Census).