Rethinking the Diffusion Models for Missing Data Imputation: A Gradient Flow Perspective

Thank you for your comments, our rebuttal organized point to point is posted as follows:

W1 & Q2: Why we use RBF Kernel Function:

• The selection of the RBF kernel function was strategically driven by the need to satisfy the following condition: $\int{\nabla_{\boldsymbol{X}^{(miss)}}[u(\boldsymbol{X}^{(miss)} ,\tau) r(\boldsymbol{X}^{(miss)})]\mathrm{d}\boldsymbol{X}^{(miss)}}=\underbrace{\int{r(\boldsymbol{X}^{(miss)})\nabla\_{\boldsymbol{X}^{(miss)}}[u(\boldsymbol{X}^{(miss)} ,\tau) ]\mathrm{d}\boldsymbol{X}^{(miss)}}}\_{\mathbb{E}\_{r(\boldsymbol{X}^{(miss)})}[ \nabla\_{\boldsymbol{X}^{(miss)}}[u(\boldsymbol{X}^{(miss)} ,\tau) ] ]} + \int{u(\boldsymbol{X}^{(miss)} ,\tau)^\top\nabla_{\boldsymbol{X}^{(miss)}}[r(\boldsymbol{X}^{(miss)}) ]\mathrm{d}\boldsymbol{X}^{(miss)}} =0$. This condition is pivotal as it allows us to circumvent the direct, explicit estimation of $r(\boldsymbol{X}^{(miss)})$ during the imputation procedure.

• A sufficient condition for $\int{\nabla_{\boldsymbol{X}^{(miss)}}[u(\boldsymbol{X}^{(miss)} ,\tau) r(\boldsymbol{X}^{(miss)})]\mathrm{d}\boldsymbol{X}^{(miss)}}=0$ is that "$r(\boldsymbol{X}^{(miss)})$ is bounded and $\lim_{\Vert \boldsymbol{X}^{(miss)} \Vert\rightarrow \infty}u(\boldsymbol{X}^{(miss)} ,\tau)=0$".

• Consequently, the key to choosing kernel function is to validate the condition as follows: For $x$, kernel function $\mathcal{K}(x,x')$ should satisfy the boundary condition: $\lim_{\Vert x \Vert\rightarrow \infty}\mathcal{K}(x,x')=0$.

• Conventional kernel functions, like linear kernel, sine kernel, polynomial kernel, cosine similarity kernel... cannot satisfy this condition. Thus we choose the RBF kernel in our manuscript similar to previous work [1].

• Furthermore, we propose the following experimental results at 0.3 missing rate (to consist with the table in main context, results for MNAR scenario are not posted due to space limit): |Scenario|Kernel|BT-MAE|BT-Wass|BCD-MAE|BCD-Wass|CC-MAE|CC-Wass|CBV-MAE|CBV-Wass|IS-MAE|IS-Wass|PK-MAE|PK-Wass|QB-MAE|QB-Wass|WQW-MAE|WQW-Wass| |-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-| |MAR|linear|0.77|0.73|0.83*|3.36*|0.85*|0.77*|0.83*|0.97*|0.72*|2.42*|0.8*|2.5*|0.7*|4.58*|0.76*|1.02*| ||polynomial|1.12*|1.4*|1.44*|8.36*|1.06*|1.16*|1.08*|1.55*|1.07*|5.93*|1.33*|5.93*|1.28*|13.63*|1.02*|1.76*| ||sigmoid|0.89*|1.02*|1.44*|10.8*|0.82|0.74*|0.81*|0.93*|1.17*|12.57|1.3*|9.87*|1.04*|7.42*|0.77*|1.04*| ||cosine similarity|0.68|0.53|0.8*|3.14*|0.82*|0.73*|0.81*|0.92*|0.71*|2.37*|0.78*|2.37*|0.74|4.56|0.74*|0.98*| ||sine|8.62*|95.45|8.74*|290.79*|14.83*|281.68*|17.93*|492.97*|10.56*|456.27*|10.25*|306.76*|8*|313.9*|15.14*|386.73*| ||RBF|0.52|0.38|0.34|0.82|0.35|0.25|0.31|0.2|0.39|1.31|0.44|1.21|0.45|3.5|0.46|0.55| |MCAR|linear|0.71*|0.45*|0.87*|5.87*|0.83*|0.84*|0.81*|1.28*|0.81*|5.65*|0.83*|4.14*|0.62*|6.31*|0.76*|1.25*| ||polynomial|0.94*|0.88*|1.31*|12.19*|0.98*|1.21*|0.99*|1.85*|1.11*|9.94*|1.27*|8.64*|1.11*|18.71*|0.92*|1.95*| ||sigmoid|0.74*|0.48*|0.97*|8*|0.81*|0.82*|0.8*|1.23*|0.92*|7.36*|0.94*|6.54*|0.77*|7.74*|0.77*|1.28*| ||cosine similarity|0.7*|0.42*|0.84*|5.51*|0.81*|0.81*|0.8*|1.24*|0.8*|5.48*|0.81*|3.96*|0.63*|6.01*|0.74*|1.2*| ||sine|9.76*|95*|7.77*|434.72*|13.53*|331.59*|11.78*|352.67*|8.27*|542.27*|8.62*|384.97*|7.36*|468.07*|10.61*|289.41*| ||RBF|0.48|0.18|0.25|0.8|0.47|0.34|0.42|0.44|0.44|3.05|0.32|1.01|0.34|3.66|0.53|0.76|

  From the results presented in the table, it is evident that the RBF kernel outperforms other kernels in scenarios where the boundary condition is not met. This observation empirically substantiates the effectiveness of the RBF kernel, thereby validating the justification for using RBF kernel.

• We decide to add the abovementioned contents to explain why we mainly consider RBF kernel in our revised manuscript.

