Research map / Image generation / Fast diffusion sampling
Rethinking the Diffusion Models for Missing Data Imputation: A Gradient Flow Perspective
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Reframes diffusion-model imputation as Wasserstein gradient flow, adds a negative-entropy regularizer to suppress diversity, and uses joint-distribution score modeling to avoid data masking (NewImp).
Problem
Directly applying diffusion models to missing data imputation yields suboptimal accuracy because sampling promotes diversity in imputed values, and training requires masking observed data whose mask choice is hard to align with the test missingness mechanism.
Approach
Show that DM-based imputation implicitly optimizes an objective with an extra positive-entropy term; replace it with a negative-entropy regularized cost, implement the WGF velocity field in an RBF RKHS, and use the joint distribution score (with masking only in the velocity computation) so that conditional modeling and data masking are unnecessary.
Key results
- Under MAR at 30% missing rate, NewImp achieves the best MAE/WASS on most of the eight datasets, e.g., CC WASS 0.82 vs. ReMasker 1.15 and CSDI_T 3.44 on BT-like baselines (MAR table).
- Ablation shows removing the NER term or joint modeling degrades results; joint modeling noticeably improves performance over conditional modeling (e.g., BT MAE 0.96 to 0.52 with both components).
- Sensitivity analysis on CC shows accuracy improves with larger negative-entropy weight λ and degrades with larger kernel bandwidth h.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add NER to NewImp (ablation) | NewImp · CC MAR 30% missing | MAE: 0.61 -> 0.35 | Helped |
| Add NER to NewImp (ablation) | NewImp · CC MCAR 30% missing | MAE: 0.62 -> 0.47 | Helped |
| Use joint rather than conditional distribution modeling (ablation) | NewImp · CC MAR 30% missing | MAE: 1.04 -> 0.35 | Helped |
| Use joint rather than conditional distribution modeling (ablation) | NewImp · CC MCAR 30% missing | MAE: 0.86 -> 0.47 | Helped |
| Add NER and joint modeling (ablation) | NewImp · CC MAR 30% missing | MAE: 1.04 -> 0.35 | Helped |
| Add NER and joint modeling (ablation) | NewImp · CC MCAR 30% missing | MAE: 0.85 -> 0.47 | Helped |
| Increase RBF bandwidth from 0.5 to 2.0 (ablation) | NewImp · CC MAR 30% missing | MAE: 0.35 to 0.82 | Hurt |
| Increase score-network hidden units from 256 to 512 (ablation) | NewImp · CC MAR 30% missing | MAE | Hurt |
| Increase NER strength (ablation) | NewImp · CC MAR 30% missing | MAE | Helped |
| Vary ODE discretization step size (ablation) | NewImp · CC MAR 30% missing | MAE | Mixed |
| Use NewImp rather than ReMasker | Missing-data imputation · BCD MAR 30% missing | MAE: 0.48 -> 0.34 | Helped |
| Use NewImp rather than ReMasker | Missing-data imputation · BT MCAR 30% missing | MAE: 0.46 -> 0.48 | Hurt |
Surprising findings
- Standard diffusion-model imputers (CSDI_T, MissDiff) perform worse than many non-diffusion baselines such as MIRACLE, MIWAE, and TDM, which the authors attribute to diversity-seeking sampling.
Limitations
- Restricted to MAR and MCAR settings; MNAR is not addressed in the main method.
- Relies on an RBF kernel with bandwidth and regularization hyperparameters that require tuning.
- Experiments limited to numerical tabular datasets.
Experimental setup
Datasets: UCI BT UCI BCD UCI CC UCI CBV UCI IS UCI PK UCI QB UCI WQW
Benchmarks: UCI BT UCI BCD UCI CC UCI CBV UCI IS UCI PK UCI QB UCI WQW
Metrics: MAE WASS (squared Wasserstein-2)
Baselines: CSDI_T MissDiff GAIN MIRACLE MIWAE Sinkhorn TDM ReMasker
Builds on: CSDI MissDiff denoising score matching Wasserstein gradient flow
Closest papers by meaning
- DiffPuter: Empowering Diffusion Models for Missing Data Imputation (ICLR 2025 · similarity 0.91)
- Diffusion Transformers for Imputation: Statistical Efficiency and Uncertainty Quantification (NeurIPS 2025 · similarity 0.90)
- Learning What Not to Impute: An Uncertainty-Aware Diffusion Framework for Meaningful Missingness (NeurIPS 2026 · similarity 0.90)
- Distribution Shift in Missing Data Imputation: A Risk-Based Perspective and Importance-Weighted Correction under MAR (NeurIPS 2026 · similarity 0.89)
- ProCTI: Prototype-Refined Global Conditioning for Diffusion-Based Time Series Imputation (NeurIPS 2026 · similarity 0.89)
- Discretizing Continuous Time Series for Imputation with Masked Diffusion Training (NeurIPS 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.