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

2026-09-14

AI Summary

本文提出Impute-EM方法，通过交替填补缺失值和重新拟合扩散模型来处理异构数据中的缺失值问题，直接建模混合状态而不使用连续替代。

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.

Problem

• Missing values
• Heterogeneous data
• Diffusion models
• Discrete variables

Innovation

• Mixed-state diffusion models
• Heterogeneous data imputation
• Expectation Maximization
• Native discrete variable handling