MissDiff: Training Diffusion Models on Tabular Data with Missing Values

## Abstract

Abstract The diffusion model has shown remarkable performance in modeling data distributions and synthesizing data. However, the vanilla diffusion model requires complete or fully observed data for training. Incomplete data is a common issue in various real-world applications, including healthcare and finance, particularly when dealing with tabular datasets. This work presents a unified and principled diffusion-based framework for learning from data with missing values under various missing mechanisms. We first observe that the widely adopted “impute-then-generate” pipeline may lead to a biased learning objective. Then we propose to mask the regression loss of Denoising Score Matching in the training phase. We prove the proposed method is consistent in learning the score of data distributions, and the proposed training objective serves as an upper bound for the negative likelihood in certain cases. The proposed framework is evaluated on multiple tabular datasets using realistic and efficacious metrics and is demonstrated to outperform state-of-the-art diffusion model on tabular data with “impute-then-generate” pipeline by a large margin.

## 1 Introduction

Diffusion models has emerged as an effective tool for modeling the data distribution and synthesize various type of data, such as images [; ; ; ], videos [], point clouds [], and tabular data [; ]. It is known that such machine learning models typically rely on high-quality training data, which are usually expected to be free of missing values. In reality, it is often challenging to obtain complete data, particularly in healthcare, finance, recommendation systems, and social networks, due to privacy concerns, high cost or sampling difficulties, and the skewed distribution of user-generated content.
For example, the respiratory rate of a patient may not have been measured, either because it was deemed unnecessary or was accidentally not recorded [; ; ]. Additionally, some information may be difficult or even dangerous to acquire, such as information obtained through a biopsy, which may not have been gathered for some reasons [].

In this work, we focus on learning a generative model from training data containing a significant amount of missing values, a problem that has been largely overlooked in the literature despite its widespread practical applications. Deep generative models, particularly diffusion modeimagels, can be used to augment training data and enhance the performance of image classification tasks [; ] and adversarial robustness [; ; ]. Following this idea, we can achieve better performance for downstream tasks by utilizing generative model learning on incomplete data for synthetic data generation.
We will primarily utilize tabular data as examples, especially for the numerical experiments, as tabular data is a commonly encountered data type and frequently contains missing values in various applications [; ]. Moreover, by considering tabular data as an example, we will simultaneously study the missing value scenarios in categorical and continuous variables, which are both contained in tabular type data.

To deal with missing values in the training data, numerous studies propose to use various imputation methods and then train the model on the imputed data. Taking tabular data as an example, some approaches involve deleting instances (rows) or features (columns) with missing data or replacing missing values with the mean of observed values for that feature (column). Other methods employ machine learning approaches [; ] or deep generative models for imputation tasks [; ; ; ; ]. It has been shown that imputation may reduce the diversity of the training data and may lead to biased performances in downstream tasks [; ].

In addition to imputation or simple deletion methods, previous work also studied learning from data with missing values and synthesizing complete data using GAN or VAE architectures [; ; ]. Compared with our proposed framework, these methods involve training additional networks, impose certain assumptions on the missing mechanisms, and the unique challenges associated with tabular data are less investigated.

In this work, we propose a diffusion-based framework, which we call MissDiff, for learning from data with missing values. We present the theoretical justifications of MissDiff on recovering the oracle score function and upper bounding the negative likelihood on the data under mild assumptions on the missing mechanisms. To the best of our knowledge, this is the first work that learns a generative model from mixed-type data containing missing values, and the missing values are used directly in the training process without prior imputation.
We conduct a suite of numerical experiments on mixed-type tabular data, comprising both continuous and categorical variables, under various missing mechanisms. Evaluated under several realistic and efficacious metrics, MissDiff consistently outperforms other baseline methods by a considerable margin.

Our contributions can be summarized as follows.

- •
We propose a diffusion-based framework, which we call MissDiff, for generative model training from incomplete data. We mainly focus on tabular data generation, which contains both categorical and continuous variables.
- •
We provide the theoretical justifications of MissDiff on recovering the oracle score function and upper bounding the negative likelihood on the data under mild assumptions on the missing mechanisms.
- •
We conduct extensive numerical experiments on multiple real tabular datasets under different missing mechanisms to demonstrate the effectiveness of MissDiff.

