# Rethinking the Diffusion Models for Numerical Tabular Data Imputation from the Perspective of Wasserstein Gradient Flow

arXiv: https://arxiv.org/abs/2406.15762 · HTML: https://mlsanity.com/paper/2406.15762/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 31 of its 56 references that are in the map.
- First posted on arXiv 22 Jun 2024.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 31)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation](https://mlsanity.com/paper/2107.03502/) (2021)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models](https://mlsanity.com/paper/2208.09399/) (2022)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [Diffusion models for missing value imputation in tabular data](https://mlsanity.com/paper/2210.17128/) (2022)
- [A Survey on Diffusion Models for Time Series and Spatio-Temporal Data](https://mlsanity.com/paper/2404.18886/) (2024)
- [Gradient Estimators for Implicit Models](https://mlsanity.com/paper/1705.07107/) (2017)
- [Deep Learning for Multivariate Time Series Imputation: A Survey](https://mlsanity.com/paper/2402.04059/) (2024)
- [MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms](https://mlsanity.com/paper/2111.03187/) (2021)
- [ESCM2: Entire Space Counterfactual Multi-Task Model for Post-Click Conversion Rate Estimation](https://mlsanity.com/paper/2204.05125/) (2022)
- [Understanding DDPM Latent Codes Through Optimal Transport](https://mlsanity.com/paper/2202.07477/) (2022)
- [A Unified Particle-Optimization Framework for Scalable Bayesian Sampling](https://mlsanity.com/paper/1805.11659/) (2018)
- [Accelerated Information Gradient flow](https://mlsanity.com/paper/1909.02102/) (2019)
- [MissDiff: Training Diffusion Models on Tabular Data with Missing Values](https://mlsanity.com/paper/2307.00467/) (2023)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods](https://mlsanity.com/paper/2305.15027/) (2023)
- [Transformed Distribution Matching for Missing Value Imputation](https://mlsanity.com/paper/2302.10363/) (2023)
- [Provably Convergent Schr"odinger Bridge with Applications to Probabilistic Time Series Imputation](https://mlsanity.com/paper/2305.07247/) (2023)
- [Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel](https://mlsanity.com/paper/2310.03054/) (2023)
- [Particle-based Variational Inference with Preconditioned Functional Gradient Flow](https://mlsanity.com/paper/2211.13954/) (2022)
- [Nonparametric Generative Modeling with Conditional Sliced-Wasserstein Flows](https://mlsanity.com/paper/2305.02164/) (2023)
- [Multiple Robust Learning for Recommendation](https://mlsanity.com/paper/2207.10796/) (2022)
- [GAD-PVI: A General Accelerated Dynamic-Weight Particle-Based Variational Inference Framework](https://mlsanity.com/paper/2312.16429/) (2023)
- [DPVI: A Dynamic-Weight Particle-Based Variational Inference Framework](https://mlsanity.com/paper/2112.00945/) (2021)

## Built on by (most cited first, 3 of 3)

- [BELM: Bidirectional Explicit Linear Multi-step Sampler for Exact Inversion in Diffusion Models](https://mlsanity.com/paper/2410.07273/) (2024)
- [Efficiently Access Diffusion Fisher: Within the Outer Product Span Space](https://mlsanity.com/paper/2505.23264/) (2025)
- [Unleashing High-Quality Image Generation in Diffusion Sampling Using Second-Order Levenberg-Marquardt-Langevin](https://mlsanity.com/paper/2505.24222/) (2025)

## Abstract (arXiv)

Diffusion models (DMs) have gained attention in Missing Data Imputation (MDI), but there remain two long-neglected issues to be addressed: (1). Inaccurate Imputation, which arises from inherently sample-diversification-pursuing generative process of DMs. (2). Difficult Training, which stems from intricate design required for the mask matrix in model training stage. To address these concerns within the realm of numerical tabular datasets, we introduce a novel principled approach termed Kernelized Negative Entropy-regularized Wasserstein gradient flow Imputation (KnewImp). Specifically, based on Wasserstein gradient flow (WGF) framework, we first prove that issue (1) stems from the cost functionals implicitly maximized in DM-based MDI are equivalent to the MDI's objective plus diversification-promoting non-negative terms. Based on this, we then design a novel cost functional with diversification-discouraging negative entropy and derive our KnewImp approach within WGF framework and reproducing kernel Hilbert space. After that, we prove that the imputation procedure of KnewImp can be derived from another cost functional related to the joint distribution, eliminating the need for the mask matrix and hence naturally addressing issue (2). Extensive experiments demonstrate that our proposed KnewImp approach significantly outperforms existing state-of-the-art methods.