Impute missing mixed data accurately with native diffusion models
Impute-EM: Native Mixed-State Diffusion Models for Heterogeneous Data Imputation
Machine Learning
Summary
Missing data is a common problem when working with datasets that mix numbers, categories, and yes-or-no answers. Many current methods try to handle categories by turning them into continuous approximations, which doesn’t perfectly match the true nature of the data. The authors introduce Impute-EM, a new approach that directly models both continuous and discrete data types without simplifying categories. This method improves the accuracy of filling in missing values, leading to better downstream analyses that rely on complete data.
What this means in practice
- •For data analytics teams: Improve accuracy of mixed-type dataset completion to enhance predictive modeling in business intelligence systems.
- •For medical records managers: Fill in missing patient data entries accurately across numerical tests and categorical diagnoses to support clinical decision making.
Authors
Sergei Kholkin, Kirill Sokolov, Dmitry Baranchuk, Evgeny Burnaev, Alexander Korotin
Abstract
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.