The rest of the paper is organized as follows. Section [2](#S2) reviews the setup of the missing data mechanism and the diffusion model. Section [3](#S3) introduces the proposed method. Section [4](#S4) theoretically characterizes the effectiveness of the proposed method. Numerical results are given in Section [5](#S5). We conclude the paper in Section [6](#S6). All proofs and additional numerical examples and details are deferred to the appendix.

### 1.1 Related Work

One line of research focuses on different types of generative models trained directly on data with missing values. These studies carefully modify the architecture and training objectives of Generative Adversarial Network (GAN) or Variational Autoencoder (VAE) to learn from incomplete data [; ; ].

Another research direction explores learning generative models for imputing missing values in observed data [; ; ; ; ; ; ; ; ]. For example, [] proposes the conditional score-based generative model for time series imputation. Moreover, [] adapt the conditional score-based diffusion model proposed in [] for imputing tabular data. Imputation methods cannot easily used for generating new complete data, which is the main difference with the first line of works.

Tabular data, as a mixed-type data that typically contains both categorical and continuous variables, has attracted significant attention in the field of machine learning. Tabular data synthesis has been a long-standing research topic in this area. The presence of mixed variable types and class imbalance for discrete variables make it a challenging task to model tabular data. Recently, several deep learning-based models have been proposed for generating tabular data [; ; ; ; ; ; ; ]. Among these methods, [] employs Gaussian transitions for continuous variables and multinomial transitions for discrete random variables, while [] proposes a self-paced learning technique and a fine-tuning strategy for score-based models and achieves state-of-the-art performance in tabular data generation. Moreover, the discrete Score Matching methods proposed in [] and [] can also be employed to handle discrete variables in tabular data.

## 2 Problem Setup and Preliminaries

### 2.1 Training with Missing Data

We aim to learn a diffusion-based generative model from training data that may contain a certain proportion of missing values.
Following the settings in [; ; ], we denote the underlying complete $d$-dimensional data as $\mathbf{x}=(x_{1},\ldots,x_{d})\in\mathcal{X}$ and assume it is sampled from the unknown true data-generating distribution $p_{0}(\mathbf{x})$. Here, each variable $x_{i}$, $i=1,\ldots,d$, can be either categorical or continuous. For each data point $\mathbf{x}$, suppose there is a binary mask $\mathbf{m}=(m_{1},\ldots,m_{d})\in\{0,1\}^{d}$ which indicates the missing entry for the current sample, i.e.,

$$ $m_{i}=\begin{cases}1&\text{ if }x_{i}\text{ is observed, }\\ 0&\text{ if }x_{i}\text{ is missing. }\end{cases}$ $$

Then, the observed data $\mathbf{x}^{\text{obs}}=\mathbf{x}\odot\mathbf{m}+\mathrm{na}\odot(\mathbf{1}-
\mathbf{m})$, where $\mathrm{na}$ indicates the missing value, $\odot$ denotes element-wise multiplication, and $\mathbf{1}$ is the all-one vector.

Suppose we have $n$ complete (unobservable) training data points $\mathbf{x}_{1},\ldots,\mathbf{x}_{n}\overset{iid}{\sim}p_{0}(\mathbf{x})$ and simultaneously $n$ corresponding masks $\mathbf{m}_{1},\ldots,\mathbf{m}_{n}$ generated from a specific missing data mechanism detailed later. Then the observed data values are $S^{\text{obs}}=\{\mathbf{x}_{i}^{\text{obs}}\}_{i=1}^{n}$ with $\mathbf{x}_{i}^{\text{obs}}=\mathbf{x}_{i}\odot\mathbf{m}_{i}+\mathrm{na}\odot
(\mathbf{1}-\mathbf{m}_{i})$.
The missing mechanisms can be categorized based on the relationships between the mask $\mathbf{m}$ and the complete data $\mathbf{x}$ [] as follows,

- •
Missing Completely At Random (MCAR): mask $\mathbf{m}$ is independent with the completed data $\mathbf{x}$.
- •
Missing At Random (MAR): mask $\mathbf{m}$ only depends on the observed value $\mathbf{x}^{\text{obs}}$.
- •
Not Missing At Random (NMAR): $\mathbf{m}$ depends on the observed value $\mathbf{x}^{\text{obs}}$ and missing value.

Compared with previous work which typically develop their algorithms and theoretical foundations under the MCAR assumption [; ; ; ], our method and theoretical guarantees aim to provide a general framework for learning on incomplete data and generate complete data.

Our objective is to train a generative model $p_{\mathbf{\phi}}$, parametrized by the neural network parameters $\phi$, using the observed data $S^{\text{obs}}$, such that $p_{\mathbf{\phi}}$ is close to the true distribution $p_{0}(\mathbf{x})$ and we can efficiently generate synthetic data from $p_{\mathbf{\phi}}$. The optimal model is the one that maximizes the likelihood of the data sampling from the true data-generating distribution $p(\mathbf{x})$, which is solved via maximizing the likelihood function on training samples $\prod_{i}p_{\phi}(\mathbf{x}_{i})$ and thus it logarithm $\sum_{i}\log p_{\phi}(\mathbf{x}_{i})$, which corresponds to the following objective in the population sense

$$ $\mathbf{\phi}=\underset{\mathbf{\phi}}{\arg\max}\mathbb{E}_{\mathbf{x}\sim p( \mathbf{x})}[\log p_{\mathbf{\phi}}(\mathbf{x})],$ $$

or equivalently, the model that minimizes the Kullback-Leibler (KL) divergence between the true data-generating distribution $p(\mathbf{x})$ and the model $p_{\mathbf{\phi}}(\mathbf{x})$. In the following, we mainly consider the score-based generative model as $p_{\mathbf{\phi}}$.

### 2.2 Score-Based Generative Model

In this work, we adopt the diffusion model as the prototype for developing our proposed method. We propose to train the model with missing values directly without the need for prior imputation. We first briefly review the key components of score-based generative models [; ].

Score-based generative models are a class of generative models that learn the score function, which is the gradient of the log-density of the data distribution. These models have gained attention due to their flexibility and effectiveness in capturing complex data distributions. Following the notation in [], the score-based generative models are based on a forward stochastic differential equation (SDE), $\mathbf{x}(t)$ with $t\in[0,T]$, defined as

$$ $\mathrm{d}\mathbf{x}(t)=\mathbf{f}(\mathbf{x}(t),t)\mathrm{d}t+g(t)\mathrm{d} \mathbf{w},$ (1) $$

where $\mathbf{w}$ is the standard Wiener process (Brownian motion), $\mathbf{f}(\cdot,t):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is a vector-valued function called the drift coefficient of $\mathbf{x}(t)$, and $g(\cdot):\mathbb{R}\rightarrow\mathbb{R}$ is a scalar function known as the diffusion coefficient of $\mathbf{x}(t)$.

The solution of a stochastic differential equation is a continuous trajectory of random variables $\{\mathbf{x}(t)\}_{t\in[0,T]}$. Let $p(\mathbf{x})$ denote the path measure for the trajectory $\mathbf{x}$ on $[0,T]$, $p_{t}(\mathbf{x})$ denote the marginal probability density function of $\mathbf{x}(t)$, and $p(\mathbf{x}(t)|\mathbf{x}(s))$ denote the conditional probability density of $\mathbf{x}(t)$ conditioned on $\mathbf{x}(s)$, where $s<t$ is a previous time point. When constructing the SDE, we let $p_{0}(\mathbf{x})$ be the true data distribution, and after perturbing the data according to the SDE, the data distribution becomes $p_{T}(\mathbf{x})$ which is close to a tractable noise distribution, usually set as the standard Gaussian distribution.

The data generation process is performed via the reverse SDE, i.e., first sampling data $\mathbf{x}_{T}$ from $p_{T}(\mathbf{x})$ and then generate $\mathbf{x}_{0}$ through the reverse of ([1](#S2.E1)). For any SDE in ([1](#S2.E1)), the corresponding backward/reverse process is

$$ $\mathrm{d}\mathbf{x}(t)=\left[\mathbf{f}(\mathbf{x}(t),t)-g(t)^{2}\nabla_{ \mathbf{x}}\log p_{t}(\mathbf{x})\right]\mathrm{d}t+g(t)\mathrm{d}\overline{ \mathbf{w}},$ (2) $$

where $\overline{\mathbf{w}}$ is a standard Wiener process when time flows backwards from $T$ to 0, and $\mathrm{d}t$ is an infinitesimal negative timestep.

We can generate new data by running backward the reverse-time SDE ([2](#S2.E2)) when the score of each marginal distribution, $\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ is known. Score Matching [; ; ] can be used for training a score-based model $\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}(t),t)$ to estimate the score:

$$ $\boldsymbol{\theta}^{*}=\underset{\boldsymbol{\theta}}{\arg\min}\mathbb{E}_{t} \left\{\lambda(t)\mathbb{E}_{p(\mathbf{x}(0))}\mathbb{E}_{\mathbf{x}(t)\mid \mathbf{x}(0)}\left[\left\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}(t),t)- \nabla_{\mathbf{x}(t)}\log p(\mathbf{x}(t)|\mathbf{x}(0))\right\|_{2}^{2} \right]\right\},$ (3) $$

where $\lambda:[0,T]\rightarrow\mathbb{R}_{>0}$ is a positive weighting function, $t$ is uniformly sampled over $[0,T]$, $\mathbf{x}(0)\sim p_{0}(\mathbf{x})$ and $\mathbf{x}(t)\sim p(\mathbf{x}(t)|\mathbf{x}(0))$. The local consistency of score mathcing is shown in [], i.e., $\mathbb{E}_{p(\mathbf{x}(0))}[\left\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{
x})-\nabla_{\mathbf{x}}\log p(\mathbf{x})\right\|_{2}^{2}]=0\Leftrightarrow
\boldsymbol{\theta}=\boldsymbol{\theta}^{*}$ under the assumption that there exists an unque $\boldsymbol{\theta}^{*}$ such that the true score function $\nabla_{\mathbf{x}}\log p(\mathbf{x})$ can be represented by $s_{\boldsymbol{\theta}^{*}}$.
[] builds the connection between Denoising Score Matching and Score Matching, and [] further proves Sliced Score Matching can learn the consistent estimator of the oracle score and the asymptotic normality for the Sliced Score Matching.

## 3 Method

In general, there are two approaches to learn a generative model from incomplete data. The first approach is to construct a complete training data set first and then learn a generative model on the complete data. We can either delete instances (rows) or features (columns) with missing data or adopt “inpute-then-generate” paradigm, i.e., we complete the data either by traditional imputation methods or training machine learning imputation models [; ] or deep generative models for imputation tasks [; ; ; ; ; ]. However, this pipeline may bring bias to the training objective. We clarify this claim in remark [3.1](#S3.Thmtheorem1). Moreover, some of the imputation methods require the ground truth for the missing values and increase the computational costs for learning additional networks, which highlights the need for alternative unbiased approaches that can handle missing data directly and more effectively.

The second way is to learn the generative model to deal with the incomplete data directly. In [; ; ], learning from incomplete data with GAN or VAE is proposed, which require training some additional networks and have assumptions on the missing mechanism. Moreover, the case of tabular data that contains mixed-type variables was not fully considered in those previous work. Therefore, a general framework for learning on missing data and generating complete data on tabular data is needed.

The diffusion model mentioned in Section [2.2](#S2.SS2) typically requires complete data for training, which is not practical for various real-world applications. To address this limitation, we propose a diffusion-based framework designed for training diffusion models on tabular data with missing values, making it more suitable for real-world scenarios where incomplete data is common.

###### Remark 3.1 (Challenges when training generative models with missing data) .

Inspired by the analysis pipeline of “inpute-then-regress” [; ] for the prediction task, we can study a corresponding framework for the generation task. The generative model $p_{\mathbf{\phi}}$ represents the probability distribution of the synthetic data $\mathbf{x}$. Under the maximum likelihood framework, $\mathbf{\phi}^{*}=\arg\max_{\mathbf{\phi}}\mathbb{E}_{\mathbf{x}\sim p_{0}(
\mathbf{x})}[\log p_{\mathbf{\phi}}(\mathbf{x})]$. When data has missing values, the general approach, known as “impute-then-generate”, may be used in practice. In this approach, the observed data $\mathbf{x}^{\text{obs}}$ is first imputed using an imputation model $f_{\mathbf{\varphi}}$. Then, the generative model is trained by maximizing the likelihood of imputed data, i.e., $\max_{\phi}\log p_{\mathbf{\phi}}(\mathbf{x}^{\mathrm{obs}},\mathbf{x}^{
\mathrm{miss}}:=f_{\mathbf{\varphi}}(\mathbf{x}^{\mathrm{obs}}))$ and with the special case of $f_{\mathbf{\varphi}}$ being the mean imputation, we have
$\max_{\phi}\log p_{\mathbf{\phi}}(\mathbf{x}^{\mathrm{obs}},\mathbb{E}_{p_{
\mathbf{\phi}}\left(\mathbf{x}^{\mathrm{miss}}\mid\mathbf{x}^{\mathrm{obs}}
\right)}\left[\mathbf{x}^{\mathrm{miss}}\right])$.
However, this pipeline is biased because with single imputation, the conditional distribution over the missing data is discarded, and the optimal single imputation can no longer capture the data variability.

Motivated by the above observation, we instead train the model parameters $\mathbf{\phi}$ by maximizing the likelihood of $p_{\mathbf{\phi}}(\mathbf{x}^{\text{obs}},\mathbf{x}^{\text{miss}})$. In general, $p_{\mathbf{\phi}}(\mathbf{x}^{\text{obs}},\mathbf{x}^{\text{miss}})\neq p_{
\mathbf{\phi}}(\mathbf{x}^{\mathrm{obs}},\mathbb{E}_{p_{\mathbf{\phi}}\left(
\mathbf{x}^{\mathrm{miss}}|\mathbf{x}^{\mathrm{obs}}\right)}[\mathbf{x}^{
\mathrm{miss}}])$. Therefore, we propose MissDiff, a diffusion-based framework for learning on missing data. This approach aims to address the limitations of traditional impute-then-generate methods by incorporating the uncertainty in missing data directly into the learning process, resulting in a more accurate generative model that better reflects the true data distribution.

We propose the following Denoising Score Matching method for data with missing values.
Instead of using Eq ([3](#S2.E3)) for learning the score-based model $\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}(t),t)$, we propose MissDiff as solution to

$$ $\displaystyle\boldsymbol{\theta}^{*}=\underset{\boldsymbol{\theta}}{\arg\min}J _{DSM}(\boldsymbol{\theta}):=$ $\displaystyle\frac{T}{2}\mathbb{E}_{t}\Big{\{}\lambda(t)\mathbb{E}_{\mathbf{x} ^{\text{obs}}(0)}\mathbb{E}_{\mathbf{x}^{\text{obs}}(t)\mid\mathbf{x}^{\text{ obs}}(0)}$ $\displaystyle\left[\left\|\left(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{ \text{obs}}(t),t)-\nabla_{\mathbf{x}^{\text{obs}}(t)}\log p(\mathbf{x}^{\text{ obs}}(t)\mid\mathbf{x}^{\text{obs}}(0))\right)\odot\mathbf{m}\right\|_{2}^{2} \right]\Big{\}},$ (4) $$

where $\lambda(t)$ is a positive weighting function, $\mathbf{m}=\mathbb{I}\{\mathbf{x}^{\text{obs}}(0)=\mathrm{na}\}$ indicated the missing entries in $\mathbf{x}^{\text{obs}}$ and $p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x}^{\text{obs}}(0))=\mathcal{N}(\mathbf{x
}^{\text{obs}}(t);\mathbf{x}^{\text{obs}}(0),\beta_{t}\mathbb{I})$ is the Gaussian transition kernel.
To make the transition $p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x}^{\text{obs}}(0))$ and the gradient $\nabla_{\mathbf{x}^{\text{obs}}(t)}\log p(\mathbf{x}^{\text{obs}}(t)\mid
\mathbf{x}^{\text{obs}}(0))$ well defined for the mixed-type data, we use 0 to replace $\mathrm{na}$ for continuous variables and a new category to represent $\mathrm{na}$ for discrete variables, which is the same operation as in [; ]. One-hot embedding is applied to discrete variables. More implementation details can be found in Appendix [B.2](#A2.SS2).

More specifically, we mainly adopt the Variance Preserving (VP) SDE in this paper although Variance Exploding (VE) SDE [] is also appliable. The forward diffusion process of the Variance Preserving SDE is defined as

$$ $\mathrm{d}\mathbf{x}=-\frac{1}{2}\beta(t)\mathbf{x}\mathrm{d}t+\sqrt{\beta(t)} \mathrm{d}\mathbf{w},$ $$

where $\left\{\beta_{t}\in(0,1)\right\}_{t\in(0,T)}$ is the increasing sequence denoting the variance schedule. Algorithm [1](#alg1) and Algorithm [2](#alg2) demonstrate the Denoising Score Matching objective on missing data and sampling procedure of MissDiff. We write $\mathbf{x}(t)$ as $\mathbf{x}_{t}$ in the algorithm box for simplicity.

Figure: Algorithm 1 MissDiff: Denoising Score Matching on Data with Missing Values

Figure: Algorithm 2 Variance Preserving Sampling of MissDiff

###### Remark 3.1 (Challenges when training generative models with missing data) .

## 4 Theory

In this section, we examine the effectiveness of MissDiff by theoretically characterizing the Score Matching objective under mild conditions on the missing mechanisms and build a further connection between Score Matching and maximizing likelihood objective for training the diffusion model.

In the following theorem, we present our first theoretical result that verifies that Denoising Score Matching on missing data can learn the oracle score, i.e, the score on complete data. Theorem [4.1](#S4.Thmtheorem1) states that the global optimal solution of Denoising Score Matching on missing data obtained by MissDiff is the same as the oracle score.
The proof can be found in Appendix [A.1](#A1.SS1).

###### Theorem 4.1 .

Denote $\rho_{i}$, $i\in\{1,2,...,d\}$ as the percentage of missing samples for the $i$-th entry in the training data. Suppose $\rho_{\text{max}}:=\max_{i=1,\ldots,d}\rho_{i}<1$. Let $\boldsymbol{\theta}^{*}$ be the solution to the training objective of MissDiff define in Eq ([3](#S3.Ex3)). Then we have

$$ $\mathbf{s}_{\boldsymbol{\theta}^{*}}(\mathbf{x}(t),t)=\nabla_{\mathbf{x}(t)} \log p_{0}(\mathbf{x}(t)).$ $$

It is well known that with careful design of the weighting function $\lambda_{t}$, Denoising Score Matching can upper bound the negative log-likelihood of the diffusion model on the complete data []. Therefore, it is straightforward to extend such connection to incomplete data scenarios, which is detailed in the following theorem. These results provide insightly connections between the training objective of MissDiff and the maximum likelihood objective of the generative model on observed data.

###### Theorem 4.2 .

The objective function of Denoising Score Matching on missing data is an upper bound for the negative likelihood of the generative model on observed data $\mathbf{x}^{\text{obs}}$ up to a constant, that is, for $\lambda_{t}=\beta_{t}$ and under the same condition of Theorem [4.1](#S4.Thmtheorem1) and mild regularity conditions detailed in Appendix [A.2](#A1.SS2), we have

$$ $-\mathbb{E}_{p(\mathbf{x}^{\text{obs}})}\left[\log p_{\mathbf{\theta}}(\mathbf {x})\right]\leq\frac{1}{1-\rho_{\text{max}}}J_{\mathrm{DSM}}\left(\mathbf{ \theta}\right)+C_{1},$ $$

where $C_{1}$ is a constant independent of $\mathbf{\theta}$.

The proof of Theorem [4.2](#S4.Thmtheorem2) can be found in Appendix [A.2](#A1.SS2). When there is missing value, Theorem [4.2](#S4.Thmtheorem2) shows that the Denoising score matching on incomplete data still upper bounds the likelihood of the incomplete data up to a constant coefficient $1/(1-\rho_{\text{max}})$.
When there is no data missing, $\mathbf{\rho}$ is all zero vector, then we have $1/(1-\rho_{\text{max}})=1$ and Theorem [4.2](#S4.Thmtheorem2) degenerates to the Corollary 1 in [], i.e.,

$$ $-\mathbb{E}_{p(\mathbf{x})}[\log p_{\mathbf{\theta}}(\mathbf{x})]\leq J_{ \mathrm{DSM}}(\boldsymbol{\theta};g(\cdot)^{2})+C_{1},$ $$

where the $J_{\mathrm{DSM}}(\boldsymbol{\theta};g(\cdot)^{2})$ is the Denoising Score Matching on complete data.

###### Theorem 4.1 .

###### Theorem 4.2 .

## 5 Experiments

In this section, we demonstrate the effectiveness of the proposed method MissDiff using simulations and two real-world tabular datasets. We introduce the experimental setup, including datasets, baseline models, and evaluation criterion, in Section [5.1](#S5.SS1). The detailed experiment results under different designs of MCAR is presented in Section [5.2](#S5.SS2). The detailed experiment results under the missing mechanisms MAR and NMAR are shown in Section [5.3](#S5.SS3).

### 5.1 Experimental Setup

##### Datasets

We present a suite of numerical evaluations of the proposed MissDiff approach on a simulated Bayesian Network data, a real Census tabular dataset [], and the MIMIC4ED tabular dataset [], with various proportions of missing values. The primary goal is to compare the effectiveness in synthetic data generation of the proposed MissDiff with the alternative baselines detailed later in this subsection.

The detailed description of the dataset can be found in Table [1](#S5.T1), which specifies the number of training data (#Train), the number of testing data (#Test), the number of categorical (discrete) variables in the tabular dataset (#Categorical), and the number of continuous variables (#Continuous). Moreover, the last column shows the evaluation task we adopted as detailed later. The details of the data generated from a Bayesian Network can be found in Appendix [B.1](#A2.SS1.SSS0.Px1).

**Table 1: Synethetic and Real-World Datasets Used in Experiments.**
| Dataset | #Train | #Test | #Categorical | #Continuous | Utility |
| --- | --- | --- | --- | --- | --- |
| Bayesian Network | 2000 | 20000 | 3 | 2 | Multi-class classification |
| Census [] | 16000 | 4000 | 9 | 6 | Binary classification |
| MIMIC4ED [] | 353150 | 88287 | 46 | 27 | Regression |

##### Choice of Masks under Different Missing Mechanisms

To evaluate the performance of MissDiff on different missing mechanisms, we give a detailed explanation of the practical implementatin of MCAR [; ; ; ], MAR, and NMAR [].

- •
MCAR: there are three types of missing mechanisms in MCAR.
–
Row Missing. For a given missing ratio $\alpha\in(0,1)$, we have the number of elements missing in each row (i.e., for each sample $\mathbf{x}_{i}$) is $\lfloor d\alpha\rfloor$, where $\lfloor z\rfloor$ is the greatest integer less than $z$, and the location/index of the missing entries is randomly chosen according to the uniform distribution.
–
Column Missing. For a given missing ratio $\alpha$, we have the number of elements missing in each column (for each feature) is $\lfloor n\alpha\rfloor$, and the location/index of the missing entries is randomly chosen according to the uniform distribution.
–
Independent Missing. Each entry in the table is masked missing according to the realization of a Bernoulli random variable with parameter $\alpha$.
- •
MAR: a fixed subset of variables that cannot have missing values is first sampled. Then, the remaining variables will have missing values according to a logistic model with random weights, which takes the non-missing variables as inputs. The outcome of this logistic model is re-scaled to attain a given missing ratio $\alpha$.
- •
NMAR: the same pipeline as MAR with the inputs of the logistic model are masked by the MCAR mechanism. We refer to [] for more detailed explanations.

###### Remark 5.1 .

Under the three missing mechanisms in MCAR, with the missing ratio parameter set as $0<\alpha<1$, condition in Theorem [4.1](#S4.Thmtheorem1) can be satisfied with probability at least $1-\delta$, where $\delta=\max\{(\frac{\alpha d-1}{d})^{n}d,\alpha,\alpha^{n}d\}$ and it will be sufficiently small when $\alpha$ is small and $n$ is sufficiently large.

Remark [5.1](#S5.Thmtheorem1) gives the guarantee that MissDiff can recover the oracle score under MCAR with high probability. In the following tables in Sections [5.2](#S5.SS2) and [5.3](#S5.SS3), we adopt the missing ratio $\alpha=0.2$ and XGBoost for the downstream tasks with no specific clarification. More experimental results can be found in Appendix [B.3](#A2.SS3).

###### Remark 5.1 .

##### Baseline Methods

We compare the proposed method with several baseline methods for synthetic data generation training on data with missing values. The baseline methods are described as follows:

- 1.
Diff-delete: Learn a vanilla diffusion model after deleting rows containing missing values.
- 2.
Diff-mean: Learn a vanilla diffusion model after imputing missing values using the mean value in that column.
- 3.
STaSy [] with the above two data completion methods. STaSy is the state-of-the-art diffusion model on tabular data, which outperforms MedGAN [], VEEGAN [], CTGAN [], TVAE [], TableGAN [], OCTGAN [], RNODE [] by a large margin.

We use the variance-preserving SDE with the time duration $T=100$ for Bayesian Network and Census dataset and $T=150$ for MIMIC4ED dataset. We use the standard pre/post-processing of tabular data to deal with mixed-type data [; ; ]. i.e., we use the min-max normalization for the continuous variables and reverse its scaler when generation. We use one-hot embedding for the discrete variables and use the rounding function after the softmax function when generation. We train the diffusion model for 250 epochs with batch size 64. For more details, please refer to Appendix [B.2](#A2.SS2).

##### Evaluation Criterion

Following [; ; ], we use two types of criterion, fidelity and utility, to evaluate the quality of the synthetic data generated. To evaluate the fidelity of synthetic data compared with real data, we adopt a model-agnostic library, SDMetrics []. The result is a float number range from 0 to 100%. The larger the score, the better the overall quality of synthetic data is.

To evaluate the utility of synthetic data, we follow the same pipeline of [], i.e., training various models, including Decision Tree, AdaBoost, Logistic/Linear Regression, MLP classifier/regressor, RandomForest, and XGBoost, on synthetic data, and validate the model on original training data, and test them with real test data. For classification tasks, we mainly use classification accuracy and also report AUROC, F1, and Weighted-F1 in Appendix [B.3](#A2.SS3). For regression tasks, we mainly use the Root Mean Squared Error (RMSE) and also report $R^{2}$ in the Appendix [B.3](#A2.SS3). All the experiments are obtained from 3 repetitions.

### 5.2 Experiment Results

#### 5.2.1 Simulation Study

Q1: How does MissDiff perform on different missing ratios against the vanilla diffusion model learned on the data completed by two baseline methods mentioned in section [5.1](#S5.SS1.SSS0.Px3)?

Figure [1](#S5.F1) summarizes the SDMetrics score on the simulated Bayesian Network dataset example. With the same diffusion model architecture and the same training hyperparameter, MissDiff achieves consistently better results against the vanilla diffusion model deleting the incomplete row or using the mean value for imputation when the missing ratio varies from 0.1 to 0.9. Moreover, the advantage of MissDiff becomes more obvious for large missing ratios. These experimental results verify the motivation of MissDiff proposed in Remark [3.1](#S3.Thmtheorem1) that the learning objective of impute-then-generate is biased. Directly learning on the missing data can significantly enhance the performance of the learned generative model.

Figure: (a) Row missing
Refer to caption: extracted/2307.00467v1/image/bayesian-row-missing.png

#### 5.2.2 Real Tabular Datasets

Q2: How does MissDiff perform on more complicated real-world data and compared with state-of-the-art generative model on tabular data?

Table [2](#S5.T2) and [3](#S5.T3) demonstrate the effectiveness of MissDiff on the Census dataset under MCAR. STaSy is a state-of-the-art generative model for tabular data, which means MissDiff achieves quite good performance on learning from incomplete data and generate complete data. More importantly, MissDiff achieves better performance than STaSy-delete and STaSy-mean even without adopting the self-paced learning technique and the fine-tuning strategy used by STaSy. Moreover, the results of STaSy-delete and STaSy-mean in Tables [2](#S5.T2) and [3](#S5.T3) are obtained by training diffusion model for 1000 epochs, compared with 250 epochs of MissDiff, Diff-delete, and Diff-mean. If we reduce the training epochs of STaSy-delete and STaSy-mean to 250 epochs, the performance will degrade significantly, which can be found in Appendix [B.3](#A2.SS3). Compared with the state-of-the-art performance of STaSy for tabular data generation, its relatively worse performance here indicates that STaSy could be susceptible to missing data.

**Table 2: Fidelity evaluation of MissDiff on Census dataset. “-” denotes the corresponding method cannot applied since no data $\mathbf{x}_{i}$ will be left after deleting the incomplete data. The larger the score, the better the overall quality of synthetic data is.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 80.59% | - | 76.92% | - | 56.75% |
| Column Missing | 82.70% | 75.03% | 76.17% | 56.90% | 51.54% |
| Independent Missing | 83.16% | 74.94% | 76.60% | 56.07% | 57.06% |

**Table 3: Utility (classification accuracy) evaluation of MissDiff on Census dataset. The larger the accuracy, the better the performance.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 79.48% | - | 78.45% | - | 70.79% |
| Column Missing | 71.68% | 72.89% | 79.60% | 68.96% | 74.47% |
| Independent Missing | 79.49% | 75.39% | 75.96% | 78.36% | 77.34% |

Q3: How does MissDiff perform on real application of large-scale Electronic Health Records data?

Table [4](#S5.T4) and [5](#S5.T5) show the performance of MissDiff on the MIMIC4ED dataset under MCAR. On this large dataset with dozens of continuous and discrete variables, MissDiff gives consistently better performance with the same training epochs (250 epochs).

**Table 4: Fidelity evaluation of MissDiff on MIMIC4ED dataset.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 84.45% | - | 75.22% | - | 82.94% |
| Column Missing | 79.24% | - | 76.57% | - | 79.03% |
| Independent Missing | 78.01% | - | 76.16% | - | 77.21% |

**Table 5: Utility (RMSE) evaluation of MissDiff on MIMIC4ED dataset. The lower the RMSE, the better the performance.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 1.826 | - | 2.166 | - | 1.894 |
| Rolumn Missing | 1.834 | - | 2.011 | - | 1.935 |
| Independent Missing | 1.852 | - | 2.483 | - | 1.972 |

### 5.3 Ablation Study

Q4: How does MissDiff perform on other missing mechanisms beyond MCAR, i.e., MAR and NMAR?

Table [6](#S5.T6) and [7](#S5.T7) demonstrate the effectiveness of MissDiff on the Census dataset beyond MCAR. The results show the great potential of learning directly on the missing data when the missing mechanism is not MCAR, which cannot be easily dealt with by previous methods [; ; ; ].

**Table 6: Fidelity evaluation of MissDiff on Census dataset under MAR, NMAR with missing ratio 0.2.**
|  | MissDiff | Diff-delete | Diff-mean |
| --- | --- | --- | --- |
| MAR | 77.45% | 73.78% | 76.08% |
| NMAR | 77.88% | 75.72% | 76.97% |

**Table 7: Utility (classification accuracy) evaluation of MissDiff on Census dataset under MAR, NMAR.**
|  | MissDiff | Diff-delete | Diff-mean |
| --- | --- | --- | --- |
| MAR | 79.95% | 69.475% | 77.425% |
| NMAR | 80.95% | 66.5% | 80.025% |

Q5: How does MissDiff perform on imputation task?

Furthermore, we validate the performance of the propose MissDiff on imputation tasks for the Census dataset. We compare MissDiff with state-of-the-art imputation approaches in Table [8](#S5.T8). The result shows that although designed for generation tasks, MissDiff also performs well for imputation tasks.

**Table 8: Imputation result comparisons on the Census dataset. The lower the RMSE, the better the performance.**
| Method | Mean /Mode | MICE(linear)[] | MissForest [] | GAIN[] | CSDI$_$T [] | MissDiff |
| --- | --- | --- | --- | --- | --- | --- |
| RMSE | 0.120 | 0.101 | 0.112 | 0.123 | 0.099 | 0.087 |

## 6 Conclusion and Discussion

We propose a diffusion-based generative framework, called MissDiff, for synthetic data generation trained on data with missing values directly. Compared with traditional methods of handling missing data, such as deletion or imputation, which may lead to reduced data diversity and biased performance, MissDiff offers a promising alternative that directly handles missing data without the need for imputation or deletion. Theoretical justification for MissDiff’s effectiveness is provided. Moreover, extensive numerical experiments demonstrate strong empirical evidence for the effectiveness of MissDiff.

##### Limitations and broader impact

Overall, this research presents a promising direction for handling missing data in generative model training. The proposed framework, MissDiff, has potential applications in a wide range of domains where missing data is a common issue. A potential limitation of this work is that it has only been empirically validated on tabular data. For future directions, it would be interesting to see how MissDiff performs empirically with more complicated data types such as video or language data. Furthermore, further research could explore the theoretical effectiveness of MissDiff on the utility perspective or differential privacy perspective.

## 7 Acknowledgement

Yidong Ouyang is partially supported by the Shenzhen Research Institute of Big Data PhD Fellowship. Yidong Ouyang and Liyan Xie are partially supported by UDF01002142 through The Chinese University of Hong Kong, Shenzhen. Guang Cheng is partially supported by Office of Naval Research, ONR (N00014-22-1-2680), NSF – SCALE MoDL (2134209), and Meta gift fund.

## Appendix A Proofs for Section 4

### A.1 Proof of Theorem 4.1

In order to show Theorem [4.1](#S4.Thmtheorem1), we aim to show that the optimal solution $\boldsymbol{\theta}^{*}$, which minimizes the objective function $J_{DSM}(\boldsymbol{\theta})$ satisfies $\mathbf{s}_{\boldsymbol{\theta}^{*}}(\mathbf{x}(t),t)=\nabla_{\mathbf{x}(t)}
\log p_{t}(\mathbf{x}(t))$, i.e., the optimal solution to the population loss function can recover the oracle score function.

For the Gaussian transition distribution that we used with the isotropic covariance matrix, the score on the incomplete data is equivalent to the score on the complete data when performing element-wise multiplication with mask, i.e., $\nabla_{\mathbf{x}^{\text{obs}}(t)}\log p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x
}^{\text{obs}}(0))\odot\mathbf{m}=\nabla_{\mathbf{x}(t)}\log p(\mathbf{x}(t)|
\mathbf{x}(0))\odot\mathbf{m}$(^1^11Assume $p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x}^{\text{obs}}(0))=\mathcal{N}(\mathbf{x }^{\text{obs}}(t);\mathbf{\mu}^{\text{obs}},\Sigma)$ and $p(\mathbf{x}(t)|\mathbf{x}(0))=\mathcal{N}(\mathbf{x}(t);\mathbf{\mu},\Sigma)$, with $\Sigma=(1-\bar{\alpha}_{t})\mathbb{I}$ and $\mathbf{\mu}^{\text{obs}}=\mathbf{\mu}\odot\mathbf{m}$. It is not hard to see $\nabla_{\mathbf{x}^{\text{obs}}(t)}\log p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x }^{\text{obs}}(0))\odot\mathbf{m}=-(\mathbf{x}^{\text{obs}}(t)-\mathbf{\mu}^{ \text{obs}})\odot\mathbf{m}=-(\mathbf{x}(t)-\mathbf{\mu})\odot\mathbf{m}= \nabla_{\mathbf{x}(t)}\log p(\mathbf{x}(t)|\mathbf{x}(0))\odot\mathbf{m}$.), where $\mathbf{m}=\mathbbm{1}\{\mathbf{x}^{\text{obs}}(0)=\mathrm{na}\}$ indicated the missing entries in $\mathbf{x}^{\text{obs}}(0)$. Therefore, under certain conditions, we may first relate the Denosing Score Matching objective on missing data to the Denosing Score Matching objective on the complete data,

$$ $\displaystyle\mathbb{E}_{p(\mathbf{x}^{\text{obs}}(0),\mathbf{m})}\mathbb{E}_{ p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x}^{\text{obs}}(0))}[\|(\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}(t),t)-\nabla_{\mathbf{x}^{\text{ obs}}(t)}\log p(\mathbf{x}^{\text{obs}}(t)|\mathbf{x}^{\text{obs}}(0)))\odot \mathbf{m}\|_{2}^{2}]$ $\displaystyle=\mathbb{E}_{p(\mathbf{x}(0),\mathbf{m})}\mathbb{E}_{p(\mathbf{x} (t)|\mathbf{x}(0))}[\|(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}(t),t)- \nabla_{\mathbf{x}(t)}\log p(\mathbf{x}(t)|\mathbf{x}(0)))\odot\mathbf{m}\|_{2 }^{2}].$ $$

Moreover, notice that we have

$$ $\displaystyle\mathbb{E}_{p(\mathbf{x}(0),\mathbf{m})}\mathbb{E}_{p(\mathbf{x}( t)|\mathbf{x}(0))}[\|(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}(t),t)-\nabla _{\mathbf{x}(t)}\log p(\mathbf{x}(t)|\mathbf{x}(0)))\odot\mathbf{m}\|_{2}^{2}]$ $\displaystyle=\mathbb{E}_{p(\mathbf{x(0)},\mathbf{x(t)})}\|(\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}(t),t)-\nabla_{\mathbf{x}(t)}\log p_{t}(\mathbf {x}(t)))\odot\sqrt{\mathbb{E}_{p(\mathbf{m}|\mathbf{x}(0))}[\mathbf{m}]}\|_{2} ^{2}],$ $$

where $\sqrt{\boldsymbol{z}}$ denotes the element-wise operation on vector $\boldsymbol{z}$. The last equation is because we take the conditional expectation of the binary mask $\mathbf{m}$ and since $\mathbf{m}_{i}\in\{0,1\}$ we have $\mathbb{E}[\mathbf{m}_{i}^{2}]=\mathbb{E}[\mathbf{m}_{i}]$ for any distribution of $\mathbf{m}$. Assuming that $\mathbb{E}_{p(\mathbf{m}|\mathbf{x}(0))}[\mathbf{m}]\equiv\mathbf{1}-
\boldsymbol{\rho}$ with $\boldsymbol{\rho}=[\rho_{1},\ldots,\rho_{d}]$ and $\rho_{i}<1$, $i\in\{1,2,...,d\}$ being the population percentage of missing samples for the $i$-th entry, we have $\mathbb{E}_{p(\mathbf{m}|\mathbf{x}(0))}[\mathbf{m}]>0$ and thus we can show the global optimal of Denoising Score Matching on missing data is the same as the oracle score.

### A.2 Proof of Theorem 4.2

The notations are defined as follows. We let $\pi$ denote the pre-specified prior distribution (e.g., the standard normal distribution), $\mathcal{C}$ denote all continuous functions, and $\mathcal{C}^{k}$ denote the family of functions with continuous $k$-th order derivatives. Consider the MCAR missing mechanism. Denote $\rho_{i}$, $i\in\{1,2,...,d\}$ as the population percentage of missing samples for the $i$-th entry in the training data. Suppose $\max_{i=1,\ldots,d}\rho_{i}<1$. In addition, we make the same mild regularity assumptions as [] in the following.

###### Assumption A.1 .

- (i)
$p(\mathbf{x})\in\mathcal{C}^{2}$ and $\mathbb{E}_{\mathbf{x}\sim p_{0}}[\|\mathbf{x}\|_{2}^{2}]<\infty$.
- (ii)
$\pi(\mathbf{x})\in\mathcal{C}^{2}$ and $\mathbb{E}_{\mathbf{x}\sim\pi}[\|\mathbf{x}\|_{2}^{2}]<\infty$.
- (iii)
$\forall t\in[0,T]:f(\cdot,t)\in\mathcal{C}^{1},\exists C>0,\forall\mathbf{x}
\in\mathbb{R}^{d},t\in[0,T]:\|f(\mathbf{x},t)\|_{2}\leq C(1+\|\mathbf{x}\|_{2})$.
- (iv)
$\exists C>0,\forall\mathbf{x},\mathbf{y}\in\mathbb{R}^{d}:\|f(\mathbf{x},t)-f(
\mathbf{y},t)\|_{2}\leq C\|\mathbf{x}-\mathbf{y}\|_{2}$.
- (v)
$g\in\mathcal{C}$ and $\forall t\in[0,T],|g(t)|>0$.
- (vi)
For any open bounded set $\mathcal{O},\int_{0}^{T}\int_{\mathcal{O}}\|p_{t}(\mathbf{x})\|_{2}^{2}+dg(t)^
{2}\|\nabla_{\mathbf{x}}p_{t}(\mathbf{x})\|_{2}^{2}\mathrm{~{}d}\mathbf{x}
\mathrm{d}t<\infty$.
- (vii)
$\exists C>0\forall\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]:\|\nabla_{\mathbf{x}}
\log p_{t}(\mathbf{x})\|_{2}\leq C(1+\|\mathbf{x}\|_{2})$.
- (viii)
$\exists C>0,\forall\mathbf{x},\mathbf{y}\in\mathbb{R}^{d}:\|\nabla_{\mathbf{x}
}\log p_{t}(\mathbf{x})-\nabla_{\mathbf{y}}\log p_{t}(\mathbf{y})\|_{2}\leq C
\|\mathbf{x}-\mathbf{y}\|_{2}$.
- (ix)
$\exists C>0\forall\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]:\|\mathbf{s}_{
\boldsymbol{\theta}}(\mathbf{x},t)\|_{2}\leq C(1+\|\mathbf{x}\|_{2})$.
- (x)
$\exists C>0,\forall\mathbf{x},\mathbf{y}\in\mathbb{R}^{d}:\|\mathbf{s}_{
\boldsymbol{\theta}}(\mathbf{x},t)-\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{y}
,t)\|_{2}\leq C\|\mathbf{x}-\mathbf{y}\|_{2}$.
- (xi)
Novikov’s condition: $\mathbb{E}[\exp(\frac{1}{2}\int_{0}^{T}\|\nabla_{\mathbf{x}}\log p_{t}(\mathbf
{x})-\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x},t)\|_{2}^{2}\mathrm{~{}d}t)]<\infty$.
- (xii)
$\forall t\in[0,T],\exists k>0:p_{t}(\mathbf{x})=O(e^{-\|\mathbf{x}\|_{2}^{k}})$ as $\|\mathbf{x}\|_{2}\rightarrow\infty$.

We mainly follow the proof strategy in [].
Consider the predefined SDE on the observed data,

$$ $\mathrm{d}\mathbf{x}^{\text{obs}}=f(\mathbf{x}^{\text{obs}},t)\mathrm{d}t+g(t) \mathrm{d}\mathbf{w},$ (5) $$

and the SDE parametrized by $\mathbf{\theta}$,

$$ $\mathrm{d}\hat{\mathbf{x}}^{\text{obs}}_{\mathbf{\theta}}=\mathbf{s}_{ \boldsymbol{\theta}}(\hat{\mathbf{x}}^{\text{obs}}_{\mathbf{\theta}},t)\mathrm {d}t+g(t)\mathrm{d}\mathbf{w}.$ (6) $$

Let $\boldsymbol{\mu}$ and $\boldsymbol{\nu}$ denote the path measure of $\{\mathbf{x}^{\text{obs}}(t)\}_{t\in[0,T]}$ and $\{\hat{\mathbf{x}}^{\text{obs}}_{\mathbf{\theta}}(t)\}_{t\in[0,T]}$, respectively. Therefore, the distribution of $p_{0}(\mathbf{x})$ and $p_{\mathbf{\theta}}(\mathbf{x})$ can be represented by the Markov kernel $K(\{\mathbf{z}(t)\}_{t\in[0,T]},\mathbf{y}):=\delta(\mathbf{z}(0)=\mathbf{y})$ as follow:

$$ $\begin{gathered}p_{0}(\mathbf{x})=\int K(\{\mathbf{x}^{\text{obs}}(t)\}_{t\in[ 0,T]},\mathbf{x})\mathrm{d}\boldsymbol{\mu}(\{\mathbf{x}^{\text{obs}}(t)\}_{t \in[0,T]}),\\ p_{\mathbf{\theta}}(\mathbf{x})=\int K(\{\hat{\mathbf{x}}^{\text{obs}}_{ \mathbf{\theta}}(t)\}_{t\in[0,T]},\mathbf{x})\mathrm{d}\boldsymbol{\nu}(\{\hat {\mathbf{x}}^{\text{obs}}_{\mathbf{\theta}}(t)\}_{t\in[0,T]}).\end{gathered}$ $$

According to the data processing inequality with this Markov kernel, the Kullback–Leibler (KL) divergence between the distribution of $p_{0}(\mathbf{x})$ and $p_{\mathbf{\theta}}(\mathbf{x})$ can be upper bounded, i.e.,

$$ $\displaystyle D_{\mathrm{KL}}(p_{0}\|p_{\mathbf{\theta}})=$ $\displaystyle D_{\mathrm{KL}}\Big{(}\int K(\{\mathbf{x}^{\text{obs}}(t)\}_{t \in[0,T]},\mathbf{x})\mathrm{d}\boldsymbol{\mu}\big{\|}\int K(\{\hat{\mathbf{x }}^{\text{obs}}_{\mathbf{\theta}}(t)\}_{t\in[0,T]},\mathbf{x})\mathrm{d} \boldsymbol{\nu}\Big{)}\leq D_{\mathrm{KL}}(\boldsymbol{\mu}\|\boldsymbol{\nu}).$ (7) $$

By the chain rule of KL divergences,

$$ $D_{\mathrm{KL}}(\boldsymbol{\mu}\|\boldsymbol{\nu})=D_{\mathrm{KL}}(p_{T}\|\pi )+\mathbb{E}_{\mathbf{z}\sim p_{T}}[D_{\mathrm{KL}}(\boldsymbol{\mu}(\cdot\mid \mathbf{x}^{\text{obs}}(T)=\mathbf{z})\|\boldsymbol{\nu}(\cdot\mid\hat{\mathbf {x}}^{\text{obs}}_{\mathbf{\theta}}(T)=\mathbf{z}))].$ (8) $$

Under assumptions (i) (iii) (iv) (v) (vi) (vii) (viii), the SDE in Eq ([5](#A1.E5)) has a corresponding reverse-time SDE given by

$$ $\mathrm{d}\mathbf{x}^{\text{obs}}=[f(\mathbf{x}^{\text{obs}},t)-g(t)^{2}\nabla _{\mathbf{x}^{\text{obs}}}\log p_{t}(\mathbf{x}^{\text{obs}})]\mathrm{d}t+g(t) \mathrm{d}\overline{\mathbf{w}}.$ (9) $$

Since Eq ([9](#A1.E9)) is the time reversal of Eq ([5](#A1.E5)), it induces the same path measure $\boldsymbol{\mu}$. As a result, $D_{\mathrm{KL}}(\boldsymbol{\mu}(\cdot\mid\mathbf{x}^{\text{obs}}(T)=\mathbf{z
})\|\boldsymbol{\nu}(\cdot\mid\hat{\mathbf{x}}^{\text{obs}}_{\mathbf{\theta}}(
T)=\mathbf{z}))$ can be viewed as the KL divergence between the path measures induced by the following two (reverse-time) SDEs:

$$ $\begin{gathered}\mathrm{d}\mathbf{x}^{\text{obs}}=[f(\mathbf{x}^{\text{obs}},t )-g(t)^{2}\nabla_{\mathbf{x}^{\text{obs}}}\log p_{t}(\mathbf{x}^{\text{obs}})] \mathrm{d}t+g(t)\mathrm{d}\overline{\mathbf{w}},\quad\mathbf{x}^{\text{obs}}(T )=\mathbf{x}^{\text{obs}},\\ \mathrm{d}\hat{\mathbf{x}}^{\text{obs}}=[f(\hat{\mathbf{x}}^{\text{obs}},t)-g( t)^{2}\mathbf{s}_{\boldsymbol{\theta}}(\hat{\mathbf{x}}^{\text{obs}},t)] \mathrm{d}t+g(t)\mathrm{d}\overline{\mathbf{w}},\quad\hat{\mathbf{x}}^{\text{ obs}}_{\mathbf{\theta}}(T)=\mathbf{x}^{\text{obs}}.\end{gathered}$ $$

Under assumptions (vii) (viii) (ix) (x) (xi), we apply the Girsanov Theorem II [[], Theorem 8.6.6], together with the martingale property of Itô integrals, which yields

$$ $\displaystyle D_{\mathrm{KL}}(\boldsymbol{\mu}(\cdot\mid\mathbf{x}^{\text{obs} }(T)=\mathbf{z})\|\boldsymbol{\nu}(\cdot\mid\hat{\mathbf{x}}^{\text{obs}}_{ \mathbf{\theta}}(T)=\mathbf{z}))$ (10) $\displaystyle=\mathbb{E}_{\boldsymbol{\mu}}[\frac{1}{2}\int_{0}^{T}g(t)^{2}\| \nabla_{\mathbf{x}^{\text{obs}}(t)}\log p_{t}(\mathbf{x}^{\text{obs}}(t))- \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}(t),t)\|_{2}^{2} \mathrm{~{}d}t]$ $\displaystyle\leq\frac{1}{2(1-\rho_{\text{max}})}\int_{0}^{T}\mathbb{E}_{p_{t} (\mathbf{x}^{\text{obs}}(t))}[g(t)^{2}\|\nabla_{\mathbf{x}^{\text{obs}}(t)} \log p_{t}(\mathbf{x}^{\text{obs}}(t))-\mathbf{s}_{\boldsymbol{\theta}}( \mathbf{x}^{\text{obs}}(t),t)\odot\sqrt{\mathbf{1}-\boldsymbol{\rho}}\|_{2}^{2 }]\mathrm{d}t$ $\displaystyle=\frac{1}{2(1-\rho_{\text{max}})}\int_{0}^{T}\mathbb{E}_{p_{t}( \mathbf{x}^{\text{obs}}(t))}[g(t)^{2}\|\nabla_{\mathbf{x}^{\text{obs}}(t)}\log p _{t}(\mathbf{x}^{\text{obs}}(t))-\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{ \text{obs}}(t),t)\odot\mathbf{m}\|_{2}^{2}]\mathrm{d}t$ $\displaystyle=\frac{1}{1-\rho_{\text{max}}}J_{\mathrm{SM}}(\mathbf{\theta};g( \cdot)^{2}),$ $$

