📊 statistics

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

The paper introduces AugMask, a plug-and-play training framework that enables standard score-based diffusion models to effectively generate incomplete tabular data by using stochastic augmentation to fill missing values as uncertain conditioning context while restricting denoising supervision to only observed coordinates.

Original authors: Jungkyu Kim, Taeyoung Park, Kibok Lee

Published 2026-06-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jungkyu Kim, Taeyoung Park, Kibok Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: Teaching a Chef with a Broken Recipe

Imagine you are trying to teach a master chef (a Diffusion Model) how to cook a specific dish (generate realistic tabular data). The chef is brilliant at following recipes, but there's a catch: the recipes you give them are incomplete. Some ingredients are missing, and the spots where they should be are just blank white spaces or marked with "NaN" (Not a Number).

Standard chefs (AI models) get confused by blank spaces. They can't cook if the recipe says "Add 2 cups of [BLANK]."

The Old Way (Zero-Filling): People used to just fill the blanks with a placeholder, like "0" or "Water." But this is like telling the chef, "Add 0 cups of sugar." The chef learns that "0" is a real ingredient, which messes up the flavor (the data distribution).
The "Guessing" Way (Imputation): Another method tries to have the chef guess the missing ingredient, write it down, and then train the chef to make that guess perfect. But if the chef guesses wrong, they keep trying to perfect a wrong guess, leading to a noisy, confused kitchen.

The Solution: AugMask

The authors propose AugMask, a new training strategy that acts like a smart sous-chef. It solves the problem by separating two jobs: Setting the Scene and Grading the Work.

1. Setting the Scene (Stochastic Augmentation)

Instead of leaving the blank spots empty or filling them with a boring "0," the sous-chef (a lightweight helper model) fills the blanks with plausible guesses.

The Twist: The sous-chef doesn't just give one guess. It gives a range of guesses.
- Analogy: If the recipe is missing "Salt," the sous-chef doesn't just say "1 teaspoon." It says, "It could be anywhere between 0.5 and 1.5 teaspoons."
Why? This teaches the main chef that the missing ingredient is uncertain. The main chef learns to look at the whole picture (the other ingredients) and understand that the missing piece is a "maybe," not a "definite."

2. Grading the Work (Masking the Loss)

This is the most important part. When the main chef tries to recreate the dish, the teacher (the training algorithm) only grades the ingredients that were originally there.

The Rule: If the recipe originally had "2 cups of flour," the teacher checks if the chef got the flour right.
The Loophole: If the recipe originally had a blank for "Salt," the teacher ignores the chef's guess for the salt. The teacher says, "I don't care what you guessed for the salt; just make sure the flour is right."
The Result: The chef learns to use the "guesses" as context (clues) to help them get the real ingredients right, but they aren't punished for getting the guesses wrong. They stop trying to memorize the guesses and start learning the relationships between the real data.

Why This Works: The "Uncertainty Penalty"

The paper uses some fancy math to explain why this works, which can be understood through a simple metaphor: The Shaky Hand.

Deterministic Guessing (Bad): If the sous-chef gives a single, confident guess (e.g., "It is exactly 1.0 teaspoon"), the main chef thinks, "Okay, I must match this exactly." If the guess is slightly wrong, the chef gets confused.
Stochastic Guessing (Good): If the sous-chef says, "It's between 0.5 and 1.5," the main chef realizes, "This is a shaky, uncertain clue."
The Penalty: The math shows that when the clue is shaky (high uncertainty), the main chef naturally stops relying on it too heavily. It's like a student ignoring a blurry hint on a test because they know it might be wrong. This prevents the model from getting "stuck" on bad guesses.

The Results: A Better Chef

The authors tested AugMask on many different datasets (like adult income data, shopping habits, etc.) with varying amounts of missing information (from 10% missing to 90% missing).

