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

Abstract: 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.

• The paper introduces AugMask, a plug-and-play framework that integrates stochastic augmentation with loss masking to train diffusion models on incomplete tabular data.
• It leverages LightGBM-based conditional sampling to generate context-aware imputations while restricting loss computation to observed entries, reducing overfitting to imputation artifacts.
• Experimental evaluations across multiple datasets show that AugMask outperforms traditional imputation methods and state-of-the-art baselines in generative fidelity and downstream utility.

AugMask: Score-Based Diffusion Training via Stochastic Augmentation and Masking for Incomplete Tabular Data

Introduction and Motivation

Score-based diffusion models have emerged as state-of-the-art generative models for complex data distributions, but their extension to tabular domains lags due to the challenge posed by missing data—a pervasive issue in real-world tabular datasets. Standard diffusion backbones require fully observed inputs, leading to a gulf between theoretical model design and practical applicability, as most tabular sources are incomplete. Traditional remedies, such as constant placeholder fills (e.g., zero/mean imputation), only partially address the problem and can introduce information leakage, sparsity bias, or deterministic artifacts correlated with the missingness pattern.

The paper "AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking" (2606.03347) introduces a minimally invasive (plug-and-play) training strategy for adapting missing-unaware score-based diffusion models to incomplete data. AugMask leverages stochastic, context-aware augmentation to construct plausible numeric inputs at missing positions and couples this with a loss masking protocol that restricts denoising supervision to truly observed entries. This separation between conditioning and supervision prevents the model from overfitting to imputation artifacts and mitigates sensitivity to mis-specified completions.

Methodology

The AugMask framework consists of two decoupled components: conditional stochastic augmentation for input construction and masking of the denoising loss for targets.

1. Conditional Stochastic Augmentation: For each missing cell, rather than filling with a deterministic statistic (mean/mode or a constant), a conditional sampler $$T(\cdot \mid \mathbf{x}^{\mathrm{obs}}, \mathbf{m})$$ produces context-aware samples. In practice, this is implemented via feature-wise conditional LightGBM regressors/classifiers that estimate both conditional mean and variance for numerical features, or softmax-probabilities for categorical features, allowing the augmented value to reflect underlying contextual uncertainty.
2. Masked Supervision: The denoising score-matching loss is restricted to observed positions. While the augmented values provide numeric context to the model, they are omitted from the loss computation, avoiding leakage of spurious correlations as learning targets.

   Figure 2: Comparison of augmentation rules and masking. AugMask (green) delivers consistent improvements over AugFull (blue), across deterministic and stochastic augmenters.

This design is justified theoretically: marginalizing the loss over the augmentation distribution produces a sensitivity regularization term weighted by the estimated conditional variance. The model is thus trained to be robust to uncertainty in missing-value imputations and only leverages them as informative context, not as pseudo-ground-truth.

Theoretical Analysis: Marginalization as Variance-Weighted Regularization

Consider the setting with a single possibly-missing coordinate $j$ and observed block $-j$. The expected reconstruction loss for the observed block, marginalized over the possible fills for $j$, is given by:

$$l^{\mathrm{RB}}(\theta; \mathbf{x}^0_{-j}) \approx g_\theta(\mu_j) + \frac{1}{2}\sigma_j^2 g_\theta''(\mu_j)$$

where $\mu_j$ and $\sigma_j^2$ are the conditional mean and variance under the true data distribution. The second term penalizes sensitivity of the denoiser to changes in the missing-value completion, scaled up as the conditional uncertainty increases.

In high dimensions, the diagonal elements of the conditional covariance control the leading regularization, with off-diagonal corrections bounded by the co-missingness probability and feature-feature coupling in the denoiser's Jacobian.

Practical Implementation

Algorithmically, AugMask constructs a single stochastic, context-augmented version of the training data (sampled once per cell per epoch for computational efficiency and stationary conditioning), using per-feature LightGBM conditional models. For categorical variables, label smoothing is added for stability. This makes conditioning highly scalable while capturing both first- and second-moment structure in missingness, thus operationalizing the variance-weighted penalty predicted by the local theoretical expansion.

The strategy is plug-compatible: existing diffusion backbones for tabular data (TabDDPM, CDTD, TabDiff) can be directly adapted without architectural modifications.