where $\rho_{\text{max}}=\max_{i=1,\ldots,d}\rho_{i}$ and $1-\rho_{\text{max}}>0$ by assumption.
Combining Eqs. ([7](#A1.E7)), ([8](#A1.E8)) and ([10](#A1.E10)), we have $D_{\mathrm{KL}}(p_{0}\|p_{\boldsymbol{\theta}})\leq\frac{1}{1-\rho_{\text{max}
}}J_{\mathrm{SM}}(\boldsymbol{\theta};g(\cdot)^{2})+D_{\mathrm{KL}}(p_{T}\|\pi)$, which further yields $-\mathbb{E}_{p(\mathbf{x}^{\text{obs}})}[\log p_{\mathbf{\theta}}(\mathbf{x})]
\leq\frac{1}{1-\rho_{\text{max}}}J_{\mathrm{DSM}}(\boldsymbol{\theta};g(\cdot)
^{2})+C_{1}$ by Lemma [A.2](#A1.Thmtheorem2), where $C_{1}$ is a constant independent of $\mathbf{\theta}$.

###### Lemma A.2 .

Denosing Score Matching on missing data is equivalent to Score Matching on missing data, i.e.,

$$ $\displaystyle\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}})}[\|(\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)-\nabla_{\mathbf{x}^{\text{ obs}}}\log p_{t}(\mathbf{x}^{\text{obs}}_{t}))\odot\mathbf{m}\|_{2}^{2}]$ (11) $\displaystyle=\mathbb{E}_{p(\mathbf{x}^{\text{obs}}_{0})}\mathbb{E}_{p(\mathbf {x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{obs}}_{0})}[\|(\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)-\nabla_{\mathbf{x}^{\text{ obs}}_{t}}\log p(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{obs}}_{0})) \odot\mathbf{m}\|_{2}^{2}]+C,$ $$

