Source-linked AI summary

Active Tabular Augmentation via Policy-Guided Diffusion Inpainting

Zheyu Zhang, Shuo Yang, Bardh Prenkaj, Gjergji Kasneci

arXiv:2605.10315v1cs.LGcs.AI

TL;DR

Tabular augmentation faces a fidelity-utility gap: plausible synthetic records do not necessarily reduce a learner’s held-out loss. TAP combines diffusion inpainting with a learner-conditioned policy, explicit gating, and conservative commitment, and consistently outperforms strong generative baselines across scarce-data settings.

  • Problem

    Distributional fidelity does not reliably indicate whether synthetic samples will reduce downstream loss on real queries, creating a fidelity-utility gap in scarce-data augmentation.

  • Method

    TAP couples diffusion inpainting with a learner-conditioned policy that steers generation and uses hard feasibility gates plus windowed commitment for injection.

  • Results

    TAP consistently outperforms strong generative baselines across seven real-world datasets and five scarcity levels, improving accuracy by up to 15.6 percentage points and reducing RMSE by up to 32%.

  • Takeaways & Limitations

    Effective tabular augmentation requires selecting samples that help the evolving learner rather than relying on distributional fidelity alone.

  • Takeaways & Limitations

    Sequential injection introduces additional overhead compared with one-shot generation, and augmentation should be validated carefully in high-stakes settings.

Abstract

from arXiv · show

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%.

1. Introduction

Tabular augmentation must produce valid records that improve downstream learners, not merely plausible records matching the observed distribution. TAP addresses this fidelity-utility gap by steering diffusion-based proposals with learner state and conservative injection decisions.

  • Tabular augmentation is fragile because heterogeneous features and strong inter-column dependencies make minor perturbations capable of invalidating samples or introducing spurious relationships.
  • High-fidelity generators can sample redundant high-density records that offer limited downstream gain, while utility depends on task-specific relevance.
  • Under scarcity, impactful samples often occur near uncertain decision boundaries or under-covered subpopulations, but uncertainty alone is insufficient.
  • TAP treats generation as a controllable proposal process whose conditions adapt to the learner as training evolves.
  • TAP combines diffusion inpainting, explicit quality gates, and a learner-conditioned policy to improve utility while enforcing feasibility.
  • TAP improves accuracy by up to 15.6 percentage points and reduces RMSE by up to 32% across seven real-world datasets and five scarcity levels.

2. Background & Motivation

The fidelity-utility gap arises because distributional plausibility does not determine whether injected samples reduce loss on real queries. The paper formalizes augmentation as sequential control and motivates feasibility checks, utility-driven selection, and conservative commitment.

  • The central question is what to inject to maximize downstream benefit rather than how to generate more realistic samples.
  • Augmentation is evaluated through loss on real queries, while the generator and injection rule jointly induce an evolving distribution over proposed and injected synthetic samples.
  • Fidelity measures plausibility under P, whereas utility depends on whether injected samples complement D and reduce loss on Qreal.
  • Directly optimizing marginal utility is intractable because it would require retraining the learner for every candidate set, so TAP learns a steering policy instead.
  • Influence-function analysis is used diagnostically to connect utility with current learner errors and motivate feasibility, gating, and diversity-aware policy design.
  • TAP follows two-stage feasibility: diffusion encourages manifold proximity, while explicit enforcement handles domain rules such as ranges and logical consistency.
  • Utility-driven selection targets samples that are informative yet learnable because boundary-adjacent or merely uncertain samples can degrade performance.
  • Conservative sequential injection accumulates admitted samples in a window and commits them only when pooled gain exceeds a minimum threshold.

3. Methodology

