AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking

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 via stochastic regularization. AugMask 1) completes inputs via conditional stochastic augmentation using lightweight auxiliary models and 2) masks the loss, using augmented missing entries for conditioning while restricting supervision to observed coordinates. We connect AugMask to a Rao-Blackwellized objective and show that marginalizing missing entries yields a variance-weighted sensitivity penalty, promoting invariance of observed-coordinate reconstruction with respect to uncertain missing entries. Across diverse datasets and missingness regimes, AugMask enables standard diffusion-based tabular generators to match or outperform specialized missing-aware baselines in both sample fidelity and downstream utility. The code will be released.

Real-world datasets, such as medical records, financial sheets, or public surveys, often contain missing information. This is a problem for modern AI models that generate realistic "synthetic" tables, because these models usually expect every cell to be filled. Common fixes, such as replacing blanks with zeros or guessing the missing numbers often force the AI to learn misleading patterns. We introduce AugMask, a training strategy for "diffusion models", the same type of AI used to generate realistic images but adapted here for spreadsheets. AugMask provides the AI with plausible, sampled values to give it context, but only grades the AI on the data points we know are real. This allows the AI to learn from the surrounding data without mistaking guesses for absolute truth. The core mechanism is simple. By feeding the AI different possible guesses for the same blank space, AugMask teaches the model when a value is highly uncertain and should be trusted less (for example, if those guesses vary wildly). Our experiments show that AugMask helps standard AI models generate more accurate and reliable synthetic datasets even from incomplete real-world data without any redesigning or engineering the existing framework. It provides a stepping stone for creating high-quality, realistic datasets for research and testing.