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

arXiv: https://arxiv.org/abs/2606.03347 · HTML: https://mlsanity.com/paper/2606.03347/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 26 of its 44 references that are in the map.
- First posted on arXiv 2 Jun 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 26 of 26)

- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Modeling Tabular data using Conditional GAN](https://mlsanity.com/paper/1907.00503/) (2019)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [TabDDPM: Modelling Tabular Data with Diffusion Models](https://mlsanity.com/paper/2209.15421/) (2022)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [Real-valued (Medical) Time Series Generation with Recurrent Conditional GANs](https://mlsanity.com/paper/1706.02633/) (2017)
- [How Faithful is your Synthetic Data? Sample-level Metrics for Evaluating and Auditing Generative Models](https://mlsanity.com/paper/2102.08921/) (2021)
- [Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space](https://mlsanity.com/paper/2310.09656/) (2023)
- [CTAB-GAN: Effective Table Data Synthesizing](https://mlsanity.com/paper/2102.08369/) (2021)
- [Handling Incomplete Heterogeneous Data using VAEs](https://mlsanity.com/paper/1807.03653/) (2018)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [Dropout Training as Adaptive Regularization](https://mlsanity.com/paper/1307.1493/) (2013)
- [STaSy: Score-based Tabular data Synthesis](https://mlsanity.com/paper/2210.04018/) (2022)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [CoDi: Co-evolving Contrastive Diffusion Models for Mixed-type Tabular Synthesis](https://mlsanity.com/paper/2304.12654/) (2023)
- [Diffusion models for missing value imputation in tabular data](https://mlsanity.com/paper/2210.17128/) (2022)
- [VAEM: a Deep Generative Model for Heterogeneous Mixed Type Data](https://mlsanity.com/paper/2006.11941/) (2020)
- [TabDiff: a Mixed-type Diffusion Model for Tabular Data Generation](https://mlsanity.com/paper/2410.20626/) (2024)
- [not-MIWAE: Deep Generative Modelling with Missing not at Random Data](https://mlsanity.com/paper/2006.12871/) (2020)
- [Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees](https://mlsanity.com/paper/2309.09968/) (2023)
- [Adversarial random forests for density estimation and generative modeling](https://mlsanity.com/paper/2205.09435/) (2022)
- [MissDiff: Training Diffusion Models on Tabular Data with Missing Values](https://mlsanity.com/paper/2307.00467/) (2023)
- [VAEs in the Presence of Missing Data](https://mlsanity.com/paper/2006.05301/) (2020)
- [Continuous Diffusion for Mixed-Type Tabular Data](https://mlsanity.com/paper/2312.10431/) (2023)
- [Why Not to Use Zero Imputation? Correcting Sparsity Bias in Training Neural Networks](https://mlsanity.com/paper/1906.00150/) (2019)
- [Score Matching With Missing Data](https://mlsanity.com/paper/2506.00557/) (2025)

## Built on by (most cited first, 1 of 1)

- [Impute-EM: Native Mixed-State Diffusion Models for Heterogeneous Data Imputation](https://mlsanity.com/paper/2609.15284/) (2026)

## Abstract (arXiv)

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.