TAP formulates tabular augmentation as sequential control under an evolving learner, combining diffusion inpainting, utility-aligned policy selection, feasibility gating, and conservative commitment. The objective is to reduce real-query loss through safe injection rather than simply match the data distribution.

  • Sequential augmentation: TAP treats effective augmentation as sequential control because the learner changes after each commitment and the utility of a generation choice evolves.The policy maps learner states to action distributions, while utility is evaluated across the resulting trajectory.
  • Manifold-constrained proposals: Diffusion inpainting generates locally coherent candidates by conditioning on real anchors, fixing selected columns, and regenerating the remainder.The label is held fixed during inpainting, and fixed coordinates are overwritten throughout reverse diffusion without retraining the backbone.
  • Manifold-constrained proposals: TAP parameterizes proposals with target condition c, mask template η, and exploration strength ρ, which jointly control target selection, locality, and diversity.Larger ρ yields more diverse proposals, whereas smaller ρ produces near-anchor samples; each action induces a proposal distribution through anchor selection, masking, and diffusion randomness.
  • Utility-aligned selection: The policy targets informative yet learnable regions using a plug-in evaluator and learner-state features covering under-covered targets, uncertainty, gate statistics, and redundancy.The plug-in evaluator enables fast repeated forward passes, ranks candidate pools during policy learning, and does not replace retraining standard downstream predictors for reported gains.
  • Safe admission and commitment: Explicit feasibility gates reject invalid candidates, while admitted samples accumulate in a window and affect the learner only when the pooled commitment rule is satisfied.Windowed commitment addresses noisy utility estimates; calibrated plug-in uncertainty supports a certified decision rule for commitment safety.
  • Safe admission and commitment: The commitment theorem requires an error bar ϵt and accepts a pool only when its estimated utility exceeds the threshold τ plus ϵt.The rule is presented as a safety guarantee contingent on calibrated plug-in uncertainty.

4. Experiments

Experiments evaluate TAP under simulated data scarcity, compare it with diverse augmentation baselines, and analyze where utility arises and how safe injection mechanisms contribute. TAP generally improves downstream performance, with gains linked to informative yet learnable samples and conservative commitment.

  • Setup: TAP is evaluated under scarcity levels nreal ∈{20, 50, 100, 200, 500}, using five random splits and a 4:1 train-validation split.The evaluator uses only the real training split during injection and never accesses validation or test labels.
  • Setup: The benchmark compares TAP with seven augmentation methods, including SMOTE, TVAE, CTGAN, ARF, SPADA, TabDDPM, and TabDiff.All methods inject the same budget of nsyn = 500 synthetic samples, with TabDiff fixed as TAP’s diffusion backbone.
  • Utility across scarcity levels: TAP achieves the best or near-best average classification accuracy and regression RMSE across scarcity levels, with its largest gains at nreal=20.Each result averages across six downstream predictors, and the smallest-data regime is especially sensitive to inconsistent injected records.
  • Utility across scarcity levels: Several generators underperform Real in some settings, whereas TAP produces consistently positive improvements across datasets and scarcity levels.The results support optimizing a surrogate of downstream utility rather than relying on distributional fidelity alone.
  • Where high-utility samples lie: Utility improves from global sampling through TAP when injection methods share the same diffusion backbone, with TAP outperforming the strong hard-inpainting reference.TAP additionally learns state-conditioned targets and templates and commits only when pooled utility is reliably positive.
  • Where high-utility samples lie: At matched informativeness, TAP achieves better learnability, supporting the characterization of useful samples as informative yet learnable.Informativeness is measured by proximity to the decision boundary, while learnability is measured through label consistency.
  • Where high-utility samples lie: Utility concentrates in middle learnability bins, while the least learnable tail produces negative utility.This interventional result indicates that boundary proximity alone is insufficient and that anchored inpainting reduces severely inconsistent samples relative to unanchored draws.
  • Conservative commitment: Removing gating sharply increases tail risk and reduces win rate, while removing windowed commitment causes a smaller but consistent negative effect.Gating filters individually invalid candidates, whereas windowed commitment rejects pools whose joint gain is uncertain.

5. Related Work