W2 & Q1: Data Distribution Property

Thank you for your comments, our responses are listed as follows:

• We have expanded the experimental results in the attached PDF to include KnewImp performance across various data distributions, such as heavy-tailed, multimodal, and skewed distributions. Please refer to Section 2.2 in the supplementary PDF provided in the common rebuttal chat window.
• The additional results demonstrate that KnewImp maintains consistent performance when transitioning from Gaussian to Skewed-Gaussian, Student's-t, and Gaussian Mixture distributions. This underscores the robustness of KnewImp across diverse data distributions.
• We will incorporate these findings into the revised manuscript to enhance its comprehensiveness and clarity.

Thank you for reading our response, we hope the above discussion fully addressed your concerns about our work, and we would really appreciate it if you could be generous in raising your score.

Refs
[1]. "Stein variational gradient descent: A general purpose bayesian inference algorithm." NeurIPS'16.

Dear Reviewer [9XXU],

Thank you very much for your supportive feedback and for taking the time to thoroughly review both our manuscript and the additional materials provided during the rebuttal process. We are grateful for your decision to increase your score and appreciate your endorsement of our paper.

Warm regards,
Authors

Thank you for your comments, our rebuttal is posted as follows:

W1: Downstream Tasks

According to your suggestions, we added the performance on downstream classification task similar to Fig. 5 in the TDM paper [1]. Please see Table 1 located in Section 2.1 from the pdf attached in common rebuttal chat window.

W2: Ablation Study

Thank you for your feedback on our ablation study.

• There seems to be some misunderstanding about the contributions of KnewImp. Our method primarily introduces WGF to analyze and enhance DM-based MDI. This introduces two main innovations: NER and Joint Distribution Modeling.
• The function of NER term is that it introduces a functionally effective modification to the existing DM -based MD—Optimizing the NER functional aligns well with optimizing the joint/conditional log-likelihood objective.
• Based on this, the ablation study aimed to show that including NER maintains an effective lower bound for MDI tasks—meaning it does not degrade model performance in general.
• Besides, we can also turn to Table E.4, when we add NER term into the model, the standard deviation of KnewImp with NER is generally smaller than those without NER term. These results further confirm that NER, although providing marginal performance gains, while enabling the KnewImp has a smaller standard deviation, is crucial for the model's theoretical robustness.

Q1: Motivation

• To clarify the trade-off between diversification and accuracy, which is a central theme of our method KnewImp, let's consider a practical example. Suppose we denote the true value by $x$ and the imputed value by $\tilde{x}$. The goal in terms of accuracy is to minimize the discrepancy $Dis(x, \tilde{x})$, where $Dis$ is discrepancy metric.
• Diversification tends to increase either the variance $\text{Var}(\tilde{x})$ or the entropy $\mathbb{H}(\tilde{x})$. These measures do not directly involve ground truth $x$.
• In DMs, where entropy is used as a term to encourage diversification (our proposition 3.1), the outcome may converge towards a uniform distribution [2]. In such a distribution, every potential value within the support is equally probable as the imputed value, which is undesirable for MDI.

Q2: Scenario Restriction

Thank you for your question regarding our focus on MAR and MCAR rather than MNAR.

• We prioritize MAR and MCAR because MNAR involves complexities due to its dependency on unobserved data, requiring in-depth knowledge of the missingness mechanism which is often challenging to determine [3,4].
• For instance, in privacy-sensitive studies like diabetes medication usage, the non-response is directly linked to the privacy nature of the participants, typical of MNAR scenarios.
• Given these challenges, we focus on MAR and MCAR for their more straightforward assumptions and applicability.
• Nevertheless, to ensure a comprehensive analysis, we include findings related to MNAR in Tables E.1 and E.3 of our manuscript, demonstrating our approach's applicability under these conditions as well.

Q3: Marking Mistake

Thank you for pointing out our problems, we will revise this table in our revised manuscript.

Thank you for reading our rebuttal, we hope our response alleviates your problem and we would really appreciate it if you could be generous in raising your score.

Refs
[1]. "Transformed distribution matching for missing value imputation." ICML'23.
[2]. "Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models" ICML'19
[3]. "StableDR: Stabilized doubly robust learning for recommendation on data missing not at random." ICLR'23
[4]. "not-MIWAE: Deep generative modelling with missing not at random data" ICLR'21

Dear Reviewer [fS2D]:

Thank you for your encouraging feedback and for raising your score, which greatly supports our work.

Regarding your additional inquiries:

• Application to Regression Tasks:

  • Currently, our focus has been primarily on classification tasks, as suggested by the framework outlined in Figure 5 of the TDM paper [1], primarily due to our time and resource constraints.
  • However, we recognize the importance of exploring the utility of KnewImp in regression scenarios. We aim to include preliminary results on this aspect in another comment chat window within the discussion period set by the NeurIPS'24 committee as much as possible.
  • Additionally, we will ensure to incorporate a detailed evaluation on downstream regression tasks in the revised version of our manuscript.

• Performance on Real-World Datasets with Inherent Missing Values:

  • Currently, due to platform constraints, we have not yet implemented our algorithm in real-world industrial settings, such as recommender systems [2], where various metrics foucus on business value, area under the curve, return on investment [3], are typically employed.
  • However, we are actively seeking opportunities to apply our methodology in industrial scenarios and plan to explore this in future work.
  • We will outline these plans and the potential for real-world applications in the future research directions section of our revised manuscript.

We appreciate your insightful questions, which guide our ongoing and future research efforts.

Refs
[1]. "Transformed distribution matching for missing value imputation." ICML'23.
[2]. "StableDR: Stabilized doubly robust learning for recommendation on data missing not at random." ICLR'23
[3]. "Kalman Filtering Attention for User Behavior Modeling in CTR Prediction" NeurIPS'20