where $\mathbf{m}=\mathbbm{1}\{\mathbf{x}^{\text{obs}}_{0}=\mathrm{na}\}$ indicated the missing entries in $\mathbf{x}^{\text{obs}}$ and $C$ is a constant that does not depend on $\boldsymbol{\theta}$. We interchange $\mathbf{x}^{\text{obs}}(t)$ with $\mathbf{x}^{\text{obs}}_{t}$.

###### Proof.

We begin with the Score Matching on the left-hand side of ([11](#A1.E11))

$$ LHS $\displaystyle=\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\|(\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)-\nabla_{\mathbf{x}^{\text{ obs}}_{t}}\log p_{t}(\mathbf{x}^{\text{obs}}_{t}))\odot\mathbf{m}\|_{2}^{2}]$ (12) $\displaystyle=\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\|\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)\odot\mathbf{m}\|^{2}]-S( \theta)+C_{2},$ $$

where $C_{2}=\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\|\nabla_{\mathbf{x}^{
\text{obs}}_{t}}\log p_{t}(\mathbf{x}^{\text{obs}}_{t})\odot\mathbf{m}\|^{2}]$ is a constant that does not depend on $\boldsymbol{\theta}$, and

$$ $\displaystyle S(\theta)$ $\displaystyle=2\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\langle\mathbf{ s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t),\nabla_{\mathbf{x}^{ \text{obs}}_{t}}\log p_{t}(\mathbf{x}^{\text{obs}}_{t})\odot\mathbf{m}\rangle]$ $\displaystyle=2\int_{\mathbf{x}^{\text{obs}}_{t}}p_{t}(\mathbf{x}^{\text{obs}} _{t})\langle\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t), \nabla_{\mathbf{x}^{\text{obs}}_{t}}\log p_{t}(\mathbf{x}^{\text{obs}}_{t}) \odot\mathbf{m}\rangle\mathrm{~{}d}\mathbf{x}^{\text{obs}}_{t}$ $\displaystyle=2\int_{\mathbf{x}^{\text{obs}}_{t}}\langle\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t),\nabla_{\mathbf{x}^{\text{ obs}}_{t}}p_{t}(\mathbf{x}^{\text{obs}}_{t})\odot\mathbf{m}\rangle\mathrm{~{}d }\mathbf{x}^{\text{obs}}_{t}$ $\displaystyle=2\int_{\mathbf{x}^{\text{obs}}_{t}}\langle\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t),\frac{\mathrm{d}}{\mathrm{ d}\mathbf{x}^{\text{obs}}_{t}}\int_{\mathbf{x}^{\text{obs}}_{0}}p_{0}(\mathbf{ x}^{\text{obs}}_{0})p(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{obs}}_{ 0})\odot\mathbf{m}\mathrm{~{}d}\mathbf{x}^{\text{obs}}_{0}\rangle\mathrm{~{}d} \mathbf{x}^{\text{obs}}_{t}$ $\displaystyle=2\int_{\mathbf{x}^{\text{obs}}_{t}}\int_{\mathbf{x}^{\text{obs}} _{0}}p_{0}(\mathbf{x}^{\text{obs}}_{0})p(\mathbf{x}^{\text{obs}}_{t}\mid \mathbf{x}^{\text{obs}}_{0})\langle\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x} ^{\text{obs}}_{t},t),\frac{\mathrm{d}\log p(\mathbf{x}^{\text{obs}}_{t}\mid \mathbf{x}^{\text{obs}}_{0})}{\mathrm{d}\mathbf{x}^{\text{obs}}_{t}}\odot \mathbf{m}\rangle\mathrm{~{}d}\mathbf{x}^{\text{obs}}_{0}\mathrm{d}\mathbf{x}^ {\text{obs}}_{t}$ $\displaystyle=2\mathbb{E}_{p(\mathbf{x}^{\text{obs}}_{t},\mathbf{x}^{\text{obs }}_{0})}[\langle\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t ),\frac{\mathrm{d}\log p(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{obs} }_{0})}{\mathrm{d}\mathbf{x}^{\text{obs}}_{t}}\odot\mathbf{m}\rangle].$ $$