Related work spans generative models for heterogeneous tabular data and methods for tabular augmentation and sample selection. These approaches commonly emphasize distributional modeling, while tabular constraints and downstream utility remain central concerns.

  • Generative modeling for tabular data: Tabular generative methods include VAEs, GANs, tree-based estimators, normalizing flows, diffusion models, and language-model-based approaches.Examples include TVAE, CTGAN, ARF, SPADA, TabDDPM, Tab-Syn, and TabDiff.
  • Generative modeling for tabular data: These generative approaches optimize distributional fidelity, but high fidelity does not guarantee downstream utility under scarcity.Passive sampling can also produce redundant samples in regions that are already covered.
  • Data augmentation and sample selection: Tabular augmentation methods must accommodate heterogeneous features and tight inter-column constraints.The literature includes SMOTE, adapted Mixup variants, TabEBM, and language-model-based augmentation approaches.

6. Conclusion

TAP addresses the fidelity–utility gap by treating tabular augmentation as sequential control of a proposal kernel. Across diverse tasks and scarcity regimes, it improves utility while limiting harmful injection, though sequential operation adds computational overhead.

  • Conclusion: TAP couples diffusion inpainting with a learner-state-conditioned policy that controls what to generate and when to inject.Hard gating and windowed commitment improve reliability under noisy utility estimates.
  • Conclusion: TAP consistently outperforms strong generative baselines across diverse tasks and scarcity regimes.The framework improves accuracy by up to 15.6 percentage points and reduces RMSE by up to 32%.
  • Conclusion: Sequential injection introduces additional overhead compared with one-shot generation.The paper mitigates this cost through window-level caching and training-free evaluators, while identifying more efficient utility estimation as future work.

Impact Statement

TAP is presented as a safer approach to low-data tabular augmentation, combining manifold-local generation with feasibility checks and conservative commitment. The supporting discussion frames these safeguards as important because synthetic samples can harm downstream performance and because high-stakes data may reproduce existing biases.

  • Impact Statement: TAP targets reliable augmentation in low-data tabular settings where each injected record can have outsized influence.The paper motivates this setting by citing cost, privacy, and domain constraints, while warning that harmful synthetic samples can degrade downstream performance.
  • Impact Statement: TAP combines manifold-local diffusion inpainting, hard feasibility gates, and conservative windowed commitment to reduce harmful injections under scarcity.The design is intended to inject samples only when utility is consistently indicated.
  • Impact Statement: High-stakes deployments require careful application because augmentation can propagate measurement noise, institutional practices, and historical inequities present in training data.The paper recommends validation on real held-out data, subgroup metrics when available, and documentation of intended use and known limitations.
  • Impact Statement: TAP frames augmentation as sequential control with a learner-conditioned policy, frozen diffusion inpainting, feasibility gating, temporary pooling, utility evaluation, and windowed commitment.The workflow uses evaluator feedback as advantage signals for preference-based policy optimization.
  • Impact Statement: The supporting theory provides safety guarantees and design intuition rather than a global optimality guarantee.The appendix discusses utility telescoping, plug-in utility error, and a sufficient condition for state-summary design.

B.5. Derivation of Equation (3)

This section develops a local, first-order lens for how injected examples affect utility and uses it to motivate TAP’s action-ranking diagnostics. Its explanatory claims are restricted by local approximations and a stated monotonicity assumption within matched informativeness buckets.

  • B.5. Derivation of Equation (3): The marginal-utility approximation is local and most accurate when adding one example induces only a small parameter shift.The paper notes that the approximation can be numerically crude in severe-scarcity regimes, so empirical utility estimation and conservative commitment remain important.
  • B.5. Derivation of Equation (3): The derivation assumes a twice-differentiable, locally strongly convex training objective and drops higher-order terms.A first-order Taylor expansion of the training gradient approximates the parameter shift after adding one example, which is then used to approximate marginal utility.
  • B.5. Derivation of Equation (3): The resulting marginal-utility approximation follows by applying a first-order Taylor expansion to the evaluation loss and substituting the parameter-shift expression.The proposition states that this approximation matches Equation (3).
  • B.5. Derivation of Equation (3): TAP’s diagnostics compare learnability within matched informativeness slices defined by decision-boundary scores.Inconsistency is used post hoc as a learnability proxy, with smaller inconsistency indicating higher learnability.
  • B.5. Derivation of Equation (3): Under the utility-monotonicity assumption, utility maximization within a fixed informativeness bucket also minimizes average inconsistency.This establishes Pareto efficiency only within the restricted family of candidate pools in that bucket.

