Title: Rethinking the Diffusion Models for Missing Data Imputation: A Gradient Flow Perspective
Venue: NeurIPS 2024
Accepted: yes
Keywords: Missing Data Imputation

Gradient Flow

Reproducing Kernel Hilbert Space

Functional Optimization
Abstract: Diffusion models have demonstrated competitive performance in missing data imputation (MDI) task. However, directly applying diffusion models to MDI produces suboptimal performance due to two primary defects. First, the sample diversity promoted by diffusion models hinders the accurate inference of missing values. Second, data masking reduces observable indices for model training, obstructing imputation performance. To address these challenges, we introduce $\underline{\text{N}}$egative $\underline{\text{E}}$ntropy-regularized $\underline{\text{W}}$asserstein gradient flow for $\underline{\text{Imp}}$utation (NewImp), enhancing diffusion models for MDI from a gradient flow perspective. To handle the first defect, we incorporate a negative entropy regularization term into the cost functional to suppress diversity and improve accuracy. To handle the second defect, we demonstrate that the imputation procedure of NewImp, induced by the conditional distribution-related cost functional, can equivalently be replaced by that induced by the joint distribution, thereby naturally eliminating the need for data masking. Extensive experiments validate the effectiveness of our method. Code is available at [https://github.com/JustusvLiebig/NewImp](https://github.com/JustusvLiebig/NewImp).

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

Abstract

Diffusion models have demonstrated competitive performance in missing data imputation (MDI) task. However, directly applying diffusion models to MDI produces suboptimal performance due to two primary defects. First, the sample diversity promoted by diffusion models hinders the accurate inference of missing values. Second, data masking reduces observable indices for model training, obstructing imputation performance. To address these challenges, we introduce Negative Entropy-regularized Wasserstein gradient flow for Imputation (NewImp), enhancing diffusion models for MDI from a gradient flow perspective. To handle the first defect, we incorporate a negative entropy regularization term into the cost functional to suppress diversity and improve accuracy. To handle the second defect, we demonstrate that the imputation procedure of NewImp, induced by the conditional distribution-related cost functional, can equivalently be replaced by that induced by the joint distribution, thereby naturally eliminating the need for data masking. Extensive experiments validate the effectiveness of our method. Code is available at https://github.com/JustusvLiebig/NewImp.

1 Introduction

Missing data is a pervasive problem for data analytics in diverse scenarios, including e-commerce [29, 30, 57], healthcare [51, 56], and process industry [33, 58]. For instance, in healthcare, patient monitoring devices may fail or lose connection, leading to missing vital signs data. Similarly, in industrial processes, sensor signals may be incomplete due to inevitable mechanical shock. These incompletenesses hamper data integrity and impede subsequent analysis. Therefore, accurate missing data imputation (MDI) is critical for enabling reliable analysis and decision in real-world applications.

Recently, diffusion models (DMs) have emerged as a powerful tool for MDI [66]. Specifically, these models first estimate the (Stein) score function of the missing data conditioned on the observed data, subsequently reformulating the imputation problem as a generative task grounded in the learned score function. These works are initiated from [51] and evolve to incorporate crafted model architecture [33] and learning objectives [38, 73] for enhancing the accuracy of score estimation [41]. Celebrated for their advantageous capability to model data distributions and generate high-quality synthetic data [41, 50, 65], diffusion models have been a prevalent approach to MDI.

Despite the successes of diffusion models, we argue that directly applying diffusion models to MDI results in suboptimal performance due to two primary limitations. First, diffusion models perform

imputation by sampling from a learned score function, which inadvertently promotes diversity in the imputed values. This increased diversity contradicts the accuracy required for precise imputation of missing data [38]. Second, the training process involves masking a portion of the observed data as labels. The selection of masking strategy significantly impacts imputation accuracy and is inherently challenging to optimize [51]. Moreover, the masked data during training often differ in missing mechanisms from those encountered during testing, resulting in a discrepancy between training and inference phases that degrades performance. Consequently, diffusion models introduce unintended diversity and impose data masking, both of which impede effective imputation.

To tackle these challenges, we introduce a novel DM-based MDI approach termed Negative Entropyregularized Wasserstein Gradient Flow Imputation (NewImp). Specifically, to handle the first issue, we revisit DM-based MDI task within the Wasserstein Gradient Flow (WGF) framework, derive the associated cost functionals, and identify that they implicitly promote diversity in the imputed values. Building on this insight, we incorporate a negative entropy-regularized (NER) cost functional to suppress imputation diversity and enhance accuracy. Furthermore, we derive a closed-form imputation procedure based on the proposed cost functional within the reproducing kernel Hilbert space (RKHS). After that, we further prove that within the WGF framework, the imputation procedure of NewImp, induced from the cost functional associated with conditional distribution, can be induced from another cost functional associated with joint distribution equivalently, within which we merely need to estimate the joint distribution during the model training stage, thereby naturally eliminating the need for data masking.

Contributions. The main contributions of this paper are summarized as follows:

We demonstrate that directly applying diffusion models to MDI causes suboptimal performance, as they prompt unintended diversity and require data masking, both impeding accurate imputation.

We propose NewImp, a novel DM-based MDI approach under the WGF framework which introduces an NER cost functional to suppress unintended diversity. Based on this, we further prove that the imputation procedure of NewImp can be induced from an equivalent joint-distribution-related functional, and consequently introduce an imputation procedure that sidesteps the data masking.

We conduct various experiments over public numerical tabular datasets to demonstrate the superiority of the NewImp method over prevalent baseline models.

2 Preliminaries

2.1 Problem Formulation

Suppose X (ideal) ∈ R N × D represents an ideal numerical tabular dataset without any missing entries, where N and D denote the number of samples and features, respectively. The observed dataset is expressed as: X (obs) = X ( ideal ) glyph[circledot] M + NaN glyph[circledot] ( ✶ N × D -M ) , where glyph[circledot] denotes the Hadamard product, ✶ N × D is a matrix of ones of size N × D , and M ∈ { 0 , 1 } N × D is a binary mask that indicates the presence (1) or absence (0) of data in each entry. The task of MDI involves imputing the missing entries in X (obs) . This is achieved by constructing a matrix ˆ X = X (obs) glyph[circledot] M + X ( imp ) glyph[circledot] ( ✶ N × D -M ) , where X ( imp ) is the matrix containing the imputed values.

The missing mechanism can be classified into three categories [44]: Missing Completely at Random (MCAR), Missing at Random (MAR), and Missing Not at Random (MNAR) (Detailed information about missing mechanisms is given in Appendix E.1). Notably, in the MNAR setting, it is generally difficult to identify the missing data distribution without additional assumptions and constraints [22]. Hence, our discussion primarily focuses on numerical tabular data with MAR and MCAR settings.

2.2 Diffusion Models and Its Application for MDI Task

Diffusion models function by gradually corrupting data towards a tractable noise distribution, such as a standard Gaussian, and subsequently reversing this corruption to generate samples [50]. Specifically, the forward corruption process is modeled as a discretization of a stochastic differential equation (SDE) over time τ : d X τ = f ( X τ )d τ + g τ d W τ , where f ( X τ ) is drift term, g τ is volatility term, and d Wτ is standard Wiener process. The solution to this SDE creates a continuous trajectory of random variables X τ | T τ =0 . The density function q τ of X τ adheres to the Fokker-Planck-Kolmogorov (FPK)

equation: ∂q τ ∂τ = -∇· ( q τ f ( X τ )) + 1 2 g 2 τ ∇· ∇ q τ (see Theorem 5.4 in reference [47]). The reverse process is governed by: d X τ = [ f ( X τ ) -g 2 τ ∇ log p ( X τ )]d τ + g τ d W τ [3], where ∇ log p ( X τ ) represents the score function, which is often parameterized by neural networks.

Diffusion models treat MDI as a conditional generation task. The score function, ∇ log p ( X ) , is defined specifically for MDI as ∇ X (miss) log p ( X (miss) | X (obs) ) [51], and the MDI task is executed by generating samples based on this conditional score function. The key challenge is to obtain an estimation ∇ X (miss) log ˆ p ( X (miss) | X (obs) ) that approximates ∇ X (miss) log p ( X (miss) | X (obs) ) . Given that the true X (miss) is unknown, existing DM-based approaches utilize a mask matrix to drop some observable data as labels. However, the specification of the mask mechanism, determining the effectiveness of ∇ log ˆ p ( X (miss) | X (obs) ) , is challenging since it should align with the data missing mechanism in the testing dataset [51], which may be unknown in practice.

2.3 Wasserstein Gradient Flow

Wasserstein space P 2 ( R D ) is defined as the set of distributions with finite second-order moments. Consider a cost functional F cost : P 2 ( R D ) → R ; the celebrated Wasserstein gradient flow (WGF) is an absolutely continuous trajectory ( q τ ) τ> 0 in this space, which evolves over time τ to minimize F cost efficiently. This dynamic is governed by the continuity equation:

where u τ : R D → R D is a time-dependent velocity field [2], whose input is sample X ∈ R D ; δ F cost δq τ denotes the first variation of F cost with respect to q τ . On this basis, the evolution of X over time τ in P 2 ( R D ) can be modeled by the ordinary differential equation (ODE):

However, simulating this ODE is challenging since u τ involves the estimation of q τ , which involves solving the differential equation ∂q τ ∂τ = -∇· ( u τ q τ ) that proves to be not analytically solvable [16].

3 Motivations

3.1 Diffusion Models Secretly Foster Diversity

Based on the notations defined in Section 2.1, we can first define the following cost functional for MDI task according to the maximum likelihood estimation principle, where we want to find the value with the highest probability:

where ˆ p ( X ( miss ) | X ( obs ) ) is the estimation of p ( X ( miss ) | X ( obs ) ) via neural network [52]. Notably, we can treat X ( miss ) as samples from a 'proposal distribution' r ( X ( miss ) ) , and formulate the following optimization problem based on variational inference [27, 70]:

0

0

0

0.5

0.5

0.5

1

1

1

0

2

4

6

8

(a) Optimal Value.

0

0

0

0.5

0.5

0.5

1

1

1

0

2

4

6

8

(b) Results by DM.

where we aim to sample some X ( miss ) samples from proposal distribution r ( X ( miss ) ) , realize the maximum log-likelihood estimation over the sampled results, and 'optimize' the proposal distribution r ( X ( miss ) ) that is represented by samples X ( miss ) . Notably, in Eq. (4), we use the spirit from previous references represented by [32], where optimizing the samples X ( miss ) is equivalent to optimizing the distribution r ( X ( miss ) ) .

Referring to Eq. (4), it is observed that the MDI task can be formulated as an optimization problem. This prompts a pertinent question: If the conditional distribution log p ( X ( miss ) i | X ( obs ) i ) is estimated accurately, what would happen if we directly apply diffusion models to solve the optimization problem corresponding to MDI task? To explore this, we consider a hypothetical scenario: Suppose we are optimizing a cost functional related to a three-dimensional Dirichlet distribution on the simplex ∆ 2 :

where a h | H 1 are variables, H is variable number, ρ k | 3 k =1 is concentration parameter, and Γ( · ) is gamma function. We compare the analytically derived optimal value with the results from diffusion models in Fig. 1. The diffusion model's results tend to surround but do not exactly reach the optimal value, suggesting that there might be implicit , diversity-encouraging terms integrated into the diffusion models' objectives that produce the observed inaccuracies. Identifying and modifying these regularization terms is crucial for enhancing the efficacy of diffusion models for MDI tasks.

3.2 Negative Entropy Regularization Term for Diversity Suppression

In this section, we identify and refine the terms in diffusion models' objectives that prompt unintended diversity and impede accurate imputation. We observe that the inference process in diffusion models adheres to the FPK equation, which is a specialized form of the continuity equation in WGF (see Sections 2.2 and 2.3). This alignment inspires us to reframe diffusion models within the WGF framework, enabling the derivation of their underlying cost functionals. By doing so, we can compare these functionals with the objective functional for MDI in Eq. (4) 2 .

Proposition 3.1. Within WGF framework, DM-based MDI approaches can be viewed as finding the imputed values X ( imp ) that maximize the following objective:

where 'const' is the abbreviation of constant, and ψ ( X ( miss ) ) is a scalar function determined by the type of SDE underlying the diffusion models.

VP-SDE: ψ ( X ( miss ) ) = 1 2 H [ r ( X (miss) )] + E r ( X ( miss ) ) { 1 4 [ X (miss) ] glyph[latticetop] [ X (miss) ] } ≥ 0

VE-SDE: ψ ( X ( miss ) ) = 1 2 H [ r ( X (miss) )] ≥ 0

sub-VP-SDE: ψ ( X ( miss ) ) = 1 2 H [ r ( X (miss) )] + E r ( X ( miss ) ) { 1 4 γ τ [ X (miss) ] glyph[latticetop] [ X (miss) ] } ≥ 0 ,

where H [ r ( X (miss) )] := -∫ r ( X (miss) ) log r ( X (miss) )d X (miss) is the entropy term, γ τ is determined by noise scale β τ : γ τ := (1 -exp( -2 ∫ τ 0 β s d s )) > 0 , 0 < β 1 < · · · < β T < 1 .

Proposition 3.1 reveals that diffusion models inherently optimize an objective functional that largely aligns with (4), but they secretly include an additional term ψ ( X ( miss ) ) > 0 . This term makes Eq. (5) an upper bound on Eq. (4), i.e., maximizing the cost functional in Eq. (5) does not guarantee to maximize the MDI objective in Eq. (4). Furthermore, the entropy term included in the models fosters sample diversity, which may compromise the accuracy required in MDI tasks [53, 38]. To address this issue, we propose incorporating a negative entropy term as ψ ( X ( miss ) ) :

where λ > 0 is a predefined regularization strength, and consequently we can define our NER cost functional for MDI task as follows:

The objective functional in Eq. (7) provides a lower bound of Eq. (4). Therefore, maximizing F NER guarantees maximizing Eq. (4). Meanwhile, F NER effectively reduces the unintended diversity term, contributing to an improvement in imputation accuracy.

4 Implementation of the NewImp

4.1 Optimizing the F NER within WGF Framework

In this section, we aim to optimize F NER within the WGF framework [46, 69]. To this end, we plug Eq. (7) into Eq. (1), producing the velocity field below that drives the ODE in Eq. (2):

However, as stated in Section 2.3, implementing this ODE in computer code is intricate due to the intractability of the density function r ( X (miss) ) . Fortunately, by restricting the velocity field within the Reproducing Kernel Hilbert Space (RKHS) defined by the kernel function u ( X (miss) ) ∈ K ( X (miss) , ˜ X (miss) ) , an alternative ODE minimizing F NER can be implemented in Proposition 4.1 [35, 31] which sidesteps the intractable r ( X (miss ) ) 3 .

Proposition 4.1. Suppose u ( X (miss) ) is a velocity field regularized by the RKHS norm under the following conditions: 1). The kernel function satisfies: lim ‖ X (miss) ‖→∞ K ( X (miss) , ˜ X (miss) ) = 0 . 2). The density r ( X (miss) ) is bounded. Then, the velocity field that minimizes the cost functional F NER = E r ( X ( miss ) ) [log ˆ p ( X ( miss ) | X ( obs ) )] -λ H [ r ( X (miss) )] can be given by:

