TabEBM: A Tabular Data Augmentation Method with Distinct Class-Specific Energy-Based Models

Abstract: Data collection is often difficult in critical fields such as medicine, physics, and chemistry. As a result, classification methods usually perform poorly with these small datasets, leading to weak predictive performance. Increasing the training set with additional synthetic data, similar to data augmentation in images, is commonly believed to improve downstream classification performance. However, current tabular generative methods that learn either the joint distribution $ p(\mathbf{x}, y) $ or the class-conditional distribution $ p(\mathbf{x} \mid y) $ often overfit on small datasets, resulting in poor-quality synthetic data, usually worsening classification performance compared to using real data alone. To solve these challenges, we introduce TabEBM, a novel class-conditional generative method using Energy-Based Models (EBMs). Unlike existing methods that use a shared model to approximate all class-conditional densities, our key innovation is to create distinct EBM generative models for each class, each modelling its class-specific data distribution individually. This approach creates robust energy landscapes, even in ambiguous class distributions. Our experiments show that TabEBM generates synthetic data with higher quality and better statistical fidelity than existing methods. When used for data augmentation, our synthetic data consistently improves the classification performance across diverse datasets of various sizes, especially small ones. Code is available at https://github.com/andreimargeloiu/TabEBM.

• The paper introduces a novel class-specific energy-based augmentation strategy that improves classifier accuracy under scarce data conditions.
• It leverages distinct binary surrogate tasks with SGLD sampling to generate synthetic samples that preserve true class distributions and mitigate mode collapse.
• Experimental evidence shows TabEBM outperforms standard generative methods, enhancing privacy and efficiency in imbalanced tabular datasets.

TabEBM: Data Augmentation via Distinct Class-Specific Energy-Based Models for Tabular Data

Introduction and Motivation

Robust generative modeling for tabular data remains an open challenge, particularly under regimes of limited data availability, class imbalance, and high class cardinality. Traditional generative approaches, such as VAEs, GANs, normalizing flows, and class-conditional generators, often fail to improve downstream task performance, especially in the low-sample setting; in some cases, their use even degrades the accuracy of classifiers as compared to models trained solely on real data. The core difficulties stem from overfitting to small data, inability to preserve the original label distributions, and collapse of conditional diversity—especially prominent when only a single shared model is used to represent all conditional distributions.

The paper "TabEBM: A Tabular Data Augmentation Method with Distinct Class-Specific Energy-Based Models" (2409.16118) addresses these issues by introducing an energy-based augmentation regime that learns separate, class-specific energy-based models (EBMs) for each class label, allowing generation of synthetic tabular samples that are distributionally aligned with each class's empirical distribution. This design mitigates mode collapse, improves class stratification, and is compatible with strong pre-trained tabular classifiers such as TabPFN, which enables training-free deployment.

Figure 1: TabEBM architecture—distinct class-specific energy models $\{E_c(x)\}$ are estimated for each class using a surrogate binary classification task.

Methodology

Distinct Class-Specific Energy Model Construction

Let $\mathcal{X}$ denote the dataset and $\mathcal{X}_c$ the set of samples labeled with class $c\in\{1,\dots,C\}$. For each class, a binary surrogate task is constructed: samples from $\mathcal{X}_c$ are treated as positives, while a set of negative surrogate samples $\mathcal{X}_c^{\text{neg}}$ is generated by sampling from a distant hypercube surrounding the data manifold. This strategy ensures that negatives are well-separated from actual data and avoids compromising the logit reliability of strong classifiers (e.g., TabPFN).

Each binary classifier $f^c_\theta$ is trained to discriminate $\mathcal{X}_c$ from $\mathcal{X}_c^{\text{neg}}$. The logits are then reinterpreted as energies:

$$E_c(x) = -\log \left( \exp(f^c(x)[0]) + \exp(f^c(x)[1]) \right)$$

yielding a distinct energy landscape per class. SGLD-based sampling is subsequently used to draw synthetic samples for each class proportional to the exponential negative energy, preserving class-stratified conditional densities.

Implementation: Leveraging Training-Free In-Context Models

A practical advantage is the use of pre-trained in-context classifiers (TabPFN), which are not updated for the surrogate tasks but used in an inference-only regime. This enables efficient and scalable construction of class-specific EBMs with no conventional training overhead.

Experimental Evidence

Data Augmentation Efficacy Across Regimes

TabEBM is evaluated against eight major tabular generative techniques (SMOTE, TVAE, CTGAN, NFLOW, TabDDPM, ARF, GOGGLE, TabPFGen) and a no-augmentation baseline, across OpenML and UCI datasets that span 2–26 classes per dataset, 7–77 features, and severe class imbalance scenarios.