C. Additional Method Details

TAP represents generation choices with a compact learner-state summary and a factorized action, then uses feasibility gates, pooled utility checks, and preference-based updates to control proposals and injection. Its policy is anchored to a conservative reference and learns from thresholded binary feedback.

  • State and action design: The state tracks target deficit, evaluator uncertainty, gate pass rates, and diversity relative to committed data and the current pool.These quantities summarize coverage, uncertainty, feasibility, and redundancy before each action is sampled.
  • State and action design: TAP’s action selects a target condition, mask template, and exploration strength, inducing a proposal distribution through anchor selection, masking, and diffusion randomness.The policy factorizes over the discrete target and template choices and the continuous exploration parameter.
  • Proposal construction: Anchor selection restricts records to the chosen target condition and mixes hard-sample preference with uniform sampling.For classification the condition is a class; for regression it is a quantile bin, with hardness based on evaluator uncertainty or error.
  • Proposal construction: Lower exploration strength produces more conservative near-anchor proposals by fixing additional numeric features during mask construction.The explore template regenerates all non-label features, whereas the conservative template also fixes important columns.
  • Safety and commitment: Pointwise gates reject candidates violating type, range, or task-dependent logical constraints before they enter the pool.Numeric values are clipped to the real-training quantile range [qmin, qmax], with qmin = 0.01 and qmax = 0.99 in the experiments.
  • Safety and commitment: A pool is committed only when its estimated utility exceeds the threshold plus an error bar, reducing tail risk from noisy estimates and capturing candidate complementarities.The commitment rule evaluates pooled rather than only per-step utility.
  • Policy learning: Preference-based policy updates convert baseline-corrected utility advantages into positive, negative, or abstaining feedback, avoiding direct regression on potentially heavy-tailed estimates.The KTO-style objective increases policy probability for desirable actions and decreases it for undesirable actions relative to a conservative reference policy.

D.1. Datasets

The experiments use seven real-world tabular datasets under repeated scarcity simulations, with frozen diffusion generation and cross-validated utility estimation. TAP is compared with established augmentation and generative baselines under shared budgets and splits, while additional analyses examine mechanisms and robustness.

  • D.1. Datasets: Seven real-world datasets span healthcare, finance, science, and operations, with six sourced from OpenML and Insurance sourced from Kaggle.Dataset statistics, task types, and feature compositions are reported in Table 3.
  • D.1. Datasets: Scarcity levels use nreal ∈{20, 50, 100, 200, 500}, with an 80:20 real train-validation split and repeated random splits.The generator and policy are trained only on the real training split, while held-out data support evaluation procedures.
  • D.1. Datasets: TAP uses TabDiff as a frozen mixed-type diffusion backbone and performs label-conditioned inpainting by fixing selected feature coordinates.The backbone is trained on the real training split and remains frozen during policy learning.
  • D.1. Datasets: Plug-in utility is estimated with cross-validation and focused query subsets, using predictive entropy for classification and residual magnitude for regression.TabPFN supplies training-free iterative evaluation without retraining at every policy step.
  • D.1. Datasets: The comparison includes SMOTE, TVAE, CTGAN, ARF, SPADA, TabDDPM, and TabDiff under the same synthetic budget nsyn = 500 and data splits.Mechanism analysis also evaluates alternative injection rules under a shared diffusion backbone.
  • D.1. Datasets: TabDDPM on Ailerons is unstable under scarce training, with high variance and occasional failure cases in the reported runs.Per-predictor results indicate that the main conclusions are not driven by a single downstream model.
  • D.1. Datasets: TAP’s downstream evaluation reports Accuracy and Macro F1 for classification and RMSE and MAE for regression across six classifiers and four regressors.Aggregate tables average performance across downstream predictors and report means and standard deviations over five splits.

