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

TL;DR

The paper tackles severe class imbalance in tabular data by proposing latent-space, tree-driven diffusion methods that learn a diffusion vector field with gradient-boosted trees under conditional flow matching. It introduces three variants—PCAForest, EmbedForest, and AttentionForest—that operate in compact latent spaces to improve minority recall while preserving tabular structure, with an end-to-end pipeline that encodes minority samples, learns diffusion, and decodes synthetic data. Across 11 real-world datasets, AttentionForest achieves the strongest average minority recall and competitive precision and calibration, while PCAForest and EmbedForest offer favorable accuracy-efficiency trade-offs; all variants demonstrate privacy-conscious behavior via NNDR and DCR and maintain reasonable distributional similarity via Wasserstein distance. Ablation studies show smaller latent embeddings can boost recall and that overly aggressive learning rates harm stability, underscoring the need for dataset-specific tuning. Overall, latent-space diffusion with tree-based vector fields provides a practical, privacy-aware approach to high-fidelity tabular augmentation under severe class imbalance, suitable for deployment in high-stakes domains like healthcare, finance, and industrial analytics.

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.

Figures (8)

• Figure 1: Synthetic Tabular Data Generation using Forest-Diffusion
• Figure 2: Flowchart for Latent Space Forest-Diffusion
• Figure 3: Our Latent Space Diffusion Models outperform Baseline Models in Downstream Machine Learning Tasks
• Figure 4: Recall Score Comparison at 100% Augmentation Ratio Across Datasets
• Figure 5: Recall Score Trends Across Augmentation Ratios