AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking
For practitioners using diffusion models on real-world tabular data with missing values, AugMask provides a simple plug-and-play solution that improves performance without modifying the backbone architecture.
AugMask adapts diffusion models to incomplete tabular data by using stochastic augmentation for conditioning and denoising supervision only on observed entries, enabling standard generators to outperform specialized missing-aware baselines across diverse datasets and missingness regimes.
Score-based diffusion models have emerged as prominent deep generative models; however, their application to tabular data remains challenging because their backbones assume fully specified inputs, whereas real-world tabular data often contain missing values. We propose AugMask, a plug-and-play training framework that adapts missing-unaware backbones to incomplete data by separating conditioning from supervision. AugMask 1) constructs numeric inputs via conditional stochastic augmentation using lightweight auxiliary models, and 2) applies denoising supervision only to observed coordinates. In effect, augmented missing entries serve as uncertain conditioning context rather than training targets. We connect this training rule to a Rao--Blackwellized objective and show that marginalizing missing entries yields a variance-weighted sensitivity penalty, discouraging over-reliance on uncertain completions. Across diverse datasets and missingness regimes, AugMask enables standard diffusion-based tabular generators to outperform specialized missing-aware baselines.