Continuously-Augmented Hybrid Masked Diffusion Model for Data Imputation
Abstract
Missing data is pervasive in real-world datasets, arising from sensor failures, image occlusions, and incomplete reporting. Imputation is a conditional generation problem that must preserve observed entries exactly while producing plausible values for missing ones. Yet most diffusion-based imputers adopt Gaussian noising whose global perturbations are structurally misaligned with this local hard constraint, so consistency is typically enforced through conditional training, clamping, projection, or repainting-style corrections. We address this mismatch with Continuously-Augmented Hybrid Masked Diffusion (CAHMD), a masked diffusion framework whose conditional sampler leaves observed coordinates unchanged by construction. To extend masked diffusion beyond purely discrete data, CAHMD attaches coordinate-wise auxiliary continuous latents to masked entries, providing a denoising channel before values are revealed in continuous and hybrid data. For incomplete training data, we introduce observed-mask gating, an observed-data masked reconstruction principle that uses available entries as pseudo-missing targets and avoids repeated EM-style imputation steps. Across tabular, image, and image--caption benchmarks, CAHMD gives a strong reconstruction--perception--efficiency trade-off while reducing the training overhead of EM-style self-imputation and the inference overhead of repainting-style inner loops.