Research map / Image generation / Fast diffusion sampling
DiffPuter: Empowering Diffusion Models for Missing Data Imputation
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.
DiffPuter combines a diffusion model with the EM algorithm for missing data imputation, training the diffusion model as the M-step and performing conditional reverse sampling as the E-step.
Problem
Missing data imputation in tabular data, where generative models must learn a joint distribution from incomplete training data and perform accurate conditional inference for missing entries.
Approach
Iteratively alternate a diffusion model's score-matching training (M-step, density estimation on current completed data) with a tailored reverse-sampling procedure that mixes forward-noised observed entries with reverse-denoised missing entries (E-step, EAP estimation of missing values).
Key results
- DiffPuter improves over the most competitive baseline by 6.94% in MAE and 4.78% in RMSE on average for in-sample continuous imputation (MCAR, 30% missing).
- On discrete columns, DiffPuter achieves average accuracy 62.82, ranking first among 19 methods (Remasker 62.06, HyperImpute 59.72).
- Out-of-sample MAE improves 13.37% over the best baseline per the appendix table caption.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DiffPuter instead of GRAPE | discrete-column imputation · Adult (MCAR, in-sample) | accuracy: 69.84 ± 0.25 -> 70.12 ± 0.17 | Helped |
| Use DiffPuter instead of Remasker | discrete-column imputation · Default (MCAR, in-sample) | accuracy: 76.88 ± 0.62 -> 77.64 ± 0.32 | Helped |
| Use DiffPuter instead of HyperImpute | discrete-column imputation · Shoppers (MCAR, in-sample) | accuracy: 59.19 ± 0.04 -> 58.82 ± 0.09 | Hurt |
| Use DiffPuter instead of Remasker | discrete-column imputation · News (MCAR, in-sample) | accuracy: 44.33 ± 0.08 -> 44.69 ± 0.13 | Helped |
| Iteratively refine imputations with additional EM iterations rather than use a one-step diffusion imputation (ablation) | DiffPuter · Missing-data imputation; datasets in Figure | imputation performance | Helped |
| Increase the number of sampled imputations per iteration (ablation) | DiffPuter · Missing-data imputation; datasets in Figure | imputation performance | Helped |
| Reduce diffusion sampling steps from 50 to 20 (ablation) | DiffPuter · Missing-data imputation | MAE: 3% | Hurt |
| Increase the missing ratio toward 99% (ablation) | DiffPuter · Missing-data imputation | MAE | Hurt |
| Combine EM with another deep generative model (ablation) | MIWAE, HIWAE, VAEM, and HH-VAEM · Missing-data imputation; Table | imputation score | Helped |
Surprising findings
- Traditional methods such as EM with Gaussian priors and KNN outperform many early deep generative imputation methods.
- Discriminative methods (Remasker, GRAPE, MOT) are stronger on discrete columns, while the generative DiffPuter excels on continuous columns.
Limitations
- Discrete features are one-hot encoded and treated as continuous, which is less suited to diffusion models.
- Performance is upper-bounded by the mean-based initialization at very high missing ratios.
- Discriminative methods remain stronger on discrete columns.
Experimental setup
Datasets: California Letter Gesture Magic Bean Adult Default Shoppers News
Benchmarks: California Letter Gesture Magic Bean Adult Default Shoppers News
Metrics: MAE RMSE accuracy
Baselines: MOT TDM GRAPE IGRM EM MICE MIRACLE SoftImpute MissForest MIWAE GAIN MCFlow
Builds on: EM algorithm score-based SDE diffusion models MCFlow TabDDPM
Closest papers by meaning
- Rethinking the Diffusion Models for Missing Data Imputation: A Gradient Flow Perspective (NeurIPS 2024 · 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.89)
- Distribution Shift in Missing Data Imputation: A Risk-Based Perspective and Importance-Weighted Correction under MAR (NeurIPS 2026 · similarity 0.88)
- Discretizing Continuous Time Series for Imputation with Masked Diffusion Training (NeurIPS 2026 · similarity 0.87)
- DiffPO: A causal diffusion model for learning distributions of potential outcomes (NeurIPS 2024 · similarity 0.86)
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.