Sincerely,
Authors

Dear Reviewer [fS2D]:

Thank you for your insightful suggestions regarding the downstream regression task. In response to your feedback, we have conducted additional experiments focusing on downstream regression task. We utilized Mean Square Error (MSE) and Mean Absolute Error (MAE) as our evaluation metrics according to references [1] on the CC dataset, specifically designed for regression analyses. The results for missing rate at 0.3 are posted as follows:

Refs:
[1]. "Deep Time Series Models: A Comprehensive Survey and Benchmark"

Thank you for reading our comments, we hope our response answer your problem. Given your busy schedule, please do not feel obliged to respond to this message.

Sincerely,

Authors

Thank you for your comments.

W1: Contributions & Novelty of This Paper

• Our primary focus is on analyzing Diffusion Model (DM)-based Missing Data Imputation (MDI) using Wasserstein Gradient Flow (WGF, initially designed for functional optimization), not merely integrating WGF into a generative model.
• The central contribution of this work is leveraging WGF as a tool to analyze and demonstrate the limitations of Diffusion Model-based MDI. This led us to redesign a novel functional and develop a unique computational approach for MDI.
• To the best of our knowledge, Proposition 3.1 has not been documented in previous literature. Specifically, for the VP-SDE model, the functional of the Ornstein–Uhlenbeck process incorporates a variance term $\mathbb{E}_{r}[\boldsymbol{X}^{(miss)}]^\top[\boldsymbol{X}^{(miss)}]$ to promote diversification. This feature is notably absent in traditional gradient field-based (GF-based) generative models as referenced in [1-6].
• In our manuscript, we want to convey the concept that: MDI should not be treated as a generative problem (in our manuscript, we want to discourage diversification, but the generative models will encourage entropy function as per reference [1]: The differential entropy term $\mathbb{H}(r)$ improves diversity and expressiveness when the gradient flow is simulated for finite time-steps. Similar entropy terms can be found in Eq. (15) in [1], Eq. (16) in [2], Eq. (3) in [3], Eq (1) in [4], Eq. (4) in [5], and Eq. (5) in [6]).
• Hence the references [1-6] you recommended about using GF to improve generative models were not considered during the preparation of our manuscript, and our related works subsections located in Section 5 are chiefly organized from the perspective of DM's application in MDI and the application of WGF works which model the conditional distribution by joint distribution.
• We plan to add another subsection to discuss the application of WGF in improving generative models to demonstrate the applicability of GF to improve generative model-related tasks and cite these references.

W2: Justification of Decomposition:

• As for the decomposition, it is not a strong assumption for MDI.
• As we mentioned in our manuscript and common rebuttal chat window, our task is to find the missing value $\boldsymbol{X}^{{(miss)}}$ and the $\boldsymbol{X}^{{(obs)}}$ will not be changed ($p(\boldsymbol{X}^{{(obs)}})$ is a constant measure).
• Based on this, when we want to sample $\boldsymbol{X}^{{(miss)}}$ from distribution $r({\boldsymbol{X}^{{(miss)}}})$, the results are the same as sampling from $r({\boldsymbol{X}^{{(joint)}}})$ since $r(\boldsymbol{X}^{{(obs)}})=p(\boldsymbol{X}^{{(obs)}})$ is a constant measure according to reference [7,8], which is unchanged.
• The key is that $p({\boldsymbol{X}^{{(joint)}}})\neq p({\boldsymbol{X}^{{(miss)}}}) p({\boldsymbol{X}^{{(obs)}}})$, but $r({\boldsymbol{X}^{{(joint)}}})= r({\boldsymbol{X}^{{(miss)}}}) p({\boldsymbol{X}^{{(obs)}}})$ is justified.
• Furthermore, similar assumption on the $r$, the 'ansatz'/variational distribution/approximate distribution/proposal distribution, can be found in mean-filed variational inference represented by reference [9].

L1: Typos on Eq. (B.2):

We regret any oversight regarding the omission of the $\nabla\cdot$ operator (divergence operator), and velocity term $v_\tau$ in our description of the continuity equation.

• This equation is given as $\frac{\partial r(\boldsymbol{X}^{{(miss)}})}{\partial \tau}=-\nabla\cdot[v_{\tau}(\boldsymbol{X}^{{(miss)}})r(\boldsymbol{X}^{{(miss)}})]$, the second equality can be obtained similar to the derivation of Fokker-Planck-Kolmogorov equation, where $v_{\tau}(\boldsymbol{X}^{{(miss)}})=-\nabla{\log{p(\boldsymbol{X}^{{(miss)}}\vert \boldsymbol{X}^{{(obs)}})}}-\lambda\nabla{\log{r(\boldsymbol{X}^{{(miss)}})}}$ and $\nabla\cdot[(\nabla{\log{r(\boldsymbol{X}^{{(miss)}})}})r(\boldsymbol{X}^{{(miss)}})] =\nabla\cdot[\frac{\nabla {{r(\boldsymbol{X}^{{(miss)}})}})}{r(\boldsymbol{X}^{{(miss)}})}r(\boldsymbol{X}^{{(miss)}}] =\nabla\cdot\nabla r(\boldsymbol{X}^{{(miss)}})$.
• We commit to revising our manuscript to rectify any ambiguities and ensure clearer presentation.

Thank you for reading our rebuttal. Given the above infos, we hope that these points could be kindly considered in the evaluation of our work, and we would really appreciate it if you could be generous in raising your score.