where the expectation term E r ( ˜ X ( miss ) ) can be efficiently estimated using Monte Carlo approximation, K ( X , ˜ X ) is set as the radial basis function (RBF) kernel.

4.2 Sidestepping Mask Matrix: Conditional Modeling via Joint Modeling

Simulating the ODE in Eq. (2) with Eq. (8) necessitates an accurate estimation of p ( X ( miss ) | X ( obs ) ) . However, this modeling is challenging due to the diverse choices of masking matrices. More specifically, the accuracy of the estimated conditional distribution p ( X ( miss ) | X ( obs ) ) heavily relies on the selection of these matrices, and these matrices should be consistent with the data missing mechanism in the testing dataset, which may be unknown in practice [51]. To bypass this difficulty, we suggest substituting the conditional distribution p ( X ( miss ) | X ( obs ) ) with the joint distribution p ( X ( joint ) ) , where X ( joint ) = ( X ( miss ) , X ( obs ) ) . Building on this substitution, the velocity field is redefined based on the estimated joint distribution ˆ p ( X (joint) ) as follows:

where ∇ ˜ X (miss) log ˆ p ( ˜ X (joint) ) can be obtained by masking the ∇ X (joint) log ˆ p ( X (joint) ) with the missing data indicator matrix M as follows:

where h is the bandwidth, the values of ˜ X (joint) and X (joint) are identical , and the tilde notation on ˜ X (joint) is merely used to distinguish the variable with respect to which the derivative is taken . On

this basis, similar to Eq. (10), the gradient term ∇ ˜ X (miss) K ( X (joint) , ˜ X (joint) ) can be given as follows:

✶ and since K ( X (joint) , ˜ X (joint) ) is a smooth function, ∇ ˜ X (joint) K ( X (joint) , ˜ X (joint) ) can be easily computed by automatic-differentiation-based deep learning backends like by PyTorch [42].

Proposition 4.2 demonstrates that the cost functional F joint-NER, associated with Eq. (10), and F NER exhibit a constant gap, indicating that optimizing F joint-NER is equivalent to optimizing F NER.

Proposition 4.2. Assume that the proposal distribution r ( X (joint) ) is factorized by r ( X (joint) ) := r ( X (miss) ) p ( X (obs) ) . The cost functional associated with the joint distribution is defined as follows:

which leads to the velocity field delineated in Eq. (9) and establishes F joint-NER as a lower bound for F NER, with the difference being a constant (i.e., F joint-NER = F NER -const , const ≥ 0 ).

The detailed justification for the factorization r ( X (joint) ) := r ( X (miss) ) p ( X (obs) ) is provided in Appendix C. Based on this proposition, the following corollary can be obtained:

Corollary 4.3. The following equation holds: u ( X (joint) ) = u ( X (miss) ) .

So far, we know that u ( X (joint) ) can reduce F NER as effectively as possible, which indicates that the velocity field defined in Eq. (9) can fully substitute for Eq. (8) in optimizing F NER without loss of accuracy. Finally, the imputed value can be obtained by simulating the following ODE:

4.3 Estimating the Joint Distribution

The remaining problem is to determine the estimation of score function ∇ X (joint) log ˆ p ( X (joint) ) . To achieve this, we employ Denoising Score Matching (DSM) [21, 52] to train the score function ∇ X (joint) log ˆ p ( X (joint) ) parameterized by a neural network. Specifically, the learning objective is designed to minimize the discrepancy between the actual score and the model's predicted score after introducing Gaussian noise to the clean X (joint) as ˆ X (joint) :

Notably, σ is variance scale, ˆ X (joint) is obtained by ˆ X (joint) = X (joint) + glyph[epsilon1], glyph[epsilon1] ∼ N ( 0 , σ 2 I ) , and ∇ ˆ X (joint) log q σ ( ˆ X (joint) | X (joint) ) = -ˆ X (joint) -X (joint) σ 2 . Once ∇ X (joint) log ˆ p ( X (joint) ) is trained, we can obtain the imputation value by simulating the ODE based on Eqs. (9) and (14).

4.4 Overall Workflow of NewImp

The computation workflow of NewImp is encapsulated in Algorithm 1. Specifically, we perform a mean imputation to the incomplete matrix X (obs) , producing a pre-imputed dataset denoted as X (imp) (step 1). After that, we iteratively conduct DSM training and ODE simulation. In DSM training (steps 3-5), we form X (joint) and conduct DSM on it to acquire a score estimator. In ODE simulation (steps 6-8), we set the starting point and perform ODE simulation, where u is calculated with the score estimator acquired in step 5. The endpoint is treated as the imputation results at the current iteration. After completing T iterations of this process, the imputed dataset ˆ X is calculated, where we preserve the observed indices in X (obs) (step 10).

5 Experiments

5.1 Experimental Setup

Datasets: We conduct the case study based on eight real-world datasets from the UCI repository. To simulate missing data, we mask the dataset using a mask matrix, which is realized with a Bernoulli random variable of fixed mean. More detailed information is provided in Appendix E.1.

Algorithm 1 The overall workflow of NewImp approach

Input: observational data X (obs) , mask matrix M .

Hyperparameter: loop time: T , simulation time: T , discretization step size η , bandwidth h , neural network learning rate lr , training epoch E , neural network hidden unit HUscore.

Output: imputed data ˆ X .

1:

Initialize

X

(imp)

2:

3:

4:

5:

6:

7:

8:

9:

10:

for

t

= 0

X

X

(miss)

(joint)

∇

X

(joint)

X

X

(miss)

0

(miss)

T

X

(imp)

end for

X

ˆ

X

(obs)

X

(

(imp)

T

to do

X

X

(miss)

(

p

X

X

X

glyph[circledot]

(obs)

(

)

glyph[circledot]

(

×

N

(joint)

(imp)

(miss)

0

X

glyph[circledot]

(miss)

T

M

+

)

✶

+

X

∫

D

✶

N

-

DSM

T

0

u

(

imp

×

D

M

-

) +

(

X

(

X

)

glyph[circledot]

(joint)

(joint)

(

M

X

)

(obs)

)

)d

τ

N

×

D

-

glyph[triangleright]

M

)

Baselines: Wecompare NewImp with DM-based MDI models: CSDI for Tabular Data (CSDI_T) [51], MissDiff [41]. In addition, we also select other well-known MDI models like Sinkhorn (Sink) [40], Transform Distribution Matching (TDM) [72], Generative Adversarial Imputation Nets (GAIN) [67], Missing Data Importance-Weighted Autoencoder (MIWAE) [39], Missing data Imputation Refinement And Causal LEarning (MIRACLE) [28], and ReMasker [15]. Details concerning experimental settings are given in Appendix E.2.

✶

Implementation Details: In this study, we employ a two-layer multi-layer perceptron to model ∇ X (joint) log ˆ p ( X (joint) ) . Each layer is configured with 256 hidden units (HUscore), and the activation function is set as the 'Swish' function [43]. For DSM training (step 5 of Algorithm 1), the variance scale σ is set as 0.1, the network is trained by the Adam optimizer [26] with the learning rate of 1.0 × 10 -3 , and the batch size is dynamically set to dataset size N . Meanwhile, for the ODE simulation part (step 7 of Algorithm 1), we specify a simulation time ( T ) of 500, a regularization strength of ( λ ) 10.0, a step size of ( η ) 0.1, and a bandwidth ( h ) of 0.5. The loop time T for NewImp is set as 2. Since only missing indices are updated, the evaluation focuses exclusively on imputation errors for these indices. To this end, we modify the mean absolute error (MAE) and squared Wasserstein-2 distance (WASS) according to reference [22] as follows:

where W 2 2 denotes the squared Wasserstein-2 distance, M 1 := { i : ∃ j, M i,j = 0 } represents the subset of M i,j with at least one missing value, m 1 is the number of data points with at least one missing value, and ∆ X is the Dirac distribution (measure) concentrated on X .

5.2 Baseline Comparison Results

Table 1 presents the imputation quality of NewImp and other imputation approaches under the MAR and MCAR scenarios. The primary observations are detailed as follows:

Models with neural architectures such as MIRACLE, MIWAE, and TDM demonstrate superior performance compared to models lacking such architectures. This observation suggests that integrating neural networks into MDI tasks can significantly enhance model performance.

DM-based imputation approaches generally perform worse than other MDI methods. This outcome indicates that despite the incorporation of complex nonlinear neural architectures to boost performance, employing diversity-oriented generative approaches may not align well with the precision requirements of MDI tasks.

glyph[triangleright]

Initialization see Algorithm 3

glyph[circledot]

glyph[triangleright]

ODE Simulation by Appendix D.1 with Eq. (9)

M

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Our proposed NewImp method consistently ranks as the best or second-best across most comparisons under most of the scenarios and datasets. Notably, NewImp significantly outperforms other DM-based MDI approaches, underscoring the effectiveness of our analytical enhancements and innovations in Sections 3.1, 3.2 and 4.2.

5.3 Ablation Study Results

In this section, we conduct the ablation study to assess the contributions of two key components in NewImp: the NER term and the joint modeling strategy (referred to as 'Joint'). The results of this study are detailed in Table 2. Analysis of the data between the second and last rows of Table 2 reveals that, for most cases, in the absence of the NER, the proposal distribution r ( X (miss) ) may become pathological, leading to diminished model performance. Additionally, when comparing results from the first, third, and last rows, it becomes evident that modeling the joint distribution directly, rather than inferring it from the conditional distribution, significantly enhances model performance. This finding underscores the effectiveness of the strategies we have implemented, as discussed in Section 4.2. Overall, the ablation study underscores the critical roles of both the NER term and the joint distribution learning strategy in promoting the performance of NewImp.

5.4 Sensitivity Analysis Results

In this section, we analyze the impact of key hyperparameters within the NewImp approach, including the bandwidth h of the RBF kernel function, the hidden units HU score in the score network, the weight λ of the NER term, and the discretization step size η for simulating the ODE defined in Eq. (9). The profound influence of these hyperparameters on learning objectives and overall performance is substantiated by the experimental results presented in Fig. 2. Initially, we explore the effects of varying

0.1 0.25 0.5

1.0

2.0

5.0 10.0

h

0.2

0.4

0.6

0.8

WASS

MAE

32

64

128

256

512

HUscore

0.2

0.3

0.4

WASS

MAE

1.0

2.0

5.0

10.0 15.0 20.0

λ

0.2

0.3

0.4

0.5

WASS

MAE

0.005

0.05

0.1

0.2

0.5

η

0.2

0.3

0.4

WASS

MAE

(a) MAR with 30% missing rate at CC dataset.

0.1 0.25 0.5

1.0

2.0

5.0 10.0

h

0.4

0.6

0.8

1.0

WASS

MAE

32

64

128

256

512

HUscore

0.35

0.40

0.45

0.50

WASS

MAE

1.0

2.0

5.0

10.0 15.0 20.0

λ

0.35

0.40

WASS

MAE

0.005

0.05

0.1

0.2

0.5

η

0.35

0.40

WASS

MAE

(b) MCAR with 30% missing rate at CC dataset.

the bandwidth h . We observe that an increase in bandwidth correlates with a decrease in imputation accuracy. For instance, as the bandwidth increases from 0.5 to 2.0, the MAE and WASS escalate from 0.35 and 0.25 to 0.82 and 0.74, respectively. This trend suggests that excessive bandwidth can lead to an over-smoothed velocity field, expanding the exploration space of the distribution r ( X (joint) ) excessively and failing to adequately 'concentrate' this distribution, ultimately diminishing performance. Subsequently, we examine changes in the score network's hidden units. Increasing the hidden units from 256 to 512 appears to decrease imputation accuracy, likely due to overfitting issues associated with larger neural networks. Next, we adjust the strength of the NER term and find that increasing its intensity generally improves imputation accuracy. This supports the necessity of the NER term, further validating its effectiveness. Lastly, we investigate the discretization step size for the ODE. We find that accuracy initially increases with smaller step sizes but then decreases. This pattern is consistent with ODE simulation behavior, where smaller step sizes require longer to converge, potentially resulting in lower accuracy within a predefined time. Conversely, larger step sizes increase discretization errors, adversely affecting accuracy as well.

6 Related Works

6.1 Diffusion Models for Missing Data Imputation

The impressive ability of diffusion models to synthesize data [54, 76, 7] has inspired extensive research into their application for MDI tasks [59, 66]. Among the pioneering efforts, the Conditional Score-based Diffusion models for Imputation (CSDI) [51] was the first to adapt diffusion models for time-series MDI, substituting the score function with a conditional distribution and pioneering a novel model training strategy by masking parts of the observational data. Building on this, to address categorical data in tabular datasets, CSDI_T [73] introduced an embedding layer within the feature extractor. To enhance inference efficiency, the conditional Schrödinger bridge method for probabilistic time series imputation proposed modeling the diffusion process as a Schrödinger bridge [10]. Meanwhile, MissDiff [41] utilizes the missing data information as the mask matrix to improve the model training procedure.

Despite these advancements from the perspective of feature extraction module [1, 64], loss function [41], and model inference approach [60], as pointed out by reference [38], the reconciliation of the inherent diversity-seeking nature of diffusion models' generative processes and the accuracycentric demands of MDI task remains underexplored. To our knowledge, this paper is the first to elucidate the relationship between diffusion models' generative processes and MDI tasks from an optimization perspective (Sections 3.1 and 3.2), which has not been discovered by previous reference [38]. Based on these insights, we further propose our NewImp approach by designing the NER term to prioritize the MDI accuracy (Section 3.2).

6.2 Modeling Conditional Distribution by Joint Distribution

Modeling conditional distribution as joint distribution remains an opening question and has a broad potential for application [68, 8, 25]. Conditional sliced WGF [14] first empirically validated that the velocity field of joint distribution and conditional distribution are identical when choosing sliced Wasserstein distance as cost functional. After that, reference [25] extended this relationship and derived the relationship between conditional and joint distribution in various discrepancy metrics like f -divergence, Wasserstein distance, and integral probability metrics. On this basis, reference [19] further theoretically proved the equivalence of velocity fields for conditional and joint distribution.

However, the objective of NewImp does not belong to any kind of discrepancy metric [25]. The most similar discrepancy metric is negative KL divergence. Notably, negative KL divergence contains diversity-encouraging 'positive' entropy as the regularization term, and the regularization term in our study is diversity-discouraging 'negative' entropy, and thus more than directly applying these results to our research is needed. On this basis, our theoretical contribution proves that this joint distribution modeling approach can still be applied when the functional is regularized by the negative entropy.

6.3 Wasserstein Gradient Flow for Generative Modeling

WGF [2, 46] has been extensively employed in various domains of machine learning, including generative modeling [17, 12, 74, 63], posterior distribution sampling [55, 35, 32, 34], and domain adaptation [75, 36, 37, 71]. In generative modeling [4, 11, 18], the problem is framed as an optimization task, with the objective functional comprising an f -divergence term that measures the discrepancy between the proposal distribution and the data distribution, alongside an entropy term that promotes diversity in generative results.

WGFis then utilized to address the optimization of this cost functional, with models being constructed during the solution process. However, as indicated by our illustrative example (Section 3.1), and further supported by our theoretical analysis (Section 3.2), pursuing diversity in accuracy-oriented tasks such as MDI may not be appropriate. Our analysis reveals that the inclusion of an entropy term in the cost functional makes the direct application of diffusion models to MDI tasks unsuitable. Based on these insights, one of our major contributions is demonstrating that WGF can be effectively used to analyze and improve the appropriateness of applying diffusion models to non-generative tasks.