E.3. Plug-in Utility Calibration

The calibration study evaluates whether the plug-in utility estimate and its error bar provide a conservative basis for commitment decisions under scarcity, alongside runtime and ablation analyses of the injection pipeline.

  • Plug-in Utility Calibration: The evaluator is used to rank candidate pools and check commitments, while final results are reported after retraining standard downstream predictors on the committed augmented set.The plug-in evaluator is not used for final reporting.
  • Plug-in Utility Calibration: At commitment checks every K steps, pooled plug-in estimates and error bars are compared with retraining-based loss reductions on a held-out real validation split.The empirical proxy averages loss reduction across downstream predictors.
  • Plug-in Utility Calibration: Coverage measures whether the plug-in error falls within ϵt, while MAE measures its absolute deviation from retraining-based utility estimates.The evaluation compares plug-in estimates with realized utility approximated by retraining downstream predictors on real and augmented data.
  • Plug-in Utility Calibration: 91.8% average coverage in classification is close to the 95% target, supporting ϵt as a useful conservative bound.Regression coverage is lower at 87.4%, reflecting higher variance in squared-error loss.
  • Plug-in Utility Calibration: All datasets maintain coverage above 80%, and MAE remains well below ϵt, supporting use of the error bar as a conservative margin across task types.Insurance has the lowest reported coverage at 84.1%, consistent with its high outcome variance.
  • Computational Cost: TAP’s wall-clock runtime is in the same order of magnitude as Hard inpainting for producing and injecting 500 samples at nreal = 50.The diffusion backbone is trained once and reused across diffusion-based mechanisms; policy updates are typically negligible relative to sampling and evaluator queries.

F.2. Results and Analysis

The ablations test TAP’s state summary, learned proposal controls, and estimator choice, finding that multiple state factors and adaptive control support performance while the mechanism remains robust to evaluator replacement.

  • State Ablation: Removing any state component degrades performance, with the largest drop occurring when the diversity score is removed.The results support jointly tracking coverage, difficulty, feasibility, and redundancy when ranking actions.
  • Action Ablation: Fixed action choices underperform learned control across tasks, while pure exploration and conservative masks each create distinct reliability or coverage trade-offs.Fixed exploration strengths also have different optima for classification and regression.
  • Estimator Ablation: TabPFN provides stable few-shot utility signals under scarcity but is not required by the method.The ablations replace it with an ensemble evaluator and a validation-based holdout estimator.
  • Estimator Ablation: The policy continues to provide gains when the online estimator is replaced, indicating that the overall mechanism does not rely on a specific evaluator.Final evaluation excludes TabPFN and averages across heterogeneous downstream predictors.
  • Evaluation Scope: Additional evaluation reports Macro-F1 for classification, MAE for regression, and per-predictor results to assess robustness across model classes.These metrics complement the main accuracy and RMSE analyses.

G.1. Other Metrics

Additional metrics broadly preserve the main performance ordering, with Macro-F1 and MAE indicating that reported gains are not confined to majority classes or a few large regression errors.

  • Additional Metrics: Macro-F1 largely follows Accuracy across datasets, suggesting improvements are not restricted to majority classes.Macro-F1 is reported as an additional classification metric beyond the main tables.
  • Additional Metrics: MAE is consistent with RMSE in most settings, indicating gains are not driven only by a small number of large errors.MAE is reported as an additional regression metric.
  • Per-Predictor Robustness: Per-predictor breakdowns are provided to check whether improvements are driven by a single downstream model.The main tables average performance across downstream predictors.
  • Per-Predictor Results: The supplementary tables report classification accuracy for LR, RF, LGBM, XGB, KNN, and MLP under varying scarcity levels.These tables provide predictor-specific classification results.
  • Per-Predictor Results: Supplementary regression tables report RMSE for RF, LGBM, XGB, and KNN under varying scarcity levels.These results complement the aggregated regression evaluation.
Loading 2605.10315v1…