# Impute-EM: Native Mixed-State Diffusion Models for Heterogeneous Data Imputation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 19 of its 25 references that are in the map.
- First posted on arXiv 14 Sep 2026.

## 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, 19 of 19)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [GP-VAE: Deep Probabilistic Time Series Imputation](https://mlsanity.com/paper/1907.04155/) (2019)
- [Handling Missing Data with Graph Representation Learning](https://mlsanity.com/paper/2010.16418/) (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)
- [TabDiff: a Mixed-type Diffusion Model for Tabular Data Generation](https://mlsanity.com/paper/2410.20626/) (2024)
- [MissDiff: Training Diffusion Models on Tabular Data with Missing Values](https://mlsanity.com/paper/2307.00467/) (2023)
- [MCFlow: Monte Carlo Flow Models for Data Imputation](https://mlsanity.com/paper/2003.12628/) (2020)
- [ReMasker: Imputing Tabular Data with Masked Autoencoding](https://mlsanity.com/paper/2309.13793/) (2023)
- [Data Imputation with Iterative Graph Reconstruction](https://mlsanity.com/paper/2212.02810/) (2022)
- [Missing Data Imputation by Reducing Mutual Information with Rectified Flows](https://mlsanity.com/paper/2505.11749/) (2025)
- [DiffEM: Learning from Corrupted Data with Diffusion Models via Expectation Maximization](https://mlsanity.com/paper/2510.12691/) (2025)
- [Token-Based Audio Inpainting via Discrete Diffusion](https://mlsanity.com/paper/2507.08333/) (2025)
- [AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking](https://mlsanity.com/paper/2606.03347/) (2026)

## Abstract (arXiv)

Missing values are ubiquitous in heterogeneous data mining, where numerical, categorical, and binary variables often coexist. Many imputation methods, especially diffusion-based ones, treat discrete variables through continuous surrogates such as one-hot relaxations rather than modeling them natively. This creates a mismatch between the model state space and the mixed discrete and continuous structure of the data. We propose Impute-EM, an Expectation Maximization style framework that alternates between imputing missing entries with the current model and refitting a diffusion backbone on completed data. We instantiate Impute-EM with native mixed-state diffusion backbones for heterogeneous data, combining Gaussian and masked categorical components without one-hot relaxations. In exact settings, we characterize the update and show that the observed mask-indexed marginals match the targets at the limit, while making explicit that the full data distribution is generally non-identifiable from incomplete observations alone. Empirically, Impute-EM delivers the best distributional fidelity on mixed-type tabular imputation, on which downstream modeling relies, with text imputation serving as a controlled validation of the native discrete backbone.