7 Conclusions

This work demonstrated that directly applying diffusion models to MDI resulted in suboptimal performance due to unintended diversity and the requirement for data masking, both of which impeded accurate imputation. To counteract this, we proposed NewImp, a novel diffusion model-based MDI approach within the Wasserstein gradient flow framework, designed to suppress unintended diversity. We developed an easy-to-implement form for realizing NewImp in computer code by constraining the velocity field within the reproducing kernel Hilbert space. Furthermore, we proved that the imputation procedure of NewImp could be derived from an equivalent joint-distribution-related functional, thereby obviating the need for data masking. Finally, extensive experiments demonstrated that NewImp effectively mitigates these issues and outperforms prevalent baseline models.

References

Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecasting with structured state space models. Trans. Mach. Learn. Res. , 2022.

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures . Springer Science & Business Media, 2005.

Brian DO Anderson. Reverse-time diffusion equation models. Stoch. Process. their Appl. , 12 (3):313-326, 1982.

Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining deep generative models via discriminator gradient flow. In Proc. Int. Conf. Learn. Represent. , pages 1-24, 2021.

Matthew James Beal. Variational algorithms for approximate Bayesian inference . University of London, University College London (United Kingdom), 2003.

John Charles Butcher. Numerical methods for ordinary differential equations . John Wiley & Sons, 2016.

Hanqun Cao, Cheng Tan, Zhangyang Gao, Yilun Xu, Guangyong Chen, Pheng-Ann Heng, and Stan Z. Li. A survey on generative diffusion models. IEEE Trans. Knowl. Data Eng. , 36(7): 2814-2830, 2024. doi: 10.1109/TKDE.2024.3361474.

Jannis Chemseddine, Paul Hagemann, Christian Wald, and Gabriele Steidl. Conditional wasserstein distances with applications in bayesian ot flow matching. arXiv preprint arXiv:2403.18705 , pages 1-42, 2024.

Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particleoptimization framework for scalable bayesian sampling. arXiv preprint arXiv:1805.11659 , pages 1-13, 2018.

Yu Chen, Wei Deng, Shikai Fang, Fengpei Li, Nicole Tianjiao Yang, Yikai Zhang, Kashif Rasul, Shandian Zhe, Anderson Schneider, and Yuriy Nevmyvaka. Provably convergent schrödinger bridge with applications to probabilistic time series imputation. In Proc. Int. Conf. Mach. Learn. , pages 4485-4513, 2023.

Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, and Yao Xie. Convergence of flow-based generative models via proximal gradient descent in wasserstein space. IEEE Trans. Inf. Theory , pages 1-1, 2024. doi: 10.1109/TIT.2024.3422412.

Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Scalable Wasserstein gradient flow for generative modeling through unbalanced optimal transport. In Proc. Int. Conf. Mach. Learn. , pages 8629-8650, 2024.

Hanze Dong, Xi Wang, LIN Yong, and Tong Zhang. Particle-based variational inference with preconditioned functional gradient flow. In Proc. Int. Conf. Learn. Represent. , pages 1-26, 2022.

Chao Du, Tianbo Li, Tianyu Pang, Shuicheng Yan, and Min Lin. Nonparametric generative modeling with conditional sliced-wasserstein flows. In Proc. Int. Conf. Mach. Learn. , pages 8565-8584, 2023.

Tianyu Du, Luca Melis Melis, and Ting Wang. Remasker: Imputing tabular data with masked autoencoding. In Proc. Int. Conf. Learn. Represent. , pages 1-23, 2024.

Lawrence C Evans. Partial differential equations , volume 19. American Mathematical Society, 2022.

Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational Wasserstein gradient flow. In Proc. Int. Conf. Mach. Learn. , pages 6185-6215, 2022.

Yuan Gao, Yuling Jiao, Yang Wang, Yao Wang, Can Yang, and Shunkang Zhang. Deep generative learning via variational gradient flow. In Proc. Int. Conf. Mach. Learn. , pages 2093-2101, 2019.

Paul Hagemann, Johannes Hertrich, Fabian Altekrüger, Robert Beinert, Jannis Chemseddine, and Gabriele Steidl. Posterior sampling based on gradient flows of the MMD with negative distance kernel. In Proc. Int. Conf. Learn. Represent. , pages 1-32, 2024.

Ya-Ping Hsieh, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored Langevin Dynamics. In Proc. Adv. Neural Inf. Process. Syst. , pages 1-10, 2018.

Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. J. Mach. Learn. Res. , 6(24):695-709, 2005.

Daniel Jarrett, Bogdan C Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperimpute: Generalized iterative imputation with automatic model selection. In Proc. Int. Conf. Mach. Learn. , pages 9916-9937, 2022.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokkerplanck equation. SIAM J. Math. Anal. , 29(1):1-17, 1998.

Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, and Ivan Oseledets. Understanding DDPM latent codes through optimal transport. In Proc. Int. Conf. Learn. Represent. , pages 1-15, 2022.

Young-geun Kim, Kyungbok Lee, and Myunghee Cho Paik. Conditional Wasserstein Generator. IEEE Trans. Pattern Anal. Mach. Intell. , 45(6):7208-7219, 2023. doi: 10.1109/TPAMI.2022. 3220965.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. Int. Conf. Learn. Represent. , pages 1-8, 2015.

Diederik P Kingma and Max Welling. Auto-encoding Variational Bayes. In Proc. Int. Conf. Learn. Represent. , pages 1-8, 2014.

Trent Kyono, Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. MIRACLE: Causallyaware imputation via learning missing data mechanisms. Proc. Adv. Neural Inf. Process. Syst. , pages 23806-23817, 2021.

Haoxuan Li, Kunhan Wu, Chunyuan Zheng, Yanghao Xiao, Hao Wang, Zhi Geng, Fuli Feng, Xiangnan He, and Peng Wu. Removing hidden confounding in recommendation: a unified multi-task learning approach. Proc. Adv. Neural Inf. Process. Syst. , pages 1-13, 2024.

Haoxuan Li, Chunyuan Zheng, Shuyi Wang, Kunhan Wu, Eric Wang, Peng Wu, Zhi Geng, Xu Chen, and Xiao-Hua Zhou. Relaxing the accurate imputation assumption in doubly robust learning for debiased collaborative filtering. In Proc. Int. Conf. Mach. Learn. , pages 2944829460, 2024.

Yingzhen Li and Richard E. Turner. Gradient estimators for implicit models. In Proc. Int. Conf. Learn. Represent. , pages 1-19, 2018.

Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In Proc. Int. Conf. Mach. Learn. , pages 4082-4092. PMLR, 2019.

Diju Liu, Yalin Wang, Chenliang Liu, Xiaofeng Yuan, Kai Wang, and Chunhua Yang. Scopefree global multi-condition-aware industrial missing data imputation framework via diffusion transformer. IEEE Trans. Knowl. Data Eng. , pages 1-12, 2024. doi: 10.1109/TKDE.2024. 3392897.

Qiang Liu. Stein variational gradient descent as gradient flow. In Proc. Adv. Neural Inf. Process. Syst. , volume 30, pages 1-15, 2017.

Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. In Proc. Adv. Neural Inf. Process. Syst. , volume 29, pages 1-13, 2016.

Weiming Liu, Jiajie Su, Chaochao Chen, and Xiaolin Zheng. Leveraging distribution alignment via stein path for cross-domain cold-start recommendation. In Proc. Adv. Neural Inf. Process. Syst. , volume 34, pages 19223-19234, 2021.

Weiming Liu, Xiaolin Zheng, Jiajie Su, Longfei Zheng, Chaochao Chen, and Mengling Hu. Contrastive proxy kernel stein path alignment for cross-domain cold-start recommendation. IEEE Trans. Knowl. Data Eng. , 35(11):11216-11230, 2023. doi: 10.1109/TKDE.2022.3233789.

Yixin Liu, Thalaiyasingam Ajanthan, Hisham Husain, and Vu Nguyen. Self-supervision improves diffusion models for tabular data imputation. In Proc. ACM Int. Conf. Inf. Knowl. Manag. , pages 1-10, 2024.

Pierre-Alexandre Mattei and Jes Frellsen. MIWAE: Deep generative modelling and imputation of incomplete data sets. In Proc. Int. Conf. Mach. Learn. , pages 4413-4423, 2019.

Boris Muzellec, Julie Josse, Claire Boyer, and Marco Cuturi. Missing data imputation using optimal transport. In Proc. Int. Conf. Mach. Learn. , pages 7130-7140, 2020.

Yidong Ouyang, Liyan Xie, Chongxuan Li, and Guang Cheng. MissDiff: Training diffusion models on tabular data with missing values. In Proc. Int. Conf. Mach. Learn. Workshop on Structured Probabilistic Inference & Generative Modeling , 2023.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Proc. Adv. Neural Inf. Process. Syst. , volume 32, pages 1-12, 2019.

Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941 , pages 1-13, 2017.

Donald B Rubin. Inference and missing data. Biometrika , 63(3):581-592, 1976.

Walter Rudin et al. Principles of mathematical analysis , volume 3. McGraw-hill New York, 1964.

Filippo Santambrogio. { Euclidean, Metric, and Wasserstein } gradient flows: an overview. Bull. Math. Sci. , 7:87-154, 2017.

Simo Särkkä and Arno Solin. Applied stochastic differential equations , volume 10. Cambridge University Press, 2019.

Jiaxin Shi, Chang Liu, and Lester Mackey. Sampling with Mirrored Stein Operators. Proc. Int. Conf. Learn. Represent. , pages 1-26, 2022.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Proc. Conf. Uncertainty in Artificial Intelligence , pages 574-584, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Proc. Int. Conf. Learn. Represent. , pages 1-36, 2020.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-based diffusion models for probabilistic time series imputation. Proc. Adv. Neural Inf. Process. Syst. , pages 24804-24816, 2021.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Comput. , 23(7):1661-1674, 2011.

Dilin Wang and Qiang Liu. Nonlinear stein variational gradient descent for learning diversified mixture models. In Proc. Int. Conf. Mach. Learn. , pages 6576-6585, 2019.

Fangyikang Wang, Hubery Yin, Yuejiang Dong, Huminhao Zhu, Chao Zhang, Hanbin Zhao, Hui Qian, and Chen Li. BELM: Bidirectional explicit linear multi-step sampler for exact inversion in diffusion models. In Proc. Adv. Neural Inf. Process. Syst. , pages 1-33, 2024.

Fangyikang Wang, Huminhao Zhu, Chao Zhang, Hanbin Zhao, and Hui Qian. GAD-PVI: A general accelerated dynamic-weight particle-based variational inference framework. In Proc. AAAI Conf. Artif. Intell. , pages 15466-15473, 2024.

Hao Wang, Jiajun Fan, Zhichao Chen, Haoxuan Li, Weiming Liu, Tianqiao Liu, Quanyu Dai, Yichao Wang, Zhenhua Dong, and Ruiming Tang. Optimal transport for treatment effect estimation. In Proc. Adv. Neural Inf. Process. Syst. , pages 5404-5418, 2023.

Hao Wang, Zhichao Chen, Zhaoran Liu, Haozhe Li, Degui Yang, Xinggao Liu, and Haoxuan Li. Entire space counterfactual learning for reliable content recommendations. IEEE Trans. Inf. Forensics Secur. , pages 1-1, 2024.

Hao Wang, Zhichao Chen, Zhaoran Liu, Licheng Pan, Hu Xu, Yilin Liao, Haozhe Li, and Xinggao Liu. SPOT-I: Similarity preserved optimal transport for industrial iot data imputation. IEEE Trans. Ind. Inform. , pages 1-9, 2024. doi: 10.1109/TII.2024.3452241.

Jun Wang, Wenjie Du, Wei Cao, Keli Zhang, Wenjia Wang, Yuxuan Liang, and Qingsong Wen. Deep learning for multivariate time series imputation: A survey. arXiv preprint arXiv:2402.04059 , pages 1-9, 2024.

Xu Wang, Hongbo Zhang, Pengkun Wang, Yudong Zhang, Binwu Wang, Zhengyang Zhou, and Yang Wang. An observed value consistent diffusion model for imputing missing values in multivariate time series. In Proc. ACM SIGKDD Int. Conf. Knowl. Discovery Data Mining , pages 2409-2418, 2023.

Yifei Wang and Wuchen Li. Accelerated information gradient flow. J. Sci. Comput. , 90:1-47, 2022.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proc. Int. Conf. Mach. Learn. , pages 681-688. Citeseer, 2011.

Chen Xu, Xiuyuan Cheng, and Yao Xie. Normalizing flow neural networks by JKO scheme. In Proc. Adv. Neural Inf. Process. Syst. , pages 47379-47405, 2023.

Jingwen Xu, Fei Lyu, and Pong C Yuen. Density-aware temporal attentive step-wise diffusion model for medical time series imputation. In Proc. ACM Int. Conf. Inf. Knowl. Manag. , pages 2836-2845, 2023.

Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Comput. Surv. , 56(4):1-39, 2023.

Yiyuan Yang, Ming Jin, Haomin Wen, Chaoli Zhang, Yuxuan Liang, Lintao Ma, Yi Wang, Chenghao Liu, Bin Yang, Zenglin Xu, et al. A survey on diffusion models for time series and spatio-temporal data. arXiv preprint arXiv:2404.18886 , pages 1-27, 2024.

Jinsung Yoon, James Jordon, and Mihaela Schaar. GAIN: Missing data imputation using generative adversarial nets. In Proc. Int. Conf. Mach. Learn. , pages 5689-5698. PMLR, 2018.

Shipeng Yu, Kai Yu, Volker Tresp, Hans-Peter Kriegel, and Mingrui Wu. Supervised probabilistic principal component analysis. In Proc. ACM SIGKDD Int. Conf. Knowl. Discovery Data Mining , pages 464-473, 2006.

Chao Zhang, Zhijian Li, Xin Du, and Hui Qian. Dpvi: A dynamic-weight particle-based variational inference framework. In Proc. Int. Joint Conf. Artif. Intell. , pages 4900-4906, 2022.

Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Trans. Pattern Anal. Mach. Intell. , 41(8):2008-2026, 2019. doi: 10.1109/ TPAMI.2018.2889774.

Yulong Zhang, Shuhao Chen, Weisen Jiang, Yu Zhang, Jiangang Lu, and James T Kwok. Domain-guided conditional diffusion model for unsupervised domain adaptation. arXiv preprint arXiv:2309.14360 , pages 1-13, 2023.

He Zhao, Ke Sun, Amir Dezfouli, and Edwin V Bonilla. Transformed distribution matching for missing value imputation. In Proc. Int. Conf. Mach. Learn. , pages 42159-42186, 2023.

Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular data. In Proc. Adv. Neural Inf. Process. Syst. Workshop on First Table Representation , 2022.

Huminhao Zhu, Fangyikang Wang, Chao Zhang, Hanbin Zhao, and Hui Qian. Neural sinkhorn gradient flow. arXiv preprint arXiv:2401.14069 , pages 1-17, 2024.

Zhan Zhuang, Yu Zhang, and Ying Wei. Gradual domain adaptation via gradient flow. In Proc. Int. Conf. Learn. Represent. , pages 1-27, 2024.