Refs
[1]. "Refining deep generative models via discriminator gradient flow." ICLR'21
[2]. "Convergence of flow-based generative models via proximal gradient descent in Wasserstein space." IEEE TIT
[3]. "Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport." ICML'24
[4]. "Deep generative learning via variational gradient flow". ICML'19
[5]. "Deep generative Wasserstein gradient flows"
[6]. "Normalizing flow neural networks by JKO scheme" NeurIPS'23
[7]. "Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel." ICLR'24.
[8]. "Nonparametric generative modeling with conditional sliced-Wasserstein flows." ICML'23.
[9]. "Variational algorithms for approximate Bayesian inference", Doctoral Thesis'03.

Dear Reviewer [833r]:
Thank you for your response and further inquiries. Here is a refined version of our rebuttal addressing the concerns raised:

1. Mean-Field Assumption: The mean-field approximation described on Page 52 of reference [1], titled The Mean Field Approximation, involves conditionally independent factorization of the approximation distribution, updated iteratively. This approach validates our factorization of $r$, aligning with established variational methods where each factor is updated independently. This operation has also been included in the text book [2] in Page 464 to 465, for variational inference. (again, $r(\boldsymbol{X}^{(obs)})$ is a constant, since $\boldsymbol{X}^{(obs)}$ remains unchange)

2. Modeling Strategy Justification: We have to emphasize that, we have not used posterior distribution in our manuscript. Thus, we think you may doubt that how we can simulate the conditional distribution modeling by joint distribution modeling about distribution $p$ (not $r$):

   • Reference [3] provides empirical validation in Appendix C of this strategy.
   • Discrepancy measurements between conditional and joint distributions are thoroughly examined in reference [4] including $f$-divergence, where supplementary material includes detailed derivations supporting our approach.
   • Reference [5] further substantiates that our modeling strategy of conditional-by-joint can be effectively applied within the WGF framework, as detailed in Remark 7 and Theorem 11 of the paper.
   • We elaborate on these validations in Section 5.2 of our manuscript (Lines 319 to 320) and provide detailed proofs on Lines 658 to 688. More specifically, our approach utilizes a discrepancy metric akin to Kullback-Leibler (KL) divergence, expressed as $-\int r(x) \frac{\log r(x)}{p(x)} \mathrm{d}x$. Unlike the traditional use of KL divergence which incorporates diversification-encouraging positive entropy $\mathbb{H}[r(x)]$, our study employs diversification-discouraging negative entropy $-\mathbb{H}[r(x)]$. We think our manuscript further extend this modeling strategy in a way theoretically.

3. Extra Example why $r$'s factorization is justified: To further clarify, let's consider an example from SVGD, which may help elucidate why, in WGF-based approaches that approximate $p(z|x)$ with $q(z)$ (represented by $r$ in our study), the process does not explicitly involve the input $x$ or the posterior $p(z|x)$.

   • Refer to Figure 4 in reference [6] for an illustrative example (it can approximate the ground truth posterior without explicitly computing posterior). In SVGD, $q(z)$ is represented as a group of particles. These particles are not directly influenced by the input $x$ or the explicit form of the posterior $p(z|x)$; instead, they are guided by the velocity field determined by the evidence $p(x|z)$ and the prior $p(z)$, with $p(x)$ effectively being 0 under the gradient operator.
   • This setup allows $q(z)$ to approximate $p(z|x)$ effectively with the help of velocity filed $v$. Similarly, in our case with $r$, there is no necessity to consider it as an input of $\boldsymbol{X}^{(obs)}$.
   • Instead, $r$ is guided by the velocity field determined by $p(\boldsymbol{X}^{(miss)}, \boldsymbol{X}^{(obs)})$, which inherently includes observation information (recall $\frac{\partial r}{\partial \tau}=-\nabla\cdot(rv)$, where $r(\boldsymbol{X}^{(miss)})$ is shaped by velocity filed $v$, $r(\boldsymbol{X}^{(obs)})$ remains unchange, the velocity term for $r(\boldsymbol{X}^{(obs)})$ is zero. Notably, and $v=-\nabla_{\boldsymbol{X}^{(miss)}}\frac{\delta \mathcal{F}\_{joint-NER}}{\delta r(\boldsymbol{X}^{(miss)})}$, the term $\mathcal{F}\_{joint-NER}$ contains information about $\boldsymbol{X}^{(obs)}$).

References:
[1] "Variational algorithms for approximate Bayesian inference", Doctoral Thesis '03.
[2] "Pattern Recognition and Machine Learning", Text Book'06
[3] "Nonparametric generative modeling with conditional sliced-Wasserstein flows", ICML '23.
[4] "Conditional Wasserstein Generator", IEEE TPAMI '23.
[5] "Posterior sampling based on gradient flows of the MMD with negative distance kernel", ICLR '24.
[6] "VAE learning via Stein variational gradient descent", NeurIPS'17.

We appreciate your detailed feedback and look forward to further discussions.

Best regards, Authors

Dear authors,

Thank you so much for your detailed rebuttal and global response. However, regarding the assumption on the distribution r, the reviewer still finds the explanation offered by the authors to be a bit untransparent. Would it be possible for the authors to write the derivations in terms of conditional probability and posterior distributions? Alternatively, would it be fine for the authors to specify which lemmas/theorems in the papers [1,2,3] lead to the desired claims?

Thanks in advance!

References:

[1] Hagemann, P., Hertrich, J., Altekrüger, F., Beinert, R., Chemseddine, J., & Steidl, G. (2023). Posterior sampling based on gradient flows of the MMD with negative distance kernel. arXiv preprint arXiv:2310.03054.

[2] Du, C., Li, T., Pang, T., Yan, S., & Lin, M. (2023). Nonparametric generative modeling with conditional sliced-Wasserstein flows. arXiv preprint arXiv:2305.02164.

[3] Beal, M. J. (2003). Variational algorithms for approximate Bayesian inference. University of London, University College London (United Kingdom).

Best regards,

Reviewer 833r

Before reading our response, we should come to following agreements, which are the settings of MDI task:

• Throughout the imputation procedure, $\boldsymbol{X}^{({obs})}$ remains invariant regardless of any modifications to $\boldsymbol{X}^{(miss)}$.
• Given this invariance, it is accurate to state that $r(\boldsymbol{X}^{({obs})})$ is constant, and consequently, $r(\boldsymbol{X}^{(obs)}|\boldsymbol{X}^{(miss)}) = r(\boldsymbol{X}^{(obs)})$, reflecting the independence of $\boldsymbol{X}^{(obs)}$ from $\boldsymbol{X}^{(miss)}$.

Based on this, according to the requirement of Reviewers [833r] and [EhUg], we should factorize the $r$ as $r(\boldsymbol{X}^{(joint)})=r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})=r(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)}) r(\boldsymbol{X}^{(obs)})$. Now let's analyze the left-hand-side of continuity equation: $\frac{\partial r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})}{\partial \tau}$.

• First, we can get: $\frac{\partial r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})}{\partial \tau} =\frac{\partial r(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)}) r(\boldsymbol{X}^{(obs)}) }{\partial \tau}$,
• Then, expand the product on the right-hand-side: $\underbrace{r(\boldsymbol{X}^{(obs)})\frac{\partial r(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)}) }{\partial \tau} }\_{ r(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)}) = \frac{r(\boldsymbol{X}^{(obs)}|\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(miss)})}{r(\boldsymbol{X}^{(obs)})} }+ \underbrace{r(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)}) \frac{\partial r(\boldsymbol{X}^{(obs)}) }{\partial \tau} }\_{0}$, where the first underbrace is the Bayesian formula, the second underbrace indicates that $r(\boldsymbol{X}^{(obs)})$ remains unchanged.
• Now expand the first underbrace, we get $\underbrace{\frac{r(\boldsymbol{X}^{(obs)})}{r(\boldsymbol{X}^{(obs)})}\frac{\partial r(\boldsymbol{X}^{(obs)}|\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(miss)}) }{\partial \tau}}_{r(\boldsymbol{X}^{(obs)}|\boldsymbol{X}^{(miss)}) = r(\boldsymbol{X}^{(obs)}) }$. The first underbrace is based on the abovementioned agreement.
• Finally, we get: $r(\boldsymbol{X}^{(obs)}) \frac{\partial r(\boldsymbol{X}^{(miss)})}{\partial \tau}$. i.e. $\frac{\partial r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})}{\partial \tau} = r(\boldsymbol{X}^{(obs)})\frac{\partial r(\boldsymbol{X}^{(miss)})}{\partial \tau}$. Notably, the factorization $r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})=r(\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(obs)})$ can also have the same result given that $r(\boldsymbol{X}^{(obs)})$ is a constant according to the agreement.

In summary, we can see that within WGF, the operation is $r(\boldsymbol{X}^{(obs)},\boldsymbol{X}^{(miss)})=r(\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(obs)})$ is justified for MDI task. The influence of $\boldsymbol{X}^{(obs)}$ is hidden in the velocity filed $v$, where the continuity equation: $\frac{\partial r(\boldsymbol{X}^{(miss)})}{\partial \tau}=-\nabla\cdot[ r(\boldsymbol{X}^{(miss)})v]$ shapes the "actor" $r(\boldsymbol{X}^{(miss)})$'s performance by the "comments" $v$ (the velocity field), given by the "critic" $p(\boldsymbol{X}^{(miss)}|\boldsymbol{X}^{(obs)})$.

We plan to add the abovementioned derivation in our revised manuscript to increase uphold the rigor of our manuscript.

We hope the above discussion will fully address your concerns about our work, and we would really appreciate it if these responses could meet with your approval. We look forward to your insightful and constructive responses to improve this work. Thank you very much!

Dear Reviewer [833r]:

Thank you for your constructive comments, which have significantly contributed to enhancing our manuscript. In response to your concerns, we would like to outline and summary how we have addressed your problem each point in detailed:

Universality of the Mean-Field Assumption on $r$'s Decomposition would it be fine for the authors to specify which lemmas/theorems in the papers [1,2,3] lead to the desired claims?

• Following your query, we have referenced specific pages in the literature to demonstrate how this assumption is commonly applied in general mean-field variational inference approaches.
• In light of your request, we have provided detailed explanations in the specific pages on the evolution of the strategy that model the joint distribution using conditional distribution, where they treat the velocity filed of the observation part "vanishes", which supports the justification of this assumption within WGF framework.

Theoretical Justification of $r$'s Decomposition within WGF framework Would it be possible for the authors to write the derivations in terms of conditional probability and posterior distributions?

• We have included a comprehensive, step-by-step derivation to substantiate the mean-field assumption of $r$ within the WGF framework, as prompted by your insightful query.

Finally, we would like to conclude with a metaphor to further illustrate the plausibility of this factorization $r(\boldsymbol{X}^{(joint)})=r(\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(obs)})$:

• Consider $r$ as an actor in a play, capable of being molded and shaped. Initially, the actor may not fully embody the role, akin to $r(\boldsymbol{X}^{(miss)})$ not containing information about $\boldsymbol{X}^{(obs)}$.
• However, just as a director shapes an actor's performance through guidance and rehearsal, all we need to do is ensure that $r$ is appropriately molded by the directorial guidance (mirrors the continuity equation $\frac{\partial r}{\partial \tau}=-\nabla\cdot(vr)$) of the velocity field $v$ and the script provided by the critic $p(\boldsymbol{X}^{(obs)} | \boldsymbol{X}^{(miss)})$/$p(\boldsymbol{X}^{(obs)} , \boldsymbol{X}^{(miss)})$.
• As long as $r$ can adapt based on this feedback (akin to the WGF framework), it can overcome the limitations of its initial portrayal (akin to $r(\boldsymbol{X}^{(joint)})=r(\boldsymbol{X}^{(miss)})r(\boldsymbol{X}^{(obs)})$).