Experimental Evaluation

Protocols

Three evaluation settings are employed:

1. TRTR: Train and test on complete real data (upper bound).
2. TSTR: Train on synthetic, test on real, measuring ML task utility.
3. TMTC: Train generator on incomplete data, synthesize fully observed data, evaluate against complete real data.

   Figure 1: Evaluation protocols for missing-aware generators, situating TMTC as the realistic regime for training on incomplete observations.

Quantitative Results

Across 90 experimental scenarios (six datasets, varying missing ratios, multiple seeds), the following are demonstrated:

• AugMask enables missing-unaware diffusion models to outperform missing-aware baselines (including ForestDiff, MIWAE, DiffPuter, and MissDiff) in both generative fidelity (measured by $\alpha$-Precision, $\beta$-Recall, correlation Trend, marginal Shape) and downstream utility (normalized AUROC/RMSE).
• Stochastic augmentation alone (AugFull) yields large gains over deterministic or imputation-based fills, but further restricting loss supervision (AugMask) is necessary to achieve state-of-the-art coverage and utility, particularly at high missingness rates.

  Figure 3: Robustness to missingness: for CDTD, a clear hierarchy emerges; AugMask consistently matches or surpasses top missing-aware baselines across increasing MCAR rates in both sample quality and ML utility.

  

Figure 4: Robustness across missing ratios (MCAR) for TabDiff reveals AugMask's advantages vis-à-vis all missing-aware baselines.

Figure 7: Robustness across missing ratios (MAR) for TabDiff extends AugMask's superior performance to missing-at-random mechanisms.

Notable claims supported by results:

• AugMask achieves best average ranks across all metrics and scenarios, with statistical significance exceeded over all prior baselines.
• Gains from stochastic augmentation and masked supervision are evident at all missingness levels, intensifying under strong missingness ($p_{\mathrm{miss}}>0.7$).

Sensitivity to Augmentation Misspecification

A detailed ablation perturbs the augmentation model both in conditional mean ($j$0; context bias) and variance ($j$1; stochasticity). AugMask retains performance even with substantial augmentation misspecification, validating the claim that its learning is less brittle due to the exclusion of augmented values from the loss.

Figure 8: Sensitivity to augmentation misspecification—AugMask remains robust to perturbations in both context and variance, outperforming AugFull across perturbation regimes.

Qualitative Analysis

Marginal and joint density visualizations (not shown here for brevity) confirm that stochastic augmentation better preserves distributional variability and structure, avoiding shrinkage and mode collapse caused by deterministic imputation.

Figure 9: UMAP visualization under heavy missingness demonstrates AugMask's restored distributional coverage (high density overlap IoU), while standard baselines suffer from severe support mismatch.

Figure 10: Preprocessing pipeline: scale-then-impute (right), the procedure advocated by AugMask, avoids distributional skewing induced by impute-then-scale stacking (left).

Implications and Future Directions

Practically, AugMask makes modern score-based diffusion generators viable for privacy-sensitive or data-limited tabular domains where missingness is the norm, facilitating the production of high-fidelity, utility-preserving synthetic datasets for ML benchmarking, simulation, or data sharing. Theoretically, AugMask formalizes a connection between marginalization under missingness and Jacobian/sensitivity regularization, suggesting deeper links between data uncertainty, sample augmentation, and model robustness in generative frameworks.

Key implications and future research avenues include:

• Extensions to latent diffusion, VAEs, and generative adversarial objectives: scaling observed-only supervision into end-to-end pipelines, especially with discrete or hybrid latent spaces.
• Generalization to block-missing, MNAR, and non-i.i.d. missingness: moving beyond factorized augmentation.
• Privacy and fairness audits for synthetic data: while AugMask limits over-reliance on observed contingency artifacts, generated data remains subject to possible memorization and bias inheritance.

Conclusion

This work provides a rigorous, theoretically-justified, and empirically validated approach for training missing-unaware diffusion models on incomplete tabular data. AugMask's combination of plug-and-play stochastic augmentation and masked supervision delivers consistently superior generative and downstream performance across broad scenarios, while establishing a new conceptual link between marginalization, uncertainty quantification, and model robustness (2606.03347).