Zhan Zhuang, Yulong Zhang, Xuehao Wang, Jiangang Lu, Ying Wei, and Yu Zhang. TimeVarying LoRA: Towards effective cross-domain fine-tuning of diffusion models. In Proc. Adv. Neural Inf. Process. Syst. , pages 1-25, 2024.

Appendix Contents

Appendix A Detailed Preliminaries of Wasserstein Gradient Flow

In this section, we want to introduce the WGF framework and its application scenarios to better understand this paper. Before introduction, the following concepts are listed to better understand the WGF framework:

Wasserstein Metric: Let P 2 ( R D ) represent the space of probability measures on R D that possess finite second moments. Formally, this is expressed as P 2 ( R D ) = { µ ∈ M ( R D ) | ∫ ‖ x ‖ 2 d µ ( x ) < ∞} , where M ( R D ) denotes the set of all probability measures on R D . Considering any two probability measures µ, ν ∈ P 2 ( R D ) , we define the Wassersteinp distance between µ and ν as follows:

Here, Γ( µ, ν ) represents the collection of all joint distributions (couplings) between µ and ν . For every joint distribution π ∈ Γ( µ, ν ) , it holds that µ ( x ) = ∫ R D π ( x, y ) d y and ν ( y ) = ∫ R D π ( x, y ) d x . The integral on the right-hand side encapsulates the transportation cost in the optimal transport (OT) problem, framed by Kantorovich's formulation, where π ∗ denotes the optimal transportation plan.

Furthermore, leveraging Jensen's inequality facilitates demonstrating the monotonicity of the Wassersteinp distance, affirming that for 1 ≤ p ≤ q , the relationship W p ( µ, ν ) ≤ W q ( µ, ν ) invariably holds. Building on this principle, we can articulate the inner product within the measurable space ( P 2 ( R D ) , W ) as delineated below:

Gradient Flow in Wasserstein Space: Consider a functional F associated with µ ∈ P 2 ( R D ) . Our objective is to identify the optimal µ that minimizes F :

To facilitate the decrease of F ( µ ) , we introduce a velocity field u µ : R D → R D designed to expedite the reduction of F ( µ ) as µ evolves under this field. Utilizing the chain rule yields:

where δ represents the first variation operator. To ensure the decrease of F ( µ ) , i.e., d F ( µ ) d τ ≤ 0 , the velocity field is defined as:

The decline of F ( µ ) aligns with the following partial differential equation (PDE) called the continuity equation:

Hence, the continuity equation Eq. (A.6), coupled with the velocity field articulated in Eq. (A.5), is recognized as the Wasserstein Gradient Flow , delineating the steepest descent direction of cost functional F ( µ ) in the Wasserstein space.

Simulation of WGF & Sampling: There are primarily two discretization techniques for the WGF: the forward scheme and the backward scheme.

Forward Scheme: The forward scheme applies gradient descent within the Wasserstein space to identify the direction of the steepest descent. For an energy functional F ( µ ) with a specified step size η , the update rule in the forward scheme is formulated as:

facilitating an intuitive and direct update mechanism that emulates the gradient flow in the Euclidean space but transposed into the Wasserstein space.

Backward Scheme: Conversely, the backward scheme, often referred to as the JordanKinderlehrer-Otto scheme [23], represents a more implicit discretization approach. It defines the subsequent distribution µ τ +1 by solving an optimization problem that balances the energy decrease and the transportation cost. This scheme is mathematically denoted as:

thereby integrating the energy minimization and transport efficiency into a single variational problem that reflects the inherent structure of the Wasserstein space.

These schemes provide distinct yet complementary approaches to discretizing the dynamics defined by WGFs, offering different perspectives and tools for the analysis and computation of these flows.

Leveraging the WGF framework, if we designate the functional F to be the KL divergence, it yields a particular formulation for the velocity field.

On this basis, plug Eq. (A.9) into Eq. (A.6), we can get the following PDE:

According to Theorem 5.4 in reference [47], denote the random sample from distribution p as x , we can obtain the following stochastic differential equation (SDE) called Langevin equation [62] for implementing this gradient flow easily:

where d W τ is the standard Wiener process (also known as Brownian motion).

Appendix B Detailed Information for Toy Cases in Section 3.1

To investigate what would happen if we directly applied diffusion models to MDI tasks, we consider the following optimization problem:

which corresponds to the density function of a Dirichlet distribution, Dir([2 . 5 , 2 . 5 , 5 . 0]) , where a h lies on the three-dimensional standard simplex ∆ 2 . The optimal value of F Dir is given by:

To optimize this cost functional, we employ the Langevin equation as presented in Eq. (A.11):

and compare the results to Eq. (B.2) to evaluate effectiveness. Additionally, since the support is on a three-dimensional standard simplex ∆ 2 , to ensure the well-definedness of our approach, we use the mirror descent technique [48, 20], where the Bregman function is defined as ψ ( a h ) = ∑ 3 k =1 a h,k log a h,k -a h,k . Moreover, the optimization of a h,k | H h =1 | 3 k =1 is conducted by simulating the Langevin equation, which is discretized by the Euler-Maruyama method as follows:

With the step size η set to 5 . 0 × 10 -3 , we repeatedly execute Eq. (B.4) 100 times, culminating in the results depicted in Fig. 1 (b).

Appendix C Theoretical Analysis

C.1 Implementation Difficulty of Velocity Field

To the best of our knowledge, the difficulty of implementing the velocity field can be given from two perspectives, namely ODE-based implementation and SDE-based implementation. In this section, we want to discuss these two implementation approaches in detail.

ODE-based Implementation:

WGFframework: According to the continuity equation, we can obtain the following velocity field:

where (i) is based on Section 2.3, and (ii) is based on Section 4.1. The expression of the velocity field involves the computation of density term r ( X (miss) ) [31, 9], which is intractable during practice as we stated in Section 2.3. Based on this, we conclude that implementing this velocity field within the WGF framework is difficult.

Probability flow ODE: According to reference [50], if we directly plug Eq. (8) into the FPK equation, we can get the following PDE:

When we set σ τ as 0 , we can find that the corresponding ODE is Eq. (C.1), where we are obliged to compute the intractable density r ( X (miss) ) .

SDE-based Implementation:

If we plug Section 4.1 into the FPK equation, the corresponding PDE can be given as follows:

where the coefficient before the Laplacian operator ∇· ∇ is -1 . To the best of our knowledge, this structure makes deriving a corresponding SDE impossible by current approaches.

C.2 Proof & Discussions of Concerning Propositions and Corollaries

Proposition (3.1) . Within WGF framework, DM-based MDI approaches can be viewed as finding the imputed values X ( imp ) that maximize the following objective:

where 'const' is the abbreviation of constant, and ψ ( X ( miss ) ) is a scalar function determined by the type of SDE underlying the DMs.

VE-SDE: ψ ( X ( miss ) ) = 1 2 H [ r ( X (miss) )] ≥ 0

where H [ r ( X (miss) )] := -∫ r ( X (miss) ) log r ( X (miss) )d X (miss) is the entropy term, γ τ is determined by noise scale β τ : γ τ := (1 -exp( -2 ∫ τ 0 β s d s )) > 0 , 0 < β 1 < · · · < β T < 1 .

Proof. Since there are various approaches for reversing the sampling procedure of DMs, for simplicity, as we emphasized in Section 3.2, we mainly consider the VP-SDE, VE-SDE, and sub-VP-SDE as analysis objects in this paper.

VP-SDE: According to reference [50], the density evolution of the generative process for VP-SDE can be delineated by the following PDE:

where β τ ∈ (0 , 1) is the time-varying noise scale. On this basis, according to [24], by changing the variable as d τ := β τ 2 d τ , we can get the following equation:

Comparing Eq. (C.6) with Eqs. (A.5) and (A.6), the cost functional to be minimized of this simulation procedure can be given as follows:

Note that 1 4 [ X (miss) ] glyph[latticetop] [ X (miss) ] ≥ 0 and -1 2 ∫ r ( X (miss) ) log r ( X (miss) )d X (miss) ≥ 0 hold, and thus the proposition for VP-SDE is proved by taking the negative of the abovementioned equation.

VE-SDE: Similarly, based on reference [50], the following PDE can be given to delineate the density evolution of the generative process for VE-SDE:

where σ 2 τ is a time varying noise scale.

As such, by chaning the variable as d τ := [ d σ 2 τ d τ ] d τ [24], Eq. (C.8) can be reformulated as follows:

Comparing Eq. (C.9) with Eqs. (A.5) and (A.6), the cost functional to be minimized of this simulation procedure can be given as follows:

Note that the entropy function -1 2 ∫ r ( X (miss) ) log r ( X (miss) )d X (miss) ≥ 0 holds, and thus the proposition for VE-SDE is proved by taking the negative of the abovementioned equation.

sub-VP-SDE: Based on reference [50], the following PDE can be given to delineate the density evolution of the generative process for sub-VP-SDE:

where γ τ := (1 -exp( -2 ∫ τ 0 β s d s )) > 0 . On this basis, by chaning the variable as d τ := β τ 2 d τ , we can get the following equation:

Comparing Eq. (C.12) with Eqs. (A.5) and (A.6), the cost functional to be minimized of this simulation procedure can be given as follows:

Note that 1 4 γ τ [ X (miss) ] glyph[latticetop] [ X (miss) ] ≥ 0 and -1 2 ∫ r ( X (miss) ) log r ( X (miss) )d X (miss) ≥ 0 hold, and thus the proposition for sub-VP-SDE is proved by taking the negative of the abovementioned equation.

In summary, the regularization term ψ ( X ( miss ) ) for VP-SDE is E r ( X (miss) ) { 1 4 [ X (miss) ] glyph[latticetop] [ X (miss) ] } + 1 2 H [ r ( X (miss) )] , for VE-SDE is 1 2 H ( r ( X (miss) )) , and for sub-VP-SDE is E r ( X (miss) ) { 1 4 γ τ [ X (miss) ] glyph[latticetop] [ X (miss) ] } + 1 2 H [ r ( X (miss) )] .

Before proving Proposition 4.1, we want to introduce the following lemma to delineate the evolution of cost functional F NER along time τ :

Lemma C.1. The evolution of F NER along time τ can be characterized by the following ODE, assuming that the boundary condition lim X (miss) →∞ [ u ( X (miss) ) r ( X (miss) )] = 0 is satisfied:

This boundary condition is achievable, for instance, when r ( X ( miss ) ) is bounded, and the limit of the velocity field as the norm of X ( miss ) approaches infinity is zero ( lim ‖ X (miss) ‖→∞ u ( X (miss) ) = 0 ).

Proof. Before proving this lemma, we should recognize that the evolution of X ( miss ) should promise the probability density function r ( X ( miss ) ) unchanged. In other words, the following continuity equation should be satisfied during the optimization of r ( X ( miss ) ) :

On this basis, the evolution of F NER along time τ , d F NER d τ , can be given as follows based on the chain rule:

where (i) and (ii) are based on integration by parts. More specifically, when condition lim X (miss) →∞ [ u ( X (miss) ) r ( X (miss) )] = 0 is satisfied, for example, r ( X ( miss ) ) is bounded, and the limit of the velocity field as the norm of X ( miss ) approaches infinity is zero ( lim ‖ X (miss) ‖→∞ u ( X (miss) ) = 0 ), we can get the following result [35, 32]:

where the left-hand-side can be further decomposed as follows based on the integration by parts:

Based on Lemma C.1, we can now start proving Proposition 4.1:

Proposition (4.1) . Suppose u ( X (miss) ) is a velocity field regularized by the RKHS norm under the following conditions: 1). The kernel function satisfies: lim ‖ X (miss) ‖→∞ K ( X (miss) , ˜ X (miss) ) = 0 . 2). The density r ( X (miss) ) is bounded. Then, the velocity field that minimizes the cost functional F NER = E r ( X ( miss ) ) [log ˆ p ( X ( miss ) | X ( obs ) )] -λ H [ r ( X (miss) )] can be given by:

where the expectation term E r ( ˜ X ( miss ) ) can be efficiently estimated using Monte Carlo approximation.

Proof. When the velocity is regularized by the RKHS norm, we can first reformulate Eq. (8) as follows to find the steepest direction for the sake of improving F NER:

Based on this, assume we have a map function φ ( x ) , the kernel function can be given as follows:

Hence, the regularization term that control the magnitude of u ( X (miss) ) can be given by 1 2 ‖ u ( X (miss) ) ‖ 2 H , and the spectral decomposition of kernel function can be given as follows:

where φ i ( · ) indicates the orthonormal basis and ξ i is the corresponding eigen-value. For any function u ( X (miss) ) ∈ H , the following decomposition is given:

where u i and ∑ ∞ i =1 ‖ u i ‖ 2 < ∞ .

The learning objective defined in Eq. (8) can be reformulated as follows:

Take the right-hand-side of (i) with-respect-to u i , and set it to 0 , we can get:

On this basis, u ∗ i can be given as follows:

and hence, u ( X (miss) ) can be given as follows:

Proposition (4.2) . Assume that the proposal distribution r ( X (joint) ) is factorized by r ( X (joint) ) := r ( X (miss) ) p ( X (obs) ) . The cost functional associated with the joint distribution is defined as follows:

which leads to the velocity field delineated in Eq. (9) and establishes F joint-NER as a lower bound for F NER, with the difference being a constant (i.e., F joint-NER = F NER -const , const ≥ 0 ).

Before proving this proposition, we want to first clarify the justification of the assumption that r ( X (joint) ) := r ( X (miss) ) r ( X (obs) ) = r ( X (miss) ) p ( X (obs) ) . In this part, we set r ( X (obs) ) = p ( X (obs) ) . Before stating the justification, we should come to the following agreements:

Throughout the imputation procedure, X (obs) remains invariant regardless of any modifications to X (miss) .

Given this invariance, it is justified to state that r ( X (obs) ) remains constant from the perspective of particle variational inference represented by reference [35, 34], and consequently r ( X (obs) | X (miss) ) = r ( X (obs) ) , reflecting the independence of X (obs) from X (miss) .

Based on this, we want to show that within the WGF framework, the factorizations r ( X (miss) ) r ( X (obs) ) and r ( X (miss) | X (obs) ) r ( X (obs) ) are equivalent. To this end, let us write down the evolution of r ( X (obs) , X (miss) ) along time τ as follows when r ( X (obs) , X (miss) ) is factorized by r ( X (obs) , X (miss) ) = r ( X (miss) | X (obs) ) r ( X (obs) ) :

where the first underbrace is the Bayesian formula, and the second underbrace is based on the abovementioned Agreement 2. Consequently, we can further obtain the following results:

Similarly, when we factorize r ( X (obs) , X (miss) ) by r ( X (obs) , X (miss) ) = r ( X (miss) ) r ( X (obs) ) , we can get the following result:

Comparing Eq. (C.28) to Eq. (C.27), we can demonstrate our justification of the factorization r ( X (obs) , X (miss) ) = r ( X (miss) ) r ( X (obs) ) . Finally, we would like to conclude with a metaphor to further illustrate the plausibility of this mean-filed factorization, which has been widely applied in variational inference [5]:

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 ( X (miss) ) not containing information about 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 ( X (miss) ) is appropriately molded by the directorial guidance (mirrors the continuity equation ∂r ∂τ = -∇· ( ur ) ) of the velocity field u and the script provided by the critic p ( X (obs) , X (miss) ) / p ( X (obs) | 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 ( X (joint) ) = r ( X (miss) ) r ( X (obs) ) ).