We remain committed to providing rigorous revisions following your suggestions irrespective of what decision you make. Given your busy schedule, please do not feel obliged to respond to this message.

Warm regards,
Authors

Dear authors,

Thank you for the detailed response. I will take a further look.

Best regards,

Reviewer 833r

Dear Reviewer [833r],

Thank you very much for your encouraging feedback and for acknowledging the clarifications provided in our response. We are grateful for your decision to increase the score and support the publication of our paper.

We also appreciate your attention to detail and your advice regarding the typos in the appendix. We assure you that we will rigorously review the manuscript again to correct all typographical errors and enhance its readability.

Thank you once again for your insightful contributions to the refinement of our work.

Best regards,
Authors

I would like to thank the authors for their detailed response, which have resolved almost all questions. I now think it is fine to support the paper to be published, so I will be increasing my score from 3 to 5. However, please make sure to correct all typos you can find in the appendix to make the whole manuscript more readable.

Thank you for your comments. Before reading our response, we think we should come to the following agreements:

1. Optimizing the instances $x$ from distribution $r(x)$ is optimizing this distribution $r(x)$, which is the basis of particle-based variational inference like Stein Variational Gradient Descent [1].
2. The velocity filed $v$ drives the optimization of the cost functional (the 'function' of function), according to Eq. (A.10), is changing and hard to be 0, until we reach the equilibrium point (nearly impossible at the beginning of the imputation procedure).
3. During imputation procedure, $\int{r(\boldsymbol{X}^{(miss)})\mathrm{d}\boldsymbol{X}^{(miss)}}=1$ is a constant (normalization constraint), but $r(\boldsymbol{X}^{(miss)})$ and $\mathbb{H}[r(\boldsymbol{X}^{(miss)})]\coloneqq-\int{r(\boldsymbol{X}^{(miss)})\log{r(\boldsymbol{X}^{(miss)})}\mathrm{d}\boldsymbol{X}^{(miss)}}$ are not constant, since $\boldsymbol{X}^{(miss)}$ is changing
4. Furthermore, $\log{p(\boldsymbol{X}^{(miss)}\vert \boldsymbol{X}^{(obs)} )}$ is not a constant when we change $\boldsymbol{X}^{(miss)}$ unless it is a uniform distribution, which is also nearly impossible in practice.

Weakness:

W1: Technical Concerns

Please see next part for detailed information.

W2: Convergence

• Our convergence discussions have already been listed in Appendix E.3 theoretically and empirically.
• For the final solution, all we can know is that it may reach the equilibrium point, where $\nabla_{\boldsymbol{X}^{(miss)}}\frac{\delta \mathcal{F}}{\delta r(X^{(miss)})}=0$ holds.

Questions:

Q1: Understanding of Eq. (2) It seems the objective in RHS is independent of X_miss ... of r, p, X_obs. Whereas objective in LHS ... X_miss, does not involve r.

The LHS is the content at the left side of $\Rightarrow$, and RHS is the content at the right side of $\Rightarrow$. Based on this, let's start to investigate Eq. (2) based on Agreement 1:

• The LHS indicates that we are finding some $\boldsymbol{X}^{{(miss)}}$ from some unknown distribution $r(\boldsymbol{X}^{{(miss)}})$ ($\boldsymbol{X}^{{(miss)}}\sim r(\boldsymbol{X}^{{(miss)}})$), such that we can maximum the likelihood function $\log{p(\boldsymbol{X}^{{(miss)}}\vert \boldsymbol{X}^{{(obs)}})}$.
• For optimizers, to the best of our knowledge whether GAMS solvers like GuROBi or Neural Network Solvers like Adam, can only handle scalar learning objectives.
• Consequently, we can convert it to $\frac{1}{M} \sum_{i=1}^{M}{\log{p(\boldsymbol{X}^{{(miss)}}\_i|\boldsymbol{X}^{{(obs)}}\_i)}}$, where $M$ is the missing value size. This is the Monte Carlo (MC) Estimation of term $\mathbb{E}_{r}[\log{p(\boldsymbol{X}^{{(miss)}}_i)\vert p(\boldsymbol{X}^{{(obs)}})}]$ (to our understanding, the integrated out you mentioned is MC integration), which is the RHS.
• Conversion in Eq. (2) is widely used in Quantum MC [2], where they sample some instances from an optimizable distribution and optimize the concerning functional based on these instances.

Q1: Constant Question to the optimization variable the objective in RHS seems to be a constant.

The functional $\mathbb{H}[r(\boldsymbol{X}^{{(miss)}})]$ is changing, not a constant since we mentioned in Eq. (A.10) $\frac{\partial r}{\partial \tau}=-\nabla\cdot(r v)$, with velocity $v$ unless $v=0$ based on Agreement 4.

Q2: r and $\boldsymbol{X}^{{(miss)}}$

Based on Agreement 1, this question is answered.

Q2: Is NER a constant?

Based on Agreements 2 to 4, this question is answered.

Q3: $r$ is missed at the LHS of Eq. (2)

The exact expression of $r$ is not our concern, the vital component is $\boldsymbol{X}^{(miss)}$, the imputed value (Agreement 1). By investigating funcional related to $r$ and optimizing concerning functional (The $\boldsymbol{X}^{{(miss)}}\sim r(\boldsymbol{X}^{{(miss)}})$ in LHS indicates that $r$ occurs in the LHS of Eq. (2)), analyze why DM-based MDI approaches do not take effect, and propose improvements is the novelty of this manuscript.

Q3: $r$ is unkown:

