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

arXiv: https://arxiv.org/abs/2409.16118 · HTML: https://mlsanity.com/paper/2409.16118/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 88 references that are in the map.
- First posted on arXiv 24 Sep 2024.

## Builds on (most cited first, 30 of 46)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [Modeling Tabular data using Conditional GAN](https://mlsanity.com/paper/1907.00503/) (2019)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [Deep Neural Networks and Tabular Data: A Survey](https://mlsanity.com/paper/2110.01889/) (2021)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One](https://mlsanity.com/paper/1912.03263/) (2019)
- [TabDDPM: Modelling Tabular Data with Diffusion Models](https://mlsanity.com/paper/2209.15421/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Language Models are Realistic Tabular Data Generators](https://mlsanity.com/paper/2210.06280/) (2022)
- [Data Augmentation Generative Adversarial Networks](https://mlsanity.com/paper/1711.04340/) (2017)
- [Data Contamination: From Memorization to Exploitation](https://mlsanity.com/paper/2203.08242/) (2022)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [Imbalanced-learn: A Python Toolbox to Tackle the Curse of Imbalanced Datasets in Machine Learning](https://mlsanity.com/paper/1609.06570/) (2016)
- [Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space](https://mlsanity.com/paper/2310.09656/) (2023)
- [Data Synthesis based on Generative Adversarial Networks](https://mlsanity.com/paper/1806.03384/) (2018)
- [OpenML Benchmarking Suites](https://mlsanity.com/paper/1708.03731/) (2017)
- [Large Language Models(LLMs) on Tabular Data: Prediction, Generation, and Understanding -- A Survey](https://mlsanity.com/paper/2402.17944/) (2024)
- [Model Evaluation, Model Selection, and Algorithm Selection in Machine Learning](https://mlsanity.com/paper/1811.12808/) (2018)
- [CTAB-GAN: Effective Table Data Synthesizing](https://mlsanity.com/paper/2102.08369/) (2021)
- [Investigating Data Contamination in Modern Benchmarks for Large Language Models](https://mlsanity.com/paper/2311.09783/) (2023)
- [Why Tabular Foundation Models Should Be a Research Priority](https://mlsanity.com/paper/2405.01147/) (2024)
- [ForecastPFN: Synthetically-Trained Zero-Shot Forecasting](https://mlsanity.com/paper/2311.01933/) (2023)
- [Investigating Data Contamination for Pre-training Language Models](https://mlsanity.com/paper/2401.06059/) (2024)
- [STaSy: Score-based Tabular data Synthesis](https://mlsanity.com/paper/2210.04018/) (2022)
- [CoDi: Co-evolving Contrastive Diffusion Models for Mixed-type Tabular Synthesis](https://mlsanity.com/paper/2304.12654/) (2023)
- [Synthcity: facilitating innovative use cases of synthetic data in different data modalities](https://mlsanity.com/paper/2301.07573/) (2023)

## Built on by (most cited first, 16 of 16)

- [LLM-TabLogic: Preserving Inter-Column Logical Relationships in Synthetic Tabular Data via Prompt-Guided Latent Diffusion](https://mlsanity.com/paper/2503.02161/) (2025)
- [TACTIC for Navigating the Unknown: Tabular Anomaly deteCTion via In-Context inference](https://mlsanity.com/paper/2603.14171/) (2026)
- [Evaluating Inter-Column Logical Relationships in Synthetic Tabular Data Generation](https://mlsanity.com/paper/2502.04055/) (2025)
- [Causal Data Augmentation for Robust Fine-Tuning of Tabular Foundation Models](https://mlsanity.com/paper/2601.04110/) (2026)
- [Harnessing Synthetic Data from Generative AI for Statistical Inference](https://mlsanity.com/paper/2603.05396/) (2026)
- [NRGBoost: Energy-Based Generative Boosted Trees](https://mlsanity.com/paper/2410.03535/) (2024)
- [How Well Does Your Tabular Generator Learn the Structure of Tabular Data?](https://mlsanity.com/paper/2503.09453/) (2025)
- [Harmonized Tabular-Image Fusion via Gradient-Aligned Alternating Learning](https://mlsanity.com/paper/2604.01579/) (2026)
- [Active Tabular Augmentation via Policy-Guided Diffusion Inpainting](https://mlsanity.com/paper/2605.10315/) (2026)
- [Guided Data Generation for Understanding Model Behavior](https://mlsanity.com/paper/2502.06658/) (2025)
- [The Prompt is Mightier than the Example](https://mlsanity.com/paper/2505.18485/) (2025)
- [Generating Logically Consistent Synthetic Supply Chain Data with LLM-Driven Knowledge Graph Reasoning](https://mlsanity.com/paper/2605.26823/) (2026)
- [Task-Conditioned Synthetic Data Generation for Improving Machine Learning Performance in Agricultural Prediction Tasks](https://mlsanity.com/paper/2607.09751/) (2026)
- [TabNSM: Neural Sparse Mixer for Tabular Regression](https://mlsanity.com/paper/2608.18026/) (2026)
- [CDMD: A Cross-Dataset Mixed-Type Diffusion Model for Tabular Data](https://mlsanity.com/paper/2609.39124/) (2026)
- [TabMosaic: Unified In-Context Diffusion for Tabular Generative Tasks](https://mlsanity.com/paper/2610.15699/) (2026)

## Abstract (arXiv)

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(x, y)  or the class-conditional distribution  p(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.