Substituting this expression for $S(\theta)$ into
Eq ([12](#A1.E12)) yields

$$ LHS $\displaystyle=\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\|\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)\odot\mathbf{m}\|^{2}]$ (13) $\displaystyle-2\mathbb{E}_{p(\mathbf{x}^{\text{obs}}_{t},\mathbf{x}^{\text{obs }}_{0})}[\langle\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t ),\frac{\mathrm{d}\log p(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{obs} }_{0})}{\mathrm{d}\mathbf{x}^{\text{obs}}_{t}}\odot\mathbf{m}\rangle]+C_{2}.$ $$

On the other hand, we also have the Denoising Score Matching objective on the right-hand side of ([11](#A1.E11)) is

$$ RHS $\displaystyle=\mathbb{E}_{p_{t}(\mathbf{x}^{\text{obs}}_{t})}[\|\mathbf{s}_{ \boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t)\odot\mathbf{m}\|^{2}]$ (14) $\displaystyle-2\mathbb{E}_{p(\mathbf{x}^{\text{obs}}_{t},\mathbf{x}^{\text{obs }}_{0})}[\langle\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}^{\text{obs}}_{t},t ),\frac{\mathrm{d}\log p_{t}(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{ obs}}_{0})}{\mathrm{d}\mathbf{x}^{\text{obs}}_{t}}\rangle\odot\mathbf{m}]+C_{3},$ $$