It is hard to compute $p$ (existence of normalization constant) and $r$, but we can still progressively increase the value of $\mathbb{E}_{r(\boldsymbol{X}^{(miss)})}[\log{p(\boldsymbol{X}^{(miss)}\vert \boldsymbol{X}^{(obs)} )}] -\lambda \mathbb{H}[r(\boldsymbol{X}^{(miss)})]$ and realize MDI task (with input $\boldsymbol{X}^{(miss)}$) throughout WGF, that's what we did in Section 3.2 and 3.3.

Q3: Relationship between $r$ and $p$

• Based on Agreement 1, we are representing $\boldsymbol{X}^{(miss)}$ by $r$.
• $r$ is 'some' proposal distribution, and $p$ is the 'evaluator', which 'evaluates wheter $r(\boldsymbol{X}^{(miss)})$/$\boldsymbol{X}^{(miss)}$ is suitable', and 'reshape $r$ (realize by reshaping $\boldsymbol{X}^{(miss)}$) to make it appropriate' based on cost functional $\mathbb{E}_r[\log{p}]-\lambda \mathbb{H}(r)$.
• $r$ is akin to an actor, and $p$ is akin to a critic. The imputation procedure is akin to improving the actors ($r$/$\boldsymbol{X}^{(miss)}$) with guidance (WGF) from the critics ($p$).

Q4: Assumption $r({\boldsymbol{X}^{{(joint)}}})\coloneqq r({\boldsymbol{X}^{{(miss)}}}) p({\boldsymbol{X}^{{(obs)}}})$

Please see point 3) in the common rebuttal chat window, where this decomposition is widely applied in mean-field variational inference [3].

Thank you for reading our rebuttal! We hope the above discussion will fully address your concerns about our work, and we would really appreciate it if you could be generous in raising your score.

Refs
[1]. "Stein variational gradient descent: A general purpose bayesian inference algorithm." NeurIPS 16.
[2]. "Ab initio solution of the many-electron Schrödinger equation with deep neural networks." Physcial Reviews Research
[3]. "Variational algorithms for approximate Bayesian inference", Doctoral Thesis'03.

Dear Reviewer [EhUg],

Thank you for your constructive feedback, which has significantly enhanced our manuscript. Below, we address your concerns point by point with detailed explanations:

Q1: It would have been nice if important portions of related appendix were moved to the main section.

• Due to page constraints, we cannot incorporate basic knowledge sections into the main content. However, we aim to provide a comprehensive background for understanding.
• To this end, we have included detailed explanations of WGF, the MDI task, and the derivation of concerning proofs in the appendix.

Q1: Also, will the proposed alternating style algorithm converge? Do we know anything about the properties of this converged solution? Some discussion around the final algorithm would have helped in understanding the methodology better.

• We have addressed the convergence properties of our proposed KnewImp approach both theoretically and empirically in Appendix E.3.
• Detailed information about the convergence behavior and properties is provided in our individual rebuttal to you, where we may reach to the equilibrium point for $\nabla_{\boldsymbol{X}^{(miss)}}\frac{\delta \mathcal{F}}{\delta r} = 0$.

W1: I am not sure how the equivalence in (2) follows.

• Equivalence of $r(\boldsymbol{X}^{(miss)})$ and $\boldsymbol{X}^{(miss)}$ in (2): Our analysis of the DM-based MDI task through the WGF framework explains that optimizing the instances from sample $x$ is essentially optimizing the sample distribution $r(x)$.
• An example for DMs:
  • Consider the initial DMs, which are used for data generation at the beginning. Starting with a group of samples drawn from white noise, we aim to progressively refine these samples until they closely approximate a target data distribution, such as an image of a dog.
  • Throughout this transformation, the data distribution evolves from random noise to the "true distribution" represented by the DM, by optimizing the sample points.
  • Consequently, optimizing individual instances from a sample $x$ can be understood as refining the sample's probablity density function $r(x)$, which encapsulates the distribution of $x$. For DMs' inference, the goal is to align $r(x)$ as closely as possible with the true data distribution.

• Our transformation is supported by Monte Carlo Estimation, where $\mathbb{E}\_{r(x)}[f(x)] \approx \frac{1}{M} \sum_{i=1}^{M}[f(x_i)],x_i\sim r(x)$, applied from left to right is used in particle variational inference [1] and Quantum MC [2]. But for theoretical analysis, it is applied from right to left (scenario in our manuscript).

W2: Then I am not sure how this regularizer, which is essentially a constant, matters.

• Please note that, the constant regularization will not effect the optimized results since $\nabla_{\boldsymbol{X}^{(miss)}}{\text{Constant}}=0$.
• Theoretically, in our model, the entropy regularization $\mathbb{H}[r]$ is dynamic and not constant during the optimization process, unless the velocity filed $v$ is zero or $r$ becomes uniform distribution.
• Practically, evidence in Section 1.2 of our attached PDF demonstrates that changing the regularization strength from negative to positive significantly alters the optimal values, confirming that entropy regularization is not a constant and does influence the optimized results.

W3: Reading line 158 gives an impression that $r$ is the unknown. However, $r$ does not seem to appear in LHS of (2). If it does, then what is the relationship with $p$ in the objective of (2)?

• $r$, represented by the particles/samples $\boldsymbol{X}^{(miss)}$, and representing $r$ by $\boldsymbol{X}^{(miss)}$ is the quintessence of particle variational inferene approaches [1].
• Obtaining the samples rather than the detailed expressions for $r$ matters.
• For a detailed explanation of the relationship between $r$ and $p$, please refer to the metaphor in our rebuttal chat with Reviewer [833r], titled Summary of the Response to Reviewer [833r].

W4: Is the assumptions of prop 3.4 meaningful?

• Detailed justification for the mean-field factorization of $r$ within the WGF framework is provided in our discussions with Reviewer [833r], including literature reviews and theoretical derivations in the rebuttal chat entitled Response to [833r]'s additional question about $r$'s factorization and Theoretical Justification for Mean-field Factorization of $r$ within WGF framework.

