Active Tabular Augmentation via Policy-Guided Diffusion Inpainting
Abstract
Generative tabular augmentation is appealing in data-scarce domains, yet the prevailing focus on distributional fidelity does not reliably translate into better downstream models. We formalize a fidelity-utility gap: common generative objectives prioritize distributional plausibility, whereas augmentation succeeds only when injected samples reduce the current learner's held-out evaluation loss. This gap motivates learning not just how to generate, but what to generate and when to inject as training evolves. We propose TAP (Tabular Augmentation Policy), which couples diffusion inpainting with a lightweight, learner-conditioned policy to steer generation toward high-utility regions and controls safe injection via explicit gating and conservative windowed commitment. Under severe data scarcity, TAP consistently outperforms strong generative baselines on seven real-world datasets, improving classification accuracy by up to 15.6 percentage points and reducing regression RMSE by up to 32%.
Lay Summary
In areas like healthcare and finance, researchers often face the same problem: there isn't enough labeled data to train a reliable AI model. A common workaround is to use AI to generate extra synthetic records that look like the real ones. But realistic-looking data doesn't always help: modern generators tend to produce examples that resemble what the model already knows, while the tricky borderline cases that would actually teach it something new get overlooked. We built TAP, which treats data generation as a guided process rather than bulk production. At each step, TAP identifies what the current model is struggling with and asks the generator to produce examples that would help it most. Each candidate is screened for validity, and a batch is only added to the training set when it passes a strict usefulness check. Across seven real-world datasets, TAP consistently outperforms existing methods, especially when very little labeled data is available. It improves classification accuracy by up to 15.6 percentage points and reduces prediction error by up to 32%. This makes it useful for high-stakes settings like rare-disease research or small clinical studies, where every labeled example is precious.