Numerical results: TabEBM exhibits the following properties:

• Consistently improves downstream balanced accuracy over both baseline and all other generators for datasets with as few as 20 real samples up to 500 samples per class. For example, TabEBM's mean balanced accuracy improvements are positive and maximize improvement especially in low-data and high-class-count regimes.
• Other modern generators (including TabPFGen and deep generative models) regularly reduce classifier accuracy when used for data augmentation (-indicated by negative or zero accuracy improvement), especially as class count or imbalance increases.

Figure 2: TabEBM yields positive mean normalized accuracy improvement across both varying sample sizes (left) and number of classes (right); other methods often worsen performance, especially as class count increases.

TabEBM maintains strong performance on highly imbalanced datasets, consistently outperforming alternatives even as the ratio between majority and minority classes increases by orders of magnitude.

Figure 3: TabEBM maintains balanced accuracy improvement across levels of data imbalance, whereas other generators' improvements collapse.

Distributional Fidelity and Privacy

Inverse KL divergence and KS/Chi-squared test statistics between real and synthetic data distributions are used to benchmark fidelity; TabEBM achieves the highest distributional similarity with real data as compared to all baselines, both for real train and extrapolated test distributions.

For privacy-preserving data sharing (where classifiers are trained solely on synthetic data), TabEBM offers the best trade-off between high downstream accuracy and high DCR/$\mathcal{X}$0-presence scores (indicating low privacy leakage/risk for real data)—a regime where all other generators yield much lower accuracy and/or higher privacy risk.

Figure 4: TabEBM achieves superior distributional fidelity (inverse KL/KS) and privacy (DCR/$\mathcal{X}$1-presence) compared to benchmark generators while maintaining or exceeding baseline classification performance.

Robustness to Noise, Class Imbalance, and Negative Sample Placement

TabEBM's use of class-specific models provides robustness under scenarios with high label overlap (noise) and heavy class imbalance, where shared-model approaches such as TabPFGen exhibit distributional failures—e.g., disastrous flipping of probability mass assignment and mode collapse.

Figure 5: Class-conditional densities under increasing noise—TabEBM reliably concentrates probability mass on true data modes, unlike TabPFGen which collapses under high noise.

Figure 6: Under increasing class imbalance, TabEBM continues to generate mode-aligned densities, whereas TabPFGen generates spurious probability islands.

Computational Efficiency and Practical Applicability

TabEBM is computationally efficient: leveraging in-context inference, generation is up to 30× faster than most other benchmark methods, with no training required and robust convergence across a wide range of SGLD and construction hyperparameters.

Practical and Theoretical Implications

Practically, TabEBM provides an augmentation tool for tabular domains (e.g., medicine, chemistry, and physics) where data scarcity and privacy are paramount. The model's ability to realize class-specific augmentation that preserves and reflects true data diversity, even with high class cardinality and imbalance, is particularly important for real-world scenarios such as rare disease detection or multi-class chemical property prediction.

Theoretically, the approach foregrounds the limitations of sharing a single conditional generative model for all classes—highlighting how overfitting, density degeneracy, and inability to maintain proper stratification can confound downstream learning. It further demonstrates the utility of using classifier logits for robust, practical energy-based generative modeling.

TabEBM is broadly extensible. While realized here using TabPFN, the formulation is generic—any classifier producing logits can, in principle, be adapted to serve as a class-specific energy estimator, suggesting a promising cross-pollination between discriminative model developments (e.g., large tabular transformers) and generative augmentation for tabular data.

Limitations and Future Directions

TabEBM inherits some scaling limitations from the underlying classifier—TabPFN is optimized for small to moderate feature counts, though the class-wise decoupling alleviates some bottlenecks. The current approach models univariate statistics very well (as evidenced by KS and KL), but is not explicitly designed to capture arbitrary high-order relationships; as such, future work on causality-aware and multivariate-correlated tabular synthesis [see, e.g., (Tu et al., 2024)] is warranted.

Integrating emerging tabular foundation models could further expand applicability to tasks with higher $\mathcal{X}$2 and $\mathcal{X}$3.

Conclusion

TabEBM establishes a new design paradigm for tabular data augmentation, centering on the explicit modeling of per-class distributions via class-specific energy-based models and leveraging powerful, training-free in-context classifiers as surrogate density approximators. Across extensive empirical benchmarks, TabEBM consistently improves classification accuracy, demonstrates high-fidelity synthesis, strong privacy preservation, and efficiency, offering a methodologically sound and practically viable approach for data-scarce tabular domains. The open-source implementation (2409.16118) invites further research, especially as scalable tabular foundation models mature.