Based on the abovementioned analysis, we can now start the proof of Proposition 4.2:

Proof. Our proof will be divided into two parts namely 'velocity field derivation' and 'upper bound acquirement'.

Velocity Field Derivation:

the following continuity equation should be satisfied during the optimization of r ( X ( miss ) ) :

where (i) is based on the fact that X (obs) remains unchanged during the imputation process. Thus, according to Eq. (C.16), the evolution of F joint-NER along time τ , d F joint-NER d τ , can be given as follows based on the chain rule:

where (i) and (ii) are based on integration by parts. More specifically, when condition lim X (joint) →∞ [ u ( X (joint) ) r ( X (joint) )] = 0 is satisfied, for example, r ( X ( joint ) ) is bounded, and the limit of the velocity field as the norm of X ( joint ) approaches infinity is zero ( lim ‖ X (joint) ‖→∞ u ( X (joint) ) = 0 ), we can get the following result [35, 32]:

where we omit the gradient operator with respect to the observed variables X (obs) , denoted as ∇ X (obs) , because X (obs) remains constant during the imputation process. This constancy implies that the divergence ∇ X (obs) · [ r ( X (joint) ) u ( X (joint) )] = 0 . Consequently, the left-hand-side of this equation can be further decomposed as follows based on the integration by parts:

Similar to the proof of proposition 4.1, we can restrict the velocity field in RKHS and find the steepest gradient boosting direction as follows according to Eqs. (C.19) to (C.21):

Take the right-hand-side of (i) with-respect-to u i , and set it to 0 , we can get:

On this basis, u ∗ i can be given as follows:

and hence, u ( X (joint) ) can be given as follows:

Lower Bound Acquirement:

Before starting the proving of this part, we should notice that given the unchanged observational data X (obs) , the distribution p ( X (obs) ) is a constant. On this basis, consider the definition of F NER (right-hand-side of Eq. (7)), the first term and the second term are denoted by 'term 1' and 'term 2' for simplicity:

For term 1, we can obtain the following derivation:

Combine Eqs. (C.36) and (C.37), we can obtain the following relationship: F -F and constant const is greater than 0 .

Corollary (4.3) . The following equation holds: u ( X (joint) ) = u ( X (miss) ) .

Proof. This corollary can be easily proven by according to Eq. (C.38):

Plugging Eq. (C.39) into Eqs. (A.5) and (A.6), we can see that the velocity fields for X (miss) within functional F NER and F joint-NER are identical.

Appendix D Detailed Explanation for the Workflow of NewImp Approach

In this section, we intend to provide detailed information about the implementation of the NewImp approach in Algorithm 1. We will focus on two primary aspects: 1) the numerical implementation of ODE simulation, and 2) the DSM algorithm.

D.1 Forward Euler's Method for ODE Simulation

During step 7 of Algorithm 1, we involve the simulation of the ODE defined by Eqs. (9) and (14). To simulate this ODE we use the forward Euler's method [6] in this paper for simplicity. Specifically, suppose we have the following ODE:

and the initial value at τ = 0 is given x 0 = x init, the value at time η can be derived as follows:

To alleviate the intergal term, the forward Euler's method attempts to approximate the integral term to summation term as follows:

On this basis, the value at time T can be obtained by repeating Eq. (D.3) from τ = 0 to τ = T , which is the forward Euler's method.

Algorithm 2 Algorithm for Forward Euler's Method

1: Input: ODE f ( x τ ) ; start point 0 ; end point T ; step size η ; initial value x 0 .

2: Output: Predicted value x T at τ = T .

4: for τ = 0 + η, 0 + 2 η, . . . , 0 + jη do

3: j ← T -0 η

5: x τ ← x τ -η + f ( x τ -η ) × η

6: end for

D.2 Detailed Information for DSM

During step 5 of Algorithm 1, we involve the training of ∇ X (joint) log ˆ p ( X (joint) ) using the DSM function. In this subsection, we aim to further elaborate on the detailed algorithm for the DSM function to uphold the completeness of this manuscript. As mentioned in Section 2.2, the score function ∇ X (joint) log ˆ p ( X (joint) ) is typically parameterized by a neural network. For simplicity, we denote the parameter set of ∇ X (joint) log ˆ p ( X (joint) ) by θ .

glyph[triangleright] Calculate the Number of Steps

Algorithm 3 DSM for ∇ X (joint) log ˆ p ( X (joint) ) Training

1: Input: joint data X (joint) .

2: Hyperparameters: neural network learning rate lr , training epoch E , and neural network hidden unit HUscore.

3: for e = 1 to E do

4: ˆ X (joint) ← X (joint) + glyph[epsilon1], glyph[epsilon1] ∼ N ( 0 , σ 2 I )

5:

(joint)

ˆ

X

(joint)

log

q

σ

( ˆ

X

6: L DSM ← Eq. (15)

∇

|

X

)

←-

7: θ ← ApplyGradient ( ∇ θ L DSM , lr )

8: end for

Appendix E Detailed Information for Experiments

E.1 Background & Simulation of Missing Data

In this paper, we consider the datasets listed in Table E.1 as our experimental datasets. Based on this, according to reference [44], missing data can be classified into three categories: Missing at Random (MAR), where the likelihood of missing data depends solely on observed data; Missing Completely at Random (MCAR), where the absence of data is completely unrelated to any observed or unobserved variables; and Missing Not at Random (MNAR), where missingness is influenced by unobserved data. In the cases of MCAR and MAR, the patterns of missing data are considered 'ignorable' because it is unnecessary to explicitly model the distribution of the missing values. Conversely, MNAR scenarios, where missing data can introduce significant biases that are not easily corrected without imposing domain-specific assumptions, constraints, or parametric forms on the missingness mechanism, present more complex challenges [40, 22]. Therefore, our discussion is primarily focused on numerical tabular data within the MCAR and MAR contexts.

To simulate missing data, we adopt the methodologies outlined in reference [22]:

MAR: Initially, a random subset of features is selected to remain non-missing. The masking of the remaining features is conducted using a logistic model, which employs the non-missing features as predictors. This model is parameterized with randomly selected weights, and the bias is adjusted to achieve the desired missingness rate.

MCAR: For each data point, the masking variable is generated from a Bernoulli distribution with a predetermined fixed mean, ensuring that the probability of missingness is the same across all data points.

MNAR: Although MNAR scenarios are not the primary focus of this manuscript, we include experiments in this context. Missingness is introduced either by additional masking of the MARselected features using a Bernoulli process with a fixed mean, or through direct self-masking of values using interval-censoring techniques. In this paper, we mainly consider the former strategy. In other words, the mechanism of MNAR we used in this paper is identical to the previously described MARmechanism, but the inputs of the logistic model are then masked by an MCAR mechanism.

(joint)

glyph[triangleright] Data Noising lr

ˆ

X

(joint)

-

σ

2

X

(joint)

glyph[triangleright]

Apply the Gradient with Learning Rate

E.2 Hyperparameter Setting of Baseline Models

In this subsection, we want to report the baseline models' hyperparameter settings to ensure the reproducibility of our paper:

Batch-Size-Related: The batch size for ReMasker is set to 64. For other baseline models, it is uniformly set at 512. (Notably, for Sink and TDM, if N < 512 , the batch size is set to 2 glyph[floorleft] N 2 glyph[floorright] . )

· Hidden-Unit-Related:

-The MIWAE model features a latent dimension of 16 and 32 hidden units.

-The TDM model includes 16 hidden units per layer with the number of layers set to 2.

-MIRACLE's hidden units are set to 32.

-For ReMasker, the embedding dimension is 32, depth is 6, mask ratio is 0.5, encoder depth is 6, decoder depth is 4, number of heads is 4, and the multi-layer perceptron ratio is 4.0.

-For MissDiff and CSDI_T, the channel size is set as 16, the embedding dimension is set to 128, and the layer number is set as 2.

-For the GAIN model, for both the generator and the discriminator, the hidden size is set to 2 × D , and the number of hidden layers is set to 3.

Diffusion-Hyperparameters-Related: The diffusion step is set at 100 and the particle number at 50 for MissDiff and CSDI_T.

Appendix F Additional Empirical Evidence

F.1 Toy Case Experiments

x

x

1

2

0.00

0.02

0.04

0.06

0.08

0.10

0.12

0.14

0.16

(a) Standard Gaussian

x

1

x

2

0.000

0.016

0.032

0.048

0.064

0.080

0.096

0.112

0.128

(b) Student's-

t

x

1

x

2

0.00

0.01

0.02

0.03

0.04

0.05

0.06

0.07

0.08

0.09

(c) Gaussian Mixture

x

1

x

2

0.084

0.072

0.060

0.048

0.036

0.024

0.012

0.000

To demonstrate the effectiveness of the NewImp method vary different type of distributions, we evaluate it across four distinct toy cases, each characterized by different distributions:

Standard Gaussian: X (ideal) ∼ N (0 , I 2 × 2 ) .

Student'st (a heavy-tailed distribution): X (ideal) ∼ Stt (0 , [ 1 0 . 5 0 . 5 1 ] ) .

Gaussian Mixture: X (ideal) ∼ 1 3 ×N ([1 , 2] , [ 0 . 5 0 0 0 . 5 ] ) + 1 3 ×N ([ -1 , -2] , [ 0 . 5 0 . 1 0 . 1 0 . 5 ] ) + 1 3 × N ([2 , -2] , [ 0 . 3 0 0 0 . 3 ] ) .

Skewed Gaussian (via exponential transformation): X (ideal) = exp( glyph[epsilon1] ) , where glyph[epsilon1] ∼ N (0 , I 2 × 2 ) .

Based on this, we display the contours of their density values in Fig. F.1, and we list the imputation accuracy comparisons for MAR, MCAR, and MNAR scenarios with a 30% missing rate in Table F.1. The results indicate that our NewImp approach generally performs better on non-standard Gaussian type data, underscoring its universality and applicability. This enhanced performance is attributable to our modeling strategy, which involves modeling the score function of the data [52, 50], which eliminates the need for normalization, and consequently results in the NewImp approach can perform well on complex data distributions, including skewed, heavy-tailed, and mixture distributions.

F.2 Additional Experimental Results with MNAR Scenario

In this section, we expand upon the results presented in Table 1 by including the MNAR scenario, as detailed in Table F.2. Additionally, we report on the outcomes of an ablation study and sensi-

tivity analysis in Tables F.4 and F.5 and Fig. F.2. These extended results lead to several pertinent observations:

Across three different missing data scenarios, the models consistently exhibit the poorest performance under the MNAR condition. For instance, in the MNAR scenario, nearly all models show a significant decrease in imputation accuracy and an increase in standard deviation. This supports the assertion made in Appendix E.1 that addressing the MNAR scenario requires the incorporation of relevant domain knowledge to mitigate biases introduced by the pattern of missing data.

The findings from the ablation study under the MNAR scenario are consistent with those observed in both MAR and MCAR scenarios in Section 5.3. This consistency underscores the importance of including the NER term and adopting the joint distribution modeling approach.

Similarly, the results from the sensitivity analysis under the MNAR scenario align with those from MAR and MCAR scenarios in Section 5.4. This alignment reinforces our interpretations of model performance across different groups of hyperparameters under MAR and MCAR scenarios.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

0.1 0.25 0.5

1.0

2.0

5.0 10.0

h

0.2

0.4

0.6

0.8

WASS

MAE

32

64

128

256

512

HUscore

0.2

0.3

0.4

WASS

MAE

1.0

2.0

5.0

10.0 15.0 20.0

λ

0.2

0.3

0.4

0.5

WASS

MAE

0.005

0.05

0.1

0.2

0.5

η

0.2

0.3

0.4

WASS

MAE

(a) MAR with 30% missing rate at CC dataset.

0.1 0.25 0.5

1.0

2.0

5.0 10.0

h

0.4

0.6

0.8

1.0

WASS

MAE

32

64

128

256

512

HUscore

0.35

0.40

0.45

0.50

WASS

MAE

1.0

2.0

5.0

10.0 15.0 20.0

λ

0.35

0.40

WASS

MAE

0.005

0.05

0.1

0.2

0.5

η

0.35

0.40

WASS

MAE

(b) MCAR with 30% missing rate at CC dataset.

0.1 0.25 0.5 1.0 2.0 5.0 10.0

h

0.4

0.6

0.8

WASS

MAE

32

64

128

256

512

HUscore

0.30

0.35

0.40

WASS

MAE

1.0

2.0

5.0

10.0 15.0 20.0

λ

0.3

0.4

0.5

WASS

MAE

0.005

0.05

0.1

0.2

0.5

η

0.30

0.35

0.40

WASS

MAE

(c) MNAR with 30% missing rate at CC dataset.

F.3 Empirical Evidence for Selecting RBF Function

In our derivation process, we specifically selected the RBF kernel to satisfy the 'zero boundary condition': lim X (joint) →∞ K ( X (joint) , ˜ X (joint) ) = 0 for the sake of avoiding the explicit density estimation of the intractable proposal distribution r ( X (joint) ) . This selection prompts an additional inquiry: What if we replaced the RBF kernel with another that does not fulfill the 'zero boundary condition'? To maintain the rigor of our analysis, we compared the performance of the NewImp with alternative kernel functions under identical settings. Consequently, we consider the following types of kernel functions:

linear kernel function (linear): K ( X (joint) , ˜ X (joint) ) = [ X (joint) ][ ˜ X (joint) ] glyph[latticetop]

polynomial kernel function (poly): K ( X (joint) , ˜ X (joint) ) = { [ X (joint) ][ ˜ X (joint) ] glyph[latticetop] } 2

sigmoid kernel function (sigmoid): K ( X (joint) , ˜ X (joint) ) = tanh { [ X (joint) ][ ˜ X (joint) ] glyph[latticetop] }

cosine similarity kernel function (cos): K ( X (joint) , ˜ X (joint) ) = { [ X (joint) ] ‖ [ X (joint) ] ‖ }{ [ ˜ X (joint) ] ‖ [ ˜ X (joint) ] ‖ } glyph[latticetop]

sine kernel function (sin): K ( X (joint) , ˜ X (joint) ) = sin( | X (joint) -˜ X (joint) | 2 2 )

The experimental results are detailed in Table F.6 (For completeness, we also report the results under the MNAR scenario). From the results, it is evident that other kernel functions, which do not meet the 'zero boundary condition', perform significantly worse compared to the RBF kernel. This demonstrates the critical importance of selecting the appropriate kernel function for achieving accurate imputation results, thereby validating the choice of the RBF kernel for our NewImp approach.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that RBF kernel significantly outperform with p -value < 0 . 05 over paired samples t -test.

F.4 Time Complexity Analysis

In this section, we present an analysis of the complexity of time for our NewImp approach. The complexity analysis is based on the algorithms described in Algorithm 1. We begin by estimating the time complexity of the score function ∇ X (joint) log ˆ p ( X (joint) ) . Assuming the number of layers in the neural network that parameterize ∇ X (joint) log ˆ p ( X (joint) ) is L and each layer has an equal number of hidden units denoted as HUscore, the time complexity for the imputation algorithm defined in Algorithm 1 is detailed as follows:

DSM Training Part (step 5): Building on the previous item, the time complexity for the DSM training algorithm defined in Algorithm 3 is given as:

where the factor of 4 comprises three distinct components: backward propagation ( 1 ), forward propagation ( 1 ), and the acquisition of the sample-wise score function ( 2 ). Note that the network parameter size is substantially smaller than the number of data points, thereby making the forward computation of the score function the primary factor in time complexity.

Imputation Part (step 7):

Score function computation: The time complexity for computing the score function is expressed as:

where the factor 2 accounts for the backward propagation needed during the score function computation.

2.25

2.50

2.75

3.00

3.25

3.50

3.75

log10N

10

1

10

2

Time (seconds)

Estimate

Impute

Total

2.25

2.50

2.75

3.00

3.25

3.50

3.75

log10N

10

1

10

2

Time (seconds)

Estimate

Impute

Total

2.25

2.50

2.75

3.00

3.25

3.50

3.75

log10N

10

1

10

2

Time (seconds)

Estimate

Impute

Total

0.6

0.8

1.0

1.2

1.4

1.6

log10D

10

1

10

2

Time (seconds)

Estimate

Impute

Total

0.6

0.8

1.0

1.2

1.4

1.6

log10D

10

1

10

2

Time (seconds)

Estimate

Impute

Total

0.6

0.8

1.0

1.2

1.4

1.6

log10D

10

1

10

2

Time (seconds)

Estimate

Impute

Total

Kernel function and its gradient: Employing the RBF kernel K ( X , ˜ X ) := exp ( -‖ X -˜ X ‖ 2 2 h 2 ) , the gradient with respect to ˜ X is analytically determined as:

The time complexities for calculating the kernel function and its gradient are specified in Eqs. (F.4) and (F.5):

Based on the abovementioned analysis, we explore how computational complexity varies with different dataset sizes N and the number of features D , as shown in Figs. F.3 (a) and (b), respectively. From these figures, it is evident that computational time increases with the dataset size N . However, changes in the number of features D do not significantly affect the computation time. This observation underscores that the primary determinant of computational complexity in our context is the dataset size, aligning with our theoretical analysis, which indicates a quadratic relationship between time complexity and the size of the dataset N for the 'Imputation' part, and N glyph[greatermuch] D for the ' DSM Training' part, aligning with our theoretical analysis.

Moreover, the data reveals that the total computational time is predominantly governed by 'Estimation' part of our NewImp approach. This suggests that the training of the score function represents a critical bottleneck in the efficiency of the NewImp algorithm. Therefore, accelerating the NewImp algorithm crucially hinges on reducing the computational demands of the 'Estimation' part.

F.5 Convergence Analysis

In this section, we want to discuss the convergence of the proposed NewImp approach, prior to delving into this discussion, it is essential to establish a clear definition of convergence:

Definition F.1. A sequence {F 1 , F 2 , ..., F T } is said to be convergent if there exists a real number G such that for any given positive number ε ( ε > 0 ), there exists a positive integer N , such that for all indices n greater than N , the corresponding terms F n , n ≥ N satisfy the inequality |F n -G| < ε .

Based on Definition F.1, if a sequence is either monotonically increasing or monotonically decreasing and bounded (either bounded above or bounded below), then it is guaranteed to converge according

to the celebrated monotone convergence theorem (Section 3.14 in reference [45]). Based on this, the convergence of the 'Imputation' part (step 7 of Algorithm 1) and DSM training part (step 5 of Algorithm 1) are proposed in the proceeding parts.

0

250

500

τ

0.6

0.7

0.8

0.9

MAE

1

(a) MAE, BT.

0

250

500

τ

0.4

0.6

0.8

1.0

WASS

1

(b) WASS, BT.

0

250

500

τ

0.6

0.8

MAE

1

(c) MAE, BCD.

0

250

500

τ

2

3

4

WASS

1

(d) WASS, BCD.

0

250

500

τ

0.4

0.6

0.8

MAE

1

(e) MAE, CC.

0

250

500

τ

0.4

0.6

0.8

WASS

1

(f) WASS, CC.

0

250

500

τ

0.4

0.6

0.8

MAE

1

(g) MAE, CBV.

0

250

500

τ

0.2

0.4

0.6

0.8

1.0

WASS

1

(h) WASS, CBV.

0

250

500

τ

0.4

0.5

0.6

0.7

MAE

1

(i) MAE, IS.

0

250

500

τ

1.5

2.0

2.5

WASS

1

(j) WASS, IS.

0

250

500

τ

0.6

0.7

0.8

0.9

MAE

1

(k) MAE, PK.

0

250

500

τ

2

3

4

WASS

1

(l) WASS, PK.

0

250

500

τ

0.5

0.6

0.7

MAE

(m) MAE, QB.

0

250

500

τ

3

4

5

WASS

(n) WASS, QB.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

(o) MAE, WQW.

0

250

500

τ

0.6

0.8

1.0

WASS

(p) WASS, WQW.

F.5.1 Convergence Analysis of the Imputation Part

In this section, we explore the convergence of the imputation part as defined in step 7 of Algorithm 1 within our NewImp approach. Based on this, we first prove the following proposition for the convergence in the 'Imputation' part:

Proposition F.1. The convergence of the imputation part can be guaranteed, given that the discretization step size η is small enough.

0

250

500

τ

0.5

0.6

0.7

MAE

1

(a) MAE, BT.

0

250

500

τ

0.2

0.3

0.4

0.5

WASS

1

(b) WASS, BT.

0

250

500

τ

0.4

0.5

0.6

0.7

MAE

1

(c) MAE, BCD.

0

250

500

τ

1

2

3

4

WASS

1

(d) WASS, BCD.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

1

(e) MAE, CC.

0

250

500

τ

0.4

0.6

0.8

WASS

1

(f) WASS, CC.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

1

(g) MAE, CBV.

0

250

500

τ

0.6

0.8

1.0

1.2

WASS

1

(h) WASS, CBV.

0

250

500

τ

0.5

0.6

0.7

MAE

1

(i) MAE, IS.

0

250

500

τ

3.0

3.5

4.0

4.5

5.0

WASS

1

(j) WASS, IS.

0

250

500

τ

0.4

0.5

0.6

0.7

MAE

1

(k) MAE, PK.

0

250

500

τ

2

3

WASS

1

(l) WASS, PK.

0

250

500

τ

0.4

0.5

0.6

MAE

(m) MAE, QB.

0

250

500

τ

4

5

6

WASS

(n) WASS, QB.

0

250

500

τ

0.55

0.60

0.65

0.70

0.75

MAE

(o) MAE, WQW.

0

250

500

τ

0.8

1.0

1.2

WASS

(p) WASS, WQW.

Proof. First, let us reformulate the velocity field as follows:

where (i) is based on integration by parts.

Based on this reformulation, the inner product can be given as follows:

where the (i) is predicated on the requirement that the kernel function, K ( · , · ) , is semi-positive definite. Consequently, according to the abovementioned derivation, we can conclude that the evolution of F joint-NER is monotonic increasing along τ . Furthermore, F joint-NER satisfies the following inequality:

which indicates that F joint-NER is upper-bounded by 0 .

According to Eqs. (F.7) and (F.8), the cost functional F joint-NER, driven by the velocity field u ( X (joint) ) along τ , converges. Building on this, employing a smaller step size η results in the iteration curve of F joint-NER more closely approximating the ODE defined in Eq. (F.7). Consequently, a smaller η leads to a sequence where F joint-NER monotonically increases, aligning with the theoretical expectations of the ODE behavior.

Unfortunately, directly obtaining F joint-NER is intractable. Nevertheless, we can still observe the changes in WASS and MAE across iteration time τ to demonstrate the convergence of the 'Impute' part. To this end, we present the convergence trends along τ in Figs. F.4 to F.6. These figures illustrate that both MAE and WASS generally decrease as the iteration epochs increase and eventually stabilize after τ = 250 . This observed behavior supports our theoretical findings regarding the convergence of the 'Imputation' part.

0

250

500

τ

0.6

0.7

0.8

MAE

1

(a) MAE, BT.

0

250

500

τ

0.4

0.6

0.8

WASS

1

(b) WASS, BT.

0

250

500

τ

0.4

0.5

0.6

0.7

0.8

MAE

1

(c) MAE, BCD.

0

250

500

τ

2

3

4

5

6

WASS

1

(d) WASS, BCD.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

1

(e) MAE, CC.

0

250

500

τ

0.4

0.6

0.8

WASS

1

(f) WASS, CC.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

1

(g) MAE, CBV.

0

250

500

τ

0.6

0.8

1.0

1.2

1.4

WASS

1

(h) WASS, CBV.

0

250

500

τ

0.5

0.6

0.7

MAE

1

(i) MAE, IS.

0

250

500

τ

3.0

3.5

4.0

4.5

5.0

WASS

1

(j) WASS, IS.

0

250

500

τ

0.5

0.6

0.7

0.8

MAE

1

(k) MAE, PK.

0

250

500

τ

2

3

4

5

WASS

1

(l) WASS, PK.

0

250

500

τ

0.45

0.50

0.55

0.60

0.65

MAE

(m) MAE, QB.

0

250

500

τ

5

6

7

WASS

(n) WASS, QB.

0

250

500

τ

0.55

0.60

0.65

0.70

0.75

MAE

(o) MAE, WQW.

0

250

500

τ

0.8

1.0

1.2

WASS

(p) WASS, WQW.

F.5.2 Convergence Analysis of the DSM Training

Similarly, we can also give the proposition of the DSM training algorithm located in step 5 of Algorithm 1, and summarized in Algorithm 3:

Proposition F.2. The convergence of the DSM training algorithm can be guaranteed, given that the learning rate lr is small enough.

Proof. In the beginning, let us reformulate the parameter learning procedure of the DSM training algorithm as follows:

which can be further reformulated as follows:

Plugging Eq. (F.10) into Eq. (F.11), we can get the following result:

which indicates that the iterative procedure for L DSM is monotonic decreasing along τ .

Meanwhile, note that:

Finally, recall Eq. (15), we can know that the following condition holds:

Building on this, employing a smaller step size lr results in the iteration curve of L DSM more closely approximating the ODE defined in Eq. (F.11). Consequently, a smaller lr leads to a sequence where L DSM monotonically decreases, aligning with the theoretical expectations of the ODE behavior.

Based on this proposition, we plot the evolution of L DSM along time τ in Fig. F.7. These figures illustrate that the L DSM generally decreases as the iteration epochs increase. This observed behavior supports our theoretical findings regarding the convergence of the DSM training algorithm.

0

100

200

τ

160

180

200

L

DSM

1

(a)

L

DSM, MAR, BT.

0

100

200

τ

1300

1400

1500

L

DSM

1

(b)

L

DSM, MAR, BCD.

0

100

200

τ

330

340

350

L

DSM

1

(c)

L

DSM, MAR, CC.

0

100

200

τ

480

490

500

L

DSM

1

(d)

L

DSM, MAR, CBV.

0

100

200

τ

1550

1600

1650

1700

L

DSM

1

(e)

L

DSM, MAR, IS.

0

100

200

τ

1000

1050

1100

1150

L

DSM

1

(f)

L

DSM, MAR, PK.

0

100

200

τ

1800

1900

2000

L

DSM

1

(g)

L

DSM, MAR, QB.

0

100

200

τ

530

540

550

L

DSM

1

(h)

L

DSM, MAR, WQW.

0

100

200

τ

180

190

200

L

DSM

1

(i)

L

DSM, MCAR, BT.

0

100

200

τ

1400

1425

1450

1475

1500

L

DSM

1

(j)

L

DSM, MCAR, BCD.

0

100

200

τ

335

340

345

350

355

L

DSM

1

(k)

L

DSM, MCAR, CC.

0

100

200

τ

480

490

500

L

DSM

1

(l)

L

DSM, MCAR, CBV.

0

100

200

τ

1600

1650

1700

L

DSM

1

(m)

L

DSM, MCAR, IS.

0

100

200

τ

1100

1150

L

DSM

1

(n)

L

DSM, MCAR, PK.

0

100

200

τ

1850

1900

1950

2000

2050

L

DSM

1

(o)

L

DSM, MCAR, QB.

0

100

200

τ

535

540

545

550

L

DSM

1

(p)

L

DSM, MCAR, WQW.

0

100

200

τ

180

190

200

L

DSM

1

(q)

L

DSM, MNAR, BT.

0

100

200

τ

1400

1450

1500

L

DSM

1

(r)

L

DSM, MNAR, BCD.

0

100

200

τ

335

340

345

350

355

L

DSM

1

(s)

L

DSM, MNAR, CC.

0

100

200

τ

480

490

500

L

DSM

1

(t)

L

DSM, MNAR, CBV.

0

100

200

τ

1600

1650

1700

L

DSM

0

100

200

τ

1050

1100

1150

L

DSM

0

100

200

τ

1900

2000

L

DSM

0

100

200

τ

535

540

545

550

L

DSM

F.6 Downstream Task Comparison

To further substantiate the rigor of our manuscript and demonstrate the efficacy of the proposed NewImp approach, we conduct downstream task comparisons as detailed in this subsection. Initially, we evaluate the classification performance on imputed data using the following protocol: 1) Selection of datasets with non-binary labels. 2) Post-imputation, we train a support vector machine equipped with an RBF kernel and an automatic kernel coefficient. We assess the model's performance using 5-fold cross-validation, reporting both the mean and standard deviation of the accuracies across 10 runs with different random seeds. In this procedure, we select classification accuracy as our evaluation metric. 3) Additionally, we include the accuracy of ground-truth data for reference. The comparative results are presented in Table F.8. From Table F.8, it can be seen that the NewImp approach generally has the best performance among all baseline models, this phenomenon reflects the superiority of the proposed NewImp approach in a further way.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Moreover, we also conduct downstream regression task comparisons as detailed in this subsection. Initially, we evaluate the regression performance on imputed data using the following protocol: 1) Selection of datasets with continuous outcome variables. 2) After imputation, we train a support vector regression model equipped with an RBF kernel and an automatic kernel coefficient. We assess the model's performance using 5-fold cross-validation. In this procedure, we report both the mean and standard deviation of the mean squared errors (MSE) and mean absolute error (MAE) across 10 runs with different random seeds. 3) Additionally, we include the MSE and MAE on ground-truth data for reference. The comparative results are presented in Table F.9. As indicated in these results, the NewImp approach consistently outperforms most of the baseline models, further validating its superiority.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

F.7 Baseline Comparison Vary Different Missing Rates and Scenarios

In this section, we further present the extended analysis of model performance across varying missing data rates, as detailed in Tables G.1, G.3, G.5 and G.7, and the corresponding standard deviation error results are presented in Tables G.2, G.4, G.6 and G.8. From the comparison results, it can be seen that our NewImp approach generally perform well compared to most of baseline models. This phenomenon reflects that the proposed NewImp approach is robust to various missing rates, and further proves its applicability.

Appendix G Limitations & Future Directions and Broader Impact

G.1 Limitations & Future Directions

The limitations and future research directions of this work can be summarized as follows:

Utilization of Kernel Function: During the derivation of the velocity field, we employ RKHS to ensure implementation easiness. However, this regularization term may impose restrictions on the velocity field's direction, potentially limiting imputation accuracy in high-dimensional settings. Additionally, the computational complexity tends to scale quadratically with dataset size increases. Exploring alternative regularization terms [13] to replace RKHS presents a promising direction.

Training of Score Function: As discussed in Section F.4, the runtime of NewImp is predominantly governed by the DSM function. Investigating techniques to reduce the training costs of this part, such as employing sliced score matching [49], represents an intriguing area for future exploration.

