[Paper Review] MissDiff: Training Diffusion Models on Tabular Data with Missing Values

Yidong Ouyang, Liyan Xie+2 more|arXiv (Cornell University)|Jul 2, 2023

Bayesian Methods and Mixture Models4 citations

TL;DR

MissDiff proposes a novel diffusion-based generative framework that trains directly on tabular data with missing values without prior imputation, using a masked regression loss in Denoising Score Matching to ensure consistent score learning. It outperforms state-of-the-art methods—including 'impute-then-generate' pipelines—by large margins in fidelity and utility across multiple real-world datasets under diverse missing mechanisms.

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.

Motivation & Objective

• To address the challenge of training generative models on incomplete tabular data, which is common in healthcare, finance, and social systems.
• To overcome the limitations of the 'impute-then-generate' pipeline, which can introduce bias and reduce data diversity.
• To develop a principled, unified framework that learns from missing data without deletion or imputation.
• To provide theoretical justification for score function recovery and negative log-likelihood upper bounding under mild missing mechanisms.
• To empirically validate MissDiff’s superiority in data fidelity and downstream utility across multiple tabular datasets.

Proposed method

• Proposes a masked regression loss in Denoising Score Matching (DSM) to handle missing values during training, avoiding imputation.
• Introduces a modified training objective that directly incorporates missing data patterns into the score matching process.
• The framework is theoretically grounded, proving consistency in learning the true score function and providing an upper bound on negative log-likelihood under mild missingness assumptions.
• Applies the method to mixed-type tabular data with both continuous and categorical features, handling missingness in both variable types.
• Uses a denoising score matching objective where the loss is masked over missing entries, ensuring the model learns from observed data without bias.
• Employs a unified architecture that supports both generation and imputation, enabling dual-purpose utility.

Experimental results

Research questions

• RQ1 Can a diffusion model be effectively trained on tabular data with missing values without prior imputation?
• RQ2 Does the proposed masked loss in score matching lead to consistent estimation of the true data score function?
• RQ3 How does MissDiff compare to 'impute-then-generate' baselines in terms of data fidelity and downstream utility?
• RQ4 How well does MissDiff generalize under non-MCAR missing mechanisms such as MAR and NMAR?
• RQ5 Can MissDiff serve as an effective imputation model, even though it is designed for generation?

Key findings

• On the Census dataset under MCAR, MissDiff achieved 83.16% fidelity under independent missingness, outperforming the next best method (STaSy-mean) by 26.1%.
• Under row missingness, MissDiff achieved 80.59% classification accuracy, significantly outperforming Diff-mean (76.92%) and STaSy-mean (56.75%).
• On the large-scale MIMIC-4ED dataset, MissDiff achieved 84.45% fidelity under row missingness, surpassing Diff-delete (no performance reported) and Diff-mean (75.22%).
• In RMSE-based utility evaluation on MIMIC-4ED, MissDiff achieved 1.826 under row missingness, outperforming Diff-mean (2.166) and STaSy-mean (1.894).
• Under MAR and NMAR mechanisms on Census, MissDiff achieved 77.45% and 77.88% fidelity respectively, outperforming Diff-delete and Diff-mean across all settings.
• For imputation, MissDiff achieved an RMSE of 0.087 on the Census dataset, outperforming CSDI-T (0.099) and GAIN (0.123), demonstrating strong imputation capability.