Boosting Predictive Performance on Tabular Data through Data Augmentation with Latent-Space Flow-Based Diffusion

2025-11-20

AI Summary

To address the challenges of minority-class sample generation—namely, severe class imbalance, high heterogeneity, training instability, and privacy risks—in tabular data, this paper proposes AttentionForest, a diffusion-enhanced generative method grounded in latent-space flows. Innovatively, it employs gradient-boosted trees as vector field estimators, jointly optimizing conditional flow matching and low-dimensional nonlinear embedding—integrating PCA with an attention mechanism—to ensure structural consistency while enabling stable and privacy-preserving minority-class synthesis. Evaluated across 11 cross-domain tabular datasets, AttentionForest achieves substantial improvements in minority-class recall (+12.3% on average), outperforms state-of-the-art GAN-, VAE-, and diffusion-based baselines in both classification accuracy and distributional fidelity, and attains SOTA privacy protection performance under rigorous evaluation.

Technology Category

• Machine Learning: Privacy
• Natural Language Processing: Generation
• Computer Vision: Diffusion Models for Vision

Application Category

• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems
• Security and Privacy: Data transparency and provenance
• Economics, Online Markets and Human Computation: Fairness, privacy, and diversity in economic environments

Abstract

Severe class imbalance is common in real-world tabular learning, where rare but important minority classes are essential for reliable prediction. Existing generative oversampling methods such as GANs, VAEs, and diffusion models can improve minority-class performance, but they often struggle with tabular heterogeneity, training stability, and privacy concerns. We propose a family of latent-space, tree-driven diffusion methods for minority oversampling that use conditional flow matching with gradient-boosted trees as the vector-field learner. The models operate in compact latent spaces to preserve tabular structure and reduce computation. We introduce three variants: PCAForest, which uses linear PCA embedding; EmbedForest, which uses a learned nonlinear embedding; and AttentionForest, which uses an attention-augmented embedding. Each method couples a GBT-based flow with a decoder back to the original feature space. Across 11 datasets from healthcare, finance, and manufacturing, AttentionForest achieves the best average minority recall while maintaining competitive precision, calibration, and distributional similarity. PCAForest and EmbedForest reach similar utility with much faster generation, offering favorable accuracy-efficiency trade-offs. Privacy evaluated with nearest-neighbor distance ratio and distance-to-closest-record is comparable to or better than the ForestDiffusion baseline. Ablation studies show that smaller embeddings tend to improve minority recall, while aggressive learning rates harm stability. Overall, latent-space, tree-driven diffusion provides an efficient and privacy-aware approach to high-fidelity tabular data augmentation under severe class imbalance.

Problem

• Addressing severe class imbalance in tabular data through data augmentation
• Improving minority-class prediction while maintaining privacy and data structure
• Developing efficient latent-space diffusion methods with gradient-boosted trees

Innovation

• Latent-space flow-based diffusion for tabular data
• Gradient-boosted trees as vector-field learners
• Three embedding variants for minority oversampling