where $C_{3}=\mathbb{E}_{p(\mathbf{x}^{\text{obs}}_{t},\mathbf{x}^{\text{obs}}_{0})}[
\|\frac{\mathrm{d}\log p_{t}(\mathbf{x}^{\text{obs}}_{t}\mid\mathbf{x}^{\text{
obs}}_{0})}{\mathrm{d}\mathbf{x}^{\text{obs}}_{t}}\odot\mathbf{m}\|^{2}]+C$ is a constant that does not depend on $\boldsymbol{\theta}$.

Comparing equations ([13](#A1.E13)) and ([14](#A1.E14)), we thus show that the two optimization objectives are equivalent up to a constant.
∎

###### Assumption A.1 .

###### Lemma A.2 .

###### Proof.

## Appendix B More Details on Experiments

### B.1 Datasets

##### Details of the Bayesian Network

Figure [2](#A2.F2) demonstrates the Bayesian Network for generating the tabular data. It contains two continuous variables C1, C2, and three discrete random variables D1, D2, and D3.
The distribution of these variables is set as follows. The marginal distribution of C1 is $\mathcal{N}(25,2)$, the conditional distribition of C2 given C1 is $\text{C2}|\text{C1}\sim\mathcal{N}(0.1\cdot\text{C1}+50,5)$, and the marginal distribution of D1 is $Bernoulli(0.3)$, where $Bernoulli(\xi)$ stands for the Bernoulli distribution with mean equal to $\xi$.
The conditional distribution of D2, given C1, C2 and D1, is set as

$$ $\text{D2}|\text{C1},\text{C2},\text{D1}\sim\begin{cases}Ca(0.3,0.6,0.1)&\text{ C1}>26,\text{C2}>55,\text{D1}=1;\\ Ca(0.2,0.3,0.5)&\text{C1}>26,\text{C2}\leq 55,\text{D1}=1;\\ Ca(0.7,0.1,0.2)&\text{C1}\leq 26,\text{C2}>55,\text{D1}=1;\\ Ca(0.1,0.2,0.7)&\text{C1}\leq 26,\text{C2}\leq 55,\text{D1}=1;\\ Ca(0.05,0.05,0.9)&\text{D1}=0,\\ \end{cases}$ $$

where $Ca(p1,p2,1-p1-p2)$ denotes the categorical (discrete) distribution for three pre-specified categories. The conditional distribution of D3 given D2 is

$$ $\text{D3}|\text{D2}\sim\begin{cases}Bernoulli(0.2)&\text{D2}=0;\\ Bernoulli(0.4)&\text{D2}=1;\\ Bernoulli(0.8)&\text{D2}=2.\end{cases}$ $$

Figure: Figure 2: The demonstration of the Bayesian Network for generating the tabular data. “C1” and “C2” denote the continuous variables and “D1”, “D2”, “D3” denotes the discrete random variables. The marginal/conditional distributions for each node are detailed in Section [B.1](#A2.SS1).
Refer to caption: x1.png

### B.2 Implementation Details

We adopt four layers residual network as the backbone of the diffusion model. The dimension of the diffusion embedding is 128 with channels as 64. We set the minimum noise level $\beta_{1}=0.0001$ and the maximum noise level $\beta_{T}=0.5$ in Algorithm [1](#alg1) and Algorithm [2](#alg2) with quadratic schedule

$$ $\beta_{t}=\left(\frac{T-t}{T-1}\sqrt{\beta_{1}}+\frac{t-1}{T-1}\sqrt{\beta_{T} }\right)^{2}.$ $$

We mainly follow the hyperparameter in the previous works that train the diffusion model on tabular data [; ].
We use the Adam optimizer with MultiStepLR with 0.1 decay at $25\%,50\%,75\%$, and $90\%$ of the total epochs and with an initial learning rate as 0.0005.

With regard to the baselines of STaSy, we adopt the same setting of its open resource implementation (^2^22[https://openreview.net/forum?id=1mNssCWt_v](https://openreview.net/forum?id=1mNssCWt_v)), i.e., Varaince Exploding SDE with six layers ConcatSquash network as the backbone of the diffusion model and Fourier embedding, the adam optimizer with learning rate as 2e-03, training with batch size 64 and 250 epochs/1000 epochs with additional 50 finetuning epochs.

For the downstream classifier/regressor, we adopt the same base hyperparameters in [[], Table 26].

### B.3 Additional Experiental Results

#### B.3.1 Additional Results for Other Criteria for Utility Evaluation

Table [9](#A2.T9), [10](#A2.T10), and [11](#A2.T11) provide the additional experimental results for other criteria under Utility evaluation for Table [3](#S5.T3), [5](#S5.T5), and [7](#S5.T7) in the main paper, i.e., the F1, Weighted-F1, AUROC for the classification task and $R^{2}$ for the regression task. A detailed explanation of the above-mentioned criteria can be found in []. To make our paper self-contained, we briefly restate it here.

- 1.
Binary F1 for binary classification: sklearn.metrics.f1_score with ‘average’=‘binary’.
- 2.
Macro F1 for multi-class classification: sklearn.metrics.f1_score with ‘average’=‘macro’.
- 3.
Weighted-F1: $=\sum_{i=0}^{K}w_{i}s_{i}$, where $K$ denotes the number of classes, the weight of $i$-th class $w_{i}$ is $\frac{1-p_{i}}{K-1},p_{i}$ is the proportion of $i$-th class’s cardinality in the whole dataset, and score $s_{i}$ is a per-class F1 of $i$-th class (in a One-vs-Rest manner).
- 4.
AUROC: sklearn.metrics.roc_auc_score.

From the results in Table [9](#A2.T9), [10](#A2.T10), and [11](#A2.T11), it can be seen that the proposed MissDiff consistently outperforms the compared methods in most instances. For the column missing case, MissDiff tends to perform worse, which indicates the potential limitations of the proposed method for future investigations.

**Table 9: Utility evaluation of MissDiff on Census dataset with other criteria. “-” denotes the corresponding method cannot applied since no data $\mathbf{x}_{i}$ will be left after deleting the incomplete data.**
| Criterian | Missing Mechanism | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- | --- |
| Binary F1 | Row Missing | 0.344 | - | 0.280 | - | 0.314 |
| Column Missing | 0.141 | 0.063 | 0.413 | 0.509 | 0.383 |  |
| Independent Missing | 0.291 | 0.045 | 0.225 | 0.274 | 0.241 |  |
| Weighted-F1 | Row Missing | 0.470 | - | 0.423 | - | 0.488 |
| Column Missing | 0.305 | 0.249 | 0.523 | 0.571 | 0.490 |  |
| Independent Missing | 0.431 | 0.237 | 0.375 | 0.416 | 0.389 |  |
| AUROC | Row Missing | 0.772 | - | 0.685 | - | 0.731 |
| Column Missing | 0.539 | 0.469 | 0.757 | 0.750 | 0.637 |  |
| Independent Missing | 0.650 | 0.474 | 0.655 | 0.621 | 0.613 |  |

**Table 10: Utility evaluation of MissDiff on MIMIC4ED dataset with $R^{2}$ criterion.**
| Missing mechanism | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 0.088 | - | 0.057 | - | 0.067 |
| Rolumn Missing | 0.095 | - | 0.023 | - | 0.073 |
| Independent Missing | 0.156 | - | 0.062 | - | 0.142 |

**Table 11: Utility evaluation of MissDiff on Census dataset under MAR, NMAR with other criteria.**
| Criterian | Missing Mechanism | MissDiff | Diff-delete | Diff-mean |
| --- | --- | --- | --- | --- |
| Binary F1 | MAR | 0.346 | 0.108 | 0.224 |
| NMAR | 0.464 | 0.233 | 0.383 |  |
| Weighted-F1 | MAR | 0.473 | 0.276 | 0.376 |
| NMAR | 0.564 | 0.364 | 0.501 |  |
| AUROC | MAR | 0.833 | 0.441 | 0.774 |
| NMAR | 0.834 | 0.499 | 0.746 |  |

#### B.3.2 Experiment Results for Different Classifiers/Regressors

As mentioned in section [5.1](#S5.SS1.SSS0.Px4), we train various models, including Decision Tree, AdaBoost, Logistic/Linear Regression, MLP classifier/regressor, RandomForest, and XGBoost, on synthetic data. Table [12](#A2.T12) to [16](#A2.T16) present the corresponding results on different classifiers/regressors, from which we can see that MissDiff still performs well under most cases.

**Table 12: Utility evaluation of MissDiff on Census dataset by Decision Tree.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 78.08% | - | 74.55% | - | 60.74% |
| Column Missing | 62.65% | 69.10% | 78.88% | 65.38% | 66.31% |
| independent | 80.68% | 72.68% | 67.70% | 76.35% | 55.99% |

**Table 13: Utility evaluation of MissDiff on Census dataset by AdaBoost.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 80.38% | - | 79.28% | - | 73.23% |
| Column Missing | 72.18% | 76.30% | 80.65% | 69.60% | 42.24% |
| independent | 78.70% | 76.13% | 75.96% | 76.55% | 78.39% |

**Table 14: Utility evaluation of MissDiff on Census dataset by Logistic Regression.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 79.20% | - | 77.08% | - | 71.04% |
| Column Missing | 73.50% | 76.30% | 77.45% | 66.91% | 69.08% |
| independent | 76.20% | 76.30% | 76.25% | 77.13% | 69.68% |

**Table 15: Utility evaluation of MissDiff on Census dataset by Multi-layer Perceptron (MLP).**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 77.70% | - | 75.13% | - | 49.78% |
| Column Missing | 68.33% | 65.75% | 75.00% | 70.97% | 58.83% |
| independent | 75.33% | 72.18% | 74.30% | 76.81% | 37.59% |

**Table 16: Utility evaluation of MissDiff on Census dataset by Random Forest.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 80.10% | - | 77.13% | - | 72.68% |
| Column Missing | 73.68% | 76.33% | 79.88% | 74.70% | 71.58% |
| independent | 79.33% | 76.30% | 76.38% | 76.31% | 76.98% |

#### B.3.3 Additional Results for STaSy-delete and STaSy-mean

In section [5.2.2](#S5.SS2.SSS2), we mentioned if we train STaSy-delete and STaSy-mean as the same training epochs (250 epochs) on the Census dataset under MCAR as MissDiff, their performance is significantly worse, which are demonstrated in Table [17](#A2.T17) and [18](#A2.T18). This observation highlights that the proposed MissDiff requires considerably fewer training epochs compared to STaSy in order to achieve satisfactory results when handling data with missing values.

**Table 17: Fidelity evaluation of MissDiff on Census dataset with 250 training epochs.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 80.59% | - | 76.92% | - | 50.08% |
| Column Missing | 82.70% | 75.03% | 76.17% | 52.49% | 49.63% |
| independent | 83.16% | 74.94% | 76.60% | 53.7% | 50.11% |

**Table 18: Utility evaluation of MissDiff on Census dataset with 250 training epochs.**
|  | MissDiff | Diff-delete | Diff-mean | STaSy-delete | STaSy-mean |
| --- | --- | --- | --- | --- | --- |
| Row Missing | 79.48% | - | 78.45% | - | 60.96% |
| Column Missing | 71.68% | 72.89% | 79.60% | 56.19% | 61.46% |
| independent | 79.49% | 75.39% | 75.96% | 49.78% | 70.68% |

### B.4 Computational Time

All the experiments are conducted on NVIDIA A100 Tensor Core GPUs. It takes around 30 minutes for each experiment on Bayesian Network, around 5 hours for each experiment on the Census dataset, and around one day for each experiment on the MIMIC4ED dataset.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Why does the imputation work? 

Algorithm 2 describes the imputation algorithm which is similar to the reverse process described in eq 2 using a **joint score model**. However, the training process only add noise to the observed coordinates, and the unobserved coordinates does not contribute to the loss (due to "odot M"), so the learned score model seems to be a "marginal score".

This is different from training a conditional score model then using it to sample from a conditional probability distribution. What is the justification of authors' imputation process? 

Could authors please write down: 

What is the optimal solution of equation 5? what is the corresponding forward process and backward process that give rise to the imputation algorithm? 

2. equation (12) in Li et al., 2019 defines the imputer of misGAN. I can just set m = [0, 0, 0... ] in (12) to generate new sample, right?

3. The same question above. Equation (2) and (3) in Yoon et al., 2018 defines the imputer of GAIN. Can I just set m = [0, 0, 0, ...] to generate new sample straightforwardly from GAIN? 

4. If yes to question 2 and 3 , what is the unique advantage of the proposed unified framework for imputing/generating samples comparing to misGAN and GAIN?


-------------------------------------------------------------

After reading the final rebuttal from the authors, I am convinced that the imputation procedure is mathematically correct.
soundness: 3
strengths: The propose method is novel to my best knowledge. 

It has nice intuition and theoretical justifications (minimizing an upper bound of  KL[p_0|p_theta])

Indeed, some existing methods do have to perform two stage inference to generate new samples. The one-short approach is not only algorithmically simpler but also avoids biases introduced in imputation step. 

The proposed method avoids unstable adversarial training which is what many existing works (such as Yoon et al., 2018 and Li et al., 2019) rely on.
summary: In this paper, authors propose a diffusion generative model that learns the complete score function from missing observations. The proposed method not only impute missing components, but also generate new samples from the data distribution in "one-shot". The proposed method is justified through minimizing a likelihood upper bound of the observed data. Experiments demonstrate that the proposed method outperforms the state-of-the-art methods on tabular datasets.
weaknesses: This paper lacks a proper motivation for generating new samples in addition to imputing missing data. In the second paragraph in the Introduction, there is one sentence on "deep generative models can be used to... enhance the performance of image classification tasks". However, I suggest authors discuss why imputation alone is not enough, and why additional samples need to be generated and back it up with specific examples. 

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

The idea of misGAN is also very similar to the proposed algorithm. On the high level, it tries to generate samples looks like the observed data  after masking. It trains a generator by minimizing a divergence between observed data (p_obs) and generated samples (p_theta). The KL divergence between p_obs and p_theta is the same as the cross entropy in Theorem 3.3 up to a constant. I would appreciate discussions on the similarities and differences between the proposed method and misGAN. 
   - There is another work that similarly maximizes the observed likelihood using normalizing flow:  https://arxiv.org/pdf/2003.12628 which I think could also be straightforwardly extended to generating new samples.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See above, also

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

2. Results in table 4 and appendix C2 provide some analysis on Q1. But why MICE is not compared? From my experience, MICE is a pretty strong baseline for these two scenarios (MAR and NMAR).
soundness: 3
strengths: Strengths: 
1. The paper is well-written and motivated. 
2. Experimental results demonstrate MissDiff's performance advantages over existing approaches.
3. Authors also provide promising simulation results.
summary: The authors present MissDiff, a diffusion based model to handle datasets with missing data. In addition to the generation task, it can also impute missing values.  Theoretical results is provided showing the connection between the training objective of MissDiff and the maximum likelihood objective of observed data. Experiments applied on multiple tabular datasets demonstrate its usefulness.
weaknesses: Weakness
1.  The paper's theoretical claims appear overstated, particularly regarding its potential as a "general framework" for handling incomplete data. Author claim "our method and theoretical guarantees aim to provide a general framework for learning on incomplete data and generate complete data."
Missing data assumptions are crucial for statistical efficiency.  And theoretical results presented in the paper ignoring these observations. Ignoring observations with missing elements may lead to invalid inference (imputation) and loss of statistical efficiency. As a result, MissDiff might have better experimental results and can't derive the identifiability result. 
My questions are: 
- Could authors clarify how do the theoretical guarantees change under different missing data mechanisms (e.g. MCAR, MAR, MNAR)?
It seems to me MissDiff is a general imputation framework, what's the strong assumption under the hood? 
- Can any connections be built between (partial) identifiability and diffusion? If not, can authors demonstrate why the "general framework" considered in the paper is better than assuming missing data mechanisms (from a theoretical and utility perspective)? For example, when we don't have a user-defined assumption, how MissDiff can benefit users?

2. The technical novelty is limited. The algorithm is not that much different from score-based diffusion. This is a minor concern. 

Overall,  W1 is my main concern. While MissDiff shows promising empirical results, its theoretical foundations and technical innovations worth more development. It would be very exciting to see an **identifiable** diffusion based models for imputation, e.g., [1], or the 

[1] Identifiable Generative Models for Missing Not at Random Data Imputation, NeurIPS 2021

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See weakness.
soundness: 1
strengths: see summary.
summary: This paper introduces *MissDiff*, a novel imputation method designed to address missing values in tabular datasets. Motivated by recent advances in using generative models like VAEs and GANs for imputation, *MissDiff* incorporates a diffusion model to handle incomplete data. Experimental results demonstrate the effectiveness of this approach, highlighting its potential to improve imputation accuracy.
weaknesses: - The authors claim that the imputation-then-gernerate method may lead to inconsistent estimator. However, MISSDIFF the author proposed is trying to fill 0 to the observed data (the i-th element of $x_{obs}$ equals to 0, when the i-th element is missing) and use this imputed data to train the model. The only difference is that the author use a mask to relieve the effect of the loss function. As such, since MISSDIFF is also kind of  imputation-then-gernerate method, why MISSDIFF does not suffer from the issue of imputation method? Is there any intuition to explain the superiority of MISSDIFF?
- The proof of Thereom 3.2 is not convincing. In the footnote 8, page 15, the author claim the score function of gaussian product the mask with the partial-observed data can recover the one with fully observed data. However, the property does not hold for $s_{\theta}(x(t),t)$. If one impute the i-th element of x(t) by 0, then possibly every element of $s_{\theta}(x(t),t)$ will change. In that case, how to show the equation line at 773-778?
- I did not see any challenge in extending the proposed method to scale up for addressing missing data issues in high-dimensional cases, such as image imputation. Could you provide relevant results to demonstrate the method's effectiveness in these scenarios?

---

confidence: 3
contribution: 4
initial_rating: 6
presentation: 3
questions: - [UPDATE: addressed] Line 351 "MICE based on random forest (MissForest) (Stekhoven, 2015)". MissForest is not exactly following the MICE framework using random forest. Are you comparing to MICE-Forest (presumably the implementation of AnotherSamWilson/miceforest) or truly MissForest (presumably the implementation of epsilon-machine/missingpy)? Please clarify such details in the manuscript.

- [UPDATE: addressed] While leaving the architectural details to the manuscript appendix is fine, it would be nice if  the main text indicated that a neural network model architecture is used.
soundness: 3
strengths: 1. This paper offers a general, rigorous framework for learning diffusion models from data with missingness. I think the tabular ML subfield, and other ML subfields that face missingness will benefit from this analysis. Furthermore, future multimodal models that combine textual and tabular data will surely benefit from this framework. 

2. The paper comes with extensive experiments, for example with varying missingness rates, incomplete-columns, and incomplete rows. And the experiments evaluate not only MCAR, but also the more realistic MAR and NMAR scenarios, showing benefits for these settings as well.
summary: Motivated by missingness in tabular data, this paper proposes a way to learn diffusion models via score-matching with missing values. The paper includes experiments on both real and synthetic data demonstrating the advantages of this approach versus baseline approaches.
weaknesses: 1. [UPDATE: addressed] Missing related work: There are works that utilize the ability of gradient-boosted decision trees (GBDTs) to handle missing features in order to avoid the impute-then-generate framework. Examples include [ForestDiffusion 1, DiffPuter 2] that use diffusion modeling and [UnmaskingTrees 3] that uses autoregression modeling. In particular, it is not clear whether the diffusion modeling approaches in [1, 2] can be seen as instances of the authors' proposed framework. (Granted, I realize that [1, 2, 3] may be parallel works, and that these rely on GBDTs for missingness, which is not a very general solution. So while I think these are relevant related works, I would like to emphasize that these do not diminish the potential impact of MissDiff, which is the first to enable the leap to neural networks, which will be important for modeling textual fields, and which is the first to be rigorously shown to not rely on assumptions like MCAR.)

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


[1] Jolicoeur-Martineau, Alexia, Kilian Fatras, and Tal Kachman. "Generating and imputing tabular data via diffusion and flow-based gradient-boosted trees." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

[2] Zhang, Hengrui, Liancheng Fang, and Philip S. Yu. "Unleashing the Potential of Diffusion Models for Incomplete Data Imputation." arXiv preprint arXiv:2405.20690 (2024).

[3] McCarter, Calvin. "Unmasking Trees for Tabular Data." arXiv preprint arXiv:2407.05593 (2024).

Meta-review: This paper proposes a diffusion-based framework—MissDiff—for simultaneously handling missing-value imputation and data generation. Reviewers found its unification of the two tasks promising and noted improvements over the state of the art on various tabular datasets. However, even after the rebuttal, they remained concerned that the paper’s theoretical foundations do not fully mesh with the proposed method in practice. Moreover, they suggested that MissDiff is similar to earlier work by Jolicoeur-Martineau et al., with the main distinction being that MissDiff’s objective is adapted for neural networks.

Decision: Reject