Refs:
[1]. "Stein variational gradient descent: A general purpose bayesian inference algorithm.".
[2]. "Ab initio solution of the many-electron Schrödinger equation with deep neural networks."

Thank you for taking the time to read our rebuttal. We are committed to implementing thorough revisions based on your suggestions. Considering your busy schedule, please feel no obligation to respond to this message.

Best regards,
Authors

Thank you for your insightful advice and valuable questions, we will respond to your concerns point by point.

Weaknesses

W1: Dense Mathematical

• We acknowledge that the theoretical concepts and mathematical formulations in our manuscript could be challenging for readers not extensively familiar with advanced optimization and diffusion models.
• To enhance the accessibility and readability of our work, we will include a new subsection in the appendix of our revised manuscript. This subsection will detail the derivation and meaning of each equation, ensuring that the mathematical underpinnings are more comprehensible.

W2: Extra Basline

• For CSDI_T [1]:
  • Its primary contribution involves the introduction of one-hot encoding, analog bits encoding, and feature tokenization to manage categorical variables in MDI tasks on tabular data. Our research, however, specifically focuses on numerical tabular data and assumes the absence of categorical data. Due to this distinction, we did not consider CSDI_T as a relevant baseline for our study.
  • Nevertheless, the categorical feature extraction module in CSDI_T, once removed, essentially transforms it into the CSDI model, which we have indeed utilized and referenced in our manuscript.
• For ReMasker [2]:
  • We have included an additional baseline model named ReMasker in our comparisons. Please see the common rebuttal chat window. We will add the results of this baseline model in the revised manuscript.

W3: Convergence

• We have included discussions on the convergence located in Appendix E.3 with theoretical proofs and empirical validations on all datasets in our manuscript.
• We will add a footnote in our revised manuscript to demonstrate this point.

Questions:

Q1: Mixed Type Data:

• The key to applying our approach to such datasets involves decomposing the missing data distribution, $r(\boldsymbol{X}^{({miss})})$, into a product of distributions for the dense and categorical components under a mean-field assumption common in variational inference [3]. Specifically, we express it as $r(\boldsymbol{X}^{({miss})}) = r(\boldsymbol{X}^{({miss, dense})})r(\boldsymbol{X}^{({miss, cate})})$.
• For the dense data component, we continue to employ our KnewImp method.
• For the categorical part, we can initially model it using a Dirichlet distribution, which naturally supports simplex spaces, the steps are summarized as follows.
  • First, we can implement mirror descent [4] with the operator $\nabla_x\psi(x) = \log{x}$, which maps the distribution’s support from the simplex $\Delta^{\text{C}-1}$ (where $\text{C}$ represents the number of categories) onto $\mathbb{R}^{\text{C}}$.
  • Subsequently, we apply KnewImp in this transformed space ($\mathbb{R}^{\text{C}}$).
  • After that, we can revert the distribution back to the simplex using the inverse operator $[\nabla_x\psi(x)]^{-1} = \text{Softmax}(x)$.
• Besides, our example located in Section 1.2 of the attached PDF has realized this scheme, where the constraint of variables is a three-dimensional standard simplex.

Q2: Requirement

• Our computational requirement is given in lines 743 to 746: all experiments are conducted on a workstation equipped with an Intel Xeon E5 processor with four cores, eight Nvidia GTX 1080 GPUs, and 128 GB of RAM., and we found that this configuration is enough for WQW dataset with 4898 items.
• In addition, we further provided detailed time complexity analysis in Appendix E.2 theoretically and empirically, which can mirror its scalability with larger datasets.

Thank you for reading our rebuttal! We hope the above discussion will fully address your concerns about our work, and we would really appreciate it if you could be generous in raising your score.

Refs
[1] "Diffusion models for missing value imputation in tabular data." NeurIPS'22 First Table Representation Workshop.
[2] "ReMasker: Imputing Tabular Data with Masked Autoencoding." ICLR'24.
[3] "Variational algorithms for approximate Bayesian inference", Doctoral Thesis'03.
[4] "Sampling with Mirrored Stein Operators" ICLR'22

Dear Reviewer [Gtbe],

Thank you for your valuable feedback and encouraging comments, which have greatly motivated us throughout the rebuttal process. Out of respect for your review efforts, we would like to summarize how we have addressed your inquiries:

Clarifying the Paper:

• In response to insights gained from discussions with other reviewers, we recognize that some of our mathematical formulations may appear dense. We plan to include metaphors to better illustrate concepts like our $r$ and $p$ dynamics [EhUg, 833r], explain transformations such as Monte Carlo estimation [EhUg], and discuss the selection of RKHS [9XXU] to enhance the manuscript's readability.

Incorporating Additional Baseline Models:

Following your recommendation, we have added the baseline model ReMasker. We will detail its integration and relevance to the DMs used in our study, tailored to our specific scenarios (numerical tabular data).

Enhanced Discussion on Convergence, Limitations, and Scenarios:

• We will augment our manuscript with a detailed proof of the convergence for the "estimate" part, furthering our discussions with Reviewer [EhUg].

• Additional experiments and analyses on toy case datasets will be included, especially considering data properties like multi-modality, heavy tails, and skewed distributions as suggested by Reviewer [9XXU].

Handling Mixed-Type Data

We plan to outline strategies for managing mixed-type data in our future research directions, building on our discussions with you.

Computational Resources

We will clearly indicate the computational resources used in our study in the revised manuscript to ensure transparency according to your guidance.

Thank you once again for your supportive and constructive review. We are grateful for your continued positive assessment. Given your busy schedule, please do not feel obliged to respond to this message.

Sincerely,

Authors