Wasserstein Gradient Flow Framework: The WGF framework currently operates as a first-order system where each sample is equally weighted. A critical advancement would be the incorporation of second-order systems, such as Hamiltonian dynamics [55, 61], and other gradient flows like Fisher-Rao gradient flow [69] that assign variable weights to samples. These adaptations aim to decrease computational times inherently.

R D Support Assumption: In our manuscript, for ease of derivation, we assume that the distribution we model has support on the real number domain R D , which limits the direct application of NewImp to tabular data with categorical variables. This limitation can be alleviated by employing the mirror descent approach. Specifically, for a categorical variable with D categories, the distribution belongs to the Dirichlet distribution whose support lies on the simplex ∆ D -1 . On this basis, we can define the Bregman function as the entropy function: Ψ( X ) := ∑ D j =1 ( X j log X j -X j ) and apply mirror descent using this Bregman function to handle the categorical domain effectively. Notably, similar approaches have been successfully applied in works focusing on constrained domain sampling, as exemplified in [48]. We have implemented a comparable scheme in Section 3.1 and Appendix B.

G.2 Broader Impact Statement

MDI and DMs are pivotal areas within machine learning, each boasting a wide array of real-world applications. While numerous applications exist, this paper does not single out any specific ones; instead, it focuses on addressing fundamental challenges in these fields. This study significantly advances the application of DMs for MDI by tackling prevalent issues such as inaccurate imputation and challenging training processes. We believe that the insights garnered here can be applied to related domains, such as probabilistic time-series forecasting and image inpainting, where accuracy is often more critical than diversity in results. A common challenge across these domains is the nuanced need for precision over variety, which can lead to overlooked opportunities in model application and development. Our proposed method provides a fresh perspective on these tasks through an optimization lens. It evaluates the appropriateness of directly applying existing diffusion models to these tasks and, where necessary, proposes the derivation of novel algorithms. This approach not only enhances the understanding of the underlying mechanisms but also paves the way for more targeted and effective solutions in the future.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Scenario

Model

BT

BCD

CC

CBV

IS

PK

QB

WQW

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAR

CSDI_T

1.45E-1

5.95E-1

3.96E-2

2.78E0

6.42E-2

2.35E-1

3.95E-3

2.24E-1

5.04E-3

2.03E0

9.93E-2

3.40E0

1.84E-2

9.01E-1

2.76E-2

1.28E-1

MissDiff

1.35E-1

7.18E-1

7.42E-3

3.48E0

1.84E-2

1.67E-1

1.38E-2

1.02E0

1.29E-2

8.67E-1

3.32E-2

3.61E0

7.57E-3

5.95E0

3.28E-2

8.13E-1

GAIN

3.18E-1

5.29E-1

5.32E-2

2.81E-1

3.98E-2

5.34E-2

2.87E-2

6.18E-2

4.32E-2

1.25E-1

2.46E-3

3.84E-2

5.71E-2

3.76E-1

6.93E-2

9.70E-2

MIRACLE

1.44E-1

1.75E-1

2.22E-2

6.20E-2

1.89E-2

1.43E-2

1.49E-2

2.46E-2

3.01E-1

4.73E0

3.73E-1

3.82E0

2.56E-2

2.59E-1

4.69E-2

3.77E-2

MIWAE

1.62E-1

2.25E-1

2.11E-2

1.76E-1

1.31E-2

4.51E-2

3.24E-2

6.88E-2

2.55E-2

4.75E-2

3.67E-2

1.36E-1

2.82E-2

2.78E-1

6.03E-2

5.11E-2

Sink

1.38E-1

2.45E-1

9.11E-3

1.75E-1

3.48E-2

7.35E-2

1.23E-2

3.27E-2

3.42E-3

4.49E-2

3.73E-2

2.47E-1

1.08E-2

3.28E-1

3.56E-2

6.35E-2

TDM

1.53E-1

2.59E-1

4.48E-3

1.16E-1

3.30E-2

7.34E-2

8.41E-3

2.25E-2

2.85E-2

3.93E-2

3.31E-2

2.15E-1

2.29E-2

3.44E-1

3.24E-2

5.11E-2

ReMasker

1.27E-1

1.32E-1

9.08E-2

3.42E-1

9.42E-2

1.75E-1

1.90E-2

2.42E-2

3.78E-2

8.97E-2

4.59E-2

1.29E-1

1.54E-1

1.38E0

4.45E-2

4.02E-2

NewImp

1.33E-1

2.16E-1

1.29E-2

1.18E-1

1.87E-2

1.23E-2

3.36E-3

5.46E-3

4.84E-3

4.37E-2

2.55E-2

1.49E-1

1.98E-2

1.76E-1

6.04E-2

6.49E-2

MCAR

CSDI_T

4.54E-2

5.09E-2

1.05E-2

1.33E0

1.73E-2

1.46E-1

5.36E-3

3.03E-2

8.51E-3

7.25E-1

1.69E-2

1.37E0

4.54E-3

5.07E-1

2.54E-3

2.50E-2

MissDiff

3.37E-2

5.02E-1

9.14E-3

5.36E-1

1.75E-2

2.94E-1

7.44E-3

1.06E0

9.22E-3

9.56E-2

1.49E-2

2.37E0

3.06E-3

2.59E0

7.27E-3

1.13E0

GAIN

5.67E-2

5.60E-2

7.51E-3

3.57E-2

6.21E-3

2.45E-2

2.02E-2

3.50E-2

7.40E-3

3.03E-2

9.46E-3

4.17E-2

1.42E-2

1.77E-1

1.67E-2

1.82E-2

MIRACLE

2.35E-2

2.34E-2

1.35E-2

4.06E-2

5.66E-3

6.01E-3

5.36E-3

1.14E-2

2.36E-2

1.98E-1

1.21E-1

1.83E0

4.58E-3

1.09E-1

7.64E-3

2.30E-2

MIWAE

2.76E-2

1.85E-2

1.10E-2

5.57E-2

2.79E-2

4.10E-2

1.68E-2

5.13E-2

2.73E-2

1.16E-1

1.64E-2

7.49E-2

1.61E-2

2.68E-1

5.22E-3

1.72E-2

Sink

2.04E-2

2.24E-2

5.92E-3

7.02E-2

1.33E-2

2.09E-2

1.16E-2

3.50E-2

2.84E-3

5.01E-2

8.93E-3

8.83E-2

7.60E-3

2.46E-1

6.25E-3

2.48E-2

TDM

1.48E-2

2.87E-2

6.57E-3

7.20E-2

1.42E-2

1.93E-2

1.38E-2

3.52E-2

2.61E-3

7.96E-3

1.84E-2

9.92E-2

1.01E-2

1.93E-1

4.91E-3

1.32E-2

ReMasker

3.25E-2

9.47E-3

1.41E-2

1.29E-1

1.07E-1

1.96E-1

1.09E-2

1.09E-2

4.00E-2

1.63E-1

4.79E-2

3.06E-1

1.13E-1

1.56E0

8.09E-3

1.52E-2

NewImp

3.25E-2

2.44E-2

1.42E-3

2.83E-2

3.12E-3

1.66E-2

8.22E-3

1.95E-3

1.13E-2

4.72E-2

6.48E-3

2.83E-2

1.97E-3

1.52E-1

4.82E-3

1.22E-2

MNAR

CSDI_T

8.79E-2

3.69E-1

2.23E-2

1.85E0

1.92E-2

1.35E-1

1.44E-2

1.07E-1

1.17E-2

1.28E0

1.33E-2

1.89E0

5.62E-3

6.07E-1

7.60E-3

3.31E-2

MissDiff

5.88E-2

7.69E-1

8.45E-3

1.82E0

1.03E-2

2.02E-1

1.97E-2

8.06E-1

2.19E-2

2.28E-1

1.39E-2

4.05E0

8.29E-3

3.70E0

8.84E-3

2.05E-1

GAIN

3.27E-2

8.42E-2

2.23E-2

1.90E-1

4.37E-2

6.57E-2

2.17E-2

3.99E-2

7.77E-3

8.22E-2

2.45E-2

2.00E-1

2.07E-2

2.61E-1

1.65E-2

5.25E-2

MIRACLE

2.42E-2

2.71E-2

5.33E-2

8.42E-1

1.65E-2

1.82E-2

2.34E-2

3.99E-2

7.48E-2

2.21E0

6.79E-2

1.25E0

2.25E-2

3.69E-1

5.58E-3

1.39E-2

MIWAE

1.04E-2

3.27E-2

9.20E-3

4.25E-2

3.77E-2

6.09E-2

1.96E-2

4.44E-2

2.51E-2

1.70E-1

5.47E-2

3.37E-1

3.40E-2

4.35E-1

6.92E-3

2.52E-2

Sink

1.93E-2

6.83E-2

2.07E-2

1.43E-1

8.90E-3

2.11E-2

2.06E-2

3.74E-2

1.97E-2

1.37E-1

3.13E-2

1.86E-1

3.57E-3

3.95E-1

7.94E-3

2.33E-2

TDM

2.20E-2

6.10E-2

2.29E-2

1.27E-1

1.02E-2

2.07E-2

2.02E-2

3.68E-2

2.99E-2

1.34E-1

2.51E-2

1.59E-1

1.03E-2

2.80E-1

5.78E-3

1.86E-2

ReMasker

2.99E-2

1.14E-1

3.00E-2

1.38E-1

8.40E-2

1.71E-1

2.04E-2

2.98E-2

3.94E-2

2.77E-1

3.44E-2

2.01E-1

6.90E-2

9.91E-1

1.76E-2

2.88E-2

NewImp

2.04E-2

1.70E-2

4.70E-3

1.96E-2

1.38E-2

1.78E-2

1.45E-2

2.77E-2

3.86E-3

8.63E-3

7.47E-3

8.12E-2

2.18E-2

6.56E-1

4.31E-3

1.59E-2

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Kindly Note : The best results are bolded and the second best results are underliend. '*' marks the results that NewImp significantly outperform with p -value < 0 . 05 over paired samples t -test.

Scenario

Model

BT

BCD

CC

CBV

IS

PK

QB

WQW

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAE

WASS

MAR

CSDI_T

8.31E-2

4.42E-1

2.83E-2

2.31E0

3.66E-2

1.29E-1

1.34E-2

5.51E-2

2.39E-2

9.07E-1

3.55E-2

2.14E0

2.05E-2

5.84E-1

3.11E-2

2.52E-2

MissDiff

4.50E-2

5.54E-1

5.12E-3

1.98E0

1.24E-2

1.09E-1

8.44E-3

8.50E-1

3.78E-2

4.65E-1

2.17E-2

2.02E0

2.51E-2

3.72E0

2.49E-2

1.10E0

GAIN

5.10E-2

1.42E-1

6.41E-2

5.67E-1

4.97E-2

9.93E-2

3.91E-2

9.48E-2

2.46E-2

2.48E-1

7.03E-2

8.68E-1

8.40E-2

1.77E0

1.01E-1

3.67E-1

MIRACLE

1.50E-1

2.05E-1

9.31E-2

1.94E0

3.74E-2

6.26E-2

1.45E-2

2.64E-2

5.95E-1

1.68E1

5.56E-1

1.26E1

4.19E-2

6.07E-1

2.96E-2

5.79E-2

MIWAE

6.94E-2

1.35E-1

1.75E-2

1.16E-1

2.18E-2

4.43E-2

1.89E-2

1.02E-1

3.21E-2

2.02E-1

3.10E-2

2.21E-1

4.53E-2

3.92E-1

2.57E-2

5.02E-2

Sink

4.60E-2

1.38E-1

8.04E-3

8.61E-2

2.40E-2

3.24E-2

1.17E-2

6.38E-2

3.09E-2

1.35E-1

1.16E-2

4.24E-2

2.72E-2

7.38E-2

3.42E-2

8.69E-2

TDM

5.43E-2

1.61E-1

1.16E-2

7.27E-2

2.37E-2

4.04E-2

1.32E-2

7.26E-2

1.92E-2

9.41E-2

1.34E-2

9.02E-2

2.84E-2

6.32E-2

3.02E-2

7.58E-2

ReMasker

7.39E-2

1.44E-1

2.99E-2

1.41E-1

9.30E-2

2.22E-1

2.35E-2

3.19E-2

2.98E-2

2.94E-1

7.54E-2

6.08E-1

6.79E-2

7.58E-1

2.26E-2

4.20E-2

NewImp

1.07E-1

2.34E-1

3.02E-2

2.10E-1

2.83E-2

5.71E-2

1.19E-2

2.15E-2

1.51E-2

1.69E-1

1.42E-2

8.97E-2

1.69E-2

1.36E-1

4.13E-2

9.67E-2

MCAR

CSDI_T

1.28E-2

1.58E-1

5.14E-3

1.23E0

8.83E-3

9.46E-2

2.93E-3

3.67E-2

7.87E-3

9.01E-2

6.38E-3

6.10E-1

5.04E-3

1.69E-1

1.11E-2

7.61E-2

MissDiff

2.62E-3

3.37E-1

1.01E-2

3.62E-1

4.33E-3

1.06E-1

3.15E-3

3.18E-1

9.26E-3

9.37E-2

4.25E-3

5.98E-1

4.81E-3

1.26E0

1.71E-3

3.02E-1

GAIN

1.71E-2

3.47E-2

2.42E-2

2.75E-1

4.53E-2

1.44E-1

2.62E-1

1.40E0

1.09E-1

2.43E0

2.01E-2

3.35E-1

1.38E-1

7.19E0

1.41E-1

6.71E-1

MIRACLE

2.19E-2

2.36E-2

7.61E-2

1.93E0

1.32E-2

2.67E-2

1.55E-2

3.61E-2

3.58E-2

2.18E0

4.86E-2

2.18E0

2.15E-2

1.96E-1

2.25E-3

5.85E-3

MIWAE

6.77E-3

4.45E-2

2.86E-2

3.40E-1

1.15E-2

3.66E-2

9.27E-3

4.98E-2

1.19E-2

1.97E-1

1.72E-2

3.34E-1

1.09E-2

3.17E-1

1.98E-3

1.68E-2

Sink

9.45E-3

2.19E-2

5.14E-3

7.93E-2

4.00E-3

1.70E-2

2.62E-3

2.17E-2

1.92E-3

5.94E-2

3.52E-3

1.46E-1

4.69E-3

7.05E-2

1.96E-3

2.86E-2

TDM

7.85E-3

2.24E-2

7.53E-3

9.35E-2

7.82E-3

2.75E-2

4.80E-3

2.12E-2

3.76E-3

1.02E-1

9.26E-3

1.67E-1

3.48E-3

1.87E-1

2.66E-3

5.92E-3

ReMasker

5.17E-2

6.46E-2

6.65E-4

4.77E-2

3.95E-2

1.45E-1

2.96E-3

1.48E-2

1.13E-2

1.97E-1

1.18E-2

2.05E-1

1.34E-2

4.30E-1

1.06E-3

2.08E-2

NewImp

1.22E-2

2.48E-2

3.72E-3

1.03E-1

9.87E-3

3.21E-2

5.17E-3

3.13E-2

8.34E-3

1.56E-1

7.22E-3

6.63E-2

6.53E-3

2.43E-1

1.57E-3

9.71E-3

MNAR

CSDI_T

5.67E-3

1.31E-1

1.04E-2

7.15E-1

2.05E-2

9.53E-2

1.23E-2

1.30E-1

6.97E-3

1.06E-1

1.36E-2

9.70E-1

4.20E-3

2.68E-1

6.23E-3

5.50E-2

MissDiff

7.33E-3

3.18E-1

6.20E-3

7.94E-1

1.48E-2

1.30E-1

6.44E-3

3.30E-1

9.96E-3

1.17E-1

1.01E-2

8.14E-1

8.28E-3

1.97E0

4.54E-3

1.73E-1

GAIN

3.37E-2

8.08E-2

3.52E-2

8.61E-1

7.13E-2

2.53E-1

8.83E-2

5.11E-1

1.55E-1

4.35E0

4.63E-2

1.21E0

1.62E-1

1.30E1

7.71E-2

5.63E-1

MIRACLE

4.62E-2

1.54E-1

1.55E-1

2.46E1

1.56E-2

4.66E-2

1.35E-2

4.39E-2

2.21E-2

1.78E0

9.72E-2

3.67E0

3.60E-2

1.71E0

1.24E-2

3.05E-2

MIWAE

2.88E-2

7.04E-2

3.13E-2

2.81E-1

6.52E-3

5.64E-2

1.06E-2

5.53E-2

2.19E-2

3.10E-1

1.69E-2

2.92E-1

1.75E-2

3.51E-1

4.18E-3

3.20E-2

Sink

4.90E-3

1.42E-2

6.82E-3

2.35E-1

7.10E-3

3.41E-2

4.21E-3

3.28E-2

6.06E-4

9.24E-2

1.26E-2

3.20E-1

7.64E-3

1.55E-1

7.06E-3

4.95E-2

TDM

2.87E-3

1.46E-2

1.06E-2

2.95E-1

7.99E-3

3.36E-2

5.73E-3

3.48E-2

1.56E-2

1.77E-1

1.51E-2

3.73E-1

6.81E-3

1.42E-1

6.82E-3

5.11E-2

ReMasker

2.52E-2

4.11E-2

2.60E-2

4.54E-1

5.75E-2

2.40E-1

1.06E-2

4.72E-2

3.77E-2

4.46E-1

3.96E-2

5.44E-1

1.13E-2

1.19E-1

4.54E-3

1.92E-2

NewImp

4.81E-3

1.42E-2

8.08E-3

1.42E-1

1.86E-2

8.53E-2

1.61E-2

7.67E-2

2.97E-3

8.78E-2

1.42E-2

1.95E-1

6.32E-3

2.48E-1

1.11E-2

4.15E-2

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main claims made in the abstract and introduction accurately reflect our paper's contributions and scope. We restrict our application in missing value imputation task in numerical tabular, and our analysis is mainly focused on diffusion models, where the score function is required.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Our limitations are listed in Appendix G.1.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: To uphold the rigor of our manuscript, we provide all proofs of our proposition as outlined in Appendix C. Besides, all theorems are properly cited in the manuscript.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We attempt to list all hyperparameters in Appendices D and E to ensure reproducibility.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We used the open access UCI datasets, and we uploaded our algorithm in this github link https://github.com/JustusvLiebig/NewImp.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have included all detailed information in Appendices D and E.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: In Tables F.3 to F.9 and G.1 to G.8, and Figs. F.4 to F.7 , we report standard deviation errors suitably and correctly defined or other appropriate information about the statistical significance and error bar of the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The required information is given in Appendices E.2 and F.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Since it is an algorithm-oriented research, there is no societal impact of the work performed.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not involve the safeguards issue.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The creators or original owners of assets (e.g., code, data, models), used in the paper, are properly credited. The license and terms of use are explicitly mentioned and properly respected.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not not release new assets in this manuscript.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our experiments did not involve 'Crowdsourcing and Research with Human Subjects'.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our answer is NA since our paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses two primary issues in Missing Data Imputation (MDI) using Diffusion Models (DMs): inaccurate imputation due to sample diversification and difficult training caused by the complexity of designing the mask matrix. The authors propose a novel approach, Kernelized Negative Entropy-regularized Wasserstein Gradient Flow Imputation (KnewImp), which aims to resolve these issues for numerical tabular datasets.\n\nSTRENGTHS:\n- The paper addresses two issues in the domain of DM-based MDI—sample diversification leading to inaccurate imputation and the complex training process due to mask matrix design. By identifying and explicitly addressing these issues, the paper provides a fresh perspective on improving MDI techniques. The introduction of the negative entropy-regularized cost functional is a creative and innovative approach to discourage diversification, aligning the generative model’s objectives more closely with the needs of MDI tasks. The integration of the WGF framework with RKHS to derive an imputation procedure is an original and elegant solution.\n- The paper provides thorough theoretical analyses and proofs, ensuring the soundness of the proposed approach. The extensive experiments conducted on multiple real-world datasets from the UCI repository underscore the robustness and effectiveness of the KnewImp approach. The detailed ablation studies and sensitivity analyses enhance the quality of the research by thoroughly examining the contributions of different components and the impact of key hyperparameters.\n- The paper is well-organized, with a clear delineation of problems, proposed solutions, theoretical foundations, experimental setup, and results. Each section logically flows into the next, making it easy to follow the authors’ arguments.\n\nWEAKNESSES:\n- The choice of kernel in the Reproducing Kernel Hilbert Space (RKHS) can significantly impact the performance of the method. The paper primarily uses the radial basis function (RBF) kernel but does not explore the effects of using different types of kernels or the rationale behind choosing the RBF kernel. A detailed analysis of how different kernel choices affect the imputation quality and computational efficiency would strengthen the technical rigor of the study.\n- The theoretical foundations of KnewImp rely on certain assumptions about the data distribution, such as the smoothness of the underlying density functions. The robustness of the method to violations of these assumptions is not thoroughly investigated. Including experiments that test the method’s performance on datasets with varying statistical properties (e.g., heavy-tailed distributions, multimodal distributions) would provide a more comprehensive assessment of its robustness.\n\nQUESTIONS:\n1. How robust is KnewImp to variations in data distributions, such as heavy-tailed, multimodal, or skewed distributions?\n2. Why was the radial basis function (RBF) kernel chosen for the RKHS, and how does the choice of kernel affect the performance of KnewImp?\n\nLIMITATIONS:\nYes. Discussed in Appendix F.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents KnewImp, a kernelized negative entropy-regularized Wasserstein gradient flow imputation approach to numerical tabular data imputation. The authors argue that existing missing data imputation frameworks based on diffusion models suffer from two major limitations. Firstly, diffusion models primarily focus on sample diversification rather than accuracy, which results in discrepancy between the training objective of the diffusion models and the aim of tabular data imputation. Secondly, existing approaches are trained by masking parts of the observed data and then predicting the masked entries. This results in training difficulty due to the need of designing complex mask matrix selection mechanisms. To address the limitations of existing models, the authors propose a Wasserstein gradient flow based framework, which employs a novel cost functional with diversification-discouraging negative entropy as regularization. KnewImp is derived within the Wasserstein gradient flow framework, reproducing the kernel Hilbert space. To bypass the need of the mask matrix and to make the model easier to train, the authors further develop a novel cost functional based on joint distribution. Experimental results on a variety of real-world datasets show that KnewImp achieves state-of-the-art, outperforming a number of state-of-the-art baseline alternatives.\n\nSTRENGTHS:\n- KnewImp is theoretically sound and achieves state-of-the-art results on a number of real-world datasets.\n\n- The paper is in general well-written and easy-to-follow.\n\nWEAKNESSES:\n- My main concern is that, the authors only presented results in terms of MAE and Wass, but did not show results on downstream tasks. It is therefore unclear whether KnewImp would be effective in real-world scenarios. I would consider raising my score if the authors can provide the result of KnewImp on downstream tasks. The authors can refer to Figure 5 of the TDM paper in terms of downstream results.\n\n- From the ablation study, I think the majority of the performance improvement comes from modeling the joint distribution directly. The diversification-discouraging negative entropy regularization only marginally improves the performance. The main novelty of KnewImp is from the NER objective. However, ablation study suggests that it does not add much to the final performance.\n\nQUESTIONS:\n- In line 150, the authors claim that \"This diversification is fundamentally at odds with the precision required in MDI tasks.\" I understand that the need of discouraging diversification in favor of precision is the main motivation of KnewImp. However, I wonder if the authors can provide a more intuition explanation about this claim, or if the discrepancy between accuracy and diversification can be somehow quantified.\n\n- Why did the authors primarily focus on MAR and MCAR rather than MNAR? Is this because of the setting of KnewImp is not particularly suitable for MNAR?\n\n- In Table 2, not all results in bold are best results.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations and potential negative societal impact of their work, and have provided sufficient justification.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considered tackling the Missing Data Imputation (MDI) problem via diffusion models, which treats MDI as an generative problem. As DM-based methods focus on sample diversification rather than accuracy, which is the primary evaluation metric for MDI, the authors proposed one cost functional to discourage diversification in sample generation based on the Wasserstein Gradient Flow framework. Moreover, given that the true values of the missing data are unknown, the authors proposed to replace the joint distribution with the conditional distribution throughout the learning procedure.\n\nSTRENGTHS:\nThis paper focuses on two important questions faced by the diffusion model based solver of the missing data imputation problem. Extensive numerical experiments are provided to validate the effectiveness of the proposed methodology.\n\nWEAKNESSES:\nThe main issue here is that incorporating Wasserstein Gradient Flow (WGF) with generative modeling doesn't seem to be a new idea. However, it seems that the authors didn't include any related references in section 2.3, which provides a brief review of WGF. For instance, it might be necessary to cite and discuss the following articles [1-6].\n\nQUESTIONS:\nIt seems to the reviewer that the authors have made a strong assumption throughout this paper. Specifically, the assumption in Proposition 3.4 says that the joint distribution $r(X^{\\text{joint}})$ can be factorized as $r(X^{(\\text{miss})})r(X^{\\text{obs}})$, where $r$ denotes the probability density function. Would it be possible for the authors to discuss whether such assumption is realistic or not? Intuitively, it seems that the missing and observed entries can't be utterly uncorrelated.\n\nLIMITATIONS:\nIt seems to the reviewer that there are too many typos in the current version of this paper, especially for the theoretical derivation part in the appendix. For instance, derivation of the equations from line 583 to line 584 seem to have lots of issues - the first step should contain some dot product? \n\nReferences:\n\n[1] Ansari, A. F., Ang, M. L., & Soh, H. (2020). Refining deep generative models via discriminator gradient flow. arXiv preprint arXiv:2012.00780.\n\n[2] Cheng, X., Lu, J., Tan, Y., & Xie, Y. (2024). Convergence of flow-based generative models via proximal gradient descent in Wasserstein space. IEEE Transactions on Information Theory.\n\n[3] Choi, J., Choi, J., & Kang, M. (2024). Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport. arXiv preprint arXiv:2402.05443.\n\n[4] Gao, Y., Jiao, Y., Wang, Y., Wang, Y., Yang, C., & Zhang, S. (2019, May). Deep generative learning via variational gradient flow. In International Conference on Machine Learning (pp. 2093-2101). PMLR.\n\n[5] Heng, A., Ansari, A. F., & Soh, H. Deep generative Wasserstein gradient flows.\n\n[6] Xu, C., Cheng, X., & Xie, Y. (2024). Normalizing flow neural networks by JKO scheme. Advances in Neural Information Processing Systems, 36.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new algorithm for data imputation. The idea is to estimate the score function corresponding to the posterior p(x_miss/x_obs) using DSM and then infer the missing values using a WGF equivalence argued in this paper itself. These alternating steps are repeated until convergence. Simulations on benchmark datasets illustrate the efficacy of the proposal.\n\nSTRENGTHS:\n1. In empirical comparisons, the proposal seems to beat existing baselines.\n\nWEAKNESSES:\n1. There are some technical concerns I have raised in the next section\n\n2. I found description in sec34 very cryptic. It would have been nice if important portions of related appendix were moved to the main section. Also, will the proposed alternating style algorithm converge? Do we know anything property of this converged solution? Some discussion around the final algorithm would have helped in understanding the methodology better.\n\n3.\n\nQUESTIONS:\n1. I am not sure how the equivalence in (2) follows. It seems the objective in RHS is independent of X_miss as it is integrated out. objective in RHS is a function of r, p, X_obs. Whereas objective in LHS is a function of X_miss, does not involve r. It will be helpful if this equivalence is clarified. Wrt. to the optimization variable the objective in RHS seems to be a constant. Is there any notation that I am missing? Since it is a basic step and subsequent derivations depend on this critically, I was unable to check the correctness of soem of the steps. \n\n2. Morevoer, (4) seems to clearly show that the regularizer is not a function of X_miss. Then I am not sure how this regularizer, which is essentially a constant, matter.\n\n3. 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, when what is the relation with p in the objective of (2) ?\n\n4. is the assumptions of prop 3.4 meaningful? If X_Obs are X_miss are independent, then what information will the observations provide for imputation? will the imputation problem remain meaningful ? Please clarify this.\n\n5.\n\nLIMITATIONS:\nna", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Kernelized Negative Entropy-regularized Wasserstein Gradient Flow Imputation (KnewImp), a novel approach for imputing missing data in numerical tabular datasets. The proposed method addresses two significant challenges in diffusion model-based missing data imputation (MDI): inaccurate imputation and difficult training. KnewImp integrates the Wasserstein gradient flow (WGF) framework with a negative entropy-regularized (NER) cost functional to enhance imputation accuracy and simplify the training process by eliminating the need for complex mask matrix designs. The method's efficacy is demonstrated through extensive experiments, showing superior performance compared to state-of-the-art imputation techniques.\n\nSTRENGTHS:\n1. The paper presents a unique integration of diffusion models with the Wasserstein gradient flow framework, incorporating a novel negative entropy regularization to address specific challenges in missing data imputation.\n\n2. The work is grounded in solid theoretical foundations, providing clear proofs and propositions that establish the effectiveness and validity of the proposed approach.\n\n3.  Extensive experiments on real-world datasets validate the method's superiority, with significant improvements in both mean absolute error (MAE) and Wasserstein distance (Wass) metrics.\n\nWEAKNESSES:\n1. The theoretical concepts and mathematical formulations presented are quite dense and may be challenging for readers not well-versed in advanced optimization and diffusion models. Simplifying these explanations or providing more intuitive descriptions could improve accessibility.\n\n2. While the method is compared to several models, including a wider range of baseline methods, particularly more recent advancements in diffusion-based MDI (e.g., [1,2]), would provide a more comprehensive evaluation.\n\n3. The paper could benefit from a more detailed discussion on the convergence properties of the proposed algorithm, including potential limitations and scenarios where the method might struggle.\n\n[1] Du, Tianyu, Luca Melis, and Ting Wang. \"ReMasker: Imputing Tabular Data with Masked Autoencoding.\" International Conference on Learning Representations, 2024.\n\n[2] Zheng, Shuhan, and Nontawat Charoenphakdee. \"Diffusion models for missing value imputation in tabular data.\" NeurIPS 2022 First Table Representation Workshop.\n\nQUESTIONS:\n1. Can the authors elaborate on how the proposed method could be adapted or extended to handle other types of data, such as categorical or mixed-type datasets?\n\n2. What are the computational requirements for implementing KnewImp in practice, and how does it scale with larger datasets?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-30"}