Impute-EM: Native Mixed-State Diffusion Models for Heterogeneous Data Imputation
Sergei Kholkin, Kirill Sokolov, Dmitry Baranchuk, Evgeny Burnaev, Alexander Korotin
cs.LG
Quick verdict
EM-style imputation with native mixed-state diffusion backbones; strong tabular fidelity gains but incremental over DiffEM/DiffPuter and no error bars.
Summary
The paper proposes Impute-EM, an Expectation-Maximization-style framework for missing-data imputation whose distinguishing feature is a generative backbone operating natively on mixed continuous-discrete state spaces: Gaussian diffusion for numerical columns and masked absorbing diffusion for categorical/binary/text variables, avoiding one-hot relaxations. Each EM iteration samples missing entries from the current conditional diffusion model (E-step) and refits it on the completed dataset (M-step), starting from a model trained with incomplete (observed-only) likelihood. The authors formulate learning as matching mask-indexed observed marginals, prove observed-marginal consistency of the exact update (Theorem 1), and explicitly acknowledge that the full data distribution is non-identifiable from incomplete observations. Experiments cover text8 character-level imputation with MDLM (perplexity substantially recovers toward the clean model under 20-50% MCAR masking) and four tabular datasets (Adult, Shoppers, Default, Beijing) with a TabDiff backbone: under MCAR, average Overall Density error is 0.071 versus 0.108 for second-best HyperImpute, and the method attains the best average ML-efficiency rank (first on three of four datasets); advantages persist under 70% MNAR and single-run MAR. It does not lead on pointwise discriminative metrics (accuracy, MAE). Limitations: theory assumes exact E/M steps with MCAR-only guarantees; main tables report no variance or seeds; MAR is single-run; and no code release is stated.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Learning from incomplete data can be cast as matching mask-indexed observed marginals, and the exact EM update achieves observed-marginal consistency (Theorem 1) while the full data distribution remains non-identifiable.
- Impute-EM is, per Table I, the first EM-style imputation framework instantiated with native discrete and mixed-state diffusion backbones, avoiding one-hot relaxations.
- TabDiff + Impute-EM achieves the best distributional fidelity on four mixed-type tabular datasets under MCAR (average Overall Density error 0.071 vs 0.108 for second-best HyperImpute) and the best average ML-efficiency rank (first on three of four datasets).
- Gains persist under MNAR (p=0.7) and single-run MAR, and most of the EM improvement accrues within the first iterations, converging by step 5 (Fig. 5).
- On text8, MDLM + Impute-EM substantially recovers clean-model perplexity under 20-50% per-token MCAR masking, validating the native discrete backbone in isolation.
Strengths
Native mixed-state backbone removes the one-hot mismatch
Section IV-B combines conditional DDPM noise-prediction and masked categorical (MDM) cross-entropy losses in one multimodal network (Eq. 15; App. B Eqs. 27-28). Table I marks Impute-EM as the only entry with a general-diffusion backbone and native (✓) discrete handling among EM-style methods, where DiffPuter/MCFlow use one-hot workarounds (✓✗) and DiffEM/MIRI handle discrete not at all (✗).
Honest, complete theory with correctly scoped claims
Proposition 1 and Theorem 1 (Appendix A) prove observed-marginal consistency of the exact update via a KL chain-rule plus log-sum-inequality telescoping argument, and the text explicitly concedes non-identifiability: the limit lies in C (Eq. 17) and 'is not necessarily p*' (Sec. IV-C).
Large margins on the main benchmark with an internal ablation
Evaluation breadth and transparency on unfavorable results
Reproducibility hygiene for the experimental setup
Weaknesses
Limited theoretical novelty
No variance or seed reporting on headline tables
Underperforms on pointwise recovery
Missing empirical comparison to AugMask
No code-release statement
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
This method paper (arXiv:2609.15284, submitted to ICDM 2026 per the authors' disclaimer in Sec. III) targets generative imputation of heterogeneous data with missing values. The central claim is threefold: (1) learning from incomplete data should be cast as matching mask-indexed observed marginals (Eq. 8), for which an EM-style update (Prop. 1) achieves observed-marginal consistency at the limit (Theorem 1); (2) the EM template can be instantiated with native mixed-state diffusion backbones - Gaussian diffusion for continuous coordinates plus masked absorbing diffusion (MDM) for categorical ones (Sec. IV-B, Eq. 15) - avoiding the one-hot surrogates of prior continuous-diffusion imputers (MissDiff, TabCSDI); (3) empirically, TabDiff + Impute-EM attains the best distributional fidelity on four mixed-type tabular datasets under MCAR/MNAR/MAR, with text8 as a controlled validation of the discrete backbone. The authors are commendably explicit that the full data distribution is non-identifiable from incomplete observations (Secs. II-A, IV-A), so guarantees are marginal-level, not distribution recovery.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is Impute-EM and how does it work?
Impute-EM is an expectation-maximization-style framework for imputing missing values in mixed continuous-discrete data. Each iteration samples the missing entries with the current conditional diffusion model (E-step), then retrains that model on the completed dataset (M-step). It uses Gaussian diffusion for numeric variables and a masked categorical diffusion for categorical/binary/text variables, instead of the one-hot encodings common in prior diffusion imputers, and starts from a model trained only on observed coordinates.
How is Impute-EM different from DiffPuter and DiffEM?
All three alternate conditional reconstruction with model refitting. DiffPuter and DiffEM use continuous (score-based) diffusion, forcing categorical variables through one-hot or continuous surrogates; Impute-EM swaps in native mixed-state backbones (Gaussian plus masked categorical, e.g., TabDiff or an MDLM). The underlying theory - observed-marginal consistency without full-distribution identifiability - is in the same family as DiffEM's analysis, which the authors acknowledge covers general state spaces.
How well does Impute-EM perform and on what data?
On four UCI tabular datasets (Adult, Shoppers, Default, Beijing) with a TabDiff backbone, it reports the lowest average Overall Density error under MCAR (0.071 vs 0.108 for second-best HyperImpute) and the best average downstream ML-efficiency rank, leading on three of four datasets; the lead holds under 70% MNAR and single-run MAR. On text8 character imputation it substantially recovers clean-model perplexity at 20-50% masking. It does not lead on pointwise accuracy/MAE metrics.
Does Impute-EM recover the true data distribution?
No - and the paper says so explicitly. The full joint distribution is generally not identifiable from incomplete observations alone. The theory guarantees only that the observed mask-indexed marginals converge to their targets at the limit; which distribution among all marginals-compatible ones is selected is not determined by the procedure.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 1-10 appear as captions only in the extraction; per-mask-rate values, iteration curves, and MAR/MNAR bar heights could not be verified against the claims made about them.
- The text appears complete through Appendix D and references, but extraction artifacts exist, e.g., the affiliation placeholder 'AXXX' and typos ('datset', 'deliveres', 'ReduceOnPlateou'), suggesting imperfect OCR.
- The number of training seeds behind Tables II-III is not stated anywhere in the text; 'single-run' is confirmed only for the MAR appendix.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.