The Outcome: By using this "Set the Scene, Ignore the Guesses" strategy, standard AI models (which usually can't handle missing data) became better than specialized models designed specifically for missing data.
The Takeaway: You don't need to build a special, complex kitchen for incomplete recipes. You just need a smart way to fill the blanks with "uncertain guesses" and then tell the chef to ignore those guesses when grading the final dish.

Summary in One Sentence

AugMask teaches AI models to handle missing data by filling the gaps with "uncertain guesses" to provide context, but then strictly grading the model only on the data that was actually there, preventing it from getting confused by its own guesses.

Technical Summary: AugMask

Problem Statement

Score-based diffusion models have emerged as powerful deep generative models but face significant challenges when applied to tabular data containing missing values. Standard diffusion architectures assume fully specified inputs, whereas real-world tabular data often contain missing entries (NaNs).

Existing approaches to bridge this gap typically fall into two categories, both of which have limitations:

Constant Placeholder Imputation: Replacing missing values with fixed constants (e.g., zeros, column means) before training. This forces the model to treat missing entries identically to valid zeros, inducing sparsity bias and making predictions dependent on the missingness level. Furthermore, even if the loss is masked to observed coordinates, the deterministic placeholder remains in the input, acting as a spurious cue correlated with missingness.
Self-Imputation (e.g., DiffPuter): Iteratively updating missing entries during training via an Expectation-Maximization (EM) approach. While this avoids fixed placeholders, it creates a "moving target" optimization problem where the model is trained on its own evolving predictions. This can lead to noisy gradients, non-stationary dynamics, and overly dispersed completions, as diffusion objectives prioritize diversity over point-wise accuracy.

The core challenge is to provide stable numeric inputs for the denoiser without treating filled-in missing values as ground-truth targets, thereby avoiding both the bias of constant placeholders and the instability of self-imputation.

Methodology: AugMask

The authors propose AugMask, a plug-and-play training framework that adapts missing-unaware backbones to incomplete data by separating conditioning from supervision. The method operates in two distinct steps:

1. Stochastic Conditional Augmentation (Input Construction)

Instead of using a deterministic placeholder, AugMask constructs a fully numeric input by sampling missing entries from a context-dependent distribution $T(\cdot | x_{obs}, m)$ .

Mechanism: For each missing feature $j$ , a value $Z_j$ is drawn from a distribution conditioned on the observed features and the missingness mask.
Implementation: The authors utilize a lightweight, factorized approach using per-feature LightGBM models. For continuous features, they fit a conditional mean regressor and a log-variance regressor to sample from a heteroskedastic Gaussian. For categorical features, they use a conditional classifier with label smoothing.
Role: These sampled values serve as uncertain conditioning context for the denoiser, not as ground-truth targets. The stochasticity encodes the inherent uncertainty of the missingness.

2. Observed-Only Supervision (Loss Masking)

The denoising loss is applied only to the coordinates that were originally observed.

Objective: The model minimizes the reconstruction error for observed coordinates $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$ , where $x_{A,0}$ is the augmented input.
Effect: The augmented (previously missing) entries enter the model through the noisy input $x_{A,t}$ to provide context, but the model is never penalized for failing to reconstruct them. This prevents the model from over-relying on potentially inaccurate completions.

Theoretical Justification

The paper connects this training rule to a Rao–Blackwellized objective. By marginalizing over plausible completions, the authors show that the objective yields a variance-weighted sensitivity penalty.

Interpretation: The second-order term in the local approximation of the objective is proportional to the conditional variance of the augmentation. High conditional variance penalizes the model's sensitivity to that specific completion, effectively discouraging reliance on uncertain fills. Low variance allows the completion to be used more directly as context.
Contrast with AugFull: The authors compare AugMask to "AugFull" (which applies loss to all coordinates, including augmented ones). They demonstrate that AugFull introduces bias if the augmentation rule is misspecified, as the model is forced to fit potentially incorrect pseudo-targets. AugMask avoids this by treating completions strictly as context.

Key Contributions

A Training Principle for Incomplete Tabular Diffusion: The introduction of AugMask, which decouples conditioning (via stochastic augmentation) from supervision (via observed-only loss masking).
A Marginalization View: A theoretical analysis showing that stochastic augmentation induces a variance-weighted sensitivity penalty, regularizing the model against over-reliance on uncertain completions.
A Plug-and-Play Recipe: A lightweight framework that enables standard, missing-unaware score-based diffusion models to train on incomplete data without architectural changes.

Experimental Results

The authors evaluated AugMask across diverse datasets (Adult, Default, Shoppers, Magic, Beijing, News) and missingness regimes (MCAR and MAR) with missing ratios ranging from 0.1 to 0.9.

Performance: AugMask enables standard missing-unaware diffusion models (e.g., CDTD, TabDiff) to outperform specialized missing-aware baselines (e.g., MIWAE, MissDiff, DiffPuter, ForestDiff).
Generative Quality: Models trained with AugMask achieved superior sample-level fidelity (measured by $\alpha$ -Precision, $\beta$ -Recall) and feature-level statistics (Trend, Shape).
Machine Learning Utility: Downstream tasks (classification and regression) trained on synthetic data generated by AugMask models showed improved performance compared to baselines.
Robustness: The benefits of AugMask became more pronounced as the missingness ratio increased. At high missingness (e.g., 90%), AugMask avoided the severe degradation observed in other methods.
Sensitivity Analysis: The method proved robust to augmentation misspecification. Even when the augmentation rule was perturbed (e.g., reducing stochasticity or introducing bias), AugMask maintained a performance advantage over AugFull, confirming that treating completions as context rather than targets is the critical factor.

Significance and Claims

The paper claims that AugMask provides a principled solution to the mismatch between fully-specified diffusion backbones and incomplete real-world data. By treating stochastic completions as uncertain context rather than training targets, the method:

Avoids the sparsity bias and artificial patterns introduced by constant placeholders.
Avoids the optimization instability and moving-target dynamics of self-imputation methods.
Leverages the variance of the augmentation to regularize the model, making it less sensitive to uncertain inputs.

The authors position AugMask as a general strategy that allows existing, high-performing missing-unaware models to be deployed in settings with missing data without requiring complex architectural modifications or specialized loss functions for imputation. The work focuses on generation; the authors note that extending this to latent diffusion pipelines or imputation-specific